字节跳动Seed团队提出PMA:通过模型合并提升大模型预训练效率

字节跳动提出PMA技术,通过合并预训练过程中的检查点,显著提升大模型性能并优化训练流程,或将改变大模型预训练范式。

原文标题:字节Seed新作:模型合并如何改变大模型预训练范式

原文作者:机器之心

冷月清谈:

字节跳动Seed团队提出了预训练模型平均(PMA)技术,通过合并训练过程中的检查点,显著提升模型性能,并能预测学习率衰减阶段的性能表现。该技术的核心发现包括:在学习率稳定阶段合并效果最佳;简单平均(SMA)策略优于复杂加权;合并间隔与模型规模正相关。PMA还具备训练稳定性与初始化优化的能力,例如在损失激增时,合并故障前的检查点作为初始化权重,可使训练恢复稳定。该研究从理论层面解释了模型合并的有效性,并通过实验验证了其可行性。PMA的提出为大模型开发者提供了一种低成本、高效能的新路径,有望改变未来大模型开发的范式。

怜星夜思:

1、PMA在实际应用中,如何平衡合并检查点的数量和计算成本?有没有更智能的策略来选择哪些检查点进行合并?
2、文章提到PMA可以作为大模型训练的“急救方案”,那么除了损失激增外,PMA还能应对哪些其他的训练问题?
3、PMA主要在预训练阶段进行模型合并,那么在下游任务的微调阶段,PMA是否还有应用价值?例如,能否通过合并微调不同领域的模型来提升泛化能力?

原文内容


字节跳动 Seed 团队近期在 arXiv 上发表的论文得到了 ViT 作者,前 Google Brain 最近跳去 OpenAI 的 Lucas Beyer 的亲自解读,Lucas 直言:「这是一篇简洁的论文,不知怎的让我回忆起美好的在 Google Brain 的旧时光。(This is a neat paper that somehow made me reminisce good old Brain times. )」



Seed 团队在这篇论文提出的预训练模型平均(PMA)技术,通过合并训练过程中的检查点(Checkpoint),不仅实现了模型性能的显著提升,还能精准预测学习率衰减阶段的性能表现。这一成果被视为大模型训练领域的重要突破,甚至可能改变未来大模型开发的范式。




  • 论文标题:Model Merging in Pre-training of Large Language Models

  • 论文地址:https://arxiv.org/pdf/2505.12082


模型合并:从「后训练」到「预训练」的跨越


后训练合并:任务能力的「拼图游戏」

模型合并并非全新概念,此前主要应用于后训练阶段,即通过合并多个领域微调模型的权重,构建一个多任务能力更强的统一模型。例如,DARE 方法将 WizardLM(通用对话模型)与 WizardMath(数学推理模型)合并后,在 GSM8K 数学推理基准上的得分从 2.2 跃升至 66.3,展现了任务能力融合的强大潜力。


相比之下,预训练阶段的模型合并研究仍较为匮乏。此类预训练合并通常涉及合并单一训练轨迹中的检查点,如 LAWA 中通过模型合并加速 LLM 训练的探索。然而,随着模型规模和数据量的急剧增长,社区研究者难以评估模型合并对大规模模型的影响,主要原因在于难以获取大规模预训练过程中的中间检查点。尽管 DeepSeek 和 LLaMA 均表明其在模型开发中使用了模型合并技术,但这些技术的详细信息尚未公开披露。


预训练合并:训练效率的「时光机」

字节跳动的研究将模型合并引入预训练阶段,提出了Pre-trained Model Averaging(PMA)框架。简单来说,PMA 就是在预训练过程中,定期将不同训练阶段的模型权重进行平均,生成一个「合并模型」。这是因为:预训练后期的模型权重往往在参数空间中探索了不同的局部最优解,通过平均化可以抵消单个模型的偏差,逼近更优的全局解例如,在稳定训练阶段(Constant LR Phase)合并 10 个检查点后,Seed-MoE-10B/100B 模型在 HumanEval 代码生成任务上的得分从 54.3 提升至 61.6,涨幅超过 13%。


PMA 技术的三大核心发现


合并时机:稳定期合并效果最佳

研究团队通过实验发现,在学习率稳定阶段(Warmup-Stable-Decay 中的 Stable Phase)进行模型合并效果最佳此时模型处于「高效学习期」,权重更新尚未进入衰减阶段,不同检查点之间的参数差异既能保证多样性,又不会因过度震荡导致合并后性能下降。 



