Training Dynamics of the Cooldown Stage in Warmup-Stable-Decay Learning Rate Scheduler
Authors: Aleksandr Dremov
Alexander Hägele
Atli Kosson
Martin Jaggi
Venue: Accepted by TMLR ()

Abstract: Learning rate scheduling is essential in transformer training, where the final annealing plays a crucial role in getting the best performance. However, the mechanisms behind this cooldown phase, with its characteristic drop in loss, remain poorly understood. To address this, we provide a comprehensive analysis focusing solely on the cooldown phase in the Warmup-Stable-Decay (WSD) learning rate scheduler. Our analysis reveals that different cooldown shapes reveal a fundamental bias-variance trade-off in the resulting models, with shapes that balance exploration and exploitation consistently outperforming alternatives. Similarly, we find substantial performance variations — comparable to those from cooldown shape selection — when tuning AdamW hyperparameters. Notably, we observe consistent improvements with higher values of $\beta_2$ during cooldown. From a loss landscape perspective, we provide visualizations of the landscape during cooldown, supporting the river valley loss perspective empirically. These findings offer practical recommendations for configuring the WSD scheduler in transformer training, emphasizing the importance of optimizing the cooldown phase alongside traditional hyperparameter tuning.

Published in Transactions on Machine Learning Research (08/2025)
Training Dynamics of the Cooldown Stage in Warmup-Stable-Decay Learning Rate Scheduler
Aleksandr Dremov, Alexander Hägele, Atli Kosson, Martin Jaggi EPFL, Switzerland
Reviewed on OpenReview:
https: https:///penrevier. net/ forum? id= 2nSYEc2od3
Abstract
Learning rate scheduliigg is essential in transformer training, where the final annealing plays a crucial role in getting the best performance. However, the mechanisms behind this cooldown phase, with its characteristic drop in loss, remain poorly understood. To address this, we provide a comprehensive analysis focusing solely on the cooldown phase in the Warmup- Stable-Decay (WSD) learning rate scheduler. Our analysis reveals that different cooldown shapes reveal a fundamental bias-variance trade-off in the resulting models, with shapes that balance exploration and exploitation consistently outperforming alternatives. Similarly, we find substantial performance variations comparable to those from cooldown shape selec- tion - when tuning AdamW hyperparameters. Notably, we observe consistent improvements with higher values of B2 during cooldown. From a loss landscape perspective, we provide visualizations of the landscape during cooldown, supporting the river valley loss perspective empirically. These findings offer practical recommendations for configuring the WSD sched- uler in transformer training, emphasizing the importance of optimizing the cooldown phase alongside traditional hyperparameter tuning.
1 Introduction
Learning rate scheduling remains critical component in training transformer-basee language models. Since the introduction of transformers by Vaswani et al. (2017). learning rate scheduling has become standard practice in transformer training pipelines and remains prevalent in recent works such as Meta Al (2024) and DeepSeck-AI ct. all 2024). Among various proposed strategies, cosine decay (Loshchilov Hutter 2017 remains one of the most commonly employed methods.
An alternative approach, the Warraup-Stable--eee (WSD) scheduler, has recently gained popularity. This method features three phases: an initial warmup phase with linear learning rate increase, a stable phase maintaining constant rates, and final decay phase reducing rates to zero. Hägele et al. (2024); Hu let al. (2024) demonstrated that WSD achieves performance comparable to, or better than, cosine decay in wransformer-basee language modeling. Importantly, it allows training without a predefined training length and enables continuation from any point in the stable phase.
Notably, WSD-trained models exhibit loss curves that remain relatively high during the stable phase but drop sharply during cooldown. Previous work has provided extensive analysis of WSD parameters and their impact on final performance (Hägele al.] 2024) well as the role of the warmup phase (Kosson al. 2024b). However, the cooldown phase dynamics remain under-explored.
This work aims to address this gap by investigating the behavior and impact of the cooldown stage in the WSD learning rate scheduler, providing insights into its effects. We propose bias-variance framework for cooldown learning rate scheduling shapes, through which we explain why some shapes tend to perform better than others. The schematics for the proposed framework are shown in Figure [ The study examines hyperparameter effects, including AdamW's beta parameters, weight decay, and batch size configurations, with experimental results verifying alignment between observed outcomes and the proposed framework.
1

---

Published in Transactions Machine Learning Research (08/2025)
Experimental bias-variance
Schematic loss surface
for cooldown shapes
contribution, [10- 2
square
Variance loss
2
mirror cosine
linear
X
var + bias min
0.5
0.3
0.1
5
7
8
optimum
x trade-off
Bias loss contribution, [10-2]
high variance,
low variance,
X low bias
X high bias
Figure 1: Conceptual visualization of the relationship between different learning rate cooldown shapes and their bias and variance. On the left, we show a conceptual visualization of loss surface points in two dimensions. Colors refer to different approaches, with crosses denoting the result of independent runs and the circles marking the mean for each method. The methods exhibit different bias-varianes trade- offs with respect to the optimum in the middle, where the variance measures the spread in the individual solutions and the bias is the error in the mean solution. On the right, we present different cooldown shapes based on their bias- variance characteristics. Here the bias computed as the loss difference between the average model experiments with different data orderings (shuffling) and a better reference model (obtained by longer training on the same data), whereas the variance the difference between the average loss of experiments and the loss of the average model of these experiments. Coordinates left plot where bias can be understood the loss distance between the circle dot (cluster mean) and star (optimum). Similarly, variance is the difference between the average loss for individual models (crosses) and the loss of the mean model (circle).
be connected to the
2 Background: Learning Rate Scheduler
Learning rate scheduling is a technique used in neu- ral network training where the optimization algo- rithm's learning rate adjusted during training ac-
100
LR, max LR %
cording to a predefined rule. Figure 2 illustrates
75
various scheduling techniques.
50
A warmup phase is commonly employed at the be- ginning of training, during which the learning rate is gradually increased from small values to the desired level. This approach allows the model to tolerate a
linear
25
cosine
wsd
higher peak learning rate during training, improves stability the initial phase, and positively affects final performance (oosson al. 2024b Kalra &
0
25
50
75
100
Training steps, %
Barkeshli 2024).
Cosine learning rate scheduling is widely used technique. This method adjusts the learning rate following a cosine curve, starting at a maximum value after a warmup period and gradually decreas- ing minimum value, often near zero (Bergsma)
Figure 2: Different learning rate schedules. We show the cosine, WSD, and linear schedules. A warmup stage of 5% of the training length added to each schedule.
et al. 2025). The key parameter for the cosine schedule is the training duration. As observed in Hoffmann et al.
2022 achicving optimal model performance for spocific token count requires matching the training
2

---

Published in Transactions on Machine Learning Research (08/2025)
duration to the cosine scheduler length. This finding was also verified by Hägele et al. (2024). Therefore, cosine learning rate scheduling has significant drawback: training sessions resumed after the cosine curve ends result in worse performance compared to training sessions with cosine schedule of the required length.
2.1 Warmup Stable Decay Scheduler
An alternative to the cosine learning rate sched- uler is the warmupsttable-deaay (WSD) scheduler (Figure 2 It consists three stages: warmup, a constant learning rate, and decay (cooldown). The cooldown stage is relatively short, typically around 20% of the total training duration. A notable fea- ture of WSD that it avoids the primary drawback
26
Validation perplexity
24
22
pre-cooldown
of the cosine scheduler. Training can be resumnd seamlessly from any part of the constant learning rate stage without performance degradation.
cooldown, different LR scheduling
The performance of this scheduler for language mod- eling, particularly the context of the Llama-like model (Touvron et al. 2023), is analyzed in depth by Hägele al. (2024). Their key finding is that the WSD learning rale scheduler can match even surpass the performance of the cosine scheduler.
20
40
60
80
100
Training steps, %
Figure 3: The drop in perplexity during the cooldown stage. The cooldown stage begins at 80% of the training, which leads to a dramatic drop in per- plexity during the final 20%.
critical part of WSD performance is the cooldown stage, As seen in Figure 3 a dramatic drop in vali- dation perplexity occurs during the cooldown stage. Such performance improvement (2024). Gaining insights into the cooldown stage may reveal ways to further improve performance. For this reason, this work focuses on the specifics of the cooldown stage.
also observed in downstream tasks, as demonstrated by Hägele et al.
Experimental Setup
As noted, this work focuses exclusively on the cooldown stage. Therefore, the pre-cooldown model is kept consistent across all experiments, and only the hyperparameters during the cooldown stage are modified, unless otherwise noted.
The model used is standard decoder-onl only transformer with 210 million parameters (Vaswani et al. 2017
identical to Llama (Touvron et al. 2023). Unless otherwise specified, we use the AdamW optimizer with Kingma & Ba 2017) and standard LLM training parameters. Training is conducted on a subset of Slim Pajamaa (Soboleva et al. 2023). The dataset is split into training and validation parts, and validation perplexity is used to evaluate performance. We provide a detailed description of the model architecture and used techniques in App.
weight decay (Loshchilov & Hutter 2019
The pre-cooldown model is trained for 26,400 iterations (2.7B tokens) with 300 warmup steps. The cooldown stage length is commonly set to 20%, which has been identified as an effective fraction for achieving strong final performance, shown by Hägele et al. (2024). During the cooldown stage, the learning rate is always reduced to zero.
Cooldown Shape: Does It Matter?
As noted by Hägele al. (2024), the shape of the cooldown phase significantly impacts the performance of WSD. Figure illustrates the validation perplexity for various cooldown shapes, well as the shapes them-
3

---

Published in Transactions on Machine Learning Research (08/2025)
selves, Interestingly, validation perplexity closely follows the shape of the learning rate and reaches different final values. However, Hägele et al. (2024) did not explore the reasons behind the different performance observed for different cooldown shapes.
100
Validation perplexity
LR, max LR %
21
75
50
20
25
19
80
85
90
95
100
80
90
95
100
Training steps, %
Training steps, %
linear
cosine
mirror cosine
square
sqri
Figure 4: The surprising effect of cooldown shapes on final performance. On the left,, different shapes of learning rate scheduling during the cooldown stage. On the right, the corresponding validation perplexities. Validation perplexity follows a pallern similar to that of the learning rate, but the final perfor- mance varies. It is clear that cooldown shapes affect the course of training as well the final perplexity.
4.1 Bias- Variance Point of View
To understand why different cooldown shapes yield varying performance, we examine the underlying op- timization dynamics. Specifically, we propose the following: different cooldown shapes create fun- damental bias-variance trade-off that explains their varying performance. Figure [ illustrates this con- cept schematically. Essentially, the cooldown shape controls critical balance: more aggressive explo- ration (high learning rates) of the loss surface in- creases variance
100
1.0
LR, max LR %
75
0.9
0.7
50
0.5
0.3
25
0.1
models (if trained with dif- ferent data orderings), while potentially yielding su- perior performance; in contrast, focusing on exploit- ing (low learning rates) the current region of the loss surface produces more consistent solutions but may sacrifice overall quality.
0
80
85
90
95
100
Training steps, %
Figure 5: Lowered linear cooldown shapes. The parameter controls the starting learning rate, where 100% is the learning rate of a constant stage.
To investigate this, we propose the following frame- work. For each cooldown shape, train N models from the same constant stage model with identical hyperparameters but different data orders. Each model is trained on random permutation of the same dataset portion for 33,000 steps in total (3.4B tokens). In addition to the cooldown shapes mentioned earlier, we use the "lowered lincar" shapes family displayed in Figure
We train models with a much longer cooldown stage (59,400 steps in total, 6B tokens) using the sqrt cooldown shape, which represents a good solution in the optimization landscape and therefore referred to as reference models. The dataset portion used for the reference models is the same that for the shorter experiments,
4

---

Published in Transactions Machine Learning Research (08/2025)
but is permuted (with repetitions) to match the extended training duration. In total, models are trained for different data permutations.
For each cooldown shape, we calculate two metrics: first, the performance the validation dataset of the average model (in weight space) over different data permulations; and second, the performance of each distinct experiment (specific data permutation). The same procedure is applied to the reference models. In other words, we obtain the following: the performance of each specific cooldown shape for all data permutations, the performance of the cooldown shape of the averaged model across different data permutations, and the averaged reference model performance.
Formally, let m, be the model with fixed cooldown shape and the i-th data permutation, m² be the average of reference models (in weight space), and L(-) the validation loss of a specified model. With this notation, can formulate the bias-variance decomposition of model performance in relation to better model m²:
E, [[m,,) I(m*)) = I [T(m=)] - I(m") )= I. [ms]) - I(m²) E. [T(ma)] L( [ms]],
(1)
bins
variance
where E, is the expectation with respect to the different data permutations. Concretely, E, [m:]] is the average model in weight space across different data permutations, and E, [L(m)) is the average loss of different models across different data permutations. We believe that the proposed framework captures information about optimization stability through the variance of performance across different experiments, and captures the distance a better solution achievable by learning on the same data. Surprisingly, similar results to those in the following can be achieved in weight space. We show how different measures of the bias and variance affect the results in App.
0.9
2.5
square
contribution [10 21
contribution [10- 2
1.5
Variance loss
0.7
Variance loss
2.0
mirror cosine
cosine
linear
1.0
0.9
1.5
0.77
sqrt
0.5
0.3
1.0
0.5
0.1
0.3
0.5
6
5.0
5.5
6.0
Bias loss contribution, [10-2]
Bias loss contribution, [10-]]
Bias + Variance min
Figure 6: Bias- plot for different cooldown shapes. We measure bias the difference in validation loss of the reference model (average of models trained for longer the same data but different permutations) and the loss of the average of experimental models (different data permutations). Variance is the difference between the average loss of experiments and the loss of the average model of these experiments (Equation D We also display a line where the minimum bias + variance for experiments is achicved. On the left, we compare only lowered linear shapes (Figure 5), which exhibit an expected decrease variance and increase in bias with decreasing parameter value. On the right, the full range of nonlinear shapes along with selected lowered linear shapes for reference. The lowered linear shape with parameter 0.7 and the sqrt shape occupy favorable positions, achieving balanced trade-off between bias and variance.
5

---

Published in Transactions on Machine Learning Research (08/2025)
Plot interpretation. The experimental results are presented in Figure 6 From the figure, we observe a clear relationship between bias and vari- ance, which we believe explains why some cooldown shapes lend lo perform belter than others. In this setting, the optimal cooldown shape is one that achieves the minimum of bias + variance. We note that the lowered linear shape with parameter 0.7 and the sqrí shape occupy an optimal position, achieving balance between bias and variance. This can be in- terpreted from exploration exploitation perspec- live: shapes with high variance and low bias tend to explore the loss surface more extensively, leading to siguificantly different solutions. Conversely, shapes with low variance and high bias tend to exploit the current solution by descending into the loss basin.
100
lincar 0.7 sqrt
LR, max LR %
75
50
25
0
80
85
90 95
100
Training steps, %
Figure 7: Comparison of the sqrt cooldown shape and lowered linear 0.7. The shapes are similar, achieving comparable performance, although the low- ered linear 0.7 exhibits slightly better perplexity.
One notable feature of Figure i is that sqrí lies close to the lowered linear shape with parameter of 0.7. We observe that the lowered linear 0.7 cooldown achieves performance comparable to, even better than, sqrt. This result shows that there is nothing special about the shape of sqri - is just one of the shapes that achieves a good trade-off between bias and variance; both shapes are presented in Figure
The bias-variance framework reveals that to perform well, a scheduler must minimize both the bias and variance loss contributions. However, we empirically observe trade-off between the two. We reproduce the bias-variance plot for smaller model and different dataset. The results align with the discussion in this section and are presented in App. D
4.2 Data Points Bias
We also explored the possibility of using data points bias as a basis for the bias-variance plot. This idea. stems from the notion that different cooldown schedulers can affect bias towards recent data points differently. From this perspective, the best cooldown scheduler would be the one that improves overall performance the most while keeping the bias towards recent data, points low. Naturally, trade-off emerges that is similar to the one we observed before.
To calculate the bias towards data points, we retrospectively calculate the perplexity for previously observed batches the original order after model training is completed (Bergsma et al., 2025). The results of this evaluation are displayed in Figure g Recency bias is evident both the pre-cooldown model and the post- cooldown models. However, the structure of the bias differs: the pre-cooldown model shows strong recency bias towards recent points, while the post-cooldown models exhibit U-shaped bias, which is consistent with the findings of Bergsma et al. (2025). This can be explained by the fact that the final cooldown steps are performed with a near-zero learning rate, making the impact of the final data points negligible.
Overall, the structure of the recency bias is greatly controlled by the cooldown shape. We can quantify it using the following framework: Let p(b;) be the pre-cooldown model perplexity on batch b, and e(b.) be the post-cooldown model perplexity on batch b, If the number of training batches before the cooldown stage starts and N the total number of training batches, we can define the shift and deviation for specitic cooldown shape
6

---

Published in Transactions Machine Learning Research (08/2025)
C = fiEN: P P<iS N) (cooldown batches indices),
difference; = c(bi) - p(b),,
shift 1
difference,
|C|
(2)
N
deviation
|ifferennce; - -shfft"
N
The idea behind this is that shift represents the improvement of the post-cooldown model over the pre- cooldown model, while deviation represents the consistency of this improvement across diffcrent data points. We empirically found that shift calculation better done only over the cooldown portion of the data, as it leads to an unskewed shift-deviation plot and high correlation with the previously observed results. We do not have a firm justification for this, but believe it is due the fact that the pre-cooldown model was not trained cooldown data at all, and therefore the shift is not skewed by the pre-cooldown model's recency bias.
The results of the evaluation are displayed in Figure S The structure of the plot is similar to what we observed before. We present additional visualizations in App.
24
square
4
cosine
Batch perplexity
22
Deviation, bias uniformity
0.9
casine
3
0.7
20
sqrt
2
16
0.1
0
25
50
75
100
-5.0
-4.5
-4.0
-3. 335
Step index, %
Shift, improvement over pre-cooldown
pre-cooldown model
cosine
lowered linear 0.1
square
sqrt
Figure 8: Shift-deviation result plots. On the left, evaluation of models on training data after training was completed (100-step window-averaged), It is clear that the pre-cooldown model has pronounced recency bias that ends abruptly at 80% of the total steps (the training limit of the pre-cooldown model). The post- cooldown models have lower overall mean perplexity and different recency bias structure. It is clear that the cooldown shape affects the recency bias structure. On the right, the shift-doviation plot for different cooldown shapes. It clear that the pattern similar to the one we observed before in Figure D We explore this similarity further in App.
4.3
Bias Variance Plot Conclusions
We conclude that different cooldown shapes regulate a trade-off between bias and variance during cooldown, which significantly impacts performance. Moreover, the same trade-off emerges from the recency bias point of view. Furthermore, we find that there of cooldown shapes, such as lowered lineor 0.7 and sqrt, that provide optimal tradc-off. Both cooldown shapes presented in Figure 7
7

---

Published in Transactions on Machine Learning Research (08/2025)
Averaging Longer Training
Averaging of models from different runs is a tech- nique extensively used recent works. In par- ticular, recently released influential models like
20
Validation perplexity
33k runs mean
OLMo2 (OLMo et al. 2025) and Llama3 (Meta All 2024] use technique of performing multiple in- dependent cooldown stages and averaging the final models, technique also referred to model souping [Wortsman et al. 2022]. Another approach is merg- ing completely different runs technique explored in Command model training (Cohere et al. 2025
33k single run
19
52.8k sqrt
18
With this molivation, we investigale how averag- ing final models compares to longer training for dif- ferent cooldown shapes. Specifically, a cooldown stage four times on non-interseeting dataset portions from the same pre-cooldown model. Additionally, we train one model for longer on all data observed by the shorter runs. The longer run uses the sqrt cooldown shape for a total of 52.8k steps and 20% cooldown portion. In total, the longer achieves the same flops and data portion as four smaller runs combined.
mirror
linear
Sqrt
square
cosine
0.1
cosine
perform
Figure 9: Performance comparison of averaged models to longer training. High-variance cooldown shapes yield the best performance. However, simple longer run outperforms all averages.
Then, we average model weights trained on different dataset portions for each cooldown shape and calculate performance on validation data. The results are displayed in Figure g While averages perform better than their corresponding non-averaged runs, they still fall short of a simple longer run (52.8k, dashed line in Figure 9h. It is also evident that high-variance, low-bias shapes provide the best performance when averaging. We consider the possibility that averages perform worse than the longer due to the model underfitting the data, and longer training may close the gap or even surpass it.
Therefore, in the setting of model averaging, high-variance low-bias shapes like mirror cosine perform better than tradc-off shapes like sqrt. We also discover that a bias-variance plot emerges in model weight space, which we discuss App.
Cooldown Stage Hyperparameters
While tuning hyperparameters is arguably infeasible for large model runs that are extremely expensive and are often only performed once, the cooldown represents significantly shorter part of training and can potentially be tuned. Therefore, to better understand the impact of modifying the cooldown shape, examine the significance of hyperparameter selection during the cooldown stage. This investigation could provide insights into the relative importance of various factors when tuning the model's training. Tt is divided into two parts: batch size selection and Adam hyperparameter selection.
6.1 Batch Size Impact
We use the sqrt cooldown shape and vary the batch size, adjusting the number of steps to maintain the same total token count. Additionally, as described by Chiley et al. (2019): Busbridge et al. (2023 ; Pagliardini et al. (2024), we adjust the AdamW 81 = 0.9 and 32 = 0.95 parameters to match the optimizer 's statistics token half-life. In essence, the token half-life the number of successive previous steps that contribute to half of the optimizer's EMA accumulator. For new batch size B kB, the momentum is set as p* We perform simple experiments without adjusting 81 and B2 (Figure 10b For experiments with higher batch
8

---

Published in Transactions on Machine Learning Research (08/2025)
sizes, we also increase the learning rate by trying set of several values and selecting the best one based on performance. The table of learning rate selection for each batch size can be found in App. H the scaling roughly follows square root scaling factor (Malladi al. 2024
Batch size variation conclusions. We observe in Figure 10a that increasing the batch size with matching token half-life leads to deterioration in performance. Therefore, while increasing the batch size can negatively impact performance,
believe this aligns with observations about critical batch sizes described by McCan- dlish et al. 2018). Despite the unusual loss behavior during the cooldown stage, it appears that batch size variations adhere to patterns previously observed in the literature.
Interestingly, significant performance improvements are observed when no adjustments are made to 31 and 82 (Figure 10by. This indicates that the token half-life of AdamW states may not be optimal in our setting, a topic further explored in Section 6.2
18.35
Validation perplexity
Validation perplexity
18.6
18.30
18.5
18.25
18.4
18.20
500
1000
1500
2000
500
1000
1500
2000
Batch size
Batch size
(a) With AdamW 81 and 82 adjustments
(b) Without AdamW 3 and 82 adjustments
Figure 10: Impact of batch size variation during the cooldown stage on validation perplexity. Each batch sample contains 512 Lokens, and the base batch size is 200. The learning rate is scaled up with increasing batch sizes, which we discuss App. H We also plot error bars displaying the minimum and maximum performance of experiments for 5 different data permutations of the same data. On the left, besides modifying the batch size, AdamW paraneters B1 and 82 are adjusted match the token half-life, resulting in performance deterioration for high batch sizes. On the right, the batch size is increased without modifying the AdamW parameters, which leads to substantial improvement. This highlights the importance of the AdamW token half-life.
6.2 AdamW Parameters Impact
We explore how AdamW betas impact cooldown stage performance and compare the resulting performance improvements to those achieved through shape selection. For all experiments, use the sqrt cooldown shape and vary the tokens' half-life by adjusting 31 = 0.9 and 8 0.95, modifying both using the parameter 31 = 3 3g G.. The results are presented in Figure 11 In addition to modifying both B1 and 82,
we conduct experiments with &a variation while keeping B1 0.9 fixed. These results are also shown in Figure
Tokens Half-Life Variation Conclusions. We observe that varying the AdamW tokens' half-life can lead to significantly different performance outcomes. From the experiment with varying both 3 = 0.9 and 82 0.95, we see that the best cooldown shape at p 1 is almost identical, in terms of perplexity, to the worst p = 0.3. Therefore, modifying AdamW parameters during cooldown worthwhile, performance improvements be comparable to those achieved through cooldown shape sclection.
9

---

Published in Transactions Machine Learning Research (08/2025)
Validation perplexity
19.0
Validation perplexity
18.8
18.8
18.6
18.6
81 0.9689,
18.4
82 0.999995
82 0.9847
18.4
18.2
18.2
0.05
0.1
0.25
0.5
1.0
1e-05
0.001
0.1
p, B1 = 0.9P, .8 = 0.95"
P, 81 = 0.9, B2 = 0.95.
cosine
miror cosine
square
lincar
sqrt
range for different cooldown shapes
Figure 11: Evaluation of the impact of AdamW parameters cooldown performance. On the left, we vary both 81 and 82. On the right, we vary only 82. The betas are adjusted using the parameter p: 3 3= BP, with base betas B1 = 0.9 and B2 = 0.95. Examining the range for different cooldown shapes makes it clear that a good choice of hyperparameters can greatly improve performance, matching or even surpassing the importance of cooldown shape selection.
From the second plot in Figure it is noticeable that performance greatly depends on B2 variations, with surprisingly large values lending to yield the best results.
Additionally, a notable property is that the optimal cooldown shape, from a bias-variance perspective, remains unchanged with the choice of 31 and 32 within a reasonable range of values. This clearly scen in Figure II as sgrl acts as a lower bound for all the shapes tested unless p becomes too small. In the case of small p, the AdamW tokens' half-life becomes so large that the optimizer state is almost not updated, which is reasonable setting and requires higher learning rates.
7 Examining the Loss Landscape
The 'river valley perspective' is one of the intuitions behind the WSD scheduler proposed by Wen et al. (2224), suggesting that WSD optimization can be visualized descending from mountain top along river. The constant learning rate stage corresponds to descending in the general direction of the river's flow, while the cooldown stage corresponds to descending directly into the river itself. However, clear visualization the river valley is still lacking. This section focuses on visualizing the river valley during the WSD cooldown stage.
To achieve this, plot the loss landscape using the following coordinates:
61 A vector in weight space connecting the pre-cooldown checkpoint to the final model after cooldown. We refer this the global oplimization direction.
62 A vector corresponding to ten optimizer steps from a given point. We refer to this as the Adam sleps direction. We choose ten steps minimize the impact noise the obtained vector.
Loss Landscape Plots Interpretation. Figures and 30 present the results. These plots clearly show river valley along the global optimization direction at the beginning of the cooldown stage. It remains visible in the middle of the cooldown stage, though less pronounced optimization nears descent into the
10

---

Published in Transactions Machine Learning Research (08/2025)
basin. By the end of the cooldown stage, only the final basin is visible. We believe these plots provide strong evidence supporting the river valley basin concept.
Morcover, this perspective provides further insights for the bias-variance framework the observation that local optimization step (from single example) is almost orthogonal to the best global direction is exactly what causes the variance. Training cannot make progress downstream (reducing the bias) without taking large steps along the orthogonal directions (causing variance).
Cooldown start (81.8 steps)
Cooldown end (100.0% steps)
19.34
17.86
36.32
-3.0
Global optimization antimizaaion
-0.3
Adam steps
-1.5
0.0
0.0
19.94
-3.0
-1.5
0.3
1.5
-0.3
0.0
3.0
Global optimization
0.0
Adam steps
1.5
0.3
3.0
current point
plot min point
Loss landscape plots at different cooldown points. We plot the loss landscape using two vectors: the global optimization direction (the line between the pre-cooldown checkpoint and the final model) and the Adam stops direction (the line between the given point and ten optimizer steps from it). The Z-axis represents perplexity on a portion of validation data. We see a potential valley along the global optimization direction, which is less pronounced optimization nears descent into the basin. A larger plot is presented in Figure 30
Figure 12:
8 Related Work
Learning rate schedulers. The evolution of learning rate schedulers has been driven by the need to ad- just the lcarning rate during training to improve optimization convergence. This need can also be viewed as a maller of balancing exploration and exploitation (Tyer el al., 2020 Subramanian et al. 2024).. Tradi- tional approaches, such as step decay (reducing the learning rate at fixed intervals) and exponential decay (continuously reducing the lcarning ratc), laid the foundation for schedule-based optimization (Li & Arora
2019
Ge al.) 2019 Smith al., 2017). The cosine schedule (Loshchilov & Hutter (017) one of the most common choices in large-scale deep learning. Others include inverse square root (Raffel et al. 2020 Chowdhery al. 2023), Noam (Vaswani et al.) (017)) stepwise schedules (Bi et al. (024).. In the last year, (WSD), also called the trapezoidal schedule (Dosovitskiy et al. 2020 has become
Caarmup-Stable.
popular for its benefits of an undetermined training length, performance (Zhai et al. 2022 Shen el al. 2024a
Hägele et al. 2024 Hu et al. 2024) and continual learning (Tbrahim et al. 2024 Gupta et al. 2023
Janson
et al. 2025).
Given the significant impact of the learning rate scheduler on model performance, and considering the com- plexity of finding the optimal learning rate for LLMs due the interplay between learning rate, batch size,
11

---

Published in Transactions on Machine Learning Research (08/2025)
number training tokens, model size, and other hyperparameters, multiple works have focused on schedule- free optimization (Defazio et al. 2024 or scheduling techniques that work well with different model sizes (Shen et al. 2024b
Understanding dynamics of learning rate schedulers. The interaction between learning schedules and optimization dynamics remains an active research area. For example, for WSD, Wen et al. (2024) proposed
the 'river valley' hypothesis to cxplain WSD behavior, while Subramanian et al. (2024) analyzed phase
transitions in loss landscapes during schedule changes. Our work extends these observations through direct
visualization of the cooldown-stage loss landscape. Additionally, recent work by Schhaip et al. (2025) showed that the empirical behavior of the WSD scheduler behaves surprisingly similarly to performance bound from non-smooth convex optimization theory (Defazio et al. 2023) which stems from either non-smoothnes: or vanishngg gradient norms. We investigate such gradient norms in the cooldown stage in App. G
Conclusions
This work provides a comprehensive analysis of the cooldown stage in the WSD learning rate scheduler, offering insights into its dynamics and impact on model performance. The key findings from this study are summarized below.
Cooldown Shape Selection. The shape of the cooldown phase critically influences model performance. We introduce bias-variance framework to compare different cooldown shapes, identifying sgrt and lowered linear 0.7 as optimal choices due to their ability to balance exploration and exploitation. We observe that high-variance low-bias shapes like mirror cosine perform better when final model averaging is performed. We also provide a recency bias perspective that aligns with the bias-variance framework.
Hyperparameter Sensitivity. Our experiments show that varying the batch size during the cooldown stage has a relatively minor impact on performance, consistent with prior findings on critical batch sizes. However, adjusting the AdamW hyperparameters (3, and 32) revealed significant performance differences, comparable to those observed with cooldown shape selection. Notably, larger values of 32 consistently yielded improved results, highlighting the importance of tunettunng these parameters.
Loss Landscape Visualization. We further validate the 'river valley perspective" through loss landscape visualizations. These plots illustrate how optimization transitions from broad exploration during the constant learning rate phase to focused descent into loss basin during the cooldown phase, providing evidence for this conceptual framework.
10 Future Work
While our analysis limited to moderate model sizes, we validate the main findings across different numbers of parameters and datasets (D). For future directions, we believe there remain several important opportunities to deepen understanding:
1. Our work focuses solely on the validation loss and the optimization landscape. However, future work should explore the relevance of observed effects for downst.ream tasks and model behavior, although connection between loss and downstream performance was previously observed (Du et al. Hägele ct al. 2024
2. Related to the above, it remains unclear what mechanistic changes occur in the model over the course of the cooldown stage. As a first step, we explore internal feature quality improvement over the course of the cooldown stage with experiments in App. [ but we did not arrive at useful conclusions. Other possible approaches include the recently proposed model diffing frameworks (Bricken et. al. 2024
Lindsey et al. 2024
12

---

Published in Transactions on Machine Learning Research (08/2025)
References
Marlon Becker, Frederick Altrock, and Benjamin Risse. Momentum-sam: Sharpness aware minimization without computational overhead, 2024. URL
Shane Bergsma, Nolan Dey, Gurprect Gosal, Gavia Gray, Daria Soboleva, and Joel Hestness. Straight to zero: Why linearly decaying the learning rate to zero works best for llms. arxiv preprinl ar Xrxv:2502.119933 15938, 2025.
Xiao Bi, Deli Chen, Guanting Chen, Shanhuang Chen, Damai Dai, Chengqi Deng, Honghui Ding, Kai Dong, Qiushi Du, Zhe Fu, et al. Deepseek Ilm: Scaling open-source language models with longtermism. arXiv preprint arXiv:240100000
2024.
Trenton Bricken, Siddharth -ishra-Sharma, Jonathan Marcus, Adam Jermyn, Christopher Olah, Kelley Rivoire, and Thomas Henighan. Stage-wist model diffing, 2024. URL fttp::://bransformer-coo circurissss... circuits. pub//2024/oodellddfffii ing/ index.html
Dan Busbridge, Jason Ramapuram, Pierre Ablin, Tatiana Likhomanenko, Feshan Gunesh Dhekane, Xavier Suau, and Russ Webb. How to scale your ema. 2023. URL https://arxiv, http://aarxx...ggbbba/////a org/abs/2307 13813
Vitaliy Chiley, Πya Sharapov, Atli Kosson, Urs Koster, Ryan Recce, Sofia Samaniego de la Fuente, Vishal Subbiah, and Michael James. Online normalization for training neural networks. Advances in Newral Information Processing Systems, 32, 2019, 19X.v1995500594
Aakanksha Chowdhery, Sharan Narang, Jacob Devlin, Maarten Bosma, Gauray Mishra, Adam Roberts, Paul Barham, Hyung Won Chung, Charles Sutton, Sebastian Gehrmann, et al. Palm: Scaling language modeling with pathways. Journal of Machine Learning Rescarch, 24(240):1-113, 2023.
Team Cohere, Aakanksha, Arash Ahmadian, Marwan Ahmed, Jay Alammar, Yazeed Almumay, Sophia Al- thammer, Arkady Arkhangorodsky, Viraat Aryabumi, Dennis Aumiller, Raphaël Avalos, Zahara Aviv, Sammie Bae, Saurabh Baji, Alexandre Barbet, Max Bartolo, Björn Bebensee, Neeral Beladia, Walter Meller-Morales, Alexandre Bérard, Andrew Berneshawi, Anna Bialas, Phil Blunsom, Matt Bobkin, Adi Bongale, Sam Braun, Maxime Brunct, Samuel Cahyawijaya, David Cairuz, Jon Ander Campos, Cassic Cao, Kris Cao, Roman Castagné, Julián Cendrero, Leila Chan Currie, Yash Chandak, Diane Chang, Giannis Chatziveroglon, Hongyu Chen, Claire Cheng, Alexis Chevalier, Justin T. Chiu, Eugene Cho, Fu- gene Choi, Eujeong Choi, Tim Chung, Volkan Cirik, Ana Cismaru, Pierre Clavier, Henry Conklin, Lucas Crawhall-Stein, Devon Crouse, Andres Felipe
Ben Cyrus, Daniel D'souza, Hugo Dalla-Torre, John Dang, William Darling, Omar Darwiche Domingues, Saurabh Dash, Antoine Debugne, Théo Dehaze, Shaan Desai, Joan Devassy, Rishit Dholakia, Kyle Duffy, Ali Edalati, Ace Eldeib, Abdullah Elkady, Sarah Elsharkawy, Irem Ergün, Beyza Ermis, Marzieh Fadaee, Boyu Fan, Lucas Fayoux, Yannis Flet-Berliac, Nick Frosst, Matthias Gallé, Wojeiech Galuba, Utsav Garg, Matthieu Geist, Mohammad Gheshlaghi Azar, Scraphina Goldfarb-Tand Tarrann,, Tomas Goldsack, Aidan Gomez, Victor Machado Gonzaga, Nithya Govin- darajan, Manoj Govindassamy, Nathan Grinsztajn, Nikolas Gritsch, Patrick Gu, Shangmin Guo, Kilian Haefeli, Rod Hajjar, Tim Hawes, Jingyi He, Sebastian Hofstätter, Sungjin Hong, Sara Hooker, Tom Hosk- ing, Stephanie Howe, Eric Hu, Renjie Huang, Hemant Jain, Ritika Jain, Nick Jakobi, Madeline Jenkins, JJ Jordan, Dhruti Joshi, Jason Jung, Trushant Kalyanpur, Siddhartha Rao Kamalakara, Julia Kedrzycki, Gokce Keskin, Edward Kim, Joon Kim, Wei-Yin Ko, Tom Koemi, Michael Kozakov, Wojciech Kryš- ciński, Arnav Kumar Jain, Komal Kumar Teru, Sander Land, Michael Lasby, Olivia Lasche, Justin Lee, Patrick Lewis, Jeffrey Li, Jonathan Li, Hangyu Lin, Acyr Locatelli, Kevin Luong, Raymond Ma, Lukas Mach, Marina Machado, Joanne Magbitang, Brenda Malacara Lopez, Aryan Mann, Kelly Marchisio, Olivia Markham, Alexandre Matton, Alex McKinney, Dominic McLoughlin, Jozef Mokry, Adrien Morisot, Au- tumn Moulder, Harry Moynehan, Maximilian Mozes, Vivek Muppalla, Lidiya Murakhovska, Hemangani Nagarajan, Alckhya Nandula, Hisham Nasir, Shauna Nehra, Josh Netto-Rosen, Daniel Ohashi, James Owers- Bardsley,,,,,, Jason Ozuzu, Dennis Padilla, Gloria Park, Sam Passaglia, Jeremy Pekmez, Laura Pen- stone, Aleksandra Piktus, Case Ploeg, Andrew Poulton, Youran Qi, Shubha Raghvendra, Miguel Ramos,
13

---

Published in Transactions on Machine Learning Research (08/2025)
Ekagra Ranjan, Pierre Richemond, Cécile Robert-Michon, Aurélien Rodriguez, Sudip Roy, Laura Ruis, Louise Rust, Anubhav Sachan, Alejandro Salamanca, Kailash Karthik Saravanakumar, Isha Satyakam, Alice Schoenauer Sebag, Priyanka Sen, Sholeh Sepehri, Preethi Seshadri, Ye Shen, Tom Sherborne, Sylvie Chang Shi, Sanal Shivaprasad, Vladyslav Shmyhlo, Anirudh Shrinivason, Inna Shteinbuk, Amir Shukayev, Mathieu Simard, Ella Snyder, Ava Spalaru, Victoria Spooner, Trisha Starostina, Florian Strub, Yixuan Su, Jimin Sun, Dwarak Talupur,, Eugene Tarassov Elena Tommasone, Jennifer Tracey, Billy Trend, Evren Tumer, Ahmet Ustün, Bharat Venkitesh, David Venuto, Pat Verga, Maxime Voisin, Alex Wang, Donglu Wang, Shijian Wang, Edmond Wen, Naomi While, Jesse Willman, Marysia Winkels, Chen Xia, Jessica Xie, Minjie Xu, Bowen Yang, Tan Yi-Chern, Tvan Zhang, Zhenyu Zhao, and Zhoujie Zhao. Command a: An enterprise-ready large language model, 2025. URL ItRt:t/ts:///art httpaaawwiw.iiireeee.e org/ hts///bb/25500 /b////////25000099
DeepSock-AI ct. al. Doepscck-v3 technical report. 2024. URL
Aaron Defazio, Ashok Cutkosky, Harsh Mehta, and Konstantin Mishchenko. When, why and how much? adaptive learning rale scheduling by refinement. Oct 2023. URL http:://axiv. httpp://aaaiivo.ggabbb/////s https///311..000..
Aaron Defazio, Xingyu Alice Yang, Harsh Mehta, Konstantin Mishchenko, Ahmed Khaled, and Ashok Cutkosky. The road less scheduled, 2024. URL https:///aaaat http://axxivv.rrro wrgg/bs///05..... 15682
Alexey Dosovitskiy, Lucas Beyer, Alexander Kolesnikov, Dirk Weissenborn, Xiaohua Zhai, Thomas Un- lerthiner, Mostafa Dehghani, Matthias Minderer, Georg Heigold, Sylvain Gelly, et al. An image is worth 16x16 words: Transformers for image recognition at scale. arXiv preprint wrXiv:2010.11920 2020.
Zhengxiao Du, Aohan Zeng, Yuxiao Dong, and Jie Tang. Understanding emergent abilities of language models from the loss perspective.
Rong Ge, Sham M. Kakade, Rahul Kidambi, and Praneeth Netrapalli. The step decay schedule: A near optimal, geometrically decaying learning rate procedure. CoRR, abs/ 1es/1994.1283,, 2019. URL http:// bexxiv.org/bbb.. 1ete/1994..111 12838
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? Aug 2023. URL http://arxiv> http:///axvvvor///////.. http//bb...///11111//
Jordan Hoffmann, Sebastian Borgeaud, Arthur Mensch, Flena. Buchatskaya, Trevor Cai, Eliza Rutherford, Diego de Las Casas, Lisa Anne Hendricks, Johannes Welbl, Aidan Clark, et al. Training compute-optimal large language models. Advances in Neural Information Processing Systems, 35:30016 30030, 2022.
Shengding Hu, Yuge Tu, Xu Han, Chaoqun He, Ganqu Cui, Xiang Long, Zhi Zheng, Yewei Fang, Yuxi- ang 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. Apr 2024. URL https:/ http://axxvvvv.o/e wrw. ab/////b/../////////////////// 06395v2
Alexander Hägele, Elie Bakouch, Atli Kosson, Loubna Ben Allal, Leandro Von Werra, and Martin Jaggi. Scaling laws and compute-optinal training beyond fixed training durations. Neural Information Processing Systems, May 2024.
Adam Ibrahim, Benjamin Thérien, Kshitij Gupta, Mats Leon Richter, Quentin Gregory Anthony, Eugene Belilovsky, Timothée Lesort, and Irina Rish. Simple and scalable strategics to continually pre-train large language models. Transactions on Machine Learning Research, 2024.
Nikhil Iyer, V. Thejas, Nipun Kwatra, Ramachandran Ramjee, and Muthian Sivathanu. Wide-minima density hypothesis and the explore-exploit learning rate schedule. CoRR, abs/ //////200000777 2020. URL https: wttp:r/////../009///939
14

---

Published in Transactions on Machine Learning Research (08/2025)
Paul Janson, Vaibhav Singh, Paria Mehrbod, Adam Ibrahim, Irina Rish, Eugene Belilovsky, and Benjamin Thérien. Beyond cosine decay: On the effectiveness of infinite learning rate schedule for continual pre- training. arXiv preprint arXi:2503 02844, 2025.
Dayal Singh Kalra and Maissam Barkeshli. Why warmup the learning rate? underlying mechanisms and improvements. 2024. URL https: wrtg/aa///////0 09405
Diederik P. Kingma and Jimmy Ba. Adam: A method for stochastic optimization. 2017. URL https:
//arxiv. org/ orgg/bs///1////// abs//412..6.99 6980
Atli Kosson, Bettina Messmer, and Martin Jaggi. Rotational equilibrium: How weight decay balances learning
across neural networks. In Forty-first International Conference Machine Learning, 2024a. URL https:
//openreview www./eoooun..dt0iiinraaa...
Xiiv:2305.17212
Atli Kosson, Bettina Messmer, and Martin Jaggi. Analyzing & reducing the need for learning rale warmup in
GPT training. In The Thirty-eighth Annual Conference on Neural Information Processing Systems, 2024b. URL https://openrevieen www..netoooundd-zgnrrrrroao
(tl:2410.23922)
Zhiyuan Li and Sanjeev Arora. An exponential learning rate schedule for deep learning. CoRR, abs/1910.07454, 2019. URL attp://arxxe wwwpa/wwwww...ab/////aa/// 07445
Kaizhao Liang, Lizhang Chen, Bo Liu, and Qiang Liu. Cautious optimizers: Improving training with one line of code, 2025. URL https:://axxv.orr///a abs///b///// 2411 16085
Jack Lindsey, Adly Templeton, Jonathan Marcus, Thomas Conerly, Joshua Batson, and Christopher Olah. Sparse crosscoders for cross-layer features and model diffing, 10 2024. URL https://
transformer-circcttee http://pp///////ssccrre////eeee html
Ilya Loshchilov and Frank Hutter. SGDR: stochastic gradient descent with warm restarts. In International Conference on Learning Representations, 2017.
Ilya Loshchilov and Frank Hutter. Decoupled weight decay regularization. In International Conference on Learning Representations, 2019.
Anton Lozhkov, Loubna Ben Allal, Leandro von Werra, and Thomas Wolf. Fineweb-edu: the finest collection of educational content, 2024. URL http::/s//nggiigggeeeeee///r////
Sadhika Malladi, Kaifeng Lyu, Abhishek Panigrahi, and Sanjeev Arora. On the sdes and scaling rules for adaptive gradient algorithms, 2024. URL https://arxiv. wtt.:aawww....rraar/////......
Sam McCandlish, Jared Kaplan, Dario Amodei, and OpenAl Dota Team. An empirical model of large-batch training. 2018. CRL
06122
Meta AL. The Ilama 3 herd of models. 2024. URL https://arxii: http:///aaaaw...rg///////000// .1783
Team OLMo, Pete Walsh, Luca Soldaini, Dirk Groeneveld, Kyle Lo, Shane Arora, Akshita Bhagia, Yuling Gu, Shengyi Huang, Matt Jordan, Nathan Lambert, Dustin Schwenk, Oyvind Tafjord, Taira Anderson, David Atkinson, Faeze Brahman, Christopher Clark, Pradeep Dasigi, Nouha Dziri, Michal Guerquin, Hamish Ivison, Pang Wei Koh, Jiacheng Liu, Saumya Malik, William Merrill, Lester James V. Miranda, Jacob Morrison, Tyler Murray, Crystal Nam, Valentina Pyatkin, Aman Rangapur, Michael Schmitz, Sam Skjonsberg, David Wadden, Christopher Wilhelm, Michael Wilson, Luke Zeitlemoyer, Ali Farhadi, Noah A. Smith, and Hannaneh Hajishirzi. 2 olmo 2 furious, 2025. URL htttss://axxiii httppp///aaaaa...or///////..... 00656.
Matteo Pagliardini, Pierre Ablin, and David Grangier. The ademamix optimizer: Better, faster, older, 2024. URL. httpp:///a//xxi..... wwtp:///..rrra//////////// 03137.
Alec Radford, Jeffrey Wu, Rewon Child, David Luan, Dario Amodei, Llya Sutskever, et al. Language models are unsupervised multitask learners. 2019.
15

---

Published in Transactions on Machine Learning Research (08/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. Journal of machine learning research, (140)):1-67, 2020. URL httt:://arrdv..or//aa///////
Anton Razzhigaev, Matvey Mikhalchuk, Elizaveta Goncharova, Nikolai Gerasimenko, Ivan Oseledets, Denis Dimitrov, and Andrey Kuznetsov. Your transformer is secretly linear. In Lun-Wei Ku, Andre Martins, and Vivek Srikumar (eds.). Proceedings of the 62nd Annual Meeting of the Association for Computational Linguishics (Volume 1: Long Papers), pp. 5376 5384, Bangkok, Thailand, August 2024. Association for Computational Linguistics. doi: (118653/v122244ac/ lolllloon...9.... 293. URL http:tps///ccaatttlogoooo/////// acl acl-long. 293/
Fabian Schaipp, Alexander Hägele, Adrien Taylor, Umut Simsekli, and Francis Bach. The surprising agree- ment between convex optimization theory and learning-rate scheduling for large model training, 2025. URL 18965
Noam Shazeer. Glu variants improve transformer, 2020. URL
Yikang Shen, Zhen Guo, Tianle Cai, and Zengyi Qin. Jetmoe: Reaching llama2 performance with 0.1m abbp2444.0771388
dollars. Apr 2024a. URL http://arxiv
Yikang Shen, Matthew Stallone, Mayank Mishra, Gaoyuan Zhang, Shawn Tan, Adilya Prasad, Adriana Meza Soria, David D. Cox, and Rameswar Panda. Power scheduler: A batch size and token number agnostic learning rate scheduler. August 2024b.
Samuel L. Smith, Pieter-Jan Kindermans, and Quoc V. Le. Don't decay the learning rate, increase the batch size. CoRR, abs/1711.00489, 2017. URL
ttrg/abb/17111 11111 11tp./1011.111199991
Daria Soboleva, Faisal Al-Khateeb, Robert Myers, Jacob R Steeves, Joel Hestness, and Nolan Dey. SlimPa-
httpajaaaaaatttt///aaaaa jama: A 627B token cleaned and deduplicated version of RedPajama. https: cerebras. net/blog/
2023. URL https:
627B
Jianlin Su, Yu Lu, Shengfeng Pan, Ahmed Murtadha, Bo Wen, and Yunfeng Liu. Roformer: Enhanced
transformer with rotary position embedding, 2023. URL https: pttts:///rrivv http::://aaaii.or//aa////// arrg/abs/2110 09864
Shreyas Subramanian, Vignesh Ganapathiraman, and Corey D) Barrett. Hop, skip, jump to convergence: Dy- namics of learning rate transitions for improved training of large language models. In Yaser Al-Onaizan, Mo- hit Bansal, and Yun-Nung Chen (eds.), Findings of the Associalion for Compulational Linguistics: EMNLP 2024, pp. 16349-16366) Miami, Florida, USA, November 2024. Association for Computational Linguisties. doi: 10.18653/v1 /v00 hinga.innnnnggennl......
URL https://aaclanthology/ 2024 f andings-emmll
954/
Akiyoshi Tomihari and Issci Sato. Understanding linear probing then fine-tuning language models from NTK perspective. In The Thirly-eighth Annual Conference on Neural Information Processing Systems, 2024.
URL
Hugo Touvron, Thibaut Lavril, Gautier Izacard, Xavier Martinet, Marie-Anne Lachaux, Timothée Lacroix, Baptiste Rozière, Naman Goyal, Eric Hambro, Faisal Azhar, et al. Llama: Open and efficient foundation language models. arXiv preprint arXiu:2302.1399 2023.
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention is all you need. In Advances in Neural Information Processing Systems, volume 30, 2017.
Kaiyue Wen, Zhiyuan Li, Jason Wang, David Hall, Percy Liang, and Tengyu Ma. Understanding warmup- stable-decay learning rates: A river valley loss landscape perspective. 2024. URL https://arxiv..m abs/ abs/2410.05192.
16

---

Published in Transactions Machine Learning Research (08/2025)
Mitchell Wortsman, Gabriel Ilharco, Samir Yitzhak Gadre, Rebecca Roelofs, Raphael Gontijo-Lopes, Ari S. Morcos, Hongseok Namkoong, Ali Farhadi, Yair Carmon, Simon Kornblith, and Ludwig Schmidt. Model soups: averaging weights of multiple fine-tuned models improves accuracy without. increasing inference time, 2022. URL
Xiaohua Zhai, Alexander Kolesnikov, Neil Houlsby, and Lucas Beyer. Scaling vision transformers. Proceed- ings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition (CVPR), pp. 12104-12113, June 2022.
Biao Zhang and Rico Sennrich. Root mean square layer normalization, 2019. URL https: bttps://arxiv..om abs/ AbF/1910.07467
17

---

Published in Transactions Machine Learning Research (08/2025)
A
Experiments Context
The experimental setup mirrors that of Hägele et al. (2024). Specifically, we employ a decoder-only transformer model resembling Llama3 (Meta A1, 2024). The architecture incorporates SwiGLU activa- tions (Shazeer 2020), RoPE (Su et al.) 2023), RMSNorm Zhang & Sennrich 2019), and alternating at- tention and forard layers. Unless explicitly stated otherwise, we adhere the following methods and hyperparameters.
10 decoupled weight decay of 0.1 [Loshchilov & Hutter 2019), and gradient clipping at 1.0. All runs are trained with bfloa/16 automatic mixed precision. We generally apply a short warmup of 300 steps. The batch size is set to 200, corresponding to 0.1 million tokens for a sequence length 512. The vocabulary is derived from the GPT-2 tokenizer [Radford et al., 2019) and consists of 50,304 tokens. Hidden shapes of the used model are presented in Table
We use the AdamW optimizer with beta parameters set to (2), (0.9,0.95).
Model Size
dmodel
layers
ffw dim
head dim
Theads
60M
512
10
1536
64
210M
768
24
2048
Table 1: Model configurations with varying parameter counts. The 210M parameter model is used in the majority of the experiments. The 60M parameter model is used for experiments in App. D
B Used Cooldown Shapes Formulas
The formulas for cooldown shapes used in the article and presented in Figure are:
[0, 1],
linear -
cosine 1 + cos TT
2
+ cos cos(TT)
mirror cosine 2(1 xt)
square x²2
sqrt Vx.
C Bias- Varianer Plot Basis Selection
In addition to the basis for bias-variance plots proposed in equation 1 (loss-based space), analyze how results change with the selection of different bases.
C.1 Model Weights Space
Surprisingly, the same structure is observed when results are plotted in model weights space. We use each experiment's model weights as a vector of model parameters of size That is, if 8in E RR represents model weights trained with cooldown shape i by data permutation n, and Tn RP represents reference model weights trained on data permutation n, the complete formulas used are presented in equation (3). In essence, we measure bias as the distance between the mean reference and experimental models, and variance as the variance of model weights across different permutations.
18

---

Published in Transactions Machine Learning Research (08/2025)
N 8in
1
bias;
N
(3)
N
variance,
|| T7
Plots in this basis are presented in Figure 13 The experiments used are the same as in section 4.1 but with different coordinates. We also plot a comparison of the two spaces in Figure 14 We see that while some variations are observed, the general structure of the plot remains the same. Presumably, variations between plots should become less noticeable with increase in the number of experiments with different data permutations.
Model weights var,, [10 5
Lowered linear
Model weights var., [10 6]
Lowered linear and shapes
1.0
square
0.9
cosine
4
litrear
0.7
mirror ohee
2
0.5
0.7 sqrt
0.3
0
0.1
300
325
350
280
290
300
Bias, model weights space
Bias, model weights space
Figure 13: Bias-variance plot in model weights space. The general structure agrees with the discussion
in section 4.1
Model weights space
Model weights var., [10 s]
Loss-based space
square
square
7.5
Variance, [10- 2]
2
mirror cosine
cosimt linear
lincar
5.0
mong cosine
sqrt
agrt
2.5
0.5
0.3
0.3
0.0
0.1
0.1
275
300
325
350
6
Bias, model weights space
Bias,
[10 [10-2] 2]
Figure 14: Comparison of bias-variance plots in loss space and model weights space. While minor variations noticeable, the general structure of the plots is the same.
19

---

Published in Transactions Machine Learning Research (08/2025)
C.2 Other Loss-based Spaces
Another coordinate space used for the bias-variance plot is a simple loss-based space. If Stn E R represents the model validation loss trained with cooldown shape by data permutation n, the complete formulas used are presented equation (4) which are simply the mean and variance.
N
bias;
Siz
N
n=1
(4)
variance;
||sin
37
n=1
2.0
square
Variance, loss, [10
mirror cosine
1.5
linear
0.9
cosine
sgrt
0.7
1.0
0.5
0.3
0.1
2.91
2.92
2.93
2.94
Bias, loss
Figure 15:
Comparison of the bias-variance plot in loss coordinates. variance trade-off is noticeable, the plot appears skewed.
While the general idea of the bias-
Variance, loss, [10 -1
2.0
square
square
model weights, [10 5
cosine
m inrrrr cosine
6
linear
1.5
linear
Var.,
Miyor cosine
aEpelin
4
0. sprrt
1.0
0.5
2
0.5
0.3
0.3
0.1
0
0.1
2.90
2.91
2.92
2.93
2.90
2.91
2.92
2.93
Bias, mean model loss
Bias, mean model loss
(a) Bias-variance plot where bias represents the loss of the (b) variane plot where bias represents the loss of average model and variance is the variance of losses the average model and variance is the variance of for each specific shape and differrnt data permutations. model weights for each specific shape and different data permutations.
Figure 16: Miscellaneous loss-based bias- variance plots.
20

---

Published in Transactions Machine Learning Research (08/2025)
Plots in this basis are displayed Figure 15 The observed trade-off is not as vivid as in loss-based or model weights coordinates and appears be skewed version of the plots observed before.
We also tried using the loss of the average model as bias instead of the mean of the losses, and the variance of the losses (Figure 16a]. Additionally, we tried using the loss of the average model as bias and the variance of model weights as variance for each experiment (Figure 16b). While both of these plots produce reasonable results, we believe they are less interpretable than the approach used in the main part of the paper.
C.3 KL-based Space
We evaluate models with different cooldown shapes and data orderings on the validation dataset portion. For each of these modcls, we calculate the KL divergence to the average prediction of reference models (referred as single reference model predictions) and determine the variance of the predictions from the N models. That is, for vocabulary size V, total tokens count K, 8ink E RVV represents model predictions (soft-labels) trained with cooldown shape by data permutation for token and Tnk E RY represents reference model predictions (soft-labels) trained on data permutation n for token k the complete formulas used are presented in equation 5
K
bias;
KL
Tak:
Sink
K
k=1
(5)
N
K
variance;
Sink
Sik
NK
R=1
Lowered linear
Lowered linear and shapes
Predictions var.,, [10 -1
Predictions var., [10-5 3]
squaree
0.9
0.7
mirror cosine
cosine
lincar
0.5
2
0. sqrt
0.3
0.1
0.08
0.10
0.0675
0.0700
0.0725
Bias, KL to reference
Bias, KL to reference
Figure 17: Bias-variance plot for different cooldown shapes. We use the KI divergence measure bias as the average distance between predictions on the validation dataset portion of reference model (average predictions models trained for longer on the same data) and an experimental model's mean predictions. Variance is the variation in predictions for experiments with different data orderings (Equation 5 On the left, we compare only lowered linear shapes (Figure 5). which exhibit an expected decrease in variance and increase in bias with decreasing parameter value. On the right, the full range of nonlinear shapes along with selected lowered linear shapes for reference. The lowered linear shape with parameter 0.7 and the sqrt shape occupy favorable positions, achieving balanced trade-off between bias and variance.
The experimental results are presented in Figure 17
While the result is similar to the basis space used in the main text, we believe that coordinates choice is less interpretable and therefore we are using it the main version.
21

---

Published in Transactions Machine Learning Research (08/2025)
D Bias Variance Plot Reproduction
We reproduced the bias-variance plot for a smaller model with 60M parameters and for a different datasct. The observed results align with previous experiments presented in the paper.
D.1 Smaller Model
We reproduced the bias-variance plot for smaller model with 60M parameters. We used the same architec- ture as in the main experiments but modified the number of layers and the hidden dimension. The results are presented in Figure [8 and Figure 19] From the plots, it is evident that the general structure and relative positions of different shapes agree with the previous experiments.
1,25
contribution [10-2
contribution [10-2 2]
square
Variance loss
1.00
0.7
Variance loss
1.5
mirror cosine
0.75
0.5
0.9
linear
cosine
1.0
0.50
sgrt
0.3
0.5
0.25
0.5
3.75
4.00
4.25
4.50
Bias loss contribution, [10-2]
Bias loss contribution, [10-?]
Bias + Variance min
Figure 18: Bias-variance plot for different cooldown shapes for the smaller model in loss-based space. The
observed results agree with our earlier discussion in the paper.
5
[10-s]
Lowered linear
Model weights var., [10-5 6]
Lowered linear and shapes
Model weights var.,, [10
1.0
square
3
cosine
2
linear
0.7
mirro 6dBine
1
0.5
2
sqrt
0.3
0
0.1
200
220
240
260
190
200
210
220
Bias, model weights space
Bias, model weights space
Figure 19: Bias-variance plot for diffcrent cooldown shapes for the smaller model in model weights spacc. The observed results agree with our earlier discussion in the paper.
22

---

Published in Transactions Machine Learning Research (08/2025)
D.2
Different Dataset
Additionally, we reproduced the bias-variance plot using a different dataset and the same model. We used the fineweb-edu (Lozhkov et al., 2024 10 billion token subset. The results are presented in Figure 20 and
Figure
21
1.5
2.5
square
contribution [10-2]
contribution [10-2
Variance loss
0.7
Variance loss
2.0
mirror cusine
1.0
linear
0.5
1.5
0.9
0.7
0.5
0.3
1.0
sqrt
0.5
0.1
0.5
0.3
0.0
5
6
8
9
5.0
5.5
6.0
6.5
Bias loss contribution, [10-2
2]
Bias loss contribution, [10-?]
Bias Variance min
Figure 20: Bias variance plot for different cooldown shapes for the fineweb-edu datasot in loss- space. The observed results agree with earlier discussion in the paper.
Lowered linear
Model weights var., [10 =]
Lowered linear and shapes
Model weights var., [10 -5
square
7.5
cosine
0.7
lincar
5.0
mirrora.ggin
2
0.5
eqrt
2.5
0.3
0
300
325
350
375
300
320
Bias, model weights space
Bias, model weights space
Figure 21: Bias-variance plot for different cooldown shapes for the fineweb-edu datasct in model weights space. The observed results agree with our earlier discussion in the paper.
Data Points Bias Supplements
We present additional visualizations for the recency bias-based shift-deviation framework in this section. Figure 22 shows the unreduced components from equation 2
We also explore how well the shift-deviation plot (Figure 8 aligns with the bias-variance plot (Figure 6 The correlation presented in Figure 23 is clear that both perspectives align well. The plot in shift-deviation coordinates, with shift calculated over the entire training process, is presented in Figure 24 The correlation with bias is not strong, but the general structure remains the same, though skewed.
23

---

Published in Transactions Machine Learning Research (08/2025)
Perplexity improvement,
[(e(b:) -(bb)) - bias] 2
-2
c(b;) p(b:)
10
Deviation,
100
-4
10-1
-6
80
85
90
95
100
0
25
50
75
100
Batch index, %
Batch index, %
lowered lincar 0.1
sqrt
cosine
square
Figure 22: Bias-variance statistics from equation 2 On the left, unreduced and 100-batch window- smoothed shift from eq. / (calculated only during cooldown). On the right, the unreduced and 100-batch window-smoothed deviation from eq. 2 It is clear that low-deviation shapes such as lowered linear 0.1 lie close to zero the right plot, indicating more uniform data point bias. However, the left plot, they lic high, indicating low improvement over the pre-cooldown model (shift). Conversely, high-deviation shapes such as square have low shift and high deviation.
Predictions var., [10-3] a]
0.09
Papearman = 0.952,
Pepearman 0.988,
Bias, loss- based
20
0.08
Poatue 0.000
Poatue = 0.000
0.07
10
0.06
0
0.05
-5.0
-4.5
-..0
-3.5
1
2
3
4
Shift, improvement over pre-cooldown
Deviation, bias uniformity
Figure 23: Comparison between bias-variance and shift-deviation plots. The left plot shows the correlation between shift and bias, while the right plot shows the correlation between deviation and variance. The correlation is calculated using the Spearman rank correlation coefficient. High correlation is clearly visible, indicating that the two perspectives aligu well.
24

---

Published in Transactions Machine Learning Research (08/2025)
square
Pepearmon 0.345,
inris of eeiiite
Bias, loss- based
Deviation, bias uniformity
0.9
0.08
Poalue = 0.328
1.0
k/rt
0.5
0.06
0.5
0.3
0.1
0.04
-3.6
-3.4
-3.2
-3.0
-3.6
33.4
-3.2
-3.0
Shift, improvement over pre-cooldown
Shift, improvement over pre-cooldown
Figure 24: Shift-deviation plot with shift calculated over the entire training process. On the left, the shift calculated over the entire training process, and on the right, the resulting correlation with bias is presented. It is noticeable that the correlation with previous results is not strong, but the general structure remains the same. Also, unwanted skew is visible the left plot.
Weight Decay and Optimizer State
18.6
Validation perplexity
regular
18.54
18.54
18.52
weight decay
18.47 18.47 18.48
18.5
reset optimizer
18.42
18.42
18.39 18.40 18.40
18.40
18.4
18.34
18.35
18.35
18.3
sqrt
linear
mirror cosine
cosine
square
Comparison of cooldown performance without weight decay or with an optimizer state reset to usual runs, For the sqrt shape, disabling weight decay or resetting the optimizer slightly worsens performance, while for high-variance shapes, disabling weight decay can improve performance.
Figure 25:
In all the experiments, we use AdamW with a weight decay of 0.1 and resume training from the pro-cooldown stage, restoring the optimizer state. It worth examining how the cooldown stage performs without weight decay with reset optimizer state. Figure 25 presents the results of these experiments, grouped by different cooldown shapes.
No Weight Decay and Reset Optimizer State Performance Conclusions. Switching off weight decay or resetting the optimizer state negatively impacts bias-variance optimal shapes like sqrt, although the effect. is not dramatic, as evident from the absolute values of the difference.
Interestingly, switching off weight decay improves performance for high-variance shapes such as square or mirror cosine. This can be explained by introducing an "effective learning rate", discussed by [osson al. (2024a). For high-variance shapes, switching off weight decay acts as an additional decaying factor for the learning rate, altering the shape's position the bias-variance plot.
25

---

Published in Transactions Machine Learning Research (08/2025)
G
Momentum & Gradient Alignment
Motivated by the idea that the cooldown stage optimization exhibits some unusual properties, we explore how gradients align with the AdamW momentum state during the cooldown stage. To do this, we calculate the cosine angle between the optimizer's state and the corresponding gradient during training for each parameter Additionally, we collect gradient norms during the cooldown stage (Figure 26b
individually (Figure 26a
Gradient Alignment Conclusions. We observe that aligmment with momentum increases during the cooldown stage, indicating that gradients become more aligned between optimization steps: starting from negalive aligmment with momentum and moving lo slightly positive. This rising alignment suggests that the optimizer's is state and the gradients converge toward a more stable direction as training progresses through
the cooldown phase (Bocker et al. 2024 Liang et al. 2025).
cos (Wgrdd umooennum
Gradient norm
0.325
0.000
0.300
0.275
-0.050
0.250
0.225
80
85
90
95
100
80
85
90
95
100
Training steps, %
Training steps, %
(b) Gradient norm during the cooldown stage, calculated la norm of single vector across all gradients.
linear
square
cosine
sqrt
mirror cosine
(a) Average cosine similarity between gradients and the cor- responding AdamW momentum state (eap__vy) avg).
Figure 26: Gradient statistics during the cooldown stage. Confidence bands represent the 95% confi- dence interval. Values averaged using sliding window of 400 steps.
H Batch Sizes Sweep Learning Rate Selection
We tried several higher learning rates for larger batch sizes in experiments of Section 6.1 and selected the best one based on final performance. The Table 27 presents learning rate scaling factors (based on pro-cooldown stage learning rate) that were chosen for the presented experiments. The scaling roughly follows a square root scaling factor (for a batch size & times larger, a scaling of Vi is chosen) (Malladi et al. 2024).
Batch Size Learning Rate Scaling Factor
200
1.06
400
1.50
800
2.12
1600
3.00
2000
3.35
Figure 27: Learning rates scaling factor (based pre-cooldown stage learning rate) for different batch sizes.
26

---

Published in Transactions Machine Learning Research (08/2025)
Features Evolution Experiments
We investigate how output features from different transformer layers cvolve during training. To analyze this, we apply linear probes (Razzhigae el al. 2024 Tomihari & Sato 2024) during the cooldown stage at intervals of 1000 steps across all transformer hidden layers. The linear layers are trained for 2,000 steps (1 million tokens) using the AdamW optimizer on subset of the training data, then evaluated on 1,000 evaluation batches from held-out portion of the training datasct. Lincar probing weights are initialized with the current Im-head weight.
Results discussion. Figure 28 demonstrates progressive improvements in linear probe perplexities throughout the cooldown phase across layers. Figure 29 reveals the temporal dynamics of these changes, showing that the final layers exhibit the most significant relative improvement during cooldown. This pat- lern suggests that deeper layers undergo substantial feature refinement. in later training stages compared to earlier layers.
Perplexity improvement
0
3
6
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
Layer index
Figure 28: Perplexity improvement over the course of the cooldown stage for lincar probing of different hidden transformer layers. The zero layer corresponds to embeddings.
27

---

Published in Transactions Machine Learning Research (08/2025)
Reltive perplexity
1.000
Relative perplexity
1.000
0.975
0.975
0.950
0.950
0.925
0.925
0.900
0.900
80
85
90
95
100
80
85
90
95
100
Training steps,
%
Training steps, %
5
10
20
6
11
16
21
2
7
12
17
22
3
8
13
18
23
4
9
14
19
24
Figure 29: Dynamics of perplexity improvement over the course of the cooldown stage for linear probing of different hidden transformer layers. The zero layer corresponds to embeddings. Perplexity relative to the start of the cooldown stage is reported. Left uses sqrt cooldown shape, right uses square cooldown shape.
28

---

Published in Transactions Machine Learning Research (08/2025)
Loss Landscape Plots
Cooldown start (81.8% steps)
36.32
opt Optimizatii nioiiatioon
20.64
-0.3
20.13
optiniza
0.0
70.00
19.94
19.99
-3.0
-0.2
-1.5
0.0
0.3
Global
-0.1
0.0
0.15
Global
Adam steps
1.5
A Adam steps
0.1
3.0
0.2
Cooldown mid (90.9% steps)
22.45
18.86
18.54
18.54
Global optimization
0.3
Global optamization
-0.15
-3.0
00.2
Adam steps
-1.5
0.0
Adam steps
-0.1
0.00
0.0
0.0
1.5
0.3
0.1
0.15
3.0
0.2
Cooldown end (100.0% steps)
19.34
18.24
17.86
17.86
-3.0
-0.2
1..5
0.1
Ida steps
0.0
0.0
1.5
-0.3
0.1
-0.15
Global optimization
0.0
Global optimization
0.00
3.0
0.2
0.3
0.15
current point
plot min point
Figure 30: Plots of the loss landscape at different cooldown points. The Z-axis represents perplexity on a portion of the validation data. On the left, plots from larger scale are displayed; on the right, the same plot points are shown at smaller scale.
29


Reviews:
[{"review_id": "NTOpBoklZu", "reviewer": "Reviewer_ooiE", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "All the experiments are meaningful, rigorous, and inspiring. The authors have good intuition on what might be happening during training, are knowledgeable in the literature, and have conducted concrete experiments for proof.\n\nA high-level weakness is that the paper lacks a clear, overall story that can closely connect the 4 parts all together. Specific to each part, I would like to have some further discussions:\n\n1. One big weakness with the current \"bias-variance trade-off\" perspective, is that the paper lacks an explicit bias-variance decomposition. Without such a decomposition, the definitions of bias and variance in the paper can be arbitrary (although, I understand the intuition and the definitions seem reasonable to me) but not well justified. There might indeed be a trade-off between \"bias\" and \"variance\", but it could also be due to some artificial math trick hidden in the definitions which make it look like so -- especially, when we look at similar bias-variance plots in the model weights space, as shown in Section C.1, the lowered linear family is no longer on a Pareto front -- which suggests that things can clearly change when the definitions of bias and variance are different.\n\n2. Comparing model averaging with longer training: I was a bit confused when reading this section, because Figure 10 shows that averaged models consistently perform better than single runs, but the narration says \"all averages perform worse than a simple longer run...\"; then I realized that the paper is talking about the 52.8k sqrt run shown as a dashed line in Figure 10 (which is not very intuitive)...\n\n3. When the hyper-parameters are changed, what are the error bars of validation perplexity, due to different data shuffles?\n\n4. The river valley structure could also be explained as: because you are connecting the pre-cooldown checkpoint to the final model after cooldown as the \"global optimization direction\", at the end of the cooldown that direction is going to reach some valley because the training converges. On the other hand, a local optimization step (or 10 optimization steps) is likely orthogonal to the global direction because the parameter space is very high-dimensional. So what is special about this \"river valley\" perspective?", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "Yes", "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": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "This paper is a collection of loosely connected investigations on the training dynamics of the cool down stage of pre-training a language model. The investigations can be roughly divided into 4 parts:\n\n1. Trying different decaying schedules of the learning-rate, and connecting the performance to a bias-variance trade-off perspective. Among usual schedules such as linear, cosine, sqrt, etc., the authors also tried a family of \"lowered linear\" schedules, and show that the family demonstrates a negative correlation curve on the bias-variance plot. Moreover, this family appears to be on a Pareto front of the bias-variance plot, among all schedules being investigated. It suggests that there is indeed a trade-off between the bias and variance defined in the paper, which might control the training performance.\n\n2. Comparing model averaging with longer training.\n\n3. Investigating other Adam hyper-parameters such as batch size, beta1 and beta2 for the cool down stage. Especially, the finding that larger beta2 at the cool down stage can improve performance, might be useful empirically for future practice and study.\n\n4. A nice visualization of the loss landscape, clearly showing a \"river valley\" like structure, which supports previous conjectures.", "claims_and_evidence_yn": "Yes", "claims_explanation": "", "audience_yn": "Yes", "audience_explanation": "", "requested_changes": "I don't want to block this paper, because it contains quite a lot of meaningful, inspiring findings; but it would definitely be better if the authors can establish more theoretical justifications of the bias-variance trade-off, such as a formal bias-variance decomposition, and more discussions on the potential logical connections between the label distribution space, model weights space, and the loss space -- to me, the similar bias-variance structure observed in the model weights space is not \"surprising\"; rather, it looks like a cause of the observed phenomena.", "broader_impact_concerns": "None.", "additional_comments": "", "extra_scores": {}, "extra_text": {}}, {"review_id": "n13tXk6aYr", "reviewer": "Reviewer_4hbb", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "The article gives a good overview of the learning rate scheduling literature, as well as extensive numerical results, leading to interesting insights and practical guidelines on how to train LMM. Furthermore, a bias-variance framework is proposed to analyze the validation perplexities of LMM, aiming to correlate the improved performance with proposed bias/variance quantities. \n\nThe proposed bias and variance definitions are, however, not so clear. It is hard to understand the central equation (1), where the definition of K is not clear. The norm with axis=k is not a common mathematical definition, therefore it needs to be clarified. Conceptually, there is a lack of connection between the bias and variance, as the bias is based on KL divergence, while the variance is based on MSE. From this perspective, the proposed framework is theoretically insufficient, and further explanation of the proposed definitions is needed. \n\nIt seems that there is no direct connection between the AdamW part (section 6) with the bias-variance trade-off. Therefore, the contributions seems to be dis-connected.", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "No", "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": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "This article studies how learning rate scheduling affects the performance of large language models. By focusing on the cooldown phase to fine-tune a pretrained model, a bias-variance trade-off analysis is given to show that the optimal performance is often achieved when a bias-variance balance is achieved. This phenomenon is observed based on various decay scheduling including warmup-stable scheduler. The impact of batch size as well as AdamW hyperparameter tuning is also studied, showing the importance of optimization in this training phase. From a loss landscape perspective, the article finds a river valley shape validation loss during the cooldown phase.", "claims_and_evidence_yn": "No", "claims_explanation": "", "audience_yn": "Yes", "audience_explanation": "", "requested_changes": "- Add further theoretically justification of the proposed framework.\n- Rewrite eq (1), make the definition of each term clearer. Is k the index a sequence of length K? \n- A definition of optimal bias/variance trade-off is needed to understand better the results in Figure 7 and 8.\n- eq (2) is very hard to understand. What does that mean that “i in cooldown”? This is not mathematically meaningful.\n- The loss landscape in Figure 13 could be made clearer to illustrate its connection with the idea given in Figure 1.", "broader_impact_concerns": "NA", "additional_comments": "", "extra_scores": {}, "extra_text": {}}, {"review_id": "3bzwHPSD4S", "reviewer": "Reviewer_v6TE", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "**Strength**\n1. This paper studies an important problem and provides novel insights. The scope is tightly focused. By looking into the cooldown phase in isolation, the paper offers actionable insights for practitioners, especially those tuning large models under fixed compute.\n\n2. The bias-variance framework is intuitive and well-motivated. The authors go beyond simple performance comparisons to propose a conceptual and empirical tool that aids in understanding why certain cooldown shapes are preferable.\n\n**Weakness**\n1. The recency bias section, while conceptually interesting, feels underdeveloped. The explanation of why recency bias transitions from sharp to U-shaped is somewhat speculative \n2. It is unclear whether the difference in validation perplexity between different cooldown shapes is significant enough to study the problem or not. If the difference is too negligible, people would just choose the random one to use.\n3. Would it be possible to include downstream task performance in the evaluation?", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "Yes", "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": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "This paper investigates the role of the cooldown stage in the Warmup-Stable-Decay (WSD) learning rate scheduler, a phase often overlooked in prior work. The authors introduce a bias-variance framework to analyze how different cooldown shapes impact model performance, finding that shapes like sqrt and lowered linear 0.7 offer strong trade-offs. They complement this with experiments on optimizer hyperparameters and visualize the loss landscape to support the \"river valley\" hypothesis. The study provides practical guidance for tuning the cooldown phase to improve transformer training outcomes.", "claims_and_evidence_yn": "Yes", "claims_explanation": "", "audience_yn": "Yes", "audience_explanation": "", "requested_changes": "1. The discussion on model averaging versus longer training stops short of offering a broader conclusion. It would be helpful to include a take-home message about when averaging is or isn't worthwhile.\n2. The notion of recency bias is intriguing, but currently, the conclusions feel tentative. Consider including a deeper analysis or a concrete hypothesis to explain the emergence of the U-shaped perplexity profile during cooldown. This would sharpen the section’s contribution.\n3. Some of the figures, particularly in Section 4, are rich but a bit dense. For example, Figure 9 would benefit from more precise legends.", "broader_impact_concerns": "No ethical concerns.", "additional_comments": "", "extra_scores": {}, "extra_text": {}}]

Author rebuttal:

Decision: Accept (poster)