技術要約: BHyT – Pre-Layer Normalizationに代わる安定かつ効率的な選択肢

1. 問題提起

大規模言語モデル（LLM）は、学習を安定させ効果的な転移学習を可能にするために、通常RMSNormで実装されるPre-Layer Normalization（Pre-LN）に大きく依存しています。しかし、モデルの深さが増すにつれ、標準的なPre-LN設計は2つの決定的な制限に直面します。

1. 深さの呪い（The Curse of Depth）: 近年の分析によれば、残差接続とPre-LNの相互作用により、層を追うごとに隠れ状態の大きさと分散が急速に増大することが示されています。このエスカレーションはブロック・ヤコビアンを単位行列へと押し進め、深い層がコストの高い恒等写像（identity mapping）のように振る舞う原因となり、学習を不安定にします。
2. 効率性と安定性のトレードオフ: 既存の解決策の多くは、どちらか一方の側面のみに対処しています。Peri-LN（サブレイヤーの前後両方で正規化を行う）やLayerNorm Scaling（LNS）のような安定性重視の手法は、分散の増大を抑制しますが、統計量の再計算による大幅な計算オーバーヘッドを招きます。対照的に、Dynamic Tanh（DyT）のような効率重視の正規化フリー手法は、統計的オーバーヘッドを排除しますが、深さ方向への残差増大を明示的に制御できず、深いネットワークにおける飽和や不安定性を引き起こす可能性があります。

中心となる課題は、有界な変換による低オーバーヘッドと、残差ストリームのスケールおよび分散に対する明示的な制御を組み合わせ、スループットを犠牲にすることなく、深さ方向のドリフトを防ぐメカニメントを設計することです。

2. 手法: Bounded Hyperbolic Tanh (BHyT)

著者らは、双曲線正接（hyperbolic tangent）の非線形性と、データ駆動型の入力境界設定を組み合わせた、Pre-LNのドロップイン・リプレースメントとしてのBHyTを提案しています。

コアメカニズム

学習可能なグローバル・スカラーを用いて入力を再スケールするDyTとは異なり、BHyTはデータ依存の統計量を用いて$\tanh(\cdot)$関数への入力を明示的に制限します。

• 確率的入力境界設定: 飽和を防ぐため、BHyTは$\tanh$の前段階の引数が、高い確率で定義された区間 $[-\lambda, \lambda]$ 内に収まるよう制約をかけます。チェビシェフの不等式を用いることで、入力の平均 ($\mu_x$) と標準偏差 ($s_x$) に基づくスケーリング規則を導出します。
• 理想的な形式 ($BHyT^*$): 理論上、この変換は入力を次のように再スケールします：
  $$BHyT^*(x) = \gamma \odot \tanh\left(\frac{\lambda}{\kappa s_x + |\mu_x|} x\right)$$
  ここで、$\gamma$ は学習可能なスケールパラメータ、$\lambda$ は目標範囲、$\kappa$ は目標確率 $p$ に基づくスケーリングの保守性を制御するパラメータです。これにより、入力が高確率で$\tanh$の非飽和な線形領域内に留まることが保証されます。

実装上の工夫: 分散近似

Transformerブロック内の各Pre-LNサイトで厳密な平均と分散を計算することは、効率性のメリットを打ち消してしまいます。これに対処するため、BHyTはブロックレベルの分散近似を導入しています。

1. 第1サイト（Attentionの前）: 入力の正確なインスタンスごとの分散を計算します。
2. 第2サイト（MLPの前）: 残差出力 ($x + h_{Attn}$) の正確な分散を再計算する代わりに、先に計算された入力の分散と、モデル依存のAttention出力の分散推定値を用いて近似を行います。
   • Attention出力の分散は、一様なAttention重みとゼロ平均のガウス分布成分を仮定し、$\tilde{s}^2_{hAttn} \approx \frac{1}{Td} \|W_V W_O\|_F^2 \cdot \frac{\lambda_{Attn}^2}{\kappa^2}$ として近似されます。
   • この近似により、第2の正規化処理において統計量を再利用または事前計算されたモデルパラメータを使用することが可能になり、クリティカルなフォワードパスにおけるセカンド・リダクション操作を回避できます。

