Large language model training improves by mixing final and averaged checkpoints
Terminal Shrinkage Averaging Reveals a Schedule-Estimator Interaction in LLM Pretraining
Machine LearningArtificial Intelligence
Summary
Training large language models usually involves using the very last version of the model, but this can mix up two decisions: how the learning rate changes during training and how to pick the final model to use. The authors suggest a new method called Terminal Shrinkage Averaging (TSA) that blends the last model with an average of recent ones, aiming to keep improvements while reducing model variation. They tested this using a small-scale language model and showed it helps both the training process and the final model quality, even speeding up some benchmarks.
What this means in practice
- •For machine learning engineers: Improve efficiency and quality of large language model training by selecting better ways to combine learning rate schedules and checkpoint averaging.
- •For ai infrastructure teams: Reduce training time for LLMs by implementing Terminal Shrinkage Averaging to accelerate convergence without sacrificing model quality.
Authors
Adam Ousherovitch, Yixin Wang
Abstract
Large language model (LLM) pretraining conventionally returns the raw final iterate. This couples two design choices: the learning-rate schedule that generates the parameter trajectory and the estimator that constructs the deployed model (e.g. the raw final iterate or a checkpoint average). A schedule that promotes optimization progress may differ from one that minimizes variation in the raw final iterate. Separating these choices creates an opportunity to maintain progress late in training while reducing variation in the returned model. To this end, we propose \emph{Terminal Shrinkage Averaging (TSA)}, which interpolates between the raw final iterate and the average of recent checkpoints to balance recent progress against terminal variation. We analyze how TSA changes the preferred terminal learning-rate schedule under a local quadratic approximation and test this interaction through a sequence of controlled NanoChat experiments. Finally, we demonstrate that the resulting gains transfer to depth-22 NanoChat, where the combined schedule and estimator improve validation quality. A qualifying time-to-GPT-2 run also finishes faster than the public baseline used in our experiments, providing preliminary evidence of benchmark acceleration.