The paper presents a comprehensive paper on incorporating model merging techniques during the pre-training process of LLMs. Its key innovation is the introduction of Pre-trained Model Averaging (PMA), a strategy that merges several intermediate checkpoints from a single pre-training trajectory to improve model performance, accelerate training, and enhance overall training stability.
Main Contributions:
- Pre-trained Model Averaging (PMA):
The authors introduce PMA as a framework to merge checkpoints obtained during the stable training phase of pre-training. By averaging weights from models trained with a constant learning rate (without or before the cosine decay phase), the merged model can simulate the effects of annealing, yield performance comparable to annealed counterparts, and reduce computational costs.
- Merging Strategies and Hyperparameter Analysis:
- Simple Moving Average (SMA): All checkpoints are weighted equally.
- Exponential Moving Average (EMA): Recent checkpoints are emphasized using exponentially decaying weights.
- Weighted Moving Average (WMA): Checkpoints are assigned linearly increasing weights to prioritize later training stages.
- Experiments show that while all methods improve performance, WMA tends to deliver the best results in the early training stages; however, as training stabilizes, the differences become minimal. Ablation studies also determine optimal intervals (token consumption gaps) and the number of merged checkpoints across different model sizes.
- Downstream Benefits and Training Stability:
Beyond improving pre-training performance, the paper shows that using merged weights to initialize downstream stages—continual training (CT) and supervised fine-tuning (SFT)—can lead to smoother gradient norm trajectories and stabilize training overall. In fact, the PMA-init method is demonstrated as a recovery strategy when models encounter loss spikes or diverging behavior, thus preventing retraining from scratch.
A detailed mathematical analysis based on a second-order Taylor expansion of the loss function illustrates why merging can lead to lower loss. The paper shows that if the deviations from the optimal parameter set have complementary (negatively correlated) behaviors with respect to the Hessian of the loss function, then averaging these checkpoints can yield a model closer to the optimum.
- Extensive Empirical Validation:
The authors conducted experiments on both Mixture-of-Experts (MoE) and dense model architectures, with models ranging from hundreds of millions to over 100 billion parameters. Results across numerous open benchmarks (including MMLU, HumanEval, GSM8K, and more) consistently demonstrate that model merging improves performance. In addition, experiments comparing models trained with constant learning rates versus cosine annealing reveal that merging with constant rates can match or even surpass naturally annealed models.
Experimental Findings:
- Pre-training Improvements:
Merging during the stable phase consistently improves downstream task performance—even for very large models. In some cases, early merged models match the performance of final-stage annealed models, allowing for more efficient validation cycles.
Ablation studies confirm that the merging interval should scale with model size (e.g., smaller intervals for smaller models and larger ones for larger models) and that incorporating more checkpoints (up to a point) improves performance once training converges.
- Downstream Initialization (PMA-init):
Using PMA-initialized weights for CT and SFT results in lower initial loss values and more stable gradient norms. In scenarios of catastrophic loss spikes, PMA-init has been effective in resuming and stabilizing the training process.
Visualization of weight distributions and contour plots of performance metrics (like MMLU scores) provides an intuitive explanation: individual checkpoints explore different regions of the parameter space, and their complementary nature can lead the averaged model to reside in a region with higher performance.
Overall Impact and Practical Implications:
The paper not only offers empirical evidence for the benefits of model merging during pre-training but also provides detailed guidelines for selecting merging strategies and hyperparameters. By using PMA, practitioners can potentially avoid the expensive cosine annealing phase, accelerate model development, and improve the reliability and stability of the training process. This work is particularly valuable for teams working on large-scale LLM pre-training, as it reduces computational costs while maintaining or even enhancing model performance.
In summary, the paper makes a significant contribution by bridging the gap between post-training model merging methods and those applicable during large-scale pre-training. It equips researchers and practitioners with both theoretical insights and practical guidelines to effectively implement model merging, ultimately enabling more efficient and robust development of LLMs.