3倍の生成速度とメモリコストの削減、Medusa2を超える効率的なデコードフレームワークがついに登場

n トークン シーケンス、CLLMs+Jacobi デコード フレームワークを効率的にデコードします。

伝統的に、大規模言語モデル (LLM) は、各トークンを 1 つずつデコードする逐次デコーダーとして考えられています。

上海交通大学とカリフォルニア大学の研究チームは、事前トレーニングされた LLM が効率的な並列デコーダになるように簡単に教育できることを示し、Consistent Large Language Model (CLLM) と呼ばれる新しい並列デコーダ ファミリを導入し、各推論ステップで n トークン シーケンスを効率的にデコードすることで推論レイテンシーを軽減します。

この論文では、「人間が頭の中で完全な文章を形成した後、一語一語の表現を表現するために使用する認知プロセスを模倣することは、事前トレーニングされた LLM を微調整するだけで効果的に学習できる」ことを研究が示しています。具体的には、CLLM は、ランダムに初期化された n トークンのシーケンスをできるだけ少ないステップにマッピングすることにより、自己回帰 (AR) デコードと同じ結果のデコードされたシーケンスを生成します。このようにして、並列復号トレーニングを実行できます。

実験結果は、研究チームが提案した方法を使用して得られた CLLM が非常に効果的であることを示し、この方法が生成速度で 2.4 倍から 3.4 倍の向上を達成し、Medusa2 や Eagle などの他の高速推論技術に匹敵することを示しています。また、推論中に補助モデル コンポーネントに対応するための追加のメモリ コストは必要ありません。

論文名:「CLLMs: Consistency Large Language Models」
論文リンク: https://arxiv.org/pdf/2403.00835

図 1: GSM8K の場合上は、Jacobi デコードを使用した場合、CLLM-ABEL-7B-001 がベースライン ABEL-7B-001 よりも約 3 倍高速であることを示しています。

ヤコビ復号

大規模言語モデル (LLM) は、プログラミングから法律や健康に関するアドバイスの提供まで、人間の生活の様相を変えています。

ただし、図 1 に示すように、LLM は推論中に自己回帰デコードを使用してトークンごとに応答を生成するため、応答が長くなると遅延が長くなります。自己回帰デコードを使用すると、多くの場合、一度に複数のトークンを生成して推論を高速化するために、アーキテクチャの変更、補助コンポーネント、または最初のドラフト モデルが必要になります。図 2: 従来の自己回帰 (AR) デコードの概略図: 一度にトークンを生成します。

ヤコビ復号は、非線形方程式を解くためのヤコビおよびガウス・ザイデル固定小数点反復法から派生し、貪欲復号を使用した自己回帰生成とまったく同じであることが証明されています。

ヤコビ復号化は、逐次生成プロセスを n 個の変数を含む n 個の非線形方程式系に再構築し、ヤコビ反復に基づいて並列的に解くことができます。

各反復ステップでは、複数の正しいトークン (いわゆる「正しい」とは、貪欲なサンプリング戦略に基づく自己回帰復号結果と一致することを意味します) を予測する可能性があり、それによって自己回帰復号が高速化される可能性があります。

。

具体的には、ヤコビ復号法はまず入力プロンプトからシーケンスの次のトークンをランダムに推測します (以下、特に明記しない限り n

トークン シーケンスと呼びます)。

次に、ⁿ トークン シーケンスが、反復更新のヒントとともに LLM に供給されます。このプロセスは、n

-token のシーケンスが安定し、変更されなくなり、固定点に達するまで継続されます。

Jacobi デコードでは、自己回帰 (AR) デコードよりも LLM へのクエリが必要ないことに注目してください。最終的に、 n

トークンのシーケンスは、貪欲戦略に基づいた AR デコードによって生成された出力に収束します。最初のランダムな推測から最終的な AR 生成結果までのプロセスは、「ヤコビ軌道」と呼ばれます。

ヤコビ復号反復プロセスとヤコビの軌跡の例を図 2 に示します。 ヤコビ復号の制限:

しかし、実際には、通常のヤコビ復号では、LLM の加速効果がわずかに改善されるだけであり、たとえば、平均加速率は 1.05 倍にすぎません。これは、以前のトークンにエラーがある場合、LLM が正しいトークンを生成することが困難であるためです。

したがって、ほとんどのヤコビ反復では n トークンのシーケンスに対して 1 つの修正しか取得できず、その結果、図 3 の左側に示されているより長い軌道が得られます。

