Checkpoint Averaging in Model Training

• Checkpoint Averaging is a technique that combines several model snapshots from a training run to yield a single model with improved generalization.
• It leverages methods like uniform, weighted, and moving averages to smooth oscillatory noise and mimic learning rate decay without extra inference cost.
• This approach is widely applied in neural machine translation, LLM pre-training, and differentially private training, enhancing performance while maintaining efficiency.

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 (Gao et al., 2022). 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 (Chen et al., 2017), averaging of recent checkpoints during large-model pre-training (Kaddour, 2022, Sanyal et al., 2023), adaptive subset selection for late-stage averages (Wang et al., 14 Feb 2025), and weighted merging of parameter-efficient fine-tuning checkpoints (Yu et al., 23 Apr 2025). 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 LLMs," the averaged quantity is reward over a recent training window, not model weights (Nguyen et al., 13 Nov 2025).

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 (Gao et al., 2022). 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 (Kaddour, 2022, Sanyal et al., 2023). 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 (Shejwalkar et al., 2022).

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, $$\hat{\bm{\theta}}=\frac{1}{K}\sum_{k\in\mathcal{S}}\bm{\theta}_k$$0 (Chen et al., 2017). 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 $$\hat{\bm{\theta}}=\frac{1}{K}\sum_{k\in\mathcal{S}}\bm{\theta}_k$$1, it averages rewards of the top-$$\hat{\bm{\theta}}=\frac{1}{K}\sum_{k\in\mathcal{S}}\bm{\theta}_k$$2 most uncertain samples from a recent training window $$\hat{\bm{\theta}}=\frac{1}{K}\sum_{k\in\mathcal{S}}\bm{\theta}_k$$3, and uses that scalar as a selection score (Nguyen et al., 13 Nov 2025). 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-$$\hat{\bm{\theta}}=\frac{1}{K}\sum_{k\in\mathcal{S}}\bm{\theta}_k$$4 and top-$$\hat{\bm{\theta}}=\frac{1}{K}\sum_{k\in\mathcal{S}}\bm{\theta}_k$$5 (Gao et al., 2022). In last-$$\hat{\bm{\theta}}=\frac{1}{K}\sum_{k\in\mathcal{S}}\bm{\theta}_k$$6, one averages consecutive checkpoints moving backward from the best checkpoint by development perplexity; in top-$$\hat{\bm{\theta}}=\frac{1}{K}\sum_{k\in\mathcal{S}}\bm{\theta}_k$$7, one ranks all checkpoints by development perplexity and averages the best $$\hat{\bm{\theta}}=\frac{1}{K}\sum_{k\in\mathcal{S}}\bm{\theta}_k$$8, whether or not they are consecutive. The reported trend is that last-$$\hat{\bm{\theta}}=\frac{1}{K}\sum_{k\in\mathcal{S}}\bm{\theta}_k$$9 improves initially and then degrades as older, weaker checkpoints enter the mixture, whereas top-$$\hat{\bm{\theta}}=\sum_{k\in\mathcal{S}}w_k\bm{\theta}_k$$0 is much less sensitive to $$\hat{\bm{\theta}}=\sum_{k\in\mathcal{S}}w_k\bm{\theta}_k$$1 so long as all selected checkpoints remain strong by development perplexity (Gao et al., 2022).

A related design question is whether averaging should be uniform or weighted. The NMT paper studies a development-perplexity-based softmax weighting,

$$\hat{\bm{\theta}}=\sum_{k\in\mathcal{S}}w_k\bm{\theta}_k$$2

which interpolates between nearly uniform averaging at small $$\hat{\bm{\theta}}=\sum_{k\in\mathcal{S}}w_k\bm{\theta}_k$$3 and near single-checkpoint selection at large $$\hat{\bm{\theta}}=\sum_{k\in\mathcal{S}}w_k\bm{\theta}_k$$4 (Gao et al., 2022). 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 (Gao et al., 2022).

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 (Yu et al., 23 Apr 2025). The merged checkpoint is still a convex combination, but only over the trainable PEFT parameters while the frozen base model remains unchanged (Yu et al., 23 Apr 2025). 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 $$\hat{\bm{\theta}}=\sum_{k\in\mathcal{S}}w_k\bm{\theta}_k$$5 (Wang et al., 2023).

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 (Schmidt et al., 2023). Others maintain a moving window over the latest $$\hat{\bm{\theta}}=\sum_{k\in\mathcal{S}}w_k\bm{\theta}_k$$6 checkpoints, as in LAWA for ImageNet, WikiText-103, and LLM pre-training (Kaddour, 2022, Sanyal et al., 2023). Still others average only a learned subset, as in SeWA, which introduces binary selection variables $$\hat{\bm{\theta}}=\sum_{k\in\mathcal{S}}w_k\bm{\theta}_k$$7 over the last $$\hat{\bm{\theta}}=\sum_{k\in\mathcal{S}}w_k\bm{\theta}_k$$8 checkpoints and forms

$$\hat{\bm{\theta}}=\sum_{k\in\mathcal{S}}w_k\bm{\theta}_k$$9

with the subset learned through a probabilistic relaxation based on Gumbel-Softmax (Wang et al., 14 Feb 2025).

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 (Gao et al., 2022). Grid search over $\mathcal{S}$0 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 (Gao et al., 2022). 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 (Gao et al., 2022).

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 (Sanyal et al., 2023). 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, $\mathcal{S}$1 or $\mathcal{S}$2 worked better than $\mathcal{S}$3 or $\mathcal{S}$4 (Sanyal et al., 2023). 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 (Sanyal et al., 2023).

