Speculative decoding：あらゆるLLMで推論速度を3倍にする方法

メモリのボトルネック

ユーザーがチャットボットにメッセージを送信するたびに、モデルはその応答をトークン単位、つまりテキストの断片ごとに生成します。この仕組みは、モデルが巨大化するにつれて障害となり、特にプログラミングアシスタント、カスタマーサービス、リアルタイムエージェントなど、ほぼ即時の応答を要求される場面で顕著です。

主な制約は計算能力ではありません。NVIDIA H100のようなGPUは、行列計算において非常に高速です。問題は、モデルの数十億のパラメータがメモリに格納されており、各トークンを生成する前にプロセッサに転送する必要がある点です。その間、計算ユニットはほとんどアイドル状態となります。エンジニアはこれをメモリ帯域幅制約システムと呼び、計算能力を追加しても速度が大幅に向上しないことを意味します。

提案し、検証する

この原理は、複数の候補トークンを一度に生成し、並列に検証することにあります。システムは2つのモデルに依存します。小さなモデル、いわゆるドラフトモデル（draft model）が、通常4〜8トークンの候補シーケンスを迅速に生成します。より大きく高品質なターゲットモデル（target model）が、そのセットを一括で評価します。GPUは並列計算用に設計されているため、複数のトークンをバッチ処理する時間は、単一トークンを処理する時間とほとんど変わりません。

各候補について、ターゲットモデルは自身の確率とドラフトモデルが割り当てた確率を比較します。ターゲットモデルの確信度が同等以上であれば、トークンは受理されます。そうでない場合、その差が大きくなるにつれて確率は減少し、受理される可能性が低くなります。拒否されたトークンは、調整された分布から抽出されたトークンに置き換えられ、ループが再開されます。この手法の説明によれば、この統計的プロセスにより、出力はターゲットモデルが単独で各トークンを生成した場合と全く同じ確率分布に従うことが保証されます。

最良の場合、すべての候補が受理され、ターゲットモデルが追加のトークンを生成するため、ターゲットモデルの1回の処理コストで5〜9トークンが得られます。最悪の場合、最初の候補が拒否され、ドラフト生成の時間が依然としてかかるため、わずかなオーバーヘッドが発生します。したがって、不適切なモデルの組み合わせでは、生成がわずかに遅くなる可能性があります。

実際の利得は主に受理率、つまりドラフトの候補が検証を通過する頻度に依存します。一般的なタスクで適切に組み合わされたペアでは、70〜80％が典型的とされ、このレベルで加速効果は2倍から3倍になります。予測可能または反復的なテキストでは、より高い受理率と大きな利得が得られます。

バリエーションと実践

元のスキームでは別個のドラフトモデルが必要であり、メモリとエンジニアリングの複雑さが増します。より新しいいくつかの技術ではこれを不要にしています。一部は厳密な統計的保証を放棄し、よりシンプルな実装や高いスループットを実現しており、情報源によれば、出力品質への影響は一般的に無視できるレベルです。

Medusaは、ターゲットモデルの最終層に複数の追加予測ヘッドを追加し、それぞれが異なる将来の位置を予測します。Eagle（Extrapolation Algorithm for Greater Language-model Efficiency）は、出力トークンではなくモデルの内部状態に依存し、情報源によれば、複雑または専門的なテキストにおいてより高い受理率を促進します。Prompt Lookup Decodingは補助モデルを使用せず、プロンプトと会話履歴内で応答に再出現しそうな表現を特定し、候補として使用します。これは特にコード編集や文書要約に適しています。

Speculative decodingは、Hugging Faceのtext-generation-inference、vLLM、llama.cppなど、いくつかの推論フレームワークでオプションとして提供されています。情報源によれば、有効化はほとんどの場合、設定の変更のみで可能です。

この技術は、モデルが一度に1つのクエリを処理し、プロンプトと応答が語彙や構造を共有する場合の、長い生成シーケンスで特に効果を発揮します。質疑応答、文書要約、コード補完などは、利得が最も大きいユースケースです。一方、高スループットのバッチ処理サービスでは、ドラフトトークンがメモリと計算リソースを消費し、他のクエリに使用できる可能性があるため、スループットを低下させる可能性があります。