[論文レビュー] EAGLE: Speculative Sampling Requires Rethinking Feature Uncertainty

Yuhui Li, Fangyun Wei+2 more|arXiv (Cornell University)|Jan 26, 2024

Natural Language Processing Techniques被引用数 5

ひとこと要約

EAGLEは、推測サンプリングの不確実性に対処するために、1ステップ先のトークン情報を用いた特徴レベルの自己回帰ドラフト段階を導入し、ファインチューニングなしで出力分布を保持しつつ、MT-benchでの待機時間を大幅に短縮する速度向上を実現する（2.1x–3.8x）。

ABSTRACT

Autoregressive decoding makes the inference of Large Language Models (LLMs) time-consuming. In this paper, we reconsider speculative sampling and derive two key observations. Firstly, autoregression at the feature (second-to-top-layer) level is more straightforward than at the token level. Secondly, the inherent uncertainty in feature (second-to-top-layer) level autoregression constrains its performance. Based on these insights, we introduce EAGLE (Extrapolation Algorithm for Greater Language-model Efficiency), a simple yet highly efficient speculative sampling framework. By incorporating a token sequence advanced by one time step, EAGLE effectively resolves the uncertainty, enabling precise second-to-top-layer feature prediction with minimal overhead. We conducted comprehensive evaluations of EAGLE, including all models from the Vicuna and LLaMA2-Chat series, the MoE model Mixtral 8x7B Instruct, and tasks in dialogue, code generation, mathematical reasoning, and instruction following. For LLaMA2-Chat 70B, EAGLE achieved a latency speedup ratio of 2.7x-3.5x, doubled throughput, while maintaining the distribution of the generated text.

研究の動機と目的

• 出力分布を変更せずに、大規模言語モデル（LLMs）のより高速な自己回帰デコードを動機づける。
• 特徴レベルの自己回帰と不確実性の処理に焦点を当て、推測サンプリングを再検討する。
• 特徴レベルの不確実性を解消するために、1ステップのトークンシフトを用いる軽量なドラフトフレームワーク（EAGLE）を提案する。

提案手法

• トークンではなくセカンド・トップ・レイヤーの特徴量上で動作するドラフト段階を持つ推測サンプリングフレームワーク、EAGLEを提案する。
• 特徴量のサンプリング不確実性に対処するため、ドラフトモデルへの入力として1ステップ先行したトークン列を導入する。
• 3つのモジュールからなるドラフトモデル（Embedding、LM Head、Autoregression Head）と、ツリーアテンションを備えたツリーストラクチャのドラフトを使用する。
• Autoregression Headを、回帰と分類の結合損失で訓練し、固定データセットと軽量データ増強を用いて特徴誤差の蓄積を緩和する。
• ツリーアテンションを備えたターゲットLLMを用いる検証フェーズを実装し、再帰的な推測サンプリング手順に従って出力分布を維持する。
• 低コストの訓練体制（70BモデルでA100上1–2日）を提供し、gpt-fastなど他の加速手法との適合性を示す。

実験結果

リサーチクエスチョン

• RQ1 LLMsの推測サンプリングにおいて、特徴レベルの自己回帰がトークンレベルの自己回帰を上回ることができるか？
• RQ2 1ステップ先のトークン入力による特徴レベルの不確実性への対処が精度と速度向上にどう影響するか？
• RQ3 多様なモデルとタスクにおけるEAGLEの待機時間の利得と出力分布を保持する保証は？
• RQ4 量子化やコンパイルなど他の加速手法やバッチサイズ設定とEAGLEはどのように相互作用するか？
• RQ5 ドラフト時の訓練データ選択と特徴ノイズに対してEAGLEは頑健か？

主な発見

• EAGLEは Vicuna および LLaMA2-Chat ファミリ全体で大幅な速度向上を達成し、例えば LLaMA2-Chat 70B で 2.7x–3.5x の待機時間短縮を実現し、出力分布を保持しつつスループットを倍増させる。
• 特徴レベルでのドラフトは、1ステップ先のトークン入力を用い、トークンのみドラフトや従来法（Lookahead、Medusa）と比較して精度と速度を改善する。
• ドラフト/検証段階のツリーアテンションは、チェインドラフトより控えめな追加利益をもたらすが、最大フォワードパスを変えず、1パスあたりのトークン数を増やす。
• MT-bench で Lookahead (1.7x–2.1x faster) および Medusa (1.5x–1.6x faster) を上回る。特定の設定で 2x–3x のスループット改善を示し、グリーディー・非グリーディーモードの両方で分布を保持する。
• EAGLEはgpt-fastと組み合わせて RTX 3090 で LLaMA2-Chat 7B のデコードを 160.4 tokens/s に押し上げる。訓練コストは低く、Autoregression Headのパラメータは1B未満、70BモデルでA100s上1–2日。
• アブレーション研究は、特徴とシフトされたトークン入力によるランダム性対策の重要性を示し、固定データセット訓練がターゲット-LM生成データとほぼ同等の性能を示すことから、データ感度が低いことを強調する。