---
title: Checkpoint Averaging in Model Training
url: https://www.emergentmind.com/topics/checkpoint-averaging
type: topic
---

# Checkpoint Averaging in Model Training

Checkpoint averaging is the post-training or training-time combination of multiple model snapshots taken from a single optimization trajectory. In the narrowest and most common sense, it denotes **parameter-space averaging**, in which checkpoint parameters \(\bm{\theta}_k\) are combined into a single model, typically by a uniform mean, \(\hat{\bm{\theta}}=\frac{1}{K}\sum_{k\in\mathcal{S}}\bm{\theta}_k\), or more generally by a convex combination \(\hat{\bm{\theta}}=\sum_{k\in\mathcal{S}}w_k\bm{\theta}_k\) with nonnegative weights summing to one [2210.11803]. The practice is attractive because it is computationally cheap, preserves single-model inference cost, and often improves generalization relative to a single checkpoint. Across the literature, however, the term also encompasses related but distinct procedures, including prediction-space checkpoint ensembles from one run [1710.03282], averaging of recent checkpoints during large-model pre-training [2209.14981, 2306.03241], adaptive subset selection for late-stage averages [2502.10119], and weighted merging of parameter-efficient fine-tuning checkpoints [2504.18580]. Some recent work further uses averaging only as a **checkpoint scoring** device rather than as model merging; in "Uncertainty-Guided Checkpoint Selection for Reinforcement Finetuning of Large Language Models," the averaged quantity is reward over a recent training window, not model weights [2511.09864].

## 1. Conceptual scope and historical formulations

The basic object is a sequence of checkpoints produced by iterative optimization. In neural machine translation, checkpoint averaging is typically defined over a selected set \(\mathcal{S}\) of checkpoints \(\bm{\theta}_k\), with either uniform or weighted coefficients [2210.11803]. In large-scale vision and language training, recent work has emphasized **latest weight averaging**, a sliding-window average over the most recent \(k\) checkpoints, often saved at epoch boundaries or fixed step intervals [2209.14981, 2306.03241]. In differentially private learning, analogous constructions include **uniform past-\(k\) average (UPA)** and **exponential moving average (EMA)** over the last \(k\) checkpoints of a single run [2210.01864].

A second lineage uses checkpoints as ensemble members rather than averaging their parameters. "Checkpoint Ensembles: Ensemble Methods from a Single Training Process" defines a model sequence \(M=\{M_1,\dots,M_n\}\), ranks checkpoints by validation score, and predicts by averaging the outputs of the best \(k\) checkpoints, \(M_{CE}(x_o)=\frac{1}{k}\sum_{i=1}^k M_{(i)}(x_o)\) [1710.03282]. This is not weight-space interpolation, but it addresses the same underlying problem: how to exploit useful variation along a single training trajectory.

A third usage is terminological rather than structural. In UGCS for RL finetuning of LLMs, the method does **not** average checkpoint parameters and does **not** ensemble predictions across checkpoints. Instead, for each checkpoint \(C\), it averages rewards of the top-\(p\%\) most uncertain samples from a recent training window \([C.\mathrm{step}-\delta,C.\mathrm{step})\), and uses that scalar as a selection score [2511.09864]. This distinction is central because the literature otherwise reserves “checkpoint averaging” primarily for parameter or prediction aggregation.

## 2. Selection rules and averaging operators

Averaging efficacy depends strongly on **which checkpoints are selected**. The NMT study "Revisiting Checkpoint Averaging for Neural Machine Translation" systematically contrasts two common strategies: **last-\(K\)** and **top-\(K\)** [2210.11803]. In last-\(K\), one averages consecutive checkpoints moving backward from the best checkpoint by development perplexity; in top-\(K\), one ranks all checkpoints by development perplexity and averages the best \(K\), whether or not they are consecutive. The reported trend is that last-\(K\) improves initially and then degrades as older, weaker checkpoints enter the mixture, whereas top-\(K\) is much less sensitive to \(K\) so long as all selected checkpoints remain strong by development perplexity [2210.11803].

