Speculative Decodingとは？LLM生成を速くする方法

draft modelが候補を作りtarget modelが受理可能なprefixをまとめて検証するSpeculative Decodingの全体像

Speculative Decoding（投機的デコーディング）は、小さなdraft model（候補を先に作る補助model）と、本来使いたいtarget model（最終判断を行うmodel）を組み合わせるLLM推論の高速化手法です。draft modelが数tokenを先読みし、target modelが候補を1回のforward passでまとめて検証します。

名前に「投機的」とありますが、未確認のtokenをそのまま出力する近似手法ではありません。greedy decoding（最大確率tokenを選ぶ生成）ではtarget model単独と同じtoken列を、sampling（確率的な生成）ではtarget model単独と同じ確率分布を保つように受理・棄却します。ただし、draft生成にも時間とメモリを使うため、どの環境でも速くなるわけではありません。

draft modelが候補を作りtarget modelが受理可能なprefixをまとめて検証するSpeculative Decodingの全体像

3文要約

  1. 小さなdraft modelが候補tokenを自己回帰的に先読みし、target modelが候補位置をまとめてscoreすることで、重いtarget modelの逐次呼出し回数を減らします。
  2. greedyではtargetの最大確率tokenとの一致を確認し、samplingでは受理確率と棄却時の補正分布を使うため、target modelが定める出力を保てます。
  3. 実速度はacceptance rate（候補受理率）だけで決まらず、draft latency、検証cost、lookahead長、batch、KV Cache、memory overheadを同じ条件で測る必要があります。

通常のLLM生成はなぜ逐次処理になるのか

LLMは、これまでのtoken列\(x_{<t}\)から次のtoken \(x_t\)の確率を計算します。

\[p(x_{1:T})=\prod_{t=1}^{T}p(x_t\mid x_{<t})\]

\(x_{t+1}\)の分布は、直前に確定した\(x_t\)を入力へ加えないと決まりません。通常のdecode（prefill後に1 tokenずつ生成する段階）では、target modelを実行して1 tokenを確定し、そのtokenを入力へ足して再びtarget modelを実行します。100 tokenを生成するなら、原則として重いtarget modelを100回、順番に呼びます。

次トークン予測の解説で扱った自己回帰性そのものは変えられません。一方、Transformerは既知の複数位置を入力へ置けば、それらの位置のscoreを並列に計算できます。Speculative Decodingは、この「候補が既知ならまとめて計算できる」という性質を利用します。

通常decodeではtargetをtokenごとに呼ぶ一方、Speculative Decodingではdraft候補をtargetがまとめて検証するtimeline

draft modelとtarget modelの処理フロー

Fast Inference from Transformers via Speculative DecodingとAccelerating Large Language Model Decoding with Speculative Samplingで示されたclassicな2-model方式は、次の流れです。先読みする候補数を\(\gamma\)とします。

  1. draft modelが現在のprefixから1 tokenを生成する。
  2. 生成したtokenをdraft自身の入力へ加え、合計\(\gamma\) tokenまで繰り返す。
  3. target modelへ「元のprefix＋\(\gamma\)個の候補」を渡し、候補位置とその次の位置をまとめてscoreする。
  4. 左から順に候補を検証し、受理できるprefixだけを確定する。
  5. 最初の不一致または棄却で止め、それより後ろの候補を捨てる。
  6. 確定した新しいprefixから次のiterationを始める。

誤解しやすいのは、draft候補の生成です。classic方式のdraft modelも自己回帰modelなので、\(\gamma\) tokenを順番に生成します。並列化される中心は、target modelによる複数候補位置の検証です。draftは小さく1 stepが速いからこそ、target modelを何度も呼ぶ代わりとして成立します。

また、途中の候補だけを飛び越えて後続候補を採用することはできません。ある位置でtargetの判断と食い違うと、その後のdraft候補は異なるprefixを前提にしているためです。

greedy decodingでは一致するprefixを採用する

greedy decodingでは規則が直感的です。各位置で、draft tokenがtarget modelのargmax（確率最大のtoken）と一致するかを左から確認します。

たとえばdraftが「東京 / は / 日本 / 最大」と4 tokenを提案したとします。targetのargmaxが「東京 / は / 首都 / で」なら、最初の2 tokenは一致します。3番目で不一致になるため、draftの「日本」と後続の「最大」を捨て、targetの「首都」を採用します。

結果として、このtarget forward pass 1回で「東京 / は / 首都」の3 tokenが確定します。通常decodeならtarget modelを3回呼ぶ部分です。4候補すべてが一致した場合は、targetが同時に計算していた次位置のtokenも追加できるため、最大\(\gamma+1\) token進みます。

greedy検証で2 tokenを受理し最初の不一致でdraft suffixを捨てtarget tokenへ置き換える例

