EAGLE 3.1：LLM推論における注意ドリフトを修正する投機的デコーディングアルゴリズム

投機的デコーディングは、大規模言語モデルの推論を高速化する手法です。小型で高速なドラフトモデルが複数のトークンを提案し、大型のターゲットモデルがそれらを並列で検証します。受け入れられれば推論が高速化し、拒否されればシステムは正常にフォールバックします。

EAGLEチーム、vLLMチーム、TorchSpecチームは、これまでにEAGLE 1、2、3と続くEAGLEシリーズをリリースしてきました。このシリーズは、研究および本番システムの両方で最も広く採用され、実際にデプロイされている投機的デコーディングアルゴリズムのファミリーの1つとなっています。今回、このファミリーに信頼性を向上させるアップグレードとしてEAGLE 3.1が追加されました。

問題点：投機的デコーディングは制御された環境では良好に機能するものの、異なるチャットテンプレート、長コンテキスト入力、または分布外のシステムプロンプトの下では性能が低下することがよくあります。EAGLEチームは、この脆弱性の原因を「注意ドリフト（Attention Drift）」という現象に特定しました。投機の深さが増すにつれて、ドラフターは元のコンテキスト（固定トークン）から注意を徐々に自身の生成トークンへとシフトさせます。簡単に言えば、ドラフターは元のコンテキストではなく、自身の過去の出力に注意を向けるようになり、受け入れ長と出力の安定性が低下します。

根本原因は2つあります。第一に、融合された入力表現が不均衡になり、高層の隠れ状態がドラフター入力を支配します。第二に、正規化されていない残差パスにより、隠れ状態の大きさが投機ステップを経るごとに増大します。これらの効果が、投機の深いレベルでドラフターの安定性を徐々に低下させます。

EAGLE 3.1のソリューション：2つの主要なアーキテクチャ改善が導入されました。第一に、FC正規化（FC normalization）です。各ターゲット隠れ状態の後、FC層の前に正規化を適用し、隠れ状態の大きさを有界に保ちます。第二に、正規化後隠れ状態フィードバック（post-norm hidden-state feedback）です。正規化された隠れ状態を次のデコードステップに供給することで、ドラフターが単に追加の層を重ねるのではなく、再帰的に呼び出されるように動作します。

ベンチマーク結果：EAGLE 3と比較して、EAGLE 3.1は訓練時から推論時への外挿、長コンテキストのロバスト性、異なるチャットテンプレートやシステムプロンプトへの適応性、および多様なサービング環境での受け入れ長の安定性において優れています。長コンテキストワークロードでは、EAGLE 3.1はEAGLE 3と比較して受け入れ長が最大2倍に向上します。

トレーニングインフラ：TorchSpecは現在、EAGLE 3.1および将来の投機的デコーディングアルゴリズムの効率的なトレーニングサポートを提供しています。研究チームはTorchSpecとvLLMに基づき、Kimi K2.6用のEAGLE 3.1ドラフトモデルをトレーニングし、HuggingFaceでオープンソース化しました。これは、実際のサービングモデル上でEAGLE 3.1をTorchSpecトレーニングとvLLMサービングでデプロイする例を示しています。

vLLM統合：EAGLE 3.1は、vLLM内の既存のEAGLE 3実装の設定駆動型拡張として追加されました。FC正規化のサポート、正規化後隠れ状態フィードバック、およびターゲット隠れ状態に関するハードコードされた前提条件の除去が含まれます。既存のEAGLE 3チェックポイントとの後方互換性は完全に維持されており、EAGLE 3.1ドラフトモデルは同じ投機的デコーディングコードパスを介して直接プラグインできます。

デプロイ例（vLLM v0.22.0）：

vllm serve nvidia/Kimi-K2.6-NVFP4 \
--trust-remote-code \
--tensor-parallel-size 4 \
--tool-call-parser kimi_k2 \
--enable-auto-tool-choice \
--reasoning-parser kimi_k2 \
--attention-backend tokenspeed_mla \
--speculative-config '{"model":"lightseekorg/kimi-k2.6-eagle3.1-mla","method":"eagle3","num_speculative_tokens":3}' \
--language-model-only

ベンチマーク（Kimi K2.6、SPEED-Benchコードデータセット、GB200 TP=4）：同時実行数1でユーザーあたりの出力スループットが2.03倍、同時実行数4で1.71倍、同時実行数16で1.66倍向上しました。

主なポイント：EAGLE 3.1は注意ドリフトを修正します。2つのアーキテクチャ変更によりドラフターが安定化され、長コンテキストでの受け入れ長が最大2倍、スループットが大幅に向上しました。EAGLE 3チェックポイントとの後方互換性を維持し、vLLMメインにマージされており、v0.22.0でリリースされます。技術詳細は公式ブログを参照ください。