Mix-LN: Unleashing the Power of Deep Layers by Combining Pre-LN and Post-LN

図1の(a)と(b)は、それぞれPost-LNとPre-LNのTransformerアーキテクチャを示している。 形式的には、$x$を入力、$\mathcal{F}(x)$をFFN層または多頭注意層、$\mathrm{LN}(\cdot)$を層正規化と定義する。Post-LNは残差加算の後に$\mathrm{LN}(\cdot)$を適用する：

対照的に、Pre-LNは残差加算の前に$\mathrm{LN}(\cdot)$を適用する：

我々は式(1)と(2)の導関数を以下のように計算できる：

上記の両方の式には重要な要素、すなわち層正規化のヤコビ行列$\mathbf{J}_{LN}(x^{\prime})=\frac{\partial\text{LN}(x^{\prime})}{\partial x^{\prime}}$が含まれている。ここで$x^{\prime}$は$\mathrm{LN}(\cdot)$の入力である。Xiong et al. (2020); Takase et al. (2023)の証明に従い、$x^{\prime}$が平均0の正規分布に従うと仮定すると、以下のようになる：

ここで、$\sigma_{x^{\prime}}$は$x^{\prime}$の標準偏差であり、$d$は隠れ次元である。したがって、

ここで、$z=(x^{\prime}-\mu_{x^{\prime}})/\sigma_{x^{\prime}}$は層正規化後に得られる標準正規分布である。LLMsでは$d\gg 1$であるため、最終的に以下を得る：

実際には、訓練中に$\sigma_{x^{\prime}}$が徐々に1より大きくなることを観察している。これは、LNのヤコビ行列のスペクトルノルムが1未満であることを意味する。式(3)のPost-LNの導関数によると、この縮小係数は複数の層$L$にわたって$\prod_{l=1}^{L}\frac{1}{\sigma_{x^{\prime}}^{l}}$として蓄積され、Post-LN Transformerの初期層で勾配消失を引き起こす。

対照的に、Pre-LNの残差接続の導関数は、式(4)に示すように、$\mathrm{LN}$の導関数に関連する項から分離されている。この設計は初期層での勾配消失問題の防止に役立つ。しかし、Pre-LNは残差接続を制約しないため、層の深さが増すにつれて連続するTransformerブロックの出力が蓄積される。結果として、式(4)のPre-LNの導関数は単位行列に近づき、式(4)のPre-LN操作全体が学習に効果的に寄与しなくなることを示している。これは、Pre-LNの深層が初期層と比較してモデルの全体的な学習にあまり寄与しない理由を説明している。

方法： 我々の評価方法は、Pre-LNを使用するモデルとPost-LNを使用するモデルの2つを比較分析することを含む。各モデルの異なる深さにおける層の有効性を実証的に評価することで、Pre-LNモデルではより深い層の有効性が低下し、Post-LNモデルではより深い層でも品質が維持または向上することが予想される。

LLMモデル： 我々の仮説を厳密に評価するため、2つの異なるカテゴリのLLMで実験を行う：(i) オープンウェイトの大規模LLMと (ii) 自社開発の小規模LLM。オープンウェイトカテゴリでは、Pre-LNモデルの代表としてLLaMa2-7B (Touvron et al., 2023)を、Post-LNモデルとしてBERT-large (Devlin, 2018)を選択した。しかし、これらのオープンウェイトモデルは正規化だけでなく、訓練データ、活性化関数、コンテキスト長などの他の要因も異なるため、正規化の影響のみを分離することが複雑になる。これらの交絡変数を制御するため、我々は自ら小規模LLMをゼロから訓練する第二の実験セットを実施した。目的は、モデル間の唯一の違いが層正規化の選択であることを確実にすることである。具体的には、Lialin et al. (2023b); Zhao et al. (2024)に従い、C4データセットでPre-LNまたはPost-LNを使用してLLaMa-130Mモデルを訓練し、RMSNorm (Zhang & Sennrich, 2019)とSwiGLU活性化 (Shazeer, 2020)を使用した。訓練設定の詳細については付録Aを参照されたい。