先読みデコードと投機的デコード手法は、ヤコビ デコードと従来の自己回帰デコードの非効率性を軽減しようとしますが、推論中に追加のメモリ コストが発生します。

CLLM では、これらの追加のメモリコストは必要ありません。

一貫した大規模言語モデル (CLLM)

予備的なヤコビ復号:

プロンプト x と事前トレーニングされた LLM p(·​​|x) が与えられた場合、通常、研究者は標準的な自己回帰を使用します。 (AR) デコード方法は、貪欲戦略の下でモデルの応答を取得します。つまり、

ヤコビデコードは、LLM 推論プロセスを非線形方程式系を解くプロセスとして再構築し、デコード プロセスを実行可能な形式の並列コンピューティングに変換します。 。以下を考慮してください:

研究者は、上記の方程式を非線形方程式系として書き直すことができます:

注意すべき点:

プロセスは次のような特定の k 値で終了します:

次に、 を不動点として定義し、をヤコビ軌跡として定義します。この問題を解決するために、研究チームは、ヤコビ軌道

J 上の任意の点 y を固定点 y* に一貫してマッピングできるように、事前トレーニングされた LLM を調整することを提案しています。 驚くべきことに、彼らはそのような目標が、拡散モデルの主要な加速方法である一貫性モデルの目標と似ていることを発見しました。

チームが提案した方法では、ターゲットモデルから収集されたヤコビ軌道を使用してモデルがトレーニングされ、ヤコビ反復中にシングルステップ収束を促進する損失関数が使用されます。

CLLM に調整されるターゲット モデル

p ごとに、トレーニングは 2 つの部分で構成されます: (1) ヤコビ軌道の準備:

各プロンプトに対して、作成者は各トークンの切り捨てに対してヤコビを順番に実行します。応答シーケンス全体

l が生成されます。これは、すべての連続する固定点の連結に相当します。 軌跡に沿って生成された各シーケンスは、データ エントリとしてカウントされます。

ここで、

N

(N ≫ n) トークンを含む長い応答の場合、この切り捨てにより、長い入力でのモデル評価の遅延が回避されることに注意することが重要です。 (2) 一貫性と AR 損失を使用したトレーニング: 著者は、CLLM を調整するために 2 つの損失を共同最適化します。一方、AR 損失は、CLLM がターゲット LLM から逸脱することを防ぎます。生成品質を維持します。 ️図 4: ワンステップの収束一貫性トレーニングの概略図: ターゲットの調整LLM は、ヤコビ軌道上の任意の状態が入力として受け取られたときに常に固定点を予測します。

整合性と AR 損失:

(1) 整合性損失

p

をターゲット LLM として表します。

を、パラメータ

θ

を pに初期化したCLLMとして表すものとします。

プロンプト x と対応するヤコビ軌道 J について、 y と y*

がそれぞれ軌道上のランダムな状態と固定点を表すものとします。

は、グローバル整合性 (GC) 損失 と呼ばれる次の損失を最小限に抑えることで、入力が y* の場合に y を出力するように促すことができます。

この式では、

著者は、データセットからの一様なサンプリングを表すためにシンボルを広範囲に使用しています。

D(・||・) は 2 つの分布間の距離を表し、その選択については GKD 法で説明します。この記事では主にフォワード KL が使用されます。

あるいは、整合性モデルの式に従ってローカル整合性 (LC) 損失を使用します。

ここで、隣接する状態: ヤコビ軌道 J では、同じ出力を生成するように駆動されます:

(2) AR 損失:

ターゲット LLM の分布からの逸脱を避けるために、著者は、ターゲット LLM p の生成 l に基づいて分布を結合します。

いくつかの重み ω を使用して両方の損失を結合することにより、トレーニング CLLM の総損失は次のようになります。 :

実験

結果:

全体として、この実験は 3 つのドメイン固有のタスクをカバーします:

(1) Spider (テキストから SQL)

(2) Human-Eval (Python コード)完了) と GSM8k (数学)

(3) より広いオープン ドメイン セッション チャレンジ MT ベンチ。

報告された実験では、タスクに応じて、微調整されたエンコーダー LLM、Deepseek-coder-7B-instruct、LLaMA-2-7B または ABEL-7B-001 をターゲット モデルとして使用します。

トレーニングと評価は、NVIDIA A100 40GB サーバーで実行されます。

