ICML2024ハイスコア！魔法のように注意力を変化させ、小さなモデルでも大きなモデルの 2 倍戦えるようにします

トランスフォーマーのコアメカニズムを集中的に改善し、小さなモデルを2倍の大きさでモデル化できるようにしました。

ICML+2024 の高得点論文で、Caiyun Technology チームは、Transformer コア コンポーネント アテンション モジュール (MHA) を置き換える DCFormer フレームワークを構築し、動的に結合可能なマルチヘッド アテンション (DCMHA) を提案しました。

DCMHA は、MHA アテンション ヘッドの検索選択ループと変換ループの固定バインドを削除し、入力に基づいてそれらを動的に組み合わせることができるようにし、モデルの表現能力を根本的に向上させます。

本来の意味は、各層が固定の H アテンション ヘッドを持っているということですが、ほぼ同じ量のパラメータと計算能力で、最大で動的に結合できると理解できます。 H×H注目です。 以下のように、コンテンツを微調整すると、元のテキストの意味がより明確に表現されます。 元のモデルの各層には、固定数の H アテンション ヘッドが含まれています。これで、

DCMHA プラグアンドプレイを使用して、あらゆる Transformer アーキテクチャの MHA を置き換え、多用途で効率的でスケーラブルな新しいアーキテクチャ DCFormer を取得できるようになります。

この研究は、北京郵電大学の研究者とAIスタートアップのCaiyun Technologyによって共同で完成されました。

DCFormer に基づいて研究者によって構築されたモデル DCPythia-6.9B は、トレーニング前の複雑さと下流のタスク評価の点で、オープンソースの Pythia-12B よりも優れています。

DCFormer モデルは、1.7 ～ 2 倍の計算を必要とする Transformer モデルとパフォーマンスが同等です。

マルチヘッド アテンション モジュールの制限は何ですか?

大規模モデルのスケーリング則は、計算能力が向上するにつれて、モデルが大きくなり、より多くのデータが含まれるようになり、モデルの効果がますます向上することを示しています。この道路の天井がどのくらい高いのか、AGI に到達できるのかを明確に説明できる人はいませんが、これが現在最も一般的なアプローチです。

しかし、これに加えて、別の質問も考える価値があります。現在の大きなモデルのほとんどはトランスフォーマーに基づいており、ビルディングブロックとしてトランスフォーマー自体が組み込まれています。改善の余地はどれくらいありますか?

これはモデル構造の研究で答えるべき基本的な質問であり、Caiyun Technology と北京郵電大学が共同で完成させた DCFormer の研究の出発点でもあります。

Transformer のマルチヘッド アテンション モジュール (MHA) では、各アテンション ヘッドは互いに完全に独立して動作します。

この設計は、そのシンプルさと実装の容易さにより、実際には非常に成功していますが、同時に注意スコアマトリクスの順位が低くなり、表現力が弱まり、注意頭機能の無駄が繰り返されることになります。パラメータと計算、その他の欠点。これに基づいて、近年のいくつかの研究では、アテンションヘッド間の何らかの形の相互作用を導入しようと試みられています。

Transformer ループ理論によると、MHA では、各アテンション ヘッドの動作は 4 つの重み行列 W^Q、W^K、W^V、W^O(WO ^{によって取得されます) によって記述されます。 MHA の出力射影行列のセグメント化)}。

それらのうち、W

Q^WK^{はQKループ}(または検索選択ループ)と呼ばれ、現在のトークンからコンテキスト内のどの(一部の)トークンに焦点を当てるかを決定します。例えば:

W

O^WV^{は、OVループ}(または射影変換ループ)と呼ばれ、対象トークンからどのような情報を取得するか(またはどの属性を射影するか)を決定します。現在の位置で残差ストリームに書き込み、次のトークンを予測します。例:

研究者は、検索 (どこで取得するか) と変換 (何を取得するか) は元々 2 つの独立したものであり、これらは個別に指定し、必要に応じて自由に組み合わせる必要があることに気づきました (SQL クエリの場合と同様に、選択WHERE の後の条件と SELECT の後の属性射影は別々に記述されます)、MHA ではそれらをアテンション ヘッドを使用して QKOV に「バンドル」する必要があり、柔軟性と表現機能が制限されます。

たとえば、アテンション ヘッド A、B、C を持つモデルがあり、その QK ループと OV ループが上記の例 = を完了できると仮定します。その後、それを次のように置き換えます。既存のアテンション ヘッド QK ループと OV ループを組み合わせると、モデルは「向きを変えることができない」可能性があります (研究者のシステムによって構築された合成テスト セットによって検証されています。6B 未満の中小規模のモデルは、そのようなモデルではうまく機能しません)一見単純なタスク）。

