No More Adam: Learning Rate Scaling at Initialization is All You Need

在这项工作中，我们对在训练深度神经网络时使用自适应梯度方法的必要性提出了质疑。SGD-SaI 是对带有动量的随机梯度下降（SGDM）的一个简单而有效的增强。SGD-SaI 在初始化时对不同的参数组进行学习率缩放（SaI），这一过程由各自的梯度信噪比（g-SNR）指导。通过不依赖于自适应二阶动量来调整学习率，SGD-SaI 从第一次迭代开始就帮助防止训练不平衡，并且相比 AdamW 将优化器的内存使用减少了近一半。尽管其简单高效，SGD-SaI 在训练各种基于 Transformer 的任务中始终能够与 AdamW 匹配或超越其性能，有效克服了使用 SGD 训练 Transformer 的长期挑战。SGD-SaI 在使用视觉 Transformer（ViT）进行 ImageNet-1K 分类以及使用 GPT-2 预训练大规模语言模型（LLM，仅解码器 Transformer）方面表现出色，显示出对超参数变化的鲁棒性和在不同应用中的实用性。我们进一步测试了其在如 LLM 的 LoRA 微调和扩散模型等任务中的鲁棒性，结果表明它始终优于最先进的优化器。从内存效率的角度来看，SGD-SaI 在优化器状态上实现了显著的内存节省，在全精度训练设置下，GPT-2（1.5B 参数）的内存使用减少了 5.93 GB，Llama2-7B 的内存使用减少了 25.15 GB。

该论文旨在探讨在训练深度神经网络时自适应梯度方法的必要性，并提出了一种简单有效的SGD增强方法——SGD-SaI，以克服使用SGD训练Transformer模型的长期挑战。

SGD-SaI通过在初始化阶段对不同参数组进行学习率缩放，基于各自的梯度信噪比（g-SNR），而不依赖于自适应二阶动量。这种方法不仅简化了优化过程，还减少了内存使用，同时在多个任务上匹配或超过了AdamW的性能。

论文在ImageNet-1K分类和GPT-2预训练等任务上展示了SGD-SaI的优越性能，特别是在超参数变化下的鲁棒性和对不同应用的适用性。此外，SGD-SaI在LoRA微调和扩散模型等任务上也表现出色，显著降低了优化器状态的内存消耗，如在GPT-2 (1.5B参数) 和Llama2-7B上分别节省了5.93 GB和25.15 GB的内存。

近期在优化算法方面的相关研究包括：1.《Adam: A Method for Stochastic Optimization》提出了Adam优化器；2.《RAdam: On The Variance Of The Adaptive Learning Rate And Beyond》改进了Adam的方差问题；3.《LAMB: Large Batch Optimization for Deep Learning Training》针对大规模批处理优化进行了研究；4.《On the Convergence of Adam and Beyond》分析了Adam的收敛性。