評価指標： 我々の仮説を検証する上での重要な課題は、個々の層の有効性を捉える堅牢な指標を定義し選択することにある。本研究では、(i) 角度距離と (ii) 性能低下という2つの指標を採用し、各層の役割と貢献を意味のある形で評価する。さらに、異なる層正規化がグラデーントフローに与える影響を示すため、各層の勾配ノルムも報告する。

(i) 角度距離 $d(x^{\ell},x^{\ell+n})$はGromov et al. (2024)で使用され、中立的な事前訓練データセットにおける層$\ell$への入力と層$\ell+n$への入力の間の角度距離を測定する。形式的には、$x_{T}^{\ell}$を層$\ell$への入力、$x_{T}^{\ell+n}$を層$\ell+n$への入力とすると、層$\ell$とその後続の$n^{th}$層、つまり$\ell+n$の間の単一トークン$T$に対する角度距離は以下のように与えられる：

ここで$|\!|\cdot|\!|$は$L^{2}$ノルムを表し、 係数$1/\pi$は$d(x^{\ell},x^{\ell+n})$を[0, 1]の範囲にスケーリングする。ランダム性の影響を排除するため、本稿で報告される角度距離はC4データセットから256Kトークンにわたって平均化されている。$d(x^{\ell},x^{\ell+n})$の値が小さいほど距離が短く、2つのベクトルがより類似していることを意味する。隣接する層との表現が極めて類似している層は容易に除去でき、したがってその重みはあまり効果的ではない。理想的には、ネットワークのパラメータを最も効果的に利用するために、表現は層から層へと大きく変化すべきである(Yang et al., 2023; Gromov et al., 2024)。

(ii) 性能低下 $\Delta P^{(\ell)}$は、層$\ell$を剪定する前後のLLMの性能の差を指す。これは、その層の除去によって引き起こされる性能劣化を定量化する。形式的には、以下のように定義できる：

ここで$P_{\text{original}}$は剪定なしのモデルの性能、$P_{\text{pruned}}^{(\ell)}$は層$\ell$を剪定した後のモデルの性能である。$\Delta P^{(\ell)}$の値が小さいほど、その層を除去してもモデルの出力にほとんど変化がないことを示し、その層の重要性が低いことを示唆する。具体的には、LLaMA2-7Bについては、一般的に使用されるMMLU (Hendrycks et al., 2020)を評価タスクとして選択した。BERT-largeについては、SQuAD v1.1 (Rajpurkar, 2016)を評価タスクとして選択した。我々の自社開発LLMの限られた容量を考慮し、性能低下の評価にはMMULの代わりに教師あり微調整後のARC-e (Clark et al., 2018)を選択した。

本節では、Mix-LNの有効性を、Post-LN (Nguyen & Salazar, 2019)、DeepNorm (Wang et al., 2024)、Pre-LN (Dai, 2019)を含む様々な一般的な正規化手法と比較して検証する。Lialin et al. (2023a); Zhao et al. (2024)に従い、我々は71Mから1Bパラメータまでの様々なサイズのLLaMAベースのアーキテクチャを用いて実験を行い、RMSNorm (Shazeer, 2020)とSwiGLU活性化関数 (Zhang & Sennrich, 2019)を組み込んだ。モデルはAdam (Kingma, 2014)を用いて学習され、モデルサイズに基づいて異なる学習率を使用した。具体的には、250Mパラメータ以下のモデルには1e-3の学習率を、1Bパラメータモデルには5e-4の学習率を使用した。 同じサイズのすべてのモデルは、正規化を除いて同一の設定で学習された。Mix-LNのハイパーパラメータ$\alpha$の最適値を決定するために、表5に示すように、LLaMA-250Mを用いて小規模なハイパーパラメータ探索を行った。我々は$\alpha=0.25$が最良の性能を提供することを発見し、したがって、この値をすべてのモデルサイズに適用した。