A related design question is whether averaging should be uniform or weighted. The NMT paper studies a development-perplexity-based softmax weighting,
\[
w_k=\frac{\exp(-\tau \log \mathrm{DevPpl}_k)}{\sum_{k'\in\mathcal{S}}\exp(-\tau \log \mathrm{DevPpl}_{k'})},
\]
which interpolates between nearly uniform averaging at small \(\tau\) and near single-checkpoint selection at large \(\tau\) [2210.11803]. The central empirical conclusion is that weighted averaging can act as an automated soft selection mechanism when bad checkpoints are present, but does not clearly outperform the plain mean when the candidate checkpoints are already strong [2210.11803].

Later work generalizes this idea in more specialized settings. In PEFT checkpoint merging, "Parameter-Efficient Checkpoint Merging via Metrics-Weighted Averaging" studies normalized metric-derived weights over late LoRA or adapter checkpoints from a single run, using metrics such as **training loss** or **training step** together with a **penalty factor** controlling the sharpness of the weight distribution [2504.18580]. The merged checkpoint is still a convex combination, but only over the trainable PEFT parameters while the frozen base model remains unchanged [2504.18580]. In ASR, "ApproBiVT" replaces conventional **last-k (LK)** and **k-best validation loss (KBVL)** checkpoint averaging with **k-best ApproBiVT (KBABVT)** selection, where checkpoints are ranked by \(\mathrm{ApproBiVT}=\mathrm{SUTL}+\mathrm{validation\ loss}\) [2308.02870].

The operator itself also varies. Some methods average all saved checkpoints from a run, as in cross-lingual transfer with XLM-R, where the default strategy is to save snapshots periodically and average all of them uniformly before inference [2305.16834]. Others maintain a moving window over the latest \(k\) checkpoints, as in LAWA for ImageNet, WikiText-103, and LLM pre-training [2209.14981, 2306.03241]. Still others average only a learned subset, as in SeWA, which introduces binary selection variables \(m_i\in\{0,1\}\) over the last \(k\) checkpoints and forms
\[
\bar{w}^{K}_{T}=\frac{1}{K}\sum_{i=T-k+1}^{T}m_i w_i,
\]
with the subset learned through a probabilistic relaxation based on Gumbel-Softmax [2502.10119].

## 3. Geometry, optimization, and theoretical interpretations

A recurrent empirical explanation is that late checkpoints from a single run lie in a compatible low-loss region, so interpolation stays within a broad basin. The NMT study states that the landscape between converged checkpoints is “rather flat,” meaning BLEU or loss changes only slightly across much of the interpolation simplex, except near the individual checkpoint vertices [2210.11803]. Grid search over \(K=3\) interpolations on de-en shows a mostly flat interior; appendix analyses on en-de and zh-en show similar patterns, including cases where favoring better checkpoints over a weaker one can yield about **+0.4 absolute BLEU**, but without revealing a sharply superior interior optimum [2210.11803]. This flatness is used to explain why weighted interpolation, gradient-assisted moves, and development-set tuning of interpolation weights bring little or no consistent benefit over the simple mean [2210.11803].

For large-model pre-training, the explanation is often phrased in terms of **oscillation smoothing**. "Early Weight Averaging meets High Learning Rates for LLM Pre-training" argues that high learning rates accelerate progress but induce oscillation, especially along high-curvature directions, and that averaging checkpoints acts as a post hoc surrogate for learning-rate decay [2306.03241]. In that account, higher learning rates produce more diverse checkpoints, and averaging more separated checkpoints can better cancel oscillatory noise, although the appendix also shows that excessively stale checkpoints can hurt and that, on Pythia-1B, \(\nu=1\mathrm{K}\) or \(2\mathrm{K}\) worked better than \(5\mathrm{K}\) or \(10\mathrm{K}\) [2306.03241]. The same paper ties successful averaging to the onset of **linear mode connectivity**, noting that very early checkpoints are not linearly connected and that averaging only becomes reliable after the trajectory has entered a linearly connected basin [2306.03241].

