WSM: Decay-Free Learning Rate Schedule via Checkpoint Merging for LLM Pre-training
Changxin Tian, jiapeng wang, Qian Zhao, Kunlong Chen, Jia Liu, Ziqi Liu, Jiaxin Mao, Xin Zhao, Zhiqiang Zhang, JUN ZHOU
Abstract
Recent advances in learning rate (LR) scheduling have demonstrated the effectiveness of decay-free approaches that eliminate the traditional decay phase while maintaining competitive performance. Model merging techniques have emerged as particularly promising solutions in this domain. We present Warmup-Stable and Merge (WSM), a general framework that establishes a formal connection between learning rate decay and model merging. WSM provides a unified theoretical foundation for emulating various decay strategies—including cosine decay, linear decay and inverse square root decay—as principled model averaging schemes, while remaining fully compatible with diverse optimization methods. Through extensive experiments, we identify merge duration—the training window for checkpoint aggregation—as the most critical factor influencing model performance, surpassing the importance of both checkpoint interval and merge quantity. With the high-quality annealing data, our framework consistently outperforms the widely-adopted Warmup-Stable-Decay (WSD) approach across multiple benchmarks, achieving significant improvements of +3.5% on MATH, +2.9% on HumanEval, and +5.5% on MMLU-Pro. The performance advantages extend to supervised fine-tuning scenarios, highlighting WSM's potential for long-term model refinement.
Ask about this paper
Your agent reads all of it.
Lune indexed this paper to the last equation, along with the top-tier papers that cite it. Ask a question and the answer quotes them.
Your agent calls
Luneget_paper_fulltext
Free to start. No credit card required.
Terminal
Install the CLIlune papers fulltext c65cd045-d6be-4254-a08d-caa750121195Cited by top-tier papers5
- Pre-training LLM without Learning Rate Decay Enhances Supervised Fine-TuningKazuki Yano, Shun Kiyono, Sosuke Kobayashi, Sho Takase et al.ICLR 2026 · 13 citations
- How Learning Rate Decay Wastes Your Best Data in Curriculum-Based LLM PretrainingKairong Luo, Zhenbo Sun, Haodong Wen, Xinyu Shi et al.ICLR 2026 · 12 citations
- MergeMix: Optimizing Mid-Training Data Mixtures via Learnable Model MergingJiapeng Wang, Changxin Tian, Kunlong Chen, ziqi liu et al.ICML 2026 · 6 citations
- Extra-Merge: Tracing the Rank-1 Subspace of Model Merging in Language Model Pre-TrainingWenJie Zhou, Bohan Wang, Hongtao Zhang, Chenxi Jia et al.ICML 2026
- Optimizing Visual Generative Models via Distribution-wise RewardsRuihang Li, Mengde Xu, Shuyang Gu, Leigang Qu et al.ICML 2026
Builds on18
- Measuring Massive Multitask Language UnderstandingDan Hendrycks, Collin Burns, Steven Basart, Andy Zou et al.ICLR 2021 · 7,905 citations
- WinoGrande: An Adversarial Winograd Schema Challenge at ScaleKeisuke Sakaguchi, Ronan Le Bras, Chandra Bhagavatula, Yejin ChoiAAAI 2020 · 3,037 citations
- PIQA: Reasoning about Physical Commonsense in Natural LanguageYonatan Bisk, Rowan Zellers, Ronan Le Bras, Jianfeng Gao et al.AAAI 2020 · 2,916 citations
- Is Your Code Generated by ChatGPT Really Correct? Rigorous Evaluation of Large Language Models for Code GenerationJiawei Liu, Chunqiu Steven Xia, Yuyao Wang, Lingming ZhangNeurIPS 2023 · 2,317 citations
- Model soups: averaging weights of multiple fine-tuned models improves accuracy without increasing inference timeMitchell Wortsman, Gabriel Ilharco, Samir Yitzhak Gadre, Rebecca Roelofs et al.ICML 2022 · 1,464 citations
Related papers
- Scaling and Transferability of Annealing Strategies in Large Language Model TrainingSiqi Wang, Zhengyu Chen, Teng Xiao, Zheqi Lv et al.AAAI 2026 · 1 citation
- Through the River: Understanding the Benefit of Schedule-Free Methods for Language Model TrainingMinhak Song, Beomhan Baek, Kwangjun Ahn, Chulhee YunNeurIPS 2025 · 9 citations
- Model Merging in Pre-training of Large Language ModelsYunshui Li, Yiyuan Ma, Shen Yan, Chaoyi Zhang et al.NeurIPS 2025 · 40 citations
- The Road Less ScheduledAaron Defazio, Xingyu Yang, Ahmed Khaled, Konstantin Mishchenko et al.NeurIPS 2024 · 208 citations
- Scaling Laws and Compute-Optimal Training Beyond Fixed Training DurationsAlexander Hägele, Elie Bakouch, Atli Kosson, Loubna Ben Allal et al.NeurIPS 2024 · 168 citations
