Technical Summary: Training for the Model You Return

Problem Statement

Modern Language Model (LM) pipelines frequently return an averaged model—typically an Exponential Moving Average (EMA) of training iterates—rather than the final iterate itself. While iterate averaging is known to stabilize optimization and improve generalization, a fundamental question remains: Given that the final output will be an iterate average, how should the training process itself be modified to optimize the performance of this specific estimator?

Prior work has implicitly addressed this by observing that averaging allows for more aggressive learning rates. However, this paper takes a direct approach: it seeks to design an optimization algorithm that explicitly controls the training trajectory to minimize the error of the returned iterate-average estimator.

Methodology

Theoretical Formulation

The authors frame the problem as an optimal-control problem within a continuous-time stochastic quadratic model.

• Dynamics: The parameter trajectory $\theta_t$ follows a stochastic differential equation (SDE) representing noisy gradient descent near a quadratic minimum (Ornstein–Uhlenbeck process), augmented by a control input $u_t$.
• Objective: The goal is to find a control strategy $u_t$ that minimizes the mean squared error between the returned time-averaged estimator $\hat{\theta}_{EMA, T}$ and the true optimum $\mu^*$, subject to a quadratic penalty on the magnitude of the control intervention (regularization parameter $\lambda$).
• Derivation: Using Linear-Quadratic-Gaussian (LQG) control theory and the principle of certainty equivalence, the authors derive an optimal controller. The exact solution depends on the posterior mean of the optimum given the trajectory, which is intractable for large-scale LMs.

Practical Approximation: PACE

To make the solution implementable, the authors introduce a series of heuristic approximations:

1. Estimator Substitution: Replace the unknown posterior mean of the optimum with the current iterate average ($\hat{\theta}_{EMA, t}$).
2. Time Approximation: Assume the current time $t$ is much smaller than the total training time $T$ ($t \ll T$), simplifying the time-dependent coefficients.
3. Discretization and Preconditioning: Discretize the continuous control and replace the unknown Hessian matrix with the diagonal preconditioner used in AdamW (the second-moment estimate $v_t$).

This yields PACE (Pullback Averaging Control for Efficient Optimization), a lightweight wrapper around AdamW. The update rule adds a "pullback" term that moves the live weights toward the EMA weights:
$$\theta_{k+1} = \theta_k - \eta g_k + C(\hat{\theta}_{EMA, k} - \theta_k)$$
where $C$ is a diagonal matrix derived from the Adam preconditioner, and the pullback strength is clipped to ensure the update remains a convex combination of the standard step and the pullback.

Key Contributions

1. Optimal Control Formulation: The paper provides the first principled derivation of an optimization algorithm specifically designed to minimize the error of an iterate-average estimator, rather than the final iterate.
2. Theoretical Guarantees:
   • Convex Setting: The authors prove that a stylized version of PACE converges to the minimum of any convex loss with unbiased stochastic gradients at the standard $O(1/\sqrt{T})$ rate, up to a constant factor dependent on the EMA strength. This ensures the method is never significantly worse than standard SGD.
   • Quadratic Setting: In the quadratic setting used for derivation, the controlled dynamics strictly improve the limiting squared error of the iterate-average estimator compared to uncontrolled dynamics (standard SGD with EMA). The improvement can be arbitrarily large on certain quadratic instances.
3. Empirical Validation: The method is evaluated on:
   • Fine-tuning: SmolLM2-1.7B, Qwen3-1.7B, and Gemma3-1B models on the smol-smoltalk dataset.
   • Pretraining: GPT-2 (124M) on FineWeb at the Chinchilla-optimal token budget.
   • Baselines: PACE is compared against vanilla AdamW, AdamW with EMA evaluation, and the Schedule-Free algorithm.

Results

• Performance: PACE strictly improves upon both vanilla AdamW and AdamW with EMA evaluation across all tested models (1B–2B parameters) and datasets.
• Robustness: The improvement is robust across a wide range of learning rates, decay schedules (constant, cosine, Warmup-Stable-Decay), and hyperparameters (pullback strength $c$, EMA power $\kappa$).
• Comparison to Schedule-Free: PACE is competitive with the Schedule-Free algorithm (Defazio et al., 2024), which also uses iterate averaging to remove the need for learning rate decay. However, PACE achieves this with a different update rule.
• Learning Rate Decay: PACE performs well even with constant learning rates, suggesting it can stabilize training in regimes where learning rate decay is typically required. In token-budget comparisons, PACE's constant-LR trajectory tracks the "river floor" of the loss landscape as effectively as WSD branches, often outperforming them.

Significance and Limitations

Significance:
The paper demonstrates that optimizing for the returned model (the average) rather than the final iterate leads to tangible performance gains. By treating the training trajectory as a control problem, PACE provides a theoretically grounded method to stabilize training and improve generalization without requiring complex schedule tuning. The authors note that PACE offers improved theoretical convergence in the quadratic setting compared to standard methods, though this comes with a slightly slower theoretical guarantee in the general convex setting due to the EMA factor.

Limitations:

• Theoretical Assumptions: The derivation relies on a simplified quadratic model with additive noise. While convergence is proven for general convex functions, the rigorous incorporation of momentum and adaptive preconditioning (central to AdamW) into the theoretical analysis remains an open direction.
• Scale: Experiments were limited to models up to 2B parameters for fine-tuning and 124M for pretraining. The efficacy at larger scales (e.g., 70B+) is unverified.
• Memory Overhead: PACE requires maintaining an additional copy of model weights (the EMA state) during training, which introduces non-trivial memory overhead, although this is mitigated by large batch sizes in practice.
• Seed Sensitivity: Due to compute constraints, headline runs used a single seed, though multi-seed robustness studies confirmed the method's ordering stability.

The authors conclude that PACE represents a promising step toward aligning optimization algorithms with the actual deployment strategy of modern LMs, where the averaged model is the final product.