有趣的是,即使在学习率余弦衰减阶段(Cosine Decay Phase)的早期进行合并,PMA 模型的性能也能媲美甚至超越自然衰减到末期的模型。例如,Seed-MoE-15B/150B 模型在衰减初期合并后,其性能与训练至末期的模型相差无几。



合并策略:简单平均(SMA)胜过复杂加权

在合并策略的对比实验中,研究团队测试了三种主流方法:


  • 简单移动平均(SMA):所有模型权重等比例平均
  • 指数移动平均(EMA):近期模型权重占比更高
  • 加权移动平均(WMA):按训练步数线性加权


实验结果表明,在训练初期,EMA 和 WMA 因更关注近期权重而表现略好,但随着训练推进,三者性能差异逐渐消失考虑到 SMA 的计算简单性和稳定性,团队最终选择其作为默认策略。这一发现打破了「复杂加权必然更优」的固有认知,为工程落地提供了便利。


超参数规律:模型规模决定合并间隔

合并间隔(V)与模型规模正相关:小模型(如 1.3B 参数的 MoE)适合较小的合并间隔(8B tokens),而大模型(如 100B 参数的 MoE)则可采用更大的间隔(80B tokens)。这与大模型通常使用更大批次训练的特性一致。

合并数量(N)越多越好:当训练完成时,合并 15 个检查点的模型性能比合并 3 个的高近 1 个百分点。但需平衡计算成本,团队建议实际应用中取 N=10 作为折中方案。


PMA 的「隐藏技能」:训练稳定性与初始化优化


PMA-init:让训练「起死回生」

在大模型训练中,「损失激增」(Loss Spike)是令人头疼的问题——硬件故障、参数震荡等因素可能导致训练崩溃,不得不从头再来。PMA 为此提供了一种「急救方案」:当损失激增发生时,合并故障前的 N 个检查点作为初始化权重(PMA-init),可使训练恢复稳定


实验中,团队故意用过高的学习率(6e-3)训练一个 330M 参数的 MoE 模型,导致其损失剧烈震荡。此时采用 PMA-init 合并 3 个故障前检查点,训练曲线迅速恢复平滑,避免了从头训练的巨大浪费。



下游阶段的「热身优势」

在持续训练(CT)和监督微调(SFT)阶段,使用 PMA 合并后的模型作为初始化权重(PMA-init),能显著改善训练动态。例如,在 CT 阶段,PMA-init 模型的 GradNorm 曲线更加平稳,早期训练中的 MMLU 得分比基线模型高 1-2 个百分点。尽管最终性能与基线持平,但其「热身优势」可加速下游任务的收敛,尤其适合数据敏感型场景。



数学原理:为什么合并能「化平凡为神奇」?


从理论层面看,模型合并的有效性可通过损失函数的二阶泰勒展开解释。假设最优参数为 图片,各检查点参数图片与 图片 的偏差为 图片,则合并后参数 图片 的损失可表示为:




其中,图片为海森矩阵(Hessian Matrix),刻画损失函数的曲率。当不同 图片 在参数空间中呈现「负相关」(即方向互补)时,交叉项 图片 为负,使得合并后的损失低于单个模型的平均损失。这意味着,合并本质上是利用不同检查点在参数空间中的「探索多样性」,通过平均化抵消局部偏差,逼近更优解


可视化实验也印证了这一点:在 Seed-MoE-1.3B/13B 模型的某层参数空间中,单个检查点的权重分布在 MMLU 得分等高线的不同位置,而合并后的权重位置往往更靠近高分区。



挑战与未来方向


未解决的问题

  • 学习率的影响:当前实验默认使用缩放定律(Scaling Law)推荐的最优学习率,未深入探索高学习率下 PMA 的表现。理论上,高学习率可能增加参数探索的多样性,进一步提升合并效果,但受限于算力成本,尚未量化分析。

  • 强化学习阶段的应用:论文主要聚焦预训练,而 RLHF(强化学习从人类反馈中学习)作为大模型训练的关键环节,其检查点合并的潜力尚未挖掘。这将是未来研究的重要方向。

行业启示

