LLM を組み込んだ機能を作っていると、精度より先に「遅さ」が壁になる場面があります。
プロンプトを投げてから、答えが1文字ずつ画面に流れてくるまでのあの間。
モデルを小さくする、出力を短く切り上げる、といった対処はたいてい品質とのトレードオフになってしまいます。
出力の質をまったく落とさずに、生成そのものを速くする方法があるとしたら、どうでしょうか。
投機的デコーディング（speculative decoding）は、まさにそれを狙った手法。
生成が1トークンずつしか進まない理由
そもそも Transformer の生成がここまで直列になるのには、はっきりした理由があります。
自己回帰モデルは、次の1トークンを出すたびに、それまで出力した全部を入力として読み直す構造。
つまり100トークン欲しければモデルを100回走らせるしかありません。
論文「Fast Inference from Transformers via Speculative Decoding」も、KトークンをデコードするにはモデルをK回直列に走らせることになる、という指摘から始まっています。
厄介なのは、この1回あたりのコストの正体。
Hugging Face の「Assisted Generation の解説」によれば、ボトルネックは計算そのものではなく、モデルの重みを演算コアへ読み込む部分にあるとのこと。
GPU の演算能力が余っていても、重みを毎回メモリから運んでくる帯域のほうで足を引っ張られるわけです。
1トークン分の計算のために、モデル全体の重みをまるごと1往復させる。
この余っている計算能力こそ、投機的デコーディングが突く隙になります。
生成が進むほど読み直す文脈も伸びるため、会話が長くなるほど1回あたりが重くなるという事情も重なりますね。
このあたりはKVキャッシュの仕組みとセットで眺めると、輪郭がつかみやすくなるはずです。
先読みして、まとめて答え合わせをする
投機的デコーディングの発想は、驚くほど素朴に見えます。
本命の大きいモデルの隣に、ドラフトモデルと呼ばれる小さくて速いモデルを置くだけ。
ドラフトモデルに数トークン先まで一気に書かせておいて、本命モデルはそれを1回の前向き計算でまとめて答え合わせするわけです。
ここが効く理由は、先ほどのメモリ帯域の話に戻ります。
5トークン分をまとめて検証するのも、1トークンを生成するのも、重みを読み込む回数は同じ1回。
余っていた計算能力を使って、同じ待ち時間のうちに複数トークンを確定させるのがねらいになります。
合格したトークンはそのまま採用し、外れたところで打ち切って、そこから本命モデルが1トークン出し直す。
比喩でいえば、後輩が下書きを数行先まで書いておき、先輩がまとめて赤入れするイメージ。
赤が入った行から先は捨てて、また下書きをやり直すという段取りですね。
論文では T5-XXL を対象に、標準的な T5X 実装と比べて2倍から3倍の高速化が報告されています。
出力が変わらない、というのが肝
「速くなる代わりに答えが雑になるのでは」と身構えたくなるところ。
ところが投機的デコーディングは、元のモデル単体で生成したときと同じ分布を保ったまま速くすることを狙って設計されています。
論文のアブストラクトにある「出力に何の変更も加えずに」という一節が、この手法の核心。
仕掛けは棄却サンプリングと呼ばれる仕組みにあります。
ドラフトモデルが提案したトークンを、本命モデルが「自分ならその確率で出したか」で受理・棄却し、棄却したときは差分の分布から引き直す。
この補正が入るおかげで、ドラフトモデルの質が低くても答えのほうは劣化しないという性質が生まれます。
速いか遅いかだけが変わり、何が出てくるかは変わらないというわけですね。
厳密には、vLLM の Speculative Decoding ドキュメントも「ハードウェアの数値精度の限界までは理論上ロスレス」という言い方をしています。
浮動小数点の誤差ぶんのわずかな揺らぎまでは否定していない、という但し書き付き。
効かない場面と、裏で払っているコスト
とはいえ、置けば必ず速くなる魔法ではありません。
効き目を決めるのは、ドラフトの提案がどれだけ受理されるかという一点。
ドラフトが外れると、先読みした分の計算がまるごと無駄になるという損が出ます。
先読みを長くするほど当たったときの伸びは大きい反面、途中で1つ外れればそこから先は全部捨てるので、伸ばしすぎれば逆効果。
vLLM では、一度に何トークン先読みするかを num_speculative_tokens で指定します。
数値を上げれば当たったときの伸びは大きくなるものの、受理率のほうは下がっていくという綱引きになるわけです。
コスト面も忘れてはいけません。
ドラフトモデルは本命とは別にメモリを占有し、運用上も面倒をみる対象が2つに増えることになります。
限られた VRAM を本命モデルの文脈長に回すのか、ドラフトモデルに割くのか、という判断も避けられません。
モデル自体を軽くする方向の工夫とどう組み合わせるかは、量子化の話とも地続きです。
vLLM のドキュメントには、投機的デコーディングはまだ最適化されておらず、すべてのプロンプトデータセットやサンプリング設定でトークン間レイテンシの改善が得られるわけではない、という注意書きも置かれています。
速くなる前提で設計に組み込む前に、まず自分のワークロードで測る。
実装で渡すものは意外と少ない
使う側から見ると、設定は拍子抜けするほど簡単に見えます。
Transformers の Assisted Generation なら、生成時に assistant_model としてドラフト側のモデルを渡すだけ。
概念を示すと、以下のようなイメージです。
# 略（本命モデルと小さいドラフトモデルは読み込み済みとする）
outputs = model.generate(
**inputs,
assistant_model=assistant_model, # 先読み役の小さいモデル
)
# 略（デコードして利用）ドラフトモデルを別に用意するのが面倒な場面に向けて、プロンプト中の n-gram を探して先読みに流用する方式も用意されています。
vLLM が並べている選択肢は、ドラフトモデル方式のほかに n-gram マッチング、MLP スペキュレータ、EAGLE。
どれも狙いは同じで、本命モデルを1回動かす間に確定するトークンを増やすという一点に尽きます。
逆にいえば、受理率さえ高く保てればモデルもプロンプトも変えずに速度だけが上がるという、なかなか珍しい種類の最適化。
最後に
投機的デコーディングは、推論の詰まりが計算ではなくメモリの往復にあるという事実をうまく突いた工夫。
品質を人質に取らずに速くできる手段は、そう多くありません。
一方で、受理率が低ければただの無駄打ちに終わる点と、ドラフトモデルぶんのリソースを別に抱える点は、導入前に見ておきたいところです。
自分のワークロードで先読みトークン数を振ってみて、受理率とレイテンシがどう動くかを測るところから。
以上です。
コメントを残す