[論文レビュー] Depthwise Hyperparameter Transfer in Residual Networks: Dynamics and Scaling Limit

Blake Bordelon, Lorenzo Noci+3 more|arXiv (Cornell University)|Sep 28, 2023

Domain Adaptation and Few-Shot Learning被引用数 4

ひとこと要約

本論文は、残差ネットワークにおける新しいパrameterizationを提案する——$μ$P初期化と残差ブランチの$1/\sqrt{\text{depth}}$スケーリングを組み合わせることで、幅と深さの両方において最適なハイパーパrameter（例：学習率、モーメンタム、重み減衰）の転送を可能にする。この手法は、動的平均場理論（DMFT）を用いて無限幅および無限深さの極限における安定した学習ダイナミクスを実現し、CIFAR-10、ImageNet、Vision Transformersにおいて理論的および実験的にハイパーパrameter転送の堅牢性を示している。

ABSTRACT

The cost of hyperparameter tuning in deep learning has been rising with model sizes, prompting practitioners to find new tuning methods using a proxy of smaller networks. One such proposal uses $μ$P parameterized networks, where the optimal hyperparameters for small width networks transfer to networks with arbitrarily large width. However, in this scheme, hyperparameters do not transfer across depths. As a remedy, we study residual networks with a residual branch scale of $1/\sqrt{ ext{depth}}$ in combination with the $μ$P parameterization. We provide experiments demonstrating that residual architectures including convolutional ResNets and Vision Transformers trained with this parameterization exhibit transfer of optimal hyperparameters across width and depth on CIFAR-10 and ImageNet. Furthermore, our empirical findings are supported and motivated by theory. Using recent developments in the dynamical mean field theory (DMFT) description of neural network learning dynamics, we show that this parameterization of ResNets admits a well-defined feature learning joint infinite-width and infinite-depth limit and show convergence of finite-size network dynamics towards this limit.

研究の動機と目的

• 大規模な深層ネットワークにおけるハイパーパrameterチューニングのコスト増加に対処し、幅および深さの両方において最適なハイパーパrameterを転送可能にする。
• 既存の$\mu$Pパrameterizationの限界を克服し、残差ネットワークにおける深さにわたるハイパーパrameter転送を可能にする。
• 動的平均場理論（DMFT）を用いて、広大で深い残差ネットワークにおける安定したサイズに依存しない学習ダイナミクスの理論的基盤を確立する。
• 提案されたパrameterizationが畳み込みResNetsおよびVision Transformersにおいて一貫したハイパーパrameter転送を可能にするかを実験的に検証する。
• ネットワークの無限幅および無限深さ極限におけるダイナミクスを特徴づけ、適切に定義されたスケーリング極限への収束を示す。

提案手法

• 残差ブランチの重みを$1/\sqrt{\text{depth}}$にスケーリングすることで、深さにわたる特徴学習ダイナミクスの安定化を図る、$\mu$Pパrameterizationの変更版を導入する。
• 無限幅および無限深さ極限におけるネットワークの解析に、動的平均場理論（DMFT）を適用し、層インデックスを連続的な「層時間」として扱う。
• 特徴の進化を記述する連続時間の確率微分方程式（SDE）を導出し、特徴および予測の更新がネットワークの幅および深さに依存しないことを保証する。
• 積分因子および連続時間近似を用いて重み減衰ダイナミクスをモデル化し、初期条件の影響が時間とともに指数関数的に減衰することを示す。
• 読み出し重みをゼロ初期化した多層ブロックアーキテクチャを採用し、深さスケーリングの効果が学習ダイナミクスおよびハイパーパrameter転送に与える影響を分離する。
• 固定学習率および複雑なスケジュールを用いたCIFAR-10、Tiny ImageNet、ImageNet、およびVision Transformersにおける広範な実験により結果を検証する。

(a) ResNet, $\mu$ P

実験結果

リサーチクエスチョン

• RQ1 最適なハイパーパrameter（例：学習率）が、残差ネットワークにおいて幅および深さの両方で転送可能か。
• RQ2 残差ブランチの$1/\sqrt{\text{depth}}$スケーリングが、無限幅および無限深さ極限における安定した学習ダイナミクスを実現するか。
• RQ3 提案されたパrameterizationは、標準的な$\mu$Pおよびバッチ正則化と比較して、深さにわたるハイパーパrameter転送をどのように向上させるか。
• RQ4 DMFTフレームワークは、この新しいパrameterizationを用いた残差ネットワークの同時無限幅および無限深さ極限を記述可能か。
• RQ5 提案された手法は、学習率に加え、モーメンタムや重み減衰といったハイパーパrameterの転送をサポートするか。

主な発見

• 提案された$1/\sqrt{\text{depth}}$-ResNetに$\mu$P初期化を適用したモデルは、CIFAR-10において幅および深さの両方で最適な学習率を転送可能であり、テストされた深さの範囲で発散する実験が一切観察されなかった。
• 実験的結果から、固定学習率だけでなく、周期的およびコサインスケジュールを含む複雑なスケジュールに対しても一貫したハイパーパrameter転送が確認された。
• 提案されたパrameterizationを用いて訓練されたVision Transformersも、バッチ正則化の有無に関わらず、最適なハイパーパrameterの転送を示した。
• バッチ正則化を用いた標準的なResNetsに比べ、本手法はハイパーパrameter転送を改善しており、特に大規模な深さにおいても一貫性を保った。
• DMFTによる理論的分析により、特徴および予測の更新がネットワークサイズに依存しない well-defined な無限幅および無限深さ極限が確認され、深層線形ネットワークに対して正確な解が導出された。
• ODEスケーリング（$\alpha=1$）では深さの増加に伴い単調な改善が見られたが、SDEスケーリング（$\alpha=1/2$）では初期段階で非単調な挙動を示し、収束ダイナミクスの違いを示唆した。

(b) This work: $\frac{1}{\sqrt{L}}$ -ResNet + $\mu$ P