結果は表1に示されている。Post-LNは一般的に最悪の性能を示し、より大きなモデルでは発散さえしており、これはTransformerにおいてPost-LNが学習の不安定性に苦しむことを示す先行研究と一致している (Xiong et al., 2020; Takase et al., 2022)。Post-LNの修正版であるDeepNormは、小規模モデルではPre-LNと同等の性能を達成するが、1Bパラメータモデルでは学習中に発散を経験する。この観察は、早期層の勾配が消失し、適切なモデルの収束を妨げるPost-LNの深刻な学習不安定性を確認するものである。対照的に、Mix-LNは様々なモデルサイズにわたって一貫して最低のパープレキシティを達成している。Mix-LNは、広く普及しているPre-LNと比較して、LLaMA-71MとLLaMA-250Mでそれぞれ1.65および0.53のパープレキシティの顕著な改善を達成している。

上記の結果は、Mix-LNがPost-LNの不安定性を克服するだけでなく、Pre-LNとPost-LNの利点を組み合わせることでモデルの品質を向上させ、大規模LLMに理想的な選択肢となることを明確に示している。

Mix-LNの利点が7Bパラメータなどのより大規模なモデルにも適用できるかを評価することは、その利点が大規模モデルにも当てはまるかを示すために不可欠である。この目的のため、我々はZhao et al. (2024)のセットアップに従い、埋め込みサイズ4096と合計32層を特徴とするLLaMa-7Bアーキテクチャを用いて実験を行った。層正規化手法を除くすべてのトレーニング設定は同一に保たれた。計算資源の制約により、我々は13,000ステップのトレーニングのみを完了することができた。 我々の経験上、トレーニング初期に一貫した改善を示すモデルは、通常、後期段階でもこれらの利点を維持する傾向がある。 トレーニング曲線の比較をFigure 4に示す。ここでは、Mix-LNがLLaMa-7Bの初期トレーニング段階において、Pre-LNを一貫して上回っていることが明らかである。

我々は、Mix-LNによって生成される優れた中間層および深層が、教師あり微調整中により効果的に学習できると考えている。この利点は、これらの層がPre-LNで訓練されたものと比較して、より多様で豊かな特徴を捉えていることに起因する。複雑な下流タスクにおいて、幅広い特徴スペクトルにアクセスできることで、モデルはより細やかな予測を行うことができ、汎化性能の向上につながる。

我々の推測を検証するため、Li et al. (2024)に従い、4.1節で得られたモデルをCommonsense170K (Hu et al., 2023)で微調整し、8つの下流タスクで評価を行った。表2に示すように、Mix-LNは評価されたすべてのデータセットにおいて、他の正規化技術を一貫して上回っている。LLaMA-250Mモデルでは、Mix-LNはPre-LNと比較して平均4.26%の大幅な向上を達成し、BoolQでは17.31%の改善を示している。同様の傾向がより大規模なLLaMA-1Bモデルでも観察された。Mix-LNはPre-LNと比較してパープレキシティをわずか0.25しか減少させていないにもかかわらず、教師あり微調整において大幅な性能向上をもたらしている。

一貫して、Mix-LNの利点はRLHFにもシームレスに転用できる。Adam-mini (Zhang et al., 2024)に従い、我々はInstructGPT (Ouyang et al., 2022)からRLHFワークフローを実装し、4.1節で得られた1Bモデルをultrafeedbackデータセットで訓練し、選好報酬を最適化した。表3は、Mix-LNがPre-LNに対して顕著な報酬の向上（高いほど良い）を達成していることを示している。すなわち、1.32対0.75である。

非言語モデルにおけるMix-LNの評価のため、我々はViTモデルのPre-LNを$\alpha=0.25$のMix-LNに置き換えた。更新されたモデルをLiu et al. (2022a)に従い、ConvNeXt (Liu et al., 2022b)の設定を用いてImageNet-1Kで120エポック訓練した。結果は、Mix-LNの利点が非言語モデルにも一般化されることを明確に示している。表4の結果は、Mix-LNの利点が非言語タスクにも及ぶことを示しており、小規模なモデル（ViT-Tiny）と比較して大規模なモデル（ViT-Small）でより顕著な性能向上が見られる。

