When, Where and Why to Average Weights?
Authors: Niccolò Ajroldi
Antonio Orvieto
Jonas Geiping
Venue: ICML 2025 poster (2025)
Keywords: Weight Averaging
Checkpoint Averaging
Optimization
Learning Rate Schedule.
Primary area: optimization->large_scale_parallel_and_distributed

Abstract: Averaging checkpoints along the training trajectory is a simple yet powerful approach to improve the generalization performance of Machine Learning models and reduce training time. Motivated by these potential gains, and in an effort to fairly and thoroughly benchmark this technique, we present an extensive evaluation of averaging techniques in modern Deep Learning, which we perform using AlgoPerf, a large-scale benchmark for optimization algorithms. We investigate whether weight averaging can reduce training time, improve generalization, and replace learning rate decay, as suggested by recent literature. Our evaluation across seven architectures and datasets reveals that averaging significantly accelerates training and yields considerable efficiency gains across all considered workloads, at the price of a minimal implementation and memory cost, while mildly improving generalization. Finally, we explore the relationship between averaging and learning rate annealing and show that combining the two achieves optimal performance.

When, Where and Why to Average Weights?
Niccolò Ajroldi
Antonio Orvieto
123 Jonas Geiping
123
Abstract
(Athiwaratkun et al., 2018; Li al., 2023; Sanyal al., 2023). Recent studies have also explored the connection between learning rate decaying and weight averaging (San- dler et al., 2023; Hägele et al., 2024), and used the latter to develop schedule-free algorithms (Defazio et al., 2024b).
Averaging checkpoints along the training trajec- tory is simple yet powerful approach to improve the generalization performance of Machine Learn- ing models and reduce training time. Motivated by these potential gains, and in an effort to fairly and thoroughly benchmark this technique, we present an extensive evaluation of averaging tech- niques in modern Deep Learning, which we per- form using AlgoPerf (Dahl et al., 2023), a large- scale benchmark for optimization algorithms, We investigate whether weight averaging can reduce training time, improve generalization, and replace learning rate decay, as suggested by recent liter- ature. Our evaluation across seven architectures and datasets reveals that averaging significantly accelerates training and yields considerable effi- ciency gains across all considered workloads, at the price of a minimal implementation and mem- ory cost, while mildly improving generalization. Finally, we explore the relationship between aver- aging and learning rate annealing and show that combining the two achieves optimal performance.
Table 1: Estimated training cost for on the AlgoPerf collection of models and tasks. Even for industrial-scale tasks like benchmark workloads, weight averaging reliably reduces compute costs.
Nadam +LAWA +EMA
GPU-Hours 612
550
541
In this work, we present the largest evaluation of averaging techniques in modern Deep Learning, which we perform using AlgoPerf (Dahl et al., 2023), collection of large- scale workloads and architectures developed to provide a unified benchmark for optimization algorithms. Framing our analysis in this setting provides a carefully designed evaluation framework, enables comparisons against strong, heavily tuned haselines, and allows drawing broad and ro- bust conclusions.
1. Introduction
Building on previous work, we investigate the following questions. (i) Can weight averaging reduce training time across multiple models and tasks? (ii) Does averaging check- points improve the generalization performance of existing optimization algorithms? (iii) Is weight averaging merely a proxy for shorter learning rate decay schedule, and can fully replace learning rate decay?
Training Deep Learning models is both resource-cxpensive and comeeminsuming.... A principled and simple approach to speed up training, dating back to Polyak (1990) and Ruppert (1988), involves averaging the model's 's weights across train- ing iterations. This can either be performed online during training or post hoc by averaging checkpoints, effectively creating an ensemble of models at minimal additional cost. Previous studies have shown that weight averaging (WA) techniques can improve generalization (Merity et al., 2017: Gupta al., 2020; Kaddour, 2022; Melis, 2023), increase robustness -Morales-Brotons et al., 2024), smooth loss land- scapes (Izmailov et al., 2019), and accelerate convergence
Our contributions are follows.
We show that averaging can significantly accelerate training dynamics across seven different models and ar- chitectures and estimate a 12% reduction in GPU-hours to train the entire AlgoPerf suite up to the validation target compared to our baseline. We find this effect to be consistent across hyperparameters and observe encouraging speed-ups ups even on a more sophisticated optimizer like Distributed Shampoo (Shi et al., 2023).
'ELLIS Institute Tübingen Max Planck Institute for Intelligent Systems, Tübingen, Germany 'Tübingen Al Center. Correspon- dence to: Niccolò Ajroldi <niccolo@ ellis.euu>.
Proceedings of the 4nnd International Conference Machine Learning. Vancouver, Canada. PMLR 267, 2025. Copyright 2025 by the author(s).

---

When, Where and Why to Average Weights?
2. Beyond efficiency gains, we demonstrate how aver- aging achieves improved generalization across all the considered workloads, and show that combining WA with learning rate annealing yields optimal results.
decoder-only language models (Sanyal et al., 2023). A valu- able alternative moving window averaging techniques like LAWA is Exponential Moving Averaging (EMA) (Li et al., 2024; Morales-Brotons et al., 2024; Arpit et al., 2022), It retains similar advantages of rolling window averaging, and constitutes an indispensable technique for high-quality image synthesis models such as GANs and diffusion models (Yazıcı et al., 2019; Song et al., 2021; Karras et al., 2024).
3. Finally, we demonstrate how averaging checkpoints can act as a proxy for shorter learning rate decay, but that it cannor fully replace learning rate schedules, at least within the explored variants, addressing an important question about the role of these techniques (Hägele et al., 2024; Defazio et al., 2024b).
Connection to learning rate annealing. Classical stochastic smooth convex optimization rates showcase a tight link between WA and learning rate annealing, suggest- ing a practical interplay between these techniques (see e.g. Theorem 5.3, in Garrigos & Gower (2023)). Intuitively, av- eraging weights along the training trajectory reduces noise and might act proxy for learning rate decay. In fact, Sandler et al. (2023) proved the theoretical and empirical equivalence between WA and decaying learning rate for SGD. However, despite this appealing result, modern Deep Learning models are still predominantly trained with learn- ing rate annealing (Hägele et al., 2024), even when main- taining an EMA of model weights (DeepSeekAI, 2025). A recent study by Defazio et al. (2024b) specifically inves- tigates this connection, proposing an approach that fully replaces learning rate schedules with iterate averaging and demonstrating strong performance on the same benchmark used in this analysis. Whereas Defazio et al. (2024b) incor- porates averaging directly into the optimization procedure, we explore a different flavor of averaging, where the aver- aged weights do not influence the updates-ake to Polyak averaging. SWA, LAWA, and EMA.
The paper is organized as follows:
Section 2 reviews related work weight averaging; Section 3 discuss the methodol- ogy and experimental setup; Section 4,5, 6 and 7 present our findings; Section 8 addresses the limitation of our analysis and Section 9 concludes.
2. Related Work
The idea of averaging iterates along a stochastic optimiza- tion trajectory dates back to Polyak (1990) and Ruppert (1988), and is often referred to as Polyak-Ruppert averag- ing. It has been extensively studied in the stochastic ap- proximation framework (Polyak& & Juditsky, 1992; Bach & Moulines, 2013; Neu & Rosasco, 2018; Lakshminarayanan & Szepesvari, 2018), and is a common technique to de- rive convergence guarantees in both convex and nonconvex optimization (Garrigos& & Gower, 2023).
Deep Learning applications. Weight averaging tech- niques have seen extensive use in the Deep Learning commu- nity, often without explicit recognition or emphasis, despite their effectiveness. Influential works such as Szegedy et al. (2016), Vaswani et al. (2017), and Merity et al. (2017) have demonstrated their ability to enhance model performance and mitigate overfitting.
Model soups. Finally, a different but notable approach that leverages the benefits of averaging is model soups (Wortsman et al., 2022). In this case, multiple models are trained with different hyperparameter configurations and later aggregated, resulting in an ensemble with improved accuracy and robustness.
SWA. The work of Izmailov et (2019) sparked renewed interest in weight averaging by demonstrating how averag- ing points along the SGD trajectory leads to wider minima and improves generalization performance. Their approach, Stochastic Weight Averaging (SWA), has since been applied to semi-supervised learning (Tarvainen & Valpola, 2018; Athiwaratkun et al., 2018), 8). low- cow-precision training (Yang et al., 2019), domain generalization tasks (Cha et al., 2021; Arpit et al., 2022), and meta-optimization (Li et al., 2023).
In this work, we demonstrate the benefits of weight aver- aging techniques on a challenging optimization benchmark (Dahl al., 2023), hoping to encourage broader adoption of these methods in training large-scale Machine Learning models.
3. Experimental Setup
Comparing optimization algorithms. When evaluating optimization algorithms, two strategies are possible:
LAWA and EMA. In the original formulation of SWA, a pretrained model is trained with cyclical or constant learn- ing rate, and multiple checkpoints are collected and later averaged. Kaddour (2022) proposed Latest Weight Averag- ing (LAWA). an online algorithm that averages the latest checkpoints in rolling window, showing significant speed- ups on vision and language tasks. Further modifications of LAWA demonstrated notable boosts in pretraining modern
(A) Fixing challenging target loss value and comparing the runtime needed to reach it.
(B) Comparing generalization performance within a fixed budget, either specified in number of steps or wall- clock time.
2

---

When, Where and Why to Average Weights?
NadamW
Nadam LAWA
Nadam EMA
0.30
MAP
DLEU
Validation MAP
Validation BLEU
0.20
28
0.15
0.10
0.05
22
10000
20000
30000
40000
50000
(a) OGBG
(b) WMT
0.80
Taryet
Validation WER 4
Valiatioo Accuray
0.75
0.3
0.65
0.60
0.1
Torget WFR
0.55
10000
50000 60000 70000
80000
25000)
50000
75000
100000
125000
150000
175000
(c) Librispeech Conformer
(d) Imagenet VIT
Figure 1: Weight averaging speeds up training across all considered workloads. The averaged schemes consistently achieve better performance during training, reaching the validation score target faster than the baseline algorithm. We display the validation score against the number of iterations (optimization steps) across different workloads; the dotted line represents the target score on each workload.
In this work, investigate whether weight averaging im- proves existing optimization algorithms in either of the two described frameworks.
timization algorithms. The benchmark is composed of eight workloads, each defined by datasel, model architecture, and a predefined target metric on a held-out set, designed to represent optimal performance on such a workload. We note that AlgoPerf has been developed following option (A) and that significant effort has been dedicated to deriving challenging target scores. Thus, provides ideal setting to evaluate the impact of weight averaging techniques ap- plied to strong optimization algorithms and compare their performance against heavily tuned baselines. We also make use of the same datasets and architectures to score algo- rithms by means their generalization capabilities (when following option (B)). We consider the following workloads from the AlgoPerf suite: (i) DLRMsmall model on Criteo ITB dataset for click-through rate prediction; (ii) U-Net FastMRI for medical image reconstruction; (iii) VIT on ImageNet- 1k for image classification; (iv) GNN model on OGBG for graph-based prediction; (v) a Transformer-Bi on WMT for machine translation; (vi) a Conformer for speech recognition; (vii) DeepSpeech model on LibriSpeech for speech-to-text. We exclude ImageNet- Nesseettt from this analysis because, shown in Dahl et al. (2023), opti- mization algorithm successfully solves this task alongside
Weight averaging. We explore two flavors of weight aver- aging discussed in Section 2: LAWA (Kaddour, 2022) and EMA. For both approaches, previous information about the model history. In the first case, we update circular queue of length L,, which stores recent checkpoints, while in the latter case, we maintain an exponential moving average of the model parameters, with coefficient 7. We update the averaging buffer every iterations, and explore the effect of such hyperparameter together with the value of L and 7. In most scenarios, we save checkpoints of the baseline algorithm and run averag- ing schemes offline, but when testing collecting consecutive checkpoints, we use online versions LAWA and EMA to avoid excessive disk storage. We discuss implementation details of the algorithms, along with their practical implica- tions and considerations, in Appendix E.
update buffer that stores
AlgoPerf. We conduct our analysis on AlgoPerf (Dahl et al., 2023), a large suite of Deep Learning workloads, which provides a comprehensive benchmark for testing op-
3

