技術要約：µpscaling 小規模モデル：原理に基づいたウォームスタートとハイパーパラメータ転送

問題提起

現代のディープラーニングのワークフローでは、通常、複数のスケールでモデルのファミリーを訓練することになります。小規模なモデルは、迅速なプロトタイピングやスケーリング則の分析に使用され、一方で大規模なモデルは、多様なハードウェアやレイテンシの制約を満たすために、ゼロから（scratchから）訓練されます。このワークフローは、小規模モデルの訓練で得られた知識を破棄してしまうため、計算資源の無駄が生じています。

「モデル・アップスケーリング」（訓練済みの小規模モデルから大規模モデルを初期化すること）は、この無駄を軽減するために提案されてきましたが、既存の手法（例：Net2Net）には2つの決定的な限界があります。

1. ハイパーパラメータへの敏感さ： アップスケールされたモデルは、ベースとなるモデルやゼロから訓練されたモデルとは大きく異なるハイパーパラメータ（特に学習率やノイズレベル）を必要とすることがよくあります。ターゲットとなる大規模なスケールにおいてこれらをチューニングすることは、極めてコストがかかります。
2. 理論的根拠の欠如： 従来のアプローチは経験的なヒューリスティックに依存しています。小規模モデルでハイパーパラメータをチューニングし、スケーリング則を通じて外挿することが、アップスケーリングの文脈においても有効であるかどうかは不明なままです。なぜなら、アップスケールされたモデルの訓練ダイナミクスは、ベースモデルとも、ゼロから訓練された大規模モデルとも異なるからです。

手法

著者らは、効率的なハイパーパラメータ転送を可能にするために、関数保存型の重み変換と最大更新パラメータ化（Maximal Update Parametrization: µP）を組み合わせた、理論的に裏付けられたµpscalingフレームワークを提案しています。

1. 動的等価性（Dynamic Equivalence）

本研究の核心となる理論的貢献は、「静的等価性」（幅を広げたモデルが初期化時にベースモデルと同じ関数を出力すること）を動的等価性へと拡張したことです。

• メカニズム： 重みを複製し、特定の幅依存の再スケーリング（行列の1のクロネッカー積 $\otimes$ に基づくもの）を適用することで、ベースモデルと全く同じ関数をパラメータ化する、より広いモデルを構築できます。
• オプティマイザの調整： 極めて重要な点として、オプティマイザのハイパーパラメータ（学習率、ウェイトディケイ、および $\epsilon$ のような補助パラメータ）が、更新関数の同次性から導出される特定の規則に従って再スケーリングされる場合、より広いモデルは関数空間においてベースモデルと同一の訓練軌跡を辿ることを著者らは証明しています。
• 汎用性： この動的等価性は、一般的なアーキテクチャ（MLP、ResNet、Transformerを含む Ne⊗or⊤ プログラム・フレームワークで表現可能なもの）および一般的な要素ごとのオプティマイザ（SGD、Adam、AdamW）に対して成立します。

2. アップスケーリング・アルゴリズム

より広いモデルの余剰容量を活用するために、著者らは「コピー・アンド・パターブ（copy-and-perturb）」手順を導入しています。

1. 拡幅（Widening）： 重みとオプティマイザの状態を複製および再スケーリングすることにより、動的に等価な、より広いモデルを構築します。
2. 対称性の破れ（Symmetry Breaking）： パラメータに微小なノイズ（$\Delta$）を注入します。ノイズがない場合、モデルはベースモデルの低次元部分空間に閉じ込められたままとなります。ノイズは、モデルがこの部分空間から脱出し、増加した容量を活用することを可能にします。
3. µPスケーリング： ノイズの大きさおよびアップスケールされたモデルのオプティマイザ・ハイパーパラメータは、µPの規則に従ってスケーリングされます。これにより、無限幅の極限における訓練ダイナミクスが「最適」となり、ベースモデルからの信号を維持しつつ、注入されたノイズが消失または爆発することなく意味のある貢献をすることを保証します。

3. ハイパーパラメータ転送

このフレームワークは、ゼロショット・ハイパーパラメータ転送を可能にします。アップスケールされた訓練ダイナミクスの無限幅の極限は、修正されたテンソルプログラムを用いて厳密に特徴付けられているため、小規模な「プロキシ（代理）」アップスケールシステム（$n_0 \to kn_0$）でチューニングされた最適なハイパーパラメータ（学習率とノイズレベル）は、ターゲットとなる大規模なシステム（$n \to kn$）へ直接転送できます。

