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

2026-06-23

AI Summary

While existing language model training often adopts iterative averaging—such as exponential moving average (EMA)—as the final model, optimizers are still designed to target the last iterate rather than the averaged model’s performance. This work formulates optimizer design as an optimal control problem and, under a continuous-time stochastic quadratic model, derives a control policy that minimizes the error of the averaged model. Building on this insight, we propose PACE, a lightweight wrapper around AdamW that dynamically steers instantaneous weights toward their EMA via per-coordinate clipping of control strength. Theoretically, PACE is shown to significantly reduce the asymptotic mean squared error of the averaged estimator under certain conditions, with unbounded potential improvement. Empirically, PACE consistently outperforms AdamW and its EMA variants across diverse learning rates, decay schedules, and hyperparameter settings in both 1–2B parameter model fine-tuning and GPT-2 pretraining.

Abstract

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.

Problem

• iterate averaging
• language models
• optimization
• model averaging
• training dynamics

Innovation

• optimal control
• iterate averaging
• PACE
• stochastic optimization
• language model training