Learning Rate Transfer in Normalized Transformers

以下は、Shigida、Hanin、Gromov による論文「Learning Rate Transfer in Normalized Transformers」の詳細な技術的サマリーです。

1. 問題定義

本論文は、Loshchilov ら (2025) によって導入された Transformer アーキテクチャの変種であるNormalized Transformers (nGPT) におけるハイパーパラメータ (HP) 転送の課題に取り組んでいます。

• 背景: nGPT は、重みとアクティベーションを大規模に正規化すること（LayerNorm/RMSNorm の代わりに $\ell_2$ ノルムを使用）により、大幅なトレーニング速度向上を実現し、重み減衰と学習率ウォームアップの必要性を排除しています。
• 課題: nGPT にはモデルサイズに比例するハイパーパラメータが存在しますが、著者らはモデルの幅（次元）、深さ（層数）、およびトークンホライズン（トレーニング反復回数）にわたって学習率転送が機能しないことを観察しています。
• 既存手法の限界: $\mu$P (Maximal Update Parametrization) などの標準的な転送手法は、重みとアクティベーションの整列に関する特定の理論的仮定（具体的には、それらが特定の方法で完全に整列している、あるいは整列していないこと）に依存しています。著者らは、これらの仮定が nGPT には当てはまらないことを発見しており、スケーリングアップ時に最適でないパフォーマンスをもたらしています。

2. 手法

著者らは、$\nu$GPT (nu-GPT) と呼ばれる新しいパラメータ化を提案しています。彼らのアプローチは、整列指数 (alignment exponents, Everett ら 2024) に基づく原理的な理論的枠組みと数値実験を組み合わせています。

A. 理論的枠組み：整列指数

$\mu$P で使用される「完全整列」($\alpha=1$) または「非整列」($\alpha=1/2$) のレジームを仮定する代わりに、著者らは重み更新とアクティベーションがモデル次元 $d_{model}$ に対してどのようにスケーリングするかを記述する整列指数 ($\alpha, \omega, \nu$) を経験的に測定します。

• 行列 - ベクトル積の変化がどのようにスケーリングするかを定義する指数：$\|\Delta W h\| \sim d^{\alpha} \|\Delta W\| \|h\|$。
• 主要な発見: nGPT においては、「中間整列」レジームが観測され、$\max\{\alpha, \nu\} \approx 3/4$ となります。これは、$\mu$P の仮定である $\max\{\alpha, \nu\} = 1$（完全整列）または $1/2$（非整列）と矛盾します。

B. $\nu$GPT パラメータ化

「中間整列」仮説 ($\max\{\alpha, \nu\} = 3/4$) と経験的なスケーリング則に基づき、著者らは学習率 ($\eta$) と初期化パラメータに対する具体的なスケーリング則を導出しました。$m_{width}$、$m_{depth}$、および $m_{data}$ を、基準モデルに対する幅、深さ、およびトークン数の乗数とします。

$\nu$GPT における主要なスケーリング則:

1. グローバル学習率: $m_{data}^{-1/3}$ としてスケーリングします。これはトークンホライズンを補正するもので、一部の先行する SDE ベースの理論が示唆する $m_{data}^{-1/2}$ とは異なります。
2. 埋め込み重み ($\eta_{input}$): $m_{width}^{-1/2}$ としてスケーリングします。
3. 隠れ層/アンエンベディング重み ($\eta_{hidden}, \eta_{output}$): $m_{width}^{-3/4}$ としてスケーリングします。これは $\mu$P（バリアントによって $m_{width}^{-1}$ または $m_{width}^{-1/2}$ を使用）からの重要な逸脱です。
4. 深さ補正:
   • 学習可能な再スケーリングベクトル $\alpha_{A, init}$ および $\alpha_{M, init}$ は $m_{depth}^{-1}$ によってスケーリングされます。
   • 「Depth-$\mu$P」（隠れ層の学習率を $m_{depth}^{-1/2}$ でスケーリングする）とは異なり、$\nu$GPT は、学習可能な再スケーラーが深さの安定性を自然に処理するため、nGPT において隠れ層の学習率を深さでスケーリングすべきではないと結論付けています。
5. 初期化定数: スケールパラメータ ($\alpha_{scale}$) に対しては、$d_{model}^{-1/2}$ でスケーリングするのではなく、固定定数（例：0.03）が使用されます。

3. 主要な貢献

1. nGPT における $\mu$P の失敗の特定: 重み - アクティベーションの整列に関する標準的な $\mu$P の仮定が正規化アーキテクチャには無効であり、学習率転送の不良につながることを実証しました。
2. $\nu$GPT の提案: nGPT で観測された「中間整列」レジーム（$\approx 3/4$）を明示的にターゲットとする新しいパラメータ化を導入しました。
3. トークンホライズンスケーリング: nGPT における最適な学習率がトークン数に対して $T^{-1/3}$ としてスケーリングすることを確立し、以前の理論的予測である $T^{-1/2}$ を精緻化しました。
4. 深さ転送の洞察: 学習可能な再スケーラーを備えた正規化モデルでは、隠れ層学習率の明示的な深さスケーリングは不要であり、潜在的に有害であることを示しました。その代わり、再スケーラーの初期化をスケーリングすることで十分です。

4. 実験結果

著者らは、OLMo 2 トークナイザーを使用して FineWeb-Edu データセット上で $\nu$GPT を検証しました。

• 幅転送:
  • 基準 nGPT: 学習率転送を示さず、幅が増加するにつれてパフォーマンスが低下しました。
  • $\mu$P 風補正: 最適学習率にわずかなドリフトが見られました。
  • $\nu$GPT: 0.26B から 3.22B パラメータにわたる幅全体で完璧な学習率転送を達成し、チューニング済みの基準モデルと比較して検証損失の損失はありませんでした。
• 深さ転送:
  • 基準 nGPT はまともな深さ転送を示しましたが、極端な深さでは不安定になりました。
  • $\nu$GPT は 8 から 128 層にわたる深さ全体で安定性とパフォーマンスを維持し、大規模な深さにおける学習率の選択に対する感度が低いことを示しました。
• トークンホライズン転送:
  • 実験により、トレーニング反復回数の増加に伴う最適学習率の $T^{-1/3}$ スケーリング則が確認されました。
• 整列指数の検証:
  • 経験的な測定により、$\omega_{hidden} = \omega_{output} = 1/2$ かつ $\max\{\alpha, \nu\} \approx 0.75$（中間整列）であることが確認され、$\nu$GPT の理論的基盤が検証されました。

5. 意義

• 実用的影響: $\nu$GPT を用いることで、研究者は小規模モデルでハイパーパラメータを調整するだけで、それを直接大規模スケールに転送して大規模な Normalized Transformers をトレーニングできるようになり、計算リソースを大幅に節約できます。
• 理論的進展: 本論文は $\mu$P の普遍性に疑問を呈し、異なるアーキテクチャの選択（特に正規化スキーム）が異なる整列仮定とスケーリング則を必要とすることを示しました。
• 効率性: 重み減衰とウォームアップの必要性を排除しつつ、堅牢なスケーリングを可能にすることで、$\nu$GPT は次世代の基盤モデルのためのより簡素化され効率的なトレーニングパイプラインを提供します。

要約すると、本論文は、正規化トランスフォーマーにおけるスケーリング問題に対する厳密で経験的に裏付けられた解決策を提供し、「万能」な $\mu$P アプローチを、正規化アーキテクチャの独自のダイナミクスを尊重するカスタマイズされたパラメータ化（$\nu$GPT）に置き換えています。