[论文解读] 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

一句话总结

该论文提出了一种残差网络的新参数化方法——结合 $\mu$P 初始化与残差分支的 $1/\sqrt{\text{depth}}$ 缩放——实现了在宽度和深度之间最优超参数（如学习率、动量、权重衰减）的迁移。该方法通过动力学平均场理论（DMFT）在无限宽和无限深极限下实现了稳定的训练动态，实证和理论结果表明，超参数迁移在 CIFAR-10、ImageNet 和视觉 Transformer 上均具有鲁棒性。

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.

研究动机与目标

• 为应对大规模深度网络中日益增长的超参数调优成本，通过在宽度和深度之间实现最优超参数的迁移。
• 克服现有 $\mu$P 参数化方法在残差网络中无法在深度之间迁移超参数的局限性。
• 利用动力学平均场理论（DMFT）为宽且深的残差网络建立稳定、与规模无关的训练动态的理论基础。
• 通过实证验证，所提出的参数化方法可在卷积残差网络和视觉 Transformer 中实现一致的超参数迁移。
• 表征网络动态在无限宽和无限深极限下的特性，并展示其收敛到一个明确定义的标度极限。

提出的方法

• 提出一种改进的 $\mu$P 参数化方法，其中残差分支权重按 $1/\sqrt{\text{depth}}$ 缩放，以稳定跨深度的特征学习动态。
• 应用动力学平均场理论（DMFT）分析网络在无限宽和无限深极限下的行为，将层索引视为连续的“层时间”。
• 推导出特征演化的连续时间随机微分方程（SDE），确保特征和预测更新与网络宽度和深度无关。
• 使用积分因子和连续时间近似方法建模权重衰减动态，表明初始条件的影响随时间呈指数衰减。
• 采用每块多层结构并使用零初始化的读出权重，以隔离深度缩放对训练动态和超参数迁移的影响。
• 通过在 CIFAR-10、Tiny ImageNet、ImageNet 和视觉 Transformer 上进行广泛实验，验证结果，采用固定学习率和复杂调度策略。

(a) ResNet, $\mu$ P

实验结果

研究问题

• RQ1 在残差网络中，最优超参数（如学习率）是否可实现在宽度和深度之间的迁移？
• RQ2 对残差分支采用 $1/\sqrt{\text{depth}}$ 缩放是否能在无限宽和无限深极限下实现稳定的训练动态？
• RQ3 与标准 $\mu$P 和批量归一化相比，所提出的参数化方法在实现深度间超参数迁移方面表现如何？
• RQ4 DMFT 框架是否能够描述采用此新参数化的残差网络的联合无限宽和无限深极限？
• RQ5 所提出的方法是否支持除学习率外的其他超参数（如动量和权重衰减）的迁移？

主要发现

• 所提出的 $1/\sqrt{\text{depth}}$-ResNet 配合 $\mu$P 初始化在 CIFAR-10 上实现了学习率在宽度和深度之间的迁移，且在所有测试深度下均未出现发散训练。
• 实证结果表明，超参数迁移不仅适用于固定学习率，也适用于复杂调度策略，包括循环调度和余弦退火调度。
• 使用该参数化训练的视觉 Transformer 同样表现出最优超参数的迁移，无论是否使用批量归一化。
• 与标准残差网络（使用批量归一化）相比，该方法在超参数迁移方面表现更优，后者仅表现出部分改进，且在大深度下仍不稳定。
• 通过 DMFT 的理论分析证实，在无限宽和无限深极限下存在明确定义的极限，其中特征和预测更新与网络规模无关，并为深度线性网络推导出精确解。
• 在 ODE 缩放（$\alpha=1$）下，模型随深度增加表现出单调改进；而在 SDE 缩放（$\alpha=1/2$）下，早期时间表现出非单调行为，提示存在不同的收敛动力学。

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