# Training for the Model You Return: Improving Optimization for Iterate-Averaged Language Models

arXiv: https://arxiv.org/abs/2606.25086 · HTML: https://mlsanity.com/paper/2606.25086/

## Key facts

- Cited by 1 papers in the map (data through 14 Oct 2026).
- Builds on 28 of its 45 references that are in the map.
- First posted on arXiv 23 Jun 2026.

## Builds on (most cited first, 28 of 28)

- [Qwen3 Technical Report](https://mlsanity.com/paper/2505.09388/) (2025)
- [Training Compute-Optimal Large Language Models](https://mlsanity.com/paper/2203.15556/) (2022)
- [Decoupled Weight Decay Regularization](https://mlsanity.com/paper/1711.05101/) (2017)
- [Model soups: averaging weights of multiple fine-tuned models improves accuracy without increasing inference time](https://mlsanity.com/paper/2203.05482/) (2022)
- [Averaging Weights Leads to Wider Optima and Better Generalization](https://mlsanity.com/paper/1803.05407/) (2018)
- [The FineWeb Datasets: Decanting the Web for the Finest Text Data at Scale](https://mlsanity.com/paper/2406.17557/) (2024)
- [Tulu 3: Pushing Frontiers in Open Language Model Post-Training](https://mlsanity.com/paper/2411.15124/) (2024)
- [Shampoo: Preconditioned Stochastic Tensor Optimization](https://mlsanity.com/paper/1802.09568/) (2018)
- [2 OLMo 2 Furious](https://mlsanity.com/paper/2501.00656/) (2024)
- [Lookahead Optimizer: k steps forward, 1 step back](https://mlsanity.com/paper/1907.08610/) (2019)
- [SmolLM2: When Smol Goes Big -- Data-Centric Training of a Small Language Model](https://mlsanity.com/paper/2502.02737/) (2025)
- [Sophia: A Scalable Stochastic Second-order Optimizer for Language Model Pre-training](https://mlsanity.com/paper/2305.14342/) (2023)
- [SOAP: Improving and Stabilizing Shampoo using Adam](https://mlsanity.com/paper/2409.11321/) (2024)
- [The Road Less Scheduled](https://mlsanity.com/paper/2405.15682/) (2024)
- [Generative Modeling with Denoising Auto-Encoders and Langevin Sampling](https://mlsanity.com/paper/2002.00107/) (2020)
- [Which Algorithmic Choices Matter at Which Batch Sizes? Insights From a Noisy Quadratic Model](https://mlsanity.com/paper/1907.04164/) (2019)
- [Analysis and Design of Optimization Algorithms via Integral Quadratic Constraints](https://mlsanity.com/paper/1408.3595/) (2014)
- [On the SDEs and Scaling Rules for Adaptive Gradient Algorithms](https://mlsanity.com/paper/2205.10287/) (2022)
- [Deep learning with Elastic Averaging SGD](https://mlsanity.com/paper/1412.6651/) (2014)
- [Understanding Warmup-Stable-Decay Learning Rates: A River Valley Loss Landscape Perspective](https://mlsanity.com/paper/2410.05192/) (2024)
- [Stop Wasting My Time! Saving Days of ImageNet and BERT Training with Latest Weight Averaging](https://mlsanity.com/paper/2209.14981/) (2022)
- [The AdEMAMix Optimizer: Better, Faster, Older](https://mlsanity.com/paper/2409.03137/) (2024)
- [How to Scale Your EMA](https://mlsanity.com/paper/2307.13813/) (2023)
- [Butterfly Effects of SGD Noise: Error Amplification in Behavior Cloning and Autoregression](https://mlsanity.com/paper/2310.11428/) (2023)
- [Through the River: Understanding the Benefit of Schedule-Free Methods for Language Model Training](https://mlsanity.com/paper/2507.09846/) (2025)
- [KOALA: A Kalman Optimization Algorithm with Loss Adaptivity](https://mlsanity.com/paper/2107.03331/) (2021)
- [EMA Without the Lag: Bias-Corrected Iterate Averaging Schemes](https://mlsanity.com/paper/2508.00180/) (2025)
- [KOALA++: Efficient Kalman-Based Optimization with Gradient-Covariance Products](https://mlsanity.com/paper/2506.04432/) (2025)

## Built on by (most cited first, 1 of 1)

- [Terminal Shrinkage Averaging Reveals a Schedule-Estimator Interaction in LLM Pretraining](https://mlsanity.com/paper/2609.25482/) (2026)

## Abstract (arXiv)

Many modern Language Model (LM) pipelines return an averaged model, such as an exponential moving average of the training iterates, rather than the final iterate itself. This raises a fundamental question: given that we will return an iterate average, how should we change training to improve the performance of this average? We study this question by formulating optimizer design for the iterate-average estimator as an optimal-control problem. In a continuous-time stochastic quadratic model, we solve for the control strategy that minimizes the error of the returned average subject to a penalty on the size of the intervention. A practical approximation to this controller yields PACE, a lightweight wrapper around AdamW that pulls the live weights toward their exponential moving average with a clipped, per-coordinate control strength. We prove that a stylized version of PACE converges at the standard stochastic convex optimization rate, up to a factor depending on the averaging rule, while in the quadratic setting it can strictly improve the limiting squared error of the iterate-average estimator and can do so by an arbitrarily large factor on some instances. Empirically, our results suggest that PACE improves over AdamW and EMA-evaluated AdamW in supervised fine-tuning of 1-2B parameter LMs and in GPT-2 pretraining on FineWeb for a wide range of learning rates, decay schedules, and other hyperparameters.