The preprint highlights a core mismatch in language model workflows: optimizers are built to chase the final training checkpoint, yet pipelines often ship an exponential moving average of the weights instead, leaving performance on the table when those two targets diverge.
New preprint: Training for the model you return
Modern LM pipelines often return an averaged model (e.g. EMA), rather than the final iterate, but optimizers are still mostly designed around the final iterate.
How should we change training if we know we will return an EMA? 🧵
Block and Zhang propose BEMA, a bias-corrected averaging scheme presented as a two-line swap that preserves variance reduction while cutting lag, backed by a model deriving its optimality plus experiments focused on small-batch fine-tuning stability.
No quantitative gains, benchmark lists, or pipeline adoption details are established yet, so whether the approach scales beyond the reported fine-tuning regime or influences production training remains an open question.
Users praise the preprint's proposed optimizer changes for EMA model returns as a clean, great idea worth investigating.
No Digg Deeper questions have been answered for this story yet.
The legendary Adam Block is now on twitter!!
New preprint: Training for the model you return
Modern LM pipelines often return an averaged model (e.g. EMA), rather than the final iterate, but optimizers are still mostly designed around the final iterate.
How should we change training if we know we will return an EMA? 🧵
Very cool work with connections to the schedule free work and iterative averaging work.
Interestingly, a natural step is to derive it for vanilla shampoo because it keeps the pre-conditioner around (b2>0) around. Essentially replace diagonal with the kronecker factors
Could be cool extend the per coordinate scheme to various other variance schemes.
New preprint: Training for the model you return
Modern LM pipelines often return an averaged model (e.g. EMA), rather than the final iterate, but optimizers are still mostly designed around the final iterate.
How should we change training if we know we will return an EMA? 🧵
@adam_block66985 Very nice idea!
@GeekParkHQ Great question! The answer is the former. On a related note, although we did observe that PACE can at times substitute for lr decay (e.g. WSD at different token budgets), our best runs still used some form of decay.
@_arohan_ Yes that is a great idea and we are actively investigating similar approaches!
(3/14) If we know we will return the average, why are we still optimizing as though we will return the final iterate?
PACE is a first attempt to take this question seriously.
(10/14) These gains persist across learning-rate schedules, learning rates, and EMA powers and the gains are robust across the pullback strengths.
(11/14) The same pattern appears in pretraining.
On GPT-2 124M trained on FineWeb at the Chinchilla-optimal token budget, PACE improves over AdamW and EMA under constant LR, WSD, and cosine decay schedules.
(8/14) In the quadratic setting PACE was designed for, the returned average can be strictly better than the uncontrolled average.
On some instances, it can be arbitrarily better.
(9/14) Empirically, the effect is simple to see in fine-tuning.
Across three 1–2B LMs—SmolLM2-1.7B, Qwen3-1.7B, and Gemma3-1B—PACE improves over AdamW and EMA-evaluated AdamW.
(6/14) After the AdamW step, PACE pulls the live weights toward the EMA weights with a clipped, per-coordinate gain.
In practice, this is a small modification to a standard training loop.
(1/14) Led by my amazing master’s student Kwok Chun Au, we derive a simple “pullback” intervention that improves training empirically and in theory.
(13/14) Takeaway: If you are going to return an averaged model, the optimizer should know that.
PACE uses the average not only at the end, but as a control signal during training.
Paper: https://arxiv.org/abs/2606.25086
(2/14) We already know EMA works.
Iterate averaging is used throughout deep learning, and many modern LM pipelines return an averaged model.
But optimizer design largely ignores this.
(12/14) PACE also has a broad basin over pullback strengths.
So while it introduces a new control parameter, performance does not depend on finding a single knife-edge value.
(5/14) The solution says: pull the live iterate toward an estimate of the optimum and toward consistency with the accumulated average.
This is the conceptual origin of PACE.
(4/14) Our starting point is an idealized control problem.
In a noisy quadratic model of optimization, we ask for the control that minimizes the error of the returned average, while penalizing how much we intervene.
(7/14) PACE is not just a heuristic from a toy model.
For convex losses, a stylized version gets the standard stochastic optimization rate, up to an EMA-dependent factor.
Welcome @AdamB1438 to Twitter! His new paper rethinks how to leverage iterate averaging for much, much faster optimization. Really worth the read.
New preprint: Training for the model you return
Modern LM pipelines often return an averaged model (e.g. EMA), rather than the final iterate, but optimizers are still mostly designed around the final iterate.
How should we change training if we know we will return an EMA? 🧵
@adam_block66985 Really clean writeup, thanks for walking through the derivation! Quick question here: Pullback toward the EMA is also an implicit per-coordinate LR decay. Is the gain the average-as-control signal, or a step-size effect a tuned schedule would recover?
