ニュース

한어Русский языкEnglishFrançaisIndonesianSanskrit日本語DeutschPortuguêsΕλληνικάespañolItalianoSuomalainenLatina

マシンハートレポート

マシーンハート編集部

怠けていると仕事がうまくいきます。

Llama 3.1 がリリースされましたが、もう試しましたか?お使いの PC が最新の最高スペックであっても、最小の 8B バージョンを実行すると、依然として大幅な遅延が発生する可能性があります。モデルの推論効率を向上させるために、研究者はさまざまな方法を考案しましたが、その多くはモデルの精度をある程度犠牲にすることになります。

最近、Apple と Meta AI の研究チームは、精度が大幅に低下しないようにしながら、Llama 2 の事前充填段階の推論速度を 2 倍以上に向上させることができる新しい方法を提案しました。3.1 の高速化は、いくつかのインスピレーションを与えてくれます。彼らはこのアプローチを LazyLLM と呼んでいます。これは Lazy Large Language Model の略です。

論文のタイトル: LazyLLM: 効率的な長いコンテキスト LLM 推論のための動的トークン プルーニング

論文アドレス: https://arxiv.org/abs/2407.14057

では、彼らはどのようにして LLM を怠惰にするのでしょうか?彼らの手法を理解するには、まず標準的なプロンプトベースの LLM 推論プロセスが何であるかを知る必要があります。簡単に言うと、図 1 に示すように、プロセスはプレフィルとデコードの 2 つの段階に分かれています。

事前設定段階では、モデルはプロンプト内の各トークンの KV キャッシュを計算して保存し、最初のトークンを予測します。事前設定フェーズに費やされる時間を「最初のトークンまでの時間 (TTFT)」と呼びます。

プレフィル段階の後にはデコード段階が続きます。この段階で、モデルはキャッシュされた KV を再度使用して、停止基準が満たされるまで次のトークンを繰り返しデコードします。

事前設定フェーズでは、すべての Transformer レイヤーがプロンプト内のすべてのトークンを使用します。現在最良の Transformer ベースの LLM は深くて広いため、プロンプトが長いと TTFT が遅くなる可能性があります。注意を計算するコストはプロンプト内のトークンの数に応じて二次関数的に増加します。たとえば、Llama 2 (バージョン 7B) は 32 層のトランスフォーマーをスタックし、モデルの寸法は 4096 です。この場合、TTFT は後続の各デコード ステップのウォールタイムの 21 倍を必要とし、これは LongBench ベンチマークの総生成時間の約 23% を占めます。

したがって、LLM 推論を効率的にするには、TTFT を最適化することが非常に重要なステップになります。

LLM 推論の最適化は活発な研究分野ですが、多くの手法はデコード段階の推論速度の向上に重点を置いています。研究者たちは、TTFT の改善にはほとんど注意を払ってきませんでした。圧縮ベースの研究結果の中には、LLM のサイズを削減することで暗黙的に TTFT を改善できるものもあります。

もう 1 つの研究の方向性は、静的な Transformer アーキテクチャの下で TTFT を改善することです。この研究の方向性では、当然次のような疑問が生じます。最初のトークンを生成するとき、すべてのプロンプト トークンは必須ですか?

図 2 は、LongBench ベンチマークでの LLM 解析結果を示しています。

最初に生成されたトークンでは、入力トークンの注意スコアが非常にまばらであることがわかります。これは、入力プロンプト内の多くのトークンが冗長であり、たとえそれらが削除されたとしても、次のトークンの予測には影響しないことを示しています。トークン。この観察は、チームが提案した LazyLLM の基礎になっています。

LazyLLM の利点としては、適用範囲が広いこと、トレーニングが不要であること、優れた結果が挙げられることです。図 3 は、標準 LLM と LazyLLM を比較しています。

レイジーLLM

図 4 に LazyLLM の全体的なフレームワークを示します。

LazyLLM は完全なコンテキストから開始してトークンを徐々にプルーニングし、それによって最終モデルを取得するために使用される計算の数を徐々に減らします。 LazyLLM を使用すると、前のステップでトークンの一部が枝刈りされた場合でも、モデルがさまざまな生成ステップでトークンのさまざまなサブセットを選択できることに注意してください。静的枝刈り (すべてのトークンが一度に枝刈りされる) と比較して、動的枝刈りは各生成ステップで次のトークンの予測を最適化し、モデルのパフォーマンスの維持に役立ちます。

