大規模モデル推論の効率が損失なく3倍に向上。ウォータールー大学、北京大学などがEAGLEをリリース

大規模言語モデル (LLM) はさまざまな分野でますます使用されるようになっています。ただし、テキスト生成プロセスはコストがかかり、時間がかかります。この非効率性は、自己回帰復号化のアルゴリズムに起因します。各単語 (トークン) の生成には順方向伝播が必要であり、数十億から数千億のパラメータを持つ LLM へのアクセスが必要になります。その結果、従来の自己回帰デコードの速度が低下します。

最近、ウォータールー大学、カナダベクトル研究所、北京大学などの機関が共同でEAGLEをリリースしました。これは、モデル出力テキストの一貫した分布を確保しながら、大規模言語モデルの推論速度を向上させることを目的としています。この方法は、LLM の 2 番目のトップレベル固有ベクトルを外挿し、生成効率を大幅に向上させることができます。

技術レポート: https://sites.google.com/view/eagle-llm
コード（商用 Apache 2.0 をサポート）: https://github.com/SafeAILab/EAGLE

EAGLE には次の機能があります。

通常の自己回帰デコードよりも3倍高速（13B）
先読みデコード（13B）より2倍高速。
Medusaデコード（13B）より1.6倍高速。
生成されたテキストの分布は通常のデコードと一致していることが示されます。
RTX 3090 でトレーニング (1 ～ 2 日以内) およびテストできます。
vLLM、DeepSpeed、Mamba、FlashAttention、量子化、ハードウェア最適化などの他の並列技術と組み合わせることができます。

自己回帰デコードを高速化する 1 つの方法は、投機的サンプリングです。この手法では、より小さなスクラッチモデルを使用して、標準的な自己回帰生成によって次の複数の単語を推測します。その後、元の LLM はこれらの推測された単語を並列に検証します (検証には 1 つの前方パスのみが必要です)。ドラフトモデルが α 個の単語を正確に予測した場合、元の LLM の 1 回のフォワードパスで α+1 個の単語を生成できます。

投機的サンプリングでは、ドラフトモデルのタスクは、現在の単語シーケンスに基づいて次の単語を予測することです。大幅に少ないパラメータを持つモデルを使用してこのタスクを達成することは非常に困難であり、最適ではない結果が生成されることがよくあります。さらに、標準的な投機的サンプリング手法のドラフトモデルは、元の LLM によって抽出された豊富な意味情報を活用せずに次の単語を独自に予測するため、潜在的に非効率になります。

この制限が EAGLE の開発のきっかけとなりました。 EAGLE は、元の LLM によって抽出されたコンテキスト特徴 (つまり、モデルの 2 番目の最上位層によって出力される特徴ベクトル) を利用します。 EAGLE は次の基本原則に基づいて構築されています。

特徴ベクトルのシーケンスは圧縮可能であるため、前の特徴ベクトルに基づいて後続の特徴ベクトルを予測することが容易になります。

EAGLE は、自動回帰ヘッドと呼ばれる軽量プラグインをトレーニングします。このプラグインは、単語埋め込みレイヤーと連携して、現在の特徴シーケンスに基づいて、元のモデルの 2 番目の最上位レイヤーから次の特徴を予測します。次に、元の LLM の固定された分類ヘッドを使用して次の単語を予測します。特徴には単語のシーケンスよりも多くの情報が含まれているため、特徴を回帰するタスクは単語を予測するタスクよりもはるかに簡単になります。要約すると、EAGLE は特徴レベルで外挿し、小さな自己回帰ヘッドを使用し、次に固定分類ヘッドを使用して予測された単語シーケンスを生成します。 EAGLE は、Speculative Sampling、Medusa、Lookahead などの同様の作業と一致して、システム全体のスループットではなく、各ヒント推論のレイテンシに重点を置いています。

EAGLE - 大規模言語モデルの生成効率を高める手法

