Speculative Decodingとは？小さなモデルの予測をLLMがまとめて検証する

LLMは通常、1 tokenを生成してから次のtokenへ進む。100 tokenの回答なら、大きなmodelを少なくとも100回順番に動かす。この直列性が、modelをGPUへ載せてもtoken間の待ち時間が残る理由の一つである。

Speculative Decodingは、小さく速いdraft modelに複数tokenを先読みさせ、大きなtarget modelが候補をまとめて検証する。候補が多く受理されれば、target modelの1回の処理で複数tokenを確定できる。

ただし、小さいmodelを追加すれば自動的に速くなるわけではない。draftの予測が外れる、検証batchが重い、同時requestが多い、二つのmodelでmemoryを使うといった条件では遅くなることもある。

draftとtargetの二段階で進む

基本的な流れは次の通りである。

draft modelが現在のprefixから数tokenを逐次生成する
target modelが候補token列を一度に評価する
受理条件を満たす先頭からの候補を確定する
最初に棄却された位置はtarget modelの分布から補正する
確定したtoken列をprefixへ加え、次のroundへ進む

たとえばdraftが4 tokenを提案し、最初の3 tokenが受理されれば、target modelの1回の検証で3 token以上進める。1 token目で外れれば、draft計算の多くが無駄になる。

重要なのは、targetが候補文を後から文章として採点する仕組みではないことである。各位置のtoken分布を使い、定められた受理と補正を行うdecoding algorithmである。

正しく実装すればtarget modelの分布を変えない

ICML 2023の原論文は、確率的samplingでもtarget modelと同じ出力分布を保つ受理・補正方法を示した。draft modelの出力をそのまま採用するのではなく、draftとtargetの確率を比較し、棄却時には補正分布からsamplingする。

greedy decodingでは、targetが選ぶtokenとdraft候補が一致するかを順に確認できる。samplingでは単純な一致判定だけでは元の分布を保てない。

したがって「Speculative Decodingなら品質が少し下がる代わりに速くなる」という説明は常に正しくない。exactなalgorithmと対応samplingを使えば、理論上はtarget単独と同じ分布を得られる。一方、近似的な受理法、独自threshold、異なるrandom numberの扱い、kernelの数値差を導入すれば、同一性を別途検証する必要がある。

速さを決めるのはacceptance rateだけではない

一つのroundで受理された候補token数が多いほど、target modelの呼び出し回数を減らせる。しかし、acceptance rateだけで高速化は決まらない。

draft modelが候補を作る時間
target modelが候補列を検証する時間
受理・補正とKV Cache更新の時間
1 roundで提案するtoken数
targetとdraftの予測の近さ
GPUのmemory帯域と計算余力
同時request数とbatch構成

draftを大きくすればtargetへ近い候補を出しやすいが、候補生成が遅くなる。小さすぎれば速いが棄却が増える。提案token数も、少なすぎればtarget呼び出しを十分減らせず、多すぎれば外れた後ろの候補を無駄に計算する。

見るべき値は、1回のtarget検証で平均何token確定したかを表すaccepted tokens per verificationである。これをround全体のlatencyと組み合わせて判断する。

taskによって当たりやすさが変わる

draftとtargetが次tokenを同じように予測できるtaskでは効果が出やすい。

定型句が多い文章
syntaxの制約が強いcode
入力文から連続部分を写す要約や変換
JSONなど構造の決まった出力

反対に、専門的な推論、固有名詞の多い回答、温度の高いsampling、draftが学習していないdomainでは候補が外れやすい。日本語での一致率が英語と同じとも限らない。

prompt全体の平均でなく、言語、task、temperature、出力長ごとにacceptanceを分ける。同じrequestでも定型的な前半と難しい後半で変化する。

tokenizerとvocabularyを揃える

基本的な二model方式では、draftとtargetが同じtokenizerを使う方が扱いやすい。同じ文字列でもtoken境界が違えば、draftの候補IDをtargetへそのまま渡せない。

