muP : unit-scalingを搭載して最強になろう

Unit-Scaling μP（μPスケーリング）の理論と実践【拡充版】

大規模モデルを効率的にチューニングするための先端的パラメータスケーリング手法、Unit-Scaling μP (u-μP) を解説します。
本記事では、以下の流れで理解を深めていきます。

1. 背景：Maximal Update Parametrization (μP) とは
2. μPの理論：数式を交えた詳細
3. Unit Scaling の概念とう-μpへの発展
4. 数式で見るu-μP
5. 適用可能なモデルと応用例
6. 実験結果と性能分析
7. 結論と今後の展望
8. 参考リンク

---

背景：Maximal Update Parametrization (μP) とは

近年のディープラーニングにおいて、大規模化とともにハイパーパラメータ（学習率や初期化など）の再調整が必要となる問題が顕在化しています。とくに、隠れ次元（幅）やチャネル数を変えるたびに最適学習率を小さくしないと学習が破綻するといった経験則は、多くの研究者が直面している課題です。

Maximal Update Parametrization (μP) は、この問題に対処すべく Microsoft Research の Yang らによって提案されました。目的は次の3点です。

1. 幅を増やしても学習動態が崩れない。
2. 小規模モデルで見つけたハイパーパラメータを大規模モデルに転用できる。（ゼロショット転移）
3. 無限幅でも特徴学習が十分起きるようにする。（NTKパラメータ化のように線形化しない）

実際の応用としては、たとえば13Mパラメータ規模でハイパーパラメータをチューニングし、それを3.5億パラメータのモデルに再調整なしで適用して高精度が得られるなど、“μTransfer” と呼ばれる活用事例が報告されています。

---

μPの理論：数式を交えた詳細

1. 標準スケーリングとNTKスケーリング

まず、最も基本的な全結合層（隠れ次元 $nn$）を考えます。重み行列 $W\mathbf{W}$ の初期化や学習率が幅 $nn$ によってどう変わるか整理してみましょう。

• 標準スケーリング（SP）
  一般的なHe初期化やXavier初期化を想定します。たとえばHe初期化では、
  

  $$Var(Wij)=2n.\mathrm{Var}(W_{ij}) = \frac{2}{n}.$$
  
  これは出力の分散が大きくなりすぎないようにするためです。一方で、学習率 $η\eta$ は実験的にチューニングされるため、モデル幅が大きくなるほど最適 $η\eta$ が小さくなっていくという傾向がしばしば報告されています。

• NTKパラメータ化
  Neural Tangent Kernel (NTK) 研究で導入されたもので、幅 $n→∞n \to \infty$ の極限でニューラルネットが初期時点の線形カーネルとして振る舞うようパラメータを設定します。重み初期化は標準スケーリングと同程度にしつつ、学習率 $η\eta$ を
  

  $$η≈O(1n)\eta \approx O\Bigl(\frac{1}{n}\Bigr)$$
  
  として極端に小さくします。こうすると「学習中に重みがほとんど動かない（lazy training）」になり、線形近似的に解析できるのがメリットですが、特徴表現の更新がほぼ起きないため、幅を増やしても性能向上は限定的になることが多いです。

2. μPのスケーリング

これら2つの方法の中間的アプローチとしてμPが位置づけられます。直感的には、「モデル幅 $nn$ を増やしても各層の出力や勾配のスケールが一定に保たれ、かつ重みの更新が無視できない程度に大きく残る」設定を探ります。

例えば2層ネットワークを

$$fα(x)=α hL(x)f_\alpha(\mathbf{x}) = \alpha \, h^L(\mathbf{x})$$

と書き、幅 $nn$ が大きいときの重み更新量 $ΔW\Delta W$ を調べると、NTKスケーリングでは

$$ΔW∼O(1n),\Delta W \sim O\Bigl(\frac{1}{n}\Bigr),$$

標準スケーリングでは

$$ΔW∼O(1n)\Delta W \sim O\Bigl(\frac{1}{\sqrt{n}}\Bigr)$$

となるのが典型です。NTKの場合は小さすぎて「ほとんど更新しない」極限に陥りますが、μP では幅に依存しない一定オーダーの学習率を使いつつ、重み初期化スケールなどを少し調整することで

$$ΔW∼O(1n)\Delta W \sim O\Bigl(\frac{1}{\sqrt{n}}\Bigr)$$

を保ち、無限幅でも特徴学習が起きるリッチな極限を実現します。

論文では「abcパラメータ化」と呼ばれるテンソル単位のスケーリング定数を導入し、各レイヤーで"出力の分散" や "勾配の分散" が幅 $nn$ に依存しないように設計することが提案されています。

---

Unit Scaling の概念とう-μPへの発展

1. Unit Scaling とは

Unit Scaling は、モデルの各テンソル（活性、重み、勾配など）の分散を常に“1” 程度のユニットスケールに保つための設計手法です。
特に、FP16 や FP8 のように数値のダイナミックレンジが狭い精度で学習する際、これを守るとオーバーフローやアンダーフローが劇的に起きにくくなるため、大規模モデルの省メモリ学習や分散学習を安定化できます。

