[Paper Review] 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 Learning4 citations

TL;DR

This paper proposes a novel parameterization for residual networks—combining $μ$P initialization with a $1/\sqrt{\text{depth}}$ scaling of residual branches—that enables transfer of optimal hyperparameters (e.g., learning rates, momentum, weight decay) across both width and depth. The method achieves stable training dynamics in the infinite-width and infinite-depth limit via dynamical mean field theory (DMFT), demonstrating empirically and theoretically that hyperparameter transfer is robust across CIFAR-10, ImageNet, and Vision Transformers.

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.

Motivation & Objective

• To address the growing cost of hyperparameter tuning in large deep networks by enabling transfer of optimal hyperparameters across both width and depth.
• To overcome the limitation of existing $\mu$P parameterization, which fails to transfer hyperparameters across depth in residual networks.
• To establish a theoretical foundation for stable, size-independent training dynamics in wide and deep residual networks using dynamical mean field theory (DMFT).
• To empirically validate that the proposed parameterization enables consistent hyperparameter transfer in convolutional ResNets and Vision Transformers.
• To characterize the infinite-width and infinite-depth limit of the network dynamics and show convergence to a well-defined scaling limit.

Proposed method

• Introduces a modified $\mu$P parameterization where residual branch weights are scaled as $1/\sqrt{\text{depth}}$ to stabilize feature learning dynamics across depth.
• Applies dynamical mean field theory (DMFT) to analyze the infinite-width and infinite-depth limit of the network, treating layer index as a continuous 'layer time'.
• Derives a continuous-time stochastic differential equation (SDE) for feature evolution, ensuring feature and prediction updates scale independently of network width and depth.
• Uses integrating factors and continuous-time approximations to model weight decay dynamics, showing that initial condition effects decay exponentially over time.
• Employs a multi-layer-per-block architecture with zero-initialized readout weights to isolate the effect of depth scaling on training dynamics and hyperparameter transfer.
• Validates results through extensive experiments on CIFAR-10, Tiny ImageNet, ImageNet, and Vision Transformers with fixed learning rates and complex schedules.

(a) ResNet, $\mu$ P

Experimental results

Research questions

• RQ1 Can optimal hyperparameters such as learning rate be transferred across both width and depth in residual networks?
• RQ2 Does a $1/\sqrt{\text{depth}}$ scaling of residual branches enable stable training dynamics in the infinite-width and infinite-depth limit?
• RQ3 How does the proposed parameterization compare to standard $\mu$P and batch normalization in enabling hyperparameter transfer across depth?
• RQ4 Can the DMFT framework describe the joint infinite-width and infinite-depth limit of residual networks with this new parameterization?
• RQ5 Does the proposed method support transfer of hyperparameters beyond learning rate, such as momentum and weight decay?

Key findings

• The proposed $1/\sqrt{\text{depth}}$-ResNet with $\mu$P initialization enables transfer of optimal learning rates across both width and depth on CIFAR-10, with no divergent runs observed across tested depths.
• Empirical results show consistent hyperparameter transfer not only for fixed learning rates but also for complex schedules, including cyclical and cosine annealing schedules.
• Vision Transformers trained with the proposed parameterization also exhibit transfer of optimal hyperparameters, both with and without batch normalization.
• The method improves hyperparameter transfer over standard ResNets with batch normalization, which only shows partial improvement and remains inconsistent at large depths.
• Theoretical analysis via DMFT confirms a well-defined infinite-width and infinite-depth limit where feature and prediction updates are independent of network size, with exact solutions derived for deep linear networks.
• The model shows monotonic improvement with increasing depth under ODE scaling ($\alpha=1$), while SDE scaling ($\alpha=1/2$) exhibits non-monotonic behavior at early times, suggesting different convergence dynamics.

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