WSM:基于 checkpoint 合并的无衰减学习率计划用于 LLM 预训练
计算与语言
2025-08-12 v2 机器学习
摘要
最近的学习率(LR)调度研究表明,消除传统衰减阶段同时保持竞争性能的无衰减方法效果显著。模型合并技术在这一领域中尤为有前景。我们提出了 Warmup-Stable and Merge(WSM),一种建立在学习率衰减与模型合并之间正式联系上的通用框架。WSM 为模拟各种衰减策略——包括余弦衰减、线性衰减和倒数平方根衰减——提供了统一的理论基础,作为原则模型平均方案,同时与多样化的优化方法完全兼容。通过大量实验,我们发现合并持续时间——即 checkpoint 聚合的训练窗口——是影响模型性能最关键的因素,超过了 checkpoint 间隔和合并数量的重要性。我们的框架在多个基准测试中 consistently 超过广泛采用的 Warmup-Stable-Decay(WSD)方法,分别在 MATH、HumanEval 和 MMLU-Pro 上实现了 +3.5%、+2.9% 和 +5.5% 的显著提升。性能优势也延伸至监督微调场景,凸显了 WSM 在长期模型细化方面的潜力。
引用
@article{arxiv.2507.17634,
title = {WSM: Decay-Free Learning Rate Schedule via Checkpoint Merging for LLM Pre-training},
author = {Changxin Tian and Jiapeng Wang and Qian Zhao and Kunlong Chen and Jia Liu and Ziqi Liu and Jiaxin Mao and Wayne Xin Zhao and Zhiqiang Zhang and Jun Zhou},
journal= {arXiv preprint arXiv:2507.17634},
year = {2025}
}