たとえば、全結合層であれば「重み行列の出力直後に $1/n1/\sqrt{n}$ を掛ける」といった具合に、層間で信号の分散が常に 1 になるよう各種スケーリング因子を挿入します。

2. Unit-Scaling μP (u-μP)

u-μP は、この Unit Scaling と μP を統合したもので、以下のようなメリットを併せ持ちます。

• モデル幅（隠れ次元）を変えても最適ハイパーパラメータが不変（μPの利点）。
• 全てのテンソルのRMS値が 1 程度に保たれ、低精度学習が安定（Unit Scalingの利点）。

実装的には、μPで層ごとに設定していた学習率スケーリングや初期化ルールを、“最初から出力が 1 のスケールになる” 形で整理しており、ユーザは「ベースモデルの幅」を明示的に意識しなくても適切なスケーリングがなされるようになっています。

---

数式で見るu-μP

ここではもう少し形式的に、u-μPの要点を数式で示します。

1. 初期化スケーリング

幅（チャネル数）を表す $dℓd_\ell$ を層 $ℓ\ell$ の出力次元とします。
重み行列 $W(ℓ)∈Rdℓ×dℓ−1\mathbf{W}^{(\ell)} \in \mathbb{R}^{d_{\ell} \times d_{\ell-1}}$ の要素を

$$Wij(ℓ)∼N(0,1dℓ−1).W^{(\ell)}_{ij} \sim \mathcal{N}\Bigl(0, \frac{1}{d_{\ell-1}}\Bigr).$$

ただし、層 $ℓ\ell$ の出力に対してスケーリング定数 $α(ℓ)\alpha^{(\ell)}$ を掛けることで、初期時点で活性が分散1 となるように調整します。たとえば、

$$h(ℓ)(x)=α(ℓ) σ(W(ℓ)h(ℓ−1)(x)),h^{(\ell)}(\mathbf{x}) = \alpha^{(\ell)} \,\sigma\Bigl(\mathbf{W}^{(\ell)} h^{(\ell-1)}(\mathbf{x})\Bigr),$$

とすることで、$Var[h(ℓ)]≈1\mathrm{Var}[h^{(\ell)}] \approx 1$ を満たすように $α(ℓ)\alpha^{(\ell)}$ を設定します。
（$σ(⋅)\sigma(\cdot)$ は適当な非線形活性、たとえばReLUなど）

2. 学習率スケーリング

NTKパラメータ化では $η∼O(1/dℓ)\eta \sim O(1/d_\ell)$ が推奨される一方、μP では「重みが更新されやすい」よう

$$η∼O(1dℓ)\eta \sim O\bigl(\frac{1}{\sqrt{d_\ell}}\bigr)$$

のオーダーを保つことがしばしば用いられます。u-μPではさらに全テンソルが単位スケールになるよう、オプティマイザ内部で層ごとにスケール調整を行います。
実際には以下のような実装が存在しています。

• MuAdam / MuSGD : 学習率を自動でスケーリングし、ベース幅を意識せずに済む仕組み（microsoft/mup 参照）。

3. abcパラメータ化の単純化

μP論文で紹介された「abcパラメータ化」では、各レイヤーに対し初期化用・前向き伝搬用・勾配計算用の係数 ($aa$, $bb$, $cc$) が存在し、組み合わせることでNTK極限やμP極限を表現できます。
u-μPでは、Unit Scalingの考え方で「最初から分散1を維持する」ための定数が統合され、実際のコードではほぼ “定数1” だけで済む層も多いとのことです。これにより、ユーザ視点では大幅にスケーリング係数が減り、実装が簡素になります。

---

適用可能なモデルと応用例

1. Transformerモデル

• 埋め込み（embedding）: ベクトルの次元 $dmodeld_{\text{model}}$ が幅に相当する場合、入力層は固定次元からの射影になるため、ここをどうスケーリングするかが重要です。u-μPでは埋め込み重みに対して
  
  $$Var(Wembed)∝1fan-in\mathrm{Var}(W_{\text{embed}}) \propto \frac{1}{\mathrm{fan\text{-}in}}$$
  
  かつ出力の活性にスケーリング定数を掛けます。
• 自己注意機構: マルチヘッドのキー・クエリ・バリュー射影やスケーリング定数 $1/dhead1/\sqrt{d_{\text{head}}}$ は標準Transformerと同様ですが、Residual接続（Add & Norm）の部分に追加のスケーリングが施される場合があります。u-μP実装では Residualブロックをユニットスケールに保つために、ブロック出力を $1/depth1/\sqrt{\text{depth}}$ で割る工夫を入れることも。
• 出力層（unembedding）: 埋め込みと重みを共有する場合 (tied embedding) もあり、そこにも u-μPのスケーリング則を適用。小さいモデル (例: $dmodel=256d_{\text{model}}=256$) の最適設定が、大きいモデル ($dmodel=4096d_{\text{model}}=4096$) でもそのまま最適化に近いと報告されています。

