线性衰减学习率到零,比传统方法更省算力且效果更好。
Straight to Zero: Why Linearly Decaying the Learning Rate to Zero Works Best for LLMs
- 用线性衰减至零的调度方式替代余弦衰减。
- 610M模型用D2Z训练80TPP,loss低于10x衰减的200TPP。
- 适合追求高效训练的大模型研究者与工程团队。
大语言模型通常采用学习率预热后余弦衰减至最大值的10%(10x衰减)。在大规模实验中,我们发现,在最优峰值学习率下,简单的线性衰减至零(D2Z)调度在计算最优的数据集规模下始终优于其他方法。D2Z在不同模型尺寸、批量大小、数据集和词表下均表现更优,且优势随数据集增大而提升。基于对AdamW的新解释——其本质是权重更新的指数移动平均,我们揭示了线性D2Z如何在训练初期(远离初始状态)与后期(平均更多更新以抑制梯度噪声)之间取得最佳平衡。实验表明,610M参数模型以D2Z训练80令牌/参数(TPP),损失低于10x衰减下200 TPP的结果,相当于节省60%算力;类似地,如Llama2-7B在286 TPP下使用10x衰减,可能通过D2Z大幅减少计算消耗。
原文摘要 · Abstract (English)
LLMs are commonly trained with a learning rate (LR) warmup, followed by cosine decay to 10% of the maximum (10x decay). In a large-scale empirical study, we show that under an optimal peak LR, a simple linear decay-to-zero (D2Z) schedule consistently outperforms other schedules when training at compute-optimal dataset sizes. D2Z is superior across a range of model sizes, batch sizes, datasets, and vocabularies. Benefits increase as dataset size increases. Leveraging a novel interpretation of AdamW as an exponential moving average of weight updates, we show how linear D2Z optimally balances the demands of early training (moving away from initial conditions) and late training (averaging over more updates in order to mitigate gradient noise). In experiments, a 610M-parameter model trained for 80 tokens-per-parameter (TPP) using D2Z achieves lower loss than when trained for 200 TPP using 10x decay, corresponding to an astonishing 60% compute savings. Models such as Llama2-7B, trained for 286 TPP with 10x decay, could likely have saved a majority of compute by training with D2Z.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。