Understanding Warmup-Stable-Decay Learning Rates: A River Valley Loss Landscape View
Authors: Kaiyue Wen
Zhiyuan Li
Jason S. Wang
David Leo Wright Hall
Percy Liang
Tengyu Ma
Venue: ICLR 2025 Poster (2025)
Keywords: pretraining
language model
learning rate
loss landscape
manifold
Primary area: optimization
TL;DR: We develop a river valley landscape analogy to pretraining loss landscape to explain the nontraditional loss curve of WSD learning rate and develop a variant of WSD based on the theory.

Abstract: Training language models currently requires pre-determining a fixed compute budget because the typical cosine learning rate schedule depends on the total number of steps. In contrast, the Warmup-Stable-Decay (WSD) schedule uses a constant learning rate to produce a main branch of iterates that can in principle continue indefinitely without a pre-specified compute budget. Then, given any compute budget, one can branch out from the main branch at a proper time with a rapidly decaying learning rate to produce a strong model. Empirically, WSD generates an intriguing, non-traditional loss curve: the loss remains elevated during the stable phase but sharply declines during the decay phase. Towards explaining this phenomenon, we conjecture that pretraining loss exhibits a river valley landscape, which resembles a deep valley with a river at its bottom. Under this assumption, we show that during the stable phase, the iterate undergoes large oscillations due to the high learning rate, yet it progresses swiftly along the river. During the decay phase, the rapidly dropping learning rate minimizes the iterate’s oscillations, moving it closer to the river and revealing true optimization progress. Therefore, the sustained high learning rate phase and fast decaying phase are responsible for progress in the river and the mountain directions, respectively, and are both critical. Our analysis predicts phenomenons consistent with empirical observations and shows that this landscape can naturally emerge from pretraining on a simple bi-gram dataset. Inspired by the theory, we introduce WSD-S, a variant of WSD that reuses previous checkpoints’ decay phases and keeps only one main branch, where we resume from a decayed checkpoint. WSD-S empirically outperforms WSD and Cyclic-Cosine in obtaining multiple pretrained language model checkpoints across various compute budgets in a single run for parameters scaling from 0.1B to 1.2B.

Published as a conference paper at ICLR 2025
UNDERSTANDING WARMUP-STABLE-DECAN LEARN- ING RATES: A RIVER VALLEY LOss LANDSCAPE VIEW
Kaiyue Wen' Zhiyuan Li
Jason Wang'
David Hall'
Percy Liang'
Tengyu Ma
1 Stanford University 2 Toyota Technological Institute at Chicago
kaiyuew@stanfowe edu dlwhêcs stanford.
zhiyuanli@ttic edu jsywangestanford
plianggcs. stanford. tengyumačcs Songyumalss.staaffrd.e.. stanford.edu
ABSTRACT
Training language models currently requires pre-determining a fixed compute bud- get because the typical cosine learning rate schedule depends on the total number of steps. In contrast, the "armuu--Stalle -eelyy (WSD) schedule uses a constant learning rate to produce main branch of iterates that can in principle continue indefinitely without a pre-specified compute budget. Then, given any compute budget, one branch out from the main branch proper time with rapidly decaying learning rate to produce a strong model. Empirically, WSD generates intriguing, non-traditional loss curve: the loss remains elevated during the stable phase but sharply declines during the decay phase. Towards explaining this phenomenon, we conjecture that pretraining loss exhibits a river valley landscape, which resembles deep valley with a river at its bottom. Under this assumption, we show that during the stable phase, the iterate undergoes large oscillations due to the high learning rate, yet it progresses swiftly along the river. During the decay phase, the rapidly dropping learning rate minimizes the iterate's oscillations, moving it closer to the river and revealing true optimization progress. Therefore, the sustained high learning rate phase and fast decaying phase are responsible for progress in the river and the mountain directions, respectively, and are both critical. Our analysis predicts phenomenons consistent with empirical observations and shows that this landscape can naturally emerge from pretraining bi-gram dataset. Inspired by the theory, we introduce WSD-S, a variant of WSD that reuses previous checkpoints' decay phases and keeps only one main branch. where we resume from a decayed checkpoint. WSD-S empirically outperforms WSD and Cyclic- Cyclic-Cosine Cosine in obtaining multiple pretrained language model checkpoints across various compute budgets in single run for parameters scaling from O. to 1.2B.
simple
INTRODUCTION
Pre-training large language models (LLMs) typically involves following learning rate schedule that decreases over pre-determined number of steps, such as cosine schedule (Loshchilov & Hutter 2017 Touvron et al.) 2023), where the learning rate starts high and gradually decreases a smooth curve following the shape of cosine function. This inflexible approach makes it difficult to adapt to additional compute or data, as the learning rate schedule for all the data is not a natural continuation of the schedule used with past data. Additionally, fitting scaling laws is costly because each compute budget requires a retraining to adjust the learning rate schedule (Hoffmann et al. 2022).
In contrast the cosine learning rate, recent work Hu al. 2024) introduces the warmup-stable- decay (WSD) schedule, which does not require committing to a pre-specified total compute budget. After a standard warm-up period, the WSD schedule maintains a main "branch" using a constant learning rate indefinitely and branches off using a fast-decaying learning rate schedule to obtain intermediate checkpoints (see the second row of Figure 2b). Using the WSD schedule, one can continue training from checkpoint in the main branch by resuming with the same constant learning rate and can obtain training losses for multiple compute budgets with a single run.
Empirically, the WSD schedule produces a non-traditional loss curve (see Figure 1): during the constant learning rate phase, the loss remains higher than the loss using other schedules like the cosine schedule; but during the decay phase, it drops sharply, often leading to better final performance compared to the cosine schedule. This raises the main question the paper aims to address:

---

