AI・機械学習
上級

投機的サンプリング(トウキテキサンプリング)

DeepMindが2023年に発表した投機的デコーディングの数学的定式化。ドラフトモデルの確率分布q(x)から生成されたトークンを、ターゲットモデルの確率分布p(x)との比率min(1, p(x)/q(x))で確率的に受理・却下する。却下時は修正分布max(0, p(x)-q(x))から再サンプリングし、最終出力がターゲットモデル単体と同一分布になることを数学的に保証する。

0 回閲覧
0 いいね

投機的サンプリングとは

投機的サンプリング(Speculative Sampling)は、DeepMindのChen et al.が2023年の論文「Accelerating Large Language Model Decoding with Speculative Sampling」で発表した、投機的デコーディングの数学的に厳密な定式化である。同時期にGoogle Researchが発表した「Speculative Decoding」と本質的に同一のアルゴリズムだが、DeepMind側の論文は受理確率の導出と出力分布の同一性証明に重点を置いており、理論的な基盤として広く参照されている。

数学的定式化

記法の定義

  • p(x|x_{<t}): ターゲットモデルの条件付き確率分布(位置tでのトークンxの確率)
  • q(x|x_{<t}): ドラフトモデルの条件付き確率分布
  • K: 投機トークン数(ドラフトモデルが先行生成するトークン数)
  • α: 平均受理率(受理されるドラフトトークンの割合の期待値)

受理/却下アルゴリズム

各ドラフトトークン x_i(i=1,...,K)に対して以下の判定を行う。

  1. 一様乱数 r ∼ Uniform(0, 1) を生成
  2. r ≤ min(1, p(x_i) / q(x_i)) なら受理
  3. r > min(1, p(x_i) / q(x_i)) なら却下。修正分布 p'(x) = normalize(max(0, p(x) - q(x))) から再サンプリング
  4. 却下されたトークンの後続ドラフトはすべて破棄

出力分布の同一性証明

投機的サンプリングの核心的な定理は、上記のアルゴリズムにより生成されたトークンの分布がターゲットモデル単体のサンプリングと完全に一致することである。

ケース確率出力分布
ドラフトが受理min(1, p(x)/q(x)) × q(x) = min(q(x), p(x))p(x)の一部
ドラフトが却下 + 再サンプリング(1 - Σmin(q(x), p(x))) × p'(x)p(x)の残り
合計—p(x)(完全一致)

この証明は、受理部分 min(q(x), p(x)) と却下後の修正分布 max(0, p(x)-q(x)) の和が p(x) に一致することから導かれる。

温度付きサンプリングとの互換性

Greedy Decoding(temperature=0)

temperature=0の場合、ターゲットモデルの最大確率トークン argmax p(x) とドラフトモデルの最大確率トークン argmax q(x) が一致すれば必ず受理される。不一致の場合は即座に却下され、ターゲットモデルのargmaxトークンが出力される。出力はターゲットモデル単体のgreedy decodingと完全一致する。

Top-k / Top-p サンプリング

top-kやtop-p(nucleus sampling)との組み合わせも可能。この場合、ターゲットモデルの確率分布を先にtop-k/top-pでフィルタリングし、正規化した分布 p'(x) を使用して受理/却下判定を行う。ドラフトモデル側も同一のフィルタリングパラメータを適用する。

温度スケーリングの注意点

temperature受理率への影響備考
0.0最高(deterministic一致時100%)greedy decoding
0.1〜0.5高い(80〜95%)ほぼ確定的な出力
0.7〜1.0中程度(60〜80%)標準的な生成設定
1.0〜1.5やや低い(50〜70%)創造的な生成
2.0+低い(30〜50%)非常にランダムな出力

温度が高いほど確率分布が平坦になり、ドラフトモデルとターゲットモデルの分布が乖離しやすくなるため、受理率が低下する。

実装上の考慮点

並列検証の効率化

K個のドラフトトークンの検証は、ターゲットモデルの1回のフォワードパスで並列実行できる。これは因果的注意機構(Causal Attention)の性質を利用している。位置tの検証に必要なのはx_{<t}の情報のみであり、ドラフトトークン列全体を入力すれば各位置の確率分布を同時に計算できる。

KVキャッシュの管理

投機的デコーディングではKVキャッシュの管理が複雑になる。

  1. ドラフトモデルのKVキャッシュ: K個のトークン生成時に蓄積。却下時は該当位置以降を破棄
  2. ターゲットモデルのKVキャッシュ: 検証フェーズで受理されたトークンまでを保持。却下位置以降は無効化
  3. ロールバック処理: 却下されたトークン以降のKVキャッシュを巻き戻す処理が必要

メモリ使用量の見積もり

投機的サンプリングのメモリオーバーヘッドは以下の通り。

  • ドラフトモデルのパラメータ: モデルサイズに比例(例: 8B = 約16GB in FP16)
  • ドラフトモデルのKVキャッシュ: 小型のため通常1GB未満
  • 追加バッファ(確率分布の保存): K × vocab_size × 4 bytes(例: K=5, vocab=128K → 約2.5MB)

Google Researchとの比較

項目DeepMind (Speculative Sampling)Google Research (Speculative Decoding)
論文発表2023年2月(arXiv)2023年2月(arXiv)
受理確率min(1, p/q)同一
修正分布max(0, p-q) を正規化同一
証明の焦点出力分布の同一性の厳密証明期待高速化率の分析
バッチ処理未言及詳細な分析あり
実装の影響理論的基盤として引用実装フレームワークに直接採用

両論文は独立に同一のアルゴリズムに到達しており、後続の研究では区別なく「Speculative Decoding/Sampling」として参照される。

よくある質問(FAQ)

Q1: 投機的サンプリングと通常のリジェクションサンプリングの違いは何ですか? A: 通常のリジェクションサンプリングは却下時にサンプルを完全に破棄するが、投機的サンプリングは却下時に修正分布 max(0, p-q) から再サンプリングするため、各ステップで必ず1トークン以上を出力する。これにより「空振り」がなくなり、最悪ケースでも従来の自己回帰デコーディングと同等の速度を保証する。

Q2: 受理確率の計算で数値的な問題は発生しますか? A: p(x)/q(x) の計算でq(x)が非常に小さい場合にオーバーフローの可能性がある。実装上はlog空間で計算し、log p(x) - log q(x) として比較する。q(x)=0かつp(x)>0の場合は必ず却下(そのトークンはドラフトモデルが生成しないため、この場合は到達しない)。

Q3: 複数の独立したドラフトシーケンスを同時に検証できますか? A: 可能。Tree Speculative Decoding(木構造投機的デコーディング)がこのアプローチを実現している。ドラフトモデルが複数の分岐候補を生成し、木構造の注意機構でターゲットモデルが全分岐を同時検証する。受理トークン数の期待値がK本のうち最良パスの分だけ増加する。

まとめ

  • 投機的サンプリングはDeepMindが定式化した、ターゲットモデル分布を完全保存する受理/却下アルゴリズム
  • 受理確率 min(1, p/q) と修正分布 max(0, p-q) の組み合わせで出力分布の同一性を数学的に保証
  • temperature=0からtop-p sampling まであらゆるサンプリング設定と互換
  • 最悪ケースでも自己回帰デコーディングと同等の速度を保証(空振りなし)
  • KVキャッシュのロールバック管理が実装上の主要な複雑さ