[论文解读] GradInit: Learning to Initialize Neural Networks for Stable and Efficient Training

Chen Zhu, Renkun Ni+4 more|arXiv (Cornell University)|Feb 16, 2021

Advanced Neural Network Applications被引用 15

一句话总结

GradInit 是一种与架构无关的方法，通过在 SGD 或 Adam 的一步更新后最小化损失，学习神经网络层的最优初始化缩放系数。该方法可加速收敛，提升测试准确率，并实现无需学习率预热的 Transformer 训练——在 ImageNet 和 IWSLT 翻译基准上达到最先进性能。

ABSTRACT

Innovations in neural architectures have fostered significant breakthroughs in language modeling and computer vision. Unfortunately, novel architectures often result in challenging hyper-parameter choices and training instability if the network parameters are not properly initialized. A number of architecture-specific initialization schemes have been proposed, but these schemes are not always portable to new architectures. This paper presents GradInit, an automated and architecture agnostic method for initializing neural networks. GradInit is based on a simple heuristic; the norm of each network layer is adjusted so that a single step of SGD or Adam with prescribed hyperparameters results in the smallest possible loss value. This adjustment is done by introducing a scalar multiplier variable in front of each parameter block, and then optimizing these variables using a simple numerical scheme. GradInit accelerates the convergence and test performance of many convolutional architectures, both with or without skip connections, and even without normalization layers. It also improves the stability of the original Transformer architecture for machine translation, enabling training it without learning rate warmup using either Adam or SGD under a wide range of learning rates and momentum coefficients. Code is available at https://github.com/zhuchen03/gradinit.

研究动机与目标

• 解决因初始化不佳导致的深层神经网络训练不稳定性问题，尤其是在复杂或新型架构中。
• 克服现有架构特定初始化规则在新架构或异构架构间缺乏可移植性的问题。
• 开发一种通用、与优化器兼容的初始化方法，提升训练稳定性和收敛速度，且无需修改网络架构。
• 在不使用学习率预热的情况下，实现使用 SGD 或 Adam 训练深层模型（如 1202 层 ResNet 和后归一化 Transformer）。
• 提供一种可扩展且高效的初始化技术，其学习最优缩放因子的时间少于总训练时间的 1%。

提出的方法

• 为每个参数块（如卷积核或全连接层）引入一个标量乘数，用于重新缩放随机初始化的权重。
• 使用数值方法优化这些标量乘数，以最小化指定优化器（SGD 或 Adam）一步更新后的训练损失。
• 将优化步骤视为可微过程，允许通过第一次更新进行反向传播，从而学习最优缩放因子。
• 考虑小批量采样、学习率以及优化器动态（如动量和自适应学习率）的随机性。
• 在标准训练开始前，将学习到的标量应用于重新缩放原始随机权重。
• 支持 Adam 和 SGD 优化器，分别为每种优化器学习不同的缩放模式，以反映其在梯度方差减少策略上的差异。

实验结果

研究问题

• RQ1 是否存在一种单一、自动化的方法，能够学习到适用于多种架构的初始化缩放系数，从而提升训练稳定性和收敛性？
• RQ2 GradInit 在使用 Adam 或 SGD 训练 Transformer 时，能在多大程度上消除对学习率预热的需求？
• RQ3 与现有初始化方法（如 Xavier、He、Admin）相比，GradInit 在最终测试准确率和训练稳定性方面表现如何？
• RQ4 GradInit 是否能稳定训练极深网络（如 1202 层 ResNet），而这些网络在标准初始化下会训练失败？
• RQ5 学习到的缩放系数中的模式揭示了关于梯度方差和初始化引起的不稳定性哪些新见解？

主要发现

• GradInit 实现了在 IWSLT-14 DE-EN 数据集上无需学习率预热即可训练原始 Post-LN Transformer，使用 Adam 或 SGD 均达到平均 BLEU 得分 36.0，且在极端超参数设置下表现稳定。
• 借助 GradInit，1202 层 ResNet 的测试准确率显著高于 ResNet-110，证明了其在标准初始化失败时仍能成功训练极深网络。
• 在 ImageNet 上，GradInit 通过学习更优的初始化方式，提升了 ResNet-50 的最终测试准确率，即使不使用批量归一化层。
• GradInit 降低了大多数层的初始梯度方差，尤其在编码器和残差连接分支中效果显著，并通过降低输出投影层和前馈网络（FFN）层的权重来平衡跳跃连接。
• 该方法为 Adam 和 SGD 学习到不同的缩放模式，表明优化器特定的初始化对梯度方差控制至关重要。
• 可视化结果表明，GradInit 通过重新缩放权重，使网络在初始化时处于损失更平滑的区域，从而减少样本间梯度方差，加速收敛。