A more formal connection appears in "WSM: Decay-Free Learning Rate Schedule via Checkpoint Merging for LLM Pre-training" [2507.17634]. WSM defines a weighted merged model
\[
\hat{\theta}_{n+k}=\sum_{j=0}^{k}c_j\theta_{n+j},
\]
then rewrites each checkpoint in terms of a base checkpoint plus subsequent updates, obtaining
\[
\hat{\theta}_{n+k}=\theta_n-\sum_{i=1}^{k}w_i g_{n+i-1},\qquad w_i=\sum_{j=i}^{k}c_j.
\]
This expresses checkpoint merging as reweighting historical updates and motivates the claim that appropriately chosen checkpoint weights can emulate discrete decay schedules such as **linear**, **cosine**, and **inverse square root** decay [2507.17634]. The theorem-derived mapping
\[
c_k=w_k,\qquad c_j=w_j-w_{j+1}\ (j\in[1,k-1]),\qquad c_0=1-w_1
\]
is valid whenever \(1\ge w_1\ge \dots \ge w_k\ge 0\) [2507.17634]. In this view, checkpoint averaging is not merely smoothing but a decay-free surrogate for LR annealing.

The differential privacy literature offers a distinct theoretical perspective. "Recycling Scraps" analyzes uniform tail averaging in convex DP-SGD and shows that the averaged tail iterate can improve excess empirical risk relative to the last checkpoint by removing an extra \(\log(n)\) factor [2210.01864]. The same paper further proves, under standard assumptions on the loss function, that sample variance from the last few checkpoints provides a good approximation to the variance of the final model of a DP run, with burn-in and spacing used to justify approximate independence [2210.01864]. SeWA, by contrast, derives stability-based generalization bounds for **selective** averaging and shows sharper bounds than SGD under both convex and non-convex assumptions, with improvements attributed jointly to late-stage averaging and selection probabilities \(s_i\) [2502.10119].

## 4. Methodological families and neighboring techniques

The literature now contains several method families that are closely related but not interchangeable.

| Family | Averaged object | Representative formulation |
|---|---|---|
| Parameter-space averaging | Model weights from one run | \(\hat{\bm{\theta}}=\frac{1}{K}\sum \bm{\theta}_k\) [2210.11803] |
| Prediction-space checkpoint ensembling | Outputs of saved checkpoints | \(M_{CE}(x_o)=\frac{1}{k}\sum M_{(i)}(x_o)\) [1710.03282] |
| Online moving averages | Running weighted average during training | \(\theta^t_{ema}=(1-\beta_t)\theta^{t-1}_{ema}+\beta_t\theta^t\) [2210.01864] |

Prediction-space ensembling and parameter-space averaging often address the same practical objective but have different deployment tradeoffs. In the 2017 checkpoint ensemble study, **checkpoint ensembles (CE)** consistently outperform **minimum validation** model selection and also outperform the paper’s weight-averaging baselines—**checkpoint smoothers (CS)** and **last \(k\) smoothers (LKS)**—especially on CIFAR-10, where weight averaging is harmful [1710.03282]. The price is inference cost, since CE requires multiple forward passes, whereas CS and LKS produce a single deployable model [1710.03282].

Online moving averages likewise differ operationally from post hoc checkpoint averaging. In DP training, **EMA** is maintained during training with a warmup schedule on \(\beta_t\), while **UPA** averages the last \(k\) checkpoints after training; both are privacy-preserving post-processing of a single DP run [2210.01864]. In LLM pre-training, LAWA is explicitly contrasted with **EMA** and **SWA**. LAWA is a post hoc or periodically emitted average over the latest \(k\) checkpoints, whereas SWA is treated as a late-training procedure that changes the learning-rate schedule, and early SWA is reported to diverge across all nanoGPT-2 sizes in that study [2306.03241]. In the ImageNet and WikiText-103 LAWA work, the averaged weights are evaluated throughout training but do **not** replace the optimizer’s live parameters in subsequent updates [2209.14981].

