muP : unit-scalingを搭載して最強になろう
Unit-Scaling μP(μPスケーリング)の理論と実践【拡充版】
大規模モデルを効率的にチューニングするための先端的パラメータスケーリング手法、Unit-Scaling μP (u-μP) を解説します。
本記事では、以下の流れで理解を深めていきます。
- 背景:Maximal Update Parametrization (μP) とは
- μPの理論:数式を交えた詳細
- Unit Scaling の概念とう-μpへの発展
- 数式で見るu-μP
- 適用可能なモデルと応用例
- 実験結果と性能分析
- 結論と今後の展望
- 参考リンク
背景:Maximal Update Parametrization (μP) とは
近年のディープラーニングにおいて、大規模化とともにハイパーパラメータ(学習率や初期化など)の再調整が必要となる問題が顕在化しています。とくに、隠れ次元(幅)やチャネル数を変えるたびに最適学習率を小さくしないと学習が破綻するといった経験則は、多くの研究者が直面している課題です。
Maximal Update Parametrization (μP) は、この問題に対処すべく Microsoft Research の Yang らによって提案されました。目的は次の3点です。
- 幅を増やしても学習動態が崩れない。
- 小規模モデルで見つけたハイパーパラメータを大規模モデルに転用できる。(ゼロショット転移)
- 無限幅でも特徴学習が十分起きるようにする。(NTKパラメータ化のように線形化しない)
実際の応用としては、たとえば13Mパラメータ規模でハイパーパラメータをチューニングし、それを3.5億パラメータのモデルに再調整なしで適用して高精度が得られるなど、“μTransfer” と呼ばれる活用事例が報告されています。
μPの理論:数式を交えた詳細
1. 標準スケーリングとNTKスケーリング
まず、最も基本的な全結合層(隠れ次元
-
標準スケーリング(SP)
一般的なHe初期化やXavier初期化を想定します。たとえばHe初期化では、
\mathrm{Var}(W_{ij}) = \frac{2}{n}.
これは出力の分散が大きくなりすぎないようにするためです。一方で、学習率 は実験的にチューニングされるため、モデル幅が大きくなるほど最適\eta が小さくなっていくという傾向がしばしば報告されています。\eta -
NTKパラメータ化
Neural Tangent Kernel (NTK) 研究で導入されたもので、幅 の極限でニューラルネットが初期時点の線形カーネルとして振る舞うようパラメータを設定します。重み初期化は標準スケーリングと同程度にしつつ、学習率n \to \infty を\eta
\eta \approx O\Bigl(\frac{1}{n}\Bigr)
として極端に小さくします。こうすると「学習中に重みがほとんど動かない(lazy training)」になり、線形近似的に解析できるのがメリットですが、特徴表現の更新がほぼ起きないため、幅を増やしても性能向上は限定的になることが多いです。
2. μPのスケーリング
これら2つの方法の中間的アプローチとしてμPが位置づけられます。直感的には、「モデル幅
例えば2層ネットワークを
と書き、幅
標準スケーリングでは
となるのが典型です。NTKの場合は小さすぎて「ほとんど更新しない」極限に陥りますが、μP では幅に依存しない一定オーダーの学習率を使いつつ、重み初期化スケールなどを少し調整することで
を保ち、無限幅でも特徴学習が起きるリッチな極限を実現します。
論文では「abcパラメータ化」と呼ばれるテンソル単位のスケーリング定数を導入し、各レイヤーで"出力の分散" や "勾配の分散" が幅
Unit Scaling の概念とう-μPへの発展
1. Unit Scaling とは
Unit Scaling は、モデルの各テンソル(活性、重み、勾配など)の分散を常に“1” 程度のユニットスケールに保つための設計手法です。
特に、FP16 や FP8 のように数値のダイナミックレンジが狭い精度で学習する際、これを守るとオーバーフローやアンダーフローが劇的に起きにくくなるため、大規模モデルの省メモリ学習や分散学習を安定化できます。
たとえば、全結合層であれば「重み行列の出力直後に
2. Unit-Scaling μP (u-μP)
u-μP は、この Unit Scaling と μP を統合したもので、以下のようなメリットを併せ持ちます。
- モデル幅(隠れ次元)を変えても最適ハイパーパラメータが不変(μPの利点)。
- 全てのテンソルのRMS値が 1 程度に保たれ、低精度学習が安定(Unit Scalingの利点)。
実装的には、μPで層ごとに設定していた学習率スケーリングや初期化ルールを、“最初から出力が 1 のスケールになる” 形で整理しており、ユーザは「ベースモデルの幅」を明示的に意識しなくても適切なスケーリングがなされるようになっています。
数式で見るu-μP
ここではもう少し形式的に、u-μPの要点を数式で示します。
1. 初期化スケーリング
幅(チャネル数)を表す
重み行列
ただし、層
とすることで、
(
2. 学習率スケーリング
NTKパラメータ化では
のオーダーを保つことがしばしば用いられます。u-μPではさらに全テンソルが単位スケールになるよう、オプティマイザ内部で層ごとにスケール調整を行います。
実際には以下のような実装が存在しています。
- MuAdam / MuSGD : 学習率を自動でスケーリングし、ベース幅を意識せずに済む仕組み(microsoft/mup 参照)。
3. abcパラメータ化の単純化
μP論文で紹介された「abcパラメータ化」では、各レイヤーに対し初期化用・前向き伝搬用・勾配計算用の係数 (
u-μPでは、Unit Scalingの考え方で「最初から分散1を維持する」ための定数が統合され、実際のコードではほぼ “定数1” だけで済む層も多いとのことです。これにより、ユーザ視点では大幅にスケーリング係数が減り、実装が簡素になります。
適用可能なモデルと応用例
1. Transformerモデル
-
埋め込み(embedding): ベクトルの次元
が幅に相当する場合、入力層は固定次元からの射影になるため、ここをどうスケーリングするかが重要です。u-μPでは埋め込み重みに対してd_{\text{model}}
\mathrm{Var}(W_{\text{embed}}) \propto \frac{1}{\mathrm{fan\text{-}in}}
かつ出力の活性にスケーリング定数を掛けます。 -
自己注意機構: マルチヘッドのキー・クエリ・バリュー射影やスケーリング定数
は標準Transformerと同様ですが、Residual接続(Add & Norm)の部分に追加のスケーリングが施される場合があります。u-μP実装では Residualブロックをユニットスケールに保つために、ブロック出力を1/\sqrt{d_{\text{head}}} で割る工夫を入れることも。1/\sqrt{\text{depth}} -
出力層(unembedding): 埋め込みと重みを共有する場合 (tied embedding) もあり、そこにも u-μPのスケーリング則を適用。小さいモデル (例:
) の最適設定が、大きいモデル (d_{\text{model}}=256 ) でもそのまま最適化に近いと報告されています。d_{\text{model}}=4096
2. CNN/ResNet系
- チャネル数が “幅” に該当します。
- Convolutionフィルタも
で初期化し、出力を1/d_{\mathrm{in}} で正規化するなどの工夫で活性の分散を制御。1/\sqrt{d_{\mathrm{out}}} - Residual接続の出力を足し合わせる前に
で割るなど、深さ方向のスケーリングも考慮すると、幅を増減させても安定した学習が得られます。1/\sqrt{\text{depth}}
3. 大規模言語モデル(LLM)
- BERT, GPT, Llama系などのモデルでは 小型版 (例: 数千万パラメータ) で学習率やベータ係数 (Adam) をチューニングし、それを数億~数十億パラメータの大型版にそのまま転用しても破綻しにくい。
- 特に FP8 など極低精度学習でメリット大。Unit Scaling のおかげで数値オーバーフローせず学習でき、メモリ削減や速度向上が期待されます。
実験結果と性能分析
1. 最適学習率の不変性
Yang らの研究では、標準スケーリング (SP) ではモデル幅を大きくするほど最適
実際、隠れ次元を 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 のアイデアと実践的メリットを概観しました。
大規模モデルを扱う際の膨大なチューニングコストや数値精度トラブルが、こうしたパラメータ化手法の進化によって大きく緩和されつつあります。
興味を持った方はぜひ公式リポジトリや関連ブログ、論文を参照し、試してみてください!
Discussion