Hugging Face Transformersには異なるtokenizer間で文字列を再符号化し、token列を対応させるUniversal Assisted Decodingもある。ただし、decode、re-encode、alignmentの処理が増える。まず同じtokenizerとvocabularyを共有するmodel pairから測る方が構成は単純である。

special token、chat template、added token、normalizationもversion管理する。tokenizer名が同じでもrevision差があれば互換とは限らない。

二つのmodelと二組のCacheが必要になる

draft model方式では、targetに加えてdraftのweightを載せる。両方が過去tokenを参照するため、原則としてそれぞれのKV Cacheも必要になる。

draftが小さくても、GPU memoryの余裕が減ればtarget側のKV Cache capacityが小さくなり、同時実行数を下げる可能性がある。単一requestのITLが改善しても、fleet全体のthroughputが落ちることがある。

選択肢には、draftを別GPUへ置く、同じmodelの途中layerから候補を出すself-speculative方式、prompt内の一致を使うn-gram方式、複数tokenを予測する追加headを使う方式がある。それぞれmemory、学習、通信、対応modelの条件が異なる。

Continuous Batchingとの相性を測る

target単独のDecodeは、多数requestの次tokenをContinuous Batchingでまとめられる。すでに大きなbatchでGPUを効率よく使っている場合、候補列の検証を追加しても単一requestほどのspeedupが出ないことがある。

requestごとに受理数が違えば、進むtoken数がroundごとにばらつく。schedulerはdraft生成、target検証、通常Decodeへ戻るrequestを扱う必要がある。長い候補を持つrequestがbatchのtoken budgetを占有すれば、ほかのrequestのITLへ影響する。

低concurrencyと高concurrencyを分け、ITLだけでなくtotal token throughput、waiting time、公平性を比較する。

KV Cacheを棄却位置へ戻す

targetは候補列全体を検証するため、棄却された後ろのtokenについても一時的なKVを計算する。実際に確定した位置より後ろは次roundへ持ち越してはならない。

PagedAttentionのblock table、slot mapping、Prefix Cacheの参照数と連携し、受理されたtokenだけを有効にする。block境界をまたぐ棄却、request cancel、EOS、stop sequenceでCache leakや誤参照が起きないかを確認する。

draft側も確定prefixへ状態を合わせる必要がある。受理数のずれは、文章品質ではなくmemory corruptionや不正なtoken列につながる実装上の問題である。

監視するmetric

draftが提案したtoken数
targetが受理したtoken数とacceptance rate
target検証1回当たりの確定token数
draft生成、target検証、受理・補正の時間
target modelのforward回数
request当たりのdraft waste token数
TTFT、ITL、E2E latency
output token throughputとrequest throughput
draft・targetそれぞれのGPU memoryとKV Cache使用率
speculationをskipして通常Decodeへ戻した回数
task、言語、temperature、context長ごとの効果

speedupは、Speculative Decodingありとtarget単独の同じ出力条件で測る。速かったrequestだけを平均しない。

負荷試験

greedyとsamplingを分ける
temperatureとtop-pを段階的に変える
日本語、英語、code、JSON、要約、推論taskを混ぜる
提案token数を変え、受理数とround latencyを比較する
低concurrencyとGPU飽和時を比較する
短い回答と長い回答を混ぜる
block境界付近で候補を棄却させる
EOS、stop sequence、cancelを候補列の途中で発生させる
draft停止時にtarget単独へ戻れるか確認する
modelとtokenizerのversion不一致を起動時に拒否する

samplingでは、十分な件数を生成し、target単独と出力分布が一致するか統計的に確認する。greedyでは同じ入力と設定でtoken列の一致を確認しやすい。

目的はtarget modelの仕事を省くことではない

Speculative Decodingでは、target modelが候補を必ず検証する。小さいmodelの回答を信頼して品質を妥協する方法ではなく、targetの複数位置の計算をまとめる方法である。

効果が出るのは、draftが十分速く、targetとよく一致し、まとめた検証が逐次Decodeより安い場合である。acceptance rate、memory、batch、KV Cache、task別のlatencyを測り、効果のないrequestでは通常Decodeへ戻す設計が必要である。