2025-11-12T01:28:29.133817

Stability of Transformers under Layer Normalization

Kan, Li, Zhang et al.
Despite their widespread use, training deep Transformers can be unstable. Layer normalization, a standard component, improves training stability, but its placement has often been ad-hoc. In this paper, we conduct a principled study on the forward (hidden states) and backward (gradient) stability of Transformers under different layer normalization placements. Our theory provides key insights into the training dynamics: whether training drives Transformers toward regular solutions or pathological behaviors. For forward stability, we derive explicit bounds on the growth of hidden states in trained Transformers. For backward stability, we analyze how layer normalization affects the backpropagation of gradients, thereby explaining the training dynamics of each layer normalization placement. Our analysis also guides the scaling of residual steps in Transformer blocks, where appropriate choices can further improve stability and performance. Our numerical results corroborate our theoretical findings. Beyond these results, our framework provides a principled way to sanity-check the stability of Transformers under new architectural modifications, offering guidance for future designs.
academic

レイヤー正規化下のTransformerの安定性

基本情報

  • 論文ID: 2510.09904
  • タイトル: Stability of Transformers under Layer Normalization
  • 著者: Kelvin Kan (UCLA)、Xingjian Li (UT Austin)、Benjamin J. Zhang (UNC Chapel Hill)、Tuhin Sahai (SRI International)、Stanley Osher (UCLA)、Krishna Kumar (UT Austin)、Markos A. Katsoulakis (UMass Amherst)
  • 分類: cs.LG、cs.AI、math.OC
  • 発表日: 2025年10月10日
  • 論文リンク: https://arxiv.org/abs/2510.09904

要約

Transformerは広く使用されていますが、深いTransformerの訓練は不安定になる可能性があります。レイヤー正規化(Layer Normalization)は標準的なコンポーネントとして訓練の安定性を改善しますが、その位置の選択はしばしば恣意的です。本論文は、異なるレイヤー正規化位置下でのTransformerの前向き安定性(隠れ状態)と後向き安定性(勾配)について、原理的な研究を行います。理論分析は訓練動力学の重要な洞察を明らかにします：訓練がTransformerを規則的な解へ導くのか、それとも病的な振る舞いへ導くのかです。前向き安定性の観点から、訓練後のTransformerにおける隠れ状態増加の明示的な界を導出しました。後向き安定性の観点から、レイヤー正規化が勾配逆伝播にどのように影響するかを分析し、各レイヤー正規化位置の訓練動力学を説明しました。分析はまた、Transformerブロック内の残差ステップサイズのスケーリングを指導し、適切な選択によってさらに安定性とパフォーマンスを改善できます。

研究背景と動機

問題定義

本研究が解決する核心的な問題は：異なるレイヤー正規化位置がTransformer訓練安定性に及ぼす影響メカニズムです。具体的には：

  1. 前向き安定性問題：深いネットワークにおける隠れ状態の増加制御
  2. 後向き安定性問題：逆伝播過程における勾配の安定性
  3. アーキテクチャ設計ガイダンス：新しいTransformer変種に対する理論的指導の提供

重要性分析

  1. 実用的価値：Transformerは現代的な深層学習の基礎であり、その訓練安定性はモデルパフォーマンスと訓練効率に直接影響します
  2. 理論的欠落：既存のレイヤー正規化位置の選択は主に経験的であり、理論的根拠が不足しています
  3. 産業的需要：モデル規模の拡大に伴い、訓練安定性の問題がますます顕著になっています

既存手法の限界

  1. Post-LN：細かい最適化スケジュールが必要であり、パフォーマンスはしばしば次善的です
  2. Pre-LN：初期訓練安定性は改善されますが、過度に大きな隠れ状態を生成し、数値不安定性につながります
  3. Peri-LN：実践では良好なパフォーマンスを示しますが、理論的性質の理解が不足しています

研究動機

著者は連続時間動力学と最適制御理論の新しい視点を採用し、Transformer訓練問題を平均場制御問題としてモデル化することで、初期化時の振る舞いだけでなく、訓練収束後のモデル性質を分析できるようにしました。

核心的貢献

  1. 理論的フレームワークの革新：最適制御理論に基づく新しいフレームワークを提案し、異なるレイヤー正規化位置下でのTransformer安定性を体系的に分析
  2. 前向き安定性分析：隠れ状態増加の明示的な界を導出し、Pre-LNが無界増加をもたらすことを証明し、Peri-LNが制御された増加を維持することを証明
  3. 後向き安定性分析：レイヤー正規化が勾配逆伝播に及ぼす影響メカニズムを明らかに
  4. 残差ステップサイズスケーリング：安定性とパフォーマンスを改善する残差ステップサイズスケーリング方法を提案
  5. 実験的検証：GPT-2シリーズモデルで理論的発見を検証

