Skip to main content

Speculative Decoding

Speculative Decoding は、小さい draft model が複数 token を先読みし、大きい target model がそれらをまとめて検証する LLM inference acceleration method です。生成分布を変えないように accept / reject する設計により、品質を保ちながら decode latency を下げることを狙います。

Speculative decoding flow

自作概念図。Draft model が候補 token を生成し、target model が一括で検証して、受理できる token を prefix に追加します。

なぜ速くなるのか​

Autoregressive decoding は、token を 1 個ずつ生成するため逐次依存が強いです。特に decode phase では、batch が小さい場合に GPU の compute を使い切れず、memory bandwidth や kernel launch overhead が効きます。Speculative decoding は、安い draft model で未来 token を先読みし、target model の 1 回の forward pass で複数 token を検証することで、target model の呼び出し回数を減らします。

基本 algorithm​

  1. Draft model qq が、現在の prefix から KK 個の token を生成します。
  2. Target model pp が、その KK 個の token を teacher-forcing で一括評価します。
  3. 各 token について、qq が提案した確率と pp の確率の比に基づき accept / reject します。
  4. Reject が起きた位置では、target model の補正分布から token を sample します。
  5. Accept された token は output prefix に追加されます。

典型的な accept probability は、draft token xx に対して次のように表されます。

α(x)=min⁡(1,p(x)q(x))\alpha(x) = \min\left(1, \frac{p(x)}{q(x)}\right)

この rejection sampling 的な補正により、理想化された条件では target model 単体で sampling した分布と同じ分布を保つことができます。

何が性能を決めるか​

要素影響
Draft model の速さ小さいほど安いが、精度が低すぎると reject が増えます。
Draft model の一致度Target model と分布が近いほど accept 率が上がります。
Draft 長 KK長いほど一度に進めますが、reject された後ろは無駄になります。
Batch / serving load高負荷 serving では scheduling との相互作用が大きくなります。
Sampling 設定Temperature が高いほど accept 率が下がりやすいです。

Speculative decoding の speedup は、単純な FLOPs 比だけでは決まりません。Draft model と target model を同時に走らせる device placement、KV cache 管理、batching、kernel fusion も重要です。

Variants​

Separate draft model​

最も標準的な構成です。Small LLM や distilled model が draft を担当します。Target model と tokenizer、chat template、pretraining distribution が近いほど受理率は上がります。

Medusa​

Medusa は、base model の上に複数の decoding head を追加し、複数 token 先を同時に予測する方法です。別の draft model を持たず、single model の拡張として speculative な候補を作ります。

EAGLE​

EAGLE は、token そのものではなく feature や hidden state の extrapolation を使って、より target model に近い draft を作ることを狙います。Draft の quality を高めて accept 率を上げる方向の variant です。

Lookahead Decoding​

Lookahead Decoding は、parallel decoding や n-gram 的な candidate verification を使い、draft model に依存しない形で将来 token をまとめて確定しようとする family です。

いつ有効か​

Speculative decoding は、次の条件で特に効きます。

  • Target model が大きく、draft model が十分に安い。
  • Draft model が target model の分布をよく近似している。
  • Temperature が低〜中程度で、accept 率が高い。
  • 1 request あたりの decode token が長い。
  • Serving stack が speculative decoding を native に支えている。

逆に、high temperature の creative sampling、tool call の短い応答、すでに high-throughput batching が効いている環境では、効果が小さくなることがあります。

実装上の注意​

  • Draft と target の tokenizer が一致しないと扱いが複雑になります。
  • Chat template の差は accept 率を下げます。
  • KV cache は draft と target で別に管理されることが多く、memory footprint は増えます。
  • Deterministic greedy decoding では簡略化できますが、sampling では分布補正が必要です。
  • Speculative decoding は latency 最適化 であり、model の能力を増やす方法ではありません。

数式で見る speculative decoding の正しさ​

Draft model を qq、target model を pp とします。Draft が提案した token xx は、次の確率で accept されます。

α(x)=min⁡(1,p(x)q(x))\alpha(x)=\min\left(1,\frac{p(x)}{q(x)}\right)

Reject された場合は、target model と draft model の差分から作る補正分布で sample します。

pcorr(x)=max⁡(p(x)−q(x),0)∑x′max⁡(p(x′)−q(x′),0)p_{\mathrm{corr}}(x)=\frac{\max(p(x)-q(x),0)}{\sum_{x'}\max(p(x')-q(x'),0)}

この式の気持ちは、「draft が target よりも過大評価している token は reject されやすくし、draft が見逃した target 側の確率質量は補正分布から回収する」というものです。理想的な条件では、最終的な生成分布は target model pp だけで sampling した場合と一致します。

Speedup は、draft length KK と accept 率 aa に大きく依存します。粗く見ると、一回の target forward で進む token 数の期待値は次のようになります。

E[accepted]≈∑k=1Kak\mathbb{E}[\mathrm{accepted}]\approx \sum_{k=1}^{K} a^k

Draft が target に近いほど aa が高くなり、一回の target call で多くの token を進められます。一方で、draft が遅すぎたり reject が多すぎたりすると、逆に overhead が増えます。

関連ページ​

主なソース​