トークン検証(トークンケンショウ)
投機的デコーディングにおいて、ターゲットモデル(大型LLM)がドラフトモデルの生成したトークン候補を確率的に承認または却下するプロセス。1回のフォワードパスでK個のドラフトトークンを同時に検証し、受理確率min(1, p(x)/q(x))に基づいて各トークンの採否を決定する。却下位置以降はターゲットモデルの分布から再サンプリングされる。
トークン検証とは
トークン検証(Token Verification)は、投機的デコーディングの中核となるプロセスで、ドラフトモデルが先行生成した複数トークンの候補をターゲットモデル(大型LLM)が品質チェックする工程である。ドラフトモデルがK個のトークンを逐次生成するのに対し、ターゲットモデルはこのK個を1回のフォワードパスで並列に検証できるため、「K回分の仕事を1回で処理する」効率化が実現する。
検証プロセスの詳細
ステップバイステップ
入力コンテキスト x_{1:n} に続くK個のドラフトトークン x_{n+1},...,x_{n+K} が与えられた場合。
- ターゲットモデルのフォワードパス: 入力 [x_{1:n}, x_{n+1},...,x_{n+K}] に対して1回のフォワードパスを実行。各位置 t=n+1,...,n+K+1 での確率分布 p(x|x_{<t}) を同時計算
- 位置ごとの検証: i=1,...,K の順に以下を実行
- ドラフトトークン x_{n+i} に対するターゲットモデルの確率 p(x_{n+i}) とドラフトモデルの確率 q(x_{n+i}) を比較
- 一様乱数 r_i ∼ U(0,1) を生成
- r_i ≤ min(1, p(x_{n+i})/q(x_{n+i})) なら受理、それ以外は却下
- 最初の却下位置の特定: 却下が発生した最初の位置を j とする
- 出力の確定: 位置1〜j-1のトークンは確定出力。位置jでは修正分布 max(0, p(x)-q(x)) から再サンプリング。位置j+1〜Kは破棄
- 全受理の場合: K個すべてが受理された場合、位置K+1でターゲットモデルの分布 p(x|x_{<n+K+1}) から追加の1トークンをサンプリング(ボーナストークン)
検証の計算コスト
| 操作 | 計算量 | メモリ | 備考 |
|---|---|---|---|
| ターゲットモデル フォワードパス | O(K × d² × L) | KVキャッシュ + K位置分 | Lはレイヤー数、dは隠れ次元 |
| 確率分布の計算(softmax) | O(K × V) | K × V × 4 bytes | Vは語彙サイズ |
| 受理/却下判定 | O(K) | 無視できる | 乱数生成+比較のみ |
| KVキャッシュ更新 | O(受理トークン数 × d × L) | 受理分のKV追加 | 却下分は破棄 |
1回のフォワードパスの計算量は、通常の1トークン生成のK倍程度だが、GPUの並列処理により実時間はほぼ1トークン生成と同等(K≤8の場合)。これが投機的デコーディングの高速化の源泉である。
受理率に影響する要因
モデルペアの選択
受理率はドラフトモデルとターゲットモデルの確率分布の近さに直結する。
| 要因 | 高受理率 | 低受理率 |
|---|---|---|
| モデルファミリー | 同一(Llama→Llama) | 異種(GPT-2→Llama) |
| 訓練データ | 同一 | 異なる |
| ファインチューニング | 両モデルとも同一ドメイン | 片方のみ |
| モデルサイズ比 | 1/5〜1/10 | 1/100以下 |
タスクの種類
- コード補完: 受理率85〜95%(文法制約によりトークン空間が狭い)
- 翻訳: 受理率75〜85%(ソース文による制約)
- 要約: 受理率70〜80%(入力文書による制約)
- 自由対話: 受理率55〜70%(制約が少なく分岐が多い)
- 創作/詩: 受理率40〜55%(意外性のある選択が多い)
コンテキスト位置
文の先頭(高エントロピー)では受理率が低く、文の途中〜末尾(低エントロピー)では受理率が高い傾向がある。固有名詞や数値の後の助詞・助動詞は受理率が非常に高い。
KVキャッシュのロールバック
却下時のキャッシュ管理
ドラフトトークンが位置jで却下された場合、ターゲットモデルのKVキャッシュには位置n+1〜n+K+1の情報が含まれている。しかし位置j+1以降は却下されたため無効であり、KVキャッシュを位置j-1まで巻き戻す必要がある。
実装パターン
- コピー&リストア方式: フォワードパス前にKVキャッシュのスナップショットを保存し、却下時にリストア。メモリ使用量は2倍になるが実装がシンプル
- インプレース truncation方式: KVキャッシュのテンソルサイズを直接truncate。メモリ効率が良いがGPU上でのテンソル操作が必要
- リングバッファ方式: KVキャッシュを循環バッファで管理し、書き込み位置ポインタを巻き戻すだけで済む。vLLMのPagedAttentionとの相性が良い
ドラフトモデルのKVキャッシュ
ドラフトモデル側もK個のトークン生成時にKVキャッシュが蓄積される。ターゲットモデルでの却下位置以降のキャッシュは無効化し、再サンプリングされたトークンから新たなドラフト生成を開始する。
バッチ推論での検証
異なるシーケンス長の処理
バッチ内の各リクエストは異なる受理率を持つため、ステップごとに生成されるトークン数が異なる。これによりバッチ内のシーケンス長がばらつく問題がある。
- パディング方式: 最長シーケンスに合わせてパディング。GPU利用効率が低下
- 連続バッチ方式: vLLMやTGIのcontinuous batchingと統合。完了したリクエストのスロットに新しいリクエストを投入
- 動的K方式: リクエストごとに受理率を監視し、受理率の低いリクエストはKを小さくしてバッチ内のばらつきを抑制
検証の最適化テクニック
Early Exit
ドラフトトークンの検証中に却下が発生した時点で、後続トークンの検証を打ち切る最適化。確率分布の計算(softmax)が位置ごとに独立であるため、因果マスクを利用して計算を省略できる。
Parallel Verification
複数のドラフトシーケンス(木構造の分岐)を並列検証する場合、Tree Attention(木構造注意機構)を使用してGPUカーネルを最適化。通常の因果注意マスクを木構造に拡張した専用マスクを使用する。
よくある質問(FAQ)
Q1: 検証で全トークンが却下された場合はどうなりますか? A: 最初のドラフトトークンが即座に却下された場合、ターゲットモデルの確率分布 p(x|x_{<n+1}) から1トークンを再サンプリングして出力する。これは通常の自己回帰デコーディングと同等であり、投機的デコーディングの最悪ケースは通常デコーディングと同じ速度(ドラフトモデルのオーバーヘッド分だけわずかに遅い)。
Q2: 検証のためにターゲットモデルのフォワードパスを追加で実行するのはコストではないですか? A: ターゲットモデルのフォワードパスは通常のデコーディングでも毎トークン実行される。投機的デコーディングでは、1回のフォワードパスでK+1位置の確率分布を同時計算するため、追加コストはバッチ次元の拡張分(通常5〜15%)のみ。K個のトークンが受理されれば、K回分のフォワードパスを1回に圧縮したことになり、大幅なコスト削減となる。
Q3: 検証結果をキャッシュして再利用できますか? A: 部分的に可能。受理されたトークンまでのKVキャッシュは次のステップでそのまま再利用される。確率分布(ロジット)自体のキャッシュは、同一入力に対してのみ有効だが、投機的デコーディングでは毎ステップ入力が変化するため実用的ではない。
まとめ
- トークン検証はターゲットモデルが1回のフォワードパスでK個のドラフトトークンを同時チェックするプロセス
- 受理確率 min(1, p/q) により出力品質を完全に維持
- GPU並列処理により、K個の検証コストは1トークン生成とほぼ同等
- KVキャッシュのロールバック管理が実装上の主要課題
- コード生成やテンプレート文書で受理率85〜95%、自由対話で55〜70%