Quadratic-accurate merging improves reconstructing training checkpoints
QAM: Quadratic-Accurate Checkpoint Merging via Sequential Consistency
Machine LearningArtificial Intelligence
Summary
Training checkpoints save snapshots of a model during learning, but they often can't fully recreate what would happen if training followed a different path. The authors analyze how well checkpoints can recover the final training state if updates changed slightly. They prove a limit on accuracy achievable using only fixed recent history and introduce Quadratic-Accurate Merging (QAM), which reaches this optimal accuracy under certain conditions. They test QAM on large language model training data and find it works better than some existing methods when using longer checkpoint windows.
What this means in practice
- •For machine learning engineers: Improve accuracy when reconstructing model states from training checkpoints with Quadratic-Accurate Merging to better analyze or resume training.
- •For hyperparameter tuning teams: Use quadratic-accurate merges to better estimate effects of different training step sizes from existing checkpoints without retraining from scratch.
Authors
Shihao Wang, Rui Kong, Xinran Chen, Hui Wu, Qipeng Qian, Jinman Zhao, Jiashu Zhao, Yuchen Li, Jimmy Huang, Dawei Yin
Abstract
Saved checkpoints record states along a training trajectory, but generally do not determine the updates at states that would be visited under a different schedule. We study how accurately these checkpoints can reconstruct the endpoint of a sequential reference with prescribed update strengths. Under a common local transition model, two checkpoint-index moment conditions characterize all convex merges that agree with this reference through second order. We then prove an information limit that for nondegenerate profiles, no algorithm using only a fixed-length gradient-descent (GD) history with step size $h$ can achieve $o(h^3)$ endpoint error uniformly over a fixed class of smooth, strongly convex losses. The lower bound follows from two losses with identical GD checkpoint histories but sequential reference endpoints separated by $Ω(h^3)$. \textbf{Quadratic-Accurate Merging} (QAM) achieves a matching uniform $O(h^3)$ endpoint error bound. Its explicit coefficients also define the unique profile-dependent merge that exactly matches the sequential GD reference across all fixed quadratic objectives. Across two public Adam checkpoint trajectories (SmolLM3-3B and OpenEuroLLM-Prelude-9B), three windows and three profiles per model, and 15 tasks, QAM shows mixed results for short windows and broader advantages over \textbf{Warmup-Stable and Merge} (WSM) for longer windows. Matched-moment GSM8K diagnostics further show that local consistency alone does not fully determine downstream scores. These results characterize the reconstruction limits of saved histories, provide a coefficient rule that attains the optimal rate, and assess its practical utility.