DLite、エージェントAI向けに履歴を持たない投機的デコードを提案

エージェントの文脈で難しくなる投機的デコード

投機的デコードは、小型のドラフトモデルに複数の候補トークンを先に生成させ、その候補を大きなターゲットモデルでまとめて検証する推論高速化手法である。候補の多くが受理されれば、ターゲットモデルが一つずつトークンを生成する回数を減らせるため、出力の正しさを保ちながら遅延を抑えられる。

しかし、エージェントのワークロードではコンテキストが単純に長くなるだけではない。ユーザーとの対話、ツールの結果、環境からのフィードバック、更新された計画などが何度も追加され、履歴が継続的に変化する。論文によると、このような状況では既存ドラフトモデルの受理長が最大 80% 低下する。また、ドラフト側で過去の KV キャッシュを維持し続けると、メモリとシステム管理の負担も大きくなる。

ターゲットモデルの表現を再利用する設計

DLite は、受理長の低下とキャッシュ負担の両方の原因を、増え続ける履歴に求める。そこで、ドラフトモデルが会話全体を改めて保持・理解するのではなく、ターゲットモデルが推論中に生成したコンテキスト隠れ状態を直接利用する。これらの表現をもとに局所的な投機予測を行い、履歴に対応するドラフト側の KV キャッシュを明示的には維持しない。

ここでいう「history-free」は、履歴の情報を無視するという意味ではない。履歴はターゲットモデルがすでにエンコードした表現を通じて利用される。これにより、予測の課題は増大する会話全体を毎回モデル化することから、現在のコンテキスト表現に基づいて近傍の候補を予測することへ置き換えられる。

DLite には軽量な sequential head も組み込まれている。投機ブロックでは複数位置を並列に予測できる一方、後続トークンは同じブロック内の先行トークンに依存する。逐次ヘッドは、このブロック内の因果関係を小さな追加コストで補い、並列性を保ちながら予測の依存関係を扱う役割を担う。

学習と実験結果

DLite は、ターゲットモデルの隠れ状態、局所予測モジュール、逐次ヘッドを組み合わせる異種構成である。そのため、一般的な独立ドラフトモデルの学習方法をそのまま使うのではなく、この構造に合わせた二段階の学習戦略を導入する。提供された概要では詳細な手順までは示されていないが、履歴を明示的に持たない予測と、ブロック内部の因果性を協調させることが目的とされている。

一般用途、長文脈、エージェント向けの実験で、DLite は受理長 2〜6 を達成した。さらに、既存の先進的手法である DSpark と比較して最大 70% の性能向上が報告されている。これは、履歴状態の重複を減らすことが単なるメモリ節約にとどまらず、変化するコンテキストでのドラフト予測の安定性にも関係する可能性を示す。

推論基盤への意味

DLite の重要な点は、投機的デコードで履歴をどこに保持するかを見直したことにある。エージェントの長いセッションでは、ドラフト側が対話全体の別コピーを持つ必要がない可能性がある。ターゲット側の中間表現を再利用できれば、長時間の処理でキャッシュの圧力を抑え、履歴長に左右されにくいドラフトコストにつながる。

一方、概要だけでは、異なるハードウェアやモデル規模での実際のエンドツーエンド効果までは判断できない。受理長はスループットや遅延そのものではなく、隠れ状態の転送、検証処理、バッチング、並行実行、メモリ管理も実運用の結果を左右する。それでも、増え続ける履歴を重複保存せず、ターゲットモデルがすでに計算した情報から局所予測を行うという方向性は、エージェント推論に対する明確な設計提案と言える。

出典：Hugging Face Daily Papers