AI・機械学習
上級

トークン検証(トークンケンショウ)

投機的デコーディングにおいて、ターゲットモデル(大型LLM)がドラフトモデルの生成したトークン候補を確率的に承認または却下するプロセス。1回のフォワードパスでK個のドラフトトークンを同時に検証し、受理確率min(1, p(x)/q(x))に基づいて各トークンの採否を決定する。却下位置以降はターゲットモデルの分布から再サンプリングされる。

0 回閲覧
0 いいね

トークン検証とは

トークン検証(Token Verification)は、投機的デコーディングの中核となるプロセスで、ドラフトモデルが先行生成した複数トークンの候補をターゲットモデル(大型LLM)が品質チェックする工程である。ドラフトモデルがK個のトークンを逐次生成するのに対し、ターゲットモデルはこのK個を1回のフォワードパスで並列に検証できるため、「K回分の仕事を1回で処理する」効率化が実現する。

検証プロセスの詳細

ステップバイステップ

入力コンテキスト x_{1:n} に続くK個のドラフトトークン x_{n+1},...,x_{n+K} が与えられた場合。

  1. ターゲットモデルのフォワードパス: 入力 [x_{1:n}, x_{n+1},...,x_{n+K}] に対して1回のフォワードパスを実行。各位置 t=n+1,...,n+K+1 での確率分布 p(x|x_{<t}) を同時計算
  2. 位置ごとの検証: 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})) なら受理、それ以外は却下
  3. 最初の却下位置の特定: 却下が発生した最初の位置を j とする
  4. 出力の確定: 位置1〜j-1のトークンは確定出力。位置jでは修正分布 max(0, p(x)-q(x)) から再サンプリング。位置j+1〜Kは破棄
  5. 全受理の場合: 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 bytesVは語彙サイズ
受理/却下判定O(K)無視できる乱数生成+比較のみ
KVキャッシュ更新O(受理トークン数 × d × L)受理分のKV追加却下分は破棄

1回のフォワードパスの計算量は、通常の1トークン生成のK倍程度だが、GPUの並列処理により実時間はほぼ1トークン生成と同等(K≤8の場合)。これが投機的デコーディングの高速化の源泉である。

受理率に影響する要因

モデルペアの選択

受理率はドラフトモデルとターゲットモデルの確率分布の近さに直結する。

要因高受理率低受理率
モデルファミリー同一(Llama→Llama)異種(GPT-2→Llama)
訓練データ同一異なる
ファインチューニング両モデルとも同一ドメイン片方のみ
モデルサイズ比1/5〜1/101/100以下

タスクの種類

  • コード補完: 受理率85〜95%(文法制約によりトークン空間が狭い)
  • 翻訳: 受理率75〜85%(ソース文による制約)
  • 要約: 受理率70〜80%(入力文書による制約)
  • 自由対話: 受理率55〜70%(制約が少なく分岐が多い)
  • 創作/詩: 受理率40〜55%(意外性のある選択が多い)

コンテキスト位置

文の先頭(高エントロピー)では受理率が低く、文の途中〜末尾(低エントロピー)では受理率が高い傾向がある。固有名詞や数値の後の助詞・助動詞は受理率が非常に高い。

KVキャッシュのロールバック

却下時のキャッシュ管理

ドラフトトークンが位置jで却下された場合、ターゲットモデルのKVキャッシュには位置n+1〜n+K+1の情報が含まれている。しかし位置j+1以降は却下されたため無効であり、KVキャッシュを位置j-1まで巻き戻す必要がある。

実装パターン

  1. コピー&リストア方式: フォワードパス前にKVキャッシュのスナップショットを保存し、却下時にリストア。メモリ使用量は2倍になるが実装がシンプル
  2. インプレース truncation方式: KVキャッシュのテンソルサイズを直接truncate。メモリ効率が良いがGPU上でのテンソル操作が必要
  3. リングバッファ方式: 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%