上の図は、入力と出力に関して、EAGLE と標準の投機的サンプリング、Medusa、Lookahead の違いを示しています。次の図は、EAGLE ワークフローを示しています。元の LLM のフォワードプロセスでは、EAGLE は 2 番目の最上位レイヤーから特徴を収集します。自己回帰ヘッドは、これらの特徴と以前に生成された単語の単語埋め込みを入力として受け取り、次の単語の推測を開始します。次に、固定された分類ヘッド (LM ヘッド) を使用して次の単語の分布を決定し、EAGLE がこの分布からサンプリングできるようにします。 EAGLE は、サンプリングを何度も繰り返すことで、下の図の右側に示すように、ツリーのような生成プロセスを実行します。この例では、EAGLE の 3 つのフォワードパスによって 10 語のツリーが「推測」されました。

EAGLE は軽量の自己回帰ヘッドを使用して、元の LLM の特徴を予測します。生成されたテキスト分布の一貫性を保証するために、EAGLE は予測されたツリー構造を検証します。この検証プロセスは、単一の順方向伝播を使用して完了できます。この予測と検証のサイクルを通じて、EAGLE はテキストワードをすばやく生成できます。

自己回帰ヘッドのトレーニングは安価です。 EAGLE は、70,000 未満の対話ターンを含む ShareGPT データセットを使用してトレーニングされます。自己回帰ヘッドの学習可能なパラメータの数も少ないです。上の画像で青で示されているように、ほとんどのコンポーネントは凍結されています。唯一の追加トレーニングは、0.24B-0.99B のパラメータを持つ単層 Transformer 構造である自己回帰ヘッドです。 GPU リソースが不十分な場合でも、自己回帰ヘッドをトレーニングできます。たとえば、Vicuna 33B の自己回帰ヘッドは、8 カードの RTX 3090 サーバーで 24 時間でトレーニングできます。

特徴を予測するために単語埋め込みを使用するのはなぜですか?

次の単語、その次の単語などを予測するために 2 番目の最上位レイヤーの特徴のみを使用する Medusa とは異なり、EAGLE は予測のための自己回帰ヘッド入力の一部として、現在サンプリングされた単語埋め込みも動的に使用します。この追加情報は、EAGLE がサンプリングプロセスで避けられないランダム性に対処するのに役立ちます。キューワードが「I」であると仮定して、下の図の例を考えてみましょう。 LLM は、「I」の後に「am」または「always」が続く確率を示します。 Medusa は、「am」または「always」のどちらがサンプリングされるかを考慮せず、「I」の下の次の単語の確率を直接予測します。したがって、メデューサの目標は、「私」だけが与えられた場合に、「私は」または「私はいつも」の次の単語を予測することです。サンプリングプロセスのランダム性により、Medusa の同じ入力「I」でも、次の単語の出力が「ready」または「begin」と異なる場合があり、入力と出力のマッピングに一貫性がなくなる可能性があります。対照的に、EAGLE の入力にはサンプリングされた単語の埋め込みが含まれており、入力と出力間の一貫したマッピングが保証されます。この区別により、EAGLE はサンプリングプロセスによって確立されたコンテキストを考慮に入れ、後続の単語をより正確に予測できるようになります。

ツリー構造の生成

EAGLE は、Speculative Sampling、Lookahead、Medusa などの他の推測検証フレームワークとは異なり、「単語の推測」段階でツリーのような生成構造を採用し、より高いデコード効率を実現します。図に示すように、標準的な投機的サンプリングと Lookahead の生成プロセスは線形または連鎖的です。推測段階ではコンテキストを構築できないため、Medusa の方法では直積によってツリーが生成され、隣接するレイヤー間の完全に接続されたグラフが生成されます。このアプローチでは、多くの場合、「I am begin」などの意味のない組み合わせが発生します。対照的に、EAGLE はよりスパースなツリー構造を作成します。このスパースツリー構造により、意味のないシーケンスの形成が防止され、より合理的な単語の組み合わせにコンピューティングリソースが集中します。

複数回の推測的サンプリング

標準的な推測的サンプリング法は、「単語を推測する」プロセス中に分布の一貫性を維持します。ツリー状の単語推測シナリオに適応するために、EAGLE はこのメソッドを複数ラウンドの再帰形式に拡張します。複数ラウンドの投機的サンプリングの疑似コードを以下に示します。ツリー生成プロセス中に、EAGLE はサンプリングされた各単語に対応する確率を記録します。 EAGLE は、複数回の推測的サンプリングを通じて、最終的に生成される各単語の分布が元の LLM の分布と一致することを保証します。