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

The Mismatch Between Training and Inference

A new paper from arXiv (2606.25086v2) tackles a subtle but consequential problem in modern language model training: the disconnect between how models are optimized and how they are ultimately served. The core insight is that many production LM pipelines use an averaged model—typically an exponential moving average (EMA) of training checkpoints—as the final returned artifact, yet the optimization process itself remains tuned for the raw, un-averaged parameters.

This creates a fundamental misalignment. Standard optimization algorithms like Adam or SGD are designed to minimize loss for the current iterate, not for the smoothed version that will actually be deployed. The authors propose modifying the training objective to directly optimize for the averaged model’s performance, rather than treating averaging as an afterthought applied post-training.

Why This Matters

The practical implications are significant. EMA averaging is widely adopted because it reduces variance and often yields better generalization than any single checkpoint. However, the current approach is essentially heuristic: practitioners average first, then hope the result works well. This paper suggests that by explicitly optimizing for the averaged model during training, we can achieve better final performance without additional compute or data.

For AI practitioners, this addresses a known pain point: the “best” checkpoint during training often differs from the best averaged model, leading to wasted evaluation cycles and suboptimal deployment choices. By aligning the optimization target with the deployment artifact, this method could streamline model selection and improve out-of-the-box performance.

Implications for AI Practitioners

First, this research underscores that infrastructure choices (like model averaging) should inform algorithm design, not the other way around. Teams using EMA should consider whether their learning rate schedules, weight decay, and batch sizes are appropriate for the averaged model, not just the raw iterates.

Second, the approach is computationally lightweight—it modifies the loss signal rather than requiring additional forward passes or ensembles. This makes it practical for large-scale training where every FLOP counts.

Third, this work highlights a broader trend: as LM pipelines mature, the gap between training theory and deployment practice is narrowing. Optimizers designed for toy problems are being re-evaluated against real-world constraints like checkpoint averaging, quantization, and distillation.

Key Takeaways

• Alignment gap: Current LM training optimizes for raw iterates, but deployment often uses averaged models—creating a mismatch that leaves performance on the table.
• Practical fix: Modifying the optimization objective to account for the final averaged model can improve results without extra compute.
• Infrastructure-aware algorithms: This work signals a shift toward training methods that explicitly account for downstream deployment choices (averaging, quantization, etc.).
• Low overhead: The proposed approach is lightweight and compatible with existing training pipelines, making it immediately relevant for production teams.