WSM: Decay-Free Learning Rate Schedule via Checkpoint Merging for LLM Pre-training
Authors: Changxin Tian
jiapeng wang
Qian Zhao
Kunlong Chen
Jia Liu
Ziqi Liu
Jiaxin Mao
Xin Zhao
Zhiqiang Zhang
JUN ZHOU
Venue: ICLR 2026 Oral (2026)
Keywords: llm pre-training
learning rate schedule
checkpoint merging
decay-free approach
Primary area: foundation or frontier models, including LLMs

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.

Published as a conference paper at ICLR 2026
WSM: DECAY-FREE LEARNING RATE SCHEDULE VIA CHECKPOINT MERGING FOR LLM PRE-TRAINING
Changxin Tian's Jiapeng Wang
- Qian Zhao', Kunlong Chen', .Jia Liu',
Ziqi Liu', Jiaxin M Wayne Xin Zhao 2. Zhiqiang Zhang't Jun Zhou'
'Ling Team, Ant Group "Gooling School of Artificial Intelligence, Renmin University of China Beijing Key Laboratory of Research on Large Models and Intelligent Governance t Lanhannxiin t cxê antgrrup..p... com, wangjp1 06 ruc. www.c.du..cccc
ABSTRACT
Recent advances in learning rate (LR) scheduling have demonstrated the effec- tiveness 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), general framework that establishes formal connection be- tween learning rate decay and model merging. WSM provides a unified theo- retical foundation for emulating various decay strategies incluudin cosine de- cay, 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 win- dow 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 out- performs the widely-adopter Warmup-Stable Staluup-ttaae--aaaa Decay (WSD) approach across mul- tiple benchmarks, achieving significant improvements of +3.5% on MATH, +2.9% 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.
INTRODUCTION
In large language model (LLM) pre-training, learning rate (LR) scheduling plays pivotal role, crit- ically impacting training stability, convergence speed, and final model performance (Jin al. 2023 Gotmare et al. 2019). Conventional LR schedules dynamically adjust the LR based on training
(Ours)
Continsal maaliing
restart poier
Rall back this polnt Do traisin
Learning Rate
Regures prospeelliain of deootrpphior staiggy
Stenlating varian
of lntel maiking stes
He checkpoic werping
Training Progression
Training Progression
Training Progression
Comparison between our proposed WSM (Warmup-Stable and Merge) and mainstream learning rate scheduling strategies, Cosine (Loshchilov & Hutter 2016) and WSD (Hu al. 2024).
Figure 1:
"qqual contribution. 'Contribution during internship at Ant Group.
I 'Correpondiig authors: Wayne Xin Zhao (batman yēgma il 1, com) and Zhiqiang Zhang
(1 ngyao. zxg@anttrruup.com)

---