samplingでもtargetの確率分布を保つ

samplingでは、単純な「一致・不一致」だけでは判定できません。draft modelの分布を\(q\)、target modelの分布を\(p\)とします。draftがtoken \(x\)を提案したとき、次の確率で受理します。

\[a(x)=\min\left(1,\frac{p(x)}{q(x)}\right)\]

targetがdraft以上の確率を\(x\)へ割り当てていれば、\(x\)は必ず受理されます。draftがtargetより\(x\)を過大評価していれば、その比率だけ受理します。

候補を棄却したとき、単にtargetの\(p\)からもう一度sampleすると、target分布より多く選ばれるtokenが生じます。そこで、draftからの提案・受理だけでは足りなかった確率質量を表すresidual distribution（正の残差を正規化した補正分布）を使います。

\[p'(x)=\frac{\max(0,p(x)-q(x))}{\sum_y\max(0,p(y)-q(y))}\]

3 tokenの数値例

語彙をA、B、Cとし、分布を次のようにします。

token draft \(q\) target \(p\)
A 0.4 0.3
B 0.4 0.5
C 0.2 0.2

draftがAを提案した場合、受理確率は\(\min(1,0.3/0.4)=0.75\)です。Aが提案されて受理される確率は\(0.4\times0.75=0.3\)となり、targetの\(p(A)\)と一致します。

棄却時の正の残差は\((0,0.1,0)\)です。正規化するとBを選びます。Bはdraftから提案・受理される確率\(0.4\)に、Aが提案され棄却された確率\(0.1\)を加えて\(0.5\)になります。Cは\(0.2\)のままです。最終分布は\((0.3,0.5,0.2)\)となり、target分布\(p\)へ戻ります。

draft分布qからの提案を受理または棄却し残差分布でtarget分布pを復元する数値例

この性質は「target model単独と同じtoken列を毎回返す」という意味ではありません。samplingでは乱数を使うため、実装順序や乱数消費の違いで個々のtoken列が変わる場合があります。理論上保たれるのはtarget modelの確率分布です。有限精度やruntime実装まで含めたbitwise一致を保証する表現でもありません。

1回のtarget検証で何token進むか

先読み長を\(\gamma\)、各候補の受理率を単純化して一定の\(\alpha\)とします。1回のtarget検証で確定するtoken数\(T\)の期待値は、次の等比級数で表せます。

\[\mathbb{E}[T]=1+\alpha+\alpha^2+\cdots+\alpha^{\gamma}\]
\[\mathbb{E}[T]=\frac{1-\alpha^{\gamma+1}}{1-\alpha}\quad(\alpha\neq1)\]

\(\alpha=1\)なら\(\mathbb{E}[T]=\gamma+1\)です。たとえば\(\gamma=4\)、\(\alpha=0.8\)なら、次の計算になります。

\[\mathbb{E}[T]=1+0.8+0.8^2+0.8^3+0.8^4=3.3616\]

簡易model上は、1回のtarget検証で平均約3.36 tokenが確定します。ただし、実際の受理率はpromptや位置ごとに変わり、独立でもありません。この式は挙動を理解する近似です。

受理率が上がるほどtarget検証1回当たりの期待確定token数が増えるが上限はlookahead長で決まる曲線

3.36 token進んでも3.36倍速いとは限らない

Speculative Decodingの性能は、target modelの呼出し回数だけでは決まりません。通常decodeでtargetが1 token生成する時間を\(t_{target}\)、draftが1 token生成する時間を\(t_{draft}\)、複数位置をtargetが検証する時間を\(t_{verify}\)、制御処理を\(t_{overhead}\)とすると、概算は次のように整理できます。

\[\operatorname{speedup}\approx\frac{\mathbb{E}[T]\,t_{target}}{\gamma t_{draft}+t_{verify}+t_{overhead}}\]

この式は厳密な性能modelではありません。\(t_{verify}\)は単一tokenのtarget callと同じとは限らず、sequence length、kernel、batch、GPU utilizationで変わります。draft側にもKV Cacheが必要で、model切替やsamplingのoverheadもあります。

Fast Inference from Transformers via Speculative Decodingは、論文のT5-XXL実験条件で2〜3倍の高速化を報告しました。Accelerating Large Language Model Decoding with Speculative Samplingも、Chinchilla 70Bを用いた分散環境で2〜2.5倍を報告しています。これらは手法が成立する実例であり、任意のmodel、GPU、runtime、batchで得られる保証値ではありません。

target呼出し削減の利得からdraft生成、複数位置検証、制御、memoryのcostを差し引いて実速度が決まる図

draft modelとlookahead長の選び方

draft modelは「targetに近いほど良い」だけでは選べません。Decoding Speculative Decodingは350を超える実験を行い、draft latencyが性能へ大きく影響すること、一般的な言語model能力がspeculative performanceと強く相関するとは限らないことを報告しています。

選定軸 小さい・短い側 大きい・長い側 判断する指標
draft model 1 stepが速くmemoryも小さいが、targetとの一致率が下がりやすい 候補は当たりやすいが、draft自体がbottleneckになり得る draft latencyとaccepted tokens/sec
lookahead \(\gamma\) 無駄な候補が少ないが、1回で進める上限が小さい 高受理率なら多く進むが、棄却後に捨てる計算が増える \(\mathbb{E}[T]\)とiteration latency
tokenizer 同一なら候補位置を直接対応させやすい 異なると変換・位置合わせが必要 runtimeの対応方式とfallback
配置 同じdeviceならmodel間転送を避けやすい 別deviceなら容量を分散できるが通信が増える end-to-end latencyとpeak memory

先読みを長くすれば必ず得をするわけではありません。受理率が低い区間では、早い位置で棄却され、後続のdraft計算が無駄になります。固定\(\gamma\)だけでなく、実装が候補数を動的に調整できるかも確認対象です。

draft modelの小ささとlookahead長について受理率、逐次cost、無駄候補、memoryのtrade-offを整理したmatrix

Speculative Decodingで遅くなる条件

次の条件では、target callを減らす利得より追加costが大きくなります。

  • taskやdomainがdraftの得意分野から外れ、最初の方で候補が棄却される
  • draft modelが大きすぎる、またはdraft用kernelが遅い
  • 長い候補列のtarget検証がhardware上で効率化されない
  • 生成token数が短く、初期化やmodel切替のoverheadを回収できない
  • batch sizeや同時request数が大きく、通常decodeでもtarget deviceを十分活用できている
  • target weightに加えてdraft weightと両方のKV Cacheを置くため、memory pressureが高まる
  • tokenizerやruntimeの非対応で変換costやfallbackが発生する

とくにonline servingでは、単一requestのinter-token latency（token間遅延）が改善しても、system全体のthroughput（単位時間当たり処理token数）が改善するとは限りません。batching policyとschedulerを含めて比較します。

Hugging Face Transformersで試す入口

Hugging Face Transformersの現行Assisted Decodingでは、target modelのgenerate()へassistant_modelを渡すのが基本です。次はAPIの形を確認する最小例です。model IDは、同じtokenizerを利用できる互換な組合せへ置き換えてください。

import logging

import torch
from transformers import (
    AutoModelForCausalLM,
    AutoTokenizer,
    BatchEncoding,
    PreTrainedModel,
    PreTrainedTokenizerBase,
)

logging.basicConfig(level=logging.INFO)
logger: logging.Logger = logging.getLogger(__name__)

target_id: str = "your-org/target-model"
draft_id: str = "your-org/draft-model"

tokenizer: PreTrainedTokenizerBase = AutoTokenizer.from_pretrained(target_id)
target: PreTrainedModel = AutoModelForCausalLM.from_pretrained(
    target_id,
    torch_dtype=torch.bfloat16,
    device_map="auto",
)
draft: PreTrainedModel = AutoModelForCausalLM.from_pretrained(
    draft_id,
    torch_dtype=torch.bfloat16,
    device_map="auto",
)

inputs: BatchEncoding = tokenizer(
    "Speculative Decodingを一文で説明すると",
    return_tensors="pt",
).to(target.device)

logger.info("assisted decodingを開始します")
output_ids: torch.Tensor = target.generate(
    **inputs,
    assistant_model=draft,
    max_new_tokens=80,
    do_sample=False,
)
print(tokenizer.decode(output_ids[0], skip_special_tokens=True))

公式documentは、standardな方式ではtargetより十分小さく、同じtokenizerを持つassistant modelを使う場合を説明しています。greedyとsamplingを扱えますが、現行documentではbatch入力は未対応とされています。API、候補schedule、対応model、backend制約は更新されるため、利用するTransformers versionのdocumentで確認してください。

比較時は、まずassistant_modelなしのbaselineを同じprompt、生成設定、dtype、device、乱数条件で測ります。出力を表示するだけのsample codeと、同期を含む正確なGPU benchmarkは分けて用意してください。

派生方式は「候補を誰が作るか」で整理する

候補生成は、小さな別modelだけに限りません。現行のTransformers documentにも複数の方式があります。

方式 候補の作り方 追加cost・制約
2-model speculative decoding 小さなdraft modelが生成 draft weightとKV Cache、model互換性が必要
prompt lookup prompt内の一致するn-gramを再利用 入力に繰り返しが少ないと候補を得にくい
self-speculative decoding target modelの一部の層を飛ばして候補生成 対応architectureとlayer設計に依存する
Universal Assisted Decoding 異なるtokenizer間を再encodeして位置合わせ 変換とalignmentのcostが加わる
MTP系 modelに複数未来tokenの予測headを持たせる 専用に学習されたmodelやruntime対応が必要

それぞれ候補の作り方は違いますが、候補をtarget側で確認し、重い逐次処理を減らすという見方は共通しています。versionや実装によって「Speculative Decoding」「assisted generation」に含める範囲が異なるため、方式名だけで同一視しない方が安全です。

導入前に測る項目

benchmarkでは、平均latencyを1つだけ記録しないようにします。

指標 分かること 注意点
acceptance rate draft候補がtargetに認められる割合 高くてもdraftが遅ければ高速とは限らない
accepted tokens/iteration target検証1回で進むtoken数 wall-clock speedupではない
draft latency 候補生成の追加cost \(\gamma\)回の逐次生成を含める
verification latency targetの複数位置score cost sequence lengthと候補長別に測る
TTFT 最初のtokenまでの時間 主にprefillを含み、decode改善と分ける
ITL token間の待ち時間 streaming体感に近い
throughput system全体のtoken処理量 batch・concurrencyを明記する
peak memory 収容可能なmodel・request数 draft weightと両KV Cacheを含める
output equivalence greedy一致またはsampling分布 samplingの単一出力一致だけで判定しない

warm-up後に複数回測り、prompt長、生成長、batch size、同時request数、dtype、hardware、software versionを記録します。acceptanceは自然文、code、domain固有文など入力の種類でも変わるため、本番に近いprompt集合を使います。

acceptanceだけでなくdraft latency、verification latency、ITL、throughput、peak memory、出力同等性を並べるbenchmark dashboard

他のLLM高速化技術との違い

Speculative Decodingは、既存の高速化技術を置き換えるものではありません。最適化する場所が違うため、組み合わせられます。

技術 主に変えるもの 主な効果 Speculative Decodingとの関係
KV Cache 過去K/Vの再利用 過去tokenの再計算を避ける draftとtargetの両方がcacheを持ち得る
量子化 weightやactivationのbit幅 memoryと演算costを減らす draft/targetを軽量化できるが分布変化を測る
FlashAttention Attention内部の計算順序 memory accessを減らす target検証やprefillのkernelを速くし得る
GQA/MQA KV head共有 KV Cacheとdecode帯域を減らす model architecture側から両modelを軽くする
Speculative Decoding decode algorithmとtarget call頻度 1回のtarget検証で複数tokenを確定する 上記のmodel・kernel最適化と併用できる

組合せでは、個々の改善率を掛け算して見積もらないでください。たとえばtargetが量子化やGQAで十分速くなると、draftや制御の相対costが大きくなり、Speculative Decodingの上積みが変わります。

よくある質問

小さいdraft modelほど速くなりますか

必ずしもそうではありません。小さくするとdraft latencyは下がりますが、targetとの分布差が広がって受理率が下がる場合があります。候補の精度ではなく、draft生成とtarget検証を含むaccepted tokens/secで選びます。

出力品質は落ちませんか

正しく実装された標準方式なら、greedyではtargetの選択、samplingではtarget distributionを保ちます。ただし、近似的なverification、別の分布を混ぜる方式、量子化、有限精度、runtime bugまで自動的に保証するわけではありません。

target modelを再学習する必要はありますか

classicな2-model方式は、互換な既存modelを組み合わせて使えるため、targetの再学習は必須ではありません。一方、MTPや専用speculatorなど、追加学習を前提とする派生方式もあります。

tokenizerが違うmodelもdraftにできますか

standard方式は同じtokenizerが扱いやすく、多くの実装で前提になります。異なるtokenizerを位置合わせするUniversal Assisted Decodingのような方式もありますが、runtime対応と変換costを確認してください。

batch処理でも速くなりますか

単一requestでの効果をそのまま一般化できません。batchやconcurrencyが高いとtarget deviceの利用率、sequence長のばらつき、schedulerとの相互作用が変わります。利用するruntimeがbatched speculative decodingを支えるか確認し、本番負荷でthroughputを測ります。

まとめ

Speculative Decodingは、速いdraft modelで未来の候補を作り、重いtarget modelでまとめて検証することで、targetの逐次呼出しを減らします。greedyでは一致prefixを、samplingでは受理確率とresidual distributionを使うため、target modelの判断を保ったまま複数tokenを確定できます。

導入判断ではacceptance rateだけを見ず、draft latency、target verification、lookahead長、ITL、throughput、peak memoryを同じ条件で比較してください。target call 1回当たりの期待token数は仕組みを理解する指標であり、実速度は追加costを含むend-to-end benchmarkで決まります。

参考資料

スポンサーリンク

コメント

タイトルとURLをコピーしました