Adaptive methods add another axis. SeWA treats checkpoint averaging as a subset-selection problem over the last \(k\) checkpoints and learns which checkpoints to include through Bernoulli masks relaxed by Gumbel-Softmax [2502.10119]. ApproBiVT changes not the averaging operator but the **checkpoint ranking criterion** and **early stopping rule**, arguing that conventional LK and KBVL heuristics are too one-sided because they target low bias or low variance rather than their tradeoff [2308.02870]. WSM then reframes merging as a schedule-design problem and studies mean, EMA, and theorem-derived weighted merges as surrogates for LR decay [2507.17634].

## 5. Domain-specific uses

Checkpoint averaging is unusually pervasive in neural machine translation. The 2022 NMT re-evaluation covers six MT datasets and language pairs—**IWSLT14 de-en, ru-en, es-en; WMT16 en-ro; WMT14 en-de; WMT19 zh-en**—and concludes that simple checkpoint averaging is necessary for optimal performance, with development perplexity being a better basis for checkpoint choice than blind recency [2210.11803]. The final metric of interest is **test BLEU**, and the main practical lesson is to average several converged, high-quality checkpoints, preferably with top-\(K\) selection by development perplexity [2210.11803].

In large-model pre-training, checkpoint averaging is used to accelerate effective convergence or replace decay schedules. LAWA on ResNet50/ImageNet reaches around **66%** validation accuracy **40 epochs earlier** than baseline SGD, corresponding to about **68 GPU hours saved**, while on RoBERTa-Base/WikiText-103 it reaches Adam’s final best validation loss **45 epochs earlier**, saving about **30 GPU hours** [2209.14981]. In LLM pre-training, early weight averaging along high-learning-rate runs improves validation loss on nanoGPT-2 and held-out perplexity or downstream metrics on Pythia models, with strong gains in early-to-mid training and an explicit emphasis on starting after the unstable initial phase [2306.03241]. WSM extends this line by replacing WSD decay with constant-LR training plus late checkpoint merging; on its 16.3B total / 1.4B active MoE setup, the abstract reports improvements over WSD of **+3.5% on MATH**, **+2.9% on HumanEval**, and **+5.5% on MMLU-Pro** [2507.17634].

In cross-lingual transfer, checkpoint averaging is used primarily as a **validation-free robustness mechanism**. "Free Lunch: Robust Cross-Lingual Transfer via Model Checkpoint Averaging" averages checkpoints from a single fine-tuning run, usually one per epoch over **10 epochs**, using XLM-R\(_{\text{base}}\) on NLI, TyDiQA-GoldP, NER, and POS tasks [2305.16834]. The method is motivated by the observation that target-language validation-based checkpoint selection violates the assumptions of “true” ZS-XLT and FS-XLT. Empirically, CA improves or matches realistic baselines and reduces variance; for example, on ZS-XLT NLI, **LAST** gives **61.8 ± 0.3**, **SRC-DEV** **61.9 ± 0.3**, and **CA** **62.8 ± 0.1** [2305.16834].

In ASR, checkpoint averaging is embedded in a broader early-stopping recipe. ApproBiVT evaluates each checkpoint by **sampled unaugmented training loss (SUTL)** plus validation loss, stops when the ApproBiVT score monotonically increases for \(S=5\) epochs, and then performs **KBABVT** checkpoint averaging over the \(k\) checkpoints with lowest ApproBiVT [2308.02870]. On AISHELL-1 Conformer, the strongest ablation result is **4.39 CER** for **ApproBiVT + KBABVT averaging 100 checkpoints**, compared with **4.56** for a validation-loss-driven baseline-like recipe [2308.02870].

In differentially private learning, aggregation over intermediate checkpoints is especially attractive because all checkpoints from a single run are already private outputs and can therefore be post-processed with **no additional privacy cost** [2210.01864]. The paper reports gains on StackOverflow, CIFAR10, and CIFAR100, with particularly large improvements in periodically distribution-shifting settings; for example, on PDS CIFAR10 at \(\varepsilon=8\), the baseline is **60.74** and best **OPA** reaches **78.27** [2210.01864].

