トランスフォーマーの簡易版がここにあります、ネットユーザー：今年の論文

Transformer アーキテクチャは、ディープラーニング分野における最近の多くの成功の原動力であると言えます。深い Transformer アーキテクチャを構築する簡単な方法は、複数の同一の Transformer「ブロック」を積み重ねることですが、各「ブロック」は比較的複雑で、優れたパフォーマンスを実現するには特定の組み合わせで配置する必要がある多くの異なるコンポーネントで構成されています。

2017年にTransformerアーキテクチャが誕生して以来、研究者たちはそれを基にした派生研究を多数開始してきましたが、Transformer「ブロック」はほとんど変更されていません。

そこで質問ですが、標準の Transformer ブロックを簡素化できるでしょうか?

最近の論文では、ETH チューリッヒの研究者らが、下流のタスクの収束特性とパフォーマンスに影響を与えずに、LLM に必要な標準の Transformer ブロックを簡素化する方法について議論しました。信号伝播理論と経験的証拠に基づいて、残差接続、正規化レイヤー (LayerNorm)、投影および値パラメーター、MLP シリアル化サブブロック (並列レイアウトを容易にする) などの一部の部分を削除して、GPT のようなデコーダー アーキテクチャとエンコーダー スタイルの BERT モデルを簡素化できることを発見しました。

関係する各コンポーネントについて、トレーニングの速度を低下させることなく（更新ごとのステップと実行時間の両方の点で）削除できるかどうか、また削除するには Transformer ブロックにどのようなアーキテクチャ変更が必要になるかを検討します。

論文リンク: https://arxiv.org/pdf/2311.01906.pdf

Lightning AI の創設者であり機械学習研究者でもあるセバスチャン・ラシュカ氏は、この研究を「今年のお気に入りの論文の 1 つ」と呼んでいます。

しかし、一部の研究者は次のように疑問を呈した。「完全なトレーニング プロセスを見ない限り、コメントするのは難しい。正規化レイヤーも残余接続もない場合、1 億を超えるパラメータを持つネットワークでどのように拡張できるのか? 」

セバスチャン・ラシュカ氏も同意し、「確かに、彼らが実験したアーキテクチャは比較的小規模であり、これがトランスフォーマーの数十億のパラメータに一般化されるかどうかはまだ分からない」と述べた。しかし、彼はそれでもこの研究は印象的であり、残留接続の除去が成功したことは（初期化スキームを考慮すると）まったくあり得ることだと考えた。

この点について、チューリング賞受賞者のヤン・ルカン氏は次のようにコメントしています。「私たちはディープラーニングアーキテクチャの表面に触れただけです。これは高次元空間なので、ボリュームはほぼ完全に表面に収まりますが、私たちは表面のほんの一部に触れただけです。 」

Transformer ブロックを簡素化する必要があるのはなぜですか?

研究者らは、トレーニング速度に影響を与えずにTransformerブロックを簡素化することは興味深い研究課題であると述べた。

まず、現代のニューラル ネットワーク アーキテクチャは設計が複雑で、多くのコンポーネントで構成されており、これらのさまざまなコンポーネントがニューラル ネットワーク トレーニングのダイナミクスでどのような役割を果たし、どのように相互作用するかは明確ではありません。この質問は、ディープラーニングの理論と実践のギャップに関係するため、非常に重要です。

信号伝播理論は、ディープ ニューラル ネットワーク アーキテクチャにおける実用的な設計選択の動機となるため、影響力があることが証明されています。信号伝播は、入力全体にわたる階層化表現の内積によって捕捉される、初期化時のニューラル ネットワーク内の幾何学的情報の進化を研究し、ディープ ニューラル ネットワークのトレーニングにおいて多くの印象的な成果を達成してきました。

しかし、現在の理論では、初期化時のモデルのみ、また多くの場合は初期のフォワードパスのみを考慮しているため、残差接続がトレーニング速度に及ぼす利点など、ディープ ニューラル ネットワーク トレーニングのダイナミクスにおける多くの複雑な問題を明らかにすることができません。信号伝播は改変の動機付けに極めて重要であるが、研究者らは理論だけでは簡略化されたトランスフォーマーモジュールを導き出すことはできず、経験的洞察にも頼ったと述べている。

