新世代のアテンション メカニズム ライトニング アテンション-2: 無制限のシーケンス長、一定の計算能力オーバーヘッド、より高いモデリング精度

大規模な言語モデルの現在の適用はシーケンス長の制限によって制限されており、人工知能の分野での適用が制限されています。たとえば、複数ターンの対話、長いテキストの理解、およびマルチモーダルなデータの処理と生成には、特定の課題があります。この制限の基本的な理由は、現在、大規模な言語モデルで一般的に使用されている Transformer アーキテクチャでは、計算の複雑さとシーケンスの長さとの間に 2 次の関係があるためです。したがって、シーケンスの長さが増加すると、コンピューティング リソースの需要が指数関数的に増加します。したがって、長いシーケンスを効率的に処理する方法は、大規模な言語モデルが常に直面する課題の 1 つです。

これまでのアプローチは主に、推論段階で大規模な言語モデルをより長いシーケンスに適応させることに焦点を当ててきました。 1 つのアプローチは、Alibi または同様の相対位置エンコーディングを使用して、モデルが異なる長さの入力シーケンスに適応できるようにすることです。もう 1 つのアプローチは、RoPE または同様の相対位置エンコーディングを差分に使用し、すでにトレーニングされたモデルを簡単に微調整してシーケンスの長さを延長することです。これらの方法では、大規模なモデルに特定の長いシーケンスのモデリング機能を持たせることができますが、トレーニングと推論のオーバーヘッドは削減されていません。

OpenNLPLab チームは、Lightning Attendant-2 と呼ばれる新しい線形アテンション メカニズムをオープンソース化しました。これは、大規模な言語モデルの長いシーケンスの問題を解決するように設計されています。このメカニズムにより、長いシーケンスのトレーニングと推論のコストが 1K シーケンス長と一致するように維持され、設定して忘れるソリューションが可能になります。メモリのボトルネックに達する前であっても、シーケンスの長さを増やしてもモデルのトレーニング速度に悪影響を及ぼさないため、無制限の長さの事前トレーニングが可能になります。さらに、非常に長いテキストの推論コストも一貫しているか、1K トークンと比較してさらに低いため、現在の大規模な言語モデルの推論コストが大幅に削減されます。以下の図に示すように、モデル サイズが 400M、1B、3B の場合、系列長が増加するにつれて、FlashAttendant2 がサポートする LLaMA の学習速度は急激に低下し始めますが、Lightning Attendance-2 がサポートする TansNormerLLM の速度はほぼ低下しています。変化なし。 ＃＃＃＃＃＃＃＃＃＃＃＃＃＃＃＃＃＃＃＃＃図1＃＃＃＃＃＃＃＃＃＃＃＃＃＃＃＃＃＃＃＃＃＃＃＃＃＃＃論文: ライトニング アテンション-2: 大規模な言語モデルで無制限のシーケンス長を処理するための無料ランチ

論文のアドレス: https://arxiv.org/pdf/ 2401.04658.pdf

オープンソースのアドレス: https://github.com/OpenNLPLab/lightning-attention

• ライトニング アテンション-2 はじめに
• # 大規模モデルの事前トレーニング速度を、異なるシーケンス長にわたって一貫して維持することは、不可能な作業のように思えます。しかし、2020 年にリニア アテンションが登場して以来、研究者たちはリニア アテンションの実際の効率を理論上の線形計算の複雑さと一致させるために懸命に取り組んできました。 2023 年半ばまで、リニア アテンションに関する研究は、主に Transformer アーキテクチャとの精度調整に焦点を当てます。最後に、改良されたリニア アテンション メカニズムの出現により、精度においては最先端の Transformer アーキテクチャに匹敵するようになりました。 ただし、線形アテンションにおける「左乗算から右乗算」という最も重要な計算トリックは、実際の実装では直接左乗算アルゴリズムよりもはるかに遅くなります。これは、右乗算の実装には、多数のループ演算を含む累積加算 (cumsum) の使用が必要であり、多数の I/O 操作により、右乗算の効率が左乗算よりも大幅に低下するためです。 したがって、異なるシーケンス長にわたって大規模モデルの事前トレーニング速度を一貫して維持するという課題がまだあります。研究者は、計算効率を向上させ、I/O 操作を削減するために、リニア アテンションの実装をさらに調査および改善する必要があります。これは、事前トレーニング速度の一貫性を達成し、異なるシーケンス長のタスク要件に適切に対処するのに役立ちます。
• 図 2

