Trajectory soup improves large language model training efficiency and results
Trajectory Soup: Pushing the Compute-Scaling Frontier of LLM Mid-training via Diverse Trajectories
Machine LearningArtificial IntelligenceComputation and Language
Summary
Training large language models for special reasoning tasks can hit a limit where more training doesn’t help and might even harm the models. The authors found that splitting the training into several different paths from the same starting point produces varied models, which can then be merged to create a stronger model overall. Their method, Trajectory Soup, averages the best models from these different training paths to improve performance beyond what continuing a single training path can achieve. This approach makes better use of training effort and works across different model sizes and training setups.
What this means in practice
- •For machine learning engineers: Optimize large language model training by splitting compute into multiple training paths and merging best checkpoints to enhance model performance.
- •For cloud ai service providers: Increase efficiency and output quality of language model mid-training by using trajectory soup methods to balance compute across diverse checkpoints.$Commercial implications: Enables offering improved AI language models with less training cost, creating a better cloud-based language model service.
Authors
Zhehao Huang, Changxin Tian, Qingyuan Yang, Kunlong Chen, Ziqi Liu, Zhiqiang Zhang, Xiaolin Huang, Jun Zhou
Abstract
Mid-training equips pretrained large language models with specialized and reasoning capabilities, but the returns of this stage are bounded since additional serial compute yields little further downstream improvement and can even degrade some capabilities, which places a practical ceiling on how much compute mid-training absorbs. We revisit how this compute should be allocated to a single run or multiple similar optimizations. We find that branches forked from a shared checkpoint under various controlled recipe reaches measurably different regions of parameter space, and establish a form of compatible diversity that extending one run cannot supply. Therefore, we introduce Trajectory Soup, which distributes a mid-training budget over several independent branches, and consolidates strongest checkpoints selected on validation through intra- and inter-trajectory averaging into a single model. A local bias and variance analysis separates the two averaging levels, showing that inter-trajectory averaging removes residual error beyond the reach of averaging within a trajectory, while checkpoint selection carries a bias that bounds how many checkpoints are worth merging. Across model scales, learning-rate schedules, token budgets, and trajectory counts, Trajectory Soup improves aggregate downstream performance over the strongest single-trajectory average under matched budgets and keeps improving as budgets expand, with the advantage preserved after an identical post-training pipeline. These results position trajectory allocation and merging as a practical way to extend the compute-scaling frontier of mid-training beyond serial saturation.