カスタムEAGLE-3ヘッドを投機的デコード用に訓練する方法

投機的デコードは、出力品質を犠牲にせずにLLM推論のレイテンシを削減する最も効果的な手法の一つとして登場しました。利用可能な手法の中で、EAGLE-3はそのシンプルさ、性能、柔軟性から広く使われています。この記事では、カスタムEAGLE-3ヘッドの訓練に必要なすべてを説明します。データセットの準備からハイパーパラメータ調整、デプロイまでをカバーします。Basetenでは、EAGLEヘッドを使用したモデルで1.5〜2.5倍のレイテンシ改善を確認しており、ターゲットモデルはQwen3-4Bです。

EAGLE-3は、自己回帰LLM推論のための投機的デコード手法です。基本アイデアは、ターゲットモデルに軽量な「ドラフトヘッド」をアタッチし、一度に複数の将来トークンを予測することです。ターゲットモデルはそれらの予測を一度のフォワードパスで検証します。ドラフトヘッドが正確な場合、複数のデコードステップをスキップでき、エンドツーエンドのレイテンシが大幅に削減されます。EAGLE論文ではベンチマークで最大4〜6倍の高速化が報告されていますが、その一部はサービスフレームワークの違いによるものです。実運用では、EAGLEヘッド自体によるレイテンシ改善は通常1.5〜2.5倍です。

EAGLE-3の実用的な特性には、品質低下なしのレイテンシ改善（検証されたトークンはターゲットモデルが生成したものと同一）、メモリバウンドワークロードでの最大の恩恵（長コンテキストやシングルバッチサービスでメモリ帯域幅がボトルネックの場合）、軽量トレーニング（ドラフトヘッドは単一のトランスフォーマーデコーダ層で、パラメータはターゲットモデルの1〜5％）が含まれます。

EAGLEは、レイテンシに敏感だが高バッチではないワークロードに最適です。特に、ユーザーが最初のトークン時間よりもエンドツーエンドの応答時間を気にする場合、例えばコード生成、エージェンティックワークフロー、構造化出力、リアルタイム会話アシスタントに効果的です。

主要パラメータには、TTT長（トレーニング中にヘッドが自身の予測をコンテキストとして使用するトークン数、推奨7-9）、ドラフトトークン数（推論時にヘッドが提案するトークン数、推奨3-4）、学習率（モデルサイズに応じて調整：小規模3-7Bで1e-4、中規模7-20Bで5e-5、大規模20B+で2e-5）、サンプリングパラメータ（貪欲デコードが最適、温度>0では受理率が約15-25%低下）が含まれます。

データセットの品質はEAGLEヘッド訓練で最も重要な要素です。ヘッドはターゲットモデルのトークン分布を学習する必要があり、汎用テキスト分布ではありません。汎用タスクでは20万〜30万サンプル（小モデル）または約50万サンプル（大モデル）、特殊タスクでは約10万サンプルが必要です。各サンプルは1000〜2000トークンを推奨します。黄金律：ターゲットモデルで出力を再生成して分布を一致させること。

トレーニングワークフローは、環境設定、データ準備（チャットテンプレート適用、Hugging Faceへのアップロードまたは再生成）、トレーニング開始（ターゲットモデルパス、データセットパス、TTT長、学習率、バッチサイズ、最大シーケンス長、エポック数を設定）、モニタリング（損失曲線が急降下後安定、精度が70-80%に上昇）を含みます。収束しない場合は、TTT長を増やし、学習率を調整し、データ形式を確認します。

評価では、EAGLEヘッドありとなしでレイテンシとスループットを比較します。適切に訓練されたヘッドは、メモリバウンドワークロードで1.5〜2.5倍のレイテンシ改善を示すはずです。デプロイでは、訓練されたヘッドをHugging Faceに保存し、サービス設定でパスを参照します。サービスフレームワークが自動的にヘッドのロード、ドラフト生成、検証を行います。

一般的な問題として、トレーニングが収束しない（TTT長を増やす、学習率を確認、データ形式を検証）や推論時の受理率が低い（ドラフトトークン数を減らす、推論入力形式を訓練時と一致させる）があります。

カスタムEAGLE-3ヘッドの訓練は、レイテンシに敏感な環境でLLMをサービスするチームにとって高レバレッジな最適化です。プロセスは直接的です：代表的なデータセットを準備して出力を再生成し、少数のハイパーパラメータを設定し、軽量ヘッドを訓練します。しかし、正しいデータ分布、チャットテンプレートの一致、TTT長の調整が、有意義な高速化を提供するヘッドとそうでないヘッドの違いを生みます。結果として、出力品質をゼロ低下させて1.5〜2.5倍のレイテンシ改善が得られます。これは推論最適化分野では珍しい組み合わせです。