実際のアプリケーションに関しては、大規模な Transformer モデルのトレーニングと展開にかかる現在のコストが高いことを考慮すると、Transformer アーキテクチャのトレーニングおよび推論パイプラインの効率性が向上すると、大きな節約につながる可能性があります。不要なコンポーネントを削除して Transformer モジュールを簡素化できれば、パラメータの数を減らしてモデルのスループットを向上させることができます。

この論文では、残差接続、値パラメータ、投影パラメータ、シリアル化サブブロックを削除した後、トレーニング速度と下流タスクのパフォーマンスの両方の点で標準のTransformerに匹敵することが可能であるとも述べられています。最終的に、研究者はパラメータの数を 16% 削減し、トレーニングと推論時間の両方でスループットが 16% 増加したことを確認しました。

Transformer ブロックを簡素化する方法は?

研究者らは、信号伝播理論と経験的観察を組み合わせて、Pre-LN モジュールから最も単純な Transformer ブロックを生成する方法を紹介しました (以下を参照)。

論文の第 4 章の各セクションでは、トレーニング速度に影響を与えずにブロック コンポーネントを一度に 1 つずつ削除する方法を紹介しています。

このセクションのすべての実験では、CodeParrot データセットで 18 ブロック、768 幅の因果デコーダーのみの GPT モデルを使用します。このモデルは十分な大きさであるため、著者が単一トレーニング エポック モードの場合、一般化ギャップが非常に小さくなり (図 2 を参照)、トレーニング速度に集中できます。

残りの接続を削除する

研究者たちはまず、注意サブブロック内の残留接続を削除することを検討しました。式（１）の表記では、これはα_SAをゼロに固定することと同等である。注意残差接続を単純に削除すると、信号劣化、つまりランクの崩壊につながり、トレーニング性が低下します。論文のセクション 4.1 では、研究者らがその方法を詳しく説明しています。

投影/値パラメータを削除する

図 3 から、更新ごとのトレーニング速度の低下を最小限に抑えながら、値と投影パラメータ W^V、W^P を完全に削除できることがわかります。つまり、β_V = β_P = 0で恒等式が初期化されている場合

の場合、同じ数のトレーニングステップの後、この研究では基本的にPre-LNブロックのパフォーマンスを達成できます。この場合、トレーニングプロセス全体を通じて W^V = W^P = I となり、値と投影パラメータは一貫しています。詳細な方法は第4.2節で紹介されている。

MLPサブブロックの残りの接続を削除する

上記のモジュールと比較すると、MLP サブブロックの残留接続を削除することはより困難です。以前の研究と同様に、著者らは、MLP 残差接続なしで Adam を使用する場合、信号伝播によってアクティベーションをより線形にすると、図 22 に示すように、更新ごとのトレーニング速度が大幅に低下することを発見しました。

彼らはまた、ガウス重み、直交重み、恒等重みなど、Looks Linear 初期化のさまざまなバリエーションを試しましたが、効果はありませんでした。そのため、彼らは研究全体を通して、MLP サブブロック内で標準的なアクティベーション (ReLU など) と初期化を使用しています。

著者らは、PALM や ViT-22B などの最近のいくつかの大型変圧器モデルで人気が実証されている並列 MHA および MLP サブブロックの概念に目を向けます。並列変圧器ブロックを下図に示します。

著者らは、論文のセクション 4.3 で、MLP サブブロックの残余接続を削除する具体的な操作を詳しく紹介しています。

正規化レイヤーを削除する

最後に削除されるのは正規化レイヤーで、その結果、図 1 の右上隅にある最も単純なブロックが作成されます。信号伝播の初期化の観点から、著者はこのセクションの簡略化のどの段階でも正規化レイヤーを削除できます。考え方としては、Pre-LN ブロックでの正規化によって残差ブランチが暗黙的に軽量化され、この有益な効果は正規化レイヤーなしで別のメカニズムによって再現できるということです。つまり、残差接続を使用するときに残差ブランチを明示的に軽量化するか、注意行列をアイデンティティにバイアスするか、MLP 非線形性を「より」線形になるように変換します。

