Through the River: Understanding the Benefit of Schedule-Free Methods for Language Model Training
Authors: Minhak Song
Beomhan Baek
Kwangjun Ahn
Chulhee Yun
Venue: NeurIPS 2025 poster ()
Keywords: Language Model
Pretraining
Schedule-Free
Learning Rate Schedules
Weight Averaging
Optimization Dynamics
Loss Landscape
Primary area: deep_learning
TL;DR: We show that Schedule-Free methods effectively navigate the river structure of the loss landscape, enabling scalable language model training without decay schedules or extra memory.

Abstract: As both model and dataset sizes continue to scale rapidly, conventional pretraining strategies with fixed compute budgets—such as cosine learning rate schedules—are increasingly inadequate for large-scale training. Recent alternatives, including warmup-stable-decay (WSD) schedules and weight averaging, offer greater flexibility. However, WSD relies on explicit decay phases to track progress, while weight averaging addresses this limitation at the cost of additional memory. In search of a more principled and scalable alternative, we revisit the Schedule-Free (SF) method [Defazio et al., 2024], which has shown strong empirical performance across diverse settings. We show that SF-AdamW effectively navigates the "river" structure of the loss landscape without decay phases or auxiliary averaging, making it particularly suitable for continuously scaling training workloads. To understand this behavior, we conduct a theoretical and empirical analysis of SF dynamics, revealing that it implicitly performs weight averaging without memory overhead. Guided by this analysis, we propose a refined variant of SF that improves robustness to momentum and performs better under large batch sizes, addressing key limitations of the original method. Together, these results establish SF as a practical, scalable, and theoretically grounded approach for language model training.

Through the River: Understanding the Benefit of Schedule-Free Methods for Language Model Training
Minhak Song"
Beomhan Baek SNU & KAIST InnoCORE) LLM 0hbaek2001@smn ac
KAIST minhaksongkkaist.co ac.
Kwangjun Ahn Microsoft Research kwangjunahn Smi cwwsoft. com
Chulhee Yun KAIST
chulhee. ac kr
Abstract
As both model and dataset sizes continue to scale rapidly, conventional pretraining strategies with fixed compute budgets-such cosine learning rate schedules are increasingly inadequate for large-scale training. Recent alternatives, including warmup-stable-decay (WSD) schedules and weight averaging, offer greater flex- ibility. However, WSD relies on explicit decay phases to track progress, while weight averaging addresses this limitation at the cost of additional memory. In search more principled and scalable alternative, we revisit the Schedule-Free (SF) method [Defazio al., 2024], which has shown strong empirical performance across diverse settings. We show that SF-AdamW effectively navigates the "river" structure of the loss landscape without decay phases auxiliary averaging, making it particularly suitable for continuously scaling training workloads. To understand this behavior, we conduct a theoretical and empirical analysis of SF dynamics, revealing that it implicitly performs weight averaging without memory overhead. Guided by this analysis, we propose refined variant of SF that improves robust- ness to momentum and performs better under large batch sizes, addressing key limitations the original method. Together, these results establish SF a practical, scalable, and theoretically grounded approach for language model training.
Introduction
As both model and dataset sizes continue scale rapidly, conventional pretraining strategies with fixed training budgets- as cosine learning rate schedules [Loshchilov and Hutter, 2017]-are becoming increasingly inadequate. These static approaches scale, evolving datasets and open-ended training regimes. For example, DeepSeek-V V3 [Liu et al., 2024. $4.2] employs a sophisticated multi-phase training procedure that falls outside the scope of traditional cosine scheduling.
ill-suited to the demands of large-
To support prolonged and flexible training, practitioners have adopted more adaptive scheduling strategies. One widely used approach is the warmup-saable-eccy (WSD) schedule [Hu et al., 2024], which avoids committing learning rate (LR) and periodically branching into decaying LR trajectories to produce intermediate terspoints-enablingg flexible and continued training. Despite its advantages, WSD has notable limitations. A key challenge lies in evaluating the quality of the current model without explicitly
fixed compute budget by maintaining main "branch" with constant
"Equal contribution. Work done undergraduate intern KAIST.
39th Conference on Neural Information Processing Systems (NcurIPS 2025).

---

entering the decay phase. This lack of visibility complicates decisions around checkpointing and training continuation, leading uncertainty in training management (see Section 2.2).
One common workaround is to maintain a weight averaging, which improves generalization and provides more stable performance estimation. However, this comes at the cost of additional memory overhead and implementation complexity, especially in distributed training setups (see Section 2.3).
These challenges motivate key question:
Is there an alternative that offers better flexibility, training visibility, and minimal resource overhead?
Our main contributions. method [Defazio et al., 2024] as principled and scalable approach for language model pretraining. Our contributions are summarized as follows:
this work, we explore this question and identify the Schedule- (SF)
We revisit the river-valley loss landscape and analyze two widely used strategies- WSD and weight averaging-througgh this lens, highlighting their respective strengths and limitations for scalable pretraining (Section2) 2).
We then focus on the SF method and show that it effectively follows the "river" structure of the loss landscape without requiring a decay phase or auxiliary averaging. This makes it particularly well-suited for continuously scaling training workloads (Section 3).
We analyze its training dynamics both theoretically and empirically. Our findings reveal that SF implicitly performs a form of weight averaging, without requiring additional memory. We also show that it operates at the Edge of Stability [Cohen al., 2021] and derive its associated central flow [Cohen et al., 2025], providing a deeper understanding of its behavior (Section 4).
Based on these insights, we propose refined version of the SF method that improves robustness to momentum parameters and scale better with large batch addes-ndr key limitations of the original method (Section 5).
Candidate Strategies for Scalable Pretraining
2.1 Backgrounds: Loss Landscape of Neural Networks
Despite the remarkable success of deep learning across nu- merous domains, classical optimization theory falls short of explaining neural network training dynamics, motivat- ing growing interest in loss landscape geometry as a key factor in understanding why deep learning works and how to design more effective optimization algorithms.
River
hill direction
River- Valley Landscape. Xing et al. [2018] observe that neural network loss landscapes exhibit a convex valley structure, with SGD trajectories oscillating near the valley floor. More recently, a growing body of work-motivated by diverse goals- has converged on a common hypothesis regarding the geometry of neural network loss landscape. This hypothesis, which we refer to as the river-valley loss landscape [Wen et al., 2025], is closely related to concepts such as the ill-conditioned valley [Song et al., 2025), basin in the loss landscape [Hägele et al., 2024], and ravine in the loss landscape [Davis et al., 2024].
river direction
1e-08 Figure I: River-valley structure in toy loss landscape. Contour plot of the objective defined in Section 4.1, illus- trating the flat river direction and steep hill direction characteristic of the river- valley geometry.
As its name suggests, a river-valley loss landscape resembles a winding ravine: steep "hill" walls flank relatively flat "river" floor that snakes through parameter space. Wen et al. [2025] formalize this picture to interpret the warmup-stad cannu--stableddeayy (WSD) LR schedule, arguing that large, noisy updates during the stable phase let SGD travel quickly downstream along the river, while the fast-decay phase pulls the iterate back to the bottom of valley. A concrete illustration is provided in Figure (see also Section 4. 1).. Complementary evidence from Song et al. [2025] show that "effective" optimizer updates happen along the river: projecting out the high-curvature hill directions does not harm progress of learning, indicating that motion the hill directions is often dispensable once the iterate is near the river.

---

Motivated by these works, we explicitly decompose the loss into two orthogonal parts:
the river component, which measures progress along low-curvature valley floor, and the hill component, which penalizes deviations away from that floor.
In the remainder this section, we revisit existing strategies for scalable pretraining through the lens of the river-valley perspective.
2.2 Warmup-Stable-Dest Schedule
To address the limitations cosine schedules particularry their reliance on a pre-specified training budget-the the carmup-stable-dess (WSD) schedule has been proposed as more flexible alterna- tive [Zhai et al., 2022, Hu et al., 2024]. WSD divides into three phases: warmup, stable, and decay, with the LR controlled separately in each. Unlike traditional schedules, WSD avoids committing to fixed training horizon by maintaining a main branch with a constant L.R and periodically branching off with decaying LRs produce intermediate checkpoints. This structure enables flexible evaluation and checkpointing without the need to predefine the total number of training steps, making WSD widely adopted strategy for scalable pretraining.
Understanding WSD. The WSD LR schedule has been widely adopted in large language model (LLM) pretraining due studies have sought to understand the mechanisms underlying its effectiveness. Hägele et al. [2024] systematically explore the impact of WSD hyperparameters and propose optimal choices for both the decay schedule and timing. In parallel, Luo et al. [2025] introduce a multi- cuulli-power-lww framework that predicts final pretraining loss as function of the LR schedule. Interestingly, their learned optimal schedules closely resemble the WSD pattern, providing further evidence for its effectiveness.
strong empirical performance. Motivated by this success, recent
Wen et al. [2025] provide geometric interpretation of WSD through the lens of the river-valley loss landscape. Their key insights are:
1. During the stable phase, the high LR amplifies stochastic gradient noise, inducing oscillations along the high-curvature hill directions. Nevertheless, enables rapid progress along the low- curvature river direction, which drives long-term improvement.
The decay phase plays crucial role near convergence: it reduces oscillations in the hill directions and steers the iterate toward the valley floor, resulting in a sharp drop in loss that is not achievable during the stable phase alone.
Limitations. A key limitation of WSD is its reliance on manually initiating the decay phase. While the stable phase often yields relatively flat loss curve, sharp drop typically occurs only once decay begins, which makes difficult to assess model quality or forecast final performance in advance. This raises natural question: can we design optimizers that closely track optimal loss-by reaching the valley floor and following the river without relying on explicit learning rate decay?
2.3 Weight Averaging
Since its introduction in stochastic approximation [Ruppert, 1988, Polyak and Juditsky, 1992], parameter averaging has been widely explored for improving optimization stability and generalization in deep learning. By reducing gradient noise and smoothing the optimization trajectory. averaging schemes can often eliminate the need for explicit LR decay, making them an appealing candidate for scalable pretraining, Among them, two widely studied approaches stand out:
1. Stochastic Weight Averaging (SWA). SWA [lzmailov et al.. 2019] enhances generalization by periodically averaging model weights. While the original method uses a cyclic LR schedule and averages every steps, many subsequent works simplify it by setting = 1, performing standard rumning averaging. Hägele et al. [2024] further refine SWA by applying fixed-length window averaging under constant LR, and demonstrated improved performance.
2. Exponential Weight Averaging (EWA). EWA maintains exponential moving average of model weights, continuously smoothing the optimization trajectory. Recently, Zhang et [2025] show that combining EWA with constant LR match the performance of cosine schedulers and

---