プログレッシブトークンプルーニング

これまでの研究では、トークン プルーニングを使用して LLM 推論を最適化することに成功しています。ただし、これらの方法では、プルーニングを開始する前にプロンプ​​ト トークンの重要性を分析するために、予測された最初のいくつかのトークンの完全なアテンション マップを蓄積する必要があります。したがって、プレフィルフェーズ中にすべての KV キャッシュを計算する必要があるため、TTFT を削減するのには適していません。

比較すると、LazyLLM は「非常に怠惰」であり、推論の最初の反復 (事前入力ステップ) から開始して、次のトークンを予測するために重要なトークンのみを計算します。

最初の反復では、各トークンの重要性を判断することが重要な課題でした。トークンの隠れ状態が Transformer レイヤーを通過するにつれて進化することを示した以前の研究に触発されたチームのソリューションは、各生成ステップでレイヤーごとのトークン プルーニングを使用することです。具体的には、各レイヤーのアテンション マップを使用して、予測対象のトークンに対する入力トークンの重要性を判断します。

トークンの信頼スコアを計算した後、もう 1 つの難しい問題は、トークンをプルーニングするためのしきい値を決定することです。

具体的には、異なるレイヤおよび異なるタスクでは、注意スコアが変化するにつれて、このしきい値が変化する可能性があります。チームの解決策は、上位 k パーセンタイル選択戦略を使用することです。具体的には、トークンの信頼スコアが入力トークンの k 番目のパーセンタイル未満の場合、そのトークンはプルーニングされます。トークンがプルーニングされると、後続のすべてのレイヤーの計算に参加しなくなります。

つまり、後続のレイヤーで使用されるトークンは、前のレイヤーで使用されるトークンのサブセットです。

その後の実験では、枝刈り層の位置と枝刈りされたトークンの数が異なると、パフォーマンスも変化することがわかりました。具体的には、同じ Transformer レイヤーの場合、プルーニングによって削除されるトークンが増えると、モデルのパフォーマンスが徐々に低下します。

また、初期の層でのプルーニングと比較して、後の層でプルーニングを実行した場合にパフォーマンスが向上することもわかりました。これは、後の層がトークン プルーニングの影響を受けにくいことを示しています。速度と精度のバランスを改善するために、チームは図 4 に示すように段階的なプルーニングを使用し、初期の層により多くのトークンを保持し、その後、後の層に流れるにつれてトークンの数を徐々に減らしました。

Aux Cache (補助キャッシュ)

事前設定段階では KV キャッシュはなく、各トークンは非表示の状態で表されます。したがって、プログレッシブ トークン プルーニングは、プルーニングされたトークンの非表示状態を削除することによって実現できます。ただし、プログレッシブ トークン プルーニングを後続のデコード ステップに拡張することは簡単ではありません。その理由は、各デコード ステップでプレフィル段階で計算された KV バッファーを使用してアテンションを計算するためです。 LazyLLM は事前生成段階でプログレッシブ トークン プルーニングを実行するため、特定のレイヤーでプルーニングされたトークンの KV は、次のレイヤーの KV キャッシュには表示されません。

LazyLLM フレームワークでは、前のステップでプルーニングされているかどうかに関係なく、各生成ステップで完全な入力トークン シーケンスからトークンの異なるサブセットを選択できることに注意してください。たとえば、後続の復号化ステップでは、KV キャッシュに存在しないプルーニングされたトークンがアテンション計算のために再選択される場合があります。この場合、モデルはこれらのトークンの KV キャッシュを取得できません。

これに対する直感的な解決策は、Transformer のオリジンを介してトークンを渡すことです。ただし、これにより同じトークンが二重にカウントされることになり、最終的に全体の生成速度が遅くなります。

この問題を解決するために、チームは元の KV キャッシュに加えて別のキャッシュ、Aux Cache (補助キャッシュ) を導入しました。

トークンが除去された KV (図 4 の T4 や T7 など) が後続のレイヤーの KV キャッシュに表示されない場合、その非表示状態は後続の反復で取得できるように補助キャッシュによって保存されます。