对于大模型开发者而言,PMA 带来的不仅是成本节省,更是一种「模拟退火」的思维革命——通过合并稳定期的检查点,可快速预测衰减阶段的性能,避免盲目延长训练周期。对于中小型企业,这意味着用更少的资源实现 comparable 性能,甚至可能颠覆「大公司垄断算力」的格局。


结语:开启高效训练的新时代


从「暴力堆算力」到「智能优化训练流程」,大模型的发展正从粗放式增长转向精细化运营。字节跳动的这项研究,以模型合并为切入点,揭示了预训练过程中被忽视的「检查点价值」,为学术界和工业界提供了一条低成本、高效能的新路径。


正如论文结语所言:「PMA 不仅是一种技术,更是一个监视器——它让预训练过程变得可预测、可优化。」随着更多类似研究的涌现,我们有理由相信,大模型训练将逐步摆脱「烧钱游戏」的标签,走向更可持续、更普惠的未来。


© THE END 

转载请联系本公众号获得授权

投稿或寻求报道:[email protected]

PMA作为“急救方案”的潜力我觉得很大。除了应对损失激增,还可以考虑用于解决过拟合问题。如果在训练过程中发现模型开始过拟合,可以合并之前几个epoch的检查点,相当于做了一种模型平均,应该可以起到一定的正则化效果。此外,PMA还可以用于迁移学习,将预训练好的模型和目标任务的模型进行合并,加速目标任务的训练。

PMA的应用场景确实很广泛。不过,我更关心的是PMA的局限性。比如,如果模型本身设计有问题,或者数据集质量不高,PMA还能起作用吗?另外,PMA在多大程度上能替代传统的正则化方法,例如dropout、weight decay等?这些问题还需要进一步研究。

各位大佬说的都很有道理。不过,我从另一个角度来看,微调阶段的PMA可能更适合用于模型压缩。例如,可以先训练一个大模型,然后在不同的子集上进行微调,最后通过PMA将这些微调后的模型合并成一个小模型。这样既能保证模型的性能,又能降低模型的计算复杂度,一举两得。

关于“急救”,我想到一个场景:有时候训练中断是随机的,比如服务器突然宕机。如果有了PMA,我们就可以随时从之前的检查点恢复,并合并之前的模型,减少损失。此外,如果团队里不同的人训练了不同的模型,也可以用PMA把它们合并起来,集思广益,提高整体性能。这其实就是一种“模型民主化”的思想。

关于PMA合并检查点数量和计算成本平衡的问题,我觉得可以考虑一种动态合并策略。比如,开始时合并较少的检查点,然后逐渐增加,观察性能收益的变化,当收益递减到一定程度时就停止增加。此外,还可以结合一些指标,例如梯度范数、损失变化等,来判断检查点的质量,优先合并质量高的检查点。这样应该能更有效地利用计算资源。

我觉得在微调阶段,PMA同样大有可为。特别是对于那些需要处理多种任务的模型,例如多语言模型、多模态模型等,可以通过合并在不同任务上微调的模型,来提升模型的泛化能力。此外,还可以结合一些领域知识,设计更精细的合并策略,例如对特定领域的模型赋予更高的权重。我感觉这有点像“专家集成”的思想。

谢邀,我觉得这个问题很有意思。从理论上讲,合并的检查点越多,模型性能应该越好,但实际上会受到计算资源的限制。我的想法是,可以借鉴“重要性采样”的思想,对检查点进行评估,选择那些对模型性能提升贡献最大的检查点进行合并。具体怎么评估,可以结合一些模型诊断工具,例如查看不同检查点对特定任务的激活值差异等。另外,也可以考虑使用一些近似合并算法,例如低秩近似等,来降低计算复杂度。

楼上说的有道理,但这让我想到了一个哲学问题:是不是所有局部最优解的平均都能接近全局最优解?换句话说,如果检查点本身质量不高,合并再多可能也没啥用。所以,我觉得更重要的还是保证训练过程的稳定性,避免出现梯度爆炸、模式崩溃等问题。在此基础上,再考虑合并策略的优化。否则,就像在一堆烂苹果里挑好一点的,意义不大。

关于微调阶段的PMA,我想到一个实际问题:不同任务的数据分布可能差异很大,直接合并模型可能会导致性能下降。因此,在合并之前,需要对模型进行一些调整,例如通过对抗训练来对齐特征空间,或者使用一些领域自适应的方法。另外,还需要仔细选择合并的比例,避免某个任务的模型过度影响其他任务。