ライトニング アテンション 2 のアイデアをよりよく理解するために、まず伝統的なソフトマックス アテンションの計算式を確認しましょう: O=softmax ((QK^T)⊙M_) V、ここで Q、K、V、M、O はそれぞれクエリ、キー、値、マスク、出力行列ですここでの M は、一方向タスク (GPT など) では下三角の all-1 行列ですが、双方向タスク (Bert など) では無視できます。つまり、マスク行列はありません。双方向タスク用。

著者は、Lightning Attendance-2 の全体的な考え方を次の 3 つの点にまとめて説明します。

1. Linear Attendee の核となるアイデアの 1 つは、計算量の多いソフトマックス演算子を削除して、Attention の計算式を O=((QK^T)⊙M_) V として記述できるようにすることです。ただし、一方向タスクにはマスク行列 M が存在するため、この形式でも左乗算の計算しか実行できないため、O (N) の複雑さは得られません。しかし、双方向タスクの場合、マスク行列がないため、Linear Attendant の計算式は O=(QK^T) V にさらに簡略化できます。 Linear Attendant の微妙な点は、行列乗算の結合法則を使用するだけで、その計算式が O=Q (K^T V) にさらに変換できることです。この計算形式は右乗算と呼ばれ、対応する前者は次のようになります。左。図 2 から、リニア アテンションが双方向タスクにおいて魅力的な O (N) の複雑さを実現できることが直感的に理解できます。

2. しかし、デコーダ専用の GPT モデルが徐々に LLM の事実上の標準になるにつれて、一方向のタスクを高速化するために Linear Attendee の適切な乗算機能を使用する方法が重要になってきました。解決すべき緊急の問題です。この問題を解決するために、この記事の著者は、「分割統治」の考え方を使用して、注意行列の計算を対角行列と非対角行列の 2 つの形式に分割し、異なるものを使用することを提案しました。それらを計算する方法。図 3 に示すように、Linear Attendant-2 は、コンピューター分野で一般的に使用されるタイルの考え方を使用して、Q、K、および V 行列を同じ数のブロックに分割します。このうち、ブロック自体の計算 (ブロック内) はマスク行列の存在により左の乗算計算方法が残り、複雑度は O (N^2) ですが、ブロックの計算 (ブロック間) はマスク行列の存在により左の乗算計算方法が維持されます。 block) にはマスク行列がありませんが、 の存在により、適切な乗算計算方法を使用して O (N) の複雑さを楽しむことができます。 2 つを別々に計算した後、これらを直接加算して、i 番目のブロックに対応するリニア アテンション出力 Oi を取得できます。同時に、KV の状態がcumsum によって蓄積され、次のブロックの計算に使用されます。このように、ライトニング アテンション 2 全体のアルゴリズムの複雑さは、ブロック内トレードオフでは O (N^2)、ブロック間トレードオフでは O (N) になります。より良いトレードオフを得る方法は、タイリングのブロック サイズによって決まります。

3. 注意深い読者であれば、上記のプロセスは Lightning Attendant-2 のアルゴリズム部分にすぎないことがわかります。Lightning という名前が付けられた理由は、作者がこのアルゴリズム プロセスを十分に考慮したためです。 GPU ハードウェアで実行されるため、プロセスにおける効率の問題が発生します。 FlashAttend の一連の作業からインスピレーションを得て、実際に GPU 上で計算を実行する際、作者は分割された Q_i、K_i、V_i テンソルを、GPU 内の容量が大きく低速な HBM から容量が小さく高速な SRAM に移動しました。これにより、大量のメモリ IO オーバーヘッドが削減されます。ブロックが Linear Attendance の計算を完了すると、その出力結果 O_i は HBM に戻されます。すべてのブロックが処理されるまで、このプロセスを繰り返します。