方法論の詳細

タスク定義

異なるレイヤー正規化位置下でのTransformerの安定性を研究し、以下を含みます：

  • 入力：埋め込みと位置エンコーディング後のシーケンス X0∈Rd×nX_0 \in \mathbb{R}^{d \times n}
  • 出力：D層のTransformerブロックを通過後の隠れ状態 XDX_D
  • 目標：前向きおよび後向き伝播の安定性を分析

連続時間モデリング

Transformerの連続時間表現

標準的なTransformerブロックのスキップ接続構造をオイラー離散化の連続時間動力学として解釈：

f_{attn}(X(t), t; \theta_{attn}(t)), & t \in [t_i, t_i + \Delta t) \\ f_{ffn}(X(t), t; \theta_{ffn}(t)), & t \in [t_i + \Delta t, t_{i+1}) \end{cases}$$ ここで $\Delta t = \frac{T}{2D}$、$t_i = 2i\Delta t$です。 #### 平均場制御問題としてのモデリング 訓練問題を連続時間平均場制御問題として定式化： $$\min_\theta \mathbb{E}_{(X_0,y)} G(X(T), y)$$ $$\text{s.t. } \frac{dX(t)}{dt} = f(X(t), t; \theta(t))$$ ここで $f \in \{f_{Pre}, f_{Peri}\}$ は異なるレイヤー正規化位置に対応します。 ### レイヤー正規化の幾何学的性質 **重要補題1**：レイヤー正規化出力は楕円面上に位置 $$\mathcal{E} = \{z \in \mathbb{R}^d : (z - \beta)^T\Gamma^{-2}(z - \beta) = d\}$$ ここで $\Gamma = \text{diag}(\gamma)$です。 ### 前向き安定性分析 #### Pre-LNの無界性 **定理2**：Pre-LN訓練問題の最適解は大きさにおいて無界です。 **証明の概要**：Hamilton-Jacobi-Bellman (HJB) 偏微分方程式を分析することで、対応するハミルトニアンが存在しないことを証明し、訓練問題が退化することを示します。 **定理3**：重み減衰を使用しても、Pre-LN Transformerの隠れ状態は指数関数的に増加： $$MA(X_D) \leq (1 + C(\lambda))^D \frac{\|X_0\|_F}{\sqrt{nd}} = O(e^D)$$ #### Peri-LNの制御された増加 **定理4**：Peri-LN Transformerの隠れ状態は線形に増加： $$MA(X_D) \leq \frac{\|X_0\|_F}{\sqrt{nd}} + 2D(\gamma_{max} + \beta_{max}) = O(D)$$ 分散は二次に増加： $$\text{Var}(X_D) \leq \frac{(\|X_0\|_F + 2D\sqrt{nd}(\gamma_{max} + \beta_{max}))^2}{nd - 1} = O(D^2)$$ ### 後向き安定性分析 勾配計算式： $$\nabla_{\theta_i} G(X_D) = \nabla_{\theta_i} X_{i+1} \cdot J_{i:D} \cdot \nabla_{X_D} G(X_D)$$ ここでヤコビ行列： $$J_{i:D} = \prod_{j=i+1}^D (I + \nabla_{X_{j-1}} f(X_{j-1}; \theta_{j-1}))$$ **命題7**：Pre-LN下では、感度 $\nabla_{X_{j-1}} f_{Pre}$ は活性化値に比例して増加します。 **命題8**：Peri-LN下では、感度 $\nabla_{X_{j-1}} f_{Peri}$ は活性化値の大きさに対して不変です。 ## 実験設定 ### データセット - **OpenWebTextデータセット**：約90億の訓練トークン、400万の検証トークン - GPT-2シリーズアーキテクチャを使用した事前訓練 ### モデル構成 - **GPT-2** (1億2400万パラメータ) - **GPT-2 Large** (7億7400万パラメータ) - **GPT-2 XL** (15億パラメータ) ### 評価指標 - **困惑度 (Perplexity)** - **ROUGEスコア** (Rouge1、Rouge2、RougeL) - **BERTスコア** (BertP、BertR、BertF1) - **訓練安定性**：発散実行回数の統計 ### 実装の詳細 - Pre-LNに対して調整されたハイパーパラメータを使用し、Peri-LNに対して個別の最適化は行いません - 残差ステップサイズスケーリング：$\Delta t \in \{0.1, 1\}$ - ハードウェア：NVIDIA H200 GPU ## 実験結果 ### 訓練安定性の比較 | レイヤー正規化設定 | 重み減衰有効 | 重み減衰無効 | |------------------|-----------|-----------| | Pre-LN | 1/5発散 | 3/5発散 | | Peri-LN | 0/5発散 | 0/5発散 | | LNなし | 5/5発散 | — | ### パフォーマンス比較結果 **GPT-2 (1億2400万)モデル結果**： - Pre-LN ($\Delta t=1$): 検証損失5.43、困惑度247.52 - Pre-LN ($\Delta t=0.1$): 検証損失3.13、困惑度24.43 - Peri-LN ($\Delta t=1$): 検証損失3.12、困惑度24.17 - Peri-LN ($\Delta t=0.1$): 検証損失3.10、困惑度23.63 ### 隠れ状態増加分析 実験は理論的予測を検証しました： - Pre-LNはより大きな$\Delta t$で急速な増加を示す - Peri-LNはより規則的な線形増加を維持 - 残差ステップサイズスケーリングは増加率を効果的に制御 ### 残差ステップサイズスケーリング効果 1. **パフォーマンス向上**：Peri-LN + $\Delta t=0.1$ はすべての指標で最高のパフォーマンスを示す 2. **安定性改善**：Pre-LNは$\Delta t=0.1$で不安定から安定に変化 3. **増加制御**：隠れ状態の平均値と分散の増加率を効果的に低減 ## 関連研究 ### レイヤー正規化研究 - **Post-LN**：最初のTransformer設計で、細かいスケジュール調整が必要 - **Pre-LN**：訓練安定性を改善するが、大きな活性化値を生成 - **Peri-LN**：Gemma2、OLMo2などの大規模モデルで最近採用 ### 理論分析手法 - 既存の研究は初期化時の振る舞いに焦点を当てるか、経験的観察に依存 - 本論文は訓練収束後のモデル性質を革新的に分析 - 連続時間視点はアーキテクチャ分析に新しいツールを提供 ## 結論と考察 ### 主要な結論 1. **Pre-LN理論的欠陥**：訓練問題は本質的に病的であり、無界解をもたらす 2. **Peri-LN利点**：良好に定義された最適化問題と制御された隠れ状態増加を提供 3. **残差スケーリング価値**：シンプルで効果的な安定性改善方法 ### 限界 1. **簡略化された仮定**：理論分析は連続時間近似に基づく 2. **ハイパーパラメータ依存**：実験はPre-LN調整のハイパーパラメータを使用 3. **規模制限**：実験は主に中規模モデルで実施 ### 将来の方向 1. **アーキテクチャ選別フレームワーク**：新しいアーキテクチャ修正に対する理論的選別基準を提供 2. **より大規模な検証**：より大きなモデルで理論的発見を検証 3. **他の正規化方法**：RMSNormなどの変種への分析拡張 ## 深い評価 ### 利点 1. **理論的革新性が強い**：最適制御理論を使用してレイヤー正規化位置問題を初めて分析 2. **数学的厳密性**：完全な理論導出と証明を提供 3. **実用的価値が高い**：残差ステップサイズスケーリング方法はシンプルで効果的 4. **実験設計が合理的**：複数のモデル規模で理論を検証 ### 不足 1. **理論と実践のギャップ**：連続時間仮定と実際の離散実装の間に差異が存在 2. **実験範囲が限定的**：主にGPT-2シリーズで検証され、より多くのアーキテクチャ検証が不足 3. **ハイパーパラメータの公平性**：Peri-LNに対する専門的なハイパーパラメータ最適化が実施されていない ### 影響力評価 1. **学術的貢献**：Transformer安定性分析に新しい理論的フレームワークを提供 2. **実用的価値**：実際のモデル設計と訓練戦略を指導 3. **再現性**：コードとモデルの公開を約束 ### 適用シナリオ 1. **深いTransformer訓練**：特に大規模な深層モデルに適用 2. **アーキテクチャ設計ガイダンス**：新しいアーキテクチャ修正に対する理論的根拠を提供 3. **訓練安定性改善**：残差スケーリングを通じて訓練安定性を向上 ## 参考文献 論文は複数の重要な研究を引用しており、以下を含みます： - Ba et al. (2016): レイヤー正規化の原始論文 - Xiong et al. (2020): Pre-LN対Post-LN比較研究 - Kim et al. (2025): Peri-LNの経験的研究 - He et al. (2016): 残差接続の開拓的研究 --- **総合評価**：これは理論と実践が良好に結合された高品質な論文であり、Transformer安定性分析に新しい数学的フレームワークを提供し、重要な学術的価値と実用的意義を持っています。理論分析は厳密で深く、実験検証は十分であり、深層学習アーキテクチャ設計に価値のあるガイダンスを提供します。