^{図 5: さまざまなダウンストリーム タスクに対する CLLM の高速化効果。結果は、「CLLM は事前トレーニングされたモデルよりもはるかに高速で、Medusa と比較して同等の高速化を実現しますが、推論中に追加コストは発生しません。」}

^{図 6: CLLM と他のモデルの比較図特定のドメイン タスク (Spider、CSN-Python、GSM8k) および MT ベンチのベンチマーク。 CLLM は、追加の推論コスト (FLOPS とメモリ消費量から判断) を導入せずに、Medusa2 と比較して同様またはそれ以上の高速化を実現します。}

専門分野:

図 5 から、元のターゲット モデル、Medusa2、推測デコードなどの他のベースラインと比較して、CLLM が最も大幅な高速化を達成していることがわかります。

オープン ドメイン セッション チャレンジ (MT ベンチ): ShareGPT データセットを使用して LLaMA2-7B からトレーニングされた

CLLM は、先読みデコードと組み合わせると、Medusa2 とほぼ同じ高速化を達成し、MT で同等のスコアを取得します。 -ベンチ。

ただし、CLLM は、ターゲット モデルの元のアーキテクチャや補助コンポーネントを変更する必要がないため、より適応性があり、メモリ効率が高くなります。

トレーニングコスト:

CLLM の微調整にかかるコストはそれほど高くありません。

たとえば、LLaMA-7B の場合、Spider データセットでは、約 100 万個のトークンを渡すだけで 3.4 倍の高速化を達成できます。データセットのサイズが大きい場合 (CodeSearchNet-Python など)、CLLM をトレーニングするためのヤコビ軌道の生成に使用する必要があるのはデータセットの 10% だけであり、結果的に約 2.5 倍の速度向上が得られます。

トークンの総数は、次の方法で推定できます:

N = 各プロンプトの平均軌跡量 × 平均軌跡長さ × プロンプトの数。

^{図 7: Spider 上のターゲット LLM と CLLM 間のヤコビ軌道の比較。ヤコビ軌道上の各点は色分けされたシーケンスです。AR 結果との正しい一致は青でマークされ、不正確なものは赤でマークされます。 CLLM は効率が向上し、ターゲット LLM よりも 2 倍速く固定点に収束します。このように CLLM の効率が向上したのは、一貫性の損失によるものと考えられます。これにより、指定された各プレフィックスの n トークン シーケンスの構造の学習が容易になります。}

図 6 の左側は、ターゲット LLM が通常 1 回の反復で正しいトークンを 1 つだけ生成することを示しています。対照的に、CLLM では、著者らは、複数の連続するトークンが 1 回のヤコビ反復で正しく予測される、急速な進歩の現象を発見しました。

さらに、ターゲット LLM では、事前に正しく生成されたトークン (図 7 の左側のインデックス 6 と 7 の「country」や「H」など) が、後続の反復で不正確に置き換えられることがよくあります。

一方、CLLM は正しいトークンを予測する能力を示しており、以前に間違ったトークンが存在した場合でもトークンが変更されないことを保証します。

著者はこのようなトークンを「固定トークン」と呼びます。これら 2 つの現象は共に、ヤコビ復号における CLLM の迅速な収束に寄与し、その結果、生成速度が大幅に向上します。

研究チームはまた、CLLMがトレーニングを通じて重要な言語概念であるコロケーションを獲得したことも観察しました。つまり、「ランダムな偶然によって予想よりも頻繁に一緒に出現する一連の単語または用語」

言語は単語の分離だけで構成されているわけではありません。ですが、特定の単語のペアにも大きく依存します。コロケーションの例は、自然言語とプログラミング言語の両方に豊富にあります。

それらには以下が含まれます:

• 動詞 + 前置詞の組み合わせ (例: 「話す」、「... を思い出させる」)

• 動詞 + 名詞構造 (例: 「決断する」、「捕まえる」) cold")

• 多くのドメイン固有の構文構造 (プログラミングの「SELECT ... FROM ...」、「if ... else」など)。

一貫性生成目標により、CLLM はヤコビ軌道の任意の点からそのような構造を推論できるようになり、CLLM が多数の連語を習得し、複数の単語を同時に予測して反復ステップを最小限に抑えることが容易になります。

^{参考リンク：}

^{https://hao-ai-lab.github.io/blogs/cllm/}

以上が3倍の生成速度とメモリコストの削減、Medusa2を超える効率的なデコードフレームワークがついに登場の詳細内容です。詳細については、PHP 中国語 Web サイトの他の関連記事を参照してください。