さらに詳細を知りたい読者は、この記事のアルゴリズム 1 とアルゴリズム 2、および論文の詳細な導出プロセスを注意深く読むことができます。アルゴリズムと導出プロセスの両方で、Lightning Attendant-2 の順方向プロセスと逆方向プロセスが区別されており、読者の理解を深めるのに役立ちます。

#画像 3

Lightning Attendance-2 の精度比較

研究者らはまず、小規模 (400M) パラメータ モデルで Lightning Attendant-2 と Lightning Attendant-1 の精度の違いを次のように比較しました。以下に示すように、両者にほとんど違いはありません。

研究者らは次に、Lightning Attendant-2 をサポートする TransNormerLLM (TNL-LA2) を、1B および 3B の他の高度な非 Transformer アーキテクチャと組み合わせました。ネットワークと FlashAttendant2 を利用した LLaMA を同じコーパスの下で比較しました。以下の図に示すように、TNL-LA2 と LLaMA は同様の傾向を維持しており、損失パフォーマンスが優れています。この実験は、Lightning Attendant-2 が言語モデリングにおいて最先端の Transformer アーキテクチャに劣らない精度パフォーマンスを備えていることを示しています。

在大語言模型任務中，研究人員比較了 TNL-LA2 15B 與 Pythia 在類似大小下的大模型常見 Benchmark 的結果。如下表所示，在吃掉了相同 tokens 的條件下，TNL-LA2 在常識推理和多項選擇綜合能力上均略高於基於 Softmax 的注意力的 Pythia 模型。

Lightning Attention-2 速度比較

研究人員對Lightning Attention-2 與FlashAttention2 進行了單模組速度與顯存佔用比較。如下圖所示，相較於 Lightning Attention-1 和 FlashAttention2，在速度上，Lightning Attention-2 表現出了相比於序列長度的嚴格線性增長。在顯存佔用上，三者都顯示了類似的趨勢，但 Lightning Attention-2 的顯存佔用更小。這個的原因是 FlashAttention2 和 Lightning Attention-1 的顯存佔用也是近似線性的。

筆者註意到，這篇文章主要關注點在解決線性注意力網路的訓練速度上，並實現了任意長度的長序列與1K 序列相似的訓練速度。在推理速度上，並沒有過多的介紹。這是因為線性注意力在推理的時候可以無損地轉化為 RNN 模式，從而達到類似的效果，即推理單 token 的速度恆定。對於 Transformer 來說，目前 token 的推理速度與它之前的 token 數量相關。

筆者測試了 Lightning Attention-1 加持的 TransNormerLLM-7B 與常見的 7B 模型在推理速度上的比較。如下圖所示，在近似參數大小下，Lightning Attention-1 的吞吐速度是百川的 4 倍，ChatGLM 的 3.5 倍以上，顯示出了優異的推理速度優勢。

小結

#Lightning Attention-2 代表了線性注意力機制的重大進步，使其無論在精度或速度上均可以完美的替換傳統的Softmax 注意力，為今後越來越大的模型提供了可持續擴展的能力，並提供了一條以更高效率處理無限長序列的途徑。 OpenNLPLab 團隊在未來將研究基於線性注意力機制的序列平行演算法，以解決目前遇到的顯存屏障問題。

以上が新世代のアテンション メカニズム ライトニング アテンション-2: 無制限のシーケンス長、一定の計算能力オーバーヘッド、より高いモデリング精度の詳細内容です。詳細については、PHP 中国語 Web サイトの他の関連記事を参照してください。