[论文解读] Adaptive Loss Scaling for Mixed Precision Training

Rongxuan Zhao, Brian Vogel+1 more|arXiv (Cornell University)|Oct 28, 2019

Model Reduction and Neural Networks参考文献 19被引用 6

一句话总结

本文提出自适应损失缩放（adaptive loss scaling），一种在混合精度训练过程中动态计算逐层损失缩放值的方法，可在无需手动调整超参数的情况下防止梯度下溢。通过根据梯度统计信息自动调整每层的缩放系数，该方法在某些情况下实现了比固定损失缩放甚至FP32训练更快的收敛速度和更高的准确率，尤其在深层模型和大批次训练中表现更优。

ABSTRACT

Mixed precision training (MPT) is becoming a practical technique to improve the speed and energy efficiency of training deep neural networks by leveraging the fast hardware support for IEEE half-precision floating point that is available in existing GPUs. MPT is typically used in combination with a technique called loss scaling, that works by scaling up the loss value up before the start of backpropagation in order to minimize the impact of numerical underflow on training. Unfortunately, existing methods make this loss scale value a hyperparameter that needs to be tuned per-model, and a single scale cannot be adapted to different layers at different training stages. We introduce a loss scaling-based training method called adaptive loss scaling that makes MPT easier and more practical to use, by removing the need to tune a model-specific loss scale hyperparameter. We achieve this by introducing layer-wise loss scale values which are automatically computed during training to deal with underflow more effectively than existing methods. We present experimental results on a variety of networks and tasks that show our approach can shorten the time to convergence and improve accuracy compared to the existing state-of-the-art MPT and single-precision floating point

研究动机与目标

• 为解决FP16中因梯度下溢导致的混合精度训练不稳定与收敛性差的问题。
• 消除在不同模型和训练阶段中对固定损失缩放系数进行手动超参数调优的需求。
• 通过逐层和逐迭代自适应损失缩放，提升混合精度训练的可用性与可靠性。
• 在准确率和收敛速度方面，实现与全精度（FP32）训练相当或更优的训练性能。
• 通过避免固定损失缩放所需的多次超参数搜索，减少总训练时间。

提出的方法

• 该方法引入了在训练过程中基于激活梯度分布动态计算的逐层损失缩放值。
• 对于每一层，损失缩放系数计算为（0.01N）分位数最小绝对梯度值的倒数，其中N为梯度大小，以确保下溢率不超过1%。
• 损失缩放系数定期更新（例如每100次迭代更新一次），以适应训练不同阶段中梯度统计的变化。
• 该方法使用FP32进行权重更新和批量归一化，而激活值和梯度则以FP16存储，从而保持计算效率。
• 通过基于梯度幅值统计将损失缩放系数保持在安全范围内，避免溢出和过度舍入误差。
• 该方法可无缝集成到现有深度学习框架中，无需修改网络架构。

实验结果

研究问题

• RQ1 与固定损失缩放相比，动态的、逐层的损失缩放是否能提升混合精度训练的稳定性和收敛性？
• RQ2 自适应损失缩放是否能显著减少在不同模型和数据集上对超参数进行手动调优的需求？
• RQ3 自适应损失缩放能否在准确率上超越固定损失缩放和全精度（FP32）训练？
• RQ4 与多次超参数搜索相比，自适应损失缩放对训练速度和总训练时间的影响如何，尤其是在大批次训练场景下？
• RQ5 自适应损失缩放在多大程度上缓解了不同层和不同训练迭代中的梯度下溢问题？

主要发现

• 在批量大小为32的SSD512上，自适应损失缩放实现了80.31%的mAP，优于FP32基线（79.50%）和固定损失缩放（80.17%）。
• 对于ResNet-50，自适应损失缩放达到79.24%的测试准确率，优于最佳固定缩放系数（79.11%），并避免了固定缩放系数为128时出现的准确率下降。
• 通过逐层自适应调整缩放系数，该方法显著降低了下溢率，尤其是在深层网络中，如图4(b)所示。
• 在自适应缩放下，梯度的最大标准差约为固定缩放下的1/20，从而减少了累积误差，提升了训练稳定性。
• 在每100次迭代更新一次的频率下，计算损失缩放的开销仅占总训练时间的0.27%，效率极高。
• 该方法消除了为调优损失缩放而进行多次训练运行的需求，相比固定缩放配合超参数搜索，显著减少了总训练时间。