---

When, Where and Why to Average Weights?
Relative Steps Nadam (%)
+LAWA
FEMA
NadawW
100%
8076
6096
40%
CritcolTB
FastMRI
VIT
Conformer
DeepSpocch
OGBG
WMT
Figure 2: LAWA and EMA speed up training across several architectures and datasets. Both averaging schemes consistently outperform the baseline, achieving average the benchmark target score using 82% of the steps required by NadamW. We estimate 12% reduction in GPU-hours to train the entire AlgoPerf suite of workloads with respect to NadamW.
the other workloads, leaving no baseline on which to build. We refer to Appendix A for more details on the workloads.
et al., 2024), or the end of each epoch (Kaddour, 2022), and investigated the optimal horizon and update frequency on a single task (Sanyal et al., 2023). We explore combinations and interactions of these variables across architectures and objectives, and compare LAWA with EMA.
Baseline optimizers. We build on top of existing opti- mization algorithms that perform well on AlgoPerf, enhanc- ing them with LAWA or EMA. Unless otherwise specified, we use NadamW (Timothy, 2016; Loshchilov & Hutter, 2017a) as our baseline and name the resulting algorithms Nadam LAWA and Nadam EMA, respectively. Ad- ditionally, we investigate the potential benefits of combin- ing a higher-order optimizer, such as Distributed Shampoo (Shi et al., 2023), with these averaging methods, exploring how this approach might further improve training efficiency. Given the high computational cost of the benchmark, we conduct most of the analysis using the best performing hy- perparameters in Dahl et al. (2023) and Kasimbeg et al. (2025): this gives strong baselines, avoiding the burden of expensive hyperparameter tuning, Additionally, we ab- late on the role the learning rate in each analysis and study in more detail the consequences of changing the learn- ing rate schedule. Given this reference algorithm, we add LAWA or EMA on top it, tuning only the hyperparame- ters of the averaging scheme and leaving the other baseline hyperparameters fixed, including learning rate schedule and weight decay. Unless otherwise specified, or when explicitly ablating on it, we train using a cosine learning rate schedule (Loshchilov & Hutter, 2017b).
To control for variability in model initialization and data shuffling, we repeat experiments for three different seeds and report the mean and standard deviation.
4. Speeding up Training
We ask whether equipping strong baseline with LAWA EMA can reduce the number of steps required to reach predefined validation target (option A in Section 3). To evaluate this, we track the number of steps required to reach the benchmark score on the held-out set. The target scores, derived Dahl al. (2023) by heavily tuning and compar- ing different optimization algorithms, represent desirable performance each workload.
Efficiency gains of averaging. When looking for scheme that reduces training time, we observe that both LAWA and EMA can significantly speed up training with respect to the considered baseline (Figure 1). If target validation score is known in advance, weight averaging can effectively be employed to trigger early stopping and sig- nificantly reduce computational costs. We estimate a 12% reduction of GPU-hours for training AlgoPerf using LAWA top of NadamW.
Optimal averaging horizon. The impact of an averag- ing scheme largely depends on the horizon over which it applied and on the frequency at which checkpoints are collected. A long horizon may prioritize outdated model checkpoints and generalize worse, whereas short one may be suboptimal result in serious computational overhead. At the same time, updating the averaging buffers every step is only possible in the online version of LAWA and EMA, and collecting checkpoints too rarely might reduce the bene- fits averaging. Previous works have collected checkpoints each training step (Morales-Be Brotons et al., 2024; Hägele Beorales-Broooos
Averaging horizon. We further investigate the relation- ship between the optimal averaging window and training efficiency. Figure 3 reveals consistent trend across work- loads: whereas moderate weight averaging significantly reduces the number of steps required to reach the valida- tion target, both excessively small and excessively large horizons can be suboptimal, resulting in minor speed-ups Importantly, we find that LAWA is often highly stable and

---

When, Where and Why to Average Weights?
CritcolTB
Conformer SeepSpeech
OGBG
NudamW
NudbuW LAWA
Nodam EMA
FRstMRT
WMT
1.0
Proportion of Steps to Reach Target
0.8
0.6
Ceeniing
Enanding Rate
(a) OGBG
(b) CriteolTB
0.4
Figure Averaging weights of a suboptimal baseline. We train using NadamW, varying the top learning rate, and compare the number of steps needed to reach the validation target. We report mean and standard error across random seeds.
0.2
0.0
10-1
10-3
10-2
10P
Averaging Window as Proportion of Steps
Figure 3: Impact of averaging horizon on training efficiency. Moderate weight averaging generally reduces the number of steps needed to reach the validation target. However, excessive averaging (rightmost region) or overly small av- eraging windows (leftmost region) may diminish gains or hinder progress. We use LAWA for this analysis, and define the averaging window proportion ExL (æ-axis), where T denotes the total training budget, is the number of steps between consecutive checkpoints, and I. is the span of the averaging window. For each workload, we fit and display a second-order polynomial: = ax² + bæ + c.
Similarly to Sanyal al. (2023), we observe a slight advan- tage on OGBG with higher learning rates, but do not see a similar behavior for Criteol IBB; nevertheless the latter WA remains effective at larger learning rates, in contrast to Nadam W, which performance degrades faster. Interestingly, we notice that, as long as the baseline algorithm is able to reach the predefined target, weight averaging consistently speeds up convergence. instead the learning rate is too big (or too small) and NadamW does not achieve the target in the maximum number of iterations, neither do the aver- aged checkpoints. This suggests that while averaging can significantly accelerate training, it has a limited impact on improving generalization-a a point we explore further in the next section.
forgiving with respect to the choice of its hyperparameters, considering that for each workload broad range of hori- zons effective. This robustness is particularly evident in OGBG and FastMRI, where performance remains rela- tively stable across a wide range of averaging windows. We attribute this to their long training horizons, as also noted in Kasimbeg et al. (2025). Furthermore, our analysis suggests a generalizable optimal averaging horizon across workloads: an averaging window around 1 1% of the total training bud- get consistently yields optimal results. We further explore the effect of LAWA and EMA hyperparameters on training speed-ups in Appendix B.
Averaging Shampoo. Finally, we explore whether more sophisticated optimizer can benefit from weight averaging. We choose Distributed Shampoo (Shi et al., 2023) as the best scoring algorithm on the AlgoPerf inaugural competi- tion (Kasimbeg et al., 2025), which provides a significant speed-up over NadamW. We report in Figure 5 the number of steps required to reach the validation targets when train- ing an mncoder-decoder Transformer on WMT, a ViT on ImageNet and a U-Net on FastMRI. We equip Distributed Shampoo with LAWA and EMA, and observe that weight averaging consistently enhances performance, even when applied to an already highly efficient optimizer like Sham- poo, reaching the targets sooner and occasionally reducing the variability of the baseline optimizer. This emphasizes that the benefits of weight averaging extend beyond compen- sating for slower or less effective optimizers, and indicates that averaging offers inherent advantages, regardless of the adopted optimizer.
Stability to hyperparameter tuning. An important ques- tion is whether these gains are to be observed only at peak tuning of an optimizer, or if they hold in generic hyper- parameter settings. We investigate this circumstance by variating the top learning rate of NadamW W and averaging the correspondent checkpoints. We note that LAWA and EMA maintain their efficiency gains across all considered learning rates, reliably achieving the target score faster than the baseline optimizer, as shown in Figure 4. Averaging re- mains effective in improving efficiency, even when applied to a suboptimal optimization algorithm. This is encourag- ing, as it suggests potential benefits even in the absence of a strong baseline or when searching for one is too costly.
5. Improving Generalization
When a predefined target is not known a priori, or a compu- tational budget is fixed and one does not need to stop training
5

---

When, Where and Why to Average Weights?
6000
Stepe Validation Target
Stepe to Valislation Targel
Steps to Validation Target
80000
5000
100000
17000
100000
75000
4000
00000
72500
3000
Shampoo
+LAWA
+EMA
Shanpoo
+LAWA
+EMA
Shampoo
(a) WMT
(b) ImageNet VIT
FastMRI
Figure 5: Weight averaging provides significant speed-ups also when applied on top of a more sophisticated optimizer like Distributed Shampoo. Averaging checkpoints through LAWA or EMA reduces the number of iterations required to reach the validation targets on the considered workloads.
Table Validation performance across workloads. Both averaging schemes show minimal improvement over the baseline, always matching or slightly surpassing its performance. We report the mean and standard deviation across 3 seeds, for all workloads but Imagenet ViT, due to its high computational cost. We observe notable gains in the WMT and OGBG workloads. We do not find consistent significant difference between the two averaging schemes, we noice that better hyperparameter configuration for LAWA and EMA are possible and may lead to larger improvmentes.
CONFORMER WER
DEEPSPEECH
CRITEOITB OGBG
FASTMRI SSIM 1
WMT
VIT
WER 4
Loss
MAP 1
BLEU +
Acc 1
NADAMW
0.08849 10.01065
0.11894
10.00514
0.1235 1.000ss
0.3012
10000118
0.72690.
10.00021
31.
10.06373
0.79048
+LAWA
0.08845,, L0.01067
0.11613
-0.01355
0.1235
_0,0000
0.3119
00.00997
0.72705,
10.00020
31.429
10.06098
0.79112
+EMA
0.08841, +0.01073
0.11739
20.0052
0.1235_a.comm
000000
0.30S84a.00025
#0.00253
0.72707
+0.00020
31.457
+0.08598
0.79232
early, weight averaging might be appealing option to im- prove generalization. We investigate this setting (Option B in Section 3) by fixing step budget and comparing the best validation score achieved during training by the baseline algorithm and by its averaged version.
alization performance across different hyperparameter con- figurations. This is a common scenario, that may occur when an optimal baseline is not available, or when it is too expensive to search for one. We variate the baseline learn- ing rate and report the achieved validation score with and without averaging. We observe in Figure that the averaged checkpoints closely track the performance of the baseline algorithms, sometimes providing marginal gains. Unlike Sanyal et al. (2023), do not notice inherent benefits when training at higher learning rates.
WA leads to moderate performance gains. We report in Table 2 and Figure 8 the results of averaging in terms of gen- eralization performance. We find that using LAWA or EMA on top of a learning rate schedule consistently improves over the baseline optimizer, achieving slightly better valida- tion scores across all the considered workloads. We notably observe significant improvements on the WMT workload. As previously reported in Kaddour (2022) and Sanyal al. (2023), in the early stage of training, averaging schemes pro- vide considerably better performance than the underlying optimization algorithm, but this gap shrinks towards the end of training. We argue that this behavior closely linked to the use of WA on top of learning rate schedule, and that the benefits of averaging diminish later training due its similarity to learning 2023). We explore this topic in more details in the following section.
6. Averaging as a Proxy for LR Decay
To achieve optimal performance on most Deep Learning tasks, it is common practice anneal the learning dur- ing training (Defazio et al., 2024a). Despite its ubiquity, this practice introduces extra hyperparameters, complicates resuming from the latest checkpoint (Singh et al., 2025) and significantly increases training costs (Hägele et al., 2024). In contrast, averaging weights along the training trajectory in- tuitively reduces noise and might act as a proxy for learning rate decay (Sandler al., 2023; Defazio al., 2024b).
annealing (Sandler al.,
In this section, investigate how WA affects performance under various cosine LRR schedules, whether it can reliably proxy a short LR decay, and whether it can fully replace it.
Stability to hyperparameter tuning. In line with the pre- vious section, we investigate how averaging affects gener-
6

---

When, Where and Why to Average Weights?
NadamW
Nadam + LAWA
NadamW + EMA
0.31
0.1219
0.30
0.20
Valicltbon Loes
0.1218
Validation
0.1217
0.27
0.1236
0.25
0.1210
0.25
10
10
10
Learning Rate
Learnnng Rate
(b) Criteo1TB
(a) OGBG
Figure 6: Averaging a suboptimal baseline. We train with NadamW, varying the top learning rate, but still decaying it to zero, and compare the validation performance when averaging weights. We report mean and standard error across random seeds.
0.25
Nodscaw Short Scheviele
0.25
Optimiier
Final
Nodazaw Long Schodule
NationW
Valddttoon WER
0.20
Nodazow Long Schedule -EMA
Nodacmw Long Schedole +LAWA
Validation WER
0.20
Nation &KMA
NadamnW +LAWA
0.5 thear
0.15
0.15
0.10
10000
20000
30003
40000
50000
60000
70000
50000
0
10000
20000
30000
4000
50000
60000
70000
80000
Steps
Steps
(a) Weight averaging of long (anncaled) training run performs similarly to training with shorter learning rate schedules. EMA and LAWA allow to materialize better-perfomning model for free withouth the need to lower the learning rate.
(b) Averaging checkpoints of training rums with different annealing strategies. The benefits of WA diminish when annealing the learning zero. Here TJmaz 0.00175, which correspond the best configuration of Nadam on Librispeech Conformer.
Figure 7: Weight averaging on top of different learning rate schedules on Librispeech Conformer.
Averaging vs LR scheduling. In line with previous work, we compare checkpoint averaging over long traning run against training with a shorter learning rate schedule. We systematically observe that the validation performance of averaged checkpoints closely tracks training run with faster learning rate decay (Figure 7a and Figure 12). This behavior, although pointed out in earlier studies, consis- tently emerges across the diverse model architectures and datasets considered for this analysis, and explains the con- siderable speed-ups observed in Section 4. Since checkpoint averaging acts proxy for cooled-down learning rate, WA enables access better model during training, thus achieving the target score faster. We note that more sophisti- cated annealing strategies are possible (Hägele et al., 2024), and while the averaged model consistently approaches the Pareto frontier of loss versus training time (Portes et al., 2022), may not always lie on it. Nevertheless, its simplic- ity and effectiveness make averaging a highly practical tool in surge-scale training, providing access to stronger model with minimal computational overhead.
Given the promising effect of averaging as an alternative to learning rate annealing, we examine its impact under different annealing strategies. We use a cosine learning rate schedule with maximum value of Mmaa and explore three variations: no annealing, annealing the learning rate to half of 7maz, and annealing to zero (Figure 7b). Despite yielding significantly hetter validation scores during train- ing across all learning rate schedules, the final validation performance of WA is strongly influenced by the annealing strategy. When little or no annealing is applied, WA pro- vides substantial improvements; however, when the learning rate is fully annealed to zero, WA converges closely to the annealed model, suggesting that its benefits diminish op- timization naturally reaches a well-converged solution. This observation is consistent with our earlier findings, where averaging improves performance during training (Section 4) but provides minimal generalization gains when learn- ing rate schedule is used (Section 5), and further supports the previously discussed relation between averaging and learning rate scheduling.
7

---

When, Where and Why to Average Weights?
AlgoPerf Target
Vadication WiEt
Validccion WER
Validation Lots
0.12
0.12400
0.24
0.22
0.12550
Nadami
+1A decay
-LAWA
LAAAA aii doeny
NadooW
+LA dauay
+LAWA
+LAWA +tit dooay
NadumlW
+IN decne
-LAW
+LAAA
+1R deeny
(a) Conformer
(b) DeepSpeech
(c) CriteolTB I
Valiobiion Avcuricy' T
0,50
Validatonn MAP 1
f
Validation BLBU
8
00.31
0 778
30
0.75
Vallelationn
0.29
0.724
0.70
0.28
Nadwoa
+LR desay
-LAWA
LLWA +LR doons
NadamdV
+EAt dhaap
-LAWA
LLWA
NachoW
+LR docas
+LAWA
+LAWA +idt Aaeay
Nadom
+LR deces
+LAWA
LAWA
+Li decar
+Ld docay
(d) WMT
(e) VIT
(1) OGBG
(g) FastMRI
Figure 8: Generalization performance when training with Nadam W, without (+) and with (A) learning rate decay, versus training with NadamW +LAWA, without (o) and with (A) learning rate decay. Despite significantly improving the performance of the baseline, weight averaging alone cannot generally match the validation performance of using learning rate annealing schedule, and can rarely reach the target validation score. Instead, we observe that combining the two yields slightly better final validation scores throughout the considered workloads. We notice that OGBG and FastMRI the only workloads where LAWA can provide a good substitute for LR decay, and attribute this behavior to the relatively easy target, very generous step budget, and high variability of these tasks, as also noted by Kasimbeg et al. (2025). We report the mean performance across 3 random seeds. The red dotted line displays the benchmark target scores.
Can averaging replace LR decay on AlgoPerf? We thor- oughly compare averaging with learning rate scheduling across the entire AlgoPerf suite, reporting our findings in Figure 8. Despite heavily tuning the averaging algorithms, we find that fully replacing the learning rate decay with averaging consistently yields inferior performances than annealing the learning rate, with the only exception of CitcoITB and OGBG, where WA can achieve the target without lowering the learning rate. This result suggests that averaging schemes which do use the average compute optimization updates cannot fully replace learn- ing rate annealing, and more advanced techniques like the one introduced by Defazio et al. (2024b) are necessary to substitute LR decay. Nevertheless, we argue that these av- eraging schemes can still be used effectively and cheaply to materialize setter-performig models during training, as shown in Section 4, with notable applications, including re- cent work by DeepSeekAl (2025) and Geiping al. (2025). Additionally, applying them alongside a learning rate de- cay schedule can further enhance performance and improve generalization, all at minimal cost (Section 5).
7. Language Modeling
Although the AlgoPerf benchmark covers diverse work- loads, including a machine translation task with an en- coder-decoder transformer, contemporary language models predominantly adopt decoder-only architectures trained for causal language modeling (next-token prediction). We therefore interested in understanding and benchmarking the potential impact of weight averaging this widely adopted application. A previous work (Sanyal et al., 2023) has high- lighted the benefits of early weight averaging on modern language models, showing how LAWA can produce stronger models early during training, reaching solid downstream tasks performance faster than the bascline trajectory. Unlike show that an Exponential Mov-
Sanyal et al. (2023), ing Average of model parameters can successfully match the performance gains of LAWA when appropriately tuned. We further investigate the role of averaging hyperparam- eters, and comment on the connection between sampling frequency, rolling window length, and discount factor in LAWA and EMA.
8

---

When, Where and Why to Average Weights?
A
EMA
x
Trage uo roadaed
Torzet Tow.
0.1
0,3
0L00
Dot with LAWA
3.8
X
X
10c
X
Validarion Loss 4
3.6
8000
Step to Targee.
3.4
6000
3.2
4000
2000
1000
2000
3000
4000
5000
6000
7000
8000
9000
4
64
128
256
512
1021
Step
Sampling Frequency v
Figure 9: EMA speed-ups when training a 124M transformer on 5B WebEEd tokens. Many configurations outperform the baseline algorithm (AdamW, x 3200 steps to target), but extreme values of /3 or v might hinder performance gains.
Experimental Setting. Since this analysis deviates from AlgoPerf, neither a well-defined target nor strong baseline are a priori available. In an attempt to fairly evaluate WA on this task, we follow an approach similar to the recent nanogpt speed-runs (Jordan et al., 2024), and set the same target validation loss of 3.28 as our validation target score. We employ a similar 124M- -aaaamaetter model, and the same GPT-2 tokenizer; training on 5B tokens from FineWebEdu (Penedo al., 2024), and evaluating on held-out split of 5M tokens. We use AdamW as the baseline optimizer, and perform a grid search over its hyperparameters to first build a strong reference, as described in Appendix D. We then apply LAWA and EMA on top of the fastest-to-target Adam W configuration and compare against by means of step-to-target and final validation loss.
Although limited small-scale experiments, this analysis confirms our earlier findings, and we hope it provides guid- ance when adopting weight averaging techniques in larger scale model training.
8. Limitations
We limit our analysis primarily to NadamW and Adam ( as the baseline optimization algorithms. However, believe that conclusions broadly applicable, and demonstrate how averaging can provide substantial benefits also on top of a higher-order optimizer like Shampoo. Nevertheless, it would be interesting to compare averaging schemes across a broader range of optimization algorithms, and against or on top of schedule-free methods (Defazio et al., 2024b).
Results. Whereas AdamW reaches the target validation loss in about 3200 steps at best, both LAWA and EMA can shorten this time by roughly 30% (Figure 9; Figure 14a; Table 4). The magnitude of this speed-up clearly depen- dends on the predefined target, and although the validation loss of the averaged checkpoints is dominated by AdamW's, the gap between the two shrinks as training proceeds-or -or more precisely, as the learning rate decreases. Indeed, we observe only minimal benefits in terms of final validation perplexity when comparing the final model obtained with AdamW W with the one from LAWA and EMA (Figure 14b; Figure 13), similarly to Section 5. Interestingly, we find that EMA can accurately match LAWA performance when appropriately tuned. Figure displays how EMA hyperpa- rameters affect training speed-ups, showing how the optimal value of 3 varies with the checkpoint sampling frequency (V). In our setting, the choice of v 1, B 0.9 from Sanyal et al. (2023) is indeed suboptimal, and a higher value of 3 is needed when collecting checkpoint at every step. Similar conclusions hold when evaluating weight averaging by means of final generalization performance, as reported in Appendix alongside LAWA results.
9. Conclusions
Through comprehensive empirical investigation, we study the effects of weight averaging on a challenging optimiza- tion benchmark (Dahl et al., 2023), consisting of several diverse Deep Learning architectures and datasets. We eval- uate averaging methods by both the time required to reach reasonable validation target and the best validation per- formance achieved, concluding that averaging offers sub- stantial efficiency gains while achieving marginally better validation scores. Accessing the optimal model during train- ing without prematurely decaying the learning rate remains an open fascinating challenge, and our work presents weight averaging as a step forward, moving models closer to the Pareto optimality of loss versus training time. We hope that this work further encourages the adoption of averaging tech- niques. Although such methods are not directly involved in optimizing updates, their significant improvements may dramatically reduce training time. In the field of efficient training, we believe this a promising study with potential for practical applications (Shen et al., 2024).
9

---

When, Where and Why to Average Weights?
Acknowledgments
learning, and graph networks. arXiv preprint arXiv: 1806.01261,200
The authors would like to acknowledge the Hector Founda- tion for their support and funding, which made this research possible.
Bojar, O., Chatterjee, R., Federmann, C., Graham, Y., Had- dow, B., Huang, S., Huck, M., Koehn, P., Liu, Q. Lo- gacheva, V., Monz, C., Negri, M., Post, M., Rubino, R. Specia, L.,, and Turchi, M. Findings of the 2017 con- ference machine translation (wmt17). In Bojar, O., Buck, C.,, Chatterjee, R., Federmann, C., Graham, Y., Haddow, B., Huck, M., Yepes, A. J., Koehn, P., and Kreutzer, (eds.), Proceedings of the Second Conference on Machine Translation, pp. 169-214, Copenhagen, Den- mark, September 2017. Association for Computational Linguistics. doi: 10. 1L.18666////W1--4111
Impact Statement
Our paper presents a thorough analysis of weight averag- ing for large-scale Machine Learning optimization, we ac- knowledge its potential impact in improving efficiency when training complex Deep Learning systems.
References
Amodei, D., Ananthanarayanan, S., Anubhai, R., Bai, J.,
Cha, J., Chun, S., Lee, K., Cho, H.-C., Park, S., Lee, Y., and Park, S. SWAD: Domain Generalization by Seeking Flat Minima. In Advances in Neural Information Pro- cessing Systems, volume 34, pp. 22405-22418. Curran Associates, Inc., 2021.
Battenberg, E.,, Case, C., Casper, J., Catanzaro, B., Cheng, Q., Chen, G., Chen, J., Chen, J., Chen, Z., Chrzanowski, M., Coates, A., Diamos, G., Ding, K., Du, N., Elsen, E., Engel, J., Fang, W., Fan, L., Fougner, C., Gao, L.. Gong, C., Hanmun, A., Han, T., Johannes, L.,, Jiang, B., Ju, C., Jun, B., LeGresley, P., Lin, L., Liu, J., Liu, Y., Li, W., Li, X., Ma, D., Narang, S., Ng. A., Ozair, S., Peng, Y., Prenger, R., Qian, S., Quan, Z., Raiman, J., Rao, V., Satheesh, S., Seetapun, D., Sengupta, S., Srinet, K., Sriram, A., Tang, H., Tang, L., Wang, C., Wang, J., Wang, K., Wang, Y., Wang, Z., Wang, Z., Wu, S., Wei, I., Xiao, B., Xie, W., Xie, Y., Yogatama, D., Yuan, B., Zhan, J., and Zhu, Z. Deep speech 2 End-to-end speech recognition in english and mandarin. In Balcan, M. and Weinberger, K. Q. (eds.), Proceedings of The 33rd Inter- national Conference on Machine Learning, volume 48 of Proceedings of Machine Learning Research, pp. 173-182, New York, New York, USA, 20-22 Jun 2016. PMLR.
Dahl, G. E., Schneider, F., Nado, Z., Agarwal, N., Sastry, C. S., Hennig, P., Medapati, S., Eschenhagen, R., Kasim- beg, P., Suo, D., Bae, J., Gilmer, J., Peirson, A. L., Khan, B., Anil, R., Rabbat, M., Krishnan, S., Snider, D., Amid, E., Chen, K., Maddison, C. J., Vasudev, R., Badura, M., Garg, A., and Mattson, P. Benchmarking Neural Network Training Algorithms, June 2023, arXiv:2306.071 [cs. stat].
DeepSeekAL. Deepseek-v3: A strong mixture-of-experrs language model. 2025.
Defazio, A., Cutkosky, A., Mehta, H., and Mishchenko, Optimal linear decay learning rate schedules and further refinements, 2024a.
Arpit, D., Wang, H., Zhou, Y., and Xiong, C. Ensemble of averages: Improving model selection and boosting performance in domain generalization, 2022.
Defazio, A., Xingyu, Yang, Mehta, H., Mishchenko, K., Khaled, A., and Cutkosky, A. The Road Less Scheduled, May 2024b. arXiv:2405.1563 [cs, math, stat].
Athiwaratkun, B., Finzi, M., Izmailov, P., and Wilson, A. G. There are many consistent explanations of unlabeled data: Why you should average. In International Conference on Learning Representations, 2018
Deng, J., Dong, W., Socher, R., Li, L.-J,, Li, K., and Fei-Fei, L. Imagenet: A large-scale hierarchical image database. 2009 IEEE conference on computer vision and pattern recognition, 2009.
Bach, F. and Moulines, E. Non-strongly-coones smooth stochastic approximation with convergence rate O(1/n). In Advances in Neural Information Processing Systems, volume 26. Curran Associates, Inc., 2013.
Dosovitskiy, A., Beyer, L., Kolesnikov, A., Weissenborn, D., Zhai, X., Unterthiner, T., Dehghani, M., Minderer, M., Heigold, G., Gelly, S., et al. An image is worth 16x 16 words: Transformers for image recognition at seale. arXiv preprint arXiv:2010. 1929, 2020.
Battaglia, P. W., Hamrick, B., Bapst, V., Sanchez- Gonzalez, A., Zambaldi, V., Malinowski, M., Tacchetti, A., Raposo, D., Santoro, A., Faulkner, R., Gulcehre, C., Song, F., Ballard, A., Gilmer, J., Dahl, G., Vaswani, A., Allen, K., Nash, C., Langston, V., Dyer, C., Heess, N., Wierstra, D., Kohli, P., Botvinick, M., Vinyals, O., Li,
Garrigos, G. and Gower, R. M. Handbook of conver- gence theorems for (stochastic) gradient methods. arXiv preprint arXiv:2301.1 1235, 2023.
Y., and Pascanu, R. Relational inductive biases, deep
10

---

When, Where and Why to Average Weights?
Geiping, J., McLeish, S., Jain, N., Kirchenbauer, J., Singh, S., Bartoldson, B. R., Kailkhura, B., Bhatele, A., and Goldstein, T. Scaling up test-time compute with latent reasoning: A recurrent depth approach, 2025.
International Conference on Artificial Intelligence and Statistics, pp. 1347-1355. PMLR, March 2018. ISSN: 2640-3498.
Li, S., Liu, 7., Tian, J., Wang. G., Wang, Z., Jin, W., Wu, D., Tan, C., Lin, T., Liu, Y., Sun, B., and Li, S. Z. Switch EMA: A Free Lunch for Better Flatness and Sharpness, February 2024. XXiv:2402.09240 [es].
Gulati, A., Qin, J., Chiu, C.-C., Parmar, N., Zhang, Y., Yu, J., Han, W., Wang, S., Zhang, Z., Wu, Y.. and Pang, R. Conformer: annvolution-uugmente transformer for speech recognition, 2020.
Li, T. Huang, Z., Tao, Q. Wu, Y., and Huang, X. Trainable Weight Averaging: Efficient Training by Optimizing His- torical Solutions. In Proceedings of the 4th International Conference on Learning Representations, 2023.
Gupta, V., Serrano, S. A., and DeCoste, D. Stochastic weight averaging parallel: Large-batch training that generalizes well. In International Conference on Learn- ing Representations, 2020.
Loshchilov, I. and Hutter, F. Decoupled weight decay reg- ularization. In International Conference on Learning Representations, 2017a.
Hägele, A., Bakouch, E.,, Kosson, A., Allal, L. B., Werra, L. V., and Jaggi, M. Scaling Laws and Compute-On Optimall Training Beyond Fixed Training Durations. Advances in Neural Information Processing Systems, 2024.
Loshchilov, and Hutter, F. Sgdr: Stochastic gradient descent with warm restarts, 2017b.
Hu, W. Fey, M., Zitnik, M., Dong, Y., Ren, H., Liu, B., Catasta, M., and Leskovec, J. Open graph benchmark: Datasets for machine learning on graphs. Advances in neural information processing systems, 2020.
Melis, G. Two-Tailed Averaging: Anytime, Adaptive, Once- in-a- While Optimal Weight Averaging for Better Gener- alization, April 2023. arXiv:2209.1255 [cs, stat].
Merity, S., Keskar, N. S., and Socher, R. Regularizing and optimizing Istm language models, 2017.
Izmailov, P., Podoprikhin, D., Garipov, T., Vetrov, D., and Wilson, A. G. Averaging weights leads to wider optima and better generalization, 2019.
Morales-Brotons D., Vogels, T., and Hendrikx, H. Ex- ponential moving average of weights in deep learning: Dynamics and benefits, 2024.
Jordan, K., Bernstein, J., Rappazzo, B., @fern- ssu.ca Vlado, B., Jiacheng, Y., Cesista, F., Koszarsky, B., and @ Grad62304977. nadoett-anogpt: Speedrunning the nanogpt baseline, 2024.
Naumov, M., Mudigere, D., Shi, H.-J. M., Huang. J., Sun- daraman, N., Park, J., Wang, X., Gupta, U., Wu, C.-J., Azzolini, A. G., al. Deep learning recommendation model for personalization and recommendation systems. arXiv preprint 19Xivv1906.000
Kaddour, Stop Wasting My Time! Saving Days of Ima- geNet and BERT Training with Latest Weight Averaging, October 2022. arXiv:2209.1499 [cs, stat].
2019.
Neu, G. and Rosasco, L. Iterate Averaging as Regularization for Stochastic Gradient Descent. In Proceedings of the 31 Conference On Learning Theory, pp. 3222-3242. PMLR, July 2018. ISSN: 2640-3498.
Karpathy, A. Nanogpt. https: http:/ggtttthho.ooo karpathy Bachhyy/aaaa /nanoGPT, 2022.
Karras, T., Aittala, M., Lehtinen, J., Hellsten, J., Aila, T., and Laine, S. Analyzing and improving the training dynamics of diffusion models, 2024.
Panayotov, V., Chen, G., Povey, D., and Khudanpur, S. Librispeech: An ast corpus based public domain au- dio books. In 2015 IEEE International Conference on Acoustics, Speech and Signal Processing (ICASSP), pp. 2015. doi: 10. 10.11099CASSPPP205..777
Kasimbeg, P., Schneider, F., Eschenhagen, R., Bae, J., Sas- try, C. S., Saroufim, M., FENG, B., Wright, L., Yang, E. Z., Nado, Z., Medapati, S., Hennig, P., Rabbat, M., and Dahl, G. E. Accelerating neural network training: An analysis of the algoperf competition. In The Thirteenth International Conference on Learning Representations, 2025.
Paszke, A., Gross, S., Chintala, S., Chanan, G., Yang, E., DeVito, Z., Lin, Z., Desmaison, A., Antiga, L.,, and Lerer, A. Automatic differentiation in pytorch. NIPS 2017 Workshop Autodiff, 2017.
Penedo, G., Kydlfček, H.,, allal, L. B., Lozhkov, A., Mitchell, M., Raffel, C., Werra, L. V., and Wolf, T. The fineweb datasets: Decanting the web for the finest text data at scale. In The Thirty-eight Conference on Neural Informa- tion Processing Systems Datasets and Benchmarks Track, 2024.
Lab, C. I. Criteo Itb click logs dataset. 2014.
Lakshminarayanan, C. and Szepesvari, C. Linear Stochastic Approximation: How Far Does Constant Step-Size and Iterate Averaging Go? In Proceedings of the Twenty-Firs

---

When, Where and Why to Average Weights?
Polyak, B. T. New stochastic approximation type procedures. Avtomatica l'elemekhanika, 7:98-107, 1990.
Su, J., Lu, Y., Pan, S., Murtadha, A., Wen, B., and Liu, Y. Roformer: Enhanced transformer with rotary position embedding, 2023.
Polyak, B. T. and Juditsky, A. B. Acceleration of Stochastic Approximation by Averaging. SIAM Journal on Control and Optimization, 30(4):838-855, July 1992. ISSN 0363- 0129, 1095-7138, doi: 10.1137/033004..
Szegedy, C., Vanhoucke, V., loffe, S., Shlens, J., and Wojna, Z. Rethinking the inception architecture for computer vision. In 2016 IEEE Conference Computer Vision and Pattern Recognition (CVPR). pp. 2818-2826, 2016. doi:
Portes, J., Blalock, D., Stephenson, C., and Frankle, Fast benchmarking of aceuracy vs. training time with eyclic learning rates, 2022.
Tarvainen, A. and Valpola, H. Mean teachers are better role models: Weight-averaged consistency targets improve semi-supervised deep learning results, 2018.
Radford, A., Wu, J.. Child, R., Luan, D., Amodei, D., and Sutskever, I. Language models are unsupervised multitask learners. OpenAl, 2019. Accessed: 2024-11-15.
Timothy, Incorporating Nesterov Momentum into Adam. In Proceedings of the 4th International Conference on Learning Representations, pp. 1-4, 2016.
Ronneberger, O., Fischer, P., and Brox, T. U-Net: Convo- lutional Networks for Biomedical Image Segmentation. In Medical Image Computing and Computer-Assissed Intervention MICCAI 2015, Cham, 2015. Springer In- ternational Publishing.
Vaswani, A., Shazeer, N., Parmar, N., Uszkoreit, J., Jones, L., Gomez, A. N., Kaiser, E., and Polosukhin, At- tention is all you need. Advances in neural information processing systems, 2017.
Ruppert, D. Efficient Estimations from Slowly Convergent Monno Process, 1988.
Wortsman, M., Ilharco, G., Gadre, S. Y., Roelofs, R., Gontijo-Lopes, R., Morcos, A. S., Namkoong, H., Farhadi, A., Carmon, Y., Kornblith, S., and Schmidt, Model soups: averaging weights of multiple fine-tuned models improves accuracy without increasing inference time. In Proceedings of the 39th International Confer- ence on Machine Learming, pp. 23965-23999 PMLR, June 2022. ISSN: 2640-3498.
Sandler, M., Zhmoginov, A., Vladymyrov, M., and Miller, N. Training trajectories, mini-batch losses and the curious role of the learning rate, February 2023. arXiv:2301.02312 [es].
Sanyal, S., Neerkaje, A., Kaddour, J., Kumar, A., and Sanghavi, S. Early Weight Averaging meets High Learning Rates for LLM Pre-training, December 2023. arXiv:2306.03241 [cs].
Yang, G.,, Zhang, T., Kirichenko, P., Bai, J., Wilson, A. G., and Sa, C. D. SWALP Stochastic Weight Averaging in Low Precision Training. In Proceedings of the 36th International Conference on Machine Learning, pp. 7015- 7024. PMLR, May 2019. ISSN: 2640-3498.
Shazeer, N. Glu variants improve transformer. arXiv preprint arXiv: 2002.05202, 2020.
YaZIcı, Y., Foo, C.-S., Winkler, S., Yap, K.-H,, Piliouras, G., and Chandrasekhar, V. The unusual effectiveness of averaging in gan training, 2019.
Shen, I.., Sun, Y., Yu, Z., Ding, I.,, Tian, X., and Tao, D. On efficient training of large-scale deep learning models. ACM Comput. Surv. 57(3). November 2024. ISSN 0360- 0300. doi: 10. 10.111 145/3700449
Zbontar, J., Knoll, F,, Sriram, A., Muckley, M. J., Bruno, M., Defazio, A., Parente, M., Geras, K. J., Katsnelson, J., Chandarana, H., Zhang, Z., Drozdzal, M., Romero, A., Rabbat, M. G., Vincent, P., Pinkerton, J., Wang, D., Yakubova, N., Owens, F.,, Zitnick, C. I.., Recht, M. P., Sodickson, D. K., and Lui, Y. W. fastmri: An open dataset and benchmarks for accelerated MRI. CORR, abs/181 .08839, 2018.
Shi, H.-J M., Lee, T.-H., Iwasaki, S., Gallego-Posade J., Li, Z., Rangadurai, K., Mudigere, D., and Rabbat, M. A distributed data-paralle, pytorch implementation of the distributed shampoo optimizer for training neural networks at-scale, 2023.
Singh, V., Janson, P., Mehrbod, Ibrahim, A., Rish, I., Belilovsky, E., and Thérien, B. Beyond cosine decay: On the effectiveness of infinite learning rate schedule for continual pre-training, 2025.
Zhang, B. and Sennrich, R. Root mean square layer normal- ization, 2019.
Song, Y., Sohl-Dickstein, J., Kingma, D. P., Kumar, A., Er- mon, S., and Poole, B. Score- based generative modeling through stochastic differential equations, 2021.
12

---

When, Where and Why to Average Weights?
A. Experimental Details
the Distributed Shampoo implementation, we resort the original submission the AlgoPerf competition'. We select the optimal hyperparameters for NadamW and Shampoo on each workload as the ones providing faster convergence to the benchmark target in the first iteration of the AlgoPerf competition'.
A.1. Datasets and Architectures
We report further details about the workloads considered for the analysis and provide the correspondence references. We consider the following architectures and tasks:
DIRMsmall model on the Criteo ITTB dataset (Nau- mov et al., 2019; Lab, 2014); a Deep Learning recom- mendation model optimized for large-scale industrial recommendation tasks.
B. LAWA and EMA Hyperparameters
Having demonstrated the benefits of averaging weights along the training trajectory, we now investigate the impact of the hyperparameters of the averaging scheme. Specif- ically, we explore the interaction between the update fre- quency and the averaging window L for LAWA (Fig- ure 10), as well as the relationship between v and / for EMA (Figure 11), and investigate how different choices of these hyperparameters influence the effectiveness of av- eraging in reducing the number of steps required to reach the benchmark validation targets. For this analysis we use Nadam W as the underlying optimization algorithm, and the best performing hyperparameter from Dahl et al. (2023) for each workload.
U-Net (Ronneberger et al., 2015) on the FastMRI dataset (Ronneberger et al., 2015; Zbontar et al., 2018): a popular architecture for medical image segmentation.
VIT (Vision Transformer) on ImageNet- lkk (Dosovit- skiy et al., 2020; Deng et al., 2009): transformer- based model for image classification.
GNN model on OGBG dataset (Battaglia et al., 2018; Hu et al., 2020): graph neural network for learning over graph-structured data, used to test performance on graph-based machine learning tasks.
In both cases, we observe a consistent trend: more frequent updates (lower v) benefit from longer memory (higher L or f),, and vice versa: sampling checkpoints further apart (higher 2/) works better with shorter memory buffers (lower I. or (8). Several combinations of v and T. or ,3 yield strong performance, suggesting that an optimal averaging horizon might exist. We also note that slower averages (high L or 3) are more sensitive to the choice of v, requiring more careful tuning, and that EMA is usually more sensitive its hyperparameters. In practice, we find that values of in the range 64 - 128 combined with L = 20 provide good perfor- mance across workloads for LAWA. For EMA, our analysis suggests that a value of / between 64 and 128 with 3 = 0.9 is effective. We acknowledge that better configurations may exist, and that the optimal values may vary depending on the workload and other hyperparameters, particularly on the LR schedule.
Transformer-big WMT dataset (Vaswani et al., 2017; Bojar et al., 2017): large-scale transformer model applied to machine translation.
Conformer on LibriSpeech dataset (Gulati al., 2020; Panayotov al., 2015): hybrid architecture combin- ing CNNs and transformers, used for speech recogni- tion tasks.
DeepSpeech (Amodei et al., 2016) on LibriSpeech dataset (Panayotov et al., 2015): recurrent model for speech-to-text, tested on popular speech recognition dataset to assess its performance in transcribing audio.
These workloads span various domains, including recom- mendation systems, image and speech processing, machine translation, and graph-based learning. Each tests different model architectures under real- real-world data conditions.
We further observe that when the baseline algorithm (NadamW) fails to reach the validation target-possiie due to poor initialization or an unfavorable data stream applying LAWA or EMA typically does not re- cover performance and still fails to attain the target. This is evident, for instance. in the Librispeech Conformer and Lib- rispeech Deepspeech workloads, where one of the random seeds leads the baseline to consistently fail, and both LAWA
A.2. Computational Resources
We conducted experiments on 1111000XXM440GBBB and H4H100-HBM3-8000 GPUs, occasionally resorting to 8XV100-32GB. We use Distributed Data Parallel to paral- lelize training across devices. Training is performed in full precision, enabling TF32-matmul for faster computations.
Pntt //gf //ggggttt wwt.hhooo//ooo..teooooooo heae///oommnn commona/ algort thrs_reau to, 5/ ree /ma //A aogoorr__eeas Seeaaen/eeeeeeeeeeeeeeeeeee /extternall Luuntng/ /tu//ng/shampoo. shampon aubmt as ton
We make use of the original AlgoPerf repository', with minimal modifications, using the PyTorch (Paszke al., 2017) implementation of the corresponding algorithms. For
httpps //gl thub -oom com/m/l commona/ a Aigrrthhaee hhrssreeu te s/tree /ma n/ logs/ goport sooor_ssooonng wacorrngg_s////// vt5/exter/daaaeeeee tuniin// AlgoPerf wtwps//////eeeeeeeeee qualificed Ficatiion
sttpp / thu.. com/ I commons fa al glggr c-of of of ency
13

---

When, Where and Why to Average Weights?
Averaging Window Length L
Target not reached
10
20
50
Step Budget
50000
800000
X0XX X0XX
47500
Step to Target 4
75000
Step to Target
45000
42500
70000
40000
68
65000
37500
60000
1.0
64.0
128.0
256.01
512.0
10210
1.0
61.0
128.0
256.0
512.0
Sampling Prequency V
Sampling Frequency y
(a) Librispeech Conformer
(b) Librispeech DeepSpeech
11000
80000
Step to Target 1
Step to Target |
10000
0000
40000
20000
7000
1.0
32.0
64.0
126.0
256.0
1.0
16.0 Sampling Frequeney y
64.0
256.0
Sompling Frequency
v
(c) Critcol
(d) OGBG
Figure 10: Performance of Nadam +LAWA across different combinations sampling frequency v and window length I. with each point representing a different random seed (3 seeds per configuration). An "x" mark indicates the target was not reached for that specific seed within the training horizon (red dotted line). We notice that small values of require large value of L, and viceversa. On Criteo ITB, we exclude = 256 and L = 50, as the first average would occur at step 256 50- = 12800 with this conifiguration, exceeding the training horizon (10666 steps).
and EMA exhibit the same behavior (indicated by the sin- gle persistent" "x" mark across hyperparameter settings in Figure 10 and Figure 11). These observations reinforce the findings in Section 6: averaging can accelerate training by acting proxy for shorter learning rate schedule, but it offers limited benefit when the baseline optimizer produces poor checkpoints and follows a suboptimal trajectory.
C. Additional AlgoPerf Experiments
We report additional results for the analysis in Section 6, comparing averaging of a long annealed Nadam trajec- tory with shorter LR schedules across additional workloads. Figure displays the decay schedule and validation perfor- mance during training for FastMRI, OGBG, and WMT. On FastMRI and OGBG, both LAWA and EMA closely match the performance of runs with smaller budget. On OGBG, averaging generally provide better MAP scores throughout training but eventually regress to the overfitting baseline algorithm. On WMT, we observe that a short 25% schedule outperforms the averaging schemes at that point. Despite consistently approaching Pareto optimality of loss versus training time, LAWA and EMA cannot always achieve it. We argue that more sophisticated annealing strategies, pos- sibly similar to Defazio et al. (2024b), may be required to access the optimal model at any time during training.
Finally, unlike Arpit et al. (2022), we observe that EMA with performs poorly for 3 < 0.99. We find this behavior peculiar, and hypothetize that the worst performance of EMA compared to LAWA observed by Sanyal et al. (2023) might be due to updating the EMA buffer at every step, whereas for LAWA they use larger values of v, and further explore this in Section 7 and Appendix D. A similar result holds for LAWA, with the only exception OGBG, where we however acknowledge very long training horizon and very slow baseline, as noted in Kasimbeg et al. (2025).
14

---

When, Where and Why to Average Weights?
EMA B
Target not reached
0.3
0.6
0.9
0.99
Step Budget
X00x
DeepSpeech EMA performance 8, seeds)
50000
Step to Target 4
Step to Target
45000
8
40000
35000
1.0
64.0
128.0
256.0
512.0
10210
1.0
61.0
128.0 Sampling Frequeney y
256.0
512.0
Sampling Prequency V
(a) Librispeech Conformer
(b) Librispeech DeepSpeech
11000
Step to Target 1
10000
Step to Target
0000
10000
7000
g
20000
6
1.0
32.0
64.0 Sompling Prequency
128.0
258.0
1.0
16.0 Sampling Frequeney
61.0
256.0
2048.0
v
y
(c) Critcol TBB
(d) OGBG
Figure 11: Performance of NadamW +EMA across different combinations of and B, each point represents different random seed (3 seeds per configuration). Small values of generally work better with larger value of (3, and vice versa.
D. Additional Details on Language Modeling
and EMA configurations. We then compare these config- urations by means of (A) cnne--to-vaiaation-ttee target and (B) minimal validation loss achieved throughout training, akin to Section and Section respectively.
We provide additional details and results for the Language Modeling analysis in Section 7.
We perform all experiments in PyTorch using IIainIM pretraining codebase and model implementations. The em- ployed model is a decoder-only 124M transformer archi- tecture, similar GPT-2 (Radford et al., 2019; Karpathy, 2022), enhanced with Rotational Positional Embedding (Su al., 2023), RMSNorm (Zhang & Sennrich, 2019), and SwiGLU (Shazeer, 2020). Models are trained for next-token prediction on 5B tokens from FineWebEdu (Penedo et al., 2024), using a batch size of 512, context length of 1024 tokens, and GPT-2 tokenizer. To derive strong baseline algorithm, we perform grid search over weight decay, gra- dient clipping and learning rate schedule hyperparameters, exploring both decay-to-zero and 10% decay, reported in Table 3. Having identified optimal baseline across these values, average top of it using different LAWA
Table 3: Adam W Hyperparameter Search Space. Final LR is reported fraction of top LR and warmup as a fraction of the step budget. We highlight in blue the fastest-to-target configuration after the grid search.
Hyperparameter
Values
Top LR
0.0003, 0.001, 0.003, 0.01, 0.03 0, 0.1
Final LR
Warmup
0.01,
Weight decay
0.01, 0.1
Gradient norm clipping
0.01, 0.1
Results for optimal WA configurations are compared to the
attpa: //gstttttttt com/ sco woo.oooooooooooooooo Agjo mo apna
AdamW baseline Table 4, and the interplay between v, L, and is explored in Figure 9, Figure 13, and Figure 14.
15

---

256 When, Where and Why to Average Weights?
WMT Shurt
040733
6.812
Seldd
vEMA
la So Suncki 199A
40077
10100
20000
1.000
1308
Skyp
2300
3000
1010
N118
3116
2118
V018
W110
70130
4013
2200
2100
0.001
3300
22000
0.736
1:30
67744
2.2
10112
Solare Sant Sto0k3
Neliee Sovr Statte
0.18
Noture Leng Schider
810
Skz
Stop
(a) FastMRI
(b) OGBG
(c) WMT
Figure 12: Top: Learning rate schedules for different workloads. Bottom: Corresponding performance during training. Weight averaging of a long annealed training trajectory approaches the performance of runs with shorter decay schedules.
EMA 8
Bost with Adun Rest anth LAWA
Table 4: Language Modeling speed-ups and validation per- formance. WA provides substantial speed-ups. but minimal generalization gains over the last checkpoint. Results report median and standard deviation over 3 seeds. The missing standard deviation for the step-to-target is an artifact of the evaluation frequency, which we set to 40 steps, limiting the resolution of the correspondent metric, and resulting in standard deviation of zero.
0.6
0P
3.10
Valitatoon Los
3.08
3.06
Step-to-target Best Validation Loss
16
128 Sampling Frequeacy
64
512
1024
AdamW
3200
3.0425+ 003630
Figure 13: Language Modeling final validation loss of AdamW +EMA across different combinations of and B.
+LAWA
2240
3.0421 40.000002
+EMA
2200
3.0413++0.0000
We implement the online schemes naively in PyTorch, defaulting to blocking communication, but note that it is possible to asynchronously offload the model weights to the CPU, using non- non-blocking communication to update an EMA or a moving average of model weights. We acknowledge a first adoption of this approach in DeepSeckAl (2025), and hope that this work encourages wider usage of it.
E. Averaging Algorithms
We perform most of the experiments using offline versions of LAWA and EMA. This involves training with a baseline algorithm (NadamW, AdamW Distributed Shampoo), frequently saving checkpoints to disk. This allows to then average and evaluate the model weights, exploring different averaging windows and hyperparameters at minimal cost, without having to retrain the model from scratch.
Finally, we note that the poor performance of LAWA and EMA in the inaugural AlgoPerf competition (Kasimbeg et al., 2025) stemmed from the API's lack of support for switching between model parameters and averaging buffers at evaluation time, resulting in inefficient CPU-GPU trans- fers on bandwidth- does not support asynchronous transfers, further penalizing such submissions. However, as noted earlier, more efficient implementations are possible in real-world applications, where weight averaging can be used to seamlessly material- ize better performing models. We report PyTorch-sty implementation of LAWA and EMA in Algorithms and 2.
In practice, however, might be beneficial to use online averaging schemes. As demonstrated throughout this work, averaging can operate as a proxy for a shorter learning rate schedule, providing significantly better performance than the currently available model. Since this aspect has interesting implications for large-scale efficient training, it is worth discussing the possible computational overhead introduced by averaging schemes. For the online versions of EMA or LAWA, we offload the averaging buffer to CPU.
hardware. Additionally, AlgoPerf
16

---

When, Where and Why to Average Weights?
Averaging Window Length
Target reached
Adam baseline
Step budget
3.10
X000
Step Target
Validation Toes
0000
3.06
1000
000
p
3.04
16
64
256
512
1024
16
64
128
256
512
1024
Sampling Frequency
Sampling Froquency
(a) Training speed-up. Training runs that do not reach the target validation loss by the end of trainng marked with "x".
(b) Final validation loss. Training runs that achieve minimum validation loss 2 3. 10, omitted.
Figure 14: Language Modeling performance of AdamW +LAWA across different combinations of v and L. Small values of pair best with larger L, and vice versa. Beyond a critical sampling frequency v, weight averaging gains drop drastically, both training acceleration and final validation performance. We observe similar performance between LAWA and EMA (Figure 13, Figure 9), with both methods achieving comparable speed-ups and loss trajectories across most v values; minor differences emerge only at extreme hyperparameter settings. These results suggest that, provided and L (or () correctly initialized, either LAWA EMA can be chosen without significant loss of efficacy.
Algorithm 1 Latest Weight Averaging (LAWA)
Initialize LAWA hyperparameters v, Initialize circular queue Q of length L for batch in enumerate(rrainlodde)) do Forward, backward and optimization step
Algorithm Exponential Moving Averaging EMA
Initialize EMA hyperparameters v,
if t mod V = 0) then
Initialize EMA parameters OEMA = 0g for batch in enumerate(trainloade]] do
if [Q] < L then
Add current parameters 0g to Q
Forward, backward and optimization step
else
if mod 0 then
Remove oldest element from Q, add current pa- rameters O
GEMA + 90 + =3),,
end if
end if
if is eval time then Eval θEMA on held-out data end if
end if
if is eval time then
lo + Eeeo
end for
Eval 0Q on held-out data
end if
end for
17


Reviews:
[{"review_id": "fRIp6x7ulL", "reviewer": "Reviewer_5hDM", "summary": "This paper systematically explores two existing weight averaging techniques, EMA and LAWA, across the seven diverse tasks in the 2024 Algoperf competition. The authors find that weight averaging can either increase the final performance for a given compute budget or speed up training to specific performance threshold (by roughly ~20% on top of a well tuned baseline). They also explore whether weight averaging can replace learning rate decay but find that it can not. Rather it is always better to combine weight averaging with standard learning rate schedules (cosine in this paper) than using either method alone.\n\n**Post-rebuttal update:** The authors have agreed to rephrasing edits which effectively address my actionable concerns. My conclusion remains that there is little to criticize about the paper aside the limited novelty of the methods and findings. I still think it is a nice reference and validation of the methods which is valuable. Overall these two factors roughly balance out to make this a borderline paper in my opinion.", "questions": "", "limitations": "", "rating": 3, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your comments! We respond to specific comments below, but please let us know if you have any additional questions.\n\n> I do not agree that the paper shows how to optimally combine the two. What the paper does show is that combing standard weight averaging with standard learning rate schedules [...] is better than just using one of them.\n\nYou are right—we do not explore *optimal* combinations of the two, but rather demonstrate that **combining** weight averaging with learning rate schedules outperforms using either alone. The broad scope of the benchmark makes a full optimality study expensive, but we believe this is still a useful result for practitioners. We will revise the manuscript to clarify this distinction. Thank you for pointing this out!\n\n> I somewhat disagree that this is a new finding. Practical works that use e.g. EMA on top of a learning rate schedule do so exactly because it works better than not using EMA.\n\nYou are right, this is not a completely novel claim. Previous works indeed show that combining WA with learning rate schedules is beneficial. We will adjust the revised paper and clarify this point. We believe the impact of our work comes from demonstrating this at scale, unifying insights from previous work in a single competitive benchmark, and showing that WA alone does not replace learning rate decay across a wide range of settings. \n\n> Only relative savings matter, 100s of GPU hours can range from very significant to very insignificant. \n\nWe agree that relative savings are important, and we will rephrase this for clarity. We did report the **15% saving** in GPU hours, which we believe to be significant for this benchmark, and will ensure this is emphasized more clearly in the revision.\n\n> I think calling Shampoo a second order optimizer is a bit of a stretch.\n\nWe will clarify that while Shampoo is inspired by second-order principles, it may not strictly be a second-order optimizer.\n\n> \"coupling WA with learning rate annealing\": I think \"combining\" would be clearer than \"coupling\" here.\n\nThank you for the useful suggestion. We agree that \"combining\" would be clearer than \"coupling\" and will make that adjustment in the revision.\n\n> On related work: I think some types of weight averaging were used earlier than you mention in deep learning.\n\nThank you for the additional references. We are aware of this and will update the manuscript accordingly. We will add a paragraph in the Related Work section, citing earlier studies that predate SWA. Apart from the Transformer paper (Vaswani et al., 2017), we found that **other influential studies** such as Szegedy et al. (2016) and Merity et al. (2017) have also employed averaging schemes to improve model performance and reduce overfitting.", "claims_and_evidence": "Most of the claims are fine, but I have minor issues with some of them.\n\n**We explore the relationship between averaging and learning rate annealing and show how to optimally combine the two to achieve better performance:** I do not agree that the paper shows how to **optimally** combine the two. What the paper does show is that combing standard weight averaging with standard learning rate schedules (as opposed to constant learning rates) is better than just using one of them. The paper does not show e.g. some new learning rate schedule that works better with EMA than standard techniques.\n\n**(We) show that coupling WA with learning rate annealing yields optimal results— a novel finding to the best of our knowledge:** I somewhat disagree that this is a new finding. Practical works that use e.g. EMA on top of a learning rate schedule do so exactly because it works better than not using EMA (Karras 2312.02696 could be an example). Other works have shown that WA alone can not replace learning rate decay (e.g. Haegele 2405.18392). I think that the fact that the combination yields better results than one alone follows directly from these, although I am not sure if a single work has shown this.", "methods_and_evaluation_criteria": "Yes, evaluating weight averaging techniques on a competitive benchmark on top of very well tuned baselines is good. The scale and diversity of the benchmark is also convincing.", "theoretical_claims": "The paper doesn't make any theoretical claims.", "experimental_designs_or_analyses": "Yes, I looked through the experimental setup and found no significant issues with it. Overall the experiments are convincing and well executed. Error bars are given for some results, further strengthening the results.", "relation_to_broader_scientific_literature": "The paper does a good job of discussing related literature. I feel this paper primarily serves as a (decently diverse and large scale) validation of existing methods (standard weight averaging approaches) and observations about them (they improve performance, simple weight averaging can not replace lr schedules). I did not find the questions asked or the findings very novel or surprising. The empirical validation is still well done and can be of value to the community.", "essential_references_not_discussed": "None that I am aware of.", "other_strengths_and_weaknesses": "The paper is well written and the figures are clear. The experiments are well done as previously mentioned.\n\nThere is little to criticize about the paper aside the limited novelty of the methods and findings. I still think it is a nice reference and validation of the methods which is valuable. Overall these two factors roughly balance out to make this a borderline paper in my opinion.", "other_comments_or_suggestions": "\"Even for industrial-scale tasks like benchmark workloads, weight averaging reliably reduces compute costs by hundreds of GPU hours.\" - I think this sentence is a bit awkward. Only relative savings matter, 100s of GPU hours can range from very significant to very insignificant. The table below shows savings of roughly 100 hours, not 100s.\n\n\"even on second-order optimizers\": There might be differing opinions on this but I think calling Shampoo a second order optimizer is a bit of a stretch (even if it is motivated by some very rough approximation of second order phenomena).\n\n\"coupling WA with learning rate annealing\": I think \"combining\" would be clearer than \"coupling\" here. To me coupling signifies a tighter relationship than just using both methods. E.g. if you varied an EMA time horizon dynamically with the learning rate schedule or lowered the learning rate based on some signal from the EMA performance, that would better qualify as coupling.\n\nOn related work: I think some types of weight averaging were used earlier than you mention in deep learning. Maybe look into deep RL, the original transformer paper (\"for the big models, we averaged the last 20 checkpoints\"), and I think TensorFlow had support for some type of weight averaging very early on.", "questions_for_authors": "No specific questions.", "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": "syQtxNWyJU", "reviewer": "Reviewer_A5Xg", "summary": "The authors benchmark weight averaging techniques using AlgoPerf, and find that without learning rate annealing it strongly accelerates training and generalization. It composes positively with learning rate annealing with small gains but cannot fully  replace annealing.", "questions": "", "limitations": "", "rating": 4, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your review and for your insightful comments. We appreciate your recognition of the value of our study! We respond below to specific comments.\n\n> It would be useful to discuss existing application of weight averaging a little bit, especially in relation to robust finetuning [1] and OOD generalization [2].\n\nThank you for the additional references, we will include them in the revised manuscript.\n\n> It would be nice to clarify which dataset the figures are associated with in all cases. Eg. Figure 7.\n\nThank you for spotting this! We will correct the labeling. Figure 7 is obtained by training a **Conformer** on Librispeech. We note that a similar effect can be observed on **other workloads** as well, and provide **additional results** [at this URL](https://anonymous.4open.science/r/ICMLE1C2/Short_LR.pdf) for comparing Weight Averaging against shorter LR schedules, [and at this URL](https://anonymous.4open.science/r/ICMLE1C2/Annealing_EMA.png) for testing the effect of EMA and LAWA when the learning rate is not annealed to zero. As suggested by Reviewers #egdu and #fFcZ, we will incorporate these additional experiments and considerations in the revised article to provide a broader view and strengthen our findings.\n\n> I like the idea that learning rate schedules impede continual learning, since you decay to 0 at some point and restarting could be problematic from an optimization standpoint. Do the authors have any evidence that WA can be used for improved continual or transfer learning?\n\nThis is a very interesting question! Resuming training from a checkpoint obtained with standard cosine annealing can indeed present challenges, such as **potential forgetting** during the **re-warming phase** (Singh, V. et al., 2025). Therefore, it looks appealing to avoid annealing, and instead resume from the average model. We have not experimented with this point yet, but we think it would be an interesting direction for future work. An interesting question is how to optimally **resume training from the averaged model**, especially how to adapt the learning rate value in this scenario. Since averaging acts as a proxy of implicit learning rate decay, we hypothesize that using a smaller learning rate could be more beneficial for continued training from the average rather than restarting from the previous top learning rate. We believe this fits into the broader question of _how to better utilize the average for training_ (Kaddour et al., 2022; Defazio et al., 2024), a topic that remains underexplored but could offer significant benefits to the community.\n\n> Why, in Figure 8, does annealing + weight averaging help whereas in figure 7b LAWA adds no benfit relative to the annealed model? Is it a different dataset, or difference in some other setting?\n\nFigure 7b shows results on Librispeech Conformer, which, as reported in Figure 8, benefits only minimally when combined with LAWA. This small gains are not visible at the scale of the plot in Figure 7b. \n\n### References\n\nSingh, V., et al. (2025). Beyond Cosine Decay: On the effectiveness of Infinite Learning Rate Schedule for Continual Pre-training. arXiv preprint arXiv:2503.02844.\n\nKaddour, J. (2022). Stop Wasting My Time! Saving Days of ImageNet and BERT Training with Latest Weight Averaging. arXiv preprint arXiv:2209.14981.\n\nDefazio, A., et al. (2024). The Road Less Scheduled. arXiv preprint arXiv:2405.15682.", "claims_and_evidence": "* Weight Averaging is generally helpful, especially in the absence learning rate annealing. Experimental evidence supports this. \n\n* Weight Averaging Speeds Up Training. Experimental evidence supports this. \n\n* Weight Averaaging + Learning Rate Annealing works best.  Experimental evidence supports this.\n\n* Weight Avearaging cannot replace learning rate annealing.", "methods_and_evaluation_criteria": "The evaluation framework allows for extensive and diverse experiments, which is important for answering the research questions being asked in the paper.", "theoretical_claims": "There are no Theoretical Claims.", "experimental_designs_or_analyses": "Experiments using AlgoPerfm are well designed.", "relation_to_broader_scientific_literature": "The scientific iteratture implictilty takes advantage of the findings of this paper extensively (wherever weight averaging is used). Formalizing them in a large scale study, however, presents the benfits clearly and is useful for distilling the implicit knowledge in the paper.", "essential_references_not_discussed": "It would be useful to discuss existing application of weight averaging a little bit, espeically in relation to robust finetuning [1] and OOD generalization [2] \n\n[1]Wortsman, Mitchell, et al. \"Robust fine-tuning of zero-shot models.\" Proceedings of the IEEE/CVF conference on computer vision and pattern recognition. 2022.\nAPA\t\n\n[2]Arpit, Devansh, et al. \"Ensemble of averages: Improving model selection and boosting performance in domain generalization.\" Advances in Neural Information Processing Systems 35 (2022): 8265-8277.", "other_strengths_and_weaknesses": "Strengths: The paper is crystal clear in its findings, and this clarity is the basis for my recommendation of acceptance. \n\nI like the finding that this weight averaging allows for access to a better model during training, useful for long duration training runs. \n\nWeaknesses: The findings are not very surprising to the many using weight averaging. However, this is not a reason to reject IMO. We are lacking in clear presentation of simple ideas in this field.", "other_comments_or_suggestions": "It would be nice to clarify which dataset the figures are associated with in all cases. Eg. Figure 7. \n\nI like the idea that learning rate schedules impede continual learning, since you decay to 0 at some point and restarting could be problematic from an optimization standpoint. Do the authors have any evidence that WA can be used for improved continual or transfer learning?", "questions_for_authors": "Why, in Figure 8, does annealing + weight averaging help whereas in figure 7b LAWA adds no benfit relative to the annealed model? Is it a different dataset, or difference in some other setting?", "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": "3a6NxHivZB", "reviewer": "Reviewer_egdu", "summary": "This submission presents the experimental findings and evaluation of two weight-averaging (WA) techniques, LAWA and EMA, on the AlgoPerf optimization benchmark, on seven tasks. They find strong training speed improvements (15% reduction in GPU-Hours), consistent with varying hyperparameters (learning rate and averaging window) with second-order optimizers. They also show a small improvement in generalization provided by WA. Finally, they provide experiments linking WA to learning rate annealing, showing that the benefits of WA diminish when annealing the LR to zero.\n\n[Score raised from 2 to 3]", "questions": "", "limitations": "", "rating": 3, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for the thorough review! We appreciate the detailed feedback.\n> I think only studying two WA methods is a bit limited.\n\nWe agree that broader benchmarking would be valuable and hope this works encourages further study at scale. We focus on LAWA and EMA due to their success and adoption. Influential works such as Szegedy et al. (2016), Vaswani et al. (2017), and Merity et al. (2017) used **one of these two averaging flavors**. We briefly tested Polyak averaging, but found it ineffective, consistent with Defazio et al. (2024), and omitted it. Finally, AlgoPerf's high computational cost limits extensive experimentation, so we prioritized exploring LAWA and EMA in depth.\n> I'm surprised the dependence on the EMA coefficient and update frequency was not studied.\n\nThank you for mentioning this! We recognize its importance and are happy to share additional results: https://anonymous.4open.science/r/ICMLE1C2/EMA.pdf\n1. Small values of $\\nu$ require a large $\\beta$, and viceversa.\n2. Update frequency around 0.003 with $\\beta=0.9$ work well across tasks.\n3. $\\nu=1$ (first group in each plot) performs poorly for $\\beta\\leq0.95$\n\n> Many experiments only consider LAWA and not EMA.\n\nWe acknowledge this and would like to share **additional experiments**:\n- [Shampoo](https://anonymous.4open.science/r/ICMLE1C2/Shampoo.pdf) We include **more workloads** and **EMA**. Shampoo consistently hits the target faster across tasks when equipped with LAWA or EMA.\n- [LR schedule](https://anonymous.4open.science/r/ICMLE1C2/Short_LR.pdf). Results on more **workloads** and **EMA** confirm that WA closely matches a short LR schedule.\n- We test [EMA with different annealing](https://anonymous.4open.science/r/ICMLE1C2/Annealing_EMA.png), and find consistent results with LAWA.\n- EMA efficiency is preserved [across different LR's](https://anonymous.4open.science/r/ICMLE1C2/LR_Sweep.pdf), similar to LAWA.\n\n> Can the authors detail the link between WA and LR decay that they find in Thm 5.3. in Garrigos & Gower?\n\nThm 5.3 provides convergence results for SGD under decreasing stepsizes, requiring an averaging scheme adapted to the stepsize schedule. A further connection to averaging stems from comparing this result to Lemma 7.3 where the gradient is computed at the running average of the iterates (see also Thm 7.4).\n>  Following Anonymous, 2025, WA perform worse than NadamW. What is the reason for this discrepancy? Could the authors give the additional time that averaging adds to training?\n\nThis is a subtle point, thank you for raising it.\n1) The AlgoPerf competition API did not allow switching model parameters to the WA buffer before evaluation, resulting to a suboptimal WA implementation. [This plot shows the](https://anonymous.4open.science/r/ICMLE1C2/Overhead.pdf) **overhead** of EMA. On Criteo it's 10x slower.\n2) AlgoPerf does not allow **asynchronous** CPU-GPU transfers and has limited VRAM, requiring WA buffers to be stored in CPU memory and transferred frequently over bandwidth-constrained hardware.\n\nHowever, **WA can be implemented efficiently** using asynchronous non-blocking transfers, as shown by Deepseek AI et al. (2024). We used offline averaging to reduce computational costs, but an efficient implementation would have negligible overhead.\n> The difference between LAWA and EMA is not really analyzed in the article.\n\nWe do not find significant difference between the two (Fig 2). Moreover, our goal was not to find _the best averaging scheme_, so we omitted a thorough comparison. However, we agree this could be better discussed and will update the paper.\n> What is the answer to the title of the paper?\n\nWA should always be used with a LR schedule to produce better models **for free** during training. How to access *the optimal* model at any time during training remains an open fascinating question. Our study presents WA as a step in this direction, moving the model closer to Pareto optimality of loss vs training time. This is precisely the reason behind the observed efficiency gains. We will clarify this point in the article.\n\n### Additional comments\n- We exclude ResNet because neither NadamW nor Shampoo reach the target (Kasimbeg et al, 2025), leaving no baseline for comparison.\n- Fig 7 is Librispeech Conformer, $\\eta_{max}=0.00175$, which is the best configuration from Dahl et al. (2023).\n- Fig 3: the y-axis shows the proportion of step wrt training horizon. Runs with $y>1$ do not reach the target and are hence excluded\n\nFinally, while some aspects of our work were explored in prior studies, our contribution expands their scale and scope. We believe our paper helps unify existing findings, evaluating them across tasks, and providing thoroughly validated claims.\n\n### Conclusion\nWe appreciate the insightful questions. We think that these additional analyses will significantly enhance our study and reinforce its contributions. We hope to have addressed your concerns, and if so, we would appreciate your endorsement for acceptance.", "claims_and_evidence": "The authors claim that their experiments show that WA methods speed-up training improve generalization, and that WA is a proxy for shorter learning rate decay, but can not replace it completely. I think that these claims are relatively well confirmed by the experiments presented, with a few caveats:\n\nI disagree with the claim that \"show that coupling WA with learning rate annealing yields optimal results\" (l49), since they later confirm that learning rate annealing is particularly the case where the benefits of WA disappear (Fig. 7b).\n\nSimilarly, the generalization improvement provided by WA is very limited (but consistent).\n\nFinally, see my main Question regarding speeding-up considering the results of the AlgoPerf competition (Anonymous, 2025 in the submission), which show WA methods as slower than the baseline.", "methods_and_evaluation_criteria": "The authors use 7 (of the 8) benchmarks of AlgoPerf to study the effectiveness of LAWA and EMA for deep learning optimization. This provides a thorough and comprehensive optimization landscape to study their claims. They also provide experiments regarding learning rate schedule, averaging window size and an experiment on a second-order optimizer.\n\nI think only studying two WA methods is a bit limited, as this would help better understand and choose WA options through extensive benchmarking, such as the effect of overfit-aware sampling and dense sampling. I'm also surprised the dependence on the EMA coefficient and update frequency was not studied like the window size, in particular when the authors claim (l189) to explore the interaction of the frequency variables with others.\n\nI am also surprised that many of the experiments (Figures 3, 4, 5, 6, 7, 8) only consider LAWA and neither EMA. This feels like an important oversight to justify the general and consistent effectiveness of WA methods to only consider one of the two approaches.", "theoretical_claims": "There are no theoretical claims in the article, but links with previous theoretical claims which are confirmed in this work are presented.", "experimental_designs_or_analyses": "The experiments done by the authors seem valid and consistent with previous findings. I voiced the concerns I had in other sections (relative lack of variety in experiments, inconsistent results in the AlgoPerf results..)", "relation_to_broader_scientific_literature": "This work is an experimental benchmarking of the effectiveness of WA methods for deep learning optimization. It is related to the WA and optimization literature, mainly confirming pre-existing findings on large optimization tasks.", "essential_references_not_discussed": "Most essential references have been discussed. I find the Model Soups part lacking more citations. The links between WA methods and self-supervised approaches could also be discussed, as well as more modern approaches using WA methods.", "other_strengths_and_weaknesses": "This paper is well-written and clear, showing the effectiveness of WA clearly.\n\nI find these results interesting, but I think that all results (except the combination with Shampoo, and the experiment in Fig 7b and 8) were already shown theoretically or experimentally by previous works: the speed-up and generalization improvement, and the main empirical link between WA and learning rate decay. I feel that this work either warrants further experiments with various WA approaches, or at least varying the hyperparameters missing (EMA coefficient our update frequency); or some more novel insights to justify an acceptance, but I can be convinced otherwise. As said by the authors (l361), studying a more broad range of optimizers could be interesting.", "other_comments_or_suggestions": "I'm a bit unsure if a second-order polynomial is an appropriate fit for the results in Figure 3, in particular for the upper results (Criteo, Conformer, DeepSpeech etc). Can the authors also provide the values for runs that result in a proportion over 1?\n\nPlease provide the value of $\\eta_{\\text{max}}$ (line 291).\n\nCan the authors please detail the link between WA and learning rate decay that they find in Theorem 5.3. in Garrigos & Gower ?\n\n\nl67/68: additional \",\"\n\nl316: \"Citeo1TBB\"\n\nFig 7b: \"diminsh\"", "questions_for_authors": "I have an important caveat on the results of the authors. Following Anonymous, 2025 (the results of the AlgoPerf competition cited by the authors); LAWA and EMA perform worse than the NadamW baseline (Fig 1. a/b, Table 1). I'm having a hard time making sense of the authors' results in this paper and the ones reported in the AlgoPerf competition, which seems to show that averaging weights only resulted in slower training. One reason I can find for this discrepancy is that the authors' report mainly the number of optimization steps and not the training time, which matters in the benchmark competition. However, this is not the case anywhere, and many claims in the submission are related to GPU-Hours training time (Tab 1 for instance). What is the reason for this discrepancy? Could the authors give the additional time that the averaging steps add to the standard training time, to better understand this possible slowdown? What throughput difference do the offline and online versions have?\n\nWhy did the authors not also report the performance for the ImageNet image classification task with a ResNet, the last benchmark in the AlgoPerf suite?\n\nIt is not clear what is the benchmark considered in Figure 7.\n\nThe difference in results between LAWA and EMA is not really analyzed in the article if I'm not mistaken.\n\nWhat is the answer to the title of this paper? The answer seems to be \"always use WA, as a way to better do learning rate decay\" to me while reading.", "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": "fj7fIEdgTx", "reviewer": "Reviewer_fFcZ", "summary": "The paper studies weight averaging, a \"trick\" rooted from the old Polyak averaging that never disappoints practitioners from fitting logistic regression to training modern LLMs. This paper provides a relatively large-scale study on AlgoPerf, a benchmark suite for optimizer, aiming at understanding: 1. Whether you should always use weight averaging; 2. How you set the weight averaging time window; 3. How you should use weight averaging together with learning rate decay.\nThe paper shows that: Using weight averaging always gives you faster convergence (i.e. fewer iterations and wall clock time) to a targeted validation loss and using a window size of 1% of the total training steps seems to be consistently the optimal option. Lastly the paper shows that although weight averaging has some connection with learning rate scheduling and can replace scheduling in some cases (e.g. The Road Less Scheduled, Defazio 2024), weight averaging still shows better performance when combined with standard learning rate scheduling.\n\nOverall, the problem studied is important and timely, the paper is mostly clearly written but I find some parts lacking details and need more empirical evidence.\n\n# Post rebuttal: The author agreed to include more discussion for the choice of WA parameters as well as other modifications I suggested. Based on that, I have increased my score from 3 to 4.", "questions": "", "limitations": "", "rating": 4, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for the thoughtful comments. We appreciate your recognition of the value of our study! We address specific points below.\n> The results for other claims are less thoroughly evaluated, e.g. the combination with Shampoo is only evaluated on one dataset. Comparision of WA with LR schedules (Fig 7) seems to be only based on one model.\n\nWe acknowledge that the mentioned experiments were missing some data, and we are happy to share **additional results**:\n- We performed [additional experiments on Shampoo](https://anonymous.4open.science/r/ICMLE1C2/Shampoo.pdf), considering **more workloads** and including **EMA**. Shampoo consistently hits the target faster across tasks when equipped with LAWA or EMA.\n- [WA vs Short LR schedule](https://anonymous.4open.science/r/ICMLE1C2/Short_LR.pdf) We report results on additional **workloads** and on **EMA**. LAWA and EMA yield better performance throughout training, closely matching the results from shorter schedules.\n- [We also included EMA](https://anonymous.4open.science/r/ICMLE1C2/Annealing_EMA.png) in the previous analysis on Librispeech Conformer, and find results to be consistent with the one on LAWA. When little or no annealing is applied, EMA provides substantial improvements on the final validation performance; however, when the LR is fully annealed to zero, EMA converges closely to the annealed model.\n- [LR sweep on EMA](https://anonymous.4open.science/r/ICMLE1C2/LR_Sweep.pdf): We add results on EMA when testing WA at different learning rates, observing that EMA efficiency is preserved across LR's.\n\n> For the recommendation of averaging window size, I think there are a couple of factors that could be important: batch size, weight decay, learning rate warmup.\n\nWe agree that **batch size** and **weight decay** might affect the optimal averaging configuration of LAWA and EMA. We appreciate you highlighting this, and we will include this clarification in the revised version of the paper. However, we would like to emphasize that:\n\n1. Due to the high computational cost of AlgoPerf and limited resources, we were unable to sweep all hyperparameters and instead focused on the best-performing combinations for each workload.\n2. These configurations account for a variety of hyperparameters, with batch size, weight decay, and learning rate warmup varying across workloads. As such, our recipe is relatively robust to these choices, though some variability may still occur.\n3. Finally, as suggested by Reviewer #egdu, [we included additional results on the effect of $\\beta$ and $\\nu$](https://anonymous.4open.science/r/ICMLE1C2/EMA.pdf) on EMA performance. We refer to the answer to Reviewer #egdu for a discussion on these additional experiments.\n\n> The type of learning rate scheduling considered is limited only to cosine schedule, more recent advances such as WSD or learning rate schedule with warmup would be interesting.\n\nWe would like to clarify that the employed schedule is Cosine Decay **with warmup**. We apologize if this was not clear and will ensure it is explicitly stated in the revision. Regarding recent advances like **WSD**, we find this to be a very interesting question! We would be particularly interested in whether averaging can accelerate the decay phase of WSD, potentially **reducing the computational cost of cooldowns**. We believe this is a promising area for future work, and hope that our study inspires further research on this topic.\n\n> It would be much better if the averaging strategies considered can come with some formulas to help readers understand better.\n\nThank you for pointing this out! We had included the algorithms in an earlier version of the paper but excluded them for the submission. [We display them at this URL](https://anonymous.4open.science/r/ICMLE1C2/Algorithms.png) and will include them in the Appendix of the revised manuscript.\n\n> Can the authors elaborate more on the computational overhead part? I don't see why it is the case\n\nYou are correct in noting that averaging does not incur additional overhead when implemented efficiently. In the quoted sentence, we were referring to a naive implementation of WA, where the average is offloaded to CPU memory and then transferred back to CUDA memory before evaluating the model. We appreciate the question and will revise the sentence to make this distinction clearer. For more details on this topic and data about the overhead, please refer to our response to Reviewer #egdu and this analysis.\n\n**To conclude**, we sincerely thank you for your review and thoughtful questions. The discussion has been valuable, and we believe that addressing these points and providing additional evidence has strengthened our work by broadening the scope of our analysis and better supporting our claims. We hope to have clarified your concerns, and we would greatly appreciate if these improvements could be reflected in a higher evaluation. We are happy to address any further questions or doubts.", "claims_and_evidence": "The empirical results for LAWA and EMA being helpful is clear and solid: Given a known optimal hyperparameter setting, applying WA always boosts the performance: Fig.1 and 2.\n\nThe results for other claims are less thoroughly evaluated, e.g.\n- the combination with Shampoo is only evaluated on one dataset\n- Comparision of WA + long decay schedule with short decay schedule without WA (Fig. 7a) seems to be only based on one model. Plus the model used is not stated in the caption.\n\nFor the recommendation of averaging window size, I think there are a couple of factors that could be important but not taken into account:\n- Batch size (see e.g. https://arxiv.org/abs/2307.13813)\n- Weight decay\n- Learning rate warmup", "methods_and_evaluation_criteria": "AlgoPerf used in the paper seems to be a sensible evaluation benchmark.", "theoretical_claims": "There are no theoretical claims in the submission.", "experimental_designs_or_analyses": "The experiment design aligned with the claims the paper is trying to make.\n\nThe most crucial hyperparmaeter: Learning rate, is swiped (Fig. 4).\n\nHowever, the type of learning rate scheduling considered is limited only to cosine schedule, some more recent advances such as WSD or learning rate schedule with warmup (although I guess the observations will still hold true).", "relation_to_broader_scientific_literature": "The idea and the observations are mostly known, and as the authors admit in the submission, many observations confirm the conclusion in past literature. However, I am not aware of any existing papers that actually perform such large-scale studies. So I think the paper could be helpful for the community as a source of empirical evidence.\n\nHowever, on the other hand, I am not sure if there are any new \"insights\" this submission provides to the community. The observation I find most interesting is the one in Fig.7a, where the authors show that a long run with weight-averaging, at the middle of training, gives you a performance similar to a full shorter run. I think this observation is relatively novel (at least for me) and could potentially help practitioners perform e.g. continuous pre-training, or tuning the training length, easier.", "essential_references_not_discussed": "Not any I am aware of.", "other_strengths_and_weaknesses": "It would be much better if the weight averaging strategies considered can come with some formulas to help readers understand better, the current version describes EMA and LAWA with only languages, which I don't find very straightforward.", "other_comments_or_suggestions": "I think Appendix A.2 could be moved into the main text since there are still sufficient space in the main text, and this seems to be pretty critical implementation details since it is highly related to the computation of the wall clock time and practical efficiency of WA.", "questions_for_authors": "In page 4 bottom, it is mentioned that \"whereas a short one may be suboptimal or result in a serious computational overhead.\", can the authors elaborate more on the computational overhead part? I don't see why it is the case if you are using, e.g. EMA.", "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:

Meta-review: This study explores the use of weight averaging (WA) across various neural network architectures and demonstrates its effectiveness in accelerating training and providing modest gains in generalization, using the AlgoPerf baseline. The results show that the advantages of WA are most prominent when the learning rate remains relatively high, and they diminish as the learning rate is annealed to zero. This suggests that WA is particularly beneficial during the more dynamic stages of training. The method appears to be generally applicable across architectures, though its impact may vary depending on the specific training regime. 

Almost all reviewers agree that the paper merits acceptance, and after carefully reading the rebuttal and discussion, I tend to concur. Please incorporate the reviewers’ feedback into the final version of the paper, with particular attention to the following concerns:

1-  Incorporating additional experiments related to hyper-parameter tuning 

2- Incorporate comparison with  other baseline optimizers ( as in your response to the reviewers )
Decision: Accept (poster)