Published as a conference paper at ICLR 2025
WSD-5
peak
Stable
min
WSD
peak
W
Learning Rate
min
D Decay
Cosine
peak
min
Cyclic Cosine
peak
min
Init
Intermed late Checkpoint
Intermediate Checkpoint
Intermediate Checkpoint
(a) River Valley Landscape
(b) Learning rate schedules
Figure 2: We demonstrate river valley loss landscape Figure 2a explain effectiveness of Warmup- Deble (WSD) schedulo(demonnstratdd ). The stable phase adopts a large learning rate and the iterate will progress along the river while oscillating between the sharp hillsides. Due to the large oscillation caused by the large learning rate, the run will potentially show higher loss compared dropped rapidly to the oscillation of the river, revealing the optimization progress. Based on our theory, we propose WSDD-Simlliiied (WSD-S)) an effective simplification of the WSD schedule in continual learning, where we start directly using high learning rate from previous intermediate checkpoints. We visualize the learning rate schedule in Figure 2b The arrow in the second of Figure 2b indicates WSD reinitializes the checkpoint from last checkpoint the constant learning rate phase instead.
using smaller learning
in this phase. During the decay phase, the learning iterates, driving it closer
Why does WSD work, especially with such a non- traditional loss curve? Specifically, why does constant learning rate phase, characterized by slow loss improve- ments, eventually lead to superior performance?
2.3
2.2
Decay
Loss
The first contribution of this paper is theoretical frame- work to explain the underlying mechanism of WSD. We characterize type of loss landscape, called the river val- ley landscape (Definition [.I]. and theoretically show that WSD has superior performance on such loss landscapes. show that the river valley landscape can provide mul- tiple theoretical predictions matching the empirical obser- vations and hence can serve as a useful conceptual picture for understanding the pretraining optimization process.
WSD
0 10000 20000 30000 40000 50000
Step
Figure 1: The Non-traditional Loss Curve produced by WSD. A constant learning rate phase, characterized by slow loss improve- ments, eventually leads to better validation loss after learning rate decay.
As the name suggests, river valley landscape intuitively features steeply sloping hillsides with river winding through the bottom of the gorge (see Figure 2a Dur-
ing the stochastic gradient-based optimization process, the iterate bounces between the hillsides as it slowly and implicitly progresses along the river direction. The loss in this landscape can be decomposed into two components: the river component, which represents the primary loss along the river at the bottom of the hills, and the hill component, which accounts for the additional loss height from the river's course. Progress is determined primarily by the
caused by deviations river component in the long run. We demonstrate that when the loss function exhibits this type of landscape, a learning rate schedule should satisfy the following two key properties effectively minimize the loss.
Sustained high learning rate. is advantageous maintain large learning rate for as long as possible during training, even at the cost of less reduction in the loss. A large learning rate yields larger bouncing due to the stochasticity of the gradient, increasing the hill component of the loss, but it also makes faster progress in the river direction. In contrast, a small learning rate results in less bouncing, keeping the iterate close to the river, but progress along the river direction is slower. Therefore, a larger learning rate leads to faster fundamental progress in minimizing the river component, which is obscured by the oscillation in the hill component. This progress will be revealed by the decay phase discussed below.
2

---

Published as a conference paper at ICLR 2025
2. Final low learning rate. As training nears completion, it becomes essential to reduce the learning rate. This decay minimizes the oscillations in the mountain direction to decrease the hill component and ensures that the iterates converge to point close to the river, which has lower loss than any nearby points up the hills.
In Section B] we provide formal theoretical statements analyzing the trajectories of (stochastic) gradient descent on the river valley landscape, fleshing out the intuitions above. Among our synthetic and real-world studies supporting the river valley landscape hypothesis, intriguing observation in language model pretraining is that the loss on the linear interpolation of two checkpoints in the stable phase exhibits a convex and unimodal shape, resembling valley, whereas between two checkpoints in the decay phase, the loss shows a smooth monotone decay.
All the theoretical results above assume river valley landscape. How likely does the next-token prediction loss follow this pattern, and why? We hypothesize the river valley landscape can naturally arise from the heterogeneity in the stochasticity of different tokens: highly deterministic tokens (which often involve facts and knowledge) contribute to the "river" direction, while uncertain tokens (which often involve flexibility and ambiguity in the language) create the steep hillsides. We demonstrate this insight by showing in Section 4 that under bigram toy model, indeed the loss has river valley landscape, and empirically, most properties of the loss curves under various learning rates on the real datasets are still seen in this toy model. We further show that the stable learning rate phase learns the deterministic tokens, whereas the decay phase learns better the stochastic tokens.
Finally, motivated by the theoretical insights, we propose simplification and improved version of WSD, called WSD-S in continual learning. In WSD, after obtaining an intermediate checkpoint, the model and optimizer are rolled back to the end of the stable phase before continuing with constant learning rate. However, our theory predicts that the decay phase also makes progress along the river direction and thus there is no reason to discard that part of the progress. Concretely, WSD-S immediately continues training from the intermediate checkpoint with a high constant learning rate, instead of rolling the model back to a checkpoint before decaying.
We evaluate the effectiveness of WSD-S with extensive experiments on LLMs from O.IB to 1.2B parameters in a continual learning setting with 50B, 100B, and 200B tokens as the three target compute budgets. We empirically show that WSD-S has performance comparable with independent oracle runs with cosine learning rate schedules optimally tuned for each of the three budgets. Furthermore, WSD-S leads to a better validation loss than WSD) under the same compute budgets due to the re-use of the decay period. We also show through ablation studies that the performance is relatively insensitive to the precise fraction of time spent decaying as long it is near 10% and the decay does not start shortly after a coincidental loss spike.
targe
RELATED W WORK
We discuss related work in two main areas: learning rate schedules and theoretical understandings of the loss landscape. We defer the detailed discussion to Appendix A
Learning Rate Schedules. Prior research has explored various choices of learning rate sched-
ules (Smith 2017: Loshchilov & Hutter (017)) Recent studies have focused on optimizing these
schedules for language model pretraining (Hu et al. 2024 Raffel et al. 2023 Defazio et al.) 2023)
Theoretical Understanding on Loss Landscape. A substantial body of research seeks to elucidate the properties of the loss landscape in deep learning. The closest ones related to our work include the impact of gradient noise and curvature (Zhang et al. 2022aa: Pan & Li 2023), the benefits of large learning rates for finding flatter minima (Kong & Tao 2020 Wang et al. 2022). and the interplay between loss landscape geometry and feature learning dynamics (Nakkiran et al. 2019 Rosenteld & Risteski (023). Among these works, (Xing et al. (2018) 8) has presented a similar conceptual picture with us, arguing that SGD locally bounces around the valley on top of valley floor. The iterates will explore the uneven valley floor to find a more generalizable solution. In contrast, we focus on the optimization perspective and assume the existence of the river at the bottom of the hillsides, where the loss monotonously decreases. We build formal theoretical framework on top of this picture, leading to multiple quantifiable theoretical predictions.
3 THEORETICAL ANALYSIS WITH RIVER VALLEY LOss LANDSCAPES
3.1 SETTING AND ASSUMPTIONS
We prove in the theorem below that the gradient flow starting from w will eventually converge near the river and remain close to it. Subsequently, if we project the iterate w(T) onto the river,
3

---

Published as a conference paper at ICLR 2025
15
1.
1.0
1.8
0.5
0.0
0.0
-0.5
-0.5
FlOns
Itrntes 841
wilarge
-1.0
River
1.0
PrOjectons KT. +0
-1.0
(a) Gradient Flow Dynamics
(b) Gradient Descent Dynamics
(c) Stochastic Gradient Descent
Figure 3: Validation of Theory on 2D Function. We validate Theorems 3.2 to 3.4 using 2D example
with the loss function xx) = (r2 sin(21))² + 0.2/10 - xi The hlue curve represents the "river", where the gradient aligns with the minimal eigeavector of the Hessian. (1) On the left, we observe that multiple the river and follow closely thereafter, consistent with Theorem 3.2 (2) In the middle, we show that discrete gradient descent with a learning rate y 0.6 shows similar behavior: after initial oscillations, the gradient descent iterates align closely with their projections on the river. The inset illustrates that the t-th projection's 's progress along the reference flow (eq. (T) approximately equals 1yt, as predicted by Theorem 3.3 (3) On the right, further illustrate that stochastic gradient descent (SGD) also tracks the river. In contrast the discrete-step gradient descent, the iterates oscillate around the river rather than staying on it. The trajectory with larger learning rate exhibits both faster progress and greater oscillations compared to the trajectory with a smaller learning rate, as predicted by Theorem 3.4
randomly initialized gradient flows converge
the projection will move along the river at pace similar to the reference flow ar(-) (eq. 1T) This phenomenon is visualized on Figure 3al We will now formally present our theory. We use w E Rd to denote the parameters and L to denote the loss. Further, use 0 (H) and us (H) to denote the k-th largest eigenvalue and eigenvector of a matrix H, respectively. The "river" in the river valley is I-dimensional manifold M formalized below.
Assumption 1. We assume the existence af "river", which a I I-dimensional manifold M such that any point w E M has gradient V ([(w)) that is in the same direction as the minimal eigenvector direction of the Hessian, Ud (ww).
Under this assumption, at every point on the river, the gradient V((w) will align with the locally flattest direction, va (V²)(w)), which we refer as the river direction. All other directions orthogonal to the river direction are considered as the mountain directions, corresponding to the steep hillsides in our conceptual picture.
We will consider neighborhood U of the river M with the following technical assumptions.
Assumption 2 (Regularity Assumption). There exists an open set U containing M satisfying the following assumptions:
1. Analyticity. L(w) is analytic with respect to w. Bounded Hessian. There exists a constant Tmax > 0, such that Vu E U. [V2L(w)) S Tmax- Existence of Eigengap. There exist constants "iat, 7 > O, such that Yw E U, Aa-1 (V2L(w)) >
Y + 4' Dda (V²L(w)) | TBat-
Slow Spinning of va- There exist constants A > Ami 0, K E [0, 0.01), such that Vw e U, Amin [VL(w)le S A. and Vea (V²)(w) llop S k2/(2A). This means that the river direction Va changes slowly during optimization.
U -M, the gradient (((w) is not parallel va va(V2 ((V))(W).. Conservation of Gradient Flows. There exists an open subset V U and constant >> 10A for Y defined Assumption 2.3 such that Vu E V, the r-neighhorhood of the gradient flow starting from stays in U for continuous time T Tmax 2 10 10 ( (K4min))/
Uniqueness of M. For any point
Throughout the analysis, k should be treated as a small dimensionless constant, indicating the river spins slowly.
Definition 3.1 (River Valley Landscape). If a loss function L satisfies Assumptions M and then we will claim that the loss function is river valley.
One simple example of river valley landscape is the quadratic loss L(x1, x222 equals to 0. In this case, the river is simply the line x2 = 0. However, the river valley landscape can also be more complex and non-convex, see Figure 3 for an illustration. We will prove that the iterates will follow the river with a predictable pace, which is characterized by the reference flow.
22 with k
4

---

Published as a conference paper at ICLR 2025
Reference Flow. We introduce Riemannian gradient flow constrained to the river M, serving reference in the following theorems. This flow intuitively represents the dynamics of iterates during gradient flow on the loss constrained by the river. We will denote the projection the tangent space of the river as (M(W) for w E M and choose an arbitrary starting point Zo on the river. The reference flow is defined
dx(T) PM (a (T)) VI(T(T))(T, z0-
(1)
Here, we use x to represent a point on the river, distinguishing it from w; which denotes a weight in the original space. T refers the continuous time variable.
3.2 MAIN RESULTS
Gradient Flow Dynamies. We will now consider gradient flow in the river valley landscape starting from a point w V, with defined in Assumption 2 6:
dw(T) = (T) dT, w = w E V.
(2)
Theorem 3.2. If a loss is a river valley (Definition 3.1 for the gradient flow w() defined in Equation (2. the iterate will obey the following dynamics:
1. Converge neighborhood of the river after constant time Tconverge 20lgg2//////aa (nAmma))//
minr]1æ((T) -o(Teonverge)ll ll1 S 2kA/7
Track the river closely with the same pace as the reference flow. There exists a time shift To depending on w w(Teonverge). ), such that for any T E [Toonverge, Tmaxl for Tmax defined in Assump-
tion 2 6, there exists T E [(1 ()T, (1 + e)7] satisfying that, || T) w(T) S 2nA/7,
for 33k.
The proof deferred to Appendix C.5 In this theorem, the lower bound on represents the time required for the iterate converge near the river. Here xTTo + T) can be viewed as a projection of w(T) onto the river. As both the geometric error (2xA/7) and the time-alignment error (e) vanish when k is small, this projection is not only close to (T') but also moves at nearly the same rate as the reference flow. Here the term To acts shift, reflecting the dependence on the initialization, as optimization trajectories starting from different initial points will enter the river at distinct locations. The term T represents the progress made along the river, consistent in the subsequent sections.
Gradient Descent Dynamics We will proceed gradient descent with discrete learning rate. Similar to the continuous case, an iterate far from the river will converge to the river (as visualized in the first few steps of Figure (3)). To case our analysis, we will skip the convergence analysis and assume the starting point w lies on the course of the river.
Wry (k: + 1) (t) (t) = -, (w,,, (t)), wa (0) = w E M.
Here we use i to denote the discrete time step, in contrast to the continuous time variable T used in the previous section. In this case, the progress along the reference flow over steps will be approximately 1p,, as shown in the following theorem.
Theorem 3.3. If a loss L is a river valley (Definition 3.1, when r) < 2 for the gradient descent defined in Equation ( with initialization w the river, there exists time shift To depending on w and n, satisfying that any S Tmax/n, there exists T [(1- c)nt,, (1 + c)nt) satisfying that, [Ta + T)- -1,(1)l2 < 10x A for 30k + 4npHaa
Wyy (T) (T)
The proof is deferred to Appendix C.6 We observe that the distance of the iterates from the river remains on the same order as Theorem 3.2 Finally, the Theorem 3.3 predicts that larger learning rate 77 will induce higher progress 7yt down the river given the same number of steps t, which is verified in the inset of Figure 3b)
Stochastic Gradient Descent Dynamics. The above analysis holds for deterministic dynamics and we will now proceed to model the stochasticity in the optimization process. This stochasticity will stop the iterate from fully converging the river and lead oscillation in the mountain direction. To simplify the analysis, we will consider a special case where the river direction is a constant and the river reduces to straight line.
5

---

Published as a conference paper at ICLR 2025
Assumption 3 (Straight River). For U in Assumption 2 Vw E U, Vea( V L(w))ll = 0. In this case, the river is a straight line parallel to the direction of val (V
Under Assumption val (())) is constant vector for E U and we will use vg denote this vector. We will also assume that the update is deterministic in the direction of the river, which simplifies our proof while still capturing the essential dynamics of SGD, Consequently, we can express the SGD update as follows:
(k + 1) xtt) 7a V (w (t) + MADK, 9k ~ N (0, (Ia- (0)) E M. (4)
Here N (14, Σ) indicates the normal distribution with mean 1r and covariance Σ. Compared to deterministic gradient descent, the introduced noise gl causes the iterates to deviate from the river instead of fully converging to it (see the difference between Figure 33b and Figure [c). Consequently, we need to impose additional assumptions (deferred Assumption
appendix) on the loss.
Stable Phase. We start with the stable phase, where the learning rate 17s = remains constant. This theorem provides formal basis for decomposing the loss into its river and hill components.
Theorem 3.4. Suppose loss L is river valley (Definition 3.1 and satisfies Assumptions
and 5 Then, for any constants 5 E 1) and T S maxx. for su sufficently small learning rate
depending on the regularity constants, (Deferred Assumption 7 in Appendix) the SGD iterates
(defined in Equation (4)) with Tpp- = 1 satisfies that for any integer t, E {1/ [1/m7,, T/n), there exists a
T E [(1 (11-c))1 ntntt, (1 + E ())tt satisfying that, E[L (w((t) - L(x(T)) = (d 1)no² /2 cL where
& = 4myaat and |ez] < (d- 1mm0² (defined Appendix C.7.
The proof is deferred to Appendix C.7 In Theorem 3.4 the error term in the approximation of the pace of the projection remains the same as in the deterministic case (Theorem 3.3). However, the stochasticity introduces an additional hill component (d - 1)ma²/2 /2 to the expected loss at the iterate. The hill component increases linearly with the learning rate. We conjecture that the theorem can be extended to a general setting and verify this conjecture on toy loss (see Figure 3c
Decay Phase. Finally, we will consider the decay phase in training and will show that a proper decaying schedule can reduce the hill component of the loss rapidly. We will first define our decaying schedule, starting from step fg- [T/n] [ = (+(( to I S 1.1f,. We choose this schedule to maximize the loss decrease rate on a quadratic function (see Appendix because we perform quadratic approximations of the loss near the river in our analysis. Our theorem predicts that the hill component of the loss will decrease linearly with the learning rate under this learning rate schedule, consistent with the empirical findings in Hu et al. (2024).
"
Theorem 3.5. Under the setting of Theorem 3.4 the SGD iterates (defined in Equation ④ with the decaying learning rate schedule satisfies that for any integer t E [ta, 1.1t,), there exists a T E [[1 -64--,11 (1) (t))))( satisfying that, EE[ (1i())) - (xx(T)) I (d 1)m 2/2 EL with
T(I) T + Σ Tk-
WSD Small LR
The formal proof is deferred to Appendix C.8 Compared with Theorem C.32 the hill component is now dominated by (d - 1111))002// /2, scaling linearly with the decaying learning rate. When the oscillation level a is large compared to the loss changes along the river, the loss decrease can then appear faster in the decay phase than in stable phases (see Figure 4 Further, the decaying phase also makes progress along the river, which corresponds to the term Ei=,, 7i in the theorem. Finally, the terms used our theorem match the scaling law formulation in the concurrent work (Tissue et al. 2024)..
Predicted Loss
Lr Decay
Iteration
Figure 4: Predicted Loss Curve of
SGD By Theorems 34 and 3.5 on
Loss L(Z1, = 72 F1
3.3 VISUALIZING THE RIVER VALLEY.
We use a direct probing method to verify our theory. Our theory suggests that when the learning rate is large, the model will bounce back and forth between the sharp valleys. However, in the decay phase, the model will move downwards the hillside to approach the river. This suggests that if we connect two checkpoints in the stable phase, we should expect to see a projection of the valley, and connect two checkpoints in the decay phase, we should expect to see smooth decreasing curves.

---

Published as a conference paper at ICLR 2025
3.04
W2:20B tokens
Ws: 20B tokens
Figure 5: Probing Loss Landscape. We validate the river valley analogy by interpolating stable and decay phases GPT-2 pretraining experiments. We observe that loss resembles a valley when constrained on the segment con- necting two models during the stable phase and smoothly decreases when connecting two models during the de- cay phase.
3.02
:25 tokens
Loss
rWo (1 - tWW
2.96
tWo+ + (1 -
2.92
W, :55 tokens
0.0
0.2 0.4 0.6
0.8
1.0
0.0
0.2 0.4 0.6 0.8 1.0
(a) Stable Phase
(b) Decay Phase
To verify this, we pretrain a 124M GPT-2 model on OpenWebText. In the first run, we train the model with constant learning rate for 25B tokens and interpolate between two checkpoints at 20B and 25B tokens (Figure 5ap. In the second run, we branch off from the first run at 20B tokens and decay the learning rate for 5B tokens, and we interpolate between two checkpoints 20B and 25B tokens (Figure The interpolation results closely resemble our theory. This observation is also consistent with Sanyal et al. (2023) which shows weight averaging improves model performance in the earlier part of the cosine training runs, where the learning rates are higher. Additionally, the smooth decreasing curves we observed when connecting two checkpoints in the decay phase are consistent with the findings in Hägele et al. (2024).
4 UNCERTAINTY VARIATION IN DATA DISTRIBUTION SHAPES THE RIVER VALLEY LANDSCAPE
What causes the loss landscape to resemble a river valley structure? In this section, we propose and validate the hypothesis that variations in next-token uncertainty shape the loss landscape. When predicting deterministic fact, a large learning rate can boost the model's 's confidence, accelerating learning. However, when the next token is inherently ambiguous such as the continuation of a phrase like "I ahe-the model must learn calibrated distribution, which may necessitate smaller step size. This variation in uncertainty leads to differences in sharpness across the loss landscape, resulting in the river valley structure.
Alex
Toy Bigram Language. We formalize this intuition using a synthetic language composed of cities and names, where each city corresponds to a unique distribution of its citizens' names. For instance, one city might have highly deterministic distri- bution, with most residents named "Ken", while another city may have more diverse distribution of names. This synthetic language follows the structure in Allen-Zhu & Lii (2024). The goal learn the distribution of names conditioned on each city. We show that cities with more deterministic name distributions align with flatter regions in the loss landscape (the "river"). In contrast, cities with more diverse name distributions correspond to sharper regions (the "hillsides").
Stochastic
n.4
Bob
0.1
0.3
Determinisfic
0.9
Ken
Figure 6: Visualization of Toy Bi- gram Language. We design a syn- thetic dataset where each city has unique name distribution. The left shows the name distributions for two cities, one deterministic and one stochastic.
Formally, let the set of cities be represented by (1...., n) and the set of names by (1,...,m) m). Data generated by first selecting city uniformly at random, then sampling a name j according to the city's 's' name distribution. The name distribution for city is parameterized by a categorical distribution ([P..a
where Pt, represents the probability of selecting name j in city i, and each P., > 0. To quantify the uncertainty in each city's name distribution, we compute the Gini impurity of the distribution as: U,- Ig (name city i) =1 Σm= Pł [0.1- 1 -h)
The value of U, reflects the uncertainty of city i's name distribution. When the distribution is close to deterministic-ia. there exists a j such that P.,y is near U-U; approaches its lower bound of Conversely, for a nearly uniform distribution, U, approaches its upper bound of 1. Given this setup, we parameterize our model with θ E Ra where each row corresponds to a city and each column to name. The model estimates the probability of name j for city i using the softmax function exp(Or.)
m
Σh- " We use sampled data to train this model with cross entropy loss. The population
loss is given by: L(O)- = Σ. ((O.)))(().)))))) = - -m Pi. log or e This loss is separable across different cities, meaning that the contribution from each city is independent. The loss component (O...) captures the contribution from city i, and different name distributions across cities lead to different forms of f, Considering parameter O" that minimizes the loss L,, will
exp(O...)
7

---

Published as a conference paper at ICLR 2025
Parameters
Parameters
Perameter
2.3
2.2
2.0
2.2
WSD-5 Casine
2.1
WSD-5
2.9
WSD-5 Cosine
Cosine
2.0
25000
50000
25000
25000
2.2
1.28 Parameters
peak
Leaminn Rate
1.9
WSD-5
Cosine
cos/me
1.8
25000
50000
10000 20000 30000
40000
50000
Step
Stop
Figure 8: Comparison with the Cosine Oracles. We show that the WSD-S schedule perform similarly the Cosine schedules in single run. The in graphs visualize the terminating validation loss different Cosine runs. The largest validation loss gap between the WSD-S and the Cosine schedules is 6e-3. The lower right figure plots the learning rate curves used in this experiment.
show that cities with more stochastic name distributions correspond to sharper components in the loss landscape, as reflected by the average-direction sharpness of lj.
Lemma 4.1. The average-dired direction sharpness of loss component & O* equals the uncertainty of the name distribution (Uz). (r(V²(((0)) l e;,- Us.
Lemma demonstrates that at the global minimum, the sharpness associated with a city decreases as the city's name distribution becomes more deterministic. This aligns with the intuition that a deterministic token corresponds to a flatter loss direction. We can further establish the existence of a generalized river (Assumption ( this loss landscape under appropriate assumptions about P (see Theorem C.35). Along the river, the gradient remains nonzero only for the cities with more deterministic name distributions, reinforcing the connection between determinism and flatness in the loss landscape.
Empirical Verification. We empirically verify that the loss curve of WSD can be reproduced in our synthetic setting. The dataset used contains two types of cities: (1) deterministic type with name distribution's 's entropy less than 0.2, and (2) a stochastic type with name distribution's 's entropy greater than 1. Each type contains 1.8k cities and there are 10 possible names. We train the toy model defined previously on this synthetic data and replicate the non-traditional loss curve of WSD (Figure
We continue to show that the difference in uncertainty also shapes the loss landscape for Transformers. We convert the data into a synthetic language in the format "The resident of [CITY]: [NAME]" and fine-tune O.IB GPT-2 model, pretrained on Open WebText, using this synthetic data. We experiment with two different learning rate schedules: a constant schedule (stable) and a decaying schedule (decay). We then calculate the difference in loss between the two models' predictions for the first token "[NAME]". A significant Spearman correlation 0.388 is observed between the loss difference and the ground truth entropy per city. This correlation indicates that the loss decrease is greater during the decay phase for more stochastic populations. Furthermore, although the decay phase achieves lower overall loss, the mean loss for the deterministic sub- population is higher than in the stable run, suggesting that the stable run better learns the deterministic sub-population.
1.30
18 1.25
Decay Stable
Cosine
1.2800
1000 Steps
1200
Figure 7: Reproducing the Non- traditional Loss Curve. We repro- duce the non non-tradition traditional loss curve of WSD this synthetic language.
5 WSD-S: A SIMPLIFICATION OF THE WSD SCHEDULE
The goal of continual pretraining is to generate checkpoints that exhibit good performance at multiple compute budgets in run. Formally, our goal is to achieve multiple intermediate checkpoints 0ny. each corresponding to a computing budget (number of steps) Ti for k E (I,...,K).
A strong baseline to measure the performance of 0T, would be running cosine learning rates (Figure 8 lower right) for each budget Tr separately, decay the learning rate linearly to the cosine function between [0, . We will dub this oracle method as Cosine-Oracle. However, Cosine-Orac done in a single run and will incur a high total compute budget Σn T₃. A simple modification to Cosine-Oracle is to multiple consecutive cosine learning rates between Tk-1 and Tj (Figure 2h
can't be
8

---

Published as a conference paper at ICLR 2025
2.2
Cyclic-Cosine
peak
min
Cosine-WSD
2.1
Learning Rate
peak
Loss
min
WSD-Cosine
peak
2.0-
Cyclic-Cosine
Cosine-WSD
min
WSD-5
WSD-Cosine
peak
WSD-5
1.9
min
Init
Intermediate Checkpoint
Intermediate
Init
Intermediate Checkpoint
Intermediate Checkpoint 2
Checkpoint 2
Figure 9: Cosine Learning Rate Implicitly Hurts the Models for Future Continual Learning. We show that while WSD and the cosine learning rate schedule may produce similar validation loss in single run, model trained with the cosine learning rate schedule is implicitly hurt compared to the model trained with WSD for future Continual learning. On the 0.6B models, after training the models for 50B tokens using both WSD-S and the cosine learning rate schedule, we continually train two models for another 50B tokens using both learning rates. We observe that the model trained with WSD-S consistently outperforms the model trained with the cosine learning rate when used as the starting point for further training.
last row), which we will dub as Cyclic-Cosine. Cyclic-Cosine only requires a total compute budgets Tk but it leads to non-negligible performance loss compared to Cosine- Oracle (Hu et al. (2024)
Warmup-Stable-Dece Decay (WSD) addresses this issue by maintaining a main branch that keeps using constant learning rate after warmup process and branch off using a decaying learning rate to achieve intermediate checkpoints. One can then continue pretraining from checkpoint in the main branch by resuming with the same constant learning rate. Formally, WSD introduces decay starting points D D3 such that T:-1 < D; < T, WSD will then correspond to the following process (Figure 2b second r1) Get a main branch of checkpoints faain by running a constant learning rate schedule for Dr steps, and (2) For each k, run decaying learning rate schedule for Ti D₁ steps starting from 0 to get fr.- The above process reutilizes the main branch checkpoints Jmain for each Tj and hence reduces the total compute budget to Tx- + ET - DD)-
Recall that in the river valley landscape model, the Sarmmu--Stable-Dese be viewed as a combination of a large learning rate phase to speed up progress down the river and a rapid learning rate drop at the end to reduce the oscillation. Because the decay phase also makes progress along the river (see Theorem 3.5), we propose a simplified version of WSD, called Warmup-Stte Staupp-ttall--DDe Stable-Decay-i Decay-Simpliffe
(WSD) algorithm can
(WSD-S), that continues with another stable phase leaving off the end of the previous decay phase (see the first row of Figure 26) without separating the training process into two branches. Formally, the WSD-S learning rate schedule defined as follows:
/ecayy(T; Dd, Mmax. Mriin) D.] if if, D, <ts T3: otherwise.
7/k
(5)
7max
The key difference from our methods is the choice of initialization point when retraining starts. In WSD, the second stable phase uses the model before the decay phase, whereas we use the model after it. This process is more convenient to implement because it does not require rolling back the main branch after each decay phase. Here the learning rate decay function decay can take many forms that decay the learning rate from Maax Mmin over T. D3 steps. In this paper, we will use the
following decay function decay( TTimiin:haa
Tinin (1 +
I (0, 1, .T} for all
Timax
experiments (visualized Figure 25 first two rows). This function motivated by the analysis on quadratic functions in Theorem 3.5 The inverse of the learning rate linearly interpolates from the inverse of the maximum to the inverse of the minimum.
5.1 EXPERIMENTS
Architecture and data. We adopt the LLaMA architecture from Touvron et al. (2023) adjusting the hyperparameters to create four model sizes: O.IB, O.3B, 0.6B, and I.2B. The exact hyperparameters are deferred to Appendix These models are trained on the Pile dataset (Gao et al. 2020) with a context length of 4096 and batch size of 4M tokens.

---

Published as a conference paper at ICLR 2025
5:15 Paiarees
2.4
D35 Tauancees
0.15 Pananess
2.2
2.4
2,3
2.1
1.33
2000
2.2
2:3
3.5
25002
30020
2,0
50000
25092
59320
25070
50300
Figure 10: Comparison With WSD, We show that WSD-S performs favorably compared with WSD when the total computes is fixed, achieving a consistent improvement over WSD on all the model sizes when trained for approximately 200B tokens.
Implementation. We use a standard Adam optimizer. We set the batch size to 1024 and fixed the peak learning rate for the same model size for all the methods. For the O. and 0.3B models, use a peak learning rate of 6c-4, and for the 0.6B and 1.2B models, we use a peak learning rate of 4c-4. These values are chosen following current empirical practice (e.g. see (rroeevelld et al. (2024). We set the minimal learning rate to 0.1 of the peak learning rate. We use a TPU v3-256 model to train the model with the Levanter framework in Jax (Bradbury et al., 2018 CRFM 2024 The fraction of time spent decaying is chosen to be 10%. The only exception is that when running WSD on the 0.3B models, we encounter a loss spike after training for 22.5B tokens and decay at the checkpoint trained for 22B tokens instead. This change is in favor of WSD in our comparison between WSD-S and WSD. The detailed hyperparameters are deferred to Appendix D 5.1.1 RESULTS
WSD-S performs competitively with Cosine-Oracle. The three endpoints of WSD-S are set at 50B, 100B, and 200B tokens for all models. As shown in Figure g. WSD-S delivers competitive results compared to Cosine-Orat Oracle in a single run.
WSD-S significantly outperforms Cyclic- Crelic-Cosine. We compare the Cyclic-Cosine and the WSD-S on 0.6B models with a total token budget of 100B tokens. Both schedules reduce to a minimal learning rate at 50B tokens to obtain intermediate checkpoint. Our results show that WSD-S outperforms PeelicCCosine with a significant performance gap of 4c-2 (Figure (9). A common belief is that loss spiking after increasing the learning rate is the main cause of the performance loss in Cyclic- Cosiie. However, this belief does not explain the advantage of WSDD-S. We hypothesize that a model trained with small learning rate for too long, as with Cosine, is implicitly hurt compared model trained with a large learning rate for the majority of the run, as with WSD or WSD-S.
To show that the model trained with WSD is more suitable for continual training, we conducted ablation studies by interchanging the schedules in the latter half of the runs to create two new learning rates (Cosine-Wo WSD and WSD-Cosine). Among the four runs, the model trained using WSD for the first half consistently achieved lower loss in continual learning, indicating that WSD produces models more suitable for continual learning, even after learning rate decay.
WSD-S matches (and slightly outperforms) WSD given the same total compute. For WSD, we adopt the following comparison methodology: assuming 10% decay portion, to get three checkpoints at 12.5k, 25k, and 50k steps, WSD then requires corresponding total steps of 12.5k, 26.25k, and 53.75k. Hence, we examine whether WSD-S can output three models of matching or better performance in the same corresponding steps (see Figure TOp. Our results suggest that WSD-S consistently outperforms WSD when trained on 200B tokens and underperforms WSD only on the smallest scale experiments when we trained 0. O. models for 25k steps. As this is the smallest scale experiment, we conclude that WSDDS has slight advantage over WSD when the total compute is fixed. This matches our intuition that WSD-S can reuse the decay phases of previous checkpoints, leading to a more efficient use of the total compute. As a simpler version of WSD, WSD-S is more user-friendly for open-source pretrained models, allowing users continue training the final checkpoint without needing intermediate ones given that the pretrained models are trained with WSD or WSD-S.
WSD and WSD-S are not sensitive the fraction of time spent decaying We conclude with an ablation study on the fraction of time spent decaying, and the result is shown in Figure [5] The final performance matches tightly within the range of 8% to 12%, showing a small sensitivity to the choice of the decay portion. However, in our experiments, we observe that decaying near a loss spike can lead to a significant performance loss (Figure [5ighth With the large learning rate, the training runs tend to be very volatile and there are multiple loss spikes in the training (see Figure 8). If decay happens closely after a loss spike and the loss has not yet decreased to its original level. it is typical that the final validation loss will be worse by le-2 or even more. We observe the same phenomenon for WSD, and when such a scenario happens, we suggest either running longer till the loss stabilizes or rolling back to slightly earlier checkpoint before the loss spikes and decays from there.
10

---

Published as a conference paper at ICLR 2025
ACKNOWLEDGEMENT
The authors would like to thank the support of NSF 221 178 and also would like to thank the Google TPU Research Cloud for the computing resources that enabled these experiments

---

Published as a conference paper at ICLR 2025
REFERENCES
Rahaf Aljundi, Lucas Caccia, Eugene Belilovsky, Massimo Caccia, Min Lin, Laurent Charlin, and Tinne Tuytelaars. Online continual learning with maximally interfered retrieval, 2019.
Zeyuan Allen-Zhu and Yuanzhi Li. Physics of language models: Part 3.3, knowledge capacity scaling laws, 2024. /ttp:tttss///xxi.ooror
05405
Maksym Andriushchenko, Aditya Varre, Loucas Pillaud-Vivien, and Nicolas Flammarion. Sgd with large step sizes learns sparse features, 2023, URL /RRl:ttpees//ax....rr//a///////
Zhangir Azerbayev, Hailey Schoelkopf, Keiran Paster, Marco Dos Santos, Stephen McAleer, Albert Q. Jiang, Jia Deng, Stella Biderman, and Sean Welleck. Llemma: An open language model for mathematics, 2024. httpt:tttt///aatttt
org/ mtds//231.11550
Guy Blanc, Neha Gupta, Gregory Valiant, and Paul Valiant. Implicit regularization for deep neural networks driven by an rmmtein-uhlenbeer like process, 2020.
James Bradbury. Roy Frostig. Peter Hawkins, Matthew James Johnson, Chris Leary, Dougal Maclaurin, George Necula, Adam Paszke, Jake VanderPlas, Skye Waaneerann-Mibee and Qiao Zhang. JAX: composable transformations of Python+NumPy programs, 2018. URL
Yuhang Cai, Jingfeng Wu, Song Mei, Michael Lindsey, and Peter Bartlett. Large stepsize gradient descent for non-homogeneoos two-layer networks: Margin improvement and fast optimization, 2024. /tt.:tttee//aavvv..oo./aaaaa/a/a
08654
Shouyuan Chen, Sherman Wong, Liangjian Chen, and Yuandong Tian. Extending context window of
large language models via positional interpolation, 2023. URL
2306. 15595
Andrea Cossu, Tinne Tuytelaars, Antonio Carta, Lucia Passaro, Vincenzo Lomonaco, and Davide Bacciu. Continual pre-training mitigates forgetting in language and vision, 2022.
Stanford CRFM. Levanter. httpss//gaaaabbcco///ttttaaaaddt wtw.de/rrreeeeeeeeeeeeeeeeeeeeee 2024.
Yiming Cui, Ziqing Yang. and Xin Yao. Efficient and effective text encoding for chinese llama and alpaca, 2024. URL /tt/ 08177734444777777
Damek Davis, Dmitriy Drusvyatskiy, and Liwei Jiang. Gradient descent with adaptive stepsize converges (nearly) linearly under fourth-order growth, 2024. URL pictps:/aarxiv.or
abs/2409.1979
DeepSeek-AI, :, Xiao Bi, Deli Chen, Guanting Chen, Shanhuang Chen, Damai Dai, Chengqi Deng, Honghui Ding, Kai Dong, Qiushi Du, Zhe Fu, Huazuo Gao, Kaige Gao, Wenjun Gao, Ruiqi Ge, Kang Guan, Daya Guo, Jianzhong Guo, Guangbo Hao, Zhewen Hao, Ying He, Wenjie Hu, Panpan Huang, Erhang Li, Guowei Li, Jiashi Li, Yao Li, Y. K. Li, Wenfeng Liang, Fangyun Lin, A. X. Liu, Bo Liu, Wen Liu, Xiaodong Liu, Xin Liu, Yiyuan Liu, Haoyu Lu, Shanghao Lu, Fuli Luo, Shirong Ma. Xiaotao Nie, Tian Pei, Yishi Piao, Junjie Qiu, Hui Qu, Tongzheng Ren, Zehui Ren, Chong Ruan, Zhangli Sha. Zhihong Shao, Junxiao Song, Xuecheng Su, Jingxiang Sun, Yaofeng Sun, Minghui Tang, Bingxuan Wang, Peiyi Wang. Shiyu Wang. Yaohui Wang, Yongji Wang, Tong Wu, Y. Wu, Xin Xie, Zhenda Xie, Ziwei Xie, Yiliang Xiong, Hanwei Xu, R. X. Xu, Yanhong Xu, Dejian Yang, Yuxiang You, Shuiping Yu, Xingkai Yu, B. Zhang, Haowei Zhang, Lecong Zhang, Liyue Zhang, Mingchuan Zhang, Minghua Zhang, Wentao Zhang, Yichao Zhang, Chenggang Zhao, Yao Zhao, Shangyan Zhou, Shunfeng Zhou, Qihao Zhu, and Yuheng Zou. Deepseck I1m: Scaling open-source language models with longtermism. 2024.
Aaron Defazio, Ashok Cutkosky, Harsh Mehta, and Konstantin Mishchenko. When, why and how much? adaptive learning rate scheduling by refinement, 2023, URL abs/2310 0 07831
Aaron Defazio. Xingyu Alice Yang, Harsh Mehta, Konstantin Mishchenko, Ahmed Khaled, and Ashok Cutkosky. The road less scheduled, 2024, httLhttts:///att orq/abs waa./2// ht///////4 2405
15682
12

---

Published as a conference paper at ICLR 2025
Abhimanyu Dubey, Abhinav Jauhri, Abhinav Pandey, Abhishek Kadian, Ahmad Al-Dahle, Aiesha Letman, Akhil Mathur, Alan Schelten, Amy Yang. Angela Fan, Anirudh Goyal, Anthony Hartshom, Aobo Yang, Archi Mitra, Archie Sravankumar, Artem Korenev, Arthur Hinsvark, Arun Rao, Aston Zhang, Aurelien Rodriguez, Austen Gregerson, Ava Spataru, Baptiste Roziere, Bethany Biron, Binh Tang, Bobbie Chern, Charlotte Caucheteux, Chaya Nayak, Chloe Bi, Chris Marra, Chris McConnell, Christian Keller, Christophe Touret, Chunyang Wu, Corinne Wong, Cristian Canton Ferrer, Cyrus Nikolaidis, Damien Allonsius, Daniel Song, Danielle Pintz, Danny Livshits, David Esiobu, Dhruv Choudhary, Dhruv Mahajan, Diego Garcia- Diego Perino, Dieuwke Hupkes, Egor Lakomkin, Ehab AlBadawy, Elina Lobanova, Emily Dinan, Eric Michael Smith, Filip Radenovic, Frank Zhang. Gabriel Synnaeve, Gabrielle Lee, Georgia Lewis Anderson, Graeme Nail, Gregoire Mialon, Guan Pang. Guillem Cucurell, Hailey Nguyen, Hannah Korevaar, Hu Xu, Hugo Touvron, Ilyann Zarov, Imanol Arrieta Ibarra, Isabel Kloumann, Ishan Mísra, Ivan Evtimov, Jade Copet. Jacwon Lee, Jan Geffert, Jana Vranes, Jason Park, Jay Mahadcokar, Jeet Shah, Jelmer van der Linde, Jennifer Billock, Jenny Hong. Jenya Lee, Jeremy Fu, Jianfeng Chi, Jianyu Huang, Jiawen Liu, Jie Wang, Jiecao Yu, Joanna Bitton, Joe Spisak, Jongsoo Park, Joseph Rocca, Joshua Johnstun, Joshua Saxe, Junteng Jia, Kalyan Vasuden Alwala, Kartikeya Upasani, Kate Plawiak, Ke Li, Kenneth Heafield, Kevin Stone, Khalid El-Arini, Krithika Iyer, Kshitiz Malik, Kuenley Chiu, Kunal Bhalla, Lauren Rantala- Yeary, Laurens van der Maaten, Lawrence Chen, Liang Tan, Liz Jenkins, Louis Martin. Lovish Madaan, Lubo Malo, Lukas Blecher, Lukas Landzaat, Luke de Oliveira. Madeline Muzzi, Mahesh Pasupuleti, Mannat Singh, Manohar Paluri, Marcin Kardas, Mathew Oldham, Mathieu Rita, Maya Pavlova, Melanie Kambadur, Mike Lewis, Min Si, Mitesh Kumar Singh, Mona Hassan, Naman Goyal, Narjes Torabi, Nikolay Bashlykov, Nikolay Bogaychev, Niladri Chatterji, Olivier Duchenne, Onur Çelebi, Patrick Alrassy, Pengchuan Zhang, Pengwei Li, Petar Vasic, Peter Weng. Prajjwal Bhargava, Pratik Dubal, Praveen Krishnan, Punit Singh Koura, Puxin Xu, Qing He, Qingxiao Dong, Ragavan Srinivasan, Raj Ganapathy, Ramon Calderer, Ricardo Silveira Cabral, Robert Stojnic, Roberta Raileanu, Rohit Girdhar, Rohit Patel, Romain Sauvestre, Ronnie Polidoro, Roshan Sumbaly. Ross Taylor, Ruan Silva, Rui Hou, Rui Wang. Saghar Hosseini, Sahana Chennabasappa, Sanjay Singh. Sean Bell, Scohyun Sonia Kim, Sergey Edunov, Shaoliang Nie, Sharan Narang, Sharath Raparthy, Sheng Shen, Shengye Wan, Shruti Bhosale, Shun Zhang, Simon Vandenhende, Soumya Batra, Spencer Whitman, Sten Sootla, Stephane Collot, Suchin Gururangan, Sydney Borodinsky, Tamar Herman, Tara Fowler, Tarck Sheasha, Thomas Georgiou, Thomas Scialom, Tobias Speckbacher, Todor Mihaylov, Tong Xiao, Ujjwal Karn, Vedanuj Goswami, Vibhor Gupta, Vignesh Ramanathan, Viktor Kerkez, Vincent Gonguet, Virginie Do Vish Vogeti, Vladan Petrovic, Weiwei Chu, Wenhan Xiong, Wenyin Fu, Whitney Meers, Xavier Martinet, Xiaodong Wang, Xiaoqing Ellen Tan, Xinfeng Xie, Xuchao Jia, Xuewei Wang, Yaelle Goldschlag, Yashesh Gaur, Yasmine Babaei, Yi Wen, Yiwen Song, Yuchen Zhang, Yue Li, Yuning Mao, Zacharie Delpierre Coudert, Zheng Yan, Zhengxing Chen, Zoe Papakipos. Aaditya Singh, Aaron Grattafiori, Abha Jain, Adam Kelsey, Adam Shajnfeld, Adithya Gangidi, Adolfo Victoria, Ahuva Goldstand, Ajay Menon, Ajay Sharma, Alex Boesenberg, Alex Vaughan, Alexei Baevski, Allie Feinstein, Amanda Kallet, Amit Sangani, Anam Yunus, Andrei Lupu, Andres Alvarado, Andrew Caples, Andrew Gu, Andrew Ho, Andrew Poulton, Andrew Ryan, Ankit Ramchandani, Annie Franco, Aparajita Saraf, Arkabandhu Chowdhury, Ashley Gabriel, Ashwin Bharambe, Assaf Eisemman, Azadeh Yazdan, Beau James, Ben Maurer, Benjamin Leonhardi, Bernie Huang, Beth Loyd, Beto De Paola, Bhargavi Paranjape, Bing Liu, Bo Wu, Boyu Ni, Braden Hancock. Bram Wasti, Brandon Spence, Brani Stojkovic, Brian Gamido, Britt Montalvo, Carl Parker, Carly Burton, Catalina Mejia, Changhan Wang. Changkyu Kim, Chao Zhou, Chester Hu, Ching- Chu, Chris Cai, Chris Tindal, Christoph Feichtenhofer, Damon Civin, Dana Beaty, Daniel Kreymer, Daniel Li, Danny Wyatt, David Adkins, David Xu, Davide Testuggine, Delia David, Devi Parikh, Diana Liskovich, Didem Foss, Dingkang Wang. Due Le, Dustin Holland, Edward Dowling, Eissa Jamil, Elaine Montgomery, Eleonora Presani, Emily Hahn, Emily Wood, Erik Brinkman, Esteban Arcaute, Evan Dunbar, Evan Smothers, Fei Sun, Felix Kreuk, Feng Tian, Firat Ozgenel, Francesco Caggioni, Francisco Guzmán, Frank Kanayet, Frank Seide, Gabriela Medina Florez, Gabriella Schwarz, Gada Badeer, Georgia Swee, Gil Halpern, Govind Thattai, Grant Herman, Grigory Sizov, Guangyi, Zhang, Guna Lakshminarayanan, Hamid Shojanazeri, Han Zou, Hannah Wang, Hanwen Zha, Haroun Habeeb, Harrison Rudolph, Helen Suk, Henry Aspegren, Hunter Goldman, Ibrahim Damlaj. Igor Molybog. Igor Tufanov, Irina- Elena Veliche, Itai Gat, Jake Weissman, James Geboski, James Kohli, Japhet Asher, Dean-Baptiste Gaya, Jeff Marcus, Jeff Tang, Jennifer Chan, Jenny Zhen, Jeremy Reizenstein, Jeremy Teboul, Jessica Zhong, Jian Jin, Jingyi Yang, Joe Cummings, Jon Carvill, Jon Shepard, Jonathan McPhie,
13

---

Published as a conference paper at ICLR 2025
Jonathan Torres, Josh Ginsburg, Junjie Wang, Kai Wu, Kam Hou U, Karan Saxena, Karthik Prasad, Kartikay Khandelwal, Katayoun Zand, Kathy Matosich, Kaushik Veeraraghavan, Kelly Michelena, Keqian Li, Kun Huang, Kunal Chawla, Kushal Lakhotia, Kyle Huang, Lailin Chen, Lakshya Garg. Lavender A, Leandro Silva, Lee Bell, Lei Zhang, Liangpeng Guo, Licheng Yu, Liron Moshkovich, Luca Wehrstedt, Madian Khabsa, Manav Avalani, Manish Bhatt, Maria Tsimpoukelli, Martynas Mankus, Matan Hasson, Matthew Lennie, Matthias Reso, Maxim Groshev, Maxim Naumov, Maya Lathi, Meghan Keneally, Michael L. Seltzer, Michal Valko, Michelle Restrepo, Mihir Patel, Mik Vyatskvv, Mikayel Samvelyan, Mike Clark, Mike Macey, Mike Wang. Miquel Jubert Hermoso, Mo Metanat, Mohammad Rastegari, Munish Bansal, Nandhini Santhanam, Natascha Parks, Natasha White, Navyata Bawa, Nayan Singhal, Nick Egebo, Nicolas Usunier, Nikolay Pavlovich Laptev, Ning Dong, Ning Zhang, Norman Cheng, Oleg Chernoguz, Olivia Hart, Omkar Salpekar, Ozlem Kalinli, Parkin Kent, Parth Parekh, Paul Saab, Pavan Balaji, Pedro Rittner, Philip Bontrager, Pierre Roux, Piotr Dollar, Polina Zvyagina, Prashant Ratanchandani, Pritish Yuvraj. Qian Liang, Rachad Alao, Rachel Rodriguez, Rafi Ayub, Raghotham Murthy, Raghu Nayani, Rahul Mitra, Raymond Li, Rebekkah Hogan, Robin Battey, Rocky Wang, Rohan Maheswari, Russ Howes, Ruty Rinott, Sai Jayesh Bondu, Samyak Datta, Sara Chugh, Sara Hunt, Sargun Dhillon, Sasha Sidorov, Satadru Pan, Saurabh Verma, Seiji Yamamoto, Sharadh Ramaswamy, Shaun Lindsay, Shaun Lindsay, Sheng Feng, Shenghao Lin, Shengxin Cindy Zha, Shiva Shankar, Shugiang Zhang, Shuqiang Zhang. Sinong Wang, Sneha Agarwal, Soji Sajuyigbe, Soumith Chintala, Stephanie Max. Stephen Chen, Steve Kehoe, Steve Satterfield, Sudarshan Govindaprasad, Sumit Gupta, Sungmin Cho, Sunny Virk, Suraj Subramanian. Sy Choudhury, Sydney Goldman, Tal Remez, Tamar Glaser, Tamara Best, Thilo Kohler, Thomas Robinson, Tianhe Li, Tianjun Zhang. Tim Matthews, Timothy Chou, Tzook Shaked, Varun Vontimitta, Victoria Ajayi, Victoria Montanez, Vijai Mohan, Vinay Satish Kumar, Vishal Mangla, Vitor Albiero, Vlad Ionescu, Vlad Poenaru, Vlad Tiberiu Mihailescu, Vladimir Ivanov, Wei Li, Wenchen Wang, Wenwen Jiang, Wes Bouaziz, Will Constable, Xiaocheng Tang. Xiaofang Wang. Xiaojian Wu, Xiaolan Wang, Xide Xia, Xilun Wu, Xinbo Gao, Yanjun Chen, Ye Hu, Ye Jia, Ye Qi. Yenda Li, Yilin Zhang, Ying Zhang, Yossi Adi, Youngjin Nam, Yu, Wang, Yuchen Hao, Yundi Qian, Yuzi He, Zach Rait, Zachary Deiito, 7ef Rosnbrick, Zhaoduo Wen, Zhenyu Yang, and Zhiwei Zhao. The llama 3 herd of models, 2024.
org/ abs 22t.//27872288
Ethan Dyer, Aitor Lewkowyez, and Vinay Ramasesh, Effect of scale on catastrophic forgetring in neural networks. In ICLR, 2022.
yPeEa
J. Falconer. Differentiation of the limit mapping in dynamical system. Journal of the London Mathematical Society, $2-27(2):356-372 1983, doi: 10.11122jmmm//2-222.33666 URL https
7/academic oup. http:///ccoooeeaaatatreee2222//2
C. Daniel Freeman and Joan Bruna. Topology and geometry of half-rectified network optimization, 2017.
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. 2020.
Timur Garipov, Pavel Izmailov, Dmitrii Podoprikhin, Dmitry Vetrov, and Andrew Gordon Wilson Loss surfaces, mode connectivity, and fast ensembling of dnns, 2018. URI. httptss/s////xiv
org Bog/abs//ee /802. 10026
Priya Goyal, Piotr Dollár, Ross Girshick, Pieter Noordhuis, Lukasz Wesolowski, Aapo Kyrola, Andrew Tulloch, Yangqing Jia, and Kaiming He. Accurate, large minibatch sgd: Training imagenet in 1 hour, 2018.
Dirk Groeneveld, Iz Beltagy, Pete Walsh, Akshita Bhagia, Rodney Kinney, Oyvind Tafjord, Ananya Harsh Jha, Hamish Ivison, Ian Magnusson, Yizhong Wang, Shane Arora, David Atkinson, Russell Authur, Khyathi Raghavi Chandu, Arman Cohan, Jennifer Dumas, Yanai Elazar, Yuling Gu, Jack Hessel, Tushar Khot, William Merrill, Jacob Morrison, Niklas Muennighoff. Aakanksha Naik, Crystal Nam, Matthew E. Peters, Valentina Pyatkin, Abhilasha Ravichander, Dustin Schwenk, Saurabh Shah, Will Smith, Emma Strubell, Nishant Subramani, Mitchell Wortsman, Pradeep
14

---

Published as a conference paper at ICLR 2025
Dasigi, Nathan Lambert, Kyle Richardson, Luke Zettlemoyer, Jesse Dodge, Kyle Lo. Luca Sol- daini, Noah A. Smith, and Hamnanch Hajishirzi. Olmo: Accelerating the science of language models, 2024. httthttt///aaat wwwpaxiww......aaomm abs 2402. 00838
Kshitij Gupta, Benjamin Thérien, Adam Ibrahim, Mats L. Richter, Quentin Anthony, Eugene Belilovsky, Irina Rish, and Timothée Lesort. Continual pre-training of large language models: How to (re)warm your model?, 2023.
Md Yousuf Harun, Jhair Gallardo, Tyler L. Hayes, Ronald Kemker, and Christopher Kanan. Siesta: Efficient online continual learning with sleep, 2023.
Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun. Deep residual learning for image recognition, 2015.
Danny Hernandez, Jared Kaplan, Tom Henighan, and Sam McCandlish. Scaling laws for transfer, 2021.
Elad Hoffer, Itay Hubara, and Daniel Soudry. Train longer, generalize better: closing the generaliza- tion gap in large batch training of neural networks, 2018.
Jordan Hoffmann, Sebastian Borgeaud, Arthur Mensch, Elena Buchatskaya, Trevor Cai, Eliza Rutherford, Diego de Las Casas, Lisa Anne Hendricks, Johannes Welbl, Aidan Clark, Tom Hennigan, Eric Noland, Katie Millican, George van den Driessche, Bogdan Damoc, Aurelia Guy, Simon Osindero, Karen Simonyan, Erich Elsen, Jack W. Rae, Oriol Vinyals, and Laurent Sifre. Training compute-optimal large language models, 2022.
Shengding Hu, Yuge Tu, Xu Han, Chaoqun He, Ganqu Cui, Xiang Long, Zhi Zheng, Yewei Fang, Yuxiang Huang, Weilin Zhao, Xinrong Zhang, Zheng Leng Thai, Kaihuo Zhang, Chongyi Wang, Yuan Yao, Chenyang Zhao, Jie Zhou, Jie Cai, Zhongwu Zhai, Ning Ding, Chao Jia, Guoyang Zeng, Dahai Li, Zhiyuan Liu, and Maosong Sun. Minicpm: Unveiling the potential of small language models with scalable training strategies, 2024.
Alexander Hägele, Elie Bakouch, Atli Kosson, Loubna Ben Allal, Leandro Von Werra, and Martin Jaggi. Scaling laws and compute-optimal training beyond fixed training durations, 2024. URL
org/ http://////255 18392
Adam Ibrahim, Benjamin Thérien, Kshitij Gupta, Mats L. Richter, Quentin Anthony, Timothée Lesort, Eugene Belilovsky, and Irina Rish. Simple and scalable strategies to continually pre-train large language models, 2024.
Yiding Jiang, Behnam Neyshabur, Hossein Mobahi, Dilip Krishnan, and Samy Bengio. Fantastic generalization measures and where to find them, 2019.
Tosio Kato. Perturbation Theory for Linear Operators. Classics in Mathematics. Springer, Berlin, Heidelberg, 2nd edition, 1995.
Lingkai Kong and Molei Tao. Stochasticity of deterministic gradient descent: Large learning rate for multiscale objective function, 2020.
Timothée Lesort Oleksiy Ostapenko, Diganta Misra, Md Rifat Arefin, Pau Rodríguez, Laurent Charlin, and Irina Rish. Challenging common assumptions about catastrophic forgetting, 2023.
Jeffrey Li, Alex Fang, Georgios Smyrnis, Maor Ivgi. Matt Jordan, Samir Gadre, Hritik Bansal, Etash Guha, Sedrick Keh, Kushal Arora, Saurabh Garg, Rui Xin, Niklas Muennighoff, Reinhard Heckel, Jean Mereat, Mayee Chen, Suchin Gururangan, Mitchell Wortsman, Alon Albalak, Yonatan Bitton, Marianna Nezhurina, Amro Abbas, Cheng- Yu Hsieh, Dhruba Ghosh, Josh Gardner, Maciej Kilian, Hanlin Zhang, Rulin Shao, Sarah Pratt, Sunny Sanyal, Gabriel Ilharco, Giannis Daras, Kalyani Marathe, Aaron Gokaslan, Jieyu Zhang, Khyathi Chandu, Thao Nguyen, Igor Vasiljevie, Sham Kakade, Shuran Song, Sujay Sanghavi, Fartash Faghri, Sewoong Oh, Luke Zettlemoyer, Kyle Lo, Alaaeldin El-Nouby. Hadi Pouransari, Alexander Toshev, Stephanie Wang, Dirk Groeneveld, Luca Soldaini, Pang Wei Koh, Jenia Jitsev, Thomas Kollar, Alexandros G. Dimakis, Yair Carmon, Achal Dave, Ludwig Schmidt, and Vaishaal Shankar. Datacomp-Ina In search of the next generation of training sets for language models, 2024. httpst rxiv org/ abs 2206 11794
15

---

Published as a conference paper at ICLR 2025
Yuanzhi Li, Colin Wei, and Tengyu Ma. Towards explaining the regularization effect of initial large learning rate in training neural networks, 2020.
Zhiyuan Li, Sadhika Malladi, and Sanjeev Arora. On the validity of modeling sgd with stochastic differential equations (sdes), 2021.
Zhiyuan Li, Tianhao Wang, and Sanjeev Arora. What happens after sgd reaches zero loss? -a mathematical framework, 2022.
Hong Liu, Sang Michael Xie, Zhiyuan Li, and Tengyu Ma. Same pre-training loss, better downstream: Implicit bias matters for language models, 2022.
Hong Liu, Zhiyuan Li, David Hall, Percy Liang, and Tengyu Ma. Sophia: A scalable stochastic second-order optimizer for language model pre-training. 2024.
Ilya Loshchilov and Frank Hutter. Sgdr: Stochastic gradient descent with warm restarts, 2017.
Kaifeng Lyu, Zhiyuan Li, and Sanjeev Arora. Understanding the generalization benefit of normaliza- tion layers: Sharpness reduction, 2023.
Chao Ma, Daniel Kunin, Lei Wu, and Lexing Ying. Beyond the quadratic approximation: the multiscale structure of neural network loss landscapes, 2022.
Sadhika Malladi, Kaifeng Lyu, Abhishek Panigrahi, and Sanjeev Arora. On the sdes and scaling rules for adaptive gradient algorithms, 2023.
Sanket Vaibhav Mehta, Darshan Patil, Sarath Chandar, and Emma Strubell. An empirical investigation of the role of pre-training in lifelong learning, 2023.
Preetum Nakkiran, Gal Kaplun, Dimitris Kalimeris, Tristan Yang. Benjamin L. Edelman, Fred Zhang, and Boaz Barak. Sgd neural networks learns functions of increasing complexity. 2019.
Matteo Pagliardini, Pierre Ablin, and David Grangier. The ademamix optimizer: Better, faster, older, 2024. (ttps:/::///ax wwg. abs abs//240..0317
Rui Pan, Haishan Ye, and Tong Zhang. Eigencurve: Optimal learning rate schedule for sgd on quadratic objectives with skewed hessian spectrums, 2022. htt.:tttts//xiii//ett
2110.14109
Yan Pan and Yuanzhi Li. Toward understanding why adam converges faster than sgd for transformers, 2023.
Bowen Peng, Jeffrey Quesnelle, Honglu Fan, and Enrico Shippole Yar:: Efficient context window extension of large language models, 2023.
00071
Jack W. Rae, Sebastian Borgeaud, Trevor Cai, Katie Millican, Jordan Hoffmmann, Francis Song, John Aslanides, Sarah Henderson, Roman Ring, Susannah Young, Eliza Rutherford, Tom Hennigan, Jacob Menick, Albin Cassirer, Richard Powell, George den Driessche, Lisa Anne Hendricks, Maribeth Rauh, Po-Sen Huang, Amelia Glaese, Johannes Welbl, Sumanth Dathathri, Saffron Huang, Jonathan Uesato, John Mellor, Irina Higgins, Antonia Creswell, Nat McAleese, Amy Wu, Erich Elsen, Siddhant Jayakumar, Elena Buchatskaya, David Budden, Esme Sutherland, Karen Simonyan, Michela Paganini, Laurent Sifre. Lena Martens, Xiang Lorraine Li, Adhiguna Kuncoro, Aida Nematzadeh, Elena Gribovskaya, Domenic Donato, Angeliki Lazaridou, Arthur Mensch, Jean-Baptiste Lespiau, Maria Tsimpoukelli, Nikolai Grigorev, Doug Fritz, Thibault Sottiaux, Mantas Pajarskas, Toby Pohlen, Zhitao Gong, Daniel Toyama, Cyprien de Masson d' Autume, Yujia Li, Tayfun Terzi, Vladimir Mikulik, Igor Babuschkin, Aidan Clark, Diego de Las Casas, Aurelia Guy, Chris Jones, James Bradbury, Matthew Johnson, Blake Hechtman, Laura Weidinger, lason Gabriel, William Isaac, Ed Lockhart, Simon Osindero, Laura Rimell, Chris Dyer, Oriol Vinyals. Kareem Ayoub, Jeff Stanway, Lorrayne Bennett, Demis Hassabis, Koray Kavukcuoglu, and Geoffrey Irving. Scaling language models: Methods, analysis & insights from training gopher, 2022.
16

---

Published as a conference paper at ICLR 2025
Colin Raffel, Noam Shazeer, Adam Roberts, Katherine Lee, Sharan Narang, Michael Matena, Yanqi Zhou, Wei Li, and Peter J. Liu. Exploring the limits of transfer learning with a unified text-to-text transformer, 2023.
Elan Rosenfeld and Andrej Risteski. Outliers with opposing signals have an outsized effect on neural network optimization, 2023.
Baptiste Rozière, Jonas Gehring, Fabian Gloeckle, Sten Sootla, Itai Gat, Xiaoqing Ellen Tan, Yossi Adi, Jingyu Liu, Romain Sauvestre, Tal Remez, Jérémy Rapin, Artyom Kozhevnikov, Ivan Evtimov, Joanna Bitton, Manish Bhatt, Cristian Canton Ferrer, Aaron Grattafiori, Wenhan Xiong, Alexandre Défossez, Jade Copet, Faisal Azhar, Hugo Touvron, Louis Martin, Nicolas Usunier, Thomas Scialom, and Gabriel Synnaeve. Code Ilama: Open foundation models for code, 2024. URL https: wtg/
12950
Sunny Sanyal, Atula Neerkaje, Jean Kaddour, Abhishek Kumar, and Sujay Sanghavi. Early weight averaging meets high learning rates for Ilm pre-training, 2023.
Leslie N. Smith. Cyclical learning rates for training neural networks, 2017.
Samuel L. Smith, Erich Elsen, and Soham De. On the generalization benefit of noise in stochastic gradient descent, 2020.
Minhak Song, Kwangjun Ahn, and Chulhee Yun. Does sgd really happen in tiny subspaces?, 2024.
wttp:////////.....000...... 16002]
Howe Tissue, Venus Wang, and Lu Wang. Scaling law with learning rate annealing, 2024. URL
https::/erxxv..rrg/ org/abs 2408 11029
Hugo Touvron, Louis Martin, Kevin Stone, Peter Albert, Amjad Almahairi, Yasmine Babaei, Nikolay Bashlykov, Soumya Batra, Prajiwal Bhargava, Shruti Bhosale, Dan Bikel, Lukas Blecher, Cris- tian Canton Ferrer, Moya Chen, Guillem Cucurull, David Esiobu, Jude Fernandes, Jeremy Fu, Wenyin Fu, Brian Fuller, Cynthia Gao, Vedanuj Goswami, Naman Goyal, Anthony Hartshorn, Saghar Hosseini, Rui Hou, Hakan Inan, Marcin Kardas, Viktor Kerkez, Madian Khabsa, Isabel Kloumann, Artem Korenev, Punit Singh Koura, Marie-Anne Lachaux, Thibaut Lavril, Jenya Lee, Diana Liskovich, Yinghai Lu, Yuning Mao, Xavier Martinet, Todor Mihaylov, Pushkar Mishra, Igor Molybog, Yixin Nie, Andrew Poulton, Jeremy Reizenstein, Rashi Rungta, Kalyan Saladi, Alan Schelten, Ruan Silva, Eric Michael Smith, Ranjan Subramanian, Xiaoqing Ellen Tan, Binh Tang, Ross Taylor, Adina Williams, Jian Xiang Kuan, Puxin Xu, Zheng Yan, Iliyan Zarov, Yuchen Zhang, Angela Fan, Melanie Kambadur, Sharan Narang, Aurelien Rodriguez, Robert Stojnic, Sergey Edunov, and Thomas Scialom. Llama 2: Open foundation and fine-tuned chat models, 2023.
Szymon Tworkowski, Konrad Staniszewski, Mikolaj Pacek, Yuhuai Wu, Henryk Michalewski, and Piotr Miloš. Focused transformer: Contrastive training for context scaling, 2023.
Tom Veniat, Ludovic Denoyer, and Marc' Aurelio Ranzato. Efficient continual learning with modular networks and task-driven priors, 2021.
Mingze Wang, Jinbo Wang, Haotian He, Zilin Wang. Guanhua Huang, Feiyu Xiong, Zhiyu Li, Weinan E. and Lei Wu. Improving generalization and convergence by enhancing implicit regularization, 2024.
20763
Yuqing Wang, Minshuo Chen, Tuo Zhao, and Molei Tao. Large learning rate tames homogeneity: Convergence and balancing effect, 2022.
Jingfeng Wu, Peter L. Bartlett, Matus Telgarsky, and Bin Yu. Large stepsize gradient descent for logistic loss: Non-monotonicity of the loss improves optimization efficiency, 2024. URL
Chen Xing, Devansh Arpit, Christos Tsiigootis, and Yoshua Bengio. A walk with sgd, 2018. URL https::/arxiv.. http://aaaiiv.or///////2222222/
Kaichao You, Mingsheng Long, Jianmin Wang, and Michael I. Jordan. How does learning rate decay help modern neural networks?, 2019.
17

---

Published as a conference paper at ICLR 2025
Yang You, Jing Li, Sashank Reddi, Jonathan Hseu, Sanjiv Kumar, Srinadh Bhojanapalli, Xiaodan Song, James Demmel, Kurt Keutzer, and Cho-Jui Hsich. Large batch optimization for deep learning: Training bert in 76 minutes, 2020.
Xiaohua Zhai, Alexander Kolesnikov, Neil Houlsby, and Lucas Beyer. Scaling vision transformers, 2022.
Jingzhao Zhang, Tianxing He, Suvrit Sra, and Ali Jadbabaie. Why gradient clipping accelerates training: theoretical justification for adaptivity, 2020a.
Jingzhao Zhang, Sai Pranceth Karimireddy, Andreas Veit, Seungyeon Kim, Sashank Reddi, San- jiv Kumar, and Suvrit Sra. Why adaptive methods good for attention models? In H. Larochelle, M. Ranzato, R. Hadsell, M.F. Balcan, and H. Lin (eds.), Advances in Neu- ral Information Processing Systems, volume 33, pp. 15383-15392 Curran Associates, Inc., 2020b. URL https::/pproceednnss neur cttp://ss.n.urrs.....o..ooooooc wctppppprrrpppppppppppppp www.ffies//eee ft..ess/pppeee//2/////////////// Flle:7555577...d0000077076666160
18

---

Published as a conference paper at ICLR 2025
A RELATED WORK
Learning Rate Schedules. Learning rate schedules are crucial in deep learning, with previous studies exploring various options. [Smith (2017) was the first to propose cyclic triangular learning rate schedule that interleaves decreasing and increasing learning rates. Loshchilov & Hutter (0017) extended the idea to a cyclic cosine learning rate schedule. He et al. (2015) introduced the notion of warmup, which gradually increases the learning rate in the earlier training phase. [Goyal et al.
(2018); Hoffer et al. 2018); You et al. (2020) concluded that the learning rate should scale linearly with the batch size, which is further theoretically examined in Smith et al. (2020); [Li et al. (2021); Malladi et al. (2023). You et al. (2019) performed an analysis on why learning rate schedules are helpful and suspected that the large learning rate at the beginning phase is mostly useful for avoiding
memorization of noisy data, which is consistent with our analysis in Section 4
In the LLM era, works including Hoffmann al. (2022); DeepSeek-Al et al. (2024); Hu al. (0024)
examined how to choose learning rate schedules for pretraining. In particular, Hu et al. (0224)
introduced learning rate schedule called Warmup-Stable-Decan (WSD) that remains constant for the majority of the runs before decaying in language model pretraining, which were studied independently
in Zhai et al. (2022; Ibrahim al. (2024); ); | et al.] (2024). p. Raffel et al. (2023); Ibrahim et al.]
(2024) explored another possibility of using an inverse square root schedule to pretrain the language models. Defazio et al. (2023) proposes to use linear decay for the entire training run. DDfazii et al. (2024) shows that with appropriate iterate averaging, constant learning rate schedule can reach better performance than the cosine learning rate schedule. |Rae et al. (2022); Gupta et al. (2023) ; Hu
ct al. (2024); Ibrahim et al. (2024) examined how to choose learning rate schedule in continual learning setting and verified that rewarming-ut up cosine learning rate brings performance drops that costly to recover. A common belief is that the performance drop is due to the sudden increase in learning rate during rewarming-up. However, our work shows that increasing the learning rate after a short decay in WSD does cause a similar performance drop as seen with the cosine learning rate, challenging the previous hypothesis. Instead, we suggest that the performance loss associated with rewarming-up cosine learning rate due the implicit damage causes to the model, making unsuitable for continual training. On the contrast, WSD avoids such damage by maintaining a high learning rate during the stable phase, hence the sudden increase in learning rate does not lead to performance drops in continual training.
Continual Learning. Continual learning. the process of updating the model with newly collected data, can improve the models' knowledge and capability. Previous continual learning research
(Aljundi et al., 2019) Veniat et al., 202T; Cossu et al. 2022] Dyer et al. 2022 Harun et al. 2023]
Mchta et al. (023) assumed significant domain shift and aimed to avoid forgetting old know ledge
while learning new knowledge. Recent works including Hernandez et al.] (20T)): Lesort et al. (2023)
suggested that optimizers including SGD and Adam have a knowledge accumulation effect and the effect of catastrophic forgetting may be less significant than expected, especially when replay is applied. Our work mainly focuses on continual pre-training without necessarily strong domain shift and hence does not touch upon the effect of covariance shift. Continual learning is also extensively employed in large language models such as LLaMA extend their capabilities, such handling longer contexts (c.g., see Tworkowski et al. 2023) ; Peng et al. (2023); Chen et al. (2023;; Dubey
al. (2024) and references therein) or dealing with new languages and domains (e.g., see Azerbayey et al.
2024) Rozière et al. (2024): Cui et al. 2024) and references therein).
Theoretical Understanding on Loss Landscape. A long line of research aims to better understand the loss landscape in deep learning (e.g., see Freeman & Bruna (2017): Garipov et al. (2018); Li et al. (2020) and references therein). We will highlight several phenomena that are related to our findings.
(1) Ill-conditioned directional sharpness and heavy-tailed noise: Zhang et al. (2020a b) examined the gradient noise in language modeling and observed that the noise is heavy-tailed in multiple dimensions. Pan & Li (2023); Liu et al. (2024) showed that the loss has vastly different curvatures in different dimensions. Pan et al. (0022) analyzes optimizing a quadratic function with skewed curvature theoretically. Our river valley landscape is consistent with these findings.
(2) Benefit of large learning rates: Large learning rates have provable regularizing effect in finding
flatter minima (Kong & Tao. 2020 Wang al. (022), and flatter minima typically have better generalization effect, even in the pretraining setting (Jiang et al. 2019 Blanc et al. 2020 [iu et al.
2022 Li et al.
2022 Ma al. 2022 Lyu al. 2023 Andriushchenko al. 2023
19

---

Published as a conference paper at ICLR 2025
(3) Connecting loss landscape with feature learning: Some recent works (Nakkiran et al. 2019 Rosenfeld & Risteski 2023) tried to understand how the loss landscape is formed through the lens of feature learning. Rosenfeld & Risteski (2023) showed that large learning rate will cause oscillation in learning subtle classification rules while continuing to learn other more deterministic features. [ang et a (2024) studied how to improve generalization and convergence by amplifying the update provided by the optimizer in the flat direction of the loss landscape. Wu et al. (2024); Cai et al. (2024) studied gradient descent dynamics on logistic regression, showing that large learning rate will cause oscillation in the earlier phase but will lead to higher progress later in training. Pagliardini al. (2024) developed modification of the Adam optimizer based on optimization analysis on the
Rosenbrock function, which is special case of the river valley landscape. [Song et al. (2024) shows that when SGD update is projected to the dominant subspace of the Hessian, the model's s optimization progress slows down and they conjecture the existence of ill-conditioned valley in the landscape, which can be viewed as similar and simpler version the river valley landscape discussed this paper.
(4) Ravines in the Loss Landscape. Concurrently with our work, Davis et al. (0024) identified the existence of ravine in the loss landscape- a manifold where every point has vanishing gradient within the sharp eigenspace of the Hessian. This feature appears in any smooth loss function exhibiting fourth-order growth near minimizers. They also demonstrate the advantages of using adaptive step sizes in this context. The concept of a ravine aligns closely with the river structure described in our paper and can be considered a specific instance of it.
The landscape analysis described in these previous works matches our river valley picture at a high level.
B ADDITIONAL EXPERIMENT RESULTS
B.1 PRETRAINING EXPERIMENTS DCLM
WSD-S outperforms WSD. We reran our experiments on another dataset called DCLM (Li et al. (2024) with the O.IB and 0.6B models for both WSD and WSD-S. We use learning rates 6e-4 and 5e-4 respectively for both models and a linear learning rate decay in the decay phase. The rest of the no longer suffer from
hyperaarameters is the same before. We observed that on this dataset, loss spikes and our results continue to hold. We also tested our final models on a sampled validation set of Penn Treebank, RedPajama, RefinedWeb, and the English subset of C4. The models trained with WSD-S continue to outperform the models trained with WSD.
0.1B Parameters
0.6B Parameters
3.2
2.8
3.1
2.7
Loss
2.9
WSD
2.5
WSD
WSD-S
WSD-S
2.8
2.4
25000
50000
25000
50000
Step
Step
Figure 11: Comparison With WSD DCLM. We show that WSD-S performs favorably compared with WSD when the total compute is fixed the DCLM dataset, achieving consistent improvement over WSD all the model sizes.
A learning rate sweep for WSD-S and Cyclic-cosine We perform a learning rate sweep for the Cyclic-cosine method and WSD-S method on the DCLM dataset. Both methods are trained for 25000 steps and are decayed to minimal learning rate 12500 steps. The peak learning rate and corresponding final loss are shown in Table I We observe that WSD-S outperforms Cosine- Rewarmup for most choices of the learning rate and the best performance of WSD-S is also better.
20

---

Published as a conference paper at ICLR 2025
LR
5E.4
1E-3
2E-3
4E-3
Cyclic-Cosine
2.54674
2.51853
2.49672
2.50063
WSD-S
2.52739
2.50944
2.49565
2.51052
Table I: Comparison of methods across learning rates.
B.2
ADDITIONAL. MODE CONNECTIVITY RESULTS
Ablations on the experiments in Section 3.3 We ablate the experiment results presented in Sec-
tion 3.3
varying the starting point and the duration used for decay and stable phase in Figures [2
and
13
Our results continue to hold.
3.14
R6 55.08 tokens
3.05-
W.
10 /B takens
3.13
3.04-
3.12-
3.03
3.11-
3.30
J0. OWL
3.08
(W 11 0I%
2.01-
WW + tWW
W..55.5 tooenns
3.00
Decay at tokens for 0.25B (b) Decay at tokens 0.5B (c) Decay at tokens 0.5B tokens
tokens
tokens
Figure 12: Loss smoothly decreases in decay phases We vary the starting point of the decaying phase and the duration decaying phase and find that loss generally follows the smooth decreasing trend when connecting two models during the decay phase. The experiment setting is the same as Figure 5
Wi 14-- tekens
3.14
tokens
3.13
Wi 5.58 tokens
tokens
3.12-
W; 6.08 tokeins
3.11
3.11
2.03
3.10
3.62.
(WA (1W:
119
3.09
3.08
Stable SB tokens for 0.SB (b) Stable 5B tokens for IB to-(c) Stable 10B tokens for 0.5B tokens
kens
tokens
Figure 13: Losses exhibit valley shape in stable phase We vary the starting point of the stable phase and the duration of the stable phase and find that loss generally exhibits valley-like shape when connecting models during stable phase. The experiment setting is the same as Figure
2-dimensional visualization of loss. Given a checkpoint trained using a constant learning rate, decay the learning rate to obtain a decayed checkpoint A'. We then continue train the checkpoint using a constant learning rate to obtain checkpoint B and corresponding decayed checkpoint B'. Our assumption states that the loss much sharper along the line AA' (the sharp hillsides), then along the line A' B' (the flat river). We present visualization of the loss in this section, validating this assumption.
C OMITTED PROOFS
C.I NOTATION.
To denote aTb for two vectors, we will (a, b). We will use the following function to denote the directional derivative of mapping F Rdd R.
F(r ov) F(x)
VF(x)(o]
lim a=+
21

---

Published as a conference paper at ICLR 2025
Figure 14: 2-dimensional Probing of Loss Landscape. We choose A to be a 0.1B GPT-2 model trained WeenWebText for 5B tokens ing constant learning rate 6e-4 and A' being the model after decaying learning rate on 0.25B tokens. B a model trained on another 0.5B token using the constant learning rate after A and B' is the model after decaying learning rate on an- other 0.25B tokens. Our visualiza- tion shows that the loss is much flat- ter in the direction of A' A'B' com- pared with the loss in the direction of AA'.
2.0
15-
417
19
1.
10
0.5
HI Direction
55
0.0
431
10
-2.0
River Direction
C.2 A WARMUP ON THE QUADRATIC FUNCTION
We will first motivate the decaying function we choose using a simple example on quadratic function. Lemma C.1. Assuming that considering the following gradient descent
yk: Vna (<v) /2) "99ks 9t E o0,
yk++1
Suppose тh = Tmax and yo yo follows normal distribution N (0, Tmmax 2-y 37mae? Then the following
two statements hold,
Vf, T7k 1jo yk will follow the same distribution as yo- 2. Consider all the learning rate schedule the following is the optimal
VI 2 1, 1 7k
Tinax
in the sense that it yields the fastest expected loss decrease. Suppose Tee corresponds to iterates variables Vio for any The and its corresponding iterates variables yk.
B(ruk /2] 2 (al)2<((zzze BB(((N))
Tlk-
We will denote ox = -Ey? and assume WLOG we start decayping at step 0. Then we will have
Proof.
σ = - (1-117))0:11 12.02
we choose all 77 7max we can directly verify that σы = o.
If we choose 7
04-17 to minimize the right hand side, we will have that
06 17 407
0R-102
Th 0h-17
+
2
1
2k
+
0k-1
0²
00
0²
This implies 0g
and plugging into Mi
0k-17
we have that
+ 26
at
72-17
7
7
o²
02
a
The optimality of vi can be easily inferred from the proof.
22

---

Published as a conference paper at ICLR 2025
C.3 LANDSCAPE ANALYSIS
We will parameterize (r(x) as ua (a(x))a(x) and (s( to denote I- Pr(x). Throughout this section, we will assume vz(2)) is continuous and pointing towards the direction of the gradient for all the x on the river. The following technical lemmas will be used repetitively in the proof.
Lemma C.2. Under Assumptions and 2
the directional derivative of (s(x) and Pr(x) exist and
satisfied that
[Pp(2)]u] = - -VP(())j) = Voa(<)(ooaatt) + ==(())
Further
V [s(r))u)Ppp = (va, Pra) Vra()))))
Ye
||1|/2-
A
Proof. As "itat. is an unigue eigenvalue of 2 + ut) and V [( + vt) is analytical with respect to f, by Theorem 1 of Kato (195)), we know that 6a(a - ut) is analytical with respect to t. Hence, the directional derivative exists.
The proof is by applying the chain rule and noticing that (1a(2), [ea(<)(u]} 0 because va(x) is always unit vector.
We will now define the projection of iterate to the river as the progress measure of the optimization dynamics.
Definition C.3. For U in Assumption and any w E U, we define the following ODE as the projection flow:
o(w, 0) w, do(w,t) -
Ps- (w) VL (ø(w, ()) dt.
(6)
When lim,-+oo Φ(w, 1) is well defined, we will define (ww) = lim o(w, 1) as the projection of w the river.
The following lemma ensures that the projection function is well-defined and is close to w: Lemma C.4. Under Assumptions and for any w satisfying that B(w, 2 U, ((w) M exists and w|l ((w) Il2 S exponentially, ( 1) (o t))l S exp(->t/2)) Ps (x)
Ps(w L Moreover, movement along the projection flow decays 7 2n
Proof. We will track ([ t))LL along the projection flow before o(w, leaves U.
d) VL((
dt
v)) t)), dPs()(w,t)
(ol( t)) + Ps(ó(w, t)) dVL(d(w, t))
dt
dt
By Lemma C.2 and assumption the first term can be bounded as
(Ps(ó(m, ()) ()), dos( 1)) (((((w,))) ()))
dt (Ps(o(w, ()) VPs 74
< (V (d(, t)) < <w//Psdddee))))/(((.....)
The second term is always negative
(Ps(o(w, ()) Ps(0(w, 1)) dVL(o(w,t)) t),
di
(Ps(o(u, Ps(o)
Ps(d(w, 11 ((((e.
23

---

Published as a conference paper at ICLR 2025
Summing up the two terms,
d| ()()(w,))) VL) S 5- [= + m (|al
dt
By Lemma C.36 we have Hence vt >> 0,
( - wwl22 ([Ps( 12)
[|Ps(w)VL(()l>>> / exp(- 2 11/2)
VL
( mas)
As (ww, 2) C U, the analysis hold along the trajectory and this shows that +(w) lim ((w, t)
(-00
exists and that (((w) wl12 < 21 Lww Further (w) satisfies that (3( VL)(() = 0, and by Assumption 2 ((w) M.
(7+2
The following lemmas focus on the properties of 0D.
Lemma C.5. Under Assumptions and for any satisfying that B(w, 2 U. (D(w) is well-defined.
Proof. Recall that (ww) lim (oo (w), as ø is differentiable (Lemmaa C.2 and M is n--00
the fixed point of Q by Theorem 5.1 of Falconer 1983 we have that (b is well-define.. n
Lemma C.6. Under Assumptions I and 2 for any satisfying that B(w, 2 ) U, it holds 7 that @ Ps (L()VLLLL))) = Further for any w = axt) E _.M it holds that (w(w) = da(t)) and that for any (W)( aligns with da(t)
0, 0 (w) dx(()
Proof. According to Lemmas C.4 and C.6. 4 OD is well-defined when B(w, 2 U. Based on Definition have that
VI, ((((w,)) ()) (wx).
Hence,
d (#(((w,)) t)) 1e=0= O,
dt
Therefore,
do(w,t) t) 10=0= (Ss
(w(w)
di
For any w M and any v RRd, it holds that
0 d00(w + oz)) Ps(w + V +
da
a=
( + / (PP
77400)/|B88077777000000000000010 + Ps(w))
= @w)w) (
Define (w) as the projection from v to 0Ps VL + Vs(
Lemma C.7. (w((V) is linear projection and the range Jmn is the range of Ps.
24

---

Published as a conference paper at ICLR 2025
Hence the range of Ju is a subspace the range Ps. When u + 0, based on Assumption
Proof. Based on Lemma
C.2
PsOPs(w)]o]VLee) =
[1)(2)) 2 (w) (1000 2 >/|u||2 - 7x||21122 >>
[(P'(w))
Hence the range of Ju has a dimension no smaller than the dimension the range of (s(w). This concludes that the range of Jw is the range of Ps(w).
Hence by Equation (7) and lemma C.7 it holds that for w E M, (s(w))((w) = 0. This shows that the range of ( ) has dimension 1.
Finally for any E M, ((w)= w. Hence,
dd(x(()) dx(l) dt
(<(()()) dx(() dx(l) Lemma C.8. Under Assumptions
dt. dt.
Hence the range of (ww) contains dz(t) this concludes the proof.
and for any w satisfying that Biw, 2 C U, holds that
PF (e(())))ee w ae w VL(w) Pp(w) S 5k,
PF V L(w) (o) 5n.
Proof. First, by Lemma C.6 it holds that
Define
st)) Ps(ó(w, (()ow, [)],
f(I) Pr (o(w, 1)) [b),
holds that s0) 0 and f(0) as o(w, 0) w.
We will bound the changes of st)) and ftt.. We will begin with calculating the time derivative of ae(w, [c).
ddp(w, D)[
(d(ww,) 1)
=0
[v]
dt
dt
à (Ps(ó(w, ()) f))) [v]
(Ps( 344.6/cc./100000/////1//44444410
Ps(((w,, P
=
Ps(olw.))VV-((((w.1)))
(8)
We will now bound d||s(r)la
dl|s(t)]13 =s)((t), ds(t), dt
-(((()), (dw(, t))l do(w, 1)
ddibww ()]v]
Ja t) Ps(ó(w, t)) - 2(s(t). (Ps(o(u, ()) () dt
dt
2(s(t), Ps(((w, ddo(w, (o))
dt
By Lemmas C.2 and C.4 and assumption] 2 the first term satisfies that, s (x(1s()))alls(()) f()]l2.
25

---

Published as a conference paper at ICLR 2025
By Equation g and assumption 2 the second term satisfies that,
(s(1), Ps(o(w, ddipwr, 1) dt t)))
(s(t), ( (w)
Hence 2|[s(t)]l2
dlle
+ (0--22//(00033---mmaaaaa00000000 and we can conclude
de
de
that
dl|s(t)ll2
< <=
dt
Similarly, can provide a bound for ( 1l.
(f(t) dt
-2 P= 1)) do(w,t) t)
ddo(w, t)]
16 4)0] PF(6(w, t))
dt
dt
2/ 0 w. P)) (ol,))
ddp(w, (v) 1l.
+ Pr(o(w, t))
dt
By Lemmas
C.2 and C.4 and assumption 2 the first term satisfies that,
(PP(o(w,t))( V (o((,, t))] (dd.)t))))) [u]l
((, ()) t)]ll (2)
By Lemmas
C.2 and C.4 and assumption 2
the second term satisfies that,
dōp(w, t)
[PF( t))
dl
|Pr(0(w, (0,, (Ww (oo ())l
= [P (ow,)) (dww,,
())I
Syxexp(-vi ())))l))((())) f()ll2--
Hence we can conclude that
df 22 < 22xmspttt11//1(1000011 1600)) dt
(10)
By Equation 99. holds that d(exp(<t//)]ls<()lla) exp()t/2)]s((())/// (xp(1/)))))))))))//2) dl|s(t)ll2
dt s 22xxcsp(nt/))]](((0
dt
Integrating the above equation from 0 to 4, and we have
[|s()]]2 < 2yx I exp()(F - (t)(fr-))))))(()llll))
(1I)
By Equations (10) and
we have that
d| (t)ll
< 2/xexp(-n/2)110002 44 1( exxp(((7 - (
dt
This suggests that
[f(())l2 < 1+ 2ynx 1' exp(->t/2)]l(())ldd +4222 1" [-) exp(>(r- 221)222(())(()))))
26

---

Published as a conference paper at ICLR 2025
Define (t) supusrse f(t), then it holds that
|| 1 [| 27// / exp( exp(->t/2)dt+ 4 /-" expl-1t/2) I exp(>(r - t)/2)dr
11+ (| (4k + 16x²).
This implies that vt, [2 & [+ S 1-4x-16k2 1 5k. By Equation 11 this suggests that (|s(()]l2 2 4x(1 5k) s 5x. Finally, returning to Equation (10). we have that
f(1)
dt 2l S 2y 2yxccop(---/////(000000) + (ss(t)
s 2yx cxxx 10k).
Hence [I](() /(0)l2 2 27xexp(=74/2) 10x) v 2x(1 10s) < 5n.
We have that
|Pr (Pr(o(w))00 (V)(()) ( (r lim [I((t) f(0) E [0,5k]
1-00
and
|Ps [ (
lim 500
|s(1]l2 E [0, 5x)]
|Pr
The proof is then complete.
The following lemma generalizes Lemma
C.8 general direction instead of VL(w).
Lemma C.9. Under Assumptions 1 and 2 for any satisfying that B(w, 23) C U, holds that
Pr [ (w)((()))))))))) (w) (o))0((w))))))) (wwwulle 5k,
PF (w) V (w)))l
Ps (w) (W)))ll 5R.
(P)lll
Proof. We only need notice that Pr(w)u aligns with va( [W))) Hence, it always holds that (Pr(w)V((w) (w)))
PF(w)u = (L)
[F 12
(w) [Pr(w))L((w)B
The proof is then complete.
The following lemma states that the angle between the gradient and the tangent direction is small for any point on the river.
Lemma C.10. For any E M, it holds that
4 [PM(()VL(w)> - A/t
Proof. Assume xTT), we will denote (M by
It holds that
V(Pg(w)VL(w)))))> =
which can be simplified to
Ps(w) VZ)
27

---

Published as a conference paper at ICLR 2025
The first term satisfies that (L (w)o) and the second term satisfies that [VVPs()))\\L(u)) I y<||el]. This then suggests || (P) S x||u||2-
Therefore ||Pr(w))l|2 (1 -)]lllll2 As
dx(t)
10=T= -PM (w) VL(w)
dt
We know that In (u)( 2 (1 x) PM (w) valle. which suggests that o1 Par(aris k). Hence we can conclude that |P 2 - x). Hence, we know that
2
2x
+ VL(w)l2 < ( 2x| VL())l S
|0112 S 4kc|u||2.
This concludes the proof.
The next lemma states that Vr(w)VL(w) and V (L()(w)) is always close.
Lemma C.11. Under Assumptions and for any w satisfying that B(w, 2 C U, it holds that
[PP](c)VL(u) - V - (ye + ww (2w)ll2.
Proof. By Lemma C.4 the line segment from (w) to w lies in U. By Assumption 2 and lemma C.4
[PPp(w)V(())) (
- =|| / VPF (4 (w) + +t 2 ( VL + + 4 (w)))dt
[ Pf ( (w) + (w- - (w)) L (w) + (w 2 (w)))) (w))ddl112
+
S(70 + Mhs ) || 1 (w)]l2
This concludes the proof.
The final theorem states that when is near the river, the movement of its projection has a similar value as the inherent speed at the river.
Lemma C.12. Under Assumptions
and when
| #(w)ll2 S
[P VL(ww) + dx(() [00(w)VL(u)) + dr
dx(r)
b-= 1l2 $ 16x|
dr 1==7 1|2-
dr
dx(r)
dr(T)
I== ||2- I==T 1|2 5 30x|
dr
Proof. By Lemma C.8
[a VLw) Pr VL))Vl & 10<[|Pp([)VV[(()))
Combining Lemma C.11 and ((w)ll2 <
10 1001|Po(ur(0))///(00)) we have that
74%
[Pr(w) VL(((w))llz - (yx ww +(w))l 10sl
By Lemma 10 let
dm(r) dr
I-
|| VL(((w))l < 4nlloll
Combining the three inequalities, we have that
10k
+ 4x
[P - <
1- 10x [VL((()))]
<
10x||e//2 S 12x 12x|[o||2.
1- 10k
28

---

Published as a conference paper at ICLR 2025
This suggests that
[PP(())LL()) - wl12 S (Pr(w))L(w)) (//00)))) ||+ (z((w))) S 16///o/12
Hence
1/ + (
< PP - w12 [
16k||U||2 10x ( (w))VL))
<30k||u|12-
This concludes the proof.
C.4 RIVER EXISTS UNDER MILD ASSUMPTIONS.
In this subsection, we will provide two results stating the local existence of rivers under the existence of eigengap (Assumption 4. Recall we define the river smooth manifold of points with vanishing gradients in sharp directions.
1. River exists and every point in V is close to some part of the river (Lemma C. 13
River is a I-dimensional manifold. (Lemma C.14).
Combining the two statements, we can conclude that there is always river near every point under the existence of eigengap.
Assumption 4. There exists an open set satisfying the following assumptions:
Analyticity. L(w) is analytic with respect w.
2. Existence of Eigengap. There exist constants "(Bat, Y > O, such that Vw E U. Ad (V*L(w))
Y 47maa \da [V²)(w)) | < Yiat
3. Slow Spinning of va- There exist constants A Amin 0, 00,0.01), such that Vu
U, Amin V S A, and V/ (V+L(w)) 1lop S N2/(2A). This means that the flat direction Ua changes slowly during optimization.
4. Conservation of Gradient Flows. There exists open subset V U and constant > 102 for defined Assumption 2.3 such that Vu E V, the r-neighhorhood of the gradient flow starting from w stays in U for continuous time Tmax 2 10 10loo(22/ (2nA(/AAAAiii m/7)
We note that Assumption 4 is a strict subset of Assumption 2
Lemma C.13. Under Assumption 4 for every us E V, there exists wt E U, such that ||w - w/ 1
2A
and || Ps (w]] M VWw')ll O.
Proof. We will define w' (ww) for T defined in the same way in Definition proof goes in the same line as Lemma C.4. We note that in the proof of Lemma our deduction does not depend the existence of the river until the last line.
C.3
The rest of the
Lemma C.14. Under Assumption 4 for every E U satisfying |Ps[s] = 0, there exists smooth I-dimensional manifold M passing through x, such that for every point u E M, the projected gradient onto the sharp directions vanishes, i.e.
LL())
Proof. To establish the existence of the river M as smooth I-dimensional manifold passing through xr we apply the Implicit Function Theorem to the system of equations defined by the vanishing of the projected gradient onto the sharp directions.
Fixing the coordinate vector as c....., Edt we will assume the rotation rotating from e1 to v, and keeping all the vectors orthogonal to eι and constant R(v). We then have R(v) is a smooth function of e as long as v is not close to é₁. We can then assume D(v) = Ei. ((())e)) (ddd-)xxd with ef being the (i - 1)-th coordinate vector in R D(v) is then also smooth function in u.
29

---

Published as a conference paper at ICLR 2025
We will now assume without loss of generality does not align with é1 and define K(u) ((va(u)), then ((u) is a smooth function in u that maps the sharp component of each vector to IR
Define the constraint function F R + Rd-1pp by
F(u) K (u)))s(
We aim to show that the solution set F F= (0) near is a smooth I-dimensional manifold. To apply the implicit function Theorem, we need verify that:
Smoothness. The function continuously differentiable in a neighborhood of æ. Given that L(u) is analytic (and hence smooth) by Assumption 4 and (s(u) is defined in terms of the continuously differentiable projection Pr(u), it follows that F is smooth.
Rank of the Jacobian d 1.
at æ. The Jacobian matrix DF (d-(())) must have rank
The Jacobian of F at point is given by:
0K DF =K() (<<))))s((())Vo)) + OKK))
=K(x) (P)() Ps())V>L(()))
Consider Ju defined Lemma C.7
DF(x)= = K (xx),z, Applying the same argument shows that Ja has rank d 1 and the range of J, is the range of Ps(x), i.e. the sharp space. Therefore D has rank d
Since F smooth and the Jacobian D DF( has rank 1, the implicit function Theorem guarantees that the solution set (0) near ((w) smooth manifold of dimension d (d- 1. Thus, there exists smooth 1-dimensional manifold M passing through ((w) where the projected gradient vanishes:
Ps(u)VL(u) = 0 Vu E M.
Furthermore, the smoothness of F ensures that the manifold M is not only locally 1-dimensional but also smoothly parameterized.
C.5 PROOF OF THEOREM
We will consider the following gradient flow:
dw(l) (V[((()) w(0) V.
(12)
We will first prove that along the gradient flow trajectory, it holds that ([Ps(w) is bounded.
Lemma C.15. Under Assumptions and along the gradient flow Equation 112 it holds that for
t 2 2
Proof. We will first compute how fast ||PS(((t))VL(((()))) change
(l]Pssu(t)V 
dt
By Lemma C.2 and assumption 2 the first term satisfies,
= +2440000000000000000000000000040
By Assumption 2 the second term satisfies.
-2( + Yam) (P0
30

---

Published as a conference paper at ICLR 2025
Hence,
d Vll
dt
(13)
We then consider the corresponding Pc(w(t))VL(u(t))
dI L(
dt
By Lemma and assumption 2 the first term satisfies,
<2ny|
By Assumption 2 the second term satisfies,
3224000h)))kkcch)))e)))44)444444 < 27ma 2mm (P(al(PP(((())))((()))))
Hence, we have that
d|| (()P((()))VVL((()))
(14)
di
Choose
VI-44 Ax
1.5k as the solution to the quadratic equation ka? - a + k = 0.
2x
Then combining Equations
L3 and [4 holds that
d VL -a aullPP(lt(t)VVLLttt))ll
-
dt
ST 41
Mha Wwt.://chht////0000000------1110
Notice that
(-1 + kax) -1
K
ax
Hence,
dl
S (7(1 k0x))) + That)
By Lemma C.36 this suggests that
< ccxp(-(((1-an) naw) +taa)) ) (1
Hence,
[Ps((((t)VL(((())) 5 1.5k| LLxx +oop: (PP(0)/))))))))/((((((()())())))
31

---

Published as a conference paper at ICLR 2025
Lemma C.16. Under Assumptions w(t)= U and [[w(t) (wu(t))l22
7 andd along the gradient flow Equation
12 it holds that
Proof. This a direct combination of Lemmas
C.4 and
C.15
Lemma C.17. Under Assumptions 7 and
2 along the gradient flow Equation 12 ifT(t) satisfies
a(T(t)) = (wo(t), then
dT(())
E [1 30k, + 30k].
dt
Proof. As Tt) satisties (T((t)) = ((w(t)), taking derivative on both sides yield,
dT((t) dx((T) |--T(=== - ( dt dr
By Lemmas C.1 and C. 16
it holds that
[0 (w))) - dx(T)
dr (T=T(1) lle S 30x dx(r)
dr l-=T(()
We then have that
dT((1) 1| - 30k,
dt
which concludes the proof.
Proof of Theorem
3.2 The proof is direct combination of Lemmas C.I6 and C.17
C.6 PROOF OF THEOREM 3.3
We will consider the following gradient descent:
Wk+1 Wr Lw)V wa M.
(15)
We will track the changes of PF (L(t)))L(((()) and (((w()))VL(w(t))) for simplicity, will denote them us fg(k) and sk(k). Further, we will use the following denotation
Wk,7 (1 T) + Tw++1
We will first prove some lemmas bounding the difference between gradient and projections at different points.
Lemma C.18. Under Assumptions and when w. E V. VT (0,1).. Wk,r U.
Proof. It holds that,
lw,, wk 1/2 S 57 S
27
Lemma C.19. Under Assumptions and when we E V. VT, r' [0, 1], it holds that (P(w.))))) - Ps(we,r)) 17K..
Proof. According to Lemma C.18 it holds that WK,7, Wk,7 E U. Assume without loss of generality T/
Ps(wk,)) - Ps(wkr+)) =11 [" (VV
<
SmyR.
32

---

Published as a conference paper at ICLR 2025
Lemma C.20. Under Assumptions and when we E V. VT (), it holds that
(L - (L 77/aa (/g)allsg(((]]2 + 217nr 7mxx 7K (L(w))ll
Proof. According to Lemma C.18 holds that wh Wk,r' E U. Define g(7') (V Ps(Pss wk. T', such that
JVL ( then by Lagrange's 's Mean Value Theorem, there exists
[[s((wRr)) VL(WE,r) (Ps(wn,.)VL((w.))
Ps(WErr VL - g(T) 900)
=rg((T) d d WK. (k,, VV(W.,,)VL(u.)))
dr'
dPs(wkk) V (( Ps VVL())
<
dr'
=n
Ps(WE.r L(Wx,r V [L(w))ll2
577maxx ( (L(())))))))))))
- (, (L
By Lemma C.19 it holds that
Tmax SLa (mo 1y V
[I(ss(w,,)) - Ps(wE. PP(N S < (V
Summing up and the proof is complete.
Lemma C.21. VT E (0, 1). holds that
- sg(k) S + VL
+ 3
Proof. This a direct combination of Lemmas C.19 and C.20 with
[Ps(werr (L(w,,)) (w, (Lw)))
Ps (wk,r [w...= Ps (wk,r VL + [[(Ps(w..)) PL(w))))V(())ll
The proof is then complete.
Lemma C.22. VT E (0, it holds that [P
VL(Wk..)) V
PF(wk, VV S man [lg + 27)² VL(w))]l2
Proof. Define g(7') |P VLL.)VLLWWWnnnn,,, - (F(w..))
VV:))ll2 P then by Lagrange's Mean
Value Theorem, there exists R', such that
(Pr(wx,)VLLwr. PPr VLtuil = g(T) -(0)
=TD((r) d VL
d7'
d (Pr(wk.r) (.W.r..)) ( 1/2
<1
dr'
=n Pr(ww, V
Smpmall (F(w,r))) VV (P - (r(ww.))) (oo
By Lemma
C.19 it holds that
Mit [[(p)(w..) -Pp(:,,)) LL (p 777maxx (V Summing up and the proof is complete.
S (macl 7/m [[,
33

---

Published as a conference paper at ICLR 2025
Lemma C.23. YT E (0, 1), it holds that
|P (rr[un:)VL((wee)))) f'((k)ll2 S [[5 3 LLwp:/l
Proof. This a direct combination of Lemmas C.19 and C.22
with
VL(wer)) (w (w:)))
<| [L(WEF)) PO (wk,r [L la + (w,,)) Px(w)))VL(())ll
The proof is then complete.
We will prove a discrete version of Lemma
C.15
Lemma C.24. Under Assumptions aand when nj <1/ <1/Tmmx along the gradient flow Equation 15
holds that || VL((v(())VL(((() S 10k]
as long as w(r) E V, ST.
Proof. We will first consider sg(k), By Lagrange's Mean Value Theorem, there exists т, such that, [|sg(k + 1)12 - [
Pl(Ps dr
Pss (L...))L(......l.l Wttp:///WLLLaaD----PP/P/////loo
-7Pss VL (wk,, /Ltcg + Ps(wk,r ((ww,.)VVLu)
The first term satisfies that
7 \/
5 |
The second term satisfies that
(w V Ps(w.,r [(wkrr \ VLw.))
- (((we))
-nny - 4 47ma) Ps ) [L( 7/mmax //s -L
As we have that ||a - 6 2 ||a|| 4 by Lemma C.19 it holds that
(Pss ((w,) + (Ps(wk.)) Ps(w)) VL(())1
(PPs(w))VL(v)) 4m (Ps(wx,r) (u:)) (L:))) S m
2
S [Ps((w)) VL(w))l1
(s(w)) V [w::)l1
+
(<mm)
2 [sg(())] + 2m7xx
[
Hence
M + 4ymaa) P
- 4mm +4 + 4r V
+
5-n(y m [Igg(()))]]]])))))))))))}})))) 2n7 [L(w))) [
By Lemmas
C.
and C.20 and m S 7/2, it holds that
Mimax
Pa - vL(w>,,)))l
< 3
http:///ooallllaaaooaoo//////lla 21yy ( (L(w)))l2)
34

---

Published as a conference paper at ICLR 2025
Hence, we can conclude that
|| + 1)112 - [
- + += [
Let - =l VL and - as 6(2a 22 +3 () +b) s - 100
holds that
(k)
[sg(( + 1)12 [1ag())]3
10r;=> 4 | (16)
Similarly, we can control the fg(k) changes. By Lagrange's Mean Value Theorem, there exists T', such that,
llfg(k + 1)112 [
=I [PPr(w,))VL( (Ww d) (
dr
-- -/ + Pr(weR v LV
The first term satisfies that
Stmp <>m/ 1wwp://///ooolllllldddddl0l11111 + 27 Imaxx 7/6 ((ww))l2) <4>ne(|VVL<u)]]]1
Similarly, the second term satisfies that
(L (L(w))
Sma
<
http:://hhhclllloo///////////111 (nalll] + 7<l
(2 (1o0////1((2
[Armmall/g())] + [,"">>>\|V[[()))
Summarizing and we have
ll 1)ll ( 2 5yyx||VV[(w]11 4 [lfg(k)
(17)
Let a₂ be the smaller positive solution of
5ka + (10x² + 5k
4 70 +
10x² = 0.
Then 0g (-10x² -5x+ +) (-11-1111221--10 6x+ -x 200m
100x?
10k
Then combining Equations (16) and 11
[|sg(k + 1)112 f. + (k 1)1l
<(1 4 (1l a, (l 1)113) 1113))mmm)) + 10n² 5x 770000² 55xaa (
=(1 41yYaat - - + 10x² 5xoa)))(]lsg())))) -
As |sg(0)]13 a (l)) 12 < O. we have that [k>g < a ll - 1)l12 < 100 (1 for
all the f.
Then show that gradient descent will also track the river closely.
35

---

Published as a conference paper at ICLR 2025
Lemma C.25. Under Assumptions w(t) U and [[w(t) (wu(t))l2 ((
7 and along the gradient flow Equation 12 it holds that
Proof. This a direct combination of Lemmas
C.4 and C.24
Finally, we will show that the movement of the projection of the gradient flow moves approximately at the same rate as the river, a discrete version of Lemma C.16
Lemma C.26. Under Assumpttons [and)
and 2 along the gradient Equation (12). along the gradient
descent Equation (5)). if) satisfies x (T)) ((( where [4] is the integer part of I, then
for any that not inleger
dT(())
E in (30k + am77)))) (30k + 4ry/aat)))
dt
Proof. Let [1] k, [4] fracda(r)dt I==T((): taking derivative on both sides yield,
T()) satisfies s<((()) $(w(k,, (4)) let
dT((t)
di u = -ww VL((u).
As the proof of Lemma C.22 there exists 7'
[[pr( [[(w)) Pr VV
+ /wt. Pston,
Smmall Pp VL + 770 V ((w)))
By Lemma C.19
it holds that
[ VL(())) Pp(wk, L (Lkk( )l12
Symmall (F(w,, [ 21 7mmax VL(wx)]la
(18)
By Lemma C.24
[VL(w))l2 -
[P ( VL(wx)ll
10k
5
10m
This shows that
1
[VL(wx]la 1 10x
VP (1 + (2xx))
1778
Combining with Equation
18 have that
|| (wk,r) (urk) Pr (u VL
Sipmat ( + x(1 12x)
(ww))ll
S(M +
This shows that
(myaat + 2xx
[P VL((e)
Pp (w(,,rr) (L(Wk, 5 -m + +2k VL
<(2 3x) [x))[PP((WH, (x
(19)
By Lemma C.12
|P VL + 0ll2 S 16k//u||2
(20)
36

---

Published as a conference paper at ICLR 2025
Combining Equations
19 and 20
[P((w..)) [(wx) + u112
(r(wk.r) ) (T(x))
Pr(wr. VL
| + 0/12
< (21/mat + ( ( 16x|
<(2 + 3x) 16k) 16k3]
$(19 + 317 2/21/22
(21)
By Lemma
C.9
|0p (k,r [L(wx) VV(wll.)VVLL(>>>l & (0s())p(uu,,))L(ee)))
VL(w))ll
(22)
Combining Equations
21 and 22 it holds that
(wx, v (wwx) + ull2
< (w (ww)) - VV(W,, + || (PPP(w,r)VLLuut + u112
10 VL + (19k + 37yla:))loll2
(300 + 4777max) ||0122
Hence
dT(())
E (30k 477maa) (30k + 4n7mat)))
dt
This concludes the proof.
Proof of Theorem
3.3 The proof is direct combination of Lemmas C.25 and C.26
C.7 PROOF THEOREM 3
Assumption 5 (Regularity Assumption for SGD). In the setting of Assumptions 2, assume in addition the following:
Bounded Hessian. There exists a constant 0, such that for any weight w E U, the nuclear
norm the Hessian is bounded. L(W)l- = Σ > (V²L(w)) ST.
2. Bounded Third Order Information. There exist constants p 0, [0,0.01]. such that
||V [(w) loop S p, Ap S k'n?. 3. Bounded Loss. There exists constant M O such that Yw, L(w) < M.
In this assumption, we treat k as small constant, indicating that the influence of the third-order gradient is minimal. This suggests that the overall shape of the loss landscape is predominantly governed by the first and second-order information.
The error term in loss term satisfies that |eL| < T17202 p(C'doo"/>)" //)//////////22000/0/ + Cn' + 8(2M - do < (d 2)mo² with C 200 log(64-TT/8 /8) and the error can be decomposed into three parts: (1) 77 2²= + p(Cnoo//)) /2)3//))/))) 3//2 are higher order discretization effects of learning rate 77; (2) Cr' dno is caused by the change of the Hessian in the valley dimensions and will diminish when K is small; (3) 5(2M + =jodd) accounts for the small chances that the iterate will escape the neighborhood of the river due to the stochastic updates. While the theorem only considers the case where 2a is constant vector,
We will first show that under Assumption the loss is separable within U. Lemma C.27. Under Assumptions To the river a straight line parallel va-
Proof. In this case, the x in Assumption 2 is and this is direct corollary of of Lemma C.12
Lemma C.28. Under Assumptions T there exists functions and ht such that for any w E U satisfying that B(w, 2A U, holds that
L(w) = (()(w) ((W))) + h(w - (
Furthermore, h is ry-strongly convex when constrained the range of Ps-
37

---

Published as a conference paper at ICLR 2025
Proof. We will choose g the constraint of L on M. Now ((w) will always fall in the range of Ps. Consider any y in the range of Ps and as Vua(w)[o] = 0, we have that
y V2 L(w)eu 0.
This then suggest that
V VVVL(w), O.
We then have for any E M, by Lemma C.27
/ ((w (w))), VL(((w) + (ww))
L(w) - (ww)-((Quu))
= / ((w ((w)), VL(a + -(
- (
We will then define liw - -((w)) = ((w ((w)), (( dr. and this concludes the proof.
Now, as h(w (ww) L(w) I(((w)))) <hy when constrained on the range of Ps has an
eigenvalue greater than
We will first consider the mixing dynamics of the current SGD iterates on a strongly convex loss with minimizer O.
ylk + y4 17Vh(yk) 19k. (00) 0, gk N((0,
(23)
We will define coupling process Vk ijk + 1) ÿs 17 H yik, w() = 0, ~ NN(0 = 0.
(24)
Here H (0) positive definite.
Assumption 6 (Regularity of h). We will assume the following for the function h, constant E (0, 1].
learning rate 13.
1. The smallest eigenvalue of Vihly) within B(0,r) is at least 7 0 and the largest eigenvalue for
H is at most "fmax-
Vy, h(y) [0, M).
3. Vy E B(0,r), r), lV s & <A,lV S
T > 1/7.
5. < 1/(27max).
np"a²= (79/(1600llg((77/88)
7. 10 Via V7 /dlog(88TT/D
Sr.
We will first show that HL will be bounded with a high probability for T/n steps.
Lemma C.29. For any 8 (0, 1). with probability 1 for ün defined in Equation 24 under As-
sumption 6 it holds that for any k < T/m.
10 no (dlog(87T/8).
10 12 <
Proof. For integer K [YT We first have that for k < K
Ver+11 = (1 1-77 1 1 M(k 1) 1 + 1
(1 - my) [1b 1- -(--1))
1 1+
Denote 9s = 11 " E (1 m77)
1 - then g a normal vector with variance
1)
n
I no²
17
(1 my
)
20 1
22II <
2y m72
38

---

Published as a conference paper at ICLR 2025
Further, denote YE Vel 1 and By (1 -11-)77) ///
1 then
Ys ey x + 9k-1
Σ
9k-i
isk-1
Then each variable YE is also Gaussian variable with variance smaller than
Σ
Blgk-
1
no²
I I.
I
2no²
1/e²
7
7
Hence, by Lemma C.37 for each k, it holds that
2/
P( Yx >
Vdlog(4K/0)) < 6/2K.
V7
Using union bound,
2/1
P(ak < K, Yx
//log(4K/0) < 8/2.
We now proceed to bound the distance of yk compared with close Ys. without loss of generality, considering je will define process called my satisfying that
(1-my m7)
1 9k.
m&
Then my is martingale and each my is Gaussian vector. In particular, m 91. This further suggests that || || is super martingale
E Blllme m []
m
By Doob's lemma (Lemma C.38
P( sup
| C²) S PP( sup exp(Allm4413) exp(AC²)) 1
k
s Elexp(A| m
112 AC
Elexp 1 XC
Following the same line of proof as Lemma C.37 7
we have that
2 2no
P( sup 1
|| 2|2 >
v Vdlog(1K/o)) S 6/2K
Vr
We further note that yk
Yo s (1 - m77
m& I 4mg We have that for any & <
8/10 Vdlog(4K/0)) S 88/2K Vx
P( sup
Yoll2 >
Combining with the bound Ys. have that
10/770
P sup bkl
Vdlog(8+T/0)) <
USIST
V7
The proof is then complete.
The following lemma states that y/ and in close with high probability.
39

---

Published as a conference paper at ICLR 2025
Lemma C.30. Assume function h(y) is y-sstrong convex in B(0, r) and has minimizer at 0,then for E (0, 1), under Assumption it holds that with probability 6,
Vk T/n, 10x
yull2
400mpo² (oghnnoodogg(8TT/00/ Uu E ()0,r), Un E ()),))
Proof. By Lemma C.29 with probability 1-6,
100 dloo(8yyT/6)
vk T/n. 1
We will use C as a shorthand for 100d (8g(()/ Under such scenario, define u Ün Vk- we will prove by induction for s T/n that
|aall2 & C/pp yy ((.,r).
(25)
Clearly V/ = 0, satisfies the induction hypothesis. Assuming Equation (25 hold for t, then
y(k + 1) yv 17 (())) - 79k
= yk - -) + - 79k
Uz - (11)) - -99 - 12 - L) + Ex.
Here ||ex|| = -| (v()) - -, 7) that
S 11 < 2nd( Uu 11 + (1)). Hence we have
11x+1ll2 (1- 2mp + || 112)
As || S 2 72
S y,
Ap. we have that 2pl 12 = m 1111/22
Hence
| 212 S (1-17/2)||px/)) + 20p0
7
= (1 1//r/)/lluulll> + 122p0"
7
By induction ||va|22 S Anpo2C//2 It is then easy check || S 4npo
The following lemma tracks the changes of E[(hyy))
Lemma C.31. Assume function h(y) is y-strong convex in B(0, r) and has minimizer at 0,then for E (0, 1), denote 100 log(8yTT (8) as C, under Assumption 6 holds that E (L/r:TT/n)
ApC 72
dno² + (dno"C)
3/2
E[M(jo)] por 2/2 < (H) +
+ 25M + 8n02 d/2
9
Proof. By Lemma C.30 with probability 1-
1i yul12 I 2 yn E B),,r), In E B()
Define this event as E
Hence
E[h(yao - [[)))
E[h(Tn) h(y))) &()((E) + E(h(ūe) - h(y))) (TT)(
ApC <4
72 ?do² + 8M.
40

---

Published as a conference paper at ICLR 2025
For |lvll r, it holds that
(I((y) y² (7 S pluv11
By Lemma C.29 with probability - 6. for 1 < //axx holds that for any S T/n.
dno C
10 -
< <t?
Define this event Eg-
Denote H V²)(0), we have that,
[[h(ū=)] B((i)) )" H
(((()n)) E2])((E) - B((ūx) ) H + ((M(T)) ES)((E))
<ÓM + E[h(ūk) Eg]P((2) B((bk) HiN
<6M + (nnd
3/2
(Uk) "H Eg]
Combining the both and we have that
S 4 ApC 7 ndoz +p mddoo)
3/2
E[h(ww)) - (ūx)) Huk]
+ 25M + (((O)) Hús Eg]
Here the covariance of Ts. denoted as Dg satisfies that
= > o I
Therefore
k - - 11211 - = 10 ー
2i = 0211 - 1 (I 15--111 2k
Hence assuming the cigenvalues of H is "....., Vd
d
Tl Tr(2+H])= no²
Σ
1 (1 (1 )2))
2 177
1=1
When 1 2 1 n5.
m7 ,
1/2, it holds that
no² 2 - 77%
(1 (1 m2))22k) noo /2
(1
- (177) )2)
=no?
oali /22
2 1/9i
1
1
2no²
2 17's 2
Sn² 22 02/11/2.
Hence
B( T Hūu)
< &TH(H))?2222
Further, let wy = Ex Nk. under Eg, we have that
Ju 12 2 Amin Amin(27 11ūi 12 Amin ((211 n21 (I- (1)
11 do?C.
41

---

Published as a conference paper at ICLR 2025
As Up is isometric Gaussian, E((b))* ((ok)"H' His
E5] - E Elu [ 1/22/2
" H k
1/2
ui 1 2 do²c]
E | 1
|| 2 do C
Elug [>x 1/2
H
12 ug
114 12
Bllluull
2
do2C)
$ dno
Bl
E
Plugging in the density function of ||u ll. we have that
E
uk
24
2
c
Svaco
, e- 1 (20") dr
E 1
do adtle-r21(200
Let V 2+1" r. then
p+7 p 102 dr = d+1 1-2 pd+1 e (19(44112211)
Valaa
d
p ld+1 (220)) -('')
<
vdCo
S 4e C/2 (0 pdt1 - (200()))) dr.
Hence, we have that
((() Hd Hi
Eg] < 4e 0/2 dno" ddno0 12.
Putting together, have that,
ApC
722 dno² + d
3/2
E[h(ūu)] - nodd/2 s 1 +
+ 26M + 8no2
The proof is then complete.
We will now state the complete version of Theorem 3.4 Assumption 7 (Sufficient Small Learning Rate). We will assume the following for constant 5 E (0, and learning rate m
1. < 1/(2
17 S 10/ 011660002222
3. 10 10 /00
VT /dlog(8)T/5))
sr.
Theorem C.32 (Complete version of Theorem 3.4. If a loss L is a river valley (Definition 3.1 and satisfies Assumptions 3 for any constants å E (0, 1) and T > 1/7, for sufficiently small
learning rate satisfying Assumption 7 the iterate defined in Equation (4) with T}s = satisfies that for any integer t E [1/m7, T/n), there exists T satisfying that,
(T(())) ((x(()))((((())) = (d 1)m0² EI.
where 6, Amyiaa and | o2 + p ( + < (d-
with C' 200 log(ll4gg000 (64-/T/0).
Proof. By Lemma C.28
we can write
L(w) = (-ww((wwww))))) + L(((w)).
Hence we can separate the dynamics of Equation 4 into two parts, namely w: = + (w (ww)) It is easy check that when constrained range of Ps. hyy) satisfies Assumption 6 Hence, can use Lemma C.31 to control h(w, (hww)) ))) For (ww)), the iterates is running a gradient descent
with learning rate on M and we can use proof analogous to the proof of Theorem 3.3 to show that if ((w)) x(f m)), then there exists To. such that
T(t, n) E [To + (1 417maa )tt, To (1 + 4symma))t)]]
This completes the proof.
42

---

Published as a conference paper at ICLR 2025
C.8 PROOF OF THEOREM 3.5
We will first state the complete version of Theorem 3.5
Theorem C.33. Under the setting of Theorem C.32 the SGD iterates (defined in Equation +
with the decaying learning rate schedule satisfies that for any integer E [ts, 1.16a], there exists T [[1 (t) (t)) satisfying that,
E[L (xo(t)) ((x(T)) (d 1) 2 +2+
with T(t) - + Σ nK-
Proof. The proof analogous to Theorem C.32 and Lemma C.31
We will omit the detail derivation
and only focus on deriving the variance of corresponding Uk
(k + 1) = ok M Tkgk, (0) = 0, gk ~ 0(0,02I).000 O.
Here the covariance of Tbk, denoted Eg satisfies that
Ex+1 = (I - nh EB + 0² I.I
we consider i-th eigenvector of H as U; and denote Tk.s = v/ Eve.
Analogous to proof of Theorem C.32 Tku.
no
2n20²
We further have that
1/2 (2 12+rmp r
11
71) 7))
20 tr,i to²
2+ rry
Then by induction, we can prove that for r 20 0
0² 17
212202
a²
4n202
ne, tr+1
7/ 2 2 rmy
%
%
The rest follows the proof of Theorem C.32
C.9 PROOF OF LEMMA AND THEOREM C.35
In this section, we will denote
Σy= (x,( Qi,
(xp(
We will study this loss
m
Paus log exp(O..)
L( (O..), ((O).)) -- -
(26)
Σ exp(Or..)
j=1
Lemma C.34. The loss defined L in Equation (26) satisfies that
(VI(O)) Pi.j Qi.--
(VLL())) 1(i (0)==')((0.. (') =j' - 2.s))
Proof. The loss satisfies that
m
L(
Ps O(.) - logg P.j.).
Hence,
(VL())) = P Qi,5
Taking differentiation for another time yields the desired result.
43

---

Published as a conference paper at ICLR 2025
Proof of Lemma 4.1 This can be done by directly summing diagonal entries in Lemma C.34
Assumption 8. We will assume there exists constant " and positive integer n' n such that P satisfies the following assumption,
For any S n', Ps, > 8y. 2. For any > n', there exists Js. Piy>1-9
A Assumption 9. We assume the existence of "generalized river", which is p-dimensional manifold M such that any point w E M has a gradient L(w) lies in the eigenspace spanned by the last eigenvectors' direction of the Hessian, {u, (V²L(W)) i (d -dd-p--- p--pp12 1.,1
Theorem C.35. Under Assumption 8 generalized river with dimension n'm + (n - r(') exists in the loss landseape defined by L in Equation 26).
Proof. According to Lemma C.34 the Hessian for L is block-diagonal. Now fixing a city we will analyze the eigenvalue distribution in this block. Let q [f then this block is diag(1) qa².
For all non-zero eigenvalue for this block, there exists u such that diag(g)v 9² uq Ao.
Hence, we have that
4 0
= 4j
m This implies > j=1 d5-c ing to À O. For the rest nonzero eigenvalue, we have that À > minq.
1. We then have À 2 0 and there exists only one eigenvector correspond-
Now we consider the manifold M defined as
= (0 vi vi Q., = P n', Qi,, >1-7).
Then for all θ E M, we have that the gradient is zero for all dimensions (i.j) with i < n'. Further, we know the nonzero eigenvalues for these dimensions are at least 8y by Assumption For the rest of dimensions (i,j) with i > n'. by Lemma 4.1 the largest eigenvalue is bounded by (1 - 7) < 2yy This shows that the gradient falls in the eigenspace spanned by the last n'm + (n- n²) eigenvectors, which concludes the proof.
C.10 TECHNICAL LEMMA
Lemma C.36. If function F(t) satisfies that dF((t)
dt < AF((t)
then F(t) < c
Proof of Lemma C.30
Consider G(() F(I)eAt then dG(() = At (F(()() + AC()(( 50.
Hence G(t) < G(0).
Lemma C.37. If random vector g NN(0 Σ) and (0, then it holds that
P([la]2 2 2, Tr() <
44

---

Published as a conference paper at ICLR 2025
Assume Σ QA?QT with Q being orthonormal matrix and A being diagonal with diagonal X; for E [d], further let g being a standard gaussian random vector, then follows the same distribution as that of AQ" g', which is further identical to Ag'.
Proof of Lemma C.37
P(1ll 2 C) A( 2 C) P
A 2 C²) S E(exp(t Σ \2 1C
well known that the moment-generating function of (g/1² is
Elexp((A7(0)*})
2
Hence P(ll9112 2 C) s 10a I 2-22
v /1-2-TYY5)
With t |MM it holds that P( 1 2 C) S 2e
18 This concludes the proof.
Lemma C.38 (Doob's Inequality). Let X1 1..... Xn positive submartingale adapted filtration F Fi..... Fr which means X; S X[ Fi), then
Xx
P(sup X; C) 5 is
D OMITTED EXPERIMENTS DETAILS
1.862
2.5
WSD-s
WSD-S
Decay at 23000 Steps
Decay at 23500 Stops
Los5
1.860
2.201
2.3
Locs sp
1.858 12% 10% 8%
6%
4%
14% 12% 10% 8%
65
244000
Decay Phase / Total Steps
Decay Phase / Total Steps
(a) 1.2B Models 200B to- (b) 0.1B Models 100B kens
(c) O.IB Models 100B Tokens. Decay
Near a Loss Spike (6%)
Figure 15: Ablation Study the Sensitivity of Fraction of Time Decaying. This study examines two settings: smaller scale with O.IB parameters trained on 100B tokens (middle figure) and larger scale with 0.6B parameters trained on 200B tokens (left figure). The results indicate that the final performance is similar when the decay phase is 8%-12% of the total training steps. However, the right figure demonstrates a significant performance loss when decaying near loss spike. It compares two training loss curves with decay phases of 8% and 6% of the total compute the O.IB models, where the latter starts immediately after loss spike, leading a validation loss increase 2e-2.
We train LLaMA models with parameter sizes using the Levanter framework for our study on WSD-S. For our theoretical study, we pretrain 124M GPT-2 using the nanoGPT framework with learning rate 6e-4 and train it with a batch size 0.5M for 100k steps with warmup steps of 2k.
We hereby provide all the hyperparameters we used for the LLaMA and GPT-2 models training.
Model
Hidden Dim
Intermediate Dim
Num Layers
Num Heads
Peak LR
O.1B LLaMa
768
3072
12
12
6e-4
0.3B LLaMa
1024
2048
24
16
6e-4
0.6B LLaMa
1536
6144
24
32
4e-4
1.2B LLaMa
2048
8096
16
32
4e-4
O. GPT-2
768
3072
12
6e-4
Table 2: Specifications for Different Sizes of LLaMa Models
45

---

Published as a conference paper at ICLR 2025
We decay the model for the last 10% of the training runs with one exception for 0.3B model using WSD method near 25k steps to avoid loss spikes. We outline the decaying and resuming point (the unit is Ik steps) choose here:
Model 0.1B LLaMa
1st Decay Starts/Resume 2nd Decay Starts/Resume
3rd Decay Starts
11.25 12.5
22.5 / 25
48.75/ 53.75
0.3B LLaMa
11.25 12.5
22.5 / 25
48.75/ 53.75
0.6B LLaMa
11,25 12.5
22.5 / 25
48.75/ 53.75
1.2B LLaMa
11.25 /12.5
22.5 / 25
48.75/ 53.75
Table 3: Specifications for Decaying Steps for WSD-S Method
Model
1st Decay Starts/Ends
2nd Decay Starts/Ends
3rd Decay Starts
Total Steps
0.1B LLaMa
11.25 12.5
22.5 25
45/ 50
53.75
0.3B LLaMa
11.25 12.5
22 25
45/ 50
54
0.6B LLaMa
11.25 12.5
22.5 25
45/ 50
53.75
1.2B LLaMa
11.25 /12.5
22.5 25
45/ 50
53.75
Table 4: Specifications for Decaying Steps for WSD Method
46


Reviews:
[{"review_id": "Jhv3pNw3DZ", "reviewer": "Reviewer_Cgph", "summary": "This work proposes a perspective on pretraining loss landscape drawing a metaphor to a river valley. The authors run toy experiments in support of this persepective, and present theoretical analysis as well. The authors present WD-S, a variant of Warmup-Stable-Decay, which they favorably compare against WSD and cosine-derived learning rate schedules across a number of model sizes.", "questions": "Nit:\nFig 7 occludes the caption of Fig 8. Please fix. Fig 7 captions includes 'textbf'", "limitations": "", "rating": 6, "confidence": 2, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "The paper is well-written, with clear prose and clean figures. The elucidation of the river-valley loss landscape is intriguing and intuitive. The empirical results are good, though they could cover a broader range of experiment setups. The work is of potential interest to model-training practitioners who may desire to train large models without pre-specifying a compute budget.", "weaknesses": "The related works is deferred to the appendix. This makes it difficult to contextualize the work presented. If space is a concern, please still present a truncated related work section and defer a more extended discussion to the appendix. Of course the optimization literature is large but the most relevant works need to be discussed in the main text to provide context.\n\nThe originality of the technical contribution of this extension of WDS is marginal, though the empirical results suggest it is a useful one.\n\nThe loss curves presented in Fig 9 are very spiky, and the field of view is highly zoomed in. These should be re-run with a larger number of training runs and averaged, otherwise the result seems plausibly artifactual. Same for Figures 10 and 11.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "bgGhwN54w6", "reviewer": "Reviewer_UHrP", "summary": "This paper compares the optimization space of language models to a river, thereby demonstrating the effectiveness of the WSD method. Based on this method, a simplified version, WSD-S, is proposed, which reduces computational complexity by reusing checkpoint weights without compromising performance.", "questions": "1. Could the authors provide more evidence that the optimization space might be a river? \n2. Regarding the valley optimization space, do the authors have any ideas for developing more optimizers? For instance, how can more steps be focused on progressing along the river rather than crossing the river? Does Adam offer better descent speed compared to SGD in such a river valley? Additionally, I did not find the optimizer used by the authors mentioned, which should be specified in section 4.1.", "limitations": "", "rating": 6, "confidence": 3, "soundness": 2, "presentation": 3, "contribution": 2, "strengths": "1. The optimization algorithms supported by theoretical foundations are a crucial research direction.\n2. Despite involving complex theoretical proofs, the writing of the article remains clear and easy to understand. \n3. The proposed WSD-S performs no worse than WSD.", "weaknesses": "1. The theoretical proof in the paper relies on many assumptions, and the validity of these assumptions in complex large models is difficult to ascertain. For example, the paper assumes that the loss function is analytic, which seems unreliable in models like GPT2 that are filled with GELU functions. \n2. The paper assumes that the optimization space of the loss function is a river and uses a toy model from Allen-Zhu's paper to verify this. However, neither the model assumptions in Zhu's paper nor the toy model in this paper are reflective of natural language situations, making it somewhat far-fetched to use such models to verify the existence of a river. \n3. Compared to WSD, WSD-S does not offer any substantial innovative improvements, neither in terms of computational complexity savings nor in the final performance.\n\nMinor comments：\n\nThe excessive length of the paper results in cramped formatting and sections. Especially in Figure 7.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "f1gwWKOkgd", "reviewer": "Reviewer_A9Vb", "summary": "In this paper, the authors explore a learning rate schedule named WSD-S, an enhancement of the previously introduced WSD model, which consists of three phases: Warm-up, Stable, and Decay. They provide a robust theoretical and empirical analysis supporting the effectiveness of WSD and propose its simplification in the new WSD-S variant.\n\nThe theoretical analysis presented in the study sets the mathematical definition of what the authors refer to as the \"river valley\" loss landscape. They demonstrate (Theorems 2.2 and 2.3) that in such landscapes, a higher learning rate yields greater progress over the same number of gradient descent steps. For stochastic gradient descent scenarios, while progress along the 'river' remains constant, an additional 'hill' component of loss emerges. The rapidly decaying learning rate phase of WSD effectively addresses this hill component (Theorem 2.4), suggesting that WSD is particularly beneficial for river valley-type losses. The authors further argue that language modeling pre-training losses resemble these river valley losses, thereby explaining the previously observed advantages of WSD.\n\nFor continual training scenarios, the authors highlight the importance of the transition from stable to decay phases. WSD continues the the training from the pre-decay checkpoint but authors argue that the progress made during the rapid decay phase is crucial and needs to be carried over. With this hypothesis, they present a simplified WSD where the same checkpoint (instead of the pre-decay checkpoint) after the decay is continued to be trained. This approach is operationalized in the WSD-S schedule, which the empirical analysis suggests outperforms other schedules like cosine, cyclic-cosine, and WSD.", "questions": "Minor edits:\nLine 290: not sure if the eta_t is red on purpose\nLine 399: missed a \\ for textbf\nLine 491: Figure 10 is difficult to read\n\n\nQuestions:\nLine 236: What does it mean by “starting point w lies on the course of the river”? Does it mean that the starting point is on the manifold M or in some neighborhood of M?\n\nFigure 8: Right bottom figure: The decay period is extended as the training progresses. If yes, why? Were there any experiments done to observe the effect of minimum LR (after the decay)? \n\nI wonder, if more decay phases can consistently keep the hill component smaller and ultimately make more progress?  Or is there any sweet spot of number decays that helps WSD-S to achieve the best results?", "limitations": "", "rating": 10, "confidence": 3, "soundness": 4, "presentation": 4, "contribution": 4, "strengths": "The paper is excellently written for the complex nature of the analysis presented in the study. The information is organized well and easy to follow. \n\nThe authors have provided solid theoretical and empirical analysis to explain the effectiveness of the WSD and how simplification further benefits the optimization process.", "weaknesses": "None", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "jS5GNyi7d4", "reviewer": "Reviewer_H6GK", "summary": "The purpose of the paper is two-fold, 1) understand the success of WSD by postulating that the loss plane is analogous to a river-valley. 2) They attempt to provide more theoretical justification and empirical support to this hypothesis, resulting in optimizing WSD to WSD-S under their assumptions.", "questions": "Have the authors examined different architectures, datasets, learning rates, or attempted to examine other possible explanations for the loss curve region?\n\nDid you evaluate the models on benchmarks post-training?\n\nProvided the hypothesis, why do we have divergence randomly during training of large models on stable plains?", "limitations": "", "rating": 3, "confidence": 3, "soundness": 2, "presentation": 3, "contribution": 2, "strengths": "They clearly articulate their idea of a river loss landscape and provide some theoretical guidance to establish it. They provide some benchmarks with various smaller models and evaluate the loss trajectories. They also provide an attempt to try and leverage their knowledge to propose WSD-S.", "weaknesses": "The main weaknesses observed are the following:\n1. The assumption of a river valley isn't guaranteed and is not strongly supported. Additional theoretical or experimental results supporting this conclusion are necessary.\n2. Results are demonstrated on very small models, with identical architectures, identical batch-sizes, and relatively small corresponding datasets. All of which could influence the loss region.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "UKsUMdFnM1", "reviewer": "Reviewer_bscX", "summary": "This paper is composed of three main sections.  In the first section, they derive some analytical results for a specific type of loss landscape, which they call the “river valley landscape.”  The river valley landscape corresponds to a kind of loss function where the loss is reduced by moving in the direction of lowest curvature (the river direction), and there are steep hills on both sides of the river.  In such a situation, with a number of assumptions, it can be shown for gradient flow (continuous time), gradient descent (full batch gradients) and stochastic gradient descent (with approximate gradients), that the parameters will follow the river within a given bound.  Moreover, in the stochastic case, the loss is also bound by a term related to the variance of the gradients multiplied by the learning rate (which they associate with movement of the parameters up the hills).  A Warmup-Stable-Decay (WSD)-style schedule that has a decay at the end can, in theory, move along the river during the constant phases, and subsequently move down the hill directions during LR annealing.  The next section describes a bigram model where different initial tokens have different amounts of uncertainty in their following token.  The degree of uncertainty is shown (at an optimum) to correspond to the degree of curvature in the loss, suggesting that the “river” direction corresponds to learning more deterministic relationships, while the hill direction corresponds to more uncertain predictions.  Decayed LRs are shown to do relatively better on the uncertain bigrams in a small experiment.  Finally, the paper considers the case where multiple checkpoints are desired at intervals during training at a given model scale (e.g., to make a data-size-based scaling law).  If using the WSD schedule, the decaying happens in separate branches off the main sequence of checkpoints.  The paper makes the observation that even during these decaying phase, progress is made along the river direction.  This motivates the WSD-Simplified schedule (WSD-S), where rather than moving back to the pre-decay checkpoint in the main branch, we continue training directly after the decay phase.  This amounts to running cycles of WSD-style training.  Some experimental results dig into whether WSD-S is effective compared to WSD and to traditional schedules like Cosine decay or cyclic Cosine decay.", "questions": "Why doesn’t the blue line go as far as the red line in the “0.3B Parameters” plot in Figure 10?  The red line seems to stop short of 50,000.\n\nThe introduction says, “WSD-S leads to a better validation loss than WSD under the same compute budgets due to the re-use of the decay period”.  Is Figure 10 validation loss? (it doesn’t say).  For the 0.6B parameter model trained to 50,000 steps, don’t they obtain basically the same loss at the end?  And by definition, they always have the same loss on the first checkpoint, right?  I mean, with the loss spikes seeming to affect the intermediate results, can we really conclude that WSD-S gets better validation loss under the same compute budgets?  Do we need to soften this statement in the introduction?\n\nDo you think WSD-S, since it has a slightly lower summed LR than WSD (lower total LR “area”) might actually do better than WSD if we increased the peak LR slightly?  Do you think Cyclic-Cosine could also improve if we increased the peak LR?", "limitations": "", "rating": 5, "confidence": 4, "soundness": 2, "presentation": 2, "contribution": 2, "strengths": "The WSD learning rate schedule has recently drawn a lot of attention in LLM training.  WSD allows us to obtain high-quality intermediate models by decaying in “branches” from checkpoints in the main training process.  Perhaps the key insight of the paper is that if using WSD, we must acknowledge that the decay portions are “wasting” a significant amount of compute in that the compute spent during these decay portions does not contribute toward the training of future checkpoints.  Prior work has recommended up to 20% of total steps in decaying, so decaying 5 times (e.g., to obtain a scaling law) could double the total compute spent training a model.  We can think of the proposed WSD-S approach as a philosophy: if you do decay, you should resume from *after* the decay, rather than from before it.  In this way, we seem to have the same amount of flexibility as WSD (with the exception that with WSD-S, we can’t run the decay in parallel with continuing training in the main branch, as you can with WSD).  Testing this philosophy is well-motivated, and of solid interest to the community.\n\nThe idea of specifying a loss landscape of a certain type (river-valley), and deriving different bounds on optimization given that landscape, and then determining whether different real-world training scenarios fit such landscapes, also seems well-motivated.\nIn terms of clarity, I found the figures fairly helpful in most cases.  For work like this, pictures of the LR schedules can tell the whole story and so Figure 2 was good, and the little arrows in Figure 2(b) (and even littler ones in Figure 10) were very helpful.  Figure 3 was helpful as well.  I also realized that after reading the paper, I wasn’t really confused about anything, so the work must have been explained somewhat clearly.\n\nI thought the material on decay after loss spikes was interesting and fairly important.", "weaknesses": "The paper in its current form seems incomplete in many ways.  For one, the paper has not undergone a proper proofreading/editing/review process.  This makes it much harder for the reviewers to assess the quality of the work – I felt like I was being asked to help prepare an early draft of a paper written by my co-author, rather than reviewing a paper that is already ready for submission.  Some typos are probably just honest mistakes, e.g., repeatedly spelling “decaying” as “decayping” – but if you pasted the paper text into ChatGPT, or ran a simple spell checker, it would find many of these typos!  It’s unsettling to see typos in the abstract (!) “one can branch out from the main branch at a proper at any time”, in main theorems “with the learning decaying learning rate schedule”, in Figure captions ‘Figure 7: textbfReproducing [sic] the Nontraditional Loss Curve.”  There are missing LaTeX references (I assume) “We will first motivate the decaying function we choose ?? [sic] using a simple example”.  Figure 7 is blocking part of the caption of Figure 6.  I feel like all of these would have been picked up if a co-author had simply reviewed the paper prior to submission (and writing could have been improved in general with such review).\n\nMoreover, the paper has a very non-traditional style: the main body of the paper is missing Related Work, Discussion, and Conclusion sections (although there is Related Work in the appendix, which reviewers are not required to read).  This makes it difficult to contextualize these results compared to prior work (missing the Related Work), and understand the key insights that we draw from the experiments and theoretical sections (missing Discussion/Conclusion).  For example, does this paper present a negative result?  I.e., is WSD just as good as WSD-S, even if we count the decay portions as part of the overall compute?  This seems to be what Figure 10 is showing.  It would be interesting to learn (from a conclusion) whether the authors feel the same way.\n\nThere is also some missing discussion of related work.  E.g., consider the following passage: “We find that the loss interpolation between parameters before and after each training iteration’s update is roughly convex with a minimum (valley floor) in between for most of the training. Based on this and other metrics, we deduce that for most of the training update steps, SGD moves in valley like regions of the loss surface by jumping from one valley wall to another at a height above the valley floor. This ’bouncing between walls at a height’ mechanism helps SGD traverse larger distance … this exploration above the valley floor allows SGD to quickly travel far away from the initialization point”.  Actually, this passage is not from the current paper, but from Xing et al., “A walk with SGD” (2018) (arXiv: 1802.08770v4), which is not cited here.  Contextualizing the findings of Xing et al (and other works) within the main body of the paper is essential.\n\nI think another weakness is that, of the three main sections of the paper, none of them really went into enough depth to convince me of the main arguments.\n\nFor the theoretical section, it’s not clear, of all the many assumptions, which apply to modern LLM training.  Like, is the 4x “eigengap” commonly found?  How important is it?  Also, what is the point of departure here from other theoretical work in SGD, for example from Bottou et al, “Optimization methods for large-scale machine learning”. SIAM review, 60(2):223–311, 2018?  The benefits of decaying LRs has a long history in optimization, and decay has previosly been motivated by similar considerations to those used here (moving from initial conditions early, reducing stochastic noise later on).  If you frame your findings in terms of this earlier work, the contrast can help us understand the benefits of your analysis.  Like, maybe the stochastic noise plays a greater role if the loss landscape has a river-valley form (motivating greater decay than we would use otherwise), and we can see this in contrast to what we’d see without all the river valley assumptions.  These other works often derive a bound from the global optimum, but here we seem to be bounded from a point x “further down the river” – what is the significance of this?  This could indeed be a paper all in itself.\n\nMoreover, I am not convinced by the probes in Figure 5, given that the two points are sampled 5 *billion* tokens apart.  Does this mean that the river has no bends?  Like, if you interpolated between two points in Figure 3, you can clearly see that you would typically travel up the hill and back down the other side.  Does this invalidate the river analogy?  If not, how does Figure 5 “validate” it?  Figure 2(a) says the iterate will “oscillate” between the hillsides – why not check this over fewer than 5 billion tokens?  Plotting Figure 5(a) at a single pair of checkpoints is not sufficiently rigorous, from my perspective.\n\nThe empirical verification in Section 3 lacks a lot of details – “we first train a toy model” – but what is this model?  A neural network?  Next we fine-tune a pre-trained GPT2 model, but of what size?  Why is it pre-trained?  What are the objectives of doing this?  There are no ablations here, nothing is varied and repeated, there are no error bars.  Basically, it’s not clear why the design choices were made for these experiments, and what the objectives were, and what the implications of these results really are.\n\nFor Section 4, my main objection is that I suspect the average learning rate to be a confounder.  Think of your Theorem 5: progress throughout training is partly controlled by the sum of the step sizes (this theorem is for decay, but it’s true during the stable phase as well).  Then look at Figure 9: WSD-S has a higher average LR than Cyclic-Cosine.  Would cyclic cosine work better than WSD-S if we simply increased the LR slightly?  The paper says, “We hypothesize that a model trained with a small learning rate for too long, as with Cosine, is implicitly hurt compared to a model trained with a large learning rate for the majority of the run, as with WSD or WSD-S.”  Exactly!  So why not train Cosine with a higher LR?", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (Poster)