WSD, particularly in large-batch settings. EWA has also been proven to be effective theoretically nonconvex optimization [Ahn and Cutkosky, 2024, Ahn et al., 2025].
Interestingly, weight averaging often regarded functionally equivalent LR decay. Sandler et al. [2023] analyze schemes like SWA and EWA, showing that their dynamics closely resemble those induced by decaying LRs. Motivated by this, several works advocate for weight averaging as a viable alternative to schedulers in scalable pretraining [e.g.. Hägele et al., 2024, $4.1J. Concurrent with our work, Li et [2025] report that training with a constant LR, when combined with weight averaging. matches the performance of models trained with decaying schedules at any point during training, without the need for LR decay. From the river-valley perspective, weight averaging serves to cancel out oscillations along hill directions, enabling the optimization trajectory to align more closely with the lieerwithout relying on explicit LR decay.
Limitations. Despite its benefits, weight averaging introduces memory overhead, as it requires maintaining an additional copy of the model parameters. This becomes a bottleneck in large-scale LLM pretraining. For instance, storing single 16-bit copy of LLaMA-88 requires over 16 GB of memory. This limits the practicality of weight averaging in cemorry-oonstainee or large-scale training environments.
2.4 Schedule-Free Methods
The Schedule-Free (SF) method [Defazio et al., 2024] provides a general framework that interpolates between two classical techniques: Polyak-Ruppert averaging, which returns the average of past iterates, and primal averaging, where gradients are evaluated at the averaged point. The abstract formulation of the SF method is given by:
(1 (x) Xt-1 + Co 2t,
Yt (1 3) Zt + Bxc.
(SF)
2r+1 2t -
where is the LR, 3 is momentum-like coefficient, = 1/t, and the initialization satisfies Z X1+ The update direction A, generic, making SF flexible framework that can be combined with any baseline optimizer. For example, in Schedule-Free SGD, A, corresponds to a stochastic gradient evaluated the y, iterate.
In this work, we focus on Schedule- Adam W (SF Add-md,,)), where A is computed using the RMSprop update along with weight decay term. The full pseudocode is provided in Algorithm 1. Here, 31 denotes the coefficient /3 in (SF), and 83 is the momentum parameter used in the second- moment of RMSprop.
SF SF-AdamW has demonstrated state-of-the-art performance across a range of deep learning tasks, including winning the Self-Tuning track in the 2024 AlgoPerf Challenge [Dahl et al., 2023, Kasimbeg et al., 2025]. Importantly, achieves this without requiring additional memory overhead compared to AdamW. However, its practical deployment reveals key limitations: sensitivity momentum hyperparameters [Hägele et al., 2024] and degraded performance under large batch sizes [Zhang et al., 2025, Morwani et al., 2025]. We revisit both limitations in later sections and propose a refined variant of SF that addresses them.
3 Schedule-Free Optimizer as a Scalable Pretraining Method
In Section 2, we discussed the limitations of WSD and weight averaging as strategies for scalable pretraining. While WSD relies on decay phase to achieve optimal performance, weight averaging avoids decay but incurs additional memory overhead. this section, we empirically investigate the SF method as an alternative. We find it to be strong candidate, as it requires neither decay nor extra memory.
Experimental Setup. We use 124M parameter LLaMA [Touvron et al., 2023a,b] style decoder- only transformer, trained with SF-Adamw using a warmup phase followed by a constant L.R The batch size is 0.5M tokens, and training is conducted on a 6B-token subset of SlimPajama [Soboleva et al.,

---

Vanishing Benefi: LR Decay
Vanishing Bonefit of EWA
Short Dessy Phase of Adam'
45
45
tou Maxi Great 3e3
Meali
Acxeli Chast En Ted
40
Perplestty
(35
sPadaam Deca LMb
Purplestyy
18
TMA ShiMaa
Perplesdiy
Adaati Dee wiy.keee
30
-6
22
1000
2000
3000
4000
5000
1000
2000
3000
4000
5000
1000
2000
3000
4000
5000
Iterations
Iterations
Iterations
While AdamW benefits from linear LR decay and EWA, SF-Adami shows no improvement from either. Right: short decay phase of AdamW (with linear LR decay from le-4 to 0) leads to a sharp loss drop for AdamN, but has no effect when applied to the SF-AdamW trajectory suggesting that SF- Adamm already tracks the river throughout training (Observation 1).
Figure 2: SF-AdamW closely follows the river, unlike AdamW. Left, Middle:
2023], with the compute budget determined by the Chinchilla scaling rule [Hoffmann et al., 2022]. We report validation loss in terms of perplexity. Additional results using 124M parameter GPT- 2 [Radford al., 2019] style decoder-only transformer trained on the OpenWebTexte dataset [Gao et al., 2020] are provided in Appendix C. Full training details are available in Appendix B.
Vanishing Benefit of Learning Rate Decay and Weight Averaging. As noted in Section 2. standard AdamW with constant LR typically yields suboptimal performance, requiring decay phase or weight averaging to reach better solutions. To test whether SF Adamm exhibits similar behavior, we first perform a grid search to identify the best hyperparameters for both AdamW and Sd-Adam under constant LR (after warmup). In our setting, the optimal hyperparameters are (31, 32) = (0.9, 0.95) with LR 1e-3 for AdamW, and ($1, 82) = (0.95, 0.99) with LR 2e-3 for SF-AdamW. Using these configurations, we train each model and periodically save checkpoints. At each checkpoint, we run LR decay phase and evaluate the resulting loss. We also track the EWA of the x, iterates throughout training. Results are shown in Figure 2.
Surprisingly, unlike Adami, neither the decay phase nor EWA provides additional benefit: SF-Adamd with a constant LR consistently reaches near-optimal solutions on its own.
Schedule-Free Optimizer Tracks the River. We next examine how closely the 3F-AdamW trajec- tory follows the river the loss landscape, building on the observation by Wen al. [2025] that AdamW with a decaying LR converges toward the river. Specifically, we run a short decay phase of AdamW (linear decay from 1e-4 0) each checkpoint from the SF-Adamw run described in the previous experiment. This decay phase-starting from small LR-is designed to make minimal progress along the river direction while substantially reducing the hill component, pulling the iterate toward the valley floor. As baseline, we apply the same procedure to AdamW. Results are shown in Figure 2.
We observe that applying a decay phase of AdamN to the SF-AdamW trajectory results in minimal additional loss reduction, in contrast the Adaml trajectory where the decay phase leads to a sharp drop in loss. This suggests that SF-Adamk already closely tracks the river throughout training, without requiring LR decay or weight averaging.
To further support this interpretation, we measure the loss along linear interpolations between the 2B and 2.5B token checkpoints under three training regimes: (1) Adam/ with constant learning rate (LR). (2) AdamN with a linear LR decay to zero, and (3) SF-Adamä with a constant LR. As shown in Figure 3, the resulting loss curves display qualitatively distinct behaviors: (1) shows a convex, valley-shaped profile, indicating oscillation across the valley; (2) shows a sharp. monotonic decline, consistent with a transition from the valley wall to the floor; and (3) shows flat, slow decline, suggesting that the trajectory is already closely aligned with the valley floor. Notably, cases (1) and (2) replicate the loss profiles reported in Figure 7 of Wen et al. [2025], while (3) offers further evidence that Adamm remains near the river throughout training.
Observation 1: SF-Adaml can follow the river without LR decay or weight averaging.
5

---

AdamW (Constant LR)
AdamW (Linear Decay LR)
SF-AdamW (Constant LR)
3.30
3.25
3.20
0.00
0.25
0.50
0.75
1.00
0.00
0.25
0.50
0.75
1.00
0.00
0.25
0.50
0.75
1.00
Figure 3: Linear interpolation between training checkpoints. We evaluate the loss along linear interpolations aw,, -(1-a)Wt, where checkpoints, respectively. We compare three training regimes: (1) Adamw with constant learning rate (LR)), (2) AdamW with linear LR decay zero, and (3) SF-Adami with constant LR. For all settings, the first 2B tokens are trained using either constant-I.R AdamW (for and 2) or constant-L.R SF-Adami/ (for 3). The resulting curves exhibit qualitatively distinct behaviors: convex (valley-shaped) for (1), sharp monotonic decay for (2), and flat, slowly declining loss for (3) (Observation 1).
[0, 1] and f1. t2 denote the 2B and 2.5B token
Sensitivity to Momentum Parameters. Despite its strong empirical performance, SF method is highly sen- sitive to the choice of momentum parameters. For ex- ample, Hägele et al. [2024] report that SF-AdamM with (2) (0.9, 0.95) performs significantly worse in their pretraining setup. even exhibiting rising loss toward the end of training, whereas (0.95, 0.99) leads strong results. This sensitivity contrasts with the thcoretical anal- ysis of Defazio et al. [2024], which shows that the SF method is worst-case optimal for any momentum setting in convex Lipschitz problems.
Short Decay Phase of AdamW, S Subopkimel P
SPAdunik R. 0.5
Mann
200
Porplectyy
AdazsW Decay
150
100
50
1000
2000
3000
4000
5000
Iterationn
Figure 4: SF-AdamW with suboptimal momentum fails to follow the river. A short decay phase of Adam' applied to SF-Adami checkpoints with 31 E {0.1, 0.5) results in sharp loss drop, unlike the case with B1 = 0.95 (Obser- vation 2).
To further investigate this gap, we repeat our experi- ment using SF-AdamW with suboptimal momentum 81 E 0..), As before, we periodically save checkpoints and apply short AdamM decay phase at each. Unlike the optimal case (3) = 0.95), we observe that the decay phase improves performance, suggesting that suboptimal momentum disrupts the optimizer's ability to follow the river. Results are shown in Figure 4.
Observation 2:
SF-AdamW is highly sensitive momentum; poor choices can prevent from reaching and following the river.
These findings lead to the following central question:
Why does well-tuned Schedule-Free method successfully follow the river, and what makes this behavior so sensitive to momentum?
We address this question in the next section by analyzing the training dynamics of the SF method.
4 Understanding the Training Dynamics of Schedule-Free Optimizer
4.1 Warmup: A Toy Model
To build intuition, we start by studying the following simple two-dimensional objective:
1 (w0 - 1)2 + log(1 (1 + exp(-w())))
f(w) =
where w= (1)), 2() E IR2. As shown in Figure 1, this objective exhibits river-valley structure: the curve ( (2) = forms river, along which the loss slowly decreases (1) increases. We use this toy model to visualize the training dynamics of SF- SF-AdamW in a controlled setting.
We run Ad-Ad with a constant LR, initialized at X; = (2, fixing 3₂ = 0.99 and varying 31 (0.1, 0.9). For each run, track the iterates (xr, Yr, Results are shown in Figure 5. We observe that the y, iterates- where the gradient valuated-oscillatt around the river across all momentum settings, with the center of oscillation remaining close to the river. In contrast, the x

---

Fallure of x: in Suboptimal B:
Oscillating Trajectories of V:
2,00
wilk 81 0.5
y) with 0.5
4.85e
4.83e-02
1.50
River
River
4.81e 02
1.25
Less tlogelloaaaa
0.332
4,756
1.00
4.77e-22
0.75
0.331
4.75 02
11-22
4.736-02
0.30
0.330
4.71e-02
0.25
3,000 3.005 3,010 3.015
3.020 3.025 3.030 3,005 3,040
Figure 5: SF-AdamW on toy model. Left:
The x, iterates fail to follow the river for B1 E {0.1, (Observation 2). Right: The y, iterates oscillate around the river but track it reliably on average, even for suboptimal values of 31 (Observation 3). As 31 increases, the oscillations shrink.
iterates fail to track the river for suboptimal values of 31 (0.1 and 0.5), whereas 31 = 0.9 results in a trajectory that remain closely aligned with the river.
This behavior aligns with Observation where we observed that SF-AdamW fails follow the river under suboptimal momentum configurations in language model training. Notably, even when 81 is suboptimal, the yt iterates continue to track the river on average, despite exhibiting oscillations (we revisit the nature of this oscillation in Section 4.3). These observations suggest that the y: sequence more robust to 31 and better aligned with the river geometry than the x iterates, making it a more reliable foundation for analyzing and guiding optimization in SF dynamics. We investigate this further in the context of language model training.
4.2 Language Model Training
We evaluate the loss at the y, iterates, as well as the EWA of yr, using the same experimental runs from Section 3 with 81 E 10.1, 0.5, 0.95). Results are shown in Figure 6. Notably, for suboptimal 31. we observe that the loss at is consistently lower than that at x,, showing that y, more faithfully follows the river geometry and remains robust to suboptimal momentum settings. This mirrors our findings in the toy model analysis (Section 4.1), where x, failed to follow the river under suboptimal momentum, while y, continued to track it. Moreover, the EWA of y, consistently achieves lower loss than the y; iterates across all momentum offers no benefit for 31 = 0.95. illustrates that the EWA of y, consistently remains closer to the river compared to the vanilla y, iterates. This observation parallels the oscillatory behavior of y, in the toy model, where the EWA would more closely align with the underlying river-like geometry.
the x; iterates, where EWA
81 0.95, Optimal
Bi 0.5, Suboptima]
8: 0.1, Suboptimal
35
Perpleait
Perpleaity
Perplesiity
30
1000
2000
3000 Iterations
4000
5000
1000
2000
3000
4000
5000
1000
2000
5000
4000
5000
[turations
Iterations
Figure 6: Performance of Xx, Ys, and the EWA of y, under varying 31. For suboptimal B1. Ye outperforms X,, and across all momentum settings, the EWA of y: achieves the lowest loss (Observation 3).
Observation 3: In SF- Adamm, the y, iterates remain well aligned with the river geometry of the loss landscape, even under suboptimal momentum settings, whereas x, may deviate.
4.3 Schedule-Free Methods at the Edge of Stability
We return to the toy model from Section 4.1 further analyze the optimization dynamics of the Ye iterates. Notably, the y, sequence exhibits period-two oscillation centered along the river
7

---

SR-AdamW Model
SF-GD on CIFAR-10
SF-AdamW CIPAR-10
Preenood Bonnee Sharggonn
Traning Linned Shu
403
300
0.6
3.01
6.4
3.03
20+06
100
0.2
202.001
084400
35:00
0300
65.00
5000
5000
100
200
300
400
500
600
250
500
750
1000
1250
Tberstoon
Figure 7: The y, iterates of Schedule-Free methods operate at the Edge of Stability. Plots of (preconditioned) sharpness during full-batch training; dashed lines indicate stability thresholds. Left: Toy model trained using SF-Adamll with ( (0.9, 0.99). Middle, Right: Fully connected network trained on a 5k subset CIFAR-10 using SF-GD with 3 = 0.9 and SF-AdamM with (81, (B2) (0.95, 0.99); (preconditioned) sharpness is evaluated the y: iterates (Observation 4).
(Figure 5). This behavior resembles the dynamics of full-batch GD the Edge of Stability (EoS), where iterates oscillate along sharp directions while remaining globally stable. The EoS phenomenon was first identified empirically by Cohen et al. [2021] and has since been studied theoretically [Arora et al., 2022, Wang et al., 2022, Ahn et al., 2023, Damian et al., 2023, Song and Yun, 2023, Zhu al., 2023]. key feature of this regime is that the sharpness-mmasureed by the largest Hessian eigenvalue stgablalue-stabilizzes near the threshold 2/7-
We now make this connection precise by analyzing SF dynamics through the lens EoS. Notation. Let ((w) be the objective and H(w) ) its Hessian at w. For brevity, write H; == H Hw,)). The largest eigenvalue X (H) called the sharpness, and for a given preconditioner X (P HH) called the preconditioned sharpness.
We first study the stability of Schedule-Free GD (SF-GD) on quadratic objectives, which serve as local Taylor approximations to neural network training losses. SF-GD is defined by (SF) with A, & fyy:)
Proposition 4.1 (Stability Threshold of SF-GD), Consider running SF-GD on a quadratic objective s(w)=
= -
Hw + 3/ww4 w If An (H) > [I-Byp then the iterates {(x,, Vt, 2x)} diverge.
Notably, this stability threshold is scaled by factor of (1 B) compared to the threshold for standard GD. This result is consistent with empirical observations from Defazio et al. [2024], which report that SF methods allow the use of larger LRs, particularly when is close to one.
Next, extend the analysis 3F-PrecondGDW, defined by (SF) with A, A P--) + Ay where P is a fixed, symmetric, and positive definite preconditioner, and À denotes the weight decay coefficient. This setting parallels the analysis of Adaptive Edge of Stability Cohen et al. [2022]. Proposition 4.2 (Stability Threshold of SF 3F-PreconddDD). Consider running SF- Pr-PrrccnnGD on f(w) = f T Hw + g w + c. If Ai (P HH) (1- >, then the iterates {(x1, Yr, 2;)3 diverge.
Proofs of Propositions 4.1 and 4.2 are deferred to Appendix D.1.
SF-AdamW can be viewed as SF-PrecondODD with with a slowly varying diagonal preconditioner Pt- Hence, its stability is governed by the preconditioned sharpness À (P, (P,1) HH). As shown in Figure 7, the preconditioned sharpness at y, iterates equilibrates near the stability threshold both the toy model and CIFAR-10 10 experiments exciiinittteeiibiiinngg typical EoS behavior.
Observation 4: full-batch settings, Schedule-Free methods operate at the Edge of Stability, with the (preconditioned) sharpness yt hovering around the stability threshold.
4.4 A Reformulation of Schedule-Free Optimizer
Motivated by Observations 3 and 4, Defazio al. [2024] and Morwani et al. [2025], define the momentum variable m, X The Schedule- update in (SF) is then equivalent
examine the dynamics of the y, iterates. Following
m, = (1 c.) m 1 + A
y:+1 = y: - 7[3041m + (1- A)]]
(SFy)
i.c. SF is momentum- based optimizer update on yc. The full derivation is provided in Appendix D.2.

---

The iterate can then be expressed as
(1 (1-ce) c+) 3)x, +aye 11-c0(11 B) + ct
X
In other words, X2 is weighted average of past yr's. Hence, arrive the following conclusion.
Observation 5: Schedule- implicitly performs weight averaging over momentum iterates without storing an extra model copy.
4.5 Central Flow Analysis
Cohen et al. [2025] observe that, at the EoS, the time-averaged optimization trajectory follows differential equation called the central flow, which characterizes the river that the dynamics trace during training. We adopt this framework to understand the magnitude of oscillation of y: iterates of SF-Adami along the river. In particular, adaptive preconditioner Vi updated as V = 32 + (1 - 32) | fyr))l?? Based on the reformulated update (SFy), and assuming that c(t) 1/t becomes negligible for sufficiently large t, the central flow equations are given by:
analyze the scalar surrogate SF-ScalarAde Admmm with an
dy 7(1 31) VD
Vf(y) + + )
du 1 - B2 vf(y)]
-o
dt
dt
B2
where S(y) = X1 [(y)) and a² the steady-state variance oscillations along the hill direction. Enforcing the stability condition S(y)/ vv = 2/[(1 - B1)1] yields
d2 (VS, vʃ) +
1-8 1
32
21
V (1-P.7P27"
1-82
S2
As B1 increases, x² decreases; thus, larger values of suppress oscillations along the hill directions, keeping the yt iterates more closely aligned with the river-consistent with the empirical observation in Figure 5. A complete derivation, including the central flow of SF-GD, is provided in Appendix D.3.
A Refined and Robust Schedule-Free Optimizer
While SF-AdamW achieves strong performance, highly sensitive momentum hyperparameters and degrades under large batch sizes [Zhang et al., 2025, Morwani et al., 2025]. Building on the insights from Section 4, we revisit these issues and propose refined variant that addresses them.
A key limitation in the vanilla (SF) setup (c: = 1/t) is that 3 simultancously controls both (i) the momentum applied y, (SFy) and (ii) the implicit averaging window that defines X,;
T
(1 [1-c.)1133) ----------1
II Y at -= a (11 B) G. s=t+1 [
XT
When 3 is small, the weights {a,} become stretched, overemphasizing early iterates and preventing X, from closely tracking the river, shown in Figure This also explains Observation 1, where we demonstrated that applying EWA to offers no benefit: since x, is already a weighted average of Yt, further averaging merely flattens the weights and weakens alignment with the river.
More fundamentally, 33 plays a dual role: it controls both the momentum update of yt and the width of the averaging window for x,. The optimal value for each may differ, and this mismatch can hinder performance. large-batch training, a narrower averaging window (i.e., larger 3 is preferred to emphasize recent iterates. However, large 3 also slows the update of yı, the (1 B) term in (SFy) becomes small, reducing the influence of recent gradients.
Averuging Weights,
Averaging Weights, $
Avaraging Weights
0.010
0.010
0.010
0.088
0.008
0.008
0.006
0.006
0.006
Figure 8: Averaging weights in SF method. Smaller values of 3 flatten the averaging weights )

---

Redined SF-AdamW, Suboptimal
Rofined SP-AdamW Optimal
Refined SF.AdamW, Large
SPAW iy
Noart Gonan
Axland 61-608.C-20
Perplexity
60
Perpleaity
Jetasd
Perplesiiy
30
40
25
20
1000
2000
3000
4000
5000
1000
2000
3000
4000
5000
1000
1,500
2,000
2500
Iterations
lterations
lterations
Performance of x and Y: iterates using the refined SF-AdamW Refined SF- Ad-md with 3n E {0.95, 0.9) and C 200 achieves improved performance over the best vanilla SF-Adamil run. Right: Under a large batch size (2M tokens), vanilla SF-AdamWi with 31 0.98 underperforms compared to AdamW with cosine schedule, while the refined -F--Adam-wwil only a sweep over C = 2at-matches its final performance.
Figure 9: Refined SF-AdamW. Left:
with 31 E {0.1, 0.5} and C' = 50. Middle:
Our Refined SF Method. We introduce an additional decoupling parameter C' and redefine Ct = 1/t in (SF) as
(1- which then leads
C t C-1 T
(
az
The full derivation is given in Appendix D.4, and pseudocode is provided in Algorithm 2. As shown in Figures 14 and 15, the weights (az) closely follow the theoretical approximation across different values of B. C, and T. Notably, both vanilla and refined SF use an averaging window that widens with T. unlike ixxhhwidth schemes such as EWA.
This modification makes the averaging weights (ax) depend solely on C, allowing /3 to independently control the momentum yı. In other words, C decouples the momentum and averaging behavior. Below, we present preliminary experiments demonstrating the benefits of this refinement.
Empirical gains. We evaluate the performance of our refinement to -F-AdamM using the experimen- tal setup as in Section Keeping all other hyperparameters fixed as the tuned vanilla configuration and varying only C' (with C' = 1/(1 31) recovering the original method), we observe:
Momentum robustness. For 81 {0.1, 0.5), the refined SF-Adami allows x, to match or outperform ye, in contrast to the underperformance of x, reported in Observation (Figure 9-1eft), Improved best-case performance. In the best vanilla setup (3) 0.95), setting C = 200 leads to further reductions in validation loss; similar gains are observed with 31 = 0.9 (Figure 9-middle). Latgh-batch setting. With 2M-token batches, vanilla SF-AdamW (3, 0.98) lags behind AdamW with cosine schedule, whereas setting C = 200 matches its final performance (Figure 9-right).
To further examine the sensitivity to the new hyperparameter C, we conduct a set of controlled sweeps (Tables 5 and 6 in Appendix C. Across both optimal and suboptimal 31. we find that refined SF-AdamW consistently outperforms the vanilla baseline over a wide range of C, demonstraing its robustness to this parameter. These results show that decoupling momentum and averaging via intro- ducing C' eliminates the tradeoff inherent in vanilla SF, yielding more robust and scalable optimizer.
6 Conclusion
We presented a principled view of Schedule-Free (SF) methods by studying its behavior through the geometry the river-valley loss landscape. Our analysis shows that SF methods naturally follow the river without requiring explicit LR decay weight averaging, making them a compelling approach for scalable pretraining. We further provided theoretical insights grounded in Edge of Stability and central flow dynamics. Building on this understanding, proposed refined variant that decouples momentum from averaging, improving both robustness and performance.
While our findings highlight the potential of SF methods, several open questions remain. Our theoretical analysis relies on simplifying assumptions, and validating the central flow approximation in deep learning is natural next step. Furthermore, extending the river-valley framework to analyze other modern optimizers, as well as exploring their integration with SF methods, a promising direction for further investigation. Lastly, our experiments are limited to small-scale language models due to computational constraints. Scaling these findings larger models and longer training durations remains an important direction for future work.
10

---

Acknowledgments and Disclosure of Funding
This work was supported by a National Research Foundation of Korea (NRF) grant funded by the Korean government (MSIT) (No. RS-2024-004421203) and the InnoCORE program of the Ministry of Science and ICT (No. N10250156).
References
Kwangjun Ahn and Ashok Cutkosky. Adam with model exponential moving average is effective for nonconvex optimization. In The Thirty-eighth Annual Conference on Neural Information Processing Systems, 2024. URL pesttps://pperre evieww..iw.... net (d=oorummiii4411L000000 4
Kwangjun Ahn, Sebastien Bubeck, Sinho Chewi, Yin Tat Lee, Felipe Suarez, and Yi Zhang. Learning threshold neurons via edge of stability. seirty-seventh Conference on Neural Information Processing Systems, 2023. URL pttttp:://pperreveee wtt.://oooommddded/oeeeeeee.
8
Kwangjun Ahn, Gagik Magakyan, and Ashok Cutkosky. General framework for mnlin-to-nonncome conversion: echedule-free SGD is also effective for nonconvex optimization. In Aarti Singh, Maryam Fazel, Daniel Hsu, Simon Lacoste-Julic Julien, Felix Berkenkamp, Tegan Maharaj. Kiri Wagstaff, and Jerry Zhu, editors, Proceedings of the 42nd International Conference Machine Learning, volume 267 of Proceedings of Machine Learning Research, pages 772-795. PMLR, 13-19 Jul 2025. httpst http::///ooceedeeeeeee mlr. press/ vv677/ ahn25a. html. 4
Sanjeev Arora, Zhiyuan Li, and Abhishek Panigrahi. Understanding gradient descent on the edge of stability in deep learning. In Kamalika Chaudhuri, Stefanie Jegelka, Le Song, Csaba Szepesvari, Gang Niu, and Sivan Sabato, editors, Proceedings of the 39th International Conference on Machine Learning, volume 162 of Proceedings Machine Learning Research, pages 148-1024. PMLR, 17-23 Jul 2022. URL https: http:////cocceeeesee/eee mlr. prwss/veevv vwtps//vve/////e/ 2tt.//rora22..aaaaaaaaaaaaaa 8
Jeremy Cohen, Simran Kaur, Yuanzhi Li, Zico Kolter, and Ameet Talwalkar. Gradient descent neural networks typically occurs at the edge of stability. In International Conference on Learning Representations, 2021. URL https: httpp::///ppeeeviieeee www.eeeeu//nneeee biooormmid=jjl wowoormmmdd-h--rrtrroa jh-rhr-rrrrr..........
Jeremy Cohen, Alex Damian, Ameet Talwalkar, Zico Kolter, and Jason D. Lee. Understanding optimization in deep learning with central flows. In The Thirteenth International Conference Learning Representations, 2025. URL https: //perrevieeerreeeeee/oeo 9, 29, 30
wwwpearrriiee.nn/fforr..m nww/terunnneeee. SI2rdiiiiiII32 2,
Jeremy M Cohen, Behrooz Ghorbani, Shankar Krishnan, Naman Agarwal, Sourabh Medapati, Michal
Badura, Daniel Suo, David Cardoze, Zachary Nado, George E Dahl, al. Adaptive gradient methods at the edge stability. arXiv preprint arXiv:2207.14444 2022. 8
George E Dahl, Frank Schneider, Zachary Nado, Naman Agarwal, Chandramouli Shama Sastry, Philipp Hennig, Sourabh Medapati, Runa Eschenhagen, Priya Kasimbeg, Daniel Suo, et al. Bench- marking neural network training algorithms. arXiv preprint arXiv:2306.07179 2023.
Alex Damian, Eshaan Nichani, and Jason D. Lee. Self-stabilization: The implicit bias of gra- dient descent the edge of stability. In The Eleventh International Conference Learning Representations, 2023. URL https:// / .ntt meet/forumidde www.:/tss............. 8
Damek Davis, Dmitriy Drusvyatskiy, and Liwei Jiang. Gradient descent with adaptive stepsize converges (nearly) linearly under fourth-order growth, 2024. URL https://arxiv.. www.aiiii..rgga////oo/aaaa 9409.197911 2
Aaron Defazio, Xingyu Yang, Ahmed Khaled, Konstantin Mishchenko, Harsh Mehta, and Ashok Cutkosky. The road less scheduled. Advances in Neural Information Processing Systems, 37: 9974-10007, 2024. 1, 2, 4, 6, 8, 20
Leo Gao, Stella Biderman, Sid Black, Laurence Golding, Travis Hoppe, Charles Foster, Jason Phang, Horace He, Anish Thite, Noa Nabeshima, Shawn Presser, and Connor Leahy. The Pile: An 800gb dataset of diverse text for language modeling. arXiv preprint arXiv:2101.00002 2020. 5, 21
11

---

Alexander Hägele, Elie Bakouch, Atli Kosson, Loubna Ben allal, Leandro Von Werra, and Martin Jaggi. Scaling laws and compute-optima training beyond fixed training durations. In The Thirty-eighth Annual Conference Neural Information Processing Systems, 2024. URL https: wwweerreviex.netfoorune (irpffrr 3, 6
Jordan Hoffmann, Sebastian Borgeaud, Arthur Mensch, Elena Buchatskaya, Trevor Cai, Eliza Ruther- ford, Diego de las Casas, Lisa Anne Hendricks, Johannes Welbl, Aidan Clark, Tom Hennigan, Eric Noland, Katherine Millican, George van den Driessche, Bogdan Damoc, Aurelia Guy, Simon Osindero, Karen Simonyan. Erich Elsen, Oriol Vinyals, Jack William Rae, and Laurent Sifre. Training compute-optimal large language models. In Alice H. Oh, Alekh Agarwal, Danielle Belgrave, and Kyunghyun Cho, editors, Advances Neural Information Processing Systems, 2022. URL https: http:///pppeeeeeeeeeeee wttp://oooooon///dB/BBaAl
5
Shengding Hu, Yuge Tu, Xu Han, Chaoqun He, Ganqu Cui, Xiang Long, Zhi Zheng, Yewei Fang. Yuxiang Huang, Weilin Zhao, et al. MiniCPM: unveiling the potential of small language models with scalable training strategies. arXiv preprint urXiv:2404.06392 2024. 1,
Pavel Izmailov, Dmitrii Podoprikhin, Timur Garipov, Dmitry Vetrov, and Andrew Gordon Wilson. Averaging weights leads to wider optima and better generalization, 2019. URL https://arxiv.
org/abs/ 110//1103..0 00407. 3
Priya Kasimbeg, Frank Schneider, Runa Eschenhagen, Juhan Bae, Chandramouli Shama Sastry, Mark Saroufim, BOYUAN FENG, Less Wright, Edward Z. Yang, Zachary Nado, Sourabh Medapati, Philipp Hennig, Michael Rabbat, and George E. Dahl. Accelerating neural network training: An analysis of the algoperf competition. In The Thirleenth International Conference on Learning Representations, 2025. URL /tstts:://penreviiese wwwp:/wwwww..eeeeeeeeeeeeeeeeeee hett/foorundddCCMM5ssr3aa
4
Yunshui Li, Yiyuan Ma. Shen Yan, Chaoyi Zhang, Jing Liu, Jianqiao Lu, Ziwen Xu, Mengzhao Chen, Minrui Wang, Shiyi Zhan, et al. Model merging in pre- pre-training of large language models, arXiv preprint arXiv:2505. /2082, 2025. 4
Aixin Liu, Bei Feng, Bing Xue, Bingxuan Wang, Bochao Wu, Chengda Lu, Chenggang Zhao, Chengqi Deng, Chenyu Zhang, Chong Ruan, et al. DeepSeck-V3 V3 Technical Report. arXiv preprint arXiv:2412.19437, 2024. 1
Ilya Loshchilov and Frank Hutter. SGDR: Stochastic gradient descent with warm restarts. In International Conference on Learning Representations, 2017. URL https://openreviiet net/ 1
Kairong Luo, Haodong Wen, Shengding Hu, Zhenbo Sun, Maosong Sun, Zhiyuan Liu, Kaifeng Lyu, and Wenguang Chen. A multi-power law for loss curve prediction across learning rate schedules. In The Thirteenth International Conference on Learning Representations, 2025. URL https://openreviie wwtp:eeeeeeoetnieteeeeeee dtte//foorrmmmmkKKSnttIIII 3
Depen Morwani, Nikhil Vyas, Hanlin Zhang, and Sham Kakade. Connections between schedule-free optimizers, ademamix, and accelerated sgd variants. arXiv preprint arXiv:2502.02433, 2025. 4, 8,
Boris T Polyak and Anatoli B Juditsky. Acceleration of stochastic approximation by averaging. SIAM journal on control and optimization, 30(4):838-88) 885, 1992. 3
Alec Radford, Jeff Wu, Rewon Child, David Luan, Dario Amodei, and Ilya Sutskever. Language models are unsupervised multitask learners. 2019. 21
David Ruppert. Efficient estimations from a slowly convergent robbins oobiios-monro monro process. Technical report, Cornell University Operations Research and Industrial Engineering, 1988. 3
Mark Sandler, Andrey Zhmoginov, Max Vladymyrov, and Nolan Miller. Training trajectories, mini- batch losses and the curious role of the learning rate, 2023. URL https://arxiv org/abs/ 2301.02312. 4
Noam Shazeer. Glu variants improve transformer. arXiv preprint arXiv:2002.05202, 2020. 21
12

---

Daria Soboleva, Faisal Al-Khateeb, Robert Myers,
Jacob R Steeves, Joel token cleaned and 627B
Hestness, and Nolan Dey. deduplicated version of RedPajama.
SlimPajama: A
a ai/blog/
s1 slimpaj aama-jama-aae ama-a. toke-token cleaned- http:///wggiggaacccccaaaaaatataa
-f
June 2023. URL
co/ htto/daaaett//eeeeee https///eeerrrbrrseaaaaaaaaaaaae 627B. 4, 21
Minhak Song and Chulhee Yun. Trajectory alignment: Understanding the edge of stability phenomenon via bifurcation theory. In A. Oh, T. Naumann, A. Glober- son, K. Saenko, M. Hardt, and S. Levine, editors, Advances in Neural Informa- tion Processing Systems, volume 36, pages 71632-71682. Curran Associates, Inc., 2023. URL https://proceeddings newrips cc/paper__a saaper_flles_ httppe///leeeaap//////////////// pttpppp/p22///ieeeee 222225562.6aa0000222222222222222
Conferecc waw.aneeee.....eee............ee 8
Minhak Song, Kwangjun Ahn, and Chulhee Yun. Does SGD really happen in tiny subspaces? In The Thirteenth International Conference on Learning Representations, 2025. URL https:
//openreview.neee htw.:/oooorrrttttee1 2
Jianlin Su, Murtadha Ahmed, Yu Lu, Shengfeng Pan, Wen Bo, and Yunfeng Liu. Roformer: Enhanced transformer with rotary position embedding. Neurocomputing. 568: 167:137063, 2024. ISSN 0925- 2312. doi: com http://aceeeetttccee////////////
mec..mm........mmm20222..3.. 127063. URI. https:::::: / //www. www.sscciieeecerce.o.co
21
Hugo Touvron, Thibaut Lavril, Gautier Izacard, Xavier Martinet, Marie-Anne Lachaux, Timothée Lacroix, Baptiste Rozière, Naman Goyal, Eric Hambro, Faisal Azhar, et al. Llama: Open and efficient foundation language models. arXiv preprint arXiv:2302.1 13971, 2023a. 4
Hugo Touvron, Louis Martin, Kevin Stone, Peter Albert, Amjad Almahairi, Yasmine Babaei, Nikolay Bashlykov, Soumya Batra, Prajjwal Bhargava, Shruti Bhosale, et al. Llama 2: Open foundation and fine-tuned chat models. arXiv preprint arXiv:2307.0922 2023b.
Zixuan Wang, Zhouzi Li, and Jian Li. Analyzing sharpness along GD trajectory: Progressive Alice H. Oh, Alekh Agarwal, Danielle Belgrave, and
sharpening and edge of stability. Kyunghyun Cho, editors, Advances in Neural Information Processing Systems, 2022. URL net/ for.middtthhIttrrr.4y.. 8
Kaiyue Wen, Zhiyuan Li, Jason S. Wang. David Leo Wright Hall, Percy Liang, and Tengyu Ma. Understanding harmup-stable-decaa learning rates: A river valley loss landscape view. In The Thirteenth International Conference on Learning Representations, 2025. URL https: /openreview wwwpeeeeewwn.eeeeeeeeeeeee tame/fooruntt fd:mnt d=m5 1BgoqvbP. 2,
Chen Xing, Devansh Arpit, Christos Tsirigotis, and Yoshua Bengio. A walk with sgd. arXiv preprint arXiv: 2018. 2
Xiaohua Zhai, Alexander Kolesnikov, Neil Houlsby, and Lucas Beyer. Scaling vision transformers. In Proceedings of the IEEF/CVF conference on computer vision and pattern recognition, pages 22104-22113 2022. 3
Biao Zhang and Rico Sennrich. Root mean square layer normalization. In H. Wal- lach, H. Larochelle, A. Beygelzimer, F. d'Alché-Bue, E. Fox, and R. Garnett, edi- tors, Advances in Neural Information Processing Systems, volume 32. Curran Associates, Inc., 2019. URL https://proceeddiee 16811998822220003000111111111113
neurips cc/ ac/paper__. lapp//pplleessss///////aaaaaeeee files/ www.fffleesleeeeeeeeee ppppp/ppp9///ii//////
pdf. 21
Hanlin Zhang, Depen Morwani, Nikhil Vyas, Jingfeng Wu, Difan Zou, Udaya Ghai, Dean Foster, and Sham M. Kakade. How does critical batch size scale in pre-training? In The Thirteenth International Conference on Learning Representations, 2025. URL 3, 4,
net/
Xingyu Zhu, Zixuan Wang, Xiang Wang, Mo Zhou, and Rong Ge. Understanding odge-of-stabiling training dynamics with minimalist example. In The Eleventh International Conference on Learning Representations, 2023. URL
13

---

NeurIPS Paper Checklist
1. Claims
Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?
Question:
Answer: [Yes]
Justification: Yes, the abstract and introduction clearly state the paper's main contributions, accurately reflecting its scope and summarizes the results.
Guidelines:
The answer NA means that the abstract and introduction do not include the claims made in the paper.
The abstract and/or introduction should clearly state the claims made, including the contributions made in the paper and important assumptions and limitations. A No or NA answer to this question will not be perceived well by the reviewers.
The claims made should match theoretical and experimental results, and reflect how much the results can be expected to generalize to other settings.
is fine to include aspirational goals motivation long it is clear that these goals not attained by the paper.
Limitations
Question: Does the paper discuss the limitations of the work performed by the authors?
Answer: [Yes]
Justification:
We discuss the limitations of our work Section 6.
Guidelines:
The answer NA means that the paper has limitation while the answer No means that the paper has limitations, but those are not discussed in the paper.
The authors are encouraged to create a separate "Limitations" section in their paper.
The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (e.g., independence assumptions, noiseless settings, model well-specification, asymptotic approximations only holding locally). The authors should reflect on how these assumptions might be violated in practice and what the implications would be.
The authors should reflect on the scope of the claims made, e.g., if the approach was only tested on a few datasets or with a few runs. In general, empirical results often depend on implicit assumptions, which should be articulated.
The authors should reflect on the factors that influence the performance of the approach. For example, facial recognition algorithm may perform poorly when image resolution is low or images are taken in low lighting. Or a speech-to-text system might not be used reliably to provide closed captions for online lectures because fails to handle technical jargon.
The authors should discuss the computational efficiency of the proposed algorithms and how they scale with dataset size.
If applicable, the authors should discuss possible limitations of their approach to address problems of privacy and fairness.
While the authors might fear that complete honesty about limitations might be used by reviewers grounds for rejection, worse outcome might be that reviewers discover limitations that aren't 't acknowledged the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play important role in developing norms that preserve the integrity of the community. Reviewers will specifically instructed to not penalize honesty concerning limitations.
3. Theory assumptions and proofs
Question: For each thcoretical result, does the paper provide the full set of assumptions and a complete (and correct) proof?
Answer: [Yes]
Justification: We state all assumptions and provide complete and correct proofs for each theoretical result in Appendix D.
14

---

Guidelines:
The answer NA means that the paper does not include theoretical results.
All the theorems, formulas, and proofs in the paper should be numbered and cross-referenced.
All assumptions should be clearly stated referenced in the statement of any theorems.
The proofs can either appear in the main paper or the supplemental material, but if they appear in the supplemental material, the authors are encouraged to provide a short proof sketch to provide intuition.
Inversely, any informal proof provided the core of the paper should be complemented by formal proofs provided in appendix or supplemental material.
Theorems and Lemmas that the proof relies upon should be properly referenced.
Experimental result reproducibility
Question: Does the paper fully disclose all the information needed to reproduce the main experi- mental results of the paper to the extent that it affects the main claims and/or conclusions of the paper (regardless of whether the code and data are provided not)?
Answer: [Yes]
Justification: We include necessary implementation details support reproducibility of the main experimental results Appendix
Guidelines:
The answer NA means that the paper does not include experiments.
the paper includes experiments, No answer this question will not be perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data are provided or not.
If the contribution is dataset and/or model, the authors should describe the steps taken to make their results reproducible verifiable.
Depending on the contribution, reproducibility can be accomplished in various ways. For example, if the contribution is novel architecture, describing the architecture fully might suffice, or if the contribution is a specific model and empirical evaluation, it may be necessary to either make it possible for others to replicate the model with the same dataset, or provide access to the model. In general. releasing code and data is often one good way to accomplish this, but reproducibility can also be provided via detailed instructions for how replicate the results, access to a hosted model (c.g., in the case of a large language model), releasing of a model checkpoint, other means that appropriate to the research performed.
While NeurIPS does not require releasing code, the conference does require all submissions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example
(a) the contribution primarily a new algorithm, the paper should make it clear how to reproduce that algorithm.
(b) If the contribution is primarily new model architecture, the paper should describe the architecture clearly and fully.
(c) the contribution is new model (e.g., a large language model), then there should either be a way to access this model for reproducing the results or a way to reproduce the model (e.g., with an open-source dataset or instructions for how to construct the dataset).
(d) We recognize that reproducibility may tricky in some cases, in which case authors are welcome to describe the particular way they provide for reproducibility. In the case of closed-source source models, it may be that access to the model is limited in some way (c.g.. to registered users). but it should be possible for other researchers to have some path to reproducing or verifying the results.
Open access data and code
Question: Does the paper provide open access the data and code, with sufficient instructions faithfully reproduce the main experimental results, as described in supplemental material?
Answer: [Yes]
Justification: We provide open access to the data and code along with instructions reproducing the main experimental results in Appendix C.
15

---

Guidelines:
The answer NA means that paper does not include experiments requiring code.
Please the NeurIPS code and data submission guidelines hattpss:/nnips.ccc/pppl
guides/ SeiideeCodeSubiii Stte//ooeeubiieseeee onsolicy) for more details.
While we encourage the release of code and data, we understand that this might not be possible, so "No" is an acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (e.g., for new open-source benchmark).
The instructions should contain the exact command and environment needed the results. See the NeurIPS code and data submission guidelines (https: / //niss...oo//// wttppp.ccpuupplp.p//eeepappptppp guides Salides//CoeSSumiiie issi onPolicy) for more details.
reproduce
cc/ cc/public/
The authors should provide instructions on data access and preparation, including how to access the raw data, preprocessed data, intermediate data, and generated data, etc.
The authors should provide scripts to reproduce all experimental results for the new proposed method and baselines. If only subset of experiments are reproducible, they should state which ones are omitted from the script and why.
At submission time, to preserve anonymity, the authors should release anonymized versions (if applicable).
Providing as much information as possible in supplemental material (appended the paper) is recommended, but including URLs to data and code is permitted.
6. Experimental setting/details
Question: Does the paper specify all the training and test details (e.g., data splits, hyperparameters, how they were chosen, type of optimizer, etc.) necessary understand the results?
Answer: [Yes]
Justification:
We provide the implementation details in Appendix B.
Guidelines:
The answer NA means that the paper does not include experiments.
The experimental setting should be presented the core of the paper to a level detail that is necessary to appreciate the results and make sense of them.
The full details can be provided either with the code, in appendix, or as supplemental material.
Experiment statistical significance
Question: Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?
Answer: [No]
Justification: We do not report error bars due to the high computational cost of repeated runs required for statistical significance analysis.
Guidelines:
The answer NA means that the paper does not include experiments.
The authors should answer "Yes" if the results are accompanied by error bars, confidence intervals, or statistical significance tests, least for the experiments that support the main claims the paper.
The factors of variability that the error bars are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, or overall run with given experimental conditions).
The method for calculating the error bars should be explained (closed form formula, call to library function, bootstrap, etc.)
The assumptions made should be given (e.g., Normally distributed errors).
It should be clear whether the error bar is the standard deviation the standard enror of the mean.
It is OK to report 1 I-sigma error bars, but one should state it. The authors should preferably report 2-sigma error bar than state that they have 96% CI, if the hypothesis of Normality of errors is not verified.
16

---

For asymmetric distributions, the authors should be careful not to show in tables or figures symmetric error bars that would yield results that are out of range (e.g. negative error rates). error bars are reported in tables or plots, The authors should explain in the text how they were calculated and reference the corresponding figures or tables in the text.
8. Experiments compute resources
Question: For each experiment, does the paper provide sufficient information on the computer resources (type of compute workers, memory, time of execution) needed reproduce the experi- ments?
Answer: [Yes]
Justification: We report the type of compute resources, memory, and runtime details needed to reproduce each experiment in Appendix B.
Guidelines:
The answer NA means that the paper does not include experiments.
The paper should indicate the type of compute workers CPU or GPU, internal cluster, or cloud provider, including relevant memory and storage.
The paper should provide the amount of compute required for each of the individual experimental runs as well as estimate the total compute.
The paper should disclose whether the full research project required more compute than the experiments reported in the paper (e.g.. preliminary or failed experiments that didn't make it into the paper).
Code of ethics
Does the research conducted in the paper conform, in every respect, with the NeurIPS Code Ethics https: http:::///euris.... httpp//p....cpccccccccccceeeeell
Question:
Answer: [Yes]
Justification: We have reviewed the NeurIPS Code of Ethics and ensured that all aspects of our research comply with its guidelines.
Guidelines:
The answer NA means that the authors have not reviewed the NeurIPS Code Ethics.
If the authors answer No, they should explain the special circumstances that require a deviation from the Code of Ethics.
The authors should make sure to preserve anonymity (c.g., if there a special consideration due laws regulations in their jurisdiction).
10. Broader impacts
Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?
Answer: [Yes]
Justification:
We discuss broader impacts at the beginning of the appendix.
Guidelines:
The answer NA means that there is no societal impact of the work performed.
If the authors answer NA or No, they should explain why their work has no societal impact or why the paper does not address societal impact.
Examples of negative societal impacts include potential malicious or unintended uses (e.g., disinformation, generating fake profiles, surveillance), fairness considerations (e.g., deploy- ment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations.
The conference expects that many papers will be foundational research and not tied par- ticular applications, let alone deployments. However, if there is a direct path to any negative applications, the authors should point out. For example, an improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, not needed to point out that generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster.
legitimate to point out that
17

---

The authors should consider possible harms that could arise when the technology is being used intended and functioning correctly, harms that could arise when the technology is being used intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology.
there are negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition to attacks, mechanisms for monitoring misuse, mechanisms to monitor how a system learns from feedback over time, improving the efficiency and accessibility of MI.).
11. Safeguards
Does the paper describe safeguards that have been put in place for responsible release of data or models that have high risk for misuse (e.g.. pretrained language models, image generators. or scraped datasets)?
Question:
Answer: [NA]
Justification: Our work focuses on foundational algorithmic design and does not involve releasing any data models, so safeguards are not applicable.
Guidelines:
The answer NA means that the paper poses no such risks.
Released models that have a high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere to usage guidelines or restrictions to access the model or implementing safety filters.
Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images.
We recognize that providing effective safeguards challenging, and many papers do not require this, but encourage authors take this into account and make best faith effort.
12. Licenses for existing assets
Question: Are the creators or original owners of assets (e.g., code, data, models), used in the paper, properly credited and are the license and terms of use explicitly mentioned and properly respected?
Answer:
[Yes]
Justification:
We acknowledge the codebases we have referred to in Appendix B.
Guidelines:
The answer NA means that the paper does not use existing assets.
The authors should cite the original paper that produced the code package dataset.
The authors should state which version of the asset used and, possible, include URL.
The name of the license (e.g., CC-BY 4.0) should included for each asset.
For scraped data from a particular source (e.g., website), the copyright and terms of service of that source should be provided.
If assets are released, the license, copyright information, and terms of use in the package should be provided. For popular datasets, paperswi wwrawwithhceeee thcode. com/ datasets has curated licenses for some datasets. Their licensing guide can help determine the license of a dataset.
For existing datasets that are re-packaged, both the original license and the license the derived asset (if it has changed) should be provided.
If this information is not available online, the authors encouraged reach out the asset's
creators.
13. New assets
Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?
Answer: [NA]
Justification:
This paper does not release new assets.
Guidelines:
The answer NA means that the paper does not release assets.
18

---

Researchers should communicate the details of the dataset/code/mmodl as part of their sub- missions via structured templates. This includes details about training, license, limitations, etc.
The paper should discuss whether and how consent was obtained from people whose asset is used.
At submission time, remember to anonymize your assets (if applicable). You can either create anonymized URL or include an anonymized zip file.
14. Crowdsourcing and research with human subjects
For crowdsoureing experiments and research with human subjects, does the paper include the full text of instructions given to participants and screenshots, if applicable, as well as details about compensation (if any)?
Question:
Answer: [NA]
Justification: This paper does not involve crowdsourcing nor research with human subjects.
Guidelines:
The answer NA means that the paper does not involve crowdsourcing research with human subjects.
Including this information in the supplemental material is fine, but if the main contribution of the paper involves human subjects, then as much detail possible should be included in the main paper.
According to the NeurlPS Code of Ethics, workers involved in data collection, curation, or other labor should paid least the minimum wage in the country of the data collector.
15. Institutional review board (IRB) approvals equivalent for research with human subjects
Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed the subjects, and whether Institutional Review Board (IRB) approvals (or an equivalent approval/review based on the requirements of your country or institution) were obtained?
Answer: [NA]
Justification: This paper does not involve crowdsourcing nor research with human subjects.
Guidelines:
The answer NA means that the paper does not involve crowdsourcing research with human subjects.
Depending on the country in which research is conducted, IRB approval (or equivalent) may be required for any human subjects research. If you obtained IRB approval, you should clearly state this in the paper.
We recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors adhere to the NeurIPS Code of Ethics and the guidelines for their institution.
For initial submissions, do not include any information that would break anonymity (if applica- ble), such as the institution conducting the review.
16. Declaration of LLM usage
Question: Does the paper describe the usage of LLMs if it is an important, original, or non- standard component of the core methods in this research? Note that if the LLM used only for writing, editing, or formatting purposes and does not impact the core methodology, scientific rigorousness, or originality of the research, declaration is not required.
Answer: [Yes]
Justification: We declare LLM usage, as our work directly studies scalable pretraining strategies and centers algorithmic designs for training L.MMs.
Guidelines:
The answer NA means that the core method development in this research does not involve LLMs any important, original, non-standard components. Please refer our LLM policy lnttps://neure should or should not be described.
for what
19

---

Appendix
A Pseudocode for Schedule-Free Adam W
20
Experimental Details
21
B.1 Language Model Experiments
21
B.2 Edge Stability Experiments on CIFAR-
22
Additional Results
23
C.1 Sensitivity to the Refinement Parameter C
25
D Omitted Derivations and Proofs
26
D.I Proof of Stability Threshold of Schedule-Free Optimizer
.t..............................
26
D.2 Deriving the Reformulation Schedule-Free Optimizer
28
D.3 Deriving the Central Flow of Schedule- Free Optimizer
28
D.4 Omitted Calculations in Section
Impact Statement
This work studies scalable strategies for language model (LM) pretraining, with a focus on foun- dational algorithmic design. While the improper application of these algorithms in downstream tasks may lead to LMs producing harmful, offensive, privacy-violating content, such applications fall outside the scope of this paper. Our contribution is limited understanding and improving the pretraining algorithm itself.
A Pseudocode for Schedule-Free AdamW
For completeness, present the pseudocode for the Schedule-Free AdamW (SF- (SF-AdamW) algorithm in Algorithm 1, along with our proposed refinement in Algorithm 2. The refinement introduces decoupling parameter C independently control the averaging window, addressing the coupling issue discussed in Section 5.
Algorithm SF-Adami [Defazio et al., 2024]
1: Input: 31, learning rate 7, decay X, warmup steps T warmmpp B1, B2,
2: 21 21
3: to = 0
4: for = to T do
5: V1 (1- 31)2 + B1
6:
E 0f(y,, Ge)
7:
82t1-1 (1
B2)9g
8: & = 0 35)
9:
Tt 7 ( warmmp))
I+1 2t 7tBe/ (V 6) 7oA
11: C--1 E
12: Fe+1 (1 c++1) 24. 02+12+11
end for
14: Return IT+1
20

---

Algorithm 2 Refined SF-AdamW (with decoupling parameter C)
1: Input: 21, learning rate "y, decay X, warmup steps T Twrnupp 31, $2, e, decoupling parameter C
2: 21 21
3: to = 0
4: for = 1 to T do
5: Y/ (1- B1)zz + 313x
6:
E 0f(y,, (t)
7: Tt 3201-11 (1 B2)gi
8: & = v 82)
9: 7t = Y min (
2++1 2t 7/9t/ (Vo, + c) Pxe
11: C-11 1 = min
(1 B1)C, 1
Dim 7
12: Zt+1 (1 Cc+1) + 1++12+1
13: end for
14: Return FT+1
B Experimental Details
B.1 Language Model Experiments
Codebase. All language model experiments are implemented using the public 3lm-baselinee ines codebase: https: //g githuu...bb.... tion of custom optimizer implementations to support the Schedule-Free and refined Schedule-Free methods. All other components (model, data pipeline, logging) remain unchanged.
halmmmaaalliees Our only modification is the addi-
Architectures. Our main experiments use a 224M-parameter LLaMA-style decoder-only trans- former with SwiGLU activations [Shazeer, 2020], RoPE embeddings [Su et al., 2024], RM- SNorm [Zhang and Sennrich, 2019], and alternating attention/MLP blocks (12 layers, 12 attention heads, hidden dimension 768). Additional results in Appendix C verify our findings with 124M- baselines with
parameter GPT-2-style transformer. Both architectures are implemented standard design choices.
Datasets. Our main experiments use the 6B-token subset of the SlimPajama dataset [Soboleva et al., 2023], available on Hugging Face.? We tokenize with the GPT-2 tokenizer [Radford et al., 2019], which has a vocabulary size 50,304. In Appendix C, we also validate our main findings on WebTeex2ex12 dataset [Gao et al., 2020],* 4 using the same setup.
Training details. We train models using AdamW and SF SF-AdamW, with short warmup phase com- prising 5% of total steps. For tage-batch runs with cosine decay, the learning rate is annealed to 10% of its peak. Main experiments use a batch size of 1,024 sequences context length 512 tokens (0.5M tokens total), trained for 5,000 steps, amounting to roughly 2.5B tokens (~ 1x Chinchilla scale). Large-batch experiments use a 2M-token batch size and 2,500 steps (~5B tokens, or ~ 2x Chinchilla scale) to evaluate efficiency in the overtrained regime. Validation is performed during training using 3,200 sequences context length 512 tokens (~1.6M tokens) to compute validation loss (perplexity) curves. For computing EWA,
a decay factor of 0.99 for all experiments.
Hyperparameters.
We fix the weight decay to 0.1 all experiments. Gradient clipping is set 1.0 for AdamW and disabled (0.0) for SF-AdamW. We perform sweeps over the learning rate, momentum parameters, and the decoupling parameter C' (for refined AAdaWW.. Full configurations are provided in Tables to 4.
Compute Resources. All experiments are conducted on a single node with 8 NVIDIA A6000 GPUs (48GB VRAM each) using data-parallel training. A typical full 5,000-step run with 0.5M-token batch size takes approximately hours.
21

---

Optimizer
Learning Rate
(81,
C (Refined SF)
AdamW
[5e-4, 1e-3, 2e-3, Se-3)
[00.9, 0.95), (0.95, 0.99)
SF-AdamW
[1e-3, 2e-3)
[(0.9, 0.99), (0.95, 0.99), (0.98, 0.99)
Refined SF-AdamW
2e-3
(0.95, 0.99)
[20, 50, 100, 200, 500)
Table 1: Hyperparameter sweep: Main experiments. Grid of hyperparameters used in our main
experiments (SlimPajama-6B, 0.5M-token batch size), including learning rates, momentum pairs (3,, B2), and the decoupling parameter C' (for Refined -F-AdamW). Bold entries indicate the best- performing configuration for each optimizer.
Optimizer
Learning Rate
(81, B2))
C (Refined SF)
AdamW (Cosine LR)
[Se-4, 1e-3, 2e-3, Se-3] [le-3, 2e-3, Se-3) 2c-3
{(0.9, 0.95). (0.95, 0.99))
SF-AdamW
((0.9, 0.99). (0.95, 0.99). (0.98, 0.99))
Refined SF AdamN
(0.98, 0.99)
(00, 50, 100, 200, 500)
Table Hyperparameter sweep: Large-batch experiments. Grid of hyperparameters used in our large-batch experiments (SlimPajama-66) 2M-token batch size), including learning rates, momentum pairs (2), and the decoupling parameter C' (for Refined SF-AdamW). Bold entries indicate the best-performing configuration.
Optimizer
Learning Rate
(81,
C (Refined SF)
AdamN
(5c-4, 1e-3, 2c-3) (Ic-3, 2e-3, 5c-3)
((0.9, 0.95), (0.95, 0.99))
SF-AdamW
(00.9, 0.99), (0.95, 0.99), (0.98, 0.99))
Refined SF-AdamW
20-3
(0.95, 0.99)
550, 100, 200, 500)
Table 3: Hyperparameter sweep: OpenWebt Text2 experiments. Grid of hyperparameters used in our additional experiments on Open WebText2 (0.5M-token batch size), including learning rates, momentum pairs (31, (2), and the decoupling parameter C' (for Refined SF-AdamW). Bold entries indicate the best-performing configuration.
Optimizer
Learning Rate
(81, .2)
C' (Refined SF)
AdamW (Cosine LR)
(Se-4, 1e-3, 2e-3, 5c-3) (Se-4, 1e-3, 2e-3, Se-3)
(00.9, 0.95), (0.95, 0.99)) (00.95, 0.99), (0.98, 0.99))
SF-AdamW
Refined SF Adamw
(2e-3, 5e-3, le-2)
(0.98, 0.99)
120, 50, 100, 200, 500, 1000)
Table 4: Hyperparameter sweep: OpenWehText2 large-batch experiments. Grid of hyperpa- rameters used in our additional large-batch experiments on OpenWebTexx (2M-token batch size), including Icarning rates, momentum pairs (B1, 32), and the decoupling parameter C' (for Refined -ddaW.... Bold entries indicate the best-performing configuration.
B.2
Edge of Stability Experiments on CIFAR-10
We provide the experimental setup for Figure 7 in Section 4.3, where we study whether Schedule-Free methods operate the Edge of Stability in a deep learning setting.
Our experiments build on the public edge- of iestabbiiity codebase," modifying only the optimizer incorporate Schedule-Free methods. The model is 3-layer MLP with hidden width 200 and tanh activations, trained on the first 5,000 samples of CIFAR- 10 using mean squared error (MSE) loss.
For SF-GD, we fix momentum B 0.9 and vary the learning rate, training each run until the loss reaches 0.02. For AddmM, we fix 3, = 0.95, 32 = 0.99, and weight decay at 0.1, and vary the learning rate, training until the loss reaches 0.05.
22

---

C Additional Results
In this section, we present additional experiments on the OpenWebText2 dataset using a 124M- parameter GPT-2 style decoder-only transformer, replicating the setups from the main text. As summarized below, the results confirm that our main findings (Observations to 3) generalize across both datasets and architectures.
As in the main experiments, we perform a grid search to determine the best hyperparameters for both AdamW and Ad-AdamW, and evaluate whether LR decay or EWQ improves performance. In this setup. the best perfforero configuration for Adamw is (51. 88) = (0.9, 0.95) with LR 1e-3, while SF-AdamW achieves optimal results with (31, .2) = (0.95, 0.99) and L.R 2e-3. We observe that neither LR decay nor EWA improves performance for SF- Adami, indicating that it already tracks the river closely (see Figure 10).
Observation 1:
Observation 2: We also assess the behavior of SF-Adaml under suboptimal momentum configura- tions, using 3 E 0.5} while keeping all other settings fixed. For both values, we observe that applying short LR decay phase using AdamW significantly improves performance, reducing the validation loss compared to the constant LR baseline. This result shows that SF-Adaml is sensitive the choice of momentum, and suboptimal settings can hinder its ability to effectively follow the river (see Figure 11).
Observation 3: Using the same runs from above, we compute the validation loss at the Yr iterates and their EWA. For suboptimal momentum settings (8, E {0.1,0.5)) the loss at yt is consistently lower than that at Xr, consistent with trends observed in both the toy model and SlimPajama experiments. Moreover, across all momentum values, applying EWA to y, further improves performance, suggesting that y, is more robust to suboptimal momentum and remains better aligned with the river trajectory (see Figure 12).
Refined SF-AdamW: W: We evaluate the refined variant of SF-AdamM by sweeping over C across multiple momentum settings. In the low-momentum regime (31 = 0.5), the refined method enables X2 match outperform yt, addressing the discrepancy observed in the vanilla formulation. In the best-performing vanilla configuration (31 = 0.95, LR 2e-3), setting C = 200 yields further performance gains. Notably, comparable results can also be achieved under suboptimal momentum setting (3) 0.9) by choosing C = 50, demonstrating improved robustness to hyperparameter choices. In the large-batch setting (2M-token batches), vanilla SF-Adam'l with a constant LR ((31, 32) = (0.98, 0.99). LR 2e-3) underperforms relative to AdamW with cosine LR schedule ((31, (3) (0.9, 0.95). LR 5e-3). However, the refined variant with C' 500 (($1, (2) = (0.98, 0.99), LR 5c-3) successfully closes this performance gap. Together. this suggest
that our refined method is robust to momentum and batch size scaling (see Figure 13).
Vanishing B Benefit of LR Decay
Vanisking B Bezefit of EWA
Short Decay Phase of Adam W
70
70
Maaat Coeet 2od
Mlalt
Alaet cland Lit
Perplexity
60
1EM Denp L11363
2k2
0
Perplesity
60
SEAM
Perpllesity
60
Acaell E. -
Crac 20c3
15
150
40
40
30
30
10000
2000
3000
4000
5000
1000
2000
3000
4000
5000
1000
2000
3000
4000
5000
Ierations
Berstions
Figure 10: WpenWee Text2 Experiment: SF-Adam closely follows the river, unlike Adam W. Left, Middle: While Adam' benefits from linear LR decay and EWA, SF-AdamM shows no improve- ment from either. Right: A short decay phase of Adaml (with linear LR decay from Ic-4 to 0) leads to sharp loss drop for AdamW, but has no effect when applied to the Adami trajectory- suggestingg that SF- Adami already tracks the river throughout training (Observation 1).
23

---

Short. Decay Phase of AdamW, Suboptimal f
SFAdamW, A: =.5
AdamW Decay LR Ie-4
300
SFAdamW, B: 0.1
Perplexity
AdamW Deasy L.R 1o-4 ++0
2.00
100
1000
4000
Iterations
Figure 11: Open WebText2 Experiment: SF-AdamW with suboptimal momentum fails to follow the river, A short decay phase of AdamW applied to SF SF-Adamli checkpoints with B1 E {0.1, results in sharp loss drop, unlike the case with B1 0.95 (Observation 2).
A1 0.95, Optimal
A 0.5, Suboptima]
0g - 0.1, Suboptimal
150
250
EWW of
125
9:
EMA of
EWA ofy.
Derglexity
4.
Perplexity
Derpleaity
150
50
25
25
5000
4000
5000
4000
Iberations
Itarations
Iteratioss
Figure 12: OpenWebText2 Experiment: Performance of X4, Yt, and the EWA of yı under varying 31. For suboptimal 31, yt outperforms X,, and across all momentum settings, the EWA of y, achieves the lowest loss (Observation 3).
Suboptimal sub 25.
EF-AdamW Optimal
Large
100
50
STJaK, di
Acaed with Couna
Reteel, di f1-c00.C-200 200
Parplosiy
80
Perpleuity
Purplesity
80
10
30
40
30
10000
2000
3000
4000
5000
1000
2000
3000
4000
5000
300
1000
1500
2000
2500
Etaratioons
Itarations
Iterations
Figure 13: Open Teb2e Experiment: Refined SF-AdamW. W. Left: Performance of x, and yr iterates using the refined SF-AdamW with B1 = 0.5 and C' = 50. Middle: Refined SF-AdamW with (B1, C) E {00995 200), (0.9, 50)} achieves improved performance over the best vanilla SF- Adamm/ run. Right: Under large batch size (2M tokens), vamilla SF Adamm with 31 = 0.98 underperforms compared AdamW with a cosine schedule, while the refined final performance.
with 500 matches its
24

---

C.1 Sensitivity to the Refinement Parameter C
To examine the sensitivity refined SF-AdamW the choice of the refinement parameter C, conduct experiments on SlimPajama (0.5M and 2M token batch sizes). We sweep a range of C values and report test perplexity over 9M held-out tokens. The results are summarized in Table 5 and Table 6.
Our results demonstrate that refined SF-Adamil is consistently robust to the choice of C. Across a broad range of settings, it outperforms vanilla SF SF-Adami (i.e., C 1/(1 - 31)). For instance, in Table values such as C E {50, 100, 200, 500} improve performance for 3I 0.9, 0.95, and C' 50, 100 improve performance across all momentum values. Similar trends hold in Table 6, with performance gains persisting even large C. These results indicate that refined SF-Adami remains effective without requiring careful tuning of C.
81
Vanilla
C =
10
20
50
100
200
500
0.1
67 20
37.06
35.80
36.62
36.15
37.27
0.5
41.01
29.32
30.87
29.96
29.57
0.9
27.70
27.70
23.97
24.64
24.93
25.11
0.95
25.12
25.12
23.98
23.60
24.37
24.83
Table 5: Refined SF-Adamt SlimPajama with 0.5M-token batch size. Test perplexity (on 9M held-out tokens) under varying C' values. LR = 2e-3, B2 0.99.
B1
Vanilla
C==3
50
200
500
1000
0.1
110.8
54.80
47.37
44.31
43.80
0.5
54.24
47.80
38.42
38.86
38.49
42.97
0,9
31.34
31.34
29.86
27.68
28.16
27.02
27.88
0.95
27.29
27.29
25.77
25.45
25.77
27.31
0.98
26.09
30.23
26.09
25.49
24.51
23.95
23.88
Table 6: Refined SF SF-AdamW on SlimPajama with 2M-token batch size. Test perplexity (on 9M held-out tokens) under varying C values. LR 2e-3, B: 0.99.
25

---

D Omitted Derivations and Proofs
D.1 Proof of Stability Threshold of Schedule-Free Optimizer
D.1.1 Proof of Proposition 4.1
Proposition 4.1 (Stability Threshold of SF-GD), Consider running SF-GD on a quadratic objective f(w) : Hw w c. If X1 (H) (1-By7' then the iterates {(x,, Vt, 2n)} diverge.
Proof. To begin with, we revisit the update rule of SF-GD, given by
X = (1 - C₁) X + ct Ze,
(1 (3) + xx:
(SF-GD)
2++1 Z. fyt))
On quadratic objective, we get Vf(w)= = Hw By substituting this and combining the last two relations, we get
2i+1 = Z: (Hy: + g)
Zt ((1 Hz; + HHx; + g) = (I - 7(1 - H)H) - B Hx, - 7g
By substituting this to the first relation, we get recurrence relation governing x; as follows:
X++1 =(1 - -c++1)x + c4+12+11
=(1 x++1)x; + c+1 - 711 3)H) Hy 7g)
(1 (++1) II 3yc++ H)x; + c++1 (I - 7(1 8)H) 7c++18
- -1-1 Xt-1
1
((1 4+11 II - Byc++1 H)x, + c++1 (I - 7(1 - H)
- X₁
- 70++1
Ct
C++1 co+1 I - --70++
(1-3 B H X₁
C2
Ct
+ ((0. [c++1
c++1 I I-7 (++1
941 (
(1-
(1
3) Xt-1 70++18.
(I)
Ct
Define (q, a) = (q, A1 (H)) to be the largest cigenvector/cigeevvaue pair of H and ã, = q x d g. Then the sequence {q T x,) diverges and only if the sequence {ø,} diverges. By multiplying q' on both sides Equation (I), we get
q x'x1+1 = I
Cc+1
(1-3
- ) a qixt
ce+1 -74++
Ca
C++1 (4+1 G+11 )0()(11))) (1 7(1 Txx1i
x,-1 774119 g,
from q H = Tqq. From the definition of T., we get
:--1 (1 C++1 C a-1 7
| 1 + (1=)+)) +))+))a)
(0t1 Ct+1 (1 (1-7(1- (11-7(113aa))) (1-7(1-3)a)E--1, -3)a),.-1,
C2
Ct
which is linear tire-iarying second order difference equation governing St. Its asymptotic behavior governed by the limiting recurrence relation:
E (2 - a(1 f))) + (a(1 5) 1)
Two roots of this recurrence relation are given by
\ 2 - 2-a(1 - 3--a(133)7+ + / (2 a(1-3))3 - B) -1) 1)
A = - a - 3) ((2 a(1 -
2
26

---

If2 < (alaa(11)11)) < 4- 422 then
|a1 = 1 (2 a(1 - (31 2
(2 (2-a(1 -37)²2 4(a(1 3)7 -)
2
4
aa 3)7 >1, 1
which implies F, diverges.
If Fa(1 -8)y 2 4 + 2, then X2 can be regarded real valued function with respect to a(1 B- Since 22 is decreasing and A2 < -1 when a(1 8)7 4 2v2, we get X < -1 when
a(l - B)y 2 4 + 2v2, which implies that also diverges.
Since we take 7 and < 1. the condition Tt.
(1-37 implies diverging कn well as
D.1.2 Proof of Proposition 4.2
To show Proposition 4.2, we first prove the following reparameterization lemma.
Lemma D.1. Define -F-PrecondGD by (SF) with A, è V Let {xi} denotes the iterates
of Dr-PorcondGD the objective [w), and let (R,) denote the iterates of SF-GD on the reparame-
terized objective f(w) (PP-1/2w) with initialization p1/221, Then, we have X, pp/22x for all steps t.
Proof. We claim that the equivalence (X,, yr. 2.) = (P1/2x,, pl/2y, p1/2z() holds for t, where (Re, Yt, žr) denotes the iterates the reparametrized objective.
For 1, holds from the definition. Assume that the equivalence holds t. Then, the update for 2++1 is given by
2c+1 2: -2--7V11 -VV((.))
2, -2p-1Pp1/2 (PP
p1/22/ -7P 1p-1/2//6
(inductive hypothesis)
-D1/2(zi xP -P--vf((.))
-P1/22+11-
Meanwhile,
X (1 - -C411) cææ+1
-P1/((11 - cx+1)x+ + cib(+1)
-21/'x41
Yt+r (1 1-8) 33++1
1112((1- - 33)2+11 + (xc+1))
11/2yc+1:
which proves the claim.
Proposition 4.2 (Stability Threshold of SF SF-PrecondCDK). Consider running SF-PrecondGDi on
f(w) w T Hw +g t If \ (P-'H) (I-B then the iterates (x, diverge.
Proof. Recall that SF-PrecondGDW defined by (SF) with A 2 p-1Vf((y)) + Ayt. which is identical
P p-IVf(y.) Xy. = P- P Vg(y)),
where g(w) s(w) + ////P1/2w|1 Therefore, UF-PreconddDNi identical to SF-PrecondDD on the objective g(w).
27

---

Let } be the iterates of SF-GD on the reparameterized objective g(w) = 9(P-1/22 with initialization X1 P1/2x, From Lemma D.I. (xr). the iterates of SF-PrecondGDV, satisfy X, = x,, which implies that if x, diverges then also diverges.
From Proposition 4.1, if A1 (P H XI)
2 then X, diverges. This proves the claim.
(1-D)y²
D.2 Deriving the Reformulation of Schedule-Free Optimizer
We begin by defining the momentum variable:
X2 Z++1
m
7
Using the update rule x, = 0z)
(x)xx X + C2 and ze+1 = -t - -7 we can express m, recursively:
Ze+1
A, (1 - c2)
Xx-1 Ze
A. (1 Co m-m + A..
7
Next, we derive an update rule for y::
Yt+1 (1 - 3) Zz+1 + 3 X
(1 3) (z: - -(s)) + +B((1 c(41)xx 11-12(11
= (1 8)2z + 3xt - (1 - -3,,,,, + Be++1 (2e+1 - X1)
Yt 7 [88t-111 im, +1 -)
Hence, the update rule (SF) can be equivalently written (SFy):
m, (1 -c)mm A., yt+1 = Yt - 7 m, 7(1 3)
Finally, we express x; as a weighted average of yv. Starting from the original definitions:
X (1 C++1 + c++12+1
2++1 = Yt+1 fx(+1
1
we substitute and obtain:
Yi+1 -3x--1
X++1 (1 - c(+1)xx+ + +041
thus we conclude that
(1 - c++1) 9)xe ct41+++
X,+1 (1 c+41 B3 + C++1
D.3
Deriving the Central Flow of Schedule-Free Optimizer
D.3.1 Deriving the Central Flow of Schedule-Free GD
We begin with the reformulated update rule for SF-GD, as derived from (SFy):
m, = (1 - ct m + (y:);),
y++1 Yt 7/844 7(1 f)V))
As in gradient descent, stable training dynamics are often well-approximated by their continuous-time analogs. We can therefore define a corresponding stable flow for SF-GD:
dy
711 (1111))Vf())) B)Vf(y) 3/ 1)m,
dt 1
dm
di 1 c(l) [Vf(y) c(t)m].
28

---

However, at the Edge of Stability, the optimization trajectory deviates from this stable flow. We now derive central flow to characterize the time-averaged behavior of SF-GD under this regime. particularly when a single top eigenvalue remains at the stability threshold. This derivation is not rigorous but follows the ansatz approach used by Cohen al. [2025].
We model the trajectory as
Y: + Pru, where Y: is the time-averaged iterate, uN is the top Hessian eigenvector at yt, and Pr is the scalar displacement along u, By construction, E[pp]] 0. Using a Taylor expansion of V f(y) around the reference point y, we obtain:
ff(y) V + pS(y)u + 1 V Op²),
where Syy) = X1 X1 ( f()) denotes the sharpness at y. Taking expectations, the time-averaged gradient norm becomes:
( [ ~ Vf(g) + Elp.)((y)u +Ele - vf(g) + +EIpAiss))))
Based on these approximations, can derive the following central flow dynamics of yr:
dy
dt 71 3) [V + 30 (t)V(()) 7/ee(t + 1)m.
dm 1 [v/(y) + 2 -c))m]]
c(I)
where (t) models E Ejg?), the instantaneous variance of the oscillations around the central flow trajectory (i.e., the river trajectory).
Recall that at the Edge of Stability, the sharpness equilibrates near the stability threshold. We therefore assume that it remains constant along the central flow trajectory, satisfying
d (S(y)) 0. dt
S(y) (1 - (1-3)r²
There exists unique value of (t) that ensures this condition holds, particularly in the regime where is large, where the coefficient c(t) = 1/t becomes negligible. To compute this value o². apply the chain rule and substitute the central flow dynamics:
dS(v) dt
dy dt
Syy)),
SV) -1 - B) (o
where we approximate c(t) 2 O. Setting
dS(y) = 0 and rearranging gives: dt
2 -
(y)
D.3.2 Deriving the Central Flow of Schedule-Free Sealar Adam
We begin with the reformulated update rule for SF SF-ScalarAdam, as derived from (SFy):
Vt = 32 (1 (2) -/)))//(
m, (1 m))m, vf(y.)
Vui
Yes = yz - 7- 7(1 31) Vf(y.)
Vui
As in gradient descent, stable training dynamics are often well-approximated by their continuous-time analogs. We therefore define corresponding stable flow for SF-ScalarAdam:
dy 7(1 Bi) vs(y) - - + ++11) 1)m,
dt Vi
dm
1
Vf(y) -(()))((t)m)
dt
1 1-((t)
vv
du
1 B2 [IVV[(y)11
dt
B2
29

---

However, at the Edge of Stability, the optimization trajectory deviates from this stable flow. We now derive a central flow to characterize the time-averaged behavior of SF-ScalarAdam under this regime, particularly when single top eigenvalue remains at the stability threshold. This derivation is not rigorous but follows the ansatz approach used by Cohen al. [2025].
We model the trajectory as = y: + prue, where Y: is the time-averaged iterate, uN is the top Hessian eigenvector at yt, and Pr is the scalar displacement along u, By construction, E[pp]] 0. Using a Taylor expansion of V f(y) around the reference point y, we obtain:
f(y) ff(()) + pS(z)u + 3 ( + O(P)),
where S(y) == X1 (v denotes the sharpness at Taking expectations, the time-averaged gradient and squared gradient norm become:
[VV)(y:)) Vf(() Elp.)S()) + | Vf(g) +
/ [V + 2 (V)()), u) + E(P])(()))) [V) Elp))
Based on these approximations, we can derive the following central flow dynamics of yc:
dy
7(1 -1)
Vf + {o*()))()))) - 78ic(t + 1)m,
di
VV
dm
+ [
c(0m) -(())
di
1 ((()
dv 1 B2
32
where (t) models E Eip?). the instantaneous variance of the oscillations around the central flow trajectory (i.e., the river trajectory).
Recall that the Edge of Stability, the preconditioned sharpness equilibrates near the stability threshold, We therefore assume that remains constant along the central flow trajectory, satisfying
S(y) V - (1-B17*
d
S(y)
di
There exists a unique value of (t) that ensures this condition holds, particularly in the regime where / is large, where the coefficient c(l) 1/L becomes negligible. To compute this value of o². we apply the chain rule and substitute the central flow dynamics:
d S(y) v0
dy s
S(y) dv
di
di
2/3/2 dl
7(1 31)
VS(y).
[V((y) + [
VD
Vv
S(y) 1 3-
+ 0² + -v]
2/3/2
32
where we approximate e(l) 0. Setting
d dt
[(y) /v
and rearranging gives:
2² VS(y),-V/it + 911-.11827
1-8,
S(y) +o S(y)
$ S + 2113311327 v2
1-82
S(y)
Using the condition S(y) obtain:
(1-9,)> we substitute Vo = 1/(1 - -Bi)7s(y) into the expression and
SV(Y), vf(y)) + 1- 1y)) (1- [f
t 3y;
2 [2 +
1-9. (1-8,18877 S(y)2
Notably, a² depends only on the current iterate y and the hyperparameters 31, Br, and 7. Moreover, unlike SF-GD, o² does depend on momentum parameters.
30

---

D.4 Omitted Calculations in Section 5
We derive the closed-form approximation of the averaging weights ar under the modified SF method, where the update coefficient is set to
4 -1-
8)C
Under this setting, we show that the induced averaging weights satisfy the approximation
C 6
(
a₂ T
T
Recall that for general {ct,, the iterates xT can be written as a weighted average of past y::
4
II
T
(1 c,)( (3) co)
XT at Yt.
a == (1 (1- a1 a1 8)
+ a
(1-ca) - () + Co
(1-D)C t
Now, substitute Ct
For large i, we approximate:
C
1 (1-B)C
00
(1-
C
C
1 1 BC
Ls=t+1
T
(1-3)0
1+
30
C
exp
BC
exp
C
T
(-9)
exp
L==++
T
C
exp
8=1-1
Using the integral approximation for the harmonic sum:
T
Σ
8=t+1
#=1
we obtain
C
T
C
Σ s=t+1
exp
t
$
C
exp
t
s=1
C
exp
(-C log
+ C Clogt)
CtC-1
TO
Thus, we conclude that
C C(4)
04
T
Figure 14 and Figure 15 show that the averaging weights (az) in refined SF method closely follow the approximation at 2 C C() across different values of 3, C. and T.
31

---

8.
8: -0.5
0.9
oe -())
200
409
600
800
1030
0.000
200
400
600
803
1000
200
400
500
500
1050
0.010
0.000
0.08
0.006
0.006
0.006
0.004
0.004
0.004
0.022
0.022
0.002
0.000
0.000
0.000
0.030
0.010
0.010
0.000
0.008
0.000
0.006
0.006
0.006
0.004
0.004
0.004
0.002
0.002
0.000
0.000
0.000
500
0.000
0.010
0.008
0.006
200
400
600
800
1030
200
400
000
503
1000
200
400
500
800
1000
Bernatons
Figure 14: Averaging weights in the refined SF method. Histogram of {ar)== over T = 1000 iterations for varying values of B and C.
0.0020
0.0020
0.0010
0.0010
0.0010
0.0005
0.0015
5,0055
0.0000
00000
1000
2030
3005
3000
10000
2030
3003
0000
5000
1010
2503
3000
0003
5000
0.0020
0.0020
0.0020
0.0005
0.0000
0.0005
0.0000
9. 0.0000
3,0000
1.000
2000
30.00
5000
0.0020
0.0020
3.0020
0.0015
0.0015
0.0005
0.005
5.0006
0.0009
9. 0.0000
3.0000
0.0929
0.0020
0.0020
0.0055
0.0015
8.0005
0.0003
5,0000
2000
Figure 15: Averaging weights in the refined SF method. Histogram of {at)T= over T 5000 iterations for varying values of 3 and C.
32


Reviews:
[{"review_id": "vfVdi4aUXK", "reviewer": "Reviewer_EzTy", "summary": "The authors revisit Schedule-Free (SF) optimization (SF-AdamW), demonstrating that it navigates the \"river-valley\" structure of loss landscapes without explicit LR decay or memory-heavy weight averaging, and also list some observations, for example, SF-AdamW is highly sensitive to momentum. A refined version of SF-AdamW is proposed, addressing two critical issues: momentum-sensitivity and scalability under large batch sizes.", "questions": "The authors have demonstrated the improved validation loss of their refined SF-AdamW optimizer compared to the original SF-AdamW. Can the refined SF optimizer lead to measurable performance gains on downstream tasks based on different models and benchmarks?", "limitations": "Yes", "rating": 5, "confidence": 2, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 3, "significance": 3, "originality": 3, "strengths_and_weaknesses": "1.\tBased on insightful observation (e.g., SF-AdamW is sensitive to the momentum parameter, an inadequate choice will lead to deviation from the river). The key contribution is the introduction of parameter C, which is used in a reformed Schedule-Free Optimizer to decouple the momentum and averaging behavior, weakening the momentum-sensitivity and resulting in further reductions of validation loss.\n2.\tThe experiments are limited to relatively small-scale models (e.g., 124M parameter transformers decoder) due to computational resource constraints. Given the claim of scalability of the optimizer as well as relevance to large-scale training, a lack of large-scale model experiments somewhat limits the strength of the claims.\n3.\tThe paper admits the lack of error bars report and analysis due to computational constraints, leaving the statistical significance of observed performance improvements uncertain.", "was_revised": true, "final_justification": "Many thanks for the author’s rebuttal comments. Most of my concerns have been addressed. I also updated my ratings for the Significance of the paper.", "rebuttal": "Thank you for your efforts and for appreciating our work. We address your concerns and questions below.\n\n\n### Response to Weaknesses \n> The experiments are limited to relatively small-scale models (e.g., 124M parameter transformers decoder) due to computational resource constraints. / The paper admits the lack of error bars report and analysis due to computational constraints.\n\nWe agree and will explicitly acknowledge this in the Limitations section of the revised manuscript. Our experiments were conducted on 124M-scale models trained up to the Chinchilla-optimal compute budget. This scale was intentionally chosen to balance interpretability of Schedule-Free dynamics with available computational resources.\n\nThe primary goal of our paper is to understand the mechanisms underlying Schedule-Free optimization, rather than to demonstrate state-of-the-art performance. Accordingly, we prioritized controlled ablations and mechanistic analysis over repeated trials or larger-scale benchmarks. In Sections 3 and 4, which focus on analyzing training dynamics rather than making performance claims, we believe error bars would provide limited additional insight. In Section 5, however—where we introduce the refined SF method and claim improvements over vanilla SF—we agree that error bars would strengthen the statistical significance of the results.\n\nDue to the limited time and resources during the rebuttal period, we could not provide these additional runs. We will, however, include error bars for the refined SF experiments in Section 5 in the revised version.\n\n\n### Response to Questions\n> The authors have demonstrated the improved validation loss of their refined SF-AdamW optimizer compared to the original SF-AdamW. Can the refined SF optimizer lead to measurable performance gains on downstream tasks based on different models and benchmarks?\n\nEvaluating downstream task performance is beyond the scope of this work. Our focus is on developing efficient and scalable pretraining algorithms. Assessing downstream performance would require fine-tuning models that have been pretrained using our method, which in turn necessitates pretraining very large models on massive datasets—a level of computation that is currently beyond our resources.\n\nWhile we acknowledge the differences between pretraining and fine-tuning loss landscapes, we believe that our refinement may, to some extent, generalize to fine-tuning tasks. Verifying the effects of decoupling and other findings from pretraining dynamics in the context of fine-tuning tasks remains a promising direction for future work."}, {"review_id": "mp8ge1Kjk4", "reviewer": "Reviewer_kaR1", "summary": "The paper \"Through the River: Understanding the Benefit of Schedule-Free (SF) Methods for Language Model Training\" proposes to study schedule-free optimizers through the lens of river-valley landscape. In particular, the authors conduct several experiments which result in several observations that could explain the loss curve in the view of river-valley landscape. A new refined SF method is also proposed to improve the SF-AdamW.", "questions": "1. First, in my opinion, the river-valley landscape is an interesting observation that is still lack of justifications. In particular, how to define this landscape mathematically? In my humble opinion, it should be that at the point in side the valley, there is an eigen direction of Hessian that decrease the function value the most (river), and an eigen direction which actually is nearly convex quadratic (valley). If we don't quantify this, how are we suppose to really understand the loss curves in Section 3?\n2. Therefore, I'm not sure about all the experiment in Section 3. Figure 2 and 3 show that SF-AdamW would not improve with LR decay or EWA, and it is shown that AdamW is benefited from lr decay, however I'm not sure how this is related to the \"river-valley\" landscape. In particular, for a convex quadratic objective, you can readily see a faster decrease of the loss with decay lr if your initial lr is too large (which leads to oscillation), and I don't see where each part of the river and the valley play a role here. It is also shown that SF method is sensitive to momentum parameter, which again could be explained without referring to the river-valley landscape...\n3. Section 4 gains some observation on the toy objective. However this is not even a stochastic optimization problem, and I would question if the observation could be carried out in large scale.\n4. For the refined optimizer, it may be less sensitive to the momentum parameters $\\beta_1$ and $\\beta_2$, but we have a new hyperparameter $C$ to tune. It might still be sensitive to $C$?\n5. I almost forgot but in line 160, the authors claim \"Importantly, it achieves this without requiring additional memory overhead compared to AdamW\", which I don't understand. To my knowledge, SF method will need three sequences x, y and z for each update iteration. You may be able to recover y by x and z, but you still have one extra sequence comparing to standard SGD or Adam, which corresponds to having another copy of the weight matrix. How come there is no additional memory overhead?\n\nIn general, I feel unconfident of the results and observations of the work to be published in NeurIPS.", "limitations": "See the questions section.\n\nNo potential negative societal impact.", "rating": 4, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 2, "clarity": 3, "significance": 1, "originality": 2, "strengths_and_weaknesses": "Strengths : The paper is clearly written, and is rich in terms of experiments and observations. A new analysis is also proposed to improve the SF trainer.\n\nWeaknesses (please response to the questions section directly): It is still unclear if the evidence is in support of a river-valley landscape, or if other landscapes could also result in the loss curves in Fig 2 and 3; The theoretical understanding is on a simple 2-dim toy objective which may not be persuasive; The experiments on the refined optimizer is questionable.", "was_revised": true, "final_justification": "The authors provides further evidences and discussions which resolve most of my concern. I'd like to increase my evaluation score.", "rebuttal": "Thank you for your efforts and valuable feedback. Let us address your concerns and questions below. We look forward to discussing further if anything remains unclear.\n\n\n### Q1 & Q2: Justification of the River-Valley Landscape\n\n> The river-valley landscape lacks mathematical justification and its connection to the experiments in Section 3 is unclear.\n\n\nWe would like to clarify that our paper does not introduce the “river-valley” landscape as a novel proposal. Rather, we build upon an emerging framework established in recent peer-reviewed literature to analyze and interpret our findings.\n\nThe river-valley model was introduced by Wen et al. (2025) and further supported by Belloni et al. (2025). It is grounded in empirical observations, including the characteristic sharp loss drop during learning rate (LR) decay and direct measurements of a convex, valley-shaped loss profile (see Fig.7 in Wen et al and Fig.2 in Belloni et al). This framework also aligns with established theoretical concepts such as the “Edge of Stability,” where iterates bounce between valley walls, and the “Central Flow,” which corresponds to the low-curvature valley floor that the optimization trajectory follows.\n\nIndependent work by Song et al. (2025) provides complementary evidence. They explicitly measured Hessian eigenvalues and eigenvectors during training to identify high-curvature (“hill”) and low-curvature (“river”) directions. They showed that projecting optimizer updates onto the river direction—by removing components along hill directions—preserves learning efficacy. Their results suggest that once the iterate reaches the valley floor, motion along hill directions is unnecessary, while progress along the river remains essential.\n\nWe use this established framework to interpret the experiments in Section 3. The reviewer’s convex quadratic intuition is consistent with this view: under the river-valley lens, the sharp loss drop after LR decay arises because the initial LR is too large for the steep \"valley walls\"; decay enables convergence to the \"valley floor\".\n\nOur key empirical finding is that SF-AdamW does not exhibit such a sharp drop upon LR decay. Within the river-valley framework, this suggests that SF iterates remain already close to the valley floor throughout training, with the SF mechanism effectively keeping the trajectory in the river.\n\nTo provide direct justification in our own experimental setting, we measured the **loss along linear interpolations between training checkpoints** under the setting of Fig.2 (Slimpajama, 0.5M batch size). We compared (1) AdamW with constant LR, (2) AdamW with linear LR decay to 0, and (3) SF-AdamW with constant LR, using checkpoints after 2B and 2.5B tokens. Results are reported in **Table 1** as below.\n\n**Table 1: Test loss evaluated at linear interploation $\\alpha w_0 + (1-\\alpha) w_1$ between two checkpoints (2B and 2.5B tokens)**\n$\\alpha$|0|0.1|0.2|0.3|0.4|0.5|0.6|0.7|0.8|0.9|1\n-|:-:|:-:|:-:|:-:|:-:|:-:|:-:|:-:|:-:|:-:|:-:\nAdamW (Constant LR)|3.322|3.295|3.274|3.259|3.248|3.242|3.240|3.243|3.250|3.261|3.278\nAdamW (Linear Decay LR)|3.322|3.297|3.275|3.255|3.239|3.224|3.212|3.204|3.198|3.195|3.195\nSF-AdamW (Constant LR)|3.263|3.253|3.245|3.237|3.231|3.226|3.222|3.220|3.219|3.218|3.219\n\nThese new results clearly demonstrate three distinct dynamics, replicating the findings of prior work and supporting our hypothesis:\n- **AdamW (constant LR)** shows a **convex, valley-shaped** loss profile between checkpoints.\n- **AdamW (linear decay to 0)** shows a **sharp, monotonic decline**, consistent with an iterate transitioning from the valley wall down to the valley floor.\n- **SF-AdamW (constant LR)** shows a **flat, slow decline**, consistent with an iterate already moving along the low-curvature valley floor.\n\nNotably, (1) and (2) replicate the loss profiles reported by Fig.7 in Wen et al. and Fig.2 in Belloni et al., while (3) further supports our claim that SF-AdamW closely tracks the river.\n\nThis combination of an established literary framework and our new, direct empirical evidence provides a solid foundation for our analysis using the river-valley landscape.\n\n**References**\n\nWen et al. (2025), *Understanding Warmup-Stable-Decay Learning Rates: A River Valley Loss Landscape View*, ICLR 2025.\n\nBelloni et al. (2025), *Universal Dynamics of Warmup Stable Decay: understanding WSD beyond Transformers*, ICML 2025 Workshop on High-dimensional Learning Dynamics.\n\nSong et al. (2025), *Does SGD really happen in tiny subspaces?*, ICLR 2025.\n\n### Q3: The Toy Objective\n\n> The observation is on a simple 2D toy objective... Can this carry over to large scale?\n\nThe toy objective in Section 4 is deterministic and simplified by design. This intentional simplification allows us to isolate and study how the optimizer interacts with the key geometric properties of the river-valley landscape without the confounding effects of stochasticity.\n\nWhat validates this approach is our finding that this minimal model reproduces surprisingly rich phenomena that are observed in deep learning experiments. For instance, the toy objective accurately captures:\n- How schedule-free dynamics track the valley floor and how this behavior depends directly on the momentum parameter.\n- The emergence of the Edge of Stability phenomenon, specifically occurring at the $y_t$ iterates, which is a highly non-trivial dynamic.\n\nThe fact that such a minimal setup can replicate these complex behaviors makes it a powerful and insightful tool for understanding the fundamental mechanisms of schedule-free optimization, confirming its relevance beyond the simplified setting.\n\n### Q4: The Refined Optimizer and Hyperparameter C\n\n> We have a new hyperparameter C to tune. It might still be sensitive to C?\n\nWe appreciate this concern and directly address it with new sensitivity experiments conducted during the rebuttal period. Building on our main experiments on Slimpajama with 0.5M and 2M batch sizes, we swept over a range of values for C and evaluated test perplexity on 9M held-out tokens. The results are shown in **Tables 2 and 3** below.\n\n**Table 2: Refined SF-AdamW, Batch Size 0.5M, $\\beta_2=0.99$, LR 2e-3, Test Perplexity**\n||Vanilla|$C=5$|10|20|50|100|200|500\n-|:-:|:-:|:-:|:-:|:-:|:-:|:-:|:-:\n$\\beta_1=0.1$|67.20|37.06|35.80|36.62|36.15|37.27 |-|-\n$\\beta_1=0.5$|41.01|-|-|29.32|30.87|29.96|29.57|-\n$\\beta_1=0.9$|27.70|-|27.70|-|23.97|24.64|24.93|25.11\n$\\beta_1=0.95$|25.12|-|-|25.12|23.98|23.60|24.37|24.83\n\n**Table 3: Refined SF-AdamW, Batch Size 2M, $\\beta_2=0.99$, LR 2e-3, Test Perplexity**\n||Vanilla|$C=5$|10|20|50|100|200|500|1000\n-|:-:|:-:|:-:|:-:|:-:|:-:|:-:|:-:|:-:\n$\\beta_1=0.1$ |110.8|54.80|47.37|44.31|43.80|-|-|-|-\n$\\beta_1=0.5$ |54.24|47.80|38.42|38.86|38.49|42.97|-|-|-\n$\\beta_1=0.9$ |31.34|-|31.34|29.86|27.68|28.16|27.02|27.88|-\n$\\beta_1=0.95$|27.29|-|-|27.29|25.77|25.45|25.77|27.31|-\n$\\beta_1=0.98$|26.09|-|-|30.23|26.09|25.49|24.51|23.95|23.88\n\nOur findings show that **refined SF-AdamW is robust to the choice of C**: it consistently outperforms vanilla SF across a broad range of values. For example, in Table 2, settings such as $C=50,100,200,500$ all improve over vanilla SF for $\\beta_1=0.9, 0.95$, and $C=50,100$ improve over vanilla SF across all momentum configurations. Similarly, in Table 3, refined SF outperforms vanilla SF over a wide range of $C$, with improvements persisting up to very large values.\n\nThese results indicate that refined SF-AdamW remains robust and effective across a wide range of C, reducing the cost of tuning this hyperparameter.\n\n\n\n\n### Q5: Memory Overhead\n\n> How can you claim no additional memory overhead compared to AdamW when SF methods require three sequences (x, y, z)?\n\nWe would like to clarify that our claim about SF-AdamW having no additional memory overhead compared to AdamW is accurate. The key point is that SF-AdamW can be interpreted as the Schedule-Free method applied over RMSProp, just as AdamW is momentum applied over RMSProp.\n\nTo be specific:\n - AdamW: Requires storing the model parameter (x_t), the 1st moment vector (m_t), and the 2nd moment vector (v_t).\n - SF-AdamW: Requires storing the model parameter (x_t), the auxillary parameter (z_t), and the 2nd moment vector (v_t). Note that y_t is computed on the fly from x_t and z_t, and does not need to be stored separately.\n\nThus, both optimizers require storage for three full copies of the model parameters, and SF-AdamW does not introduce any additional memory overhead compared to AdamW. We will add this clarification in the next revision."}, {"review_id": "CdgV9lXFZH", "reviewer": "Reviewer_JseK", "summary": "The paper analyzes Schedule-Free (SF) optimization using different theoretical and empirical frameworks, including River-Valley landscapes, Edge of Stability, central flows, and a toy model. The authors analyze two limitations of SF optimization: 1. Sensitivity to the momentum parameter $\\beta$ and (2) large batch size.\n\n**Main contributions** are summarized below:\n\n1. The authors show that SF-AdamW can effectively follow the optimal trajectory along the river direction without explicit learning rate decay or weight decay. In particular, the iterate $y_t$ faithfully follows the river direction, while $x_t$ may not follow at suboptimal $\\beta$ values.\n2. They characterize the dynamics of SF using the central flow framework. The resultant equations shed light on why suboptimal $\\beta$ values can steer training away from the river direction.\n3.  Analysis of SF on a toy model also sheds light on the effect of suboptimal $\\beta$ parameter.\n4. Derive the Edge of Stability thresholds of SF-GD and SF-AdamW.\n5. Introduce a new SF variant that aims to address momentum sensitivity and large batch performance by decoupling the momentum for $y_t$ iterate and EMA averaging of $x_t$ through an additional hyperparameter $C$.\n\nI am not very familiar with prior works on understanding SF. I will rely on other reviewers for my judgment on new contributions to the understanding of SF. \n\n**Experimental setup**: ~100 M Llama / GPT style Transformers pre-trained on SlimPajama/OpenWebText trained up to Chinchilla optimal.", "questions": "1. I think the Central Flow can be used to explain why the iterate $x_t$ deviates from the river direction at suboptimal $\\beta$. This result will complement the paper's analysis.\n2. I have a basic question about SF. If $x_t$ deviates from the river direction and $y_t$ does not, then why does it matter for the overall performance? This suggests that one can always evaluate the loss at the iterate $y_t$, and it should perform well. Have the authors tested this?\n3. What's the motivation behind the toy model? The first term can be viewed as a two-layer network with one training example. What's the intuition behind the second term? How does a term like that appear in practice?\n4. How does SF use the same memory as AdamW? It has 4 variables $x_t, y_t, z_t, v_t$. Is it because of the reformulation in Section 4.4?", "limitations": "* Experimental results are limited to two pre-training setups at ~100M scale trained up to Chinchilla optimal. Please note that I am not saying that the authors should perform extensive experiments at large scales. Rather, mentioning the setup sets the scope of the experiments during the review process.\n* (Minor) The visualizations can be improved: authors can use contrastive colors for the landscape and trajectories.\n* Please see weaknesses.", "rating": 5, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 4, "significance": 4, "originality": 4, "strengths_and_weaknesses": "**Strengths**\n\n1. Strong theoretical and empirical analysis connecting to prior work, including (1) the River-Valley Landscape, (2) Edge of Stability, (3) Central Flows.\n2. Well-motivated problem statement: improving limitations of SG optimization for a scalable alternative to conventional pretraining strategies, which use learning rate schedules.\n3. Insights into the sensitivity of the SF method to the momentum parameter and fixes.\n\nMore generally, I think the paper significantly contributes to our understanding of SF optimization. \n\n**Weaknesses**\n\n1.  I am unsure about the implications of the large batch size results. While the authors aim to address the large batch size limitation of SF, they do not perform an analysis on the effect of large batch size, and much of the paper analyzes the effect of the momentum coefficient. The only exception is Figure 8, where they show that SF-AdamW matches AdamW's performance. Furthermore, a batch size of 2M tokens is only 4 times larger compared to the default setting of 0.5M tokens.\n2. Introduces a tunable hyperparameter $C$. Its unclear how sensitive SF is wrt $C$ in different training regimes beyond pre-training. Furthermore, an analysis of the hyperparameter $C$ would be helpful.\n3. In Section 3, the authors are implicitly assuming the River Valley Landscape in SF optimization. While the weight averaging and LR decay experiments suggest the existence of the River-Valley Landscape, I believe interpolation between different iterates along the valley and river directions (as in the original River Valley Landscapes) would be helpful.", "was_revised": true, "final_justification": "The authors' rebuttal has addressed my major concerns. While a few concerns regarding the iterates and their role in dynamics remain, I hope these will motivate future works. As my original rating was already high, I will be maintaining my score.", "rebuttal": "Thank you for your valuable feedback! We appreciate that you found our work **significantly contributes to our understanding of SF optimization**. We would like to address your questions as below.\n\n### W1: Implications of the Large Batch Results\n> While the authors aim to address the large batch size limitation of SF, they do not perform an analysis on the effect of large batch size.\n\nWe acknowledge that while our paper aims to address the large-batch limitation of SF, our experimental focus is primarily on analyzing the effect of momentum rather than batch size. This choice was motivated by the fact that the role of momentum in SF methods remains relatively underexplored, and our goal was to address this gap. In contrast, the impact of batch size on SF performance has already been studied—most notably by [1], who conducted scaling experiments and identified the critical batch size at which vanilla SF methods begin to degrade.\n\nRather than replicating those experiments, we sought to build on their findings by providing intuition for **why vanilla SF fails in large-batch regimes.** Our analysis suggests that the failure stems from the entanglement between momentum and the weight averaging window, which motivated our refinement to decouple them.\n\nWe agree with the reviewer that our current large-batch results for the refined SF method are preliminary. Due to computational constraints, and because proposing the refined method is not the main focus of our paper, we did not conduct a full batch size ablation. We will revise the text to avoid overclaiming the effectiveness of the refined method in large-batch settings.\n\n[1] Zhang et al., *How Does Critical Batch Size Scale in Pre-training?*, ICLR 2025.\n\n### W2: Analysis of HP C\n> Introduces a tunable HP C. Its unclear how sensitive SF is wrt C in different training regimes beyond pre-training.\n\nWe agree that introducing a new HP entails tuning cost. To assess this, we conducted an additional sweep over C values, reported in  **Tables 2 and 3 of our response to Reviewer kaR1, Q4**. We observe that performance remains stable across a wide range of C.\n\n### W3: Assuming River Valley Landscape\n> I believe interpolation between different iterates along the valley and river directions would be helpful.\n\nWe appreciate the suggestion and would like to clarify our position. Our work builds upon the river-valley landscape as a modeling framework that has been established in recent literature. This assumption is not specific to any particular optimizer, as the river-valley refers to the geometry of the loss landscape itself, rather than the dynamics of a specific optimization method.\n\nInterpolation-based visualizations have already been presented in prior works—for example, Fig. 7 in [2] and Fig.2 in [3]—demonstrating that modern deep learning loss landscapes often exhibit a convex, valley-shaped structure. In our own setup, we independently verified this landscape structure through interpolation between iterates and observed the same trend (see **Table 1 of our response to Reviewer kaR1, Q1&Q2**).\n\nIn addition, we would like to emphasize that the river-valley landscape is further supported by several well-established empirical signatures:\n- The sharp loss drop during LR decay.\n- The EoS phenomenon, where iterates bounce between the steep walls of the valley.\n- The Central Flow, which captures the trajectory along the valley floor.\n\nFurther supporting evidence comes from [4], who show that optimization updates projected onto the river direction (i.e., high-curvature hill directions removed) preserve learning efficacy. This implies that motion in the hill directions is often unnecessary once the iterates are near the valley floor.\n\nWe will revise the text to include these clarifications and highlight that we also validated the river-valley structure in our setting.\n\n[2] Wen et al., *Understanding Warmup-Stable-Decay Learning Rates: A River Valley Loss Landscape View*, ICLR 2025.\n\n[3] Belloni et al., *Universal Dynamics of Warmup Stable Decay: understanding WSD beyond Transformers*, ICML 2025 Workshop on HiLD.\n\n[4] Song et al., *Does SGD really happen in tiny subspaces?*, ICLR 2025.\n\n### Q1: Is Central Flow Analysis of x_t Possible?\n> I think the Central Flow can be used to explain why the iterate x_t deviates from the river direction at suboptimal β.\n\nWe believe that applying central flow analysis to the x iterate is not appropriate. Central flow is intended to model the time-averaged optimization trajectory at the Edge of Stability. Our observations indicate that only the y iterates operate at the EoS, exhibiting oscillations centered around the river. Therefore, we model the central flow based on the y iterates, which allows us to capture the oscillation variance in the hill direction away from the valley floor.\n\nIn contrast, the x iterates do not oscillate around the river and do not operate at the EoS. Consequently, central flow analysis does not provide a suitable explanation for the deviation of x_t from the river at suboptimal β. We appreciate the reviewer’s insight and would welcome suggestions on alternative frameworks for modeling this behavior.\n\n### Q2: x_t vs y_t in SF method\n> If x_t deviates from the river direction and y_t does not ... one can always evaluate the loss at the iterate y_t, and it should perform well. \n\nIt is not the case that the y iterates always outperform the x iterates. For optimal values of momentum (e.g., β = 0.95), we observe that x_t achieves better performance than y_t. This is because x_t corresponds to an \"appropriate\" weighted average of the y iterates, which can enhance performance when the averaging is well-balanced.\n\nIn contrast, when β is small, the weighting scheme places too much emphasis on early y iterates, leading to degraded performance of x_t. This issue motivated our introduction of the refined SF method, which decouples the averaging mechanism from the momentum parameter β. This decoupling allows for improved averaging regardless of the momentum setting. As shown in Figure 8 (left), the refined method enables x_t to consistently outperform y_t, even at smaller β values.\n\nReturning to the original question: no, evaluating the loss at y_t alone does not always yield optimal performance. A well-designed weighted average over the y iterates is often necessary for achieving the best results.\n\n### Q3: The Toy Objective\n> What's the motivation behind the toy model?\n\nOur design of the toy model is intended to construct a minimal model that captures the river-valley geometry while replicating key behaviors observed in language model training.\n\nThe first term, $(w_1w_2-1)^2/2$, can indeed be viewed as the loss of a two-layer network on a single example. More importantly, it defines a river-shaped loss landscape, where the constraint $w_1w_2 = 1$ defines the valley floor, and the curvature in the hill direction becomes sharper as $w_1$ increases. This structure is central to capturing the loss geometry observed in modern training settings.\n\nThe second term, $\\log(1+\\exp(-w_1))$, is included to create a directional bias along the valley floor so that the optimizer progresses in a consistent direction. The specific form of this term is not critical to the main phenomena we study and could be replaced by other forms (e.g., a linear term like $-w_1$) without qualitatively changing the behavior. Its role is simply to ensure that the loss decreases along the river as training progresses.\n\nWhile the exact form of the toy objective may not correspond to a practical loss function, that is not its purpose. Instead, it serves as a minimal model that exhibits rich dynamics consistent with those observed in realistic settings. Notably, it replicates:\n- How SF dynamics remain near the valley floor and how this behavior depends on the momentum parameter β.\n- The emergence of the EoS, specifically at the y iterates, which mirrors our findings in large-scale experiments.\n\nThe fact that this toy model reproduces such nontrivial behaviors confirms its utility as a theoretical tool. It provides insight into the mechanisms behind SF optimization that would be difficult to extract from large-scale models alone.\n\n### Q4: Memory Overload in SF-AdamW\n> How does SF use the same memory as AdamW? It has 4 variables x, y, z, v. Is it because of the reformulation in Section 4.4?\n\nSF-AdamW does not require storing all four variables explicitly. In particular, y is computed on the fly from x and z, and therefore does not need to be stored as a separate tensor. This implementation detail is not specific to our reformulation and is already described in the original SF method (see Section 4.4 in Defazio et al.).\n\nTo clarify:\n- AdamW stores the model parameters x_t, the first-moment vector m_t, and the second-moment vector v_t.\n- SF-AdamW stores the model parameters x_t, the auxiliary parameters z_t, and the second-moment vector v_t.\n\nSF-AdamW can be interpreted as SF method applied over RMSProp, just as AdamW can be viewed as momentum applied over RMSProp. In both cases, the optimizer maintains three parameter-sized tensors, and SF-AdamW does not incur any additional memory overhead compared to AdamW.\n\n### Limitations\n> Experimental results are limited to two pre-training setups at ~100M scale trained up to Chinchilla optimal.\n\nWe agree and will add a \"Limitation\" section to acknowledge this. Our experiment scale was intentionally chosen to make the training dynamics of SF methods more interpretable within our compute constraints, as our primary goal is to understand the underlying mechanisms of SF.\n\nWe acknowledge that validating these findings at larger scales and across more diverse architectures would further strengthen their generality. We view this as an important direction for future work, though it is beyond the scope of the current paper.\n\n> The visualizations can be improved\n\nThank you for the suggestion. We will revise the figures for better readability."}, {"review_id": "omhZYWyhwd", "reviewer": "Reviewer_ens5", "summary": "This paper analyzes the schedule-free-optimizer (SFO) from DeFazio et al. (2024) through the perspective of the \"river-valley loss landscape.\"  Unlike with AdamW, performing a short decay phase on the SFO weights (from a well-tuned run) does not yield loss improvements, suggesting that a well-tuned SFO model is already \"close to the river.\"  Digging into the behaviour of SFO, they find that the simple weight average x_t lags behind the main y_t iterates when the SFO momentum is suboptimal.  They reformulate the SFO derivation in order to show that the \\beta momentum term is effectively doing double-duty: it controls the effective smoothing across gradients, while also controlling the shape of how prior outputs are averaged.  By decoupling these two effects via the introduction of an additional hyperparameter, some of the sensitivity of SFO to its hyperparameters appear to be ameliorated, including when the batch size is increased by a factor of 4.", "questions": "See weaknesses above.", "limitations": "I feel the paper could have benefited from a detailed limitations section, summarizing the \"simplifying assumptions\" that were made in the theory, and noting the specific items that were not explored in the paper, but that might have been given more compute:\n- other optimizers\n- larger-scale models\n(both as noted in the conclusion)\n\nBut also, maybe:\n- testing the scaling of the C hyperparameter\n- testing a range of different batch sizes\n- testing other vocabulary sizes, model architectures, datasets, etc.\n\nYou might also acknowledge the disadvantages of introducing another hyperparameter, and how these might be mitigated.\n\nIn short, I'd be curious what the *authors* think that people training large-scale models should test before adopting the refined optimizer proposed in this paper.", "rating": 4, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 3, "significance": 2, "originality": 2, "strengths_and_weaknesses": "## Strengths\n\nIdentifying the key issues with SFO and how these can be mitigated is valuable and well-motivated.\n\nFor the most part, the work carefully stepped through the theory, providing accompanying empirical results at each step, in a manner that was both rigorous and instructive.    \n\nBoth the figures and the writing in each section were fairly clear and easy-to-digest.\n\n## Weaknesses\n\nThe main issue is that the work seems to be a bit incomplete and preliminary, and it felt like the later sections were hastily written and a bit disorganized.\n\nThe paper essentially culminates in a proposal for a refined version of SFO.  But the section evaluating this refined version doesn't really provide enough details for the reader to understand whether this refined version is worth adopting.  The main disadvantage of introducing a new hyperparameter is, of course, that this introduces new tuning costs at small scale, and questions of HP transfer across scales for training large models.  If we view the paper as, rather than providing actionable insights and new tools for practitioners, just hypothesis-driven empirical insights, then I would have liked to see more experiments that evaluate the refined method across further settings.  E.g., does the optimal C change *systematically* with batch size?  Is having x_t track y_t *sufficient* for good models, or only *necessary*?\n- e.g., it’s a little hard to compare the results in Figure 8 since the y-axis has a different scale in Left vs. Middle, but I'm left wondering if C really does \"fix\" SF-AdamW?  Yes, x_t catches up to y_t, but if y_t is still below what could be achieved with an optimal β1 (or with a simple Cosine-10%-decay schedule), then SF-AdamW is apparently still sensitive to these HP values.  What should the reader take away here, and what else is worth checking?\n\n### Other empirical concerns:\n\n“We perform sweeps over the learning rate, momentum parameters, and the refinement parameter C (for refined SF-AdamW).”\n- But how are you evaluating the results across the sweep?  On the validation set?  It’s totally expected that if you introduce a new hyperparameter and you then “fit” this hyperparameter on the validation set, you can do better than if you didn’t introduce this hyperparameter.  At the very least, we should evaluate on a held-out set.  But even beyond that, we should evaluate that a fit at one scale provides a good model at a larger scale, since we can’t actually afford to fit C at every scale.\n\nIn the appendix, line 796: “For large-batch runs with cosine decay, the learning rate is annealed to 10% of its peak.”\n- I’m a little concerned that this Cosine-10x decay is a bit of a “straw man”.  E.g., Bergsma 2502.15938 showed that decaying to 0 was superior to 10% decay across a range of pre-training runs --- and gains increase at higher tokens-per-parameter.  Given that you trained to ~48TPP, I would expect decaying to zero to be superior.  DeFazio used Linear as the baseline, also see their section B.1.  Moreover, to be fair to Cosine, you probably also need to adjust some of its HPs as you scale the batch size.\n\n### Regarding the paper's main contributions:\n\n- \"We begin by revisiting two widely used strategies\" -- is this really a main contribution?  \n- \"Reveal\" that SF performs a form of weight averaging -- I think you need to be a lot more precise here because we know that SFO performs a form of weight averaging -- that's exactly what it says on the box!  Maybe you can specify exactly what you show that hasn't been shown previously.  It would have also been helpful if you could make clear exactly what Morwani et al showed versus what this paper shows -- I wasn't familiar with that work, but your paper hints that they touch on a very similar topic.\n\n### Regarding prior work\n\nFor a full description of the advantages and disadvantages of prior work, you might note that weight averaging does not necessarily need to consume extra memory. E.g., in DeepSeek-V3:\n\n> Exponential Moving Average in CPU. During training, we preserve the Exponential Moving Average (EMA) of the model parameters for early estimation of the model performance after learning rate decay. The EMA parameters are stored in CPU memory and are updated asynchronously after each training step. This method allows us to maintain EMA parameters without incurring additional memory or time overhead.\n\n### Paper organization\n\n- Section 2.2 directly re-states things about WSD that were mentioned in 2.1.\n- I felt like I had to flip between Section 4.4 and Section 5 -- I'm not sure, but shouldn't the unnumbered equation for the alphas just go back with the unnumbered equation for the x_t iterate you have above on the page?\n- Line 262: Please define P when you introduce it - or maybe introduce this term after Line 271?\n\n### Writing\n\nSection 5: I felt like this section was more confusing than it should be.  Like, could we start with, “if large β has all these benefits in terms of keeping us in the valley, why not just use a large β all the time”?  I don't have specific suggestions, I just found this hard to follow.  Even the point that “large β slows yt updates, as both βct+1 and (1 − β) in (SFy) become vanishingly small when t is large\" – isn’t the point that (1-β) becomes small and we pay less attention to recent gradients, but if batch sizes are larger, we actually want to pay more attention to recent gradients? – I.e., the scale of DATA that we pay attention to should be similar.  I don't know why it matters that later gradients count less at higher steps... or, if that *is* important (which it might be), then this sentence is carrying way too much weight to explain all of this.\n\n### Minor\n\n- WSD “allows dynamic adjustment of HPs across phases”  But doesn’t Cosine as well?  Which HPs?  E.g., it's not \"allowed\" with a Cosine schedule to change the batch size or LR halfway through, or not common?  I’m not following.\n\nLine 37 -- when you mention the \"workaround\" of weight averaging, it would be good to have some citations here or at least a forward pointer to your related work (so we know what you have in mind exactly by weight averaging when you introduce it).\n\n124M model trained on 6B tokens, so that's 48.4 TPP.  You say \"compute budget determined by Chinchilla,\" but isn't the rule-of-thumb for compute-optimal training 20 TPP?\n\nTypos:\nLine 263: “We first the stability”\nLine 276: “with with”", "was_revised": true, "final_justification": "The authors have addressed some of my concerns regarding quality (experimental rigor) and clarity, so I have revised my score.  I'm still not sure about the overall significance.  But overall, this is exactly a \"technically solid paper where reasons to accept outweigh reasons to reject.\"", "rebuttal": "We appreciate your thoughtful feedback. We believe the core of the concern arises from a slight misalignment between the primary focus of our paper and the review’s emphasis on Section 5, which treats the work as if its main goal were to propose a new optimizer.\n\nAs reflected in our title, *“Through the River: Understanding the Benefit …”*, our central contribution is a principled analysis of SF dynamics, grounded in the river-valley loss landscape framework. This focus spans the majority of the paper (Sections 2–4), where we present both theoretical and empirical insights into the behavior, strengths, and limitations of SF methods. The refined method introduced in Section 5 is intended as a concrete illustration of this framework—a proof-of-concept that addresses a diagnosed failure mode—rather than the core contribution of the work.\n\nWe acknowledge that the refinement could appear preliminary if evaluated in isolation as a standalone optimizer. However, its purpose is to demonstrate how insights from our analysis can inform practical improvements. With this context in mind, we respond to the reviewer’s specific points below.\n\n### Contribution and Evaluation of Refined SF method\n\n> the section evaluating this refined version doesn't really provide enough details\n\nThe refined SF method was introduced to empirically test a key prediction from our analysis: that the coupling between averaging and momentum in the vanilla SF update leads to degraded performance of x_t relative to y_t under small β1 (Figure 5). By introducing the decoupling parameter C, we isolate these effects, and the observed improvement of x_t over y_t (Figure 8) provides empirical support for our hypothesis. While the primary goal of the refined method is diagnostic, these results show how our analysis informs a practical path to improving SF in regimes where they currently struggle.\n\n> The main disadvantage of introducing a new HP is ... tuning costs\n\nTo address this, we conducted sensitivity analysis of C, reported in Tables 2 and 3 of our response to Reviewer kaR1, Q4. The results show that  **refined SF-AdamW is robust to the choice of C**: it consistently outperforms vanilla SF across a broad range of values, reducing the practical cost of tuning.\n\n> Does the optimal C change systematically with batch size? Is having x_t track y_t sufficient for good models, or only necessary?\n\nA systematic batch size ablation is an important future direction. Our hypothesis is that larger batch sizes, which reduce gradient noise, would favor larger optimal C values that prioritize recent iterates. However, we refrain from making strong claims without direct evidence. Due to computational constraints—and since the refined method is not the paper’s main focus—we did not perform these scaling experiments. We believe it would be valuable to conduct a study akin to Fig. 2 in [1], measuring the critical batch size of the refined SF method.\n\nHaving x_t track or outperform y_t is a necessary but not sufficient condition for good performance. A good averaging scheme may improve x_t, but the quality of the underlying momentum iterates y_t is also important.\n\n> It’s a little hard to compare the results in Fig.8. I'm left wondering if C really does \"fix\" SF-AdamW?\n\nOur intention is not to claim that C universally “fixes” SF-AdamW, but rather that it mitigates a specific failure mode where x_t lags behind y_t due to entangled averaging. The differing y-axis scales were chosen to highlight the x_t vs. y_t gap within each subplot, not to compare across momentum values.\n\n### Other empirical concerns\n\n> But how are you evaluating the results across the sweep? On the validation set?\n\nYes, our evaluation was conducted on the validation set, which is standard in LM pretraining where, under the single-pass setup, prior work has shown there is no generalization gap (see [2]).\n\nTo rigorously address this concern, we evaluated on a held-out test set (9M-token subset of SlimPajama) with HP sweep of refined method (see **Tables 2 and 3**). The performance trends on the test set closely match those on the validation set, confirming that our conclusions are not due to overfitting.\n\n> Cosine-10x decay is a bit of a “straw man”\n\nWe conducted an HP sweep for AdamW with cosine decay to 0, evaluated on a held-out test set:\n\n**Table 4: AdamW with Cosine Decay to 0 (Batch Size 2M, Test Perplexity)**\n||max LR=5e-4|1e-3|2e-3|5e-3\n:-:|:-:|:-:|:-:|:-:\n$(\\beta_1,\\beta_2)=(0.9, 0.95)$|27.26|24.29|23.21|23.58\n$(\\beta_1,\\beta_2)=(0.95, 0.99)$|27.26|24.44|23.42|24.05\n\nThese results show that well-tuned AdamW with cosine decay to 0 slightly outperforms refined SF-AdamW (Table 3). Nonetheless, our refinement significantly narrows the gap between vanilla SF-AdamW and AdamW with optimal cosine decay in the large-batch regime. We appreciate your insightful comment and will clearly discuss this result in the revised manuscript.\n\n### Regarding the paper's main contributions\n> \"We begin by revisiting two widely used strategies\" – is this really a main contribution?\n\nSection 2 is indeed a review, not a novel contribution. Its purpose is to contextualize existing strategies through the lens of scalable training and river‑valley framework, setting the stage for subsequent analysis. We will revise main contributions list to reflect this.\n\n> \"Reveal\" that SF performs a form of weight averaging – you need to be a lot more precise here\n\nThe original SF method [3] indeed frames x_t as a uniform average of the z iterates. Our contribution is to show an alternative and more informative reformulation: x_t can be viewed as a weighted average of y iterates, where y follows a momentum‑like update.\n\nThis perspective aligns with [4], who interpret SF‑SGD as a weighted average of accelerated SGD, but our work extends it in several ways:\n1. We analyze SF-AdamW, which has not been studied in this context.\n2. We connect the reformulation to the river-valley perspective, showing that y tracks the valley floor while x acts as a smoothed average.\n3. We observe that EoS arises specifically at y iterates, not at x or z.\n\nThese observations support our proposal that y should be viewed as the primary optimization path, with x serving as its weighted average. We will revise the text to clearly distinguish our contributions from prior work.\n\n### Regarding prior work\n\n> weight averaging does not necessarily need to consume extra memory, e.g., in DeepSeek-V3\n\nWe clarify as follows:\n- Even if DeepSeek‑V3’s implementation is fully optimized, tracking EWA requires extra CPU memory, so it is not strictly correct to say it “consumes no extra memory” unless the claim is limited to GPU memory.\n- While techniques exist to reduce GPU memory use, in practice they demand non‑trivial engineering efforts, depend on training setup, and still incur overhead from CPU allocations, extra compute, and CPU–GPU synchronization. The degree of overhead also varies with hardware.\n\nTherefore, we believe our original statement that weight averaging often incurs non-negligible practical overhead is a fair characterization of the trade-offs involved.\n\n### Paper Organization\n\nThank you for the suggestions. We will revise the paper accordingly. \n\n### Writing\n> Section 5: “if large β has all these benefits in terms of keeping us in the valley, why not just use a large β all the time”?\n\nOur analysis shows that larger β reduces oscillation variance, keeping y_t closer to the river floor, but does not imply that larger β always improves performance.\n- For y_t, performance depends on both oscillation magnitude and effective \"speed\" along the river, requiring a balance when tuning β.\n- For x_t, β also controls averaging weights, and the optimal weighting is task‑dependent.\n\nThus, β must be tuned to jointly account for its dual effects on x_t and y_t.\n\n> “large β slows yt updates,...\" – isn’t the point that (1-β) becomes small and we pay less attention to recent gradients, but if batch sizes are larger, we actually want to pay more attention to recent gradients?\n\nOur original argument highlighted that as $β\\approx 1$ and $t\\to\\infty$, both $β c_{t+1}$ and $(1-β)$ in (SFy) become small, slowing updates of y_t. We also agree with your perspective: large β places less emphasis on recent gradients, potentially degrading y_t performance in large-batch regimes. Conversely, a larger β is beneficial for x_t due to its narrower averaging window. We will clarify this trade-off in the revised manuscript, incorporating your insights.\n\n### Minor\n> WSD “allows dynamic adjustment of HPs across phases” But doesn’t Cosine as well? \n\nWSD decouples the constant-LR \"trunk\" from temporary \"decay\" branches, enabling one to experiment with phase-specific HPs (e.g., decay shape) on new branches without altering the main run. In contrast, a monolithic cosine schedule binds LR to the entire training horizon, so testing a different decay profile requires a full restart.\n\n> Line 37-when you mention the \"workaround\" of weight averaging ... forward pointer to your related work\n\nWe will add a pointer to Section 2.3.\n\n> 124M model trained on 6B tokens, so that's 48.4 TPP\n\nOur pre-training runs were set to match the Chinchilla-optimal compute scale. For our main experiments, 124M model was trained for 5k iterations on 2.5B tokens, consistent with this scale (see Appendix B).\n\n> Typos: Line 263, Line 276\n\nThank you for catching these. They will be corrected.\n\n### Limitations\n> I feel the paper could have benefited from a detailed limitations section\n\nWe will add a \"Limitations\" section outlining theoretical assumptions and the scope of our experiments.\n\n---\n[1] Zhang et al., *How Does Critical Batch Size Scale in Pre-training?*, ICLR 2025.\n\n[2] D’Angelo et al., *Why Do We Need Weight Decay in Modern Deep Learning?*, NeurIPS 2024.\n\n[3] Defazio et al., *The road less scheduled*, NeurIPS  2024.\n\n[4] Morwani et al., *Connections between schedule-free optimizers, ademamix, and accelerated sgd variants*, 2025."}]

Author rebuttal:

Meta-review: The paper analyzes Schedule-Free (SF) optimization (SF-AdamW), demonstrating that it implicitly performs weight averaging and follows a “river-valley” loss landscape without explicit decay schedules or extra memory. The authors identify momentum sensitivity and limited large-batch robustness as key weaknesses and propose a refined SF variant with an extra hyperparameter to address these issues. Reviewers find the work technically solid, clearly written, and insightful, particularly in its theoretical framing and analysis of SF dynamics. However, concerns remain about the small experimental scale, lack of statistical rigor, unclear formalization of the “river-valley” concept, and the practicality of the refined optimizer given its new hyperparameter. Overall, reviewers lean towards acceptance, recognizing the contribution as important for understanding and potentially scaling SF methods, although it is still preliminary in evaluation.
Decision: Accept (poster)