主な貢献

1. 動的等価性理論： 本論文は、特定の重みおよびハイパーパラメータの再スケーリングの下で、アップスケールされたモデルが初期化時だけでなく、訓練プロセス全体を通じてベースモデルと動的に等価であることを確立しました。これは、既存の理論的基盤を訓練の全軌跡にまで拡張するものです。
2. 原理に基づいたアップスケーリング・アルゴリズム： 重みの複製、対称性を破るためのノイズ注入、およびµPベースのハイパーパラメータ・スケーリングを統合した、統一されたアップスケーリング手順を提供します。これは、Net2Netのようなヒューリスティックな手法に代わる厳密な選択肢となります。
3. ハイパーパラメータ転送の保証： 著者らは、アップスケーリングの設定におけるハイパーパラメータ転送のための初の厳密なフレームワークを提供しました。テンソルプログラムの仕組みを訓練途中のアップスケーリングに対応させることで、安価な小規模プロキシシステム上で最適なハイパーパラメータをチューニングし、それを大規模なターゲットへと転送できることを証明しました。これにより、チューニングコストを劇的に削減できます。
4. 無限幅解析： 修正されたテンソルプログラム・フレームワークを導入し、アップスケールされた訓練の無限幅の極限を特徴付けました。これにより、事前活性化（pre-activations）が、標準的な非アップスケール訓練とは異なる、非自明な共分散構造を持つi.i.d.ブロックに収束することを示しています。

実験結果

本手法は、3つのアーキテクチャ（MLP、ResNet、GPT-2）、3つのデータセット（Forest Cover Type、CIFAR-100、FineWeb）、および各種オプティマイザ（SGD、AdamW）を用いて評価されました。

• パフォーマンス： 全ての条件下において、アップスケールされたモデルは、同じ幅でゼロから訓練されたモデルと比較して、収束が速く、最終的な訓練損失も低くなりました。
• ハイパーパラメータ転送： 小規模なシステム（例：幅 $n_0=100$）でチューニングされた最適な学習率とノイズレベルが、より大規模なシステム（例：幅 $n=500$ または $n=1024$）へロバストに転送されることが実験によって検証されました。
• 効率性： ターゲットのスケールではなく小規模なスケールでハイパーパラメータをチューニングすることにより、本手法はチューニングコストにおいて大幅な高速化を実現しました。
  • MLP（$n/n_0 = 5$）において 23.6倍 の高速化
  • ResNet（$n/n_0 = 4$）において 16.0倍 の高速化
  • GPT-2（$n/n_0 = 10$）において 48.2倍 の高速化
• 汎化に関する注意点： 訓練損失は改善したものの、CIFAR-100を用いたアップスケールされたResNetは、ゼロから訓練されたモデルよりも汎化性能（検証精度）が悪化しました。著者らは、これは彼らの理論（訓練ダイナミクスを制御するものであり、汎化を制御するものではない）と一致しており、限定的なデータに対する過剰パラメータ化されたCNNの既知の特性であると述べています。

重要性と主張

本論文は、モデル・アップスケーリングにおけるハイパーパラメータ転送のための初の厳密な理論的フレームワークを提供すると主張しています。その重要性は以下の点にあります。

• 理論と実践の架け橋： アップスケーリングを、無限幅の極限理論とテンソルプログラムに根ざした原理的なアプローチへと昇華させ、経験的なヒューリスティックから脱却させました。
• コスト削減： 大規模なモデルのハイパーパラメータを小規模なプロキシシステム上でチューニングすることを可能にし、大規模なアーキテクチャの探索をより計算可能なものにします。
• 普遍性： 以前の手法は特定のアーキテクチャに依存することが多かったのに対し、本フレームワークは標準的なアーキテクチャとオプティマイザに適用可能です。

著者らは、自身の理論が訓練ダイナミクスに焦点を当てていることを認め、アップスケーリング下での汎化は未解決の課題であるとして、汎化に関する主張については謙虚な姿勢を保っています。また、現在のメソッドは幅方向のアップスケーリングに適用されるものであり、深さ方向のアップスケーリングは今後の課題であるとも述べています。