Layer Normalization (LN)は、Ba (2016)によって最初に提案され、現代の言語モデルにおいて活性化を正規化するための事実上の標準となっている。これは隠れ層内のニューロンへの入力の合計から直接正規化統計を推定し、各層への入力分布が訓練全体を通じて安定していることを保証する。元のTransformer (Vaswani, 2017)では、LNは当初残差接続の後に適用されており、この構成はPost-LNとして知られている。しかし、その後の研究 (Baevski & Auli, 2018; Dai, 2019; Nguyen & Salazar, 2019)により、LayerNormを残差接続の前に配置する（Pre-LN）ことで、特に大規模言語モデルにおいてより安定したパフォーマンスが得られることが判明した (Brown, 2020; Touvron et al., 2023; Jiang et al., 2023)。 Xiong et al. (2020)は理論的に、Post-LNが出力層付近でより大きな勾配を生じさせ、不安定性を避けるためにウォームアップが不可欠であることを示した。対照的に、Pre-LNはモデルの深さに応じて勾配をスケールダウンし、初期化時により安定した勾配を確保する。我々の研究はXiong et al. (2020)の研究を基に、Pre-LNが勾配の大きさを減少させることで不安定性を防ぐ一方で、より深い層での小さな勾配が対応する重みの有効性を低下させる可能性があることを強調している。

言語モデルのより深い層の有効性を向上させるために、様々なLNの変種が提案されている。例えば、Wang et al. (2019)は、深いTransformerにおいてPost-LNが勾配消失に苦しむ一方で、Pre-LNがより多くの層の積み重ねを可能にすることを経験的に検証した。その結果、彼らは訓練可能性を向上させるために全ての前の層を接続する動的線形結合層（DLCL）を導入した。同様の技術が他の研究でも採用されている (Bapna et al., 2018; Dou et al., 2018)。Liu et al. (2020)は、Post-LNが残差分岐に強く依存し、しばしば不安定性につながることを明らかにした。これに対処するため、Adaptive Model Initialization (Admin)が導入された。これは追加のパラメータを使用してPost-LNの残差依存性を制御し、訓練を安定化させる。DeepNorm (Wang et al., 2024)は、LNを適用する前に残差接続をアップスケールすることで、深いTransformerの訓練可能性をさらに改善し、モデルの更新を減少させ、より深いアーキテクチャを可能にした。さらに、Ding et al. (2021)は、各Transformerサブレイヤーの入力と出力の両方を正規化するSandwich LayerNormを提案した。Takase et al. (2022)は、Post-LNがより深い層でより大きな勾配ノルムを保持する傾向があり、潜在的により効果的な訓練につながる可能性があることを特定した。初期層での勾配消失の問題に対処するため、彼らは各層の最後のLNを除くすべてのLNをバイパスする残差接続を使用するB2Tという方法を導入した。我々はTakase et al. (2022)からインスピレーションを得て、Pre-LNとPost-LNの両方を組み合わせることで、それぞれの限界に対処している。我々は付録BでScaled InitializationとScaled Embedを研究している。

LLMにおける深層の非効率性は、LLMの剪定の有効な指標となる。 Yin et al. (2023)は、LLaMAやMistralなどの著名なLLMの深層レイヤーが、浅層レイヤーよりも積極的に剪定できることを実証した。これは性能の大幅な低下を引き起こすことなく可能である。同様に、Gromov et al. (2024)とMen et al. (2024)はレイヤー剪定をさらに探究し、LLMの深層レイヤーが通常あまり重要でないことを特定した。Lad et al. (2024)は、PythiaやGPT-2などのモデルにおいて、深層レイヤーがレイヤーの削除や入れ替えなどの介入に対して強い耐性を示すことを観察した。我々の研究は、レイヤーの有効性を評価するために角度距離を適用する点でGromov et al. (2024)と類似している。しかし、彼らは深層レイヤーの非効率性を特定しているものの、この現象の説明や解決策を提案していない。

先行研究ではこれらの深層レイヤーの特性をモデル圧縮の機会として捉えることが多いが(Siddiqui et al., 2024; Zhong et al., 2024; Sreenivas et al., 2024)、我々はこの挙動がより深刻な学習の不足を示していると主張する。これは主にPre-LNの広範な使用によるものである。これに対応して、我々はMix-LNを導入する。これは深層レイヤーの有効性を高め、アーキテクチャ全体がより効果的に学習され、ネットワークのパラメータを十分に活用することを保証する新しい手法である。