通过动态调整学习率与批量大小,实现训练加速而不损失精度。
Seesaw: Accelerating Training by Balancing Learning Rate and Batch Size Scheduling
- 学习率减半时,批量大小翻倍,保持损失变化稳定。
- 在相同计算量下,比传统方法快36%且精度相当。
- 适用于大规模语言模型预训练,尤其适合自适应优化器场景。
在大语言模型预训练中,逐步增大批量大小(即批量增长策略)是一种有前景的加速方法。尽管对于SGD而言,批量大小加倍等价于学习率减半,但自适应优化器如Adam的最优策略尚不明确,导致批量增长常需手动调参。本文提出一个理论严谨的批量调度框架,并引入Seesaw方法:当标准调度器将学习率减半时,Seesaw改为乘以1/√2并使批量大小翻倍,从而在保持损失动态的同时减少串行步骤。理论上,我们首次给出了针对噪声线性回归的有限样本下学习率衰减与批量增长等价性的证明,并将该结论扩展至归一化SGD(Adam的可解析代理),该代理在实践中表现出方差主导特性。实验表明,在150M/300M/600M参数模型、按Chinchilla规模训练且采用恒定(临界)批量大小的情况下,Seesaw在相同浮点运算量下达到与余弦衰减相当的性能,同时将实际运行时间减少约36%,逼近分析所推导的理论极限。
原文摘要 · Abstract (English)
Increasing the batch size during training -- a ''batch ramp'' -- is a promising strategy to accelerate large language model pretraining. While for SGD, doubling the batch size can be equivalent to halving the learning rate, the optimal strategy for adaptive optimizers like Adam is less clear. As a result, any batch-ramp scheduling, if used at all, is typically tuned heuristically. This work develops a principled framework for batch-size scheduling and introduces Seesaw: whenever a standard scheduler would halve the learning rate, Seesaw instead multiplies it by $1/\sqrt{2}$ and doubles the batch size, preserving loss dynamics while reducing serial steps. Theoretically, we provide, to our knowledge, the first finite-sample proof of equivalence between learning-rate decay and batch-size ramp-up for SGD on noisy linear regression, and we extend this equivalence to normalized SGD, a tractable proxy for Adam, under a variance-dominated regime observed in practice. Empirically, on 150M/300M/600M-parameter models trained at Chinchilla scale using a constant (critical) batch size, Seesaw matches cosine decay at equal FLOPs while reducing wall-clock time by $\approx 36\%$, approaching the theoretical limit implied by our analysis.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。