Model Merging in Pre-training of Large Language Models

本文提出预训练模型平均（PMA）策略，通过融合预训练阶段的检查点显著提升大型语言模型性能、预测退火效果并增强训练稳定性，为高效模型开发提供了新方法和实用指南。

Large Language Model, Pre-training, Efficiency, Representation Learning

Yunshui Li, Yiyuan Ma, Shen Yan, Chaoyi Zhang, Jing Liu, Jianqiao Lu, Ziwen Xu, Mengzhao Chen, Minrui Wang, Shiyi Zhan, Jin Ma, Xunhao Lai, Yao Luo, Xingyan Bin, Hongbin Ren, Mingji Han, Wenhao Hao, Bairen Yi, LingJun Liu, Bole Ma, Xiaoying Jia, Zhou Xun, Liang Xiang, Yonghui Wu

ByteDance Seed, Peking University, The University of Hong Kong

Generated by grok-3

Background Problem

大型语言模型（LLMs）的预训练成本高昂，性能扩展难以预测，且大规模训练存在不稳定性。模型融合（Model Merging）作为一种新兴技术，在后训练阶段已被证明能有效整合不同任务模型的能力，但在预训练阶段的研究仍较少。本文聚焦于预训练阶段的模型融合，旨在通过融合训练过程中的检查点，解决预训练成本高、退火性能预测难以及训练不稳定等问题，为更高效的模型开发提供新路径。

Method

本文提出了一种名为预训练模型平均（PMA）的策略，用于在预训练阶段进行模型权重融合。其核心思想是通过对训练轨迹中的多个检查点进行加权平均，生成一个性能更优的模型。具体方法包括：

融合对象：选择预训练过程中同一训练轨迹上的多个检查点（Checkpoint），通常来自稳定训练阶段或退火阶段。
融合策略：采用简单移动平均（SMA）、加权移动平均（WMA）和指数移动平均（EMA）等方法对检查点权重进行加权计算，生成融合模型 $M_{\text{avg}} = \sum_{i=1}^{N} w_i M_i$ ，其中 $w_i$ 为权重， $M_i$ 为第 $i$ 个检查点的参数。
关键超参数：包括融合间隔（Interval, $V$ ）和融合模型数量（ $N$ ），通过实验确定最优配置。
创新应用：提出PMA-init，即利用融合后的权重作为后续持续训练（CT）或监督微调（SFT）的初始化权重，以提升训练稳定性和下游性能。

批判性思考：虽然PMA方法在概念上简单直观，但其理论依据（基于损失函数的二阶泰勒展开和Hessian矩阵分析）显得有些牵强，未能充分解释为何融合能带来性能提升。此外，不同融合策略的效果差异在后期趋于消失，这可能暗示方法对训练后期权重变化的敏感性不足，限制了其适用性。

Experiment

实验在Dense模型（参数规模从411M到70B）和Mixture-of-Experts（MoE）模型（激活参数从0.7B到20B，总参数至200B）上进行，使用内部预训练语料库（万亿级别token）和开源基准测试（如MMLU、GSM8K、HumanEval等）评估性能。实验设计围绕六个关键问题展开，包括融合对性能的影响、不同融合策略的效果、最优超参数选择、下游训练影响、训练稳定性以及融合机制分析。

性能影响：在稳定训练阶段，融合模型在多个下游任务上表现出显著提升，例如Seed-MoE-1.3B/13B在HumanEval上从31.1提升至36.6；在退火阶段早期，PMA效果甚至可媲美或超越最终退火模型。
融合策略比较：初期WMA表现最佳，但随着训练进展，SMA、EMA和WMA的性能差异逐渐消失，显示出方法对后期权重变化的适应性有限。
超参数优化：最优融合间隔随模型规模呈比例关系（如1.3B/13B模型为8B token，10B/100B模型为80B token），融合更多检查点（ $N$ 较大）在训练完成后带来更大提升。
下游训练与稳定性：PMA-init在CT阶段初期降低损失并略提升性能，但在SFT阶段效果不一致；对于训练稳定性，PMA-init在损失尖峰后恢复训练时表现出较好的稳定性，避免了从头训练的资源浪费。
机制分析：通过数学推导和权重分布可视化，作者认为融合效果源于检查点在损失景观中的互补性探索，但这一解释缺乏更直接的实验验证。

批判性思考：实验设置覆盖了多种模型规模和架构，显示出一定的全面性，但结果的显著性存疑，例如部分性能提升（如HumanEval上的5.5点）是否具有统计学意义未明确。此外，实验未充分探讨学习率对融合效果的影响，这可能是关键变量。PMA-init在下游任务中的不一致表现也表明其实际应用价值可能有限，需更多实验验证其普适性。

Further Thoughts

本文提出的PMA策略为预训练阶段的模型融合提供了一个有趣的视角，但其理论深度和实验验证仍有不足。未来研究可以探索学习率对融合效果的具体影响，尤其是在高学习率场景下是否能进一步提升性能，这与文献中提到的早期权重平均与高学习率的结合可能存在协同效应。此外，PMA在训练稳定性方面的应用（PMA-init）启发了一个新方向：是否可以将模型融合与其他训练动态优化技术（如梯度裁剪或学习率调度）结合，以应对大规模训练中的不稳定性？另一个值得思考的点是，模型融合是否能在强化学习（RL）阶段发挥类似作用，尤其是在RL训练周期较长、权重变化显著的场景中，这可能是连接预训练与后训练融合研究的一个桥梁。最后，考虑到PMA对后期权重变化的敏感性不足，是否可以通过引入自适应权重分配机制（如基于损失景观曲率的动态权重）来进一步优化融合效果？这些方向值得深入探索。