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

Abstract: Many modern LLM (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.

• The paper introduces PACE as a novel optimizer designed specifically for training trajectories that yield optimal iterate-averaged returns.
• It formulates the optimization process as a stochastic control problem, providing closed-form solutions and robust theoretical convergence guarantees.
• Empirical evaluations show PACE outperforms conventional optimizers across diverse language models, eliminating the need for learning rate decay.

Training for the Model You Return: Optimization Methods for Iterate-Averaged LLMs

Motivation and Problem Statement

Modern LLM training frequently uses iterate-averaged estimators—typically exponential moving average (EMA) of weights—to improve both optimization stability and downstream performance. However, nearly all existing optimization algorithms are designed and analyzed with respect to the final iterate, not the average, despite the ubiquity of weight averaging in production deployments. This paper directly addresses the fundamental question: how should the training process be modified if the goal is to return an iterate average? Rather than simply increasing learning rates or decaying less aggressively, the authors formulate optimizer design for the iterate-average estimator itself as an optimal-control problem. Theoretical analysis and practical approximations culminate in a new optimizer, PACE—a lightweight wrapper around AdamW—that explicitly controls training trajectories for optimal returned averages (2606.25086).

Theoretical Derivation via Stochastic Control

The authors build on classical stochastic convex optimization and continuous-time analysis, modeling the optimization process as an Ornstein-Uhlenbeck SDE around a quadratic minimum. Introducing an additive control term, they pose the trajectory modification as a stochastic linear-quadratic control problem. The objective is to minimize the squared error of the returned iterate average with a regularization penalty on control magnitude.

• Optimal controller: Derived in closed form for continuous-time quadratic objectives with additive noise. The controller pulls weights toward both the optimum and the current average.
• Certainty equivalence principle: When the optimum is unknown, the controller tracks the best available estimator (the EMA itself) in practice.
• Approximate discretization: The optimal continuous-time controller is heuristically approximated for practical implementation, leading to the PACE update:

$$\theta_{k+1} = \theta_k - \eta g_k + C (\text{EMA}_k - \theta_k)$$

where $C$ is a (clipped) per-coordinate control strength, and EMA is updated as usual.

Convergence Guarantees and Quadratic Improvements

Theoretical analysis extends to general convex losses, showing that:

• PACE achieves the same stochastic convex optimization rate as SGD up to a constant factor determined by the averaging rule (Theorem 3).
• Quadratic setting: In locally quadratic objectives, the controlled dynamics yield strictly lower limiting squared error for the iterate-average estimator compared to standard SGD, and this improvement can be arbitrarily large for certain instances (Propositions 2 and 3).

Thus, PACE is never substantially worse than SGD for convex objectives, while being provably much better in locally quadratic regimes.

Practical Algorithm: PACE (Pullback Averaging Control)

PACE is implemented as a wrapper around AdamW (with momentum and second-moment preconditioning), and concretely pulls live weights toward their EMA, with per-coordinate gain clipped for stability. Hyperparameters include:

• Pullback strength $c$: Controls the magnitude of correction toward EMA.
• EMA power $K$: Controls the responsiveness of EMA to recent iterates (recovering standard EMA or uniform averaging).
• Update frequency $uf$: Allows updates to be performed less frequently, as in typical EMA practice.

The per-coordinate gain leverages Adam's second-moment estimator as a Hessian-diagonal proxy. PACE's additional memory overhead (a separate copy of model weights) is noted as a practical concern at scale.

Empirical Evaluation

The paper extensively evaluates PACE in both supervised fine-tuning (SmolLM2-1.7B, Qwen3-1.7B, Gemma3-1B) and GPT-2 pretraining (FineWeb, Chinchilla-optimal token budget):

• Main empirical findings: Across all tested models and tasks, PACE strictly improves validation cross-entropy over AdamW and EMA baselines for a wide range of learning rates, decay schedules, and hyperparameters.
• The improvement is robust to choice of pullback strength $c$, EMA power $K$, and update frequency, consistently outperforming or matching finely tuned AdamW/EMA settings.
• Schedule-Free comparison: PACE is competitive with Defazio et al.'s Schedule-Free optimizer, which also removes the need for learning rate decay via averaging. However, PACE achieves improved theoretical convergence in the quadratic regime, while Schedule-Free trains slower and requires fewer additional copies of model weights.
• Learning rate decay: PACE at a constant learning rate matches or exceeds AdamW/EMA baselines with linear/cosine decay, suggesting decay schedules are unnecessary for achieving state-of-the-art performance with PACE.

Implications and Future Directions

PACE demonstrates the nontrivial benefit of designing optimization algorithms for the returned average, rather than the final iterate. Its practical utility is established across open-source LMs, and its theoretical guarantees are strong in both convex and quadratic settings.

Implications:

• Optimization paradigm shift: Designing optimizers for weight-averaged estimators rather than final iterates better aligns with LM deployment practice.
• Training stability and performance gains: PACE allows aggressive learning rates without sacrificing convergence, potentially reducing training times and compute budgets.
• Generalization to second-order optimizers and reduction of memory overhead: Future work may extend PACE to second-order Adam variants (e.g., SOAP, MuON) or replace momentum entirely with controlled pullback, reducing memory costs.
• Token-budget flexibility: With PACE, post-decay performance is continuously available at every checkpoint, enabling more flexible training horizons and model selection.

The theoretical foundation provided by the optimal-control formulation opens avenues for customized trajectory interventions targeting specific statistical properties of returned models. Further research should rigorously analyze PACE with momentum/adaptive schemes and evaluate its effectiveness at larger scales (>2B parameters) and on diverse architectures.

Conclusion

This paper introduces and analyzes PACE, a principled and practical algorithm for optimizing iterate-averaged estimators in LLM training. By explicitly controlling the training trajectory, PACE yields robust improvements in both convergence guarantees and empirical performance over AdamW and EMA across a variety of model families, training regimes, and hyperparameter settings. The results encourage further exploration into optimizers tailored for weight-averaged returns and lay the groundwork for optimization/control-theoretic advances in deep learning.