Apple は大規模モデルに遅延を学習させます: 最初のトークンをより速く吐き出し、精度を維持します

より良く働くために怠け者になりましょう。

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

最近、AppleとMeta AIの研究チームは、精度を大幅に低下させずにLlama 2の事前充填段階の推論速度を2倍以上に向上させることができる新しい方法を提案しました。 Llama 3.1 の高速化のためのインスピレーション。彼らはこのアプローチを LazyLLM (Lazy Large Language Model) と呼んでいます。 Paperタイトル：lazyllm：効率的な長いコンテキスト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 を比較しています。

LazyLLM

図 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 ベンチマークには、単一ドキュメント Q&A、複数ドキュメント Q&A、要約、少数ショット学習、合成タスク、コード補完を含む 6 つのタスクをカバーする 16 個のデータセットが含まれています。

評価指標は、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 の各段階でのプロンプト トークンの累積使用量を示しています。

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

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

。

以上がApple は大規模モデルに遅延を学習させます: 最初のトークンをより速く吐き出し、精度を維持しますの詳細内容です。詳細については、PHP 中国語 Web サイトの他の関連記事を参照してください。