In PEFT and LoRA-style fine-tuning, checkpoint merging is cheap because only adapter parameters are merged. Metrics-Weighted Averaging reports that merged late checkpoints can improve over the final checkpoint on a GSM weighted benchmark and an alignment weighted benchmark, with a best math score increase from **0.21728** to **0.22826** (**+5.05% relative**) and a best alignment increase from **0.4872** to **0.4920** (**+1.17% relative**) [2504.18580]. At the same time, the paper’s direct baseline comparisons show that plain unweighted averaging is often stronger than MWA, reinforcing the broader literature’s conclusion that simple means are difficult to beat reliably [2504.18580].

## 6. Practical considerations, misconceptions, and limits

Several points recur across domains. First, checkpoint averaging is usually most effective when checkpoints are **late**, **compatible**, and drawn from the **same optimization trajectory**. Uniformly averaging all checkpoints can be robust in some settings, such as XLT [2305.16834], but can include harmful early snapshots on small or unstable datasets, as noted for ZS-XLT TyDiQA-GoldP [2305.16834]. In NMT, performance degrades once the average begins to include clearly weaker earlier checkpoints [2210.11803]. In LLM pre-training, averaging too early is unreliable, and successful recipes explicitly begin after the unstable initial phase or after roughly **10%–15% of training** [2306.03241].

Second, more sophistication does not automatically help. The strongest single conclusion of the NMT re-evaluation is that one should do checkpoint averaging at all, and choose good checkpoints, but not expect **weighted averaging**, **gradient-assisted interpolation**, or **development-set optimization of interpolation weights** to yield meaningful, reliable improvements over the simple mean [2210.11803]. PEFT merging reaches a similar practical conclusion: metric-driven weighting is plausible, yet uniform averaging remains a remarkably strong baseline and often wins in direct comparisons [2504.18580].

Third, “checkpoint averaging” is not synonymous with all checkpoint-based combination methods. Parameter-space averaging produces one deployable model with no runtime ensemble overhead [2210.11803]. Prediction-space checkpoint ensembles can be more robust empirically but require multiple forward passes at inference [1710.03282]. Windowed reward averaging in UGCS is a selection-time statistic rather than model merging [2511.09864]. Treating these as interchangeable obscures both their mechanics and their deployment costs.

Fourth, resource tradeoffs differ by method. Post hoc uniform or weighted parameter averaging is typically cheap in FLOPs and expensive mainly in storage, because multiple checkpoints must be retained [2209.14981, 2306.03241]. Prediction-space checkpoint ensembles add inference cost proportional to the number of retained models [1710.03282]. In PEFT, checkpoint merging is especially lightweight because only adapter parameters are combined [2504.18580]. In DP learning, aggregation is unusual in that it improves utility while preserving privacy expenditure, since it operates entirely by post-processing already released private checkpoints [2210.01864].

A common misconception is that checkpoint averaging is merely a late-stage cosmetic improvement. Multiple papers directly reject that interpretation. LAWA is explicitly motivated as a way to **speed up the training progression** itself [2209.14981]; early weight averaging for LLM pre-training is presented as a way to recover high-LR generalization early and throughout training [2306.03241]; WSM makes checkpoint merging the core mechanism replacing LR decay [2507.17634]. Another misconception is that any convex combination within the checkpoint simplex hides a sharply better model. The NMT evidence argues the opposite: once one is averaging several good checkpoints, the interpolation region is already broad and forgiving, and there is simply not much left to optimize in interpolation space [2210.11803].

Overall, checkpoint averaging has evolved from an empirical heuristic into a broad family of trajectory-reuse techniques. Despite considerable variation in objective, schedule, and theoretical framing, the dominant empirical pattern is stable: averaging or otherwise aggregating compatible late checkpoints from a single run is often a low-cost route to better generalization, greater robustness, or a more favorable training trajectory, whereas the main unresolved questions concern *which* checkpoints to include, *when* averaging should begin, and *when* more elaborate weighting schemes justify their complexity [2210.11803, 2306.03241, 2502.10119, 2507.17634].

Source: https://www.emergentmind.com/topics/checkpoint-averaging