長い注意力のダイナミックな組み合わせはどのようなものですか?

これを出発点として、この記事の研究チームは MHA の 合成操作 を導入しました。

以下の図に示すように、DCMHA が取得されます。

△図 1. DCMHA の全体構造

は QW^{Q になります。} と KW^K によって計算された注意スコア行列 A^S と注意重み行列 A^W は、num_heads 次元に線形マッピングされ、新しい行列 A' が取得されます。さまざまなアテンションヘッドの組み合わせの効果を実現するために、さまざまな線形マッピング行列 ^{(合成マップ)} を介して VWV が乗算されます。

たとえば、図 2(c) では、ヘッド 3 と 7 の QK ループがヘッド 1 の OV ループと結合されて、「新しい」アテンション ヘッドが形成されます。

△ 図2. 8つの注目ヘッドの簡略化された典型的な構成マップ関数、明るい色は大きな値を表します

表現能力を最大化するために、研究者はマッピング行列が動的に駆動されることを望んでいます。入力生成、つまり、アテンションヘッドがどのように結合されるかを動的に決定します。

しかし、生成したいマッピング行列は1つではありません。そのような行列は、シーケンス内のソース位置のクエリQ_iと宛先位置のキーK_jの各ペアに対して生成する必要があります。計算オーバーヘッドとメモリ使用量の両方を受け入れるのは難しいでしょう。

この目的のために、マッピング行列を入力に依存しない静的行列 W_b、低ランク行列 w₁w₂、および対角行列 Diag(w_{g) の合計にさらに分解します。} )、それらはそれぞれ、基本的な組み合わせ、アテンション ヘッド間の限定された方法の動的組み合わせ (つまり、ランク R、およびヘッド自体の動的ゲート (図 2 (d) および図 3 を参照) (b)) 。後の 2 つの行列は、Q 行列と K 行列によって動的に生成されます。

効果を犠牲にすることなく、計算とパラメータの複雑さをほぼ無視できるレベルまで削減します (詳細については、論文の複雑さの分析を参照してください)。 JAX および PyTorch 実装レベルの最適化と組み合わせることで、DCFormer は効率的にトレーニングおよび推論できます。

△図 3. Compose の計算

はどのように行われますか?

スケールの拡張

アーキテクチャの品質を評価するために、研究者が注目する中心的な指標は、コンピューティングパワーをインテリジェンスに変換する効率 (またはパフォーマンスコンピューティングパワー比)、つまり達成可能なモデルのパフォーマンスです。ユニットのコンピューティング能力を投資することによる改善 - コンピューティング能力の消費を減らし、より良いモデルを取得します。

それは、図 4 と図 5 のスケーリング則曲線 から見ることができます (対数座標では、計算能力の変化に応じて、各モデル アーキテクチャの損失は近似直線として描くことができます。損失が低いほど優れています)モデル) 、DCFormer は 1.7 ～ 2 倍の計算能力で Transformer モデルの効果を実現できます。つまり、計算能力のインテリジェント変換率が 1.7 ～ 2 倍増加します。

△図4. TransformerとDCFormerのスケール拡張効果

△図5. PythiaとDCPythiaのスケール拡張効果

この改善をどのように理解すればよいでしょうか?

2017 年の Transformer の誕生以来、パフォーマンスと計算能力比の向上という観点から見ると、GLU MLP と回転位置エンコーディング RoPE は、普遍的に効果的であることが証明され、大規模なシステムで広く採用されている数少ないアーキテクチャの改善点の 2 つです。練習の数。

元の Transformer にこれら 2 つの改善点を追加したアーキテクチャは、Transformer++ とも呼ばれ、Llama や Mistral などの最も強力なオープンソース モデルはすべてこのアーキテクチャを使用しています。 Transformer または Transformer++ アーキテクチャに関係なく、DCMHA を通じて大幅な改善が得られます。

1.4B モデルスケールでは、DCMHA の改善は Transformer++ の 2 つの改善の合計よりも大きく、スケーラビリティも優れています (図 4 の青緑の線と黒の線の比較、DCMHA の改善)計算能力が増加するにつれて減衰します。遅くなります。図 4 と図 5 の比較)。

DCFormer は、Transformer の機能を新たなレベルに引き上げたと言えます。

ダウンストリーム タスクの評価