著者らは修正プロセス中にこれらのメカニズム（MLP β_FF と Shaped Attention の重みの削減など）を考慮に入れたため、正規化は必要ありません。著者らはセクション4.4でさらに詳しい情報を提供しています。

実験結果

深さ拡張

信号伝播理論は通常、深い深さに焦点を当てているため、このような状況では信号の劣化が頻繁に見られます。そこで当然の疑問として、簡素化されたトランスフォーマー ブロックによって達成されたトレーニング速度の向上を、より深いレベルまで拡張できるかどうかが挙げられます。

図 6 から、深さを 18 ブロックから 72 ブロックに拡張することで、モデルと Pre-LN トランスフォーマーの両方のパフォーマンスが向上していることがわかります。これは、この研究で簡素化されたモデルはトレーニングが高速であるだけでなく、深さが増すことで提供される追加機能も利用できることを示しています。実際、正規化を使用すると、この研究の簡略化されたブロックと Pre-LN の各更新された軌跡は、異なる深度ではほとんど区別がつきません。

バート

次に、著者らは、簡略化されたブロックのパフォーマンスが、自己回帰デコーダーに加えて、さまざまなデータセットやアーキテクチャ、下流のタスクにも当てはまることを示しています。彼らは、マスクされた言語モデリングに双方向エンコーダのみの BERT モデルの一般的な設定を選択し、ダウンストリーム GLUE ベンチマークを採用しました。

図 7 に示すように、簡素化されたブロックは、24 時間の実行時間内で、(Crammed) Pre-LN ベースラインと比較して、マスクされた言語モデリング タスクで同等の事前トレーニング速度を達成できます。一方、値と投影を変更せずに残余接続を削除すると、トレーニング速度が大幅に低下します。図24では、著者らはマイクロバッチステップの同等の図を示しています。

さらに、表 1 では、研究者らは、GLUE ベンチマークで微調整した後、自分たちの方法が Crammed BERT ベースラインと同等のパフォーマンスを発揮することを発見しました。

下流のタスクは表 2 に分類されます。公平な比較のために、彼らはGeiping & Goldstein (2023)と同じ微調整プロトコル（5エポック、タスクごとに一定のハイパーパラメータ、ドロップアウト正規化）を使用しました。

効率性の向上

表 1 では、研究者らは、マスクされた言語モデリング タスクでさまざまな Transformer ブロックを使用したモデルのパラメーターの数とトレーニング速度も詳細に示しています。彼らは、24 時間の事前トレーニング中に実行されたマイクロバッチ ステップの数とベースラインの Pre-LN Crammed BERT の比率として速度を計算しました。モデルでは 16% 少ないパラメータが使用され、SAS-P と SAS の各反復はそれぞれ Pre-LN ブロックよりも 16% と 9% 高速化されたという結論が出ました。

ここでの実装では、並列ブロックはPre-LNブロックよりもわずか5％高速であるのに対し、Chowdheryら（2022）は15％高速なトレーニング速度を観測しており、より最適化された実装によって全体的なトレーニング速度がさらに向上する可能性があることを示唆しています。 Geiping & Goldstein (2023) と同様に、この実装では PyTorch (Sarofeen et al., 2022) の自動演算子融合を使用します。

より長いトレーニング

最後に、より小さなモデルをより多くのデータで長期間トレーニングするという現在の傾向を考慮して、簡素化されたブロックが長いトレーニング期間を経ても Pre-LN ブロックと同じトレーニング速度を達成できるかどうかについて説明します。これを実現するために、彼らは CodeParrot の図 5 のモデルを使用し、3 倍のトークンでトレーニングしました。正確に言うと、バッチ サイズ 128、シーケンス長 128 で約 120K ステップ (40K ステップではなく) トレーニングされ、約 20 億のトークンが生成されます。

図 8 からわかるように、トレーニングにさらに多くのトークンが使用されると、簡略化された SAS および SAS-P コード ブロックのトレーニング速度は PreLN コード ブロックのトレーニング速度と同等であり、PreLN コード ブロックよりも速くなります。

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