A more formal connection appears in "WSM: Decay-Free Learning Rate Schedule via Checkpoint Merging for LLM Pre-training" (Tian et al., 23 Jul 2025). WSM defines a weighted merged model

$\mathcal{S}$5

then rewrites each checkpoint in terms of a base checkpoint plus subsequent updates, obtaining

$\mathcal{S}$6

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 (Tian et al., 23 Jul 2025). The theorem-derived mapping

$\mathcal{S}$7

is valid whenever $\mathcal{S}$8 (Tian et al., 23 Jul 2025). 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 $\mathcal{S}$9 factor (Shejwalkar et al., 2022). 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 (Shejwalkar et al., 2022). 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 $\bm{\theta}_k$0 (Wang et al., 14 Feb 2025).

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	$\bm{\theta}_k$1 (Gao et al., 2022)
Prediction-space checkpoint ensembling	Outputs of saved checkpoints	$\bm{\theta}_k$2 (Chen et al., 2017)
Online moving averages	Running weighted average during training	$\bm{\theta}_k$3 (Shejwalkar et al., 2022)

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 $\bm{\theta}_k$4 smoothers (LKS)—especially on CIFAR-10, where weight averaging is harmful (Chen et al., 2017). The price is inference cost, since CE requires multiple forward passes, whereas CS and LKS produce a single deployable model (Chen et al., 2017).

Online moving averages likewise differ operationally from post hoc checkpoint averaging. In DP training, EMA is maintained during training with a warmup schedule on $\bm{\theta}_k$5, while UPA averages the last $\bm{\theta}_k$6 checkpoints after training; both are privacy-preserving post-processing of a single DP run (Shejwalkar et al., 2022). In LLM pre-training, LAWA is explicitly contrasted with EMA and SWA. LAWA is a post hoc or periodically emitted average over the latest $\bm{\theta}_k$7 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 (Sanyal et al., 2023). 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 (Kaddour, 2022).

Adaptive methods add another axis. SeWA treats checkpoint averaging as a subset-selection problem over the last $\bm{\theta}_k$8 checkpoints and learns which checkpoints to include through Bernoulli masks relaxed by Gumbel-Softmax (Wang et al., 14 Feb 2025). 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 (Wang et al., 2023). WSM then reframes merging as a schedule-design problem and studies mean, EMA, and theorem-derived weighted merges as surrogates for LR decay (Tian et al., 23 Jul 2025).

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 (Gao et al., 2022). The final metric of interest is test BLEU, and the main practical lesson is to average several converged, high-quality checkpoints, preferably with top-$\bm{\theta}_k$9 selection by development perplexity (Gao et al., 2022).

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 (Kaddour, 2022). 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 (Sanyal et al., 2023). 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 (Tian et al., 23 Jul 2025).

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$k$0 on NLI, TyDiQA-GoldP, NER, and POS tasks (Schmidt et al., 2023). 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 (Schmidt et al., 2023).

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 $k$1 epochs, and then performs KBABVT checkpoint averaging over the $k$2 checkpoints with lowest ApproBiVT (Wang et al., 2023). 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 (Wang et al., 2023).

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 (Shejwalkar et al., 2022). The paper reports gains on StackOverflow, CIFAR10, and CIFAR100, with particularly large improvements in periodically distribution-shifting settings; for example, on PDS CIFAR10 at $k$3, the baseline is 60.74 and best OPA reaches 78.27 (Shejwalkar et al., 2022).

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) (Yu et al., 23 Apr 2025). 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 (Yu et al., 23 Apr 2025).

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 (Schmidt et al., 2023), but can include harmful early snapshots on small or unstable datasets, as noted for ZS-XLT TyDiQA-GoldP (Schmidt et al., 2023). In NMT, performance degrades once the average begins to include clearly weaker earlier checkpoints (Gao et al., 2022). 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 (Sanyal et al., 2023).

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 (Gao et al., 2022). 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 (Yu et al., 23 Apr 2025).

Third, “checkpoint averaging” is not synonymous with all checkpoint-based combination methods. Parameter-space averaging produces one deployable model with no runtime ensemble overhead (Gao et al., 2022). Prediction-space checkpoint ensembles can be more robust empirically but require multiple forward passes at inference (Chen et al., 2017). Windowed reward averaging in UGCS is a selection-time statistic rather than model merging (Nguyen et al., 13 Nov 2025). 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 (Kaddour, 2022, Sanyal et al., 2023). Prediction-space checkpoint ensembles add inference cost proportional to the number of retained models (Chen et al., 2017). In PEFT, checkpoint merging is especially lightweight because only adapter parameters are combined (Yu et al., 23 Apr 2025). 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 (Shejwalkar et al., 2022).

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 (Kaddour, 2022); early weight averaging for LLM pre-training is presented as a way to recover high-LR generalization early and throughout training (Sanyal et al., 2023); WSM makes checkpoint merging the core mechanism replacing LR decay (Tian et al., 23 Jul 2025). 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 (Gao et al., 2022).

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 (Gao et al., 2022, Sanyal et al., 2023, Wang et al., 14 Feb 2025, Tian et al., 23 Jul 2025).