図 4 に示すように、各デコード ステップで、各 Transformer レイヤーは最初に過去のトークンの KV キャッシュ (存在する場合) を取得します。 KV キャッシュにないトークンの場合、その非表示状態は、前の層を再度経由することなく、前の層の補助キャッシュから直接取得されます。 Aux Cache は、各トークンが各 Transformer レイヤーで最大 1 回計算されることを保証し、LazyLLM が最も遅い場合でも標準の LLM よりも高速であることを保証します。

実験

チームは、この新しい「遅延」アプローチを 2 つの大きな言語モデル、Llama 2 7B と XGen 7B でテストしました。比較用の標準 LLM は、追加のトレーニングを行わずに公開されている同じ事前トレーニング済みチェックポイント モデルです。

実験的なベンチマークは、長いコンテンツを理解するためのマルチタスク ベンチマークである LongBench です。 LongBench ベンチマークには 16 のデータ セットが含まれており、単一ドキュメント Q&A、複数ドキュメント Q&A、要約、少数ショット学習、合成タスク、コード補完を含む 6 つのタスクが含まれます。

評価指標は、TTFT 加速と精度のトレードオフに関する各メソッドの有効性と効率です。

結果

表 1 に、LazyLLM、標準 LLM、およびその他のベースライン手法の TTFT のスピードアップと精度の結果を示します。

この表では、ベースラインは標準の LLM 推論を指します。 ランダム トークン ドロップとは、トークンに対してランダムなプルーニングを実行することを指します。 静的トークン プルーニングとは、事前充填段階で前の Transformer レイヤーのアテンション メソッドに基づいて、入力トークンに対して 1 回限りのプルーニングを実行することを指します。 プロンプト圧縮は、LLM を使用して入力コンテキストの冗長性を削除するプロンプト圧縮方法です。

表 1 からわかるように、LazyLLM は TTFT の高速化において総合的に優れており、精度の低下は基本的に無視できます。 LLM を使用してプロンプトを圧縮するには、多くの計算が必要になることに注意してください。したがって、プロンプト圧縮により推論は高速化されますが、実際の TTFT は標準の LLM よりも長くなります。

全体的なビルド速度への影響

全体的な生成速度に対する新しい手法の影響を評価するために、チームは計算に使用されるプロンプト トークンの割合と生成の加速度を分析しました (表 2 を参照)。

LazyLLM の計算で使用されるトークンの割合は常に 100% 未満であることがわかります。これは、LazyLLM が生成の最後にプロンプ​​ト内のすべてのトークンを使い切っていないことを示していますが、理論的にはモデルはすべてのトークンを使用できます。これにより、さまざまなタスクの生成プロセス全体がさらに高速化されます。

さまざまなレイヤーでのドロップ率

チームはまた、プルーニング レイヤーの場所とプルーニングされたトークンの数の影響も分析しました。結果を図 6 に示します。

同じ Transformer レイヤーでプルーニングが実行される場合、残っているトークンが少なくなるほど、モデルのパフォーマンスが低下することがわかります。これは私たちの直観的な理解とも一致します。さらに、前の Transformer 層でプルーニングを実行する場合と比較して、後の層でプルーニングを行うとパフォーマンスが向上します。これは、後の層がトークン プルーニングの影響を受けにくいことを示しています。

これらの観察に基づいて、プログレッシブ トークン プルーニングの有効性が証明されたと言えます。

KV の漸進的な成長

最後に、チームはトークン プルーニング ロジックを使用してモデルの内部を理解しようとしました。具体的には、プロンプト トークンの累積使用率と、対応する未使用の割合を知りたいと考えています。この「累積トークン使用量」は、各ステップの KV キャッシュ サイズとして等価的に定義できます。図 7 は、LazyLLM の各段階でのこれらのプロンプト トークンの累積使用量を示しています。

この結果は、(理論的にはモデルがプロンプト内のすべてのトークンを使用できるにもかかわらず、多くのトークンがモデルによって決して選択されない) という仮説を裏付けています。

モデルがタスク実行の精度を維持できることを考慮すると、モデルは出力品質に影響を与えないトークンを効果的に破棄できると結論付けることができます。