Published as a conference paper at ICLR 2026
progress, with the cosine decay schedule emerging widely-adopted approach (Loshchilov & Hutter 2016, Kaplan et al., 2020) Hoffmann et al. 2022). This method features an initial warm-up phase followed by cosine-based decay, both constrained by predeternined total training duration. Consequently, any extension of the training process, such as incorporating new data, necessitates a complete restart from the beginning recalibrate the entire decay curve.
To address this inflexibility, the Warmup-Stable Sranuu--table-eeca (WSD) strategy (Hu et al. 2024 inserts a stable training phase with constant LR between the warmup and decay phases. This schedule pro- vides two key advantages: (1) flexible initiation of the decay phase independent of total step count, and (2) complete decoupling from fixed training durations. The WSD approach has demonstrated significant effectiveness, evidenced by its adoption in several recent LLMs including DeepSeek-
V3 (DeepSeek-Al et al. 2024) and ERNIE 4.5 JERNIH-Team 2025 However, WSD introduces new scheduling requirements: researchers must manually decide when to initiate the decay, how many tokens to allocate for it, and which decay function (e.g., cosine, linear) to employ. Moreover, if training needs to be extended after the decay has commenced, one must roll back the training to the state preceding the decay phase and re-design the decay strategy. This dependency on a manually configured schedule counteracts the goal of a fully autonomous and continuous training process.
To minimize scheduling complexity, recent research has investigated alternative approaches that completely eliminate the decay phase from LR schedules (Defazio et al., 2024 Song et al. 2025 Zhang et al., 2025). A prominent direction in this field explores weight averaging (a.k.a., model merging) techniques. Empirical results demonstrate that simply maintaining a constant LR combined with standard weight averaging strategie-ssud as exponentially weighted averaging (EWA)-can achieve performance competitive with WSD-based schedules (Li et al. 2025). Build- ing upon this line of research, we make three key extensions:
(1) We present Warmup-Star Stable and Merge (WSM), a simple yet general framework for LR schedul- ing. By formalizing the connection between LR decay and checkpoint merging, we demonstrate that WSM can be instantiated to emulate various decay strategies stratgiis--inlludnng cosine decay, linear decay and inverse square root (1-sqrt) decay. Our framework provides a principled approach to convert any LR decay method into a theoretically approximate model averaging implementation. This contrasts with prior work (Defazio et al. 20244 [Song et al., 2025 Li et al. 2025), which has largely focused on analysis of specific averaging strategies and their optimization properties. Notably, our framework is celimizer-agnostin enabling seamless integration with various optimization algorithms (e.g., SGD, Adam) without requiring modifications on the underlying training pipeline.
(2) Our extensive experiments systematically investigate key factors for instantiating this framework, including merge methods, merge frequency, duration, granularity, and the compatibility of merging and decay strategies. Our findings reveal that, aside from the introduction of high-quality annealing data during the decay (i.e., merge) phase, merge duration the training period covered by the merged eheckpoints-emerges as the most critical factor influencing model performance, with significantly greater impact than both the checkpoint interval and the number of merged models.
(3) The proposed decay-free LR schedule delivers substantial performance gains. Extensive em- pirical evaluation shows consistent improvements across multiple benchmarks: +3.5% on MATH, +2.9% on HumanEval, and +5.5% on MMLU-Pro-a significant advance over the WSD method, while prior work often only matches WSD's performance. Moreover, these benefits naturally ex- tend to the post-training stage, demonstrating our method's potential for sustained improvements in long-term model refinement.
The proposed WSM framework presents a promising direction for developing effective decay-free LR schedules. Our results demonstrate consistent performance advantages across both pre-training and fine-tuning stages. Notably, WSM enables the implementation of sophisticated decay-like meth- ods through simple and stable optimization approach, combining the benefits of strategic schedul- ing with robust training dynamics.
2 PRELIMINARY
Typically, existing mainstream LR schedulers consist of two phases: initial warm-up followed by decay. During the warm-up. the LR increases linearly from a small value to a peak value, Ir pepaaks over T Taarmup steps. This helps stabilize the optimization process in the early stages of training.
2

---

Published as a conference paper at ICLR 2026
Following the warm-up, the LR gradually decreases according to predefined function, such as cosine, linear, or inverse square root decay. These schedules share a critical limitation: they require the total number of training tokens (or steps), T Tmaxs to be known in advance. For instance, the prevalent cosine LR schedule is formulated as:
peakk
if Twormap
Twarmup
Ir(t)=
2 peak
+1 cos
x(t--varmuu) Timxx Twarmup
if 2 wormaup
The Warmup-Stable Stable-Decay (WSD) LR schedule introduces stable phase between warm-up and decay, maintaining constant LR its peak (Ir peak), which defined as:
peak:
if < T Twarmuu
Twarmup
br(t)=
peak decay function(()
if T warmup St Tdecay start if decaay.start St Tmaa
As shown, WSD facilitates flexible experimentation by supporting multiple decay attempts from the endpoint of the stable thase-withone requiring reset to the initial state. Despite this flexibility, it still requires predefined decay-phase settings-such as the decay start step Edacoogsstrrl decay function (decay function), (deccyyff and total training steps (Tmoz).
3 THE PROPOSED METHODOLOGY
In this section, first establish the theoretical connection between checkpoint merging and LR decay, formalize our proposed WSM (Warmup-Stable and Merge) schedule, and compare it with the widely used WSD schedule.
3.1 THEORETICAL CONNECTION BETWEEN LR DECAY AND CHECKPOINT MERGING
The core idea of checkpoint merging this work is to take an ordered checkpoints, [0,, (x+x)) and apply a merge function to generate single model 0ntx. Here, 0, e Rd represents the model's parameter vector at the i-th training iteration. The most general form is weighted average of the checkpoints:
0uth~ CysOn+s
j=0
where {cj} are non-negative weights that sum to one, i.e., Σj=o = This formulation obscures a deeper connection to the training dynamics. We can reveal this connection by expressing each checkpoint terms of an initial checkpoint 9, and the subsequent gradient updates. For simplicity, we assume the updates between checkpoints at different time steps are independent and ignore op- timizer states. Let di the gradient update vector (including the LR) at step i, such that the model Thus, an intermediate checkpoint fn+z can be expressed as the sum of an initial state f and the sequence of negative gradient updates that followed:
updates 0 = 0g
O On Σ 1=1
Inti-1
(2)
Furthermore, we can substitute Eq. 2 into the general merging formula (Eq, T and then rearrange the double summation by changing its order. A gradient update Jn44-1 is included in the sum for all involved checkpoints Ontj where j i.
0mth=
Σe,
Gi
0n
Σ E=1
Dati-1
0n
Σ
Σ
Cj
Inti-1
3

---

Published as a conference paper at ICLR 2026
Nit 2110-1-011-
WTO
w10=1-x7
Takking Progression
Traning Progeesioon (2)
Training Progresae ression
Training Progression 10
Velgttt
Herge Weights
Weights
0.35
Wogght Volue
120
6/1
Index
Checkpeint Index (b): Linear Decay
Checkpoint Index (c): Costine Decky
Checkpont Index (d): 1 --srrt Decay
(a): EMA
Figure 2: Visualization of checkpoint merging weight distributions and their corresponding decay functions over span of 10 checkpoints (larger checkpoint indices denote recent checkpoints). (a) Exponential Moving Average (EMA) weights, exhibiting convex decay characteristics; (b) Uni- form averaging weights, demonstrating linear decay behavior; (c) and (d) Our theoretically derived weights, designed to approximate cosine and 1-sqrt decay patterns. Note: Decay curves are ren- dered smoothly for clarity, though the underlying weights are applied discretely at each checkpoint.
This shows that weight coefficient of Σ}= Cj is applied to the gradient update In+i-1 from step n + i 1. By defining new set of weights for the gradient updates, w; = ΣS ; Cj. we arrive at the final equivalent expression for checkpoint merging:
θ-- = On Σ wi Inti-1
(4)
This equation shows that merging checkpoints with weights ej is equivalent to applying synthetic decay schedule defined by weights w; to the gradients accumulated after the base checkpoint 0,, where mapping exists between the merge weights cj and the effective learning rates w;- We there- fore propose Theorem 3.1 which can approximate monotonically decreasing decay curve. Figure 2 illustrates various checkpoint merging weights and their corresponding decay curves. Additional details are provided in Appendix B
Theorem 3.1 (Checkpoint Weight Derivation from Gradient Decay Schedule). Given a desired sequence of gradient decay coefficients {wi)r=r that is monotonically non-increasing and bounded, 1 2 Z W2 2 ... 2 wx O, the corresponding non-negative checkpoint weights {ez {cr} required to satisfy Eq. are uniquely determined by:
C = Wk
= Wi Wj+1s
for j E [1,
(5)
CD 1 13=1C== W1
3.2 EMULATING LR DECAY THROUGH CHECKPOINT MERGING
The central hypothesis of WSM that the optimization benefits of LR decay can be decoupled from the live training process and instead be effectively achieved through the merging of model checkpoints (Kingma & Ba 2015, DeepSeek-Ar Al et al. 2024 Li et al. 2025). The WSM simplifies the LR schedule by completely eliminating the decay phase:
Ir peak Twarnuu
ift T warmap
(r())=
=
(rpeak
if > T worrmup
The WSM pipeline, detailed in Algorithm [ operates in two primary phases. It begins with stan- dard warmup phase, where the learning rate linearly increases to its peak value, peak- Following this, the process enters the main stable training phase, where the learning rate is held constant.
4

---

Published as a conference paper at ICLR 2026
After specified step Tswitch, the model transition from the general pre-trainn pre-training data D to smaller, high-quality annealing dataset D Donneat, allowing the "annealing" to focus on curated data mixture. Throughout this stable phase, checkpoints chronous merging process continuously fetches the last n checkpoints from storage and combines them into model W merged. Specifically, for the Merge(.) operation, can select various decay strategies to emulate (e.g., the decay curve shape and minimum LR), calculate the corresponding gradient decay coefficients {w;}, and then derive the checkpoint merging weights {ci} based on Theorem 3.1 This merged checkpoint, which emulates the effect of decay schedule, provides a robust, annealed model without ever altering the live learning rate.
periodically saved. Concurrently, an asyn-
EXPERIMENT
Next, we present the empirical evaluation of our proposed WSM validate its effectiveness.
4.1 EXPERIMENT SETUP
The model we used for the main experiment is a standard MoE model with a total of 16.3 billion parameters and 1.4 billion active parameters. We utilized the AdamW optimizer (Loshchilov & Hutter 2019), and the hyperparameters are set to 31 = 0.9 and 3g = 0.95, with 0.1 weight decay applied. Through preliminary scaling laws experiments, we set the peak LR and batch size 4.78c- 4 and 2048, respectively. Comprehensive details regarding our model architecture, specific training parameters, dataset composition, and evaluation protocols are provided in Appendix [A]
We begin with checkpoint pretrained on 10.2 trillion tokens from the pretraining dataset using a stable, constant L.R Then, we continue training for an additional 400B tokens on a specialized high-quality annealing dataset, following two distinct strategies to evaluate their relative effective- ness: (1) We apply a conventional LR decay schedule to the model. This branch serves as our baseline, representing the standard Warmup-Stable- Decay (WSD) methodology. Note that em- ploy the standard re-warm-up strategy for continual pre-training as in a standard WSD setting. (2) We continue training with the same constant LR. The final model is then produced by merging the checkpoints saved during this stable phase. This branch represents our proposed WSM sched- ule. Unless otherwise specified, we save checkpoint every 25B tokens and mean averaging to merge the most recent checkpoints, which corresponds to a linear decay LR schedule.
4.2 OVERALL PERFORMANCE OF WSM SCHEDULE
Effectiveness in pre-training We evaluate our WSM schedule against the baseline WSD using three series of checkpoints: (1) those obtained using the standard LR decay schedule WSD, (2) checkpoints from the last stable phase of WSM before merging, and (3) our final merged checkpoints from the WSM method using various merge durations (window sizes). The category-wise average results are summarized in Figure 3 and Table [ and scores for each benchmark are provided in [
Our first and most significant finding is that the WSM method consistently outperforms WSD across the majority of tasks considered. On average, WSM achieves performance improvements across all benchmark categories. Notably, when comparing the best-performing checkpoint from each strat- egy, WSM improves upon WSD by an average of 1.3 points. We observe remarkable improvements of up 2.7 points on MATH, 2.4 points on HumanEval, 2.1 points on MMLU-Pro. These results provide compelling evidence that replacing the LR decay phase used in WSD with the checkpoint merging strategy of WSM is not only a feasible alternative but also a more effective approach for enhancing the diverse capabilities of the final pretrained model.
Long-term effectiveness for post-training To assess the long-term impact of WSM, we extended evaluation to the post-training phase. We apply supervised fine-tuning checkpoints generated by the WSM and WSD under identical settings for epochs. Results in Table 2 show that this advantage persists beyond the post-training phase.
5

---

Published as a conference paper at ICLR 2026
Overail Average
General Knowledge/ Rearoning
Language Understadding
Accurccy 150
Accurccy
Accurtey 150
Accuracy
Tokens(B)
Takensi8)
TakenstB) WSM IMerge 201
WSM (Bofore Merga)
W50
WSR (Mange 5)
WERI (Marge 12)
WSM (Merge 161
Figure 3: Comprehensive performance comparison (overall and by category) between our WSM schedule (via checkpoint merging, blue line) and standard WSD scheduling (via LR decay, red line). Notably, WSD requires a predefined decay schedule (e.g., over 400B tokens this study), whereas WSM eliminates this constraint. This flexibility enables seamless training continuation (gray regions) and allows WSM to approximate various decay curves.
Table 1: Base model performance comparison. Results reported based on the checkpoint with the highest average benchmark score.
General Knowl- edge
Professional Knowl- edge
Language Modeling
Overall Average
Math
Code
WSD
69.06
67.78 68.67
57.49 58.81
64.88 65.58
53.46 56.04
62.67
Base Model
WSM
70.22
63.95
Improv.
+1.68%
+1.31
+2.30%
+1.08%
+4.83%
+2.04%
Table 2: Instruct model performance comparison. Results are reported based on the epoch with the highest average benchmark score.
Overall Average
Language
Knowledge
Math
Code
Reason
Agent
WSD WSM
81.12
60.00
61.43
58.23 57.95
63.21 64.94
68.16 16
62.90
Model
84.78
61.73
62.28
69.33
64.07
Improv.
+4.51%
+2.88%
+1.38%
-0.48%
+2.74%
+1.72%
+1.86%
4.3 EMPIRICAL ANALYSIS OF WSM SCHEDULE
Next, we conduct comprehensive empirical study to dissect the WSM schedule. We aim to under- stand the key factors influencing its performance, its robustness across different training scenarios, its interaction with conventional decay schedule, and its broader implications on model dynamics.
4.3.1 ROBUSTNESS ACROSS PRE-TREE TRAINING PROCESS
Beyond applying WSM as a final step on a high-quality dataset, we also evaluated its utility and robustness throughout the entire pre-training lifecycle. To achieve this, conducted comparative analysis at various intermediate stages a long training run, comparing the performance of a model produced by our computationallly-ruge. WSM merging against one produced by initiating full,

---

Published as a conference paper at ICLR 2026
Mean Merging Strategy
1-sqrt Marging Strategy
Merging Strategy
64.0
Merge 4
Merge
Merge
Merge
63.5
Merge
Merge
Accuracy (%)
Merge 16
Merge 16
63.0
67.0
Merge
Merge
61.5
Marge
Merge 16
61.0
Merge 20
100
200
300
400
500
100
200
300
400
500
100
200
300
400
500
Tokens (B)
Tokens (B)
Tokens (B)
Figure 4: Merge duration analysis across algorithms. "Merge 4" indicates merging the most recent 4 checkpoints, with each checkpoint saved at 25B-toker -token intervals.
resource-intensive LR decay. As shown in Figure 5a although the performance gains over WSD not as significant as when switching to high-quality data, we find that the performance of models generated via WSM merging consistently and closely mirrors the results of true LR anneal. In the figure, the gray line represents the base model trained with a constant LR. The blue lines show the performance of WSM models, which were created by mean-merging four checkpoints from within the preceding 100B-token merge duration. The red lines represent full 100B-token decay runs initiated at the 2T, 4T, 6T, 8T, and 10T token milestones. This result demonstrates that the incorporation of high-quality annealing data is the key factor enabling WSM to outperform WSD, while simultaneously establishing WSM as a reliable, high-fidelity proxy for estimating a model's post-anneal potential at any point during training. Consequently, it can provide effective assessment throughout the pre- pre-training phase, eliminating the need to launch multiple, expensive decay runs to determine the model's 's ultimate performance.
4.3.2 IMPACT OF MERGING ALGORITHM
As derived Section 3.1 the checkpoint merging process can be viewed as an approximation of a LR decay schedule, where the weighting scheme of the merge is analogous to the functional form of the decay curve. For instance, a simple mean average is analogous to a linear decay curve. An Exponential Moving Average (EMA) would correspond to a convex exponential decay curve. Existing works (Hägele et al.) 2024) and our prior experiments (details are provided in Appendix (F) with the WSD schedule revealed led performance hierarchy among decay curves: concave schedules (e.g.. inverse square root) and linear schedules outperform convex schedules. Building on these findings, hypothesize that a merging algorithm designed to approximate decay curves that have been proven effective in WSD scheduling will similarly yield better results.
Based on Theorem 3.1 we experimentally compare three merging algorithms: one using our theorem-generte generated weights to approximate 1-sqrt decay, another using simple mean averaging, and a third using EMA. Our experimental results in Table 3 validate this hypothesis. While the merge method outperforms decay, we observe a distinct performance ranking: the 1-sgrt merge approach outperforms Mean, and both are markedly better than EMA. Crucially, this ranking (1-sqrt> > Mean > EMA) is identical to the hierarchy observed in standard LR decay schedules. This consistency demonstrates that WSM reliably reproduces the relative effectiveness of different decay schedules, confirming that checkpoint merging serves principled and effective simulation of the LR decay process.
4.3.3 IMPACT OF MERGING DURATION AND GRANULARITY
We further investigate in detail the impact of different merge durations (window sizes) on various al- gorithms. First, different merge durations essentially correspond to decaying with varying amounts of data. As shown in Figure 4 when comparing the best-performing checkpoints across the entire merging trajectory for both mean and 1-sqrt merging algorithms, larger merging windows tend to yield better results. However, this advantage gradually diminishes the window size increases. This observation aligns with previous practical LR decay experiments, where simply increasing the amount of annealing data shows diminishing returns and may eventually fail to provide further
7

---

Published as a conference paper at ICLR 2026
Average Puffomance
Overall Average Performanc
Accuracy 8
Accurayy ONL
2000
20:0
19423
1300
4032
5302
$350
2003
1990
3933
1393
1050
Decay then
(c)
Figure 5: Interplay of checkpoint merging and LR decay, (a) A comparison of checkpoint merging (WSM) and true LR decay (WSD) across long pre-training run. (b) Investigating the effect of applying merging within decay phase LR decay after checkpoint merging Seerge-then-Decyy)
(c) Investigating the effect of applying
improvement. For EMA merging, its performance is significantly inferior to other algorithms and shows no clear variation with merge durations. This may indicate that EMA not an effective merg- ing algorithm (also suggesting that the convex nature may not represent optimal decay curve).
Then we further investigate the granularity of checkpoint merging, i.c., the interval between saved checkpoints used for merging. Finer-grained merging represents more precise approximation of the true decay curve. The results in Table show that finer- mmer-grained merging tends to achieve better performance. However, frequent checkpoint saving imposes storage overhead, requiring careful trade-off considerations.
Table 3: Impact of merging algorithm.
General Knowledge
Language Modeling
Professional Knowledge
Overall
Math
Code
Average
Decay
1 1-sqrt
69.06
67.78
57.49
64.88
53.46
62.67
EMA
69.05
67.64
58.81
64.19
54.44
63.01
Merging
Mcan
70.22
68.67
58.81
65.58
56.04
63.95
1 1-sgrt
70.27
68.26
59.65
65.70
55.42
64.06
Table 4: Comparison of different saving/merging intervals within an 80B-token merge duration. For example, (6B,16) indicates that saving every 5B tokens while merging the latest 16 checkpoints.
Merging Granularity
General Knowledge
Language Modeling
Professional Knowledge
Overall
Math
Code
Average
(6)
69.46
79.95
57.07
64.68
53.82
63.63
(10B,8)
69.23
80.00
57.94
64.39
53.78
63.78
(20B,4)
69,29
80.29
56.79
63.87
54.19
63.36
(40B,2)
68.47
79.83
56.57
63.59
52.20
62.77
(80B,1)
67.47
64.98
55.07
61.69
51.61
60.33
4.3.4
ON THE COMPATIBILITY OF MERGE AND DECAY
Given that checkpoint merging effectively simulates decay, a natural question arises: can merg- ing and decay be combined to achieve synergistic performance gains? We investigate this by testing two hybrid approaches. (1) Decay-then-Merge:
first apply a standard decay schedule and then merge checkpoints selected from within the decay phase. (2) Merge-then-Decay: We further ap- ply a decay schedule to the resulting merged model. As shown in Figures (bb and 5c the hybrid approach failed to yield any improvement in either configuration, although the Merge-thent then-Decat model showed better performance at the beginning periment (Figure 5b), the blue stars represent the results of merging checkpoints selected along the
training. For the Decay-then. Becay-then-Merge ex-
8

---

Published as a conference paper at ICLR 2026
decay trajectories (red lines), which were initiated at various pre-training milestones (4T, 6T, 8T, 10T). For the Berge-then--ecay experiment (Figure 5e), we compare a decay run initiated from WSM model (blue linet-created by merging four checkpoints from the 9.8T to 10T token in- terval-against a standard decay baseline initiated from a single 10T checkpoint (red line). These results suggest that checkpoint merging and LR decay are not complementary but rather alternative pathways to similar optimization objective.
4.3.5 IMPACT ON MoE LOAD BALANCING
Table 5: Impact on MoE load balancing. The WSM strategy demonstrates improved expert uti- lization (lower load balancing violation scores) with slightly higher test language modeling loss. The mean global max violatio represents the average of the highest violations across all layers (measuring the severity of "overloaded" experts), while mean -loball global_minn.n mii mninvviolation averages the vi- olations for the least- ceast-utilized experts (measuring the risk of "routing collapse").
language modeling loss
mean global max violatio mean global min violation
WSD WSM
0.675
0.601 0.545
0.322
0.697
0.201
We provide an analysis of the implications of WSM on MoE router balance in Table Specifically, the violation for single expert is calculated as its relative deviation from the average load within its layer"] When comparing the merged checkpoint with decayed checkpoint, the merged checkpoint achieves more balanced routing, although its loss is slightly higher. We argue that this trade-off marginally higher loss for superior downstream performance- indicative of enhanced generaliza- tion rather than overfitting the training data.
RELATED WORK
5.1 LEARNING RATE SCHEDULE
Learning rate (LR) scheduling is critical for training performant models (Jin et al.,
2023 Gotmare al. (019)) Classic schedules like Cosine (Kaplan al., [2020] Ibrahim al., (024) or Lin-
ear (Defazio et al. (2023) decay adjust the LR based on predefined total training duration, which inflexible for continual training. The Warmup-Stable-Deean (WSD) schedule (Hu al.) 2024)
addresses this by introducing stable LR phase after warmup, decoupling the eventual decay from a fixed training length and offering greater flexibility for long or continuous training runs. More recently, researches have explored "schedule-free" methods to eliminate the decay phase entirely, which maintain a constant LR and instead leverage weight averaging techniques (Defazio et al, 2024| Song al. 2025 Zhang et al.) 2025). Builds upon these schedule- ffree principles, our work propose to replace WSD's decay phase with a post-hoe checkpoint merging operation instead of specific online averaging strategies. This simplifies the training pipeline and, by formalizing the connection between LR decay and checkpoint merging, allows decay strategy to be theoretically approximated, leading free offline exploration and enhanced model performance.
5.2
MODEL. MERGING
Model merging (Izmallvv et al. 2018 Wortsman et al. 2022) has emerged as efficient paradigm for model construction. This approach achieves effective knowledge transfer and performance im- provement through parameter-level integration of multiple models. Model merging is primarily utilized in two distinct scenarios: (1) The integration of knowledge and capabilities from multiple independently trained models into single parameter set, with the objective of preserving maximal performance from each specialized model (Aakanksha et al.) 2025) Ramé et al. 2024); and (2) the merging of checkpoints along a single training trajectory. This second category formally estab- lished in the literature as Stochastic Weight Averaging (S
et al. 2018). SWA func-
For a single layer, let L be the vector of token loads for its experts and μ mean L), max violatitt max (=)el and min violation [

---

Published as a conference paper at ICLR 2026
tions practical realization of Polyak Preraging(Polyakk & Juditsky 1992). acting smoothing
mechanism that reduces the noise inherent in stochastic gradient-basee based optimization (Sanyal et al,
2023
Liu al. 2024a Kaddour 2022; Li et al., 2023] Sandler et al.) 2023 Hägele et al. (024).
While Li et al. (2025) show model merging can achieve performance competitive with decay-based
schedule, these techniques have also demonstrated practical utility in industrial-scale L.M develop-
ment (Grattafiori et al. 2024 DeepSeek-AT et al. 2024 Aakanksha et al. 2025).
Concurrent with our work, the ERNIE 4.5 (ERNIE-Team 2025)
explored the relationship between EMA and LR decay. Our work generalizes this by introducing a formal theoretical framework that maps any LR decay schedule to principled averaging weights, and we empirically demonstrate that these non Mon-EMAA weights significantly outperform the sub-optimal EMA baseline.
Our work builds upon this SWA paradigm. We establish theoretical connection between this operation and LR decay and provide a principled approach to convert various LR decay strategies into a theoretically approximate model averaging implementation.
CONCLUSION
In this paper, we have presented WSM, a decay-free LR scheduling approach for LLM pre-training. Our method bridges LR decay and checkpoint merging by establishing the theoretical connection. By climinating the conventional decay phase, WSM simplifies LR scheduling while reformulat- ing various decay strategies as principled model averaging schemes. Through systematic analy- sis, we identified the incorporation of high-quality annealing data and merge duration as the most critical factors influencing model performance-outveee tensive experiments have demonstrated WSM's superiority over traditional WSD baselines in LLM pre-training, with consistent improvements across multiple benchmarks, robustness across different optimizers, and sustained benefits during fine-tuning. Consequently, WSM offers highly versa- tile approach that can be seamlessly integrated into existing training pipelines without additional complexity.
other implementation choices. Ex-
USE OF LARGE LANGUAGE MODELS
We utilized large language models (e-g., GPT-5 and Gemini-2.5-pro) during the preparation of this manuscript. Their assistance included editing the text to enhance clarity and readability, as well as generating the plotting code for the figures.
ACKNOWLEDGMENTS
This work was partially supported by the National Natural Science Foundation of China No. 92470205 and Beijing Major Science and Technology Project under Contract No. 10000084225000. This work was also supported by Ant Group Research Fund. Xin Zhao is the corresponding author.
REFERENCES
Aakanksha, Arash Ahmadian, Marwan Ahmed, Jay Alammar, Milad Alizadeh, Yazeed Alnumay, Sophia Althammer, Arkady Arkhangorodsky, Viraat Aryabumi, Dennis Aumiller, Raphaël Ava- los, Zahara Aviv, Sammie Bae, Saurabh Baji, Alexandre Barbet, Max Bartolo, Björn Bebensee, Neeral Beladia, Walter Beller-Morales. Alexandre Bérard, Andrew Berneshawi, Anna Bialas, Phil Blunsom, Matt Bobkin. Adi Bongale, Sam Braun, Maxime Brunet, Samuel Cahyawijaya, David Cairuz, Jon Ander Campos, Cassie Cao, Kris Cao, Roman Castagné, Julián Cendrero, Leila Chan Currie, Yash Chandak, Diane Chang, Giannis Chatziveroglou, Hongyu Chen, Claire Cheng, Alexis Chevalier, Justin Т. Chiu, Eugene Choi, Eujeong Choi, Tim Chung, Volkan Cirik. Ana Cismaru, Pierre Clavier, Henry Conklin, Lucas Crawhall-Stein, Devon Crouse, Fe- lipe Cruz-Salinas, Ben Cyrus, Daniel D souuz Hugo Dalla-Tore Torre, John Dang, William Dar- ling, Omar Darwiche Domingues, Saurabh Dash, Antoine Debugne, Théo Dehaze, Shaan De- sai, Joan Devassy, Rishit Dholakia, Kyle Duffy, Ali Edalati, Ace Eldeib, Abdullah Elkady, Sarah Elsharkawy, Irem Ergün, Beyza Ermis, Marzieh Fadaee, Boyu Fan, Lucas Fayoux, Yannis
10

---

Published as a conference paper at ICLR 2026
Flet-Berliac, Nick Frosst, Matthias Gallé, Wojeiech Galuba, Utsav Garg, Matthieu Geist, Mo- hammad Gheshlaghi Azar, Ellen Gilsenan- Seraphina Goldfarb-Tarrane Tomas Gold- sack, Aidan N. Gomez, Victor Machado Gonzaga, Nithya Govindarajan, Manoj Govindassamy, Nathan Grinsztajn, Nikolas Gritsch, Patrick Gu, Shangmin Guo, Kilian Haefeli, Rod Hajjar, Tim Hawes, Jingyi He, Sebastian Hofstätter, and Sungjin Hong. Command A: enterprise- large language model. CoRR, abs/2504.00655 006998 2025. doi: 10448550/RXXV22544.0000 https: http::/ooi..oogo org./.org/1000 11111..48550 48550/ arxiv 2504. 00698
URL
Joshua Ainslie, James Lee-Thorg Michiel de Jong. Yury Zemlyanskiy, Federico Lebrón, and Sumit Sanghai. GQA: training generalized multi-query transformer models from multi-head checkpoints. Houda Bouamor, Juan Pino, and Kalika Bali (eds.). Proceedings of the 2023 Conference on Empirical Methods in Natural Language Processing, EMNLP 2023, Singapore, December 6-10, 2023, pp. 4895-4901. Association for Computational Linguistics, 2023. doi: 10.18653/V EMNLP- 298. URL https:: //oii 1ttpo///o//..11////22200 em.1p p-lppppmaamnn 298
Orly Alter, Patrick O Brown, and David Botstein. Singular value decomposition for genome-wide expression data processing and modeling. Proceedings of the National Academy of Sciences, 97 (18):10101-10106, 2000.
Mohammad Bavarian, Heewoo Jun, Nikolas Tezak, John Schulman, Christine McLeavey, Jerry Tworek, and Mark Chen. Efficient training of language models fill in the middle. CoRR, abs/2207.14255. 2022. doi: 10.8555/ARXX112220..1555 URL https:: htt::://doi.oooe org/o.. 48550/arXis httpaaawwwi........... 2207 14255
Shane Bergsma, Nolan Simran Dey, Gurpreet Gosal, Gavia Gray, Daria Soboleva, and Joel Hest- ness. Straight to zero: Why linearly decaying the learning rate to zero works best for llms. In The Thirteenth International Conference on Learning Representations, ICLR 2025, Singapore, April 24-28, 2025. OpenReview.net, 2025. URL https: /ttp:::://///peeeeeeeeee net ftte/foummmii hrO1BgHaMM
Sumithra Bhakthavatsalam, Daniel Khashabi, Tushar Khot, Bhavana Dalvi Mishra, Kyle Richard- son, Ashish Sabharwal, Carissa Schoenick, Oyvind Tafjord, and Peter Clark. Think you have solved direct-answer question answering? try arc-da, the direct-answer A12 reasoning challenge.
CoRR, abs/2102.03315, 2021. URL https::: //t.xiiivv././///////aa///aaaav/ www.wwwwwvgg.....a....a... httg:/b//22000 03315
Yonatan Bisk, Rowan Zellers, Ronan Le Bras, Jianfeng Gao, and Yejin Choi. PIQA: reasoning about physical commonsense in natural language. In The Thirty-Fourty AAAI Conference on Artificial Intelligence, AAAI 2020, The Thirty-Second Innovative Applications of Artificial Intel- ligence Conference, IAAI 2020, The Tenth AAAI Symposium on Educational Advances Artifi- cial Intelligence, EAAI 2020, New York, NY, USA, February 72, 2020, pp. 7432-7439. AAAI Press, 2020, doi: 10.1609/AAAL V34105.6623 URLhttpp:://dot www.oiiorrg/0// 1609/aaai. v341 62.399
Mark Chen, Jerry Tworek, Heewoo Jun Qiming Yuan, Henrique Pondé de Oliveira Pinto, Jared Kaplan, Harri Edwards, Yuri Burda, Nicholas Joseph, Greg Brockman, Alex Ray, Raul Puri, Gretchen Krueger, Michael Petrov, Heidy Khlaaf, Girish Sastry, Pamela Mishkin, Brooke Chan, Scott Gray, Nick Ryder, Mikhail Pavlov, Alethea Power, Lukasz Kaiser, Mohammad Bavar- ian, Clemens Winter, Philippe Tillet, Felipe Petroski Such, Dave Cummings, Matthias Plap- pert, Fotios Chantzis, Elizabeth Barnes, Ariel Herbert- Voss, William Hebgen Guss, Alex Nichol, Alex Paino, Nikolas Tezak, Jie Tang, Igor Babuschkin, Suchir Balaji. Shantanu Jain, William Saunders, Christopher Hesse, Andrew N. Carr, Jan Leike, Joshua Achiam, Vedant Misra, Evan Morikawa, Alec Radford, Matthew Knight, Miles Brundage, Mira Murati, Katie Mayer, Pe- Welinder, Bob McGrew, Dario Amodei, Sam McCandlish, Ilya Sutskever, and Wojciech Zaremba. Evaluating large language models trained on code. CoRR, abs/2107.03377 2021. URL wttptttt///// htattts://axxv...t
abs /ass//110..m 03374
Karl Cobbe, Vineet Kosaraju, Mohammad Bavarian, Mark Chen, Heewoo Jun, Lukasz Kaiser, Matthias Plappert, Jemy Tworek, Jacob Hilton, Reiichiro Nakano, Christopher Hesse, and John Schulman. Training verifiers to solve math word problems. CoRR, abs/2110.14166. 2021. URL
https://erxxt wwwpeiiwwwgg/a////////oo 14168

---

Published as a conference paper at ICLR 2026
DeepSeek-AL, Aixin Liu, Bei Feng, Bing Xue, Bingxuan Wang, Bochao Wu, Chengda Lu, Cheng- gang Zhao, Chengqi Deng, Chenyu Zhang, Chong Ruan, Damai Dai, Daya Guo, Dejian Yang, Deli Chen, Dongjie Ji. Erhang Li, Fangyun Lin, Fucong Dai, Fuli Luo, Guangbo Hao, Guanting Chen, Guowei Li, H. 7hang, Han Bao, Hanwei Xu, Haocheng Wang. Haowei Zhang. Honghui Ding, Huajian Xin, Huazuo Gao, Hui Li, Hui Qu, I.. Cai, Jian Liang. Jianzhong Guo, Jiaqi Ni, Jiashi Li. Jiawei Wang, Jin Chen, Jingchang Chen, Jingyang Yuan, Junjie Qiu, Junlong Li, Junxiao Song, Kai Dong, Kai Hu, Kaige Gao, Kang Guan, Kexin Huang, Kuai Yu, Lean Wang, Lecong Zhang, Lei Xu, Leyi Xia, Liang Zhao, Litong Wang. Liyue Zhang. Meng Li, Miaojun Wang, Mingchuan Zhang, Minghua Zhang, Minghui Tang, Mingming Li. Ning Tian, Panpan Huang, Peiyi Wang. Peng Zhang, Qiancheng Wang. Qihao Zhu, Qinyu Chen, Qiushi Du, R. J. Chen, R. L.. Jin, Ruiqi Ge, Ruisong Zhang, Ruizhe Pan, Runji Wang. Runxin Xu, Ruoyu Zhang, Ruyi Chen, S. S. Li, Shanghao Lu, Shangyan Zhou, Shanhuang Chen, Shaoqing Wu, Shengfeng Ye, Shengfeng Ye, Shirong Ma, Shiyu Wang, Shuang Zhou, Shuiping Yu, Shunfeng Zhou, Shuting Pan. T. Wang, Tao Yun. Tian Pei, Tianyu Sun. W. L. Xiao, and Wangding Zeng. Deepseek- v3 technical report. CORR, abs/2412. 2024. doi: 0044850/ARXIV.2412.20 19437. URL
https:://6o..org/00.. 48550 w 2412 19437
Aaron Defazio, Ashok Cutkosky, Harsh Mehta, and Konstantin Mishchenko. Optimal linear decay learning rate schedules and further refinements. arXiv preprint arXiv:2310.07831, 2023.
Aaron Defazio, Xingyu Yang. Ahmed Khaled, Konstantin Mishchenko, Harsh Mehta, and Ashok Cutkosky. The road less scheduled. In NeurIPS, 2024.
Baidu ERNIE-Team. Ernie 4.5 technical report, 2025.
Akhilesh Gotmare, Nitish Shirish Keskar, Caiming Xiong, and Richard Socher. closer look at deep learning heuristics: Learning rate restarts, warmup and distillation. In 7th International Conference on Learning Representations, ICLR 2019, New Orleans. LA, USA, May 6-9, 2019. r11idd--1111gg
OpenReview.net, 2019.
Aaron Grattafiori, Abhimanyu Dubey, Abhinav Jauhri, Abhinav Pandey, Abhishek Kadian, Ahmad Al-Dahle, Aiesha Letman, Akhil Mathur, Alan Schelten, Alex Vaughan, al. The llama 3 herd of models. arXiv preprint arXiv:2407.21773, 2024.
Alex Gu, Baptiste Rozière, Hugh James Leather, Armando Solar-Lezama, Gabriel Synnaeve, and Sida Wang. Cruxeval: A benchmark for code reasoning, understanding and execution. In Forty-first International Conference on Machine Learning. ICML 2024, Vienna, Austria, July 21-27, 2024. OpenReview.net, 2024. URL https:: F Ffpg52ssvg
htttps://openreviiee http://ppeeevveeeeeee
net/ forum? d=
Alexander Hägele, Elie Bakouch, Atli Kosson, Loubna Ben Allal, Leandro von Werra, and Martin Jaggi. Scaling laws and compute-optimal training beyond fixed training durations. In NeurIPS, 2024.
Dan Hendrycks, Collin Burns, Steven Basart, Andy Zou, Mantas Mazeika, Dawn Song, and Jacob Steinhardt. Measuring massive multitask language understanding. In 9th International Confer- ence on Learning Representations, ICLR 2021, Virtual Event, Austria, May 3-7, 2021. OpenRe- view.net, 2021a. URL WtRLshtts:///ppp revi ew net/ forum? d=d 3Gm@
Dan Hendrycks, Collin Burns, Saurav Kadavath, Akul Arora, Steven Basart, Eric Tang, Dawn Song, and Jacob Steinhardt. Measuring mathematical problem solving with the MATH dataset. In Joaquin Vanschoren and Sai-Kit Yeung (eds.), Proceedings of the Neural Information Processing Systems Track on Datasets and Benchmarks 1, NeurlPS Datasets and Benchmarks 2021, December 2021, virtual, 2021b. URL htttpp//e.attssseeeeeeennnnaaasp becchss-bbbbmmrara proceeddigg..oooooeoeo.o.o...... wwwp:/eeeddingg..nuis........ http::/pw....arppp....cpppppe//p wwtpc//pppp//////e//e 20011 -bwtraattttt.. wow.acctrooudd.hccc.... html
Jordan Hoffmann, Sebastian Borgeaud, Arthur Mensch, Elena Buchatskaya, Trevor Cai, Eliza Rutherford, Diego de Las Casas, Lisa Anne Hendricks, Johannes Welbl, Aidan Clark, et al. Train- ing compute-optimal large language models. arXiv preprint arXiv:2203. 15556, 2022.
12

---

Published as a conference paper at ICLR 2026
Jack Hong, Shilin Yan, Jiayin Cai, Xiaolong Jiang, Yao Hu, and Weidi Xie. Worldsense: Evaluating real-world omnimodal understanding for multimodal llms. CoRR, abs/2502.003326 2025. doi: 10. 04326
48850/ARXXV.2502.4453 243266..0.
htRhhtts:: //ooi org/ 10. 48550
Shengding Hu, Yuge Tu, Xu Han, Chaoqun He. Ganqu Cui, Xiang Long, Zhi Zheng, Yewei Fang. Yuxiang Huang, Weilin Zhao, Xinrong Zhang, Zhen Leng Thai, Kai Zhang, Chongyi Wang. Yuan Yao. Chenyang Zhao, Jie Zhou, Jie Cai, Zhongwu Zhai, Ning Ding, Chao Jia, Guoyang Zeng, Dahai Li, Zhiyuan Liu, and Maosong Sun. Minicpm: Unveiling the potential of small language models with scalable training strategies. CoRR, abs/2404.06395, 2024. doi: A0....V./..XIV. 2404.06395. URL http:ttttss........t //oi: www.orggg.....o 48550/ (05550/arxi:0 2404.4046500 06395
Yuzhen Huang, Yuzhuo Bai, Zhihao Zhu, Junlei Zhang, Jinghan Zhang, Tangjun Su, Junteng Liu, Chuancheng Lv, Yikai Zhang, Jiayi Lei. Yao Fu, Maosong Sun, and Junxian He. C-eval: A multi-level multi-disciplin chinese evaluation suite for foundation models. In Alice Oh, Tristan Naumann, Amir Globerson, Kate Saenko, Moritz Hardt, and Sergey Levine (eds.), Advances in Neural Information Processing Systems 36: Annual Conference on Neural In- formation Processing Systems 2023, NeurlPS 2023, New Orleans, LA, USA, December 10 16, 2023, 2023. URL http:// cttpp//arrainaeccpapeeeaeeeee..a wtteeerffilee/pppe////// and
Ben.cmmrrs.htmm
Adam Ibrahim, Benjamin Thérien, Kshitij Gupta, Mats Richter, Quentin Gregory Anthony, Eu- gene Belilovsky, Timothée Lesort, and Irina Rish. Simple and scalable strategies to continu- ally pre-train large language models. Trans. Mach. Learn. Res., 2024, 2024. URL https:
//openzeview. aeapen/eview..et//orrmmmm
Pavel Izmailov, Dmitrii Podoprikhin, Timur Garipov, Dmitry P. Vetrov, and Andrew Gordon Wil- son. Averaging weights leads to wider optima and better generalization. In Amir Globerson and Ricardo Silva (eds.), Proceedings of the Thirty-Fourth Conference on Uncertainty in Artificial In- telligence, UAI 2018, Monterey, California, USA, August 6-10, 2018, pp. 876 885. AUAI Press, 2018. htLpttp::
Naman Jain, King Han, Alex Gu, Wen-Ding Li, Fanjia Yan, Tianjun Zhang, Sida Wang, Armando Solar-Lezama, Koushik Sen, and Ion Stoica. Livecodebench: Holistic and contamination free evaluation of large language models for code. In The Thirteenth International Conference on Learming Representations, ICLR 2025. Singapore. April 24-28, 2025. OpenReview.net, 2025. /
URL //ppnreiee/ee...ie..... net http://rrrmndddhdtae0e
Hongpeng Jin, Wenqi Wei, Xuyu Wang, Wenbin Zhang, and Yanzhao Wu. Rethinking learning rate tuning in the era of large language models. In 5th IEEE International Conference on Cog- nilive Machine Intelligence, CogMI 2023, Atlanta, GA, USA, November 1-4, 2023, pp. 112-121. IEEE, 2023, doi: 10. 1LLL109CCOMM158552222000.. 00022.. URL
1109/
CocMI58952 2023. 00025
Keller Jordan. Muon: An optimizer for hidden layers in neural networks. 2024. URL htttps://eeleerjoret
https://
ke kellerjordan. cwtrrjordanggithhee gith it/posts/mmootsooo github. httpluu...///otts/mmee muunn/ Accessed: 2024.
Mandar Joshi, Eunsol Choi, Daniel S. Weld, and Luke Zettlemoyer. Triviaqa: A large scale distantly supervised challenge dataset for reading comprehension. In Regina Barzilay and Min- Yen Kan (eds.), Proceedings of the 55th Annual Meeting of the Association for Computational Linguistics, ACI. 2017, Vancouver, Canada, July 30 August 4, Volume I: Long Papers, pp. 1601-1611. Association for Computational Linguistics, 2017. doi: (0.18655/V1/117-11447 URL https:
//doi. 18133/100..6333/1/11111
Jean Kaddour. Stop wasting my time! saving days of imagenet and BERT training with latest weight averaging. CoRR, 0tw/2209.1498/. 2022. doi: (0.4855//RXXI.2209.. 14981.. URL. https:
//doi. //doi.org/10 48550/ark3 2209. 14981
Jared Kaplan, Sam McCandlish, Tom Henighan, Tom B Brown, Benjamin Chess, Rewon Child, Scott Gray, Alec Radford, Jeffrey Wu, and Dario Amodei. Scaling laws for neural language models. arXiv preprint 2rX1v:2001.08361, 2020.
13

---

Published as a conference paper at ICLR 2026
Diederik P. Kingma and Jimmy Ba. Adam: A method for stochastic optimization. In Yoshua Bengio and Yann LeCun (eds.), 3rd International Conference on Learning Representations, ICLR 2015, San Diego, CA, USA, May 7-9, 2015, Conference Track Proceedings. 2015. //arxiv. wwrxiv..rg/abb///0
http: URLhttp:
6980
Tom Kwiatkowski, Jennimaria Palomaki, Olivia Redfield, Michael Collins, Ankur P. Parikh, Chris Alberti, Danielle Epstein, Illia Polosukhin, Jacob Devlin, Kenton Lee, Kristina Toutanova, Llion Jones, Matthew Kelcey, Ming-War Wei Chang, Andrew M. Dai, Jakob Uszkoreit, Quoe Le, and Slav Petrov. Natural questions: benchmark for question answering research. Trans. Assoc. Comput. Linguistics, 7:452-466, 2019. doi: 10.1162/TACL) A) 00276. URL https: //op::://do...oo org/ 10. 1il./tt2accc.. a 00276
Guokun Lai, Qizhe Xie, Hanxiao Liu, Yiming Yang, and Eduard H. Hovy. RACE: large-scale reading comprehension dataset from examinations. In Martha Palmer, Rebecca Hwa, and Se- bastian Riedel (eds.), Proceedings of the 2017 Conference on Empirical Methods in Natural Language Processing, EMNLP 2017, Copenhagen, Denmark, September 9-11, 2017, pp. 785- 794. Association for Computational Linguistics, 2017. doi: 10.18653/V1/D17-10333 URL https: wtt../....../....13/11111111
1082
Haonan Li, Yixuan Zhang, Fajri Koto, Yifei Yang, Hai Zhao, Yeyun Gong, Nan Duan, and Timothy Baldwin. CMMLU: measuring massive multitask language understanding in chi- nese. In Lun-Wei Ku, Andre Martins, and Vivek Srikumar (eds.), Findings of the Associa- tion for Computational Linguistics, ACL 2024, Bangkok, T'hailand and virtual meeting, August 11-16, 2024, pp. 11260-11285 Association for Computational Linguistics, 2024a. doi: 10. 18653N1/224.FNNDINNSS ACL.677.. URL findings-ace www.ings-ccl..0 671
Qintong Li, Leyang Cui, Xueliang Zhao, Lingpeng Kong, and Wei Bi. Gsm-plus: A compre- hensive benchmark for evaluating the robustness of Ilms as mathematical problem solvers. In Lun-Wei Ku, Andre Martins, and Vivek Srikumar (eds.). Proceedings of the 62nd Annual Meet- ing of the Association for Computational Linguistics (Volume 1: Long Papers), ACL 2024, Bangkok, Thailand, August 11-16, 2024, pp. 22944 Association for Computational Lin- guistics, 2024b. doi: 100.66//11222AA0--LLLOOO
http:t/::: http://ooioooooo org/10.
I8653/v1/2024 acl long. 163
Tao Li, Zhehao Huang, Qinghua Tao, Yingwen Wu, and Xiaolin Huang. Trainable weight averaging: Efficient training by optimizing historical solutions. In The Eleventh International Conference Learning Representations. ICLR 2023, Kigali, Rwanda, May 1-5, 2023. OpenReview.net, 2023. net www./eeenrttee.. worummrrmmmmmm id 1wdddddwwwww------
URL httpttttt////eeeeeee
Yunshui Li, Yiyuan Ma, Shen Yan, Chaoyi Zhang, Jing Liu, Jianqiao Lu, Ziwen Xu, Mengzhao Chen, Minrui Wang. Shiyi Zhan, Jin Ma. Xunhao Lai, Deyi Liu, Yao Luo, Xingyan Bin, Hongbin Ren, Mingji Han, Wenhao Hao, Bairen Yi, LingJun Liu, Bole Ma, Xiaoying Jia, Xun Zhou, Siyuan Qiao, Liang Xiang, and Yonghui Wu. Model merging in pre-training of large language doi: 10.48550/ARXX1.5555.222 doi www.oogg....... 1w.g///0.00 48550/550480 /ac.00raxiii arXiv. 2505.12082
models. CoRR, abs/2505. 11055005.12822 1082. 2025.
URL https://
Deyuan Liu, Zecheng Wang, Bingning Wang, Weipeng Chen, Chunshan Li, Zhiying Tu, Dianhui Chu, Bo Li, and Dianbo Sui. Checkpoint merging via bayesian optimization in LLM pretraining. CoRR, abs/2403. 19390, 2024a. doi: 10,48550/ARXIV.24220 19390. URLhttps pecttss://ee ///oi. wtt.doooooo. / 10. 48550/ awtp/ariivvvv. 2403. 9390
Hongwei Liu, Zilong Zheng, Yuxuan Qiao, Haodong Duan, Zhiwei Fei, Fengzhe Zhou, Wenwei Zhang, Songyang Zhang, Dahua Lin, and Kai Chen. Mathbench: Evaluating the theory and ap- plication proficiency of llms with a hierarchical mathematics benchmark. Lun-Wei Ku, Andre Martins, and Vivek Srikumar (eds.), Findings of the Association for Computational Linguistics, ACL 2024, Bangkok, Thailand and virtual meeting. August 11-16, 2024. pp. 6884-6915. Associ- ation for Computational Linguistics, 2024b. doi: 10. https://dof wrw.bi..orgg/... wwt.oo//..........1.....
URL
findinginnngngnne .11.411
Jiawei Liu, Chunqiu Steven Xia, Yuyao Wang, and Lingming Zhang. Is your code generated by chatgpt really correct? rigorous evaluation of large language models for code generation. In
14

---

ELRA and Published as a conference paper at ICLR 2026
Alice Oh, Tristan Naumann, Amir Globerson, Kate Saenko. Moritz Hardt, and Sergey Levine (eds.), Advances in Neural Information Processing Systems 36: Annual Conference on Neural Information Processing Systems 2023, NewrIPS 2023, New Orleans, LA, USA, December 10 16, 2023, 2023. URL http::/://pppppppp.ppppppppppppp
Jingyuan Liu, Jianlin Su, Xingcheng Yao, Zhejun Jiang, Guokun Lai, Yulun Du, Yidao Qin, Weixin Xu, Enzhe Lu, Junjie Yan, Yanru Chen, Huabin Zheng. Yibo Liu, Shaowei Liu, Bohong Yin, Weiran He, Han Zhu, Yuzhi Wang, Jianzhou Wang. Mengnan Dong, Zheng Zhang, Yongsheng Kang, Hao Zhang, Xinran Xu, Yutao Zhang, Yuxin Wu, Xinyu Zhou, and Zhilin Yang. Muon is scalable for IIm training, 2025. URL http:: http:://axivv.vo org/ abs /2002.2502222 16982
Ilya Loshchilov and Frank Hutter. Sgdr: Stochastic gradient descent with warm restarts. arXiv preprint arXiv: 1608.0398,, 2016.
Ilya Loshchilov and Frank Hutter. Decoupled weight decay regularization. 7th International Conference on Learning Representations, ICLR 2019, New Orleans, LA, USA, May 6-9, 2019. OpenReview.net, 2019.
net htt:/forummmiii i Bkg&RicqY)
Kaijing Ma, Xeron Du Yunran Wang, Haoran Zhang, Zhoufutu Wen, Xingwei Qu, Jian Yang, Jia- heng Liu, Minghao Liu, Xiang Yue, Wenhao Huang, and Ge Zhang. Kor-bench: Benchmarking language models on showledge-orthogonee reasoning tasks. The Thirteenth International Con- ference on Learning Reprexentations, ICLR 2025, Singapore, April 24-28, 2025. OpenReview.net, 2025.
net/ Stw.ummmiiiissRRRsoo00000
Todor Mihaylov, Peter Clark, Tushar Khot, and Ashish Sabharwal. Can a suit of armor conduct electricity? A new dataset for open book question answering. In Ellen Riloff, David Chiang, Julia Hockenmaier, and un'ichii Tsujii (eds.), Proceedings of the 2018 Conference on Empirical Methods in Natural Language Processing, Brussels, Belgium, October 31 November 4, 2018, pp. 2381-2391. Association for Computational Linguistics, 2018. doi: 181188/////D11 1-260000
URL https: / /oo. 11t.3/.0...6.vvvv. c18 1260
Qiwei Peng, Yekun Chai, and Xuhong Li. Humaneval-xl: A multilingual code generation bench- mark for Pross-lingual natural language generalization. In Nicoletta Calzolari, Min- Yen Kan, Véronique Hoste, Alessandro Lenci, Sakriani Sakti, and Nianwen Xue (eds.), Proceedings of the 2024 Joint International Conference on Computational Linguistics, Language Resources and Evaluation, LREC/COLING 2024, 20-25 May. 2024, Torino, Italy, pp. ICCL, 2024. URL htt.tp::/aaaaaatttoaaat/ www.onoooonnnoooooooooooooooooco org./024 www.errrccmmamaa 735
Boris Polyak and Anatoli Juditsky. Acceleration of stochastic approximation by averaging. SIAM journal on control and optimization, 30(4):838 855, 1992.
Pranav Rajpurkar, Robin Jia, and Percy Liang. Know what you don't 't know: Unanswerable questions for squad. In Iryna Gurevych and Yusuke Miyao (eds.), Proceedings of the 56th Annual Meeting of the Association for Computational Linguistics, ACL 2018, Melbourne, Australia, July 15-20, 2018, Volume 2: Short Papers, pp. 784-789. Association for Computational Linguistics, 2018. doi: 10.18653/V1/P18-22
URL htt.s: //ac/aaactttttt
ant www.gogg........ 810-P188-1 2t////1881111111 2124/
Alexandre Ramé, Johan Ferret, Nino Vieillard, Robert Dadashi, Léonard Hussenot, Pierre-Loute Cedoz, Pier Giuseppe Sessa, Sertan Girgin, Arthur Douillard, and Olivier Bachem. WARP: on the benefits of weight averaged rewarded policies. CoRR, abs/2406. 16768 2024. doi: 10.48550/ 48544/855///// 2406.2206660 16768
ARXIVV2206.16777 URL htttps:://ol.orrgg
David Rein, Betty Li Hou, Asa Cooper Stickland, Jackson Petty, Richard Yuanzhe Pang, Julien Dirani, Julian Michael, and Samuel R. Bowman. GPQA: A graduate-level google-proof q&a benchmark. CoRR, abs/2311. 2023. doi: 1044850/ARXXV22111..22 URL https:
7/doi. ot0/10.rg/1... 085.444550//// 2311. 12022
Olivier Roy and Martin Vetterli. The effective rank: A measure of effective dimensionality. In 2007 15th European signal processing conference, pp. 606-610. IEFF, 2007.
15

---

Published as a conference paper at ICLR 2026
Keisuke Sakaguchi, Ronan Le Bras, Chandra Bhagavatula, and Yejin Choi. Winogrande: an adversarial winograd schema challenge at scale. Commuun ACM, 64(9):99-106, 2021. doi: 10.1145/3474381.
Mark Sandler, Andrey Zhmoginov, Max Vladymyrov, and Nolan Miller. Training trajectories, mini- batch losses and the curious role of the learning rate. CoRR, 02//3311023111 2023.
Sunny Sanyal, Atula Neerkaje, Jean Kaddour, Abhishek Kumar, and Sujay Sanghavi. Early weight averaging meets high learning rates for I1m pre-training. arXiv preprint arXiv:2306.03241, 2023.
Freda Shi, Mirac Suzgun, Markus Freitag, Xuezhi Wang, Suraj Srivats, Soroush Vosoughi, Hyung Won Chung, Yi Tay, Sebastian Ruder, Denny Zhou, Dipanjan Das, and Jason Wei. Lan- guage models are multilingual chain-of-thought reasoners. In The Eleventh International Confer- ence on Learning Representations, ICLR 2023, Kigali, Rwanda, May 1-5, 2023. OpenReview.net, 2023. URL https: htttt:::///penreiiew.. net wttp://wwwww.eettootttamaaa d- FR3wGCk IXp
Minhak Song, Beomhan Baek, Kwangjun Ahn, and Chulhee Yun. Through the river: Understanding the benefit of schedule- free methods for language model training. In biigh-dimensionl Learning Dynamics 2025, 2025.
Jianlin Su, Murtadha H. M. Ahmed, Yu Lu, Shengfeng Pan, Wen Bo, and Yunfeng Liu. Roformer: Enhanced transformer with rotary position embedding. Neurocomputingg 568:127063, 2024. doi: 10.1016//NEEUCOM.22......00
URL httttttt////o/tt wrg/ 0000:/00010..00100.111116...... neucom. neucom.
2023. 127063
Mirac Suzgun, Nathan Scales, Nathanael Schärli, Sebastian Gehrmann, Yi Tay, Hyung Won Chung, Aakanksha Chowdhery, Quoc V. Le. Ed H. Chi. Denny Zhou, and Jason Wei. Challenging big- bench tasks and whether thain-of-thought can solve them. In Anna Rogers, Jordan L. Boyd- Graber, and Naoaki Okazaki (eds.), Findings of the Association for Computational Linguisties: ACL 2023, Toronto, Canada, July 9-14, 2023, pp. 13003-13051, Association for Computational Linguisties, 2023. doi: 10.18653/V1/0023.FII FWW.:OOO.FIIIIINGGSS ACL.88, URL https: I0.18653/v1/2000 (.t.02...inndii findings a21.824
Zhengyang Tang, Xingxing Zhang, Benyou Wang. and Furu Wei. Mathscale: Scaling instruction tuning for mathematical reasoning. In Forty-first International Conference Machine Learn- ing, ICML 2024, Vienna, Austria, July 21-27, 2024. OpenReview.net, 2024. URL https:
//openzeview.o wwt.://nttt.oouumooemmee www.roorrrii...oo.o.oo.... jww 72N47M
Ning Tao, Anthony Ventresque, Vivek Nallur, and Takfarinas Saber. Enhancing program synthesis with large language models using many-objective grammar-guided genetic programming. Algo- rithms, 17(7):287, 2024. doi: (..3390/A17070222 URI. htttps://dci wrw.oorgg//000.... 10/p/11.33300 3390/ a17070287
M.-A-P. Team, Xinrun Du, Yifan Yao, Kaijing Ma. Bingli Wang, Tianyu Zheng, Kang Zhu, Minghao Liu, Yiming Liang, Xiaolong Jin, Zhenlin Wei, Chujie Zheng, Kaixin Deng, Shian Jia, Sichao Jiang. Yiyan Liao, Rui Li. Qinrui Li, Sirun Li. Yizhi Li, Yunwen Li, Dehua Ma, Yuansheng Ni, Haoran Que, Qiyao Wang. Zhoufutu Wen, Siwei Wu, Tianshun Xing, Ming Xu, Zhenzhu Yang, Zekun Moore Wang, Jun Zhou, Yuelin Bai, Xingyuan Bu, Chenglin Cai, Liang Chen, Yifan Chen, Chengtuo Cheng. Tianhao Cheng, Keyi Ding, Siming Huang, Yun Huang, Yaoru Li, Yizhe Li, Zhaoqun Li, Tianhao Liang, Chengdong Lin, Hongquan Lin, Yinghao Ma, Tianyang Pang, Zhongyuan Peng, Zifan Peng, Qige Qi. Shi Qiu, Xingwei Qu, Shanghaoran Quan, Yizhou Tan, Zili Wang, Chenqing Wang, Hao Wang, Yiya Wang. Yubo Wang, Jiajun Xu, Kexin Yang. Ruibin Yuan, Yuanhao Yue, Tianyang Zhan, Chun Zhang, Jinyang Zhang, Xiyue Zhang, Xingjian Zhang, Yue Zhang, Yongchi Zhao, Xiangyu Zheng, Chenghua Zhong, Yang Gao, Zhoujun Li. Dayiheng Liu, Qian Liu, Tianyu Liu, Shiwen Ni, Junran Peng. Yujia Qin, Wenbo Su, Guoyin Wang, Shi Wang, Jian Yang, Min Yang, Meng Cao, Xiang Yue, Zhaoxiang Zhang, Wangchunshu Zhou, Jiaheng Liu, Qunshu Lin. Wenhao Huang, and Ge Zhang. Supergpqa: Scaling LLM evaluation doi: 10.48550/ARXIV25500.. (185500/aX... 2502. 22502.14733 14739
across 285 graduate disciplines. CoRR, abs/2502. 14739, 2025.
14739. URL
Laurens der Maaten and Geoffrey Hinton. Visualizing data using t-SNE. Journal of Ma- chine Learning Research. 9:2579-2605, 2008, URLhtpp://www jmir. owg/papers//s/
vaane rranteermettten8e html
16

---

Published as a conference paper at ICLR 2026
Xiao Wang, Colin Wei, Yuxin Xiao, Alex Zhang, and Tengyu Ma. How set AdamW's weight decay as you scale model and dataset size. arXiv preprint arXiv:2405.13382, 2024a
Yubo Wang, Xueguang Ma. Ge Zhang, Yuansheng Ni, Abhranil Chandra, Shiguang Guo, Weiming Ren, Aaran Aruiraj, Xuan He. Ziyan Jiang, Tianle Li, Max Ku, Kai Wang, Alex Zhuang, Rongqi Fan, Xiang Yue, and Wenhu Chen. Mmlu-pro: A more robust and chal- lenging multi-task language understanding benchmark. In Amir Globersons, Lester Mackey, Danielle Belgrave, Angela Fan, UIrich Paquet, Jakub M. Tomczak, and Cheng Zhang (eds.), Advances in Neural Information Processing Systems 38: Annual Conference on Neural In- formation Processing Systems 2024, NeurlPS 2024, Vancouver, BC, Canada, December 10 15, 2024, 2024b. URL http: http://apprsssse nips cc cc/paper.pppeeeeeeeceeeeeece..ee Abssttaaaa-bbbtaaaataaaa -aateetattssssst and
Berchmarks_ Track_rrcckk.. html
Tianwen Wei, Jian Luan, Wei Liu, Shuang Dong, and Bin Wang. CMATH: can your language model pass chinese elementary school math test? CoRR, abs/2306. 16636, 2023. doi: 10.48550/ARXIV. 16636 URL httts:tt::///
/10/10. 48550/ 2306. 16636.
Kaiyue Wen, Zhiyuan Li, Jason Wang, David Hall, Percy Liang, and Tengyu Ma. Understanding warmup-stable-decay learning rates: A river valley loss landscape perspective. arXiv preprint arXiv:2410.05192, 2024.
Mitchell Wortsman, Gabriel Ilharco, Samir Yitzhak Gadre, Rehecca Roelofs, Raphael Gontijo Lopes, Ari S. Morcos, Hongseok Namkoong, Ali Farhadi, Yair Carmon, Simon Kornblith, and Ludwig Schmidt. Model soups: averaging weights of multiple fine-tuned models im- proves accuracy without increasing inference time. In Kamalika Chaudhuri, Stefanic Jegelka, Le Song, Csaba Szepesvári, Gang Niu, and Sivan Sabato (eds.), International Conference on Machine Learning, ICML 2022, 17-23 July 2022, Baltimore, Maryland, USA, volume 162 of Proceedings of Machine Learning Research, pp. 239998. PMLR, 2022. URL https: //proceedings ml www.m.............mmmmmm.m..m..e press vv 62 http://otssaaa22......m.
Rowan Zellers, Ari Holtzman, Yonatan Bisk, Ali Farhadi, and Yejin Choi. Hellaswag: Can ma- chine really finish your sentence? In Anna Korhonen, David R. Traum, and Lluís Marquez (eds.), Proceedings of the 57th Conference of the Association for Computational Linguistics, ACL 2019, Florence, Italy, July 28- August 2. 2019, Volume I: Long Papers, pp. 4791- 4800. Association for Computational Linguistics, 2019. doi: 10.18653/V1//199144 1472. URL org.10. SBl.c..16333v/ee
Hanlin Zhang, Depen Morwani, Nikhil Vyas, Jingfeng Wu, Difan Zou, Udaya Ghai, Dean P. Foster, and Sham M. Kakade. How does critical batch size scale in pre-training? ICLR. OpenRe- view.net, 2025.
Jingzhao Zhang, Tianxing He, Suvrit Sra, and Ali Jadbabaie. Why gradient clipping accelerates training: A theoretical justification for adaptivity. In 8th International Conference on Learning Representations, ICLR 2020, Addis Ababa, Ethiopia, April 26-30, 2020. OpenReview.net, 2020. URL https:///pperreiie net wwt.:/wr.eeteeeeeeeeemmmm
forum? BJdddddn
Xiaotian Zhang, Chunyang Li, Yi Zong, Zhengyu Ying. Liang He, and Xipeng Qiu. Evaluating the performance of large language models on GAOKAO benchmark. CoRR, abs/2305. 12444 2023. doi: 10.48550/ARXIV220000 12474. 12474. URL https:://dott 10050///....//......oo.oo 2305.
48550/
12474
Wanjun Zhong, Ruixiang Cui, Yiduo Guo, Yaobo Liang, Shuai Lu, Yanlin Wang. Amin Saied, Weizhu Chen, and Nan Duan. Agieval: A human-centric benchmark for evaluating foun- dation models. In Kevin Duh, Helena Gómez-Adomno, and Steven Bethard (eds.). Find- ings of the Association for Computational Linguistics: NAACL 2024, Mexico Ciry, Mexico, June 16-21, 2024, pp. 2299- 2314. Association for Computational Linguisties, 2024. doi: URL https://doi. http:///oooorr////////////////// 18653/v1/ //5.
2024. f www.:ffiniingssse
17

---

Published as a conference paper at ICLR 2026
A EXPERIMENTAL SETTINGS
Model architecture The core architectures of our experimental model are detailed in Table The model is configured with 20 layers and hidden dimension size of 2048. Except for the first layer, all FFNs layers are replaced with MoE layers. We adopt the GQA attention mechanism (Ainslic et al. 2023) and integrate Rotary Position Embedding (RoPE) (Su et al. (022)), enabling the model to support sequence lengths up to 8K tokens. For parameter initialization, all learnable parameters randomly initialized using standard deviation of 0.006. Under this configuration, the model consists of a total of 16.3 billion parameters, of which approximately 1.43 billion are activated for each token during inference.
Table 6: Detailed model architectures.
Rlagers dmodel
djfn
despert
Thead
Rkkuheadd
E
Ea
E,
N
Na
2048
5120
512
16
256
16.3B
1.43B
Training hyperparameters We use the AdamW optimizer (Loshchilov & Hutter 2019 with hy- perparameters set as follows: $1 0.9, 32 = 0.95, and weight decay of 0.1. Gradient clip- ping (Zhang al., 2020) norm is set to 1.0. According to the scaling laws for MoE optimal hyper- parameters, the maximum learning rates were set to 3.74c-4. The batch size is set to 2048, and with a maximum sequence length of 8K, each training batch contains 16M tokens.
Pre-training dataset The training data is sourced from large-scale multilingual corpus, primar- ily covering English and Chinese, while also including various other languages. This corpus en- compasses web text, mathematical materials, programming scripts, published literature, and diverse textual content. validate model performance, we extracted a 10T-token subset from this corpus for training.
Evaluation setup To evaluate performance, we consider a diverse suite of downstream tasks de- signed to provide a holistic assessment of model capabilities. For base model, tasks grouped into several categories, such as: (a) General Knowledge/Reasoning (e.g.. ARC (Bhakthavatsalam et al.
2021.
AGIEval (Zhong et al.,
2024), OpenBookQA (Mihaylov al. 2018). BBH (Suzgun al.
2023
WorldSense (Hong et al. 2025), PIQA (Bisk et al. [all]]222)))) hellaswag (Zellers et al. 2019) and
TOR-Bench (Ma et al. (025) (b) Language Understanding (e.g., race (Lai al. 2017). SQUAD
2.0
(Rajpurkar et al.] 2018). TriviaQA (Joshi et al. 2017). NQ (Kwiatkowski et al. 2019) and
winogrande (Sakaguchi et al. (02T)) (c) Professional Knowledge (e.g., MMLU (Hendrycks
al.
2021a), CMMLU (Li et al. 2024a). C-Eval (Huang et al. (203)). MMLU-Pro (Wang et al. 2024b).
GPQA (Rein al. 2023) and SuperGPQA (Team al. (025 (d) Math (c.g.. GSM8K (Cobbe
et al. 2021). MATH (Hendrycks et al. 2021b), gaokao (Zhang et al. 2023) GSM-Plus (Li et al. 2024b), mgsm-zh zhh (Shi et 2003), CMATH (Wei et al.] MathBench (Liu et al. 2024b), minerva_math (Hendrycks et al. 2222 college_mas m_tth (Tang et al. 2024) and cn en_middle..m smiddle_school 24) (e) Code (e.g.. HumanEval (Chen et al. 202T). LiveCodeBench (Jain et 2025), MBPP (Tao et al., 2024), HumanEval_pian (Liu et 2023), MBPP_pluu (Liu et al. 2023). HumanEval_com (Peng et al.
2024 y HumanEval.fim (Bavarian et al. 2022) and CruxEval (Gu et al. 2024)). For SET model, the
categories and tasks are shown in Table and 10
B DETAILS OF THEORETICAL CONNECTION BETWEEN DECAY AND MERGING
The core idea of checkpoint merging in this work is to take an ordered list of checkpoints, [0,, 0at1...., .n0x)) and apply merge function to generate single model 0 Here, 0, Rd represents the model's parameter vector at the i-th training iteration.
The most general form is weighted average of the checkpoints:
On4k
Σ cgsnns
(6)
1=
18

---

Published as a conference paper at ICLR 2026
where {e,} are non-negative weights that sum to one, i.e. Σド=0 Cj =
While intuitive, this formulation obscures a deeper connection to the training dynamics. We can reveal this link by expressing each checkpoint terms of an initial checkpoint 0n and the subsequent gradient updates.
We assume that the updates between checkpoints at different time steps are independent. Let gi be the gradient update vector (including the learning rate) at step i, such that the model is updated via 8(41 = 02 g.. Any checkpoint bn+y can therefore be written as the sum of an initial state 6, and the sequence of negative gradient updates that followed:
Ontj On
Σ 9n+t-1
Substituting Eq. into the general merging formula (Eq. 6 yields:
0nth
Cj
On
In+I-1
(8)
j=0
lm1
Cj
On
Jutt-1
1=0
j=0
h
On
9n+1-1
(10)
i=1
The double summation in Eq. 10 can be rearranged by changing the order of summation. A gradient update Inti-1 is included in the sum for all checkpoints 6ntj where
k
Σ
Intl-1
Cj
Inti-1
(11)
This shows that weight coefficient of Σ Cj. is applied to the gradient update In+i-1 from step +i-1 - 1. By defining new set of weights for the gradient updates, ws = Ej= Cj, we arrive at the final equivalent expression for checkpoint merging:
0 = On Σ
W. Iati-1
(12)
1=1
This equation demonstrates that merging checkpoints with weights } is equivalent to applying synthetic decay schedule defined by weights {w;} the gradients accumulated after the base checkpoint o,, and there exists mapping between the merging weights ej and the effective learning rates Wi
B. PROOF OF THEOREM 3.1
We seek to find the unique checkpoint weights } corresponding to a given desired sequence of gradient decay coefficients {wi]e-. The relationship derived in the previous section is the starting point:
w
(13)
We can solve for the checkpoint weights {e,} by starting from the last element and working back- wards.
For i = k. the sum in Eq. 13 has only one term:
Wkk
Σ
Cj
Ck
(14)
j=k
19

---

Published as a conference paper at ICLR 2026
This gives us the value of directly.
For any E [1, h - 1]. we can write out the expressions for wi and w441:
Wi Cj + G++1 + G+42 .... Ck
W1+1 C++1 + G42 +.....
Subtracting the second equation from the first yields the expression for c;:
Wi - - =
(15)
Finally, for co, use the constraint that the checkpoint weights must sum one: Σ Cj 1.
Co
j=1
(01 + C2 + Ck-1 ck) (w w2) + (w22 W3) + ... (wk-1 wx) - wk)
(16)
The sum in the parentheses is a telescoping series which simplifies to w1-
cp WF
(17)
This completes the derivation of the unique formulas for {c} as stated in the theorem.
For the checkpoint weights {ey} } to be valid, they must be non-negative. This imposes conditions on the sequence {w.}.
From Eq.
14
Ck implies wy 2 0.
From Eq.
15 Cj 0 for E [1, A 1] implies wy #j+1 2 0, which means wy 2 Wz+t- This
shows that sequence (w;) must be monotonically non-increasing.
From Eq.
17 co 2 implies w1 2 0, which means W1 1.
Combining these conditions, arrive the requirement that the gradient decay sequence must be bounded and monotonically non-increasing: 2 W1 2 W2 2 ... 2 wk 2 0. Thii ensures that valid (non-negative) set of checkpoint weights {cy} can be derived.
C DISCUSSIONS
Choice between Offline and Online Merging We prioritize an offline, checkpoint-based ap- proach over standard online running averages (such as maintaining a recursive average every step) for two key reasons: (1) Flexibility: WSM allows us explore dozens of distinct virtual decay schedules from a single training run. By preserving checkpoints, practitioners can systematically evaluate how different factors influence the final model, including choice of merging algorithm (simulating various decay curves like linear or 1-sqrt) and the annealing duration (by varying the merge window size n).. contrast, an online average commits to single anncaling path. that trajectory proves suboptimal, the only recourse is expensive retraining. (2) Performance: Standard online running averages typically behave as an Exponential Moving Average (EMA). As shown in comparisons, EMA is suboptimal merging strategy compared to the concave schemes (e.g., 1-sqrt) enabled by WSM. A recursive online average is fundamentally incapable of simulating these more effective non-uniform schedules. However, once this offline exploration identifies superior, fixed strategy (e.g., "an average of the last 4 checkpoints"). it can be operationalized as a simple and efficient online process using sliding window.
Practicality and Complexity Compared to discovering optimal decay schedule through mul- tiple, insource-intensiie training runs, the merging operations in the WSM strategy save substantial computational resources. The main complexity introduced is storage overhead. For the offline merg- ing approach, which is ideal for exploration, one must store a history of checkpoints, with the total number being the total training tokens divided by the checkpointing interval. However, this is
20

---

Published as a conference paper at ICLR 2026
manageable trade-off, as this storage footprint represents minor fraction of a typical pre-training budget. Furthermore, once a superior strategy is identified, or in scenarios with extreme storage constraints, online merging approach by maintaining only a small, fixed-size window of n checkpoints. Our experiments confirm that a relatively small window (e.g., 12) is sufficient to achieve strong results.
be used. This method minimizes the storage footprint
Comparison with Concurrent Work. We acknowledge concurrent work that has explored model merging during pre-training, most notably from the EMA-based approach in ERNIE 4.5 (ERNIE: Team 2025) and PMA (Li et al. 2025). We highlight the key differences between our work and theirs as follows: (1) Motivation: We frame checkpoint merging as a novel LR scheduling mecha- nism, with the primary goal of discovering schedule that outperforms strong baselines like WSD. In contrast, their work treats model merging as a standalone pre-training technique, evaluating its value by comparing the performance of the merged model against its before-merge checkpoints. (2) Methodology: We propose a general theoretical framework that generalizes beyond simple EMA. While their discussion centers almost exclusively on EMA, our framework formalizes mapping from any LR decay schedule (e.g.. linear, inverse-sqrt, cosine) to specific, principled set of averag- ing weights. (3) Key Findings: Our empirical results demonstrate that merge weights corresponding to linear and 1-sqrt decay schedules yield significantly better performance, while the EMA baseline sub-optimal. Additionally, we identify the merge duration as the most critical hyperparameter in- fluencing the final model. Ultimately, our method surpasses the performance of the WSD baseline. In contrast, their work concludes that checkpoint merging can effectively match the performance of an -nnnellle model. We view the findings from our work and these concurrent studies as complementary, and together they offer valuable practical insights for LLM pre-training.
Distinction from Prior Theoretical Works While previous studies (Bergsma et al. 2025 Wang
et al. 2024a) have noted theoretical connections between learning rate schedules and weight averag- ing, our work differs fundamentally in both objective and scope. (1) Paradigm Shift vs. Analysis: Prior works primarily utilize this connection as an analytical tool to explain the convergence prop- erties of standard decay schedules. Their goal is analyze the existing paradigm. In contrast, our objective is to replace it. We propose WSM as a standalone training strategy that eliminates the need for an explicit decay phase. We are, to our knowledge, the first to systematically validate of- fline averaging as direct, efficient substitute for LR decay in large-scale pre-training. (2) General Framework vs. EMA: Previous theoretical discussions have focused almost exclusively on the re- lationship between SGD and Exponential Moving Average (EMA). Our framework generalizes this significantly by providing a constructive proof that maps any decay schedule (e.g., linear, I-sqrt) to specific set of principled averaging weights This distinction is practically vital: our experi- ments demonstrate that EMA is often sub-optimal strategy. The superior performance of WSM relies directly the precise, non-EMA weights derived from our general framework, allowing us to simulate concave schedules that standard EMA cannot approximate.
D ADDITIONAL EXPERIMENTS ON MUON OPTIMIZER
To empirically validate the applicability of WSM across different optimization algorithms, we con- ducted an additional experiment using the Muon optimizer (Jordan 2024). distinet representative optimizer compared to the Adam optimizer used in our main experiments. We performed a 100B token enhancement training session initialized from checkpoints trained on 2T tokens. The evalu- ation metrics and benchmarks remain consistent with those described in the main paper. The com- parative results between WSD and WSM are presented in Table We observe that WSM achieves significant gains over WSD across the majority of domains. These findings demonstrate that the effectiveness WSM is not limited to specific optimizers and maintains its superiority in diverse optimization landscapes.
21

---

Published as a conference paper at ICLR 2026
Table Performance comparison with Muon optimizer on 100B-token enhancement training (ini- tialized from 2T checkpoints).
General Knowledge/ Reasoning
Language Understand- ing
Professional Knowledge
Overall Average
Math
Code
WSD WSM
63.37 63.95
75.50 76.95
55.01 56.85
47.10 47.30
45.09 45.04
56.10 56.74
E
ALGORITHM
Algorithm The WSM LRS Pre-training Pipeline
Input: Initial model weights Wo. pre-training data D, high-quality annealing data Danneul (optional), peak LR Ir peakk warmup steps Tw warrmup. switch step Tswitch. checkpointing interval Tept, merging window size
Initialize:
& The merged checkpoint.
W Weerged null
& Storage for the base (un-merged) checkpoints. Cstorage +
D Phase 1: Warmup on Pre-training Data
for to T Twarmup do
(r(1) <- peak (l/Twermup)
W. < Update(W (W-1, (()r())
end for
& Phase Stable Training and Merging
for = Twarmup to ... do
Ir(t) + Ireeak
Select dataset for the current step if > Tawitch and Danneat is available then Dourrent Danneut
else
Deurrent + D
end if
W. t Deurentn, Ir(t))
ift (mod Teps) == then
b Save the checkpoint from the main training process.
W.)
Asynchronously Update the Merged checkpoint
Chatest f- CeeLLasttCheckpoinnssCccoo
n.
if (en((uut) = n then
& Update the merged checkpoint for evaluation. W Wmerged <- Merge(Ciated)
end if
end if
end for
return The stored base checkpoints in Catorage and the merged checkpoint W Weerged.
F ADDITIONAL EXPERIMENTS ON WSD
Our preliminary experiments into Warmup-Stable-Decat (WSD) learning rate schedules revealed clear performance hierarchy among different decay curves, with the 1-sqrt strategy emerging as superior. Specifically, we conducted controlled experiment initialized from base checkpoint that was pre-trained 7T tokens. We then annealed the model for additional 100B tokens using a consistent, high- dataset, where we varied only the annealing decay function. The final learning rate for all experimental runs was set zero. The results, depicted in Figure confirm that
22

---

Published as a conference paper at ICLR 2026
the 1-sqrt decay outperforms other methods in benchmarks, Based this evidence, we establish the WSD schedule with 1 -sqrt decay as a strong baseline for all subsequent experiments. As this was a preliminary study, the starting checkpoint used here differs from that used in main experiments in Section
Overall Average
General Knowledge Reasoning
Understanding
24.5
62.5
cecorintial
92.0
23.5
Bneer
Accurncy (50
$5.0
62.0
73.0
02.5
60.5
Jsert
72.5
$2.0
Logt
css
77.0
cocht
rupenertia
ixporwettia
59.5
71.3
81.5
Irmes
20
30
52
60
90
100
20
30
47
50
60
70
80
1.00
20
33
40
50
60
70
60
50
100
Takun(B)
Tokeni(8)
TobentB)
Professional Knowledge
Math
Code
Esgr
15grt modite
35art
rodre
#cconentiol Ineer
aspereitiio
58.0
57.3
56.5
56.0
52
00
Figure 6: Comprehensive performance comppori between different decay strategy of WSM sched-
ule.
G How CHECKPOINT MERGING INFLUENCES NETWORK BEHAVIOR: STABILITY & GENERALIZABILITY
The architecture composed of sequentially connected transformer blocks has been widely adopted in mainstream LLMs. such chain-like structures, variations shallow-layer outputs typically propagate layer-wise. More stable and generalizable input-output patterns exhibit greater potential for enhancing downstream task performance. To investigate how checkpoint merging influences network behavior, conduct the following analyses:
Stability In numerical analysis, the condition number quantifies function's 's sensitivity to input perturbations and the resultant output errors. Remarkably, our analysis in Figure 7 reveals that the decay process induces sharp deterioration of condition numbers, whereas checkpoint merging demonstrates superior stability. This indicates that the checkpoint merging strategy not only im- proves performance but also preserves parametric stability.
Generalizability When maintaining competitive performance, higher SVD entropy (singular val- ues entropy) correlates with greater matrix effective rank, indicating greater information capacity matrix operations. For continue pre-training and fine-tuning scenarios, higher SVD Entropy of- ten means higher potential for model plasticity (Alter et al. 2000, Roy & Vetterli 2007 Liu et al.) 2025). Figure g shows the trend of SVD Entropy during training. We observe that decay is a rapid entropy reduction process, which continuously damages the potential for future continue training of the model. In contrast, the merged models still maintain a higher generalizability, manifested as higher SVD Entropy.
23

---

Published as a conference paper at ICLR 2026
proj
72.5
80
Cenditin number
67.5
65.0
70-
60.0
55-
52.5
4209
6000
60.00
13030
4000
6000
3000
10000
4020
6000
5000
10269
Training Tokana (B)
Thatsing Tokena (8)
Training Tokene (B)
constant
Merge
Dccay
Docay
Docay
Decay-137
Figure 7: Condition number of weight matrices across different training tokens.
down pro)
proj
0.878
0.815
0.28
0.876
0.610
0.874
Entrepy
0.000
SWD
0.600
0.868
0.755
0.10
0.790
0.808
0.854
0.785
4000
6000
5000
10300
4020
5000
9000
10000
4000
6000
3000
10000
Training Tokans (8)
Training Teleens (0)
Training
(B)
Deay--1
Decay
Decay-107
Figure 8: The entropy of singular values for each weight matrix in the fced-forward network (FFN)
layers.
H PARAMETER TRAJECTORIES AND MODEL PERFORMANCE: CONSTANT LR WITH MERGE VS. LR DECAY
To visually analyze the correlation between model parameters and performance during training, we employ t-SNE (van der Maaten Hinton| 2008) dimensionality reduction project the weight ma- trix of specific layer into the 2-dimensional embedding space. This projection is combined with performance evaluations to generate the composite contour map illustrated Figure 9 The direc- tional arrows in the visualization explicitly illustrate the parameter trajectories across both constant and decay phases. Our experimental analysis revealed two principal findings: 1) During the de- cay phase, model parameters gradually converge toward the merged model solution space, which achieves superior performance compared to checkpoints with a constant LR. 2) The LR reduction in the decay phase enables more precise parameter refinement than the expansive exploration observed under constant LR. This controlled convergence facilitates localization of nearby optimal solutions. We confirm the speculation of previous work Wen et al. (2024), and formalize this dynamic with an analogy: constant LR training resembles traversing a "canyon" with oscillating steps, while merging resembles finding a "river" at the canyon floor that guides efficient convergence.
I COMPARISON WITH A LONG-COOLDOWN WSD BASELINE
To ensure a rigorous comparison and align with best practices for WSD schedules, which often recommend cooldown durations of at least 10% of total training tokens, we conducted a controlled experiment to evaluate WSM against a well-tuned WSD baseline with a significantly longer decay period. For this experiment, we started from 500B-token pre-trained checkpoint and continued training for an additional 200B tokens. For the WSD baseline, the learning rate was exponentially decayed to 10% over this entire 200B-token phase. resulting in cooldown duration of approxi- mately 28.6% (200B/ 700B total tokens). This comfortably satisfies the conventional ": 10% steps" recommendation. For the WSM baseline, applied our proposed merging strategy to checkpoints
24

---

Published as a conference paper at ICLR 2026
60.370
-5.5
Merged Model Decay-8i -T
constant
59.496
-5.0
-6.5
Weight-2
$7.748
Accuracy (%)
-7.0
7500
56.874
-7.5
56.000
55.126
-8.5
54.252
-8.5
-8.0
-7.5
-7.0
-6.5
-5.5
Weight-1
Figure 9: Visualization of Model Parameters and Performance Contour Lines. Points connected by directional arrows represent parameter trajectories during the constant and decay phases, respec- tively. The red star indicates the solution space of the merged model.
Overall Average
General Knowledge/ Reasoning
Language Understanding
Accuracy (%)
Acoracy
Accuracy
Professional Knowledge
Narth
Code
Acturacy racy (50)
Accuraay (16)
230
425
171
420
WSA (Bwlore Merge)
1550 (Merrge
Figure 10: Comprehensive performance comparison between our WSM schedule and standard WSD scheduling on larger decay duration. Starting from a 500B-token checkpoint, we continued training for additional 200B tokens, resulting cooldown ratio of 28.6% (200B 700B total),
saved during the same 200B-token training phase. The performance of both methods was evaluated on our suite of downstream benchmarks, with the results presented in Table 4. As shown, WSM achieves a higher Overall Average score (51.43 vs. 50.57) and outperforms the WSD baseline in four out of the five sub-categories: Language Understanding, Professional Knowledge, Math, and Code. While WSD shows a marginal advantage General Knowledge/R<assonnng the overall fectiveness of WSM remains superior, Figure 10 illustrates the training dynamics.
25

---

Published as a conference paper at ICLR 2026
J DETAILED EVALUATION RESULTS
We provide detailed evaluation results to compare our methods with WSD. Figure 11 shows de- tailed comparison each dataset across the various categories from the main experiments in Section select the checkpoint with
I where WSM achieves an advantage on the vast majority of datasets. the highest average score for each method, including WSS and the three WSM merging algorithms, and list them Table VablegJand(() shows the performance comparison of checkpoints generated by the WSM and WSD schedules after supervised fine-tuning (SFT).
Figure II: Detailed performance comparison between the standard WSD schedule (via learning rate decay) and WSM schedule (via checkpoint merging) under different merging durations.
26

---

Published as a conference paper at ICLR 2026
Table 8: Detailed performance comparison base models trained using WSM (with three distinct merging algorithms) versus WSD scheduling approaches.
WSM
Metric
WSD
EMA
mean
1-sqrt
ARC-c
88.8
87.8
88.1
89.8
ARC-e
AGIEval
53.2
53.8
54.4
54.2
OpenBookQA.fatt
89.2
89.2
92.2
91.8
General Knowledge
BBH
69.0
70.0
71.2
71.0
& Reasoning
BBH-zh
64.5
WorldSense
50.6
PIQA
78.6
78.2
78.5
78.5
HellaSwag
69.6
68.2
69.9
69.7
KOR-Bench
31.9
33.7
34.8
34.6
race-middle
85.7
88.7
88.8
87.9
race-high
81.9
82.2
84.2
83.9
SQuAD2.0
82.9
Language
ShviaQAASshoot
58.8
57.5
59.8
59.7
Understanding
NQ_Sshot
29.8
29.1
30.4
30.3
WinoGrande
67.6
66.7
67.3
66.6
MMLU
70.7
71.2
72.2
72.1
GPQA
24.8
30.8
33.3
30.8
Professional Knowledge
CMMLU
MMLU-Pro
43.3
43.8
45.2
44.7
SuperGPQA
23.8
24.3
25.5
25.7
C-Eval
78.5
78.3
80.9
79.9
gsm8k
83.6
82.0
82.6
82.3
MATH
48.3
50.8
50.0
49.6
gaokao
38.5
47.3
39.6
48.4
gaokao2023en
54.0
49.1
53.5
53.0
gsm_plus
71.5
69.3
71.4
70.9
Math
mgsm_ zh
72.0
73.2
73.2
72.8
CMATH
88.6
88.4
89.3
89.2
MathBench
minerva_mate math
16.9
19.9
20.6
19.1
college_math
38.9
37.8
38.3
39,0
cn-middle_choool.
62.4
67.3
68.3
70.3
HumanEval
79.9
80.5
81.7
81.1
LiveCodeBench
21.2
20.3
23.2
22.7
HumanEval_plus
75.0
HumanEval.
73.8
69.3
73.1
73.9
Code
HumanEval en
74.4
77.4
75.0
75.6
MBPP
67.2
65.0
66.8
66.8
MBPP plus
72.0
70.6
70.4
70.6
CruxEval
55.6
57.0
57.4
27

---

Published as a conference paper at ICLR 2026
Table Detailed performance comparison of checkpoints generated by the WSM and WSD sched- ule after supervised fine-tuning (SFT). Both base checkpoints are fine-tuned under identical settings for epochs. Results are reported based on the epoch with the highest average benchmark score.
Metric
WSD
WSM
ARC-c
90.85
89.49
BoolQ
84.80
85.38
Basic Knowledge
GaokaoBench
75.70
79,95
AGIEval
61.87
65.22
NQ
25.4
26.43
TriviaQA
53.93
55.52
Knowledge
Average
65.42
67.00
C-Eval
76.37
77.87
CMMLU
76.13
76.78
Professional Knowledge
MMLU
72.76
74.59
MMLU-Pro
46.09
49.67
GPQA
29.55
33.4
SuperGQPA
26.50
26.43
Average
54.57
56.46
HumanEval
85.90
86.20
HumanEval plus
81.10
80.95
HumanEval.co
78.20
76.07
Code Completion
CruxEval
59.69
58.88
Multiple
66.14
64.75
HumanEvalFix
63.01
62.50
Average
72.34
71.56
MBPP
81.03
81.97
Code
MBPP _plus
73.02
71.43
LiveCodeBench
31.31
33.31
Code Generation
BigCodeBench
33.77
33.16
CodeForces
19.60
18.07
Bird-SQL
25.95
28.10
Average
44.11
44.34
CMATH
93.08
93.99
gsm&k
87.34
87.79
Elementary Mathematics
en_middle_schee
72.28
73.27
mgsm
76.80
78.25
gsm plus
77.28
77.93
Average
81.36
82,25
MATH
74.24
77.08
MathBench
74.72
75.74
Math
college math
43.57
43.52
Intermediate Mathematics
gaokao
64.17
60.77
minerva math
55.74
57.31
gaokao2023en
63.12
65.19
MATH500
74.75
77.05
Average
64.33
65.24
OlympiadBench
41.89
41.93
Advanced Mathematics
AIME2024
11.04
11.67
AIME2025
11.46
12.71
Average
21.46
22.10
28

---

Published as a conference paper at ICLR 2026
Table 10: Detailed performance comparison of checkpoints generated by the WSM and WSD sched- after supervised fine-tuning (SFT). Both base checkpoints are fine-tuned under identical settings for 5 epochs. Results are reported based the epoch with the highest average benchmark score.
Metric
WSD
WSM
C3
84.38
87.18
WSC
69.23
77.88
Language Understanding
Language
race-high
82.93
84.65
race-middle
87.95
89.42
Average
81.12
84.78
bbh
72.87
73.97
drop
76.41
78.66
hellaswag
68.11
70.93
ocnli
$1.32
$2.75
piqa
81.50
82.92
Complex Reasoning
Reasoning
ProntoQA
37.00
38.00
Molii-LogiiEall
0.27
56.68
MuSR
48.53
$1.33
korbench
37.52
37.52
bbh-zh
69.38
71.47
Average
63.21
64.94
teval. v2 en
84.4
86.16
teval.v2 .zh
83.49
84.72
Tool-use
BFCL.AST
79.35
78.30
BFCL-Live NEXUS
68.33
71.39
Agent
30.51
29.21
Average
69.22
69.96
IFEval alignbench
71.9
75.71
Instruction Following
59.10
59.80
Average
65.5
67.75
29


Reviews:
[{"review_id": "Siap8ELz0X", "reviewer": "Reviewer_hHGj", "summary": "This paper proposes an alternative approach to WSD (Warmup-Stable-Decay) learning rate schedule. It proposes to use a constant learning rate and then apply weight merging on the saved checkpoints. The authors establish an informal correspondence between the two methods and show empirically some of the weight merging schedules can result in higher downstream performance than the WSD method.", "questions": "1. The authors mention sensitivity regarding the weight merging schedules. Is there similar sensitivity in the decay schedule in WSD? Also, is the ranking between the weight merging schedules consistent with the corresponding learning rate schedule?\n\n2. Do the authors observe differences in the ranking of weight merging schedules when choosing different downstream evaluation tasks?", "limitations": "", "rating": 10, "confidence": 4, "soundness": 3, "presentation": 4, "contribution": 3, "strengths": "1. The experiment in this paper is large-scale and detailed. The models are evaluated over a range of hard benchmarks to show the improvement of the methods.\n\n2. The proposed method is conceptually clean and can be useful in many settings.\n\n3. The theoretical justification of the method, while not rigorous, provides intuitions in guiding the design of the algorithm.", "weaknesses": "1. In Table 5, the authors show that WSM has a higher ending loss compared to WSD. This seems to be counterintuitive, especially given the higher downstream accuracy. This also brings questions regarding the validity of the correspondence discussed in Section 3.1.\n\n2. Regarding continual pretraining, the authors propose to continual pretrain from the constant learning rate checkpoints before weight merging. It is unclear how this will compare with re-warming up the decay checkpoint in a standard WSD setting.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "eb7bAoa6UC", "reviewer": "Reviewer_kALL", "summary": "The authors present a theoretically grounded approach to model merging, which they suggest could fully replace LR decay schedules for WSD. This would be great! The basic idea is that, since every theta_k is a weighted combination of the gradients so far, you can choose a particular weighted average of the theta_k to emulate any decay schedule you please. They show the algebra to justify this and then a bunch of experiments at a more-than-reasonable scale that very convincingly show that you can, on tasks, do every bit as well as a proper decay (if not better).", "questions": "As noted above, it would be nice if you were able to conduct some preliminary experiments comparing other decay schedules at a modest scale 150M/3B or something. (A different scale is fine too.)\n\nAgain, if you fully believe the theory, you should be able to approximately recover a 100% cosine schedule by starting your averaging/merging from the end of warmup. I think?\n\nI recognize that your pipeline isn't set up for this (and I think it was a good choice for your experiments), and I don't expect you to run a large experiment here, but in my experience with WSD with EMA merging, we just update the merged theta at every step. This is obviously a lot more merges, but it's cheap and online. One could also do this with mean averaging. And Figure 4 indicates that it's definitely better to merge more often! how does maintaining a running average every step compare?", "limitations": "", "rating": 10, "confidence": 4, "soundness": 3, "presentation": 4, "contribution": 4, "strengths": "Overall I really like this paper! It's well written and argued, and the empirical results are quite strong. In terms of significance, this has a non-trivial chance of becoming the standard way people produce final model checkpoints.\n\nThe theoretical derivation of the equivalence is simple in a wonderful way (though I worry a bit specious). It's of course just algebra, but I have no problem with that. The intuition is the same one that motivates EMA and SWA etc.\n\n The empirical results are very strong, with across the board accuracy gains relative to EMA (another decay-free strategy) and even over WSD itself. \n\nThey also do a fairly good job defending the proposed equivalence: averaging schemes derived from stronger decay schedules (1-sqrt, linear) outperform those derived from weaker schedules (exponential decay--> EMA); and they show little to no benefit of combining a real decay with averaging, strengthening the claim that they are in some sense equivalent.", "weaknesses": "I have two categories of concerns:\n\n### Experiments don't really demonstrate connection between theory and practice.\n\nMy biggest concern is that, algebra-aside, the theory is a bit specious. The derivation assumes that gradients would be ~the same in a decay versus non-decay setup, and that can't be true. So, really I think this is \"decay-inspired averaging\" or \"empirical substitute\" more than an equivalent/replacement.\n\nTo that end, there aren't any experiments directly comparing the supposedly equivalent decay schedules with merges, except for the linear decay schedule, and even that experiment isn't fully convincing. I get that you can't afford to do this for a model trained to 10T tokens, but I would expect it wouldn't be too hard to do this for a smallish model?\n\nFor instance, if you were to really believe the theory strongly enough, you could, after warmup, keep the LR hot the entire time and recover a cosine decay via appropriate weighting. Right? Do experiments support this? If not, where does the theory/reality connection break down?\n\nI would expect that the losses wouldn't be that similar, because this approach assumes the gradients you get are similar enough in a decay vs non-decay setup that the averaging is ~equivalent to decay.\n\nI'd expect these to show up in loss plots rather than accuracy plots. I note that loss plots are conspicuously absent from the paper, which to me suggests that WSM isn't so much \"equivalent to WSD\" but more \"just as good as\" (which is more than good enough, given the advantages!). I just think we should see these results and adjust the positioning of the paper as appropriate.\n\n### Can't avoid data scheduling\n\nAnother, smaller weakness is that, in the age of \"mid-training\", you still have to decide when to start including higher quality data. In typical WSD setups, that is also when you start the decay. So you don't really have the nice any-time property that you might want from a truly decay-free world. I'm not sure how to work around this. It just (slightly) diminishes the promise of this approach.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "Uix1J3qDod", "reviewer": "Reviewer_fGtp", "summary": "This paper proposes merging checkpoints as an alternative to learning rate annealing.\nIn particular, warmup-stable-merge (WSM) is proposed as a replacement for the warmup-stable-decay (WSD) learning rate schedule,\nwhich is common for LLM pre-training.\nThere is some theory to provide intuition. Results are good.\nSome analysis experiments investigate which factors are more important.", "questions": "Some concerns / questions / comments:\n\n* Section 5.2: What you are calling model merging in this paper is more commonly termed as Stochastic Weight Averaging (SWA). To my knowledge, \"model merging\" is used in the context where the models come from different training/fine-tuning runs. i.e. (1) in L460 is \"model merging\" but (2) in L463 is SWA. At the very least, SWA should be mentioned and discussed in the paper as this is the official name used by the 2018 paper (which you have cited under model merging in L458).\n\n* WSD is not exactly schedule-free as it has been found that optimal LR depends on the decay start time [1]. By extension WSM is also not schedule-free. Please fix this claim, or at least add a caveat and discussion.\n\n* L115: You still need / will benefit from LR tuning, so please remove this line.\n\n* L257: To help the readers, please clarify that this corresponds to a linear decay LR schedule.\n\n* What is the intuition for why WSM is better than WSD? Can you do some empirical analysis to support the intuition?\n\n* Why is there a systematic dip in all plots towards the end in Figure 4?\n\n* L348: \"this finding reinforces the theoretical connection\"... I believe you need Figure 5a style curves to rigorously establish this.\n\n* L429: Is \"language modeling loss\" column the train loss or test loss?\n\n* L483: The theory is not exactly optimizer-agnostic with a more careful analysis as in [2]. For example, the weight decay constant $\\lambda$ is an important part of the SWA weight terms.\n\n[1] Straight to zero: Why linearly decaying the learning rate to zero works best for LLMs. ICLR 2025.\n\n[2] How to set AdamW's weight decay as you scale model and dataset size. ICML 2025.", "limitations": "", "rating": 2, "confidence": 4, "soundness": 2, "presentation": 3, "contribution": 2, "strengths": "* WSM is a simple method which provides sufficient improvement over WSD and is convenient enough to implement and use that I support the paper's proposal to adopt it instead of WSD.\n\n* The experiments are comprehensive.", "weaknesses": "* This work has limited novelty.\nThe connection between LR schedules and weight averaging is well-known and has more complete treatment elsewhere [1,2].\nThis includes almost the entirety of Section 3 and forms the main contribution of the paper.\nFurther, the theoretical treatment in this paper is rather hand-wavy.\n\n* The proposed methodology is also not surprising, although maybe it hasn't appeared in this packaging before.\nWeight averaging is a common practice, often applied with annealing [3].\n\n* The authors identify the \"merge duration\" as the most significant factor and suggest an offline computation to compute the merged model. However, from the experiments it is not clear why or if merging along the entire training trajectory would hurt. That would be much easier to implement online without memory overhead by maintaining a running average. For example, this is the recommendation in [4].\n\n[1] Straight to zero: Why linearly decaying the learning rate to zero works best for LLMs. ICLR 2025.\n\n[2] How to set AdamW's weight decay as you scale model and dataset size. ICML 2025.\n\n[3] The Llama 3 Herd of Models. See Section 3.4.3.\n\n[4] Post-Hoc Reversal: Are We Selecting Models Prematurely? NeurIPS 2024. Search for SWA.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "y4woGg1sEc", "reviewer": "Reviewer_852L", "summary": "The paper introduces Warmup-Stable and Merge (WSM), a novel and compelling learning rate (LR) scheduling framework for pre-training large language models. The core idea is to replace the conventional, inflexible LR decay phase with a \"decay-free\" approach. In WSM, after a standard warmup, the model is trained with a constant peak LR (the \"stable\" phase), during which checkpoints are periodically saved. The final performant model is then produced by applying a weighted average (merging) to a series of these saved checkpoints.\nThe authors make a significant theoretical contribution by establishing a formal connection between checkpoint merging and LR decay. They demonstrate that a weighted average of checkpoints is mathematically equivalent to applying a synthetic, effective decay schedule to the gradient updates accumulated since the first checkpoint in the merge window. This framework allows for the emulation of various decay curves (e.g., linear, cosine, 1-sqrt) by simply choosing the appropriate checkpoint weights, without ever modifying the live learning rate during training.\nThrough extensive experiments on a 16.3B parameter MoE model, the paper shows that WSM consistently and significantly outperforms the strong Warmup-Stable-Decay (WSD) baseline on a wide range of benchmarks, including MATH, HumanEval, and MMLU-Pro. The key empirical finding is that the \"merge duration\" (i.e., the time window over which checkpoints are aggregated) is the most critical factor for performance. The benefits of WSM are shown to persist even after supervised fine-tuning.", "questions": "- see weakness", "limitations": "", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "- This paper is the formal connection it establishes between LR decay and checkpoint merging (Section 3.1), which is insightful. \n\n- The experiments are conducted at a scale that is highly relevant to the LLM community (a 16.3B parameter model trained on trillions of tokens). \n\n- The paper provides a thorough empirical analysis of the WSM schedule. The authors methodically investigate the impact of the merging algorithm (Table 3), merge duration (Figure 4), and merge granularity (Table 4).", "weaknesses": "- **Significant Computational and Storage Overhead:**\n  - While the paper suggests transitioning to an online, sliding-window approach after identifying an optimal strategy, the initial exploration phase itself represents a significant, non-trivial cost. Can auther give a quantified results？\n  - The merging process also incurs computational overhead, specifically in terms of I/O for loading multiple large checkpoints and memory consumption, which is not fully quantified in the paper.\n\n- **Insufficient Baselines for Model Averaging:**\n  - While the paper compares its theoretically-derived merging weights against mean averaging and EMA, it overlooks other well-established trajectory averaging methods. A direct comparison with classic techniques like Stochastic Weight Averaging (SWA) would be highly valuable. \n\n- **Shift in Hyperparameter Complexity:**\n  - The paper claims that WSM simplifies LR scheduling by eliminating the decay phase. However, it arguably replaces one set of hyperparameters (decay function, duration) with another, potentially more complex set: merge duration/window size, checkpoint interval, and merge algorithm.\n  - The finding that \"merge duration\" is the most critical factor highlights that WSM does not eliminate the need for extensive tuning; it merely shifts the tuning process from an online LR schedule design to an offline search for the optimal merging strategy. This trade-off should be more explicitly discussed.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (Oral)