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

• URL: http://arxiv.org/abs/2606.25086v2
• Date: Tue, 30 Jun 2026 16:49:44 GMT
• Title: Training for the Model You Return: Improving Optimization for Iterate-Averaged Language Models
• Abstract summary: Many modern Language Model (LM) pipelines return an averaged model, such as an exponential moving average of the training parameters.<n>This raises a fundamental question: how should we change training to improve the performance of this average?<n>We study this question by formulating a design for the iterate-average estimator as an optimal-control problem.
• Score: 10.749836992585363
• License: http://creativecommons.org/licenses/by/4.0/
• 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.

Related papers

• Zero-Order Optimization for LLM Fine-Tuning via Learnable Direction Sampling [40.94400211806987]
  We propose a policy-driven ZO framework that treats the sampling distribution over perturbation directions as a learnable policy.<n>We show that learned sampling improves quality gradient information and relax the explicit dependence on $d$ in convergence bounds.<n>Our results suggest that adaptive direction sampling is a promising route to make ZO fine-tuning viable at scale.
  arXiv  Detail & Related papers  (2026-02-14T08:01:41Z)
• Parameter Averaging in Link Prediction [7.914373583039709]
  We introduce model merging, specifically weighted averaging, in knowledge graph embedding (KGE) models.<n>Running average of model parameters from a training epoch onward is maintained and used for predictions.<n>We evaluate these two different weighted averaging approaches on link prediction tasks, comparing the state-of-the-art benchmark ensemble approach.
  arXiv  Detail & Related papers  (2025-10-29T10:32:39Z)
• PADAM: Parallel averaged Adam reduces the error for stochastic optimization in scientific machine learning [5.052293146674794]
  Averaging techniques such as Ruppert--Polyak averaging and exponential movering averaging (EMA) are powerful approaches to accelerate optimization procedures of gradient descent (SGD) optimization methods such as the popular ADAM.<n>In this work we propose an averaging approach, which we refer to as parallel averaged ADAM (PADAM) in which we compute parallely different averaged variants of ADAM and during the training process dynamically select the gradients with the smallest optimization error.
  arXiv  Detail & Related papers  (2025-05-28T08:07:34Z)
• Lean and Mean: Decoupled Value Policy Optimization with Global Value Guidance [52.65461207786633]
  Policy-based Reinforcement Learning from Human Feedback is essential for aligning large language models with human preferences.<n>It requires joint training of an actor and critic with a pretrained, fixed reward model for guidance.<n>We propose textbfDecoupled Value Policy Optimization (DVPO), a lean framework that replaces traditional reward modeling with a pretrained emphglobal value model (GVM)
  arXiv  Detail & Related papers  (2025-02-24T08:11:33Z)
• Hyperspherical Normalization for Scalable Deep Reinforcement Learning [57.016639036237315]
  SimbaV2 is a novel reinforcement learning architecture designed to stabilize optimization.<n>It scales up effectively with larger models and greater compute, achieving state-of-the-art performance on 57 continuous control tasks.
  arXiv  Detail & Related papers  (2025-02-21T08:17:24Z)
• MARS: Unleashing the Power of Variance Reduction for Training Large Models [56.67982828148859]
  We propose a unified training framework for deep neural networks.<n>We introduce three instances of MARS that leverage preconditioned gradient optimization.<n>Results indicate that the implementation of MARS consistently outperforms Adam.
  arXiv  Detail & Related papers  (2024-11-15T18:57:39Z)
• Truncating Trajectories in Monte Carlo Policy Evaluation: an Adaptive Approach [51.76826149868971]
  Policy evaluation via Monte Carlo simulation is at the core of many MC Reinforcement Learning (RL) algorithms. We propose as a quality index a surrogate of the mean squared error of a return estimator that uses trajectories of different lengths. We present an adaptive algorithm called Robust and Iterative Data collection strategy Optimization (RIDO)
  arXiv  Detail & Related papers  (2024-10-17T11:47:56Z)
• Two-Tailed Averaging: Anytime, Adaptive, Once-in-a-While Optimal Weight Averaging for Better Generalization [0.0]
  Tail Averaging improves on Polyak averaging's non-asymptotic behaviour. We propose an anytime variant of Tail Averaging intended for improving generalization not pure optimization.
  arXiv  Detail & Related papers  (2022-09-26T10:46:37Z)
• Deep Equilibrium Optical Flow Estimation [80.80992684796566]
  Recent state-of-the-art (SOTA) optical flow models use finite-step recurrent update operations to emulate traditional algorithms. These RNNs impose large computation and memory overheads, and are not directly trained to model such stable estimation. We propose deep equilibrium (DEQ) flow estimators, an approach that directly solves for the flow as the infinite-level fixed point of an implicit layer.
  arXiv  Detail & Related papers  (2022-04-18T17:53:44Z)
• MaxVA: Fast Adaptation of Step Sizes by Maximizing Observed Variance of Gradients [112.00379151834242]
  We propose adaptive learning rate principle, in which the running mean of squared gradient in Adam is replaced by a weighted mean, with weights chosen to maximize the estimated variance each coordinate. This results in faster adaptation, which leads more desirable empirical convergence behaviors.
  arXiv  Detail & Related papers  (2020-06-21T21:47:43Z)
• Extrapolation for Large-batch Training in Deep Learning [72.61259487233214]
  We show that a host of variations can be covered in a unified framework that we propose. We prove the convergence of this novel scheme and rigorously evaluate its empirical performance on ResNet, LSTM, and Transformer.
  arXiv  Detail & Related papers  (2020-06-10T08:22:41Z)
• Nested Model Averaging on Solution Path for High-dimensional Linear Regression [12.173071615025504]
  We study the nested model averaging method on the solution path for a high-dimensional linear regression problem. We propose to combine model averaging with regularized estimators (e.g., lasso and SLOPE) on the solution path for high-dimensional linear regression. A real data analysis on predicting the per capita violent crime in the United States shows an outstanding performance of the nested model averaging with lasso.
  arXiv  Detail & Related papers  (2020-05-16T18:09:57Z)

This list is automatically generated from the titles and abstracts of the papers in this site.

This site does not guarantee the quality of this site (including all information) and is not responsible for any consequences.