发现大模型预训练中SGD性能差,因难以维持大有效学习率。
Revisiting the Adam-SGD Gap in LLM Pre-Training: The Role of Large Effective Learning Rates

- 通过分析梯度与权重比,揭示大批次下需大有效学习率
- 发现输出层梯度不均且常有尖峰,限制了SGD学习率
- 简单裁剪机制让SGD接近Adam性能,验证理论有效性
普遍认为随机梯度下降(SGD)在大语言模型(LLM)预训练中表现远逊于自适应优化器如Adam。本文指出,这一差距主要源于SGD无法维持与Adam相当的大有效学习率。通过对大模型预训练动态的实证与理论分析,我们发现训练过程具有小梯度范数和大权重-梯度比特征,且该效应随典型预训练的大批次规模进一步增强,亟需大有效学习率。然而,输出层不同词元类别的梯度幅度差异显著,训练中频繁出现大梯度尖峰,严重制约了SGD可接受的学习率。基于此理解,我们证明简单的裁剪机制能稳定大学习率下的SGD,使其恢复接近Adam的性能。在大规模实验中,当使用100万令牌批次预训练10亿参数的LLaMA模型时,大学习率SGD与Adam之间的验证损失差距从超过50%缩小至约3.5%。
原文摘要 · Abstract (English)
It is widely believed that stochastic gradient descent (SGD) performs significantly worse than adaptive optimizers such as Adam in pre-training Large Language Models (LLMs). Yet the underlying reason for this gap remains unclear. In this work, we attribute a large part of the discrepancy to SGD's inability to sustain learning rates comparable to Adam's much larger effective learning rates. Through empirical and theoretical analysis of LLM pre-training dynamics, we identify that training is characterized by small gradient norms and large weight-to-gradient ratios, an effect that becomes more pronounced with larger batch sizes typical in pre-training, necessitating such large effective learning rates. However, we find that output-layer gradient magnitudes become highly uneven across token classes, and that large gradient spikes frequently occur during training. Together, these effects severely restrict the admissible learning rate of SGD. Guided by this understanding, we show that simple clipping mechanisms that stabilize SGD at large learning rates enable it to recover most of Adam's performance. In our large-scale experiments, the validation loss gap between large-learning-rate SGD and Adam shrinks from more than 50% to only about 3.5% when pre-training a 1B-parameter LLaMA model with a 1M-token batch size.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。