2. CNN/ResNet系

• チャネル数が “幅” に該当します。
• Convolutionフィルタも $1/din1/d_{\mathrm{in}}$ で初期化し、出力を $1/dout1/\sqrt{d_{\mathrm{out}}}$ で正規化するなどの工夫で活性の分散を制御。
• Residual接続の出力を足し合わせる前に $1/depth1/\sqrt{\text{depth}}$ で割るなど、深さ方向のスケーリングも考慮すると、幅を増減させても安定した学習が得られます。

3. 大規模言語モデル（LLM）

• BERT, GPT, Llama系などのモデルでは 小型版 (例: 数千万パラメータ) で学習率やベータ係数 (Adam) をチューニングし、それを数億～数十億パラメータの大型版にそのまま転用しても破綻しにくい。
• 特に FP8 など極低精度学習でメリット大。Unit Scaling のおかげで数値オーバーフローせず学習でき、メモリ削減や速度向上が期待されます。

---

実験結果と性能分析

1. 最適学習率の不変性

Yang らの研究では、標準スケーリング (SP) ではモデル幅を大きくするほど最適 $η\eta$ が下がっていくのに対し、$μ\mu$P や u-$μ\mu$P では幅を変えてもほぼ同一の $η\eta$ で最適化できると報告されています。
実際、隠れ次元を 128 から 4096 に拡大しても、u-μP であれば同じ学習率が転用可能との実験結果があります。

2. FP8など超低精度学習の安定性

Unit Scaling を組み込んだ u-μP の大きな利点は、数値レンジを狭めた場合でもテンソルの RMS 値が 1 程度に保たれることです。
実験では、従来のパラメータ初期化では FP8 での学習が発散しやすかったのに対し、u-μP モデルは FP8 でも発散せず収束し、FP32 とほぼ遜色ない性能を達成しました。

3. 探索コストの削減

u-μP では学習率以外のハイパーパラメータ（例: ベータ1・ベータ2・重み減衰など）をデフォルト値 (1相当) に固定しても、比較的最適に近い性能を得られるケースが報告されています。
小規模モデル (幅 256) で学習率を数段階スイープし、最適を探したら、大型モデル (幅 4096) にそのまま転用しても高い精度が得られるという結果が示されています。

4. 正則化との相互作用

一方、weight decay や dropout 率などの正則化ハイパーパラメータに関しては、データセットサイズやバッチサイズとも絡んでくるため、完全に “スケール不変” とは言えない場合があると報告されています。
今後、データ量やタスクの特性も踏まえた “スケーリング則” の開発が進むと期待されています。

---

結論と今後の展望

Unit-Scaling μP (u-μP) は、大規模モデル時代のチューニングパラダイムを変革する可能性を秘めた手法です。
以下のポイントで特に注目されています。

• “ゼロショット” ハイパーパラメータ転移が可能
• FP8など低精度学習での数値安定性を飛躍的に向上
• 実装簡素化：abcパラメータ化を“ユニットスケール”でまとめることで、ユーザが意識すべきスケーリング係数が減少

残る課題としては、非常に深いネットワーク（100層超えなど）や再帰型ネットワーク（RNN, LSTM）への適用、正則化ハイパーパラメータのサイズ依存問題、さらにモデルサイズ・データサイズ・コンピュート量の三要素を同時に考慮するスケーリング研究との統合が挙げられます。

とはいえ、μP から始まり、u-μP という実用的な形に発展してきたスケーリング理論は、今後ますます大規模化していくディープラーニングの基盤技術になるでしょう。

---

参考リンク

• microsoft/mup (GitHub)
  μPの公式実装（PyTorch）。MuAdamなどの最適化器を含む。

• EleutherAI/nanoGPT-mup (GitHub)
  nanoGPTにμPを適用した最小実装サンプル。

• Yang et al. "Tuning Large Neural Networks via Zero-Shot Hyperparameter Transfer"
  https://proceedings.neurips.cc/paper/2021/hash/fe160410a4e8cdef98c804a8566c91c0-Abstract.html
  NeurIPS 2021 における μP の代表的論文。

• Charlie Blake et al. "Unit-Scaled Maximal Update Parametrization"
  (ICLR 2025の予定, arXivリンク未定)。u-μPの理論と実験を詳述。

• Microsoft Research Blog :
  "μTransfer: A technique for hyperparameter tuning of enormous neural networks"
  μPによるHP転移をわかりやすく解説。

---

本記事では u-μP のアイデアと実践的メリットを概観しました。
大規模モデルを扱う際の膨大なチューニングコストや数値精度トラブルが、こうしたパラメータ化手法の進化によって大きく緩和されつつあります。
興味を持った方はぜひ公式リポジトリや関連ブログ、論文を参照し、試してみてください！