研究チームは、主流の NLP ダウンストリーム タスクで評価するために 2 つのモデル DCPythia-2.8B と DCPythia-6.9B をトレーニングし、同じスケールのオープンソース モデル Pythia と比較しました (トレーニングでは同じものを使用します) Pythia 設定としてのハイパーパラメータ) 。

△表 1. ダウンストリームタスクにおける DCFormer と Pythia のパフォーマンス

表 1 からわかるように、DCPythia-2.8B と 6.9B は Pile 検証セットの PPL が低いだけでなく、PPL も低いです。 DCPythia6.9B のダウンストリーム タスクの平均精度は、ほとんどのダウンストリーム タスクで Pythia を大幅に上回っており、Pythia-12B をも上回っています。

DCFormer++2.8B は DCPythia-2.8B と比較してさらに改良されており、DCMHA と Lllama アーキテクチャの組み合わせの有効性が検証されています。

トレーニングと推論の速度

DCMHA の導入によりトレーニングと推論のオーバーヘッドが追加されますが、表 2 から、DCFormer++ のトレーニング速度は Transformer++ の 74.5% ～ 89.2%、推論速度は 81.1% であることがわかります。 -89.7 % であり、モデル パラメーターが増加するにつれて、追加の計算オーバーヘッドは徐々に減少します。

△表 2. Transformer++ と DCFormer++ のトレーニング速度と推論速度の比較

トレーニング速度は、TPU v3 ポッド、シーケンス長が 2048、batch_size が 1k の場合で比較され、推論速度は A100 80G GPU の場合です。入力の長さは 1024 で、生成された長さは 128 です。

アブレーション実験

結果は以下のとおりです:

△表 3. DCMHA のアブレーション実験

表 3 から、次の点がわかります:

• 静的な組み合わせの重みを追加すると、PPL を減らすことができますが、しかし、動的結合重みを導入すると、PPL をさらに減らすことができ、これは動的結合の必要性を示しています。
• 低ランクの動的組み合わせは、動的ゲートよりも優れたパフォーマンスを発揮します。
• クエリ単位またはキー単位の動的組み合わせのみを使用して得られる PPL は非常に似ており、DCFormer++ との差は非常に小さいです。
• アテンションヘッドの組み合わせは、ソフトマックスの前よりもソフトマックスの後に行う方が効果的です。これはおそらく、ソフトマックス後の確率が出力に直接影響するためです。
• 動的組み合わせ重みのランクはあまり大きく設定する必要はありません。これは、組み合わせ重みのランクが低いことも示しています。

さらに、研究者らは、ローカル アテンション レイヤーの割合を増やし、クエリごとの動的組み合わせのみを使用することで、トレーニングと推論のオーバーヘッドをさらに削減しました。詳細については、論文の表 10 を参照してください。

一般的に、研究チームは2つの結論を出しました。

動的ウェイトについて: Mamba、GLA、RWKV6、HGRN などの最近の SSM およびリニア アテンション/RNN 作品は、動的 (入力依存) ウェイトを導入することで Transformer++ に追いつきましたが、DCFormer はアテンション ヘッドの動的組み合わせを使用します。この方法は、ソフトマックス アテンションを使用する場合、動的な重みを導入することで Transformer++ の効果を大幅に改善できることを示しています。

モデル アーキテクチャの革新について: この研究は、極端なコンピューティング能力とインテリジェントな変換効率を備えた「理想的なモデル アーキテクチャ」が存在する場合、現在の Transformer アーキテクチャはすでに非常に強力ではあるものの、この理想からはまだ遠い可能性があることを示しています。アーキテクチャにはまだ大きな改善の余地があります。したがって、コンピューティングパワーとデータの積み重ねによる奇跡の活発な発展に加えて、モデルアーキテクチャの革新も大きな可能性を秘めています。

研究チームはまた、Caiyun Technology が自社製品の Caiyun Weather、Caiyun Xiaoyi、Caiyun Xiaomeng に DCformer を適用する最初の企業になると述べました。

研究の詳細については、元の論文を参照してください。

ICML2024 論文リンク: https://icml.cc/virtual/2024/poster/34047。
Arxiv 論文リンク: https://arxiv.org/abs/2405.08553。
コードリンク: https://github.com/Caiyun-AI/DCFormer。

以上がICML2024ハイスコア！魔法のように注意力を変化させ、小さなモデルでも大きなモデルの 2 倍戦えるようにしますの詳細内容です。詳細については、PHP 中国語 Web サイトの他の関連記事を参照してください。