理論的安定性

本論文は、有限深さの分散伝播境界を提供しています。定理3.5によれば、ネットワークの深さを $L$ としたとき、ハイパーパラメータが $\lambda/\kappa < 1/\sqrt{L}$ を満たす場合、BHyT下での出力分散は、深さ $L$ までのすべての層においてLNSよりも厳密に小さくなります。これは、BHyTがLNSよりも効果的に深さ方向の分散増大を抑制するという理論的保証を与えています。

3. 主な貢献

1. BHyTの提案: 深さ方向の活性化増大を制御するために、有界な$\tanh(\cdot)$変換とデータ駆動型の入力境界設定を組み合わせた、Pre-LNのドロップイン・リプレースメントを提案。
2. 理論的および実用的な効率性:
   • BHyTが特定の条件下でLNSよりも小さな分散伝播境界を持つことを証明する、有限深さの分散境界の導出。
   • 繰り返しの正確な分散計算を減らし、レイテンシを大幅に低下させるブロックレベルの分散近似の導入。
3. 実証的検証: Llama-374M、Llama-1B、Llama-3Bモデルを用いた包括的な評価。BHyTが安定性と効率性のトレードオフを改善し、強力なダウンストリーム性能を維持することを実証。

4. 実験結果

実験は、C4コーパスを用いて（モデルサイズに比例したトークン予算、約20 tokens per parameter）、Llamaスタイルのモデル（374M、1B、3B）に対して行われました。

• 安定性: BHyTは、活性化の大きさおよび分散における深さ方向の増大を抑制することに成功し、RMSNormおよびDyTを上回りました。活性化の統計量は、Peri-LNのような安定性重視のベースラインと同等のレベルで、層全体にわたって制御されていました。
• 性能:
  • 事前学習: BHyTは、すべてのモデルスケールにおいて、最も低い事前学習損失とパープレキシティを達成しました。また、3-shotベンチマーク（ARC、MMLU、Winograndeなど）において、最も高い平均ダウンストリーム精度を記録し、1Bモデルでは2番目に優れた手法を最大0.6ポイント上回りました。
  • 事後SFT: Lima1Kを用いた教師あり微調整（SFT）後、BHyTモデルはRMSNorm、Peri-LN、LNS、DyTと比較して、C4およびWikiText-2におけるパープレキシティが低く（忘却が少なく汎化性能が高いことを示す）、より高いダウンストリーム精度を示しました。
  • スケール: Llama-3Bにおける20Bトークンの比較において、BHyTは強力な安定性重視のベースラインであるPeri-LNに対し、平均ダウンストリーム精度（事前学習）で3.4ポイント、事後SFTで2.5ポイント上回りました。
• 効率性:
  • 学習スループット: BHyTは、RMSNormと比較して平均1.6%高速な学習を実現しました。
  • 生成スループット: BHyTはLlama-1Bにおいて1,199.9 tokens/sを達成し、RMSNormを1.77%、Peri-LNを18.0%上回りました。DyTは（統計量がないため）依然として最速でしたが、BHyTは速度と安定性の間でより優れたトレードオフを提供しました。
  • アブレーション研究: 分散近似メカニズムにより、厳密な形式のBHyT*と比較して14.9%のスピードアップを実現しましたが、収束性やダウンストリーム性能を実質的に損なうことはありませんでした。

5. 意義と主張

本論文は、BHyTがより深く、より効率的な基盤モデルへの実用的な道筋を提供するものであると主張しています。有界な変換の低オーバーヘッドと、残差ストリームの分散に対する明示的な制御を統合することで、BHyTは二重正規化スキームによる計算ペナルティを課すことなく、「深さの呪い」に対処します。

著者らは、スケーラブルな学習は安定性を犠牲にすることなく効率を向上させることができると断言しています。BHyTは単なる理論的な代替案ではなく、訓練と推論の両方でレイテンシを削減しながら、安定した事前学習と強力な転移性能を可能にする、RMSNormに対する堅牢なドロップイン・リプレースメントとして提示されています。