Extra-Merge: 追踪语言模型预训练中模型合并的 Rank-1 子空间

# 追踪语言模型预训练中模型合并的秩一子空间 来源：https://arxiv.org/html/2605.26484 ###### 摘要 模型合并已成为增强大型语言模型（LLM）的一种轻量级范式，但其底层机制仍鲜为人知。在这项工作中，我们分析了后期预训练轨迹，并发现了一个*秩一子空间*现象：虽然原始的优化步骤剧烈震荡，但连续的*合并后*检查点却坍缩到一个稳定、近似一维的线性流形上。我们从理论层面将这一观察置于*河谷*景观分析中：平均操作充当了几何低通滤波器，抑制了高曲率噪声，从而揭示了最优下降方向。基于这一洞见，我们提出了 Extra-Merge，一种无需训练的策略，沿该子空间外推以最小化损失，无需额外梯度更新。在 GPT-2 和 LLaMA 系列（124M 到 2B）上的大量实验表明，Extra-Merge 持续优于标准合并基线。值得注意的是，它在 Pythia-12B 下游任务上产生了一致的零样本准确率提升，并有效泛化到 Muon 优化器（Jordan 等，2024 (https://arxiv.org/html/2605.26484#bib.bib28)）。机器学习，ICML ## 1 引言 预训练大型语言模型（LLM）以严格的优化调度和巨大的计算成本为特征（OpenAI 等，2024 (https://arxiv.org/html/2605.26484#bib.bib47)；Comaniciet 等，2025 (https://arxiv.org/html/2605.26484#bib.bib37)）。现代预训练运行消耗大量 GPU 小时（Kaplan 等，2020 (https://arxiv.org/html/2605.26484#bib.bib12)；Hoffmann 等，2022 (https://arxiv.org/html/2605.26484#bib.bib13)；hägele2024scalinglawscomputeoptimaltraining），其优化轨迹严格由学习率调度器决定——例如 Warmup-Stable-Decay（WSD）（Hu 等，2024 (https://arxiv.org/html/2605.26484#bib.bib15)）或余弦退火。因此，最终模型能力往往具有路径依赖性，仅在学习率衰减的最终阶段实现。鉴于这些限制，迫切需要能够利用优化轨迹中的潜在信息以最大化性能的*事后*机制，而无需承受额外梯度步骤的过高成本。 作为回应，模型合并已成为一种强大且无需训练的范式，以解决这一效率瓶颈。该方法植根于经典的 Polyak 平均（Polyak 和 Juditsky，1992 (https://arxiv.org/html/2605.26484#bib.bib16)）并由随机权重平均（SWA）（Izmailov 等，2018 (https://arxiv.org/html/2605.26484#bib.bib17)）重新焕发活力，认为对优化路径上的点求平均能够逼近平坦局部最小值的中心，从而增强泛化能力。近期，诸如最新权重平均（LAWA）（Kaddour，2022 (https://arxiv.org/html/2605.26484#bib.bib4)）和预训练模型平均（PMA）（Li 等，2025 (https://arxiv.org/html/2605.26484#bib.bib1)）等框架已成功将这一范式适配到 LLM 预训练中。通过均匀或指数移动平均聚合后期检查点，这些方法能够持续降低验证损失并增强下游推理任务的性能，为最终收敛模型提供了“免费午餐”式的改进。 尽管模型合并已被广泛采用，但支撑其在 LLM 非凸景观中应用的几何原理仍不清楚。虽然现有文献将这些增益归因于*局部平坦性*的静态属性（Izmailov 等，2018 (https://arxiv.org/html/2605.26484#bib.bib17)；Chen 等，2017 (https://arxiv.org/html/2605.26484#bib.bib19)；Li 等，2025 (https://arxiv.org/html/2605.26484#bib.bib1)），但这一视角忽略了优化轨迹本身的动态几何。合并过程如何将原始梯度的混沌震荡转化为一致的下降路径，这一机制尚不明确。我们旨在通过将焦点从静态损失景观转向*轨迹几何*来弥合这一差距，解决一个基本问题：合并检查点所形成流形的几何特性是什么？我们能否利用这种结构进行进一步改进？ 我们的贡献总结如下： - • **合并轨迹中的秩一子空间现象**。我们经验性地表明，后期*合并后*检查点集中在一个近似一维的线性流形上。具体而言，相邻*原始*检查点之间的插值呈现凸谷曲线（中点损失低于两端点），而相邻*合并后*检查点之间的插值几乎变为单调，表明路径更为平滑。通过广泛的 PCA 分析，我们证明这一*秩一子空间*捕获了主导下降方向（解释方差 >94%>94%），这与原始轨迹形成鲜明对比——原始轨迹的方差分散在多个主成分上，且其投影是非单调的。 - • **Extra-Merge 策略**。基于这一几何洞见，我们提出 Extra-Merge，一种无需训练的算法，通过估计子空间切向来外推进展，无需额外梯度更新。我们在*河谷*损失框架（Wen 等，2024 (https://arxiv.org/html/2605.26484#bib.bib18)）下提供了严格的理论证明，表明平均操作充当了几何低通滤波器，使轨迹与最优下降方向对齐，从而通过外推实现损失降低。 - • **跨尺度和优化器的普适性**。最后，我们在广泛的模型规模（GPT-2/LLaMA，124M 到 2B）和学习调度器上验证了 Extra-Merge。关键在于，我们的方法在 Pythia-12B（Biderman 等，2023 (https://arxiv.org/html/2605.26484#bib.bib67)）下游基准上产生了一致的零样本准确率提升，并有效泛化到 Muon 优化器的正交更新机制，表明秩一子空间是 LLM 训练中鲁棒的、与优化器无关的属性。 #### 利益冲突声明。作者声明与本研究无财务利益冲突。 ## 2 相关工作 #### 预训练中的模型合并。 模型合并是机器学习中的经典课题（Utans，1996 (https://arxiv.org/html/2605.26484#bib.bib20)；Chen 等，2017 (https://arxiv.org/html/2605.26484#bib.bib19)；Wortsman 等，2022 (https://arxiv.org/html/2605.26484#bib.bib21)；Yang 等，2024 (https://arxiv.org/html/2605.26484#bib.bib44)）。在本文中，我们专门使用预训练过程中保存的检查点，并对这些检查点参数求平均以提升性能，无需大量计算资源。这一范式在深度学习中的现代复兴主要归功于**随机权重平均（SWA）**（Izmailov 等，2018 (https://arxiv.org/html/2605.26484#bib.bib17)）。SWA 表明，对使用周期性或高恒定学习率的 SGD 遍历的权重求平均，能够获得更宽的最优解并实现比标准训练更好的泛化。更近期，这些技术已被专门适配到大型语言模型（LLM）中。**最新权重平均（LAWA）**（Kaddour，2022 (https://arxiv.org/html/2605.26484#bib.bib4)；Sanyal 等，2023 (https://arxiv.org/html/2605.26484#bib.bib7)）和**预训练模型平均（PMA）**（Li 等，2025 (https://arxiv.org/html/2605.26484#bib.bib1)）研究了在预训练中平均后期检查点的益处。类似地，近期研究（Tian 等，2025 (https://arxiv.org/html/2605.26484#bib.bib3)；Liu 等，2024 (https://arxiv.org/html/2605.26484#bib.bib43)；Luo 等，2025 (https://arxiv.org/html/2605.26484#bib.bib46)）将合并直接集成到预训练循环中。与先前将检查点视为静态插值点的研究不同，我们将其解释为动态序列，从而能够外推到观察到的权重凸包之外。为确保对新颖性和差异性的全面评估，我们在附录 A (https://arxiv.org/html/2605.26484#A1) 中提供了模型合并景观的详细分类学分析。 #### 损失景观的几何结构。 损失景观的几何结构是理解深度神经网络成功的基础（Fort 和 Ganguli，2019 (https://arxiv.org/html/2605.26484#bib.bib40)；Gur-Ari 等，2018 (https://arxiv.org/html/2605.26484#bib.bib41)；Li 等，2018 (https://arxiv.org/html/2605.26484#bib.bib42)）。与合并密切相关的现象是**线性模式连通性（LMC）**（Frankle 等，2020 (https://arxiv.org/html/2605.26484#bib.bib23)；Gotmare 等，2018 (https://arxiv.org/html/2605.26484#bib.bib34)；Goodfellow 等，2014 (https://arxiv.org/html/2605.26484#bib.bib35)；Vlaar 和 Frankle，2022 (https://arxiv.org/html/2605.26484#bib.bib38)；Lucas 等，2021 (https://arxiv.org/html/2605.26484#bib.bib39)），该现象表明在相同任务上独立训练的多个模型通常可以通过低损失线性路径相连。虽然 LMC 通常处理收敛后且独立模型之间的连通性，我们的工作研究了**单次**预训练运行后期阶段轨迹的局部几何结构。与我们提出的 Extra-Merge 最相关的是关于优化动态的**河谷**景观视角（Cohen 等，2021 (https://arxiv.org/html/2605.26484#bib.bib31)；Wen 等，2024 (https://arxiv.org/html/2605.26484#bib.bib18)；Song 等，2024 (https://arxiv.org/html/2605.26484#bib.bib30)；Davis 等，2025 (https://arxiv.org/html/2605.26484#bib.bib72)）。这些工作将神经网络损失景观描述为由平坦方向（“河流”）和尖锐方向（“山脉”）组成，分别对应具有不同曲率的 Hessian 特征方向。此外，近期研究利用这一直觉提出了有效的 LLM 预训练加速策略（Wang 等，2025 (https://arxiv.org/html/2605.26484#bib.bib11)；Zhou 等，2025 (https://arxiv.org/html/2605.26484#bib.bib9)；Wang 等，2024 (https://arxiv.org/html/2605.26484#bib.bib32)），为该视角的正确性提供了经验验证。 ## 3 预备知识 #### 问题设定。 我们考虑以 θ∈Rd\\theta\\in\\mathbb{R}^{d} 参数化的大型语言模型（LLM）的预训练。优化过程通过在一个数据集上最小化损失函数 L(θ)\\mathcal{L}(\\theta) 生成参数轨迹 {θ0,θ1,...,θT}\\{\\theta_{0},\\theta_{1},\\dots,\\theta_{T}\\}，通常使用 AdamW 等优化器。我们将训练步骤 tt 保存的检查点记为 θt\\theta_{t}。 参考图注 图 1：合并策略比较。我们将原始优化轨迹的验证损失与 PMA（均匀）和基于 EMA 的合并方案（γ∈{0.1,0.2}\\gamma\\in\\{0.1,0.2\\}）进行比较。实验在 GPT-2 Small（124M）上进行，训练数据为 FineWeb（τ=500,n=8\\tau=500,n=8）。 参考图注 图 2：从凸性到单调性的几何转变。我们可视化沿着相邻检查点（τ=500\\tau=500）之间线性插值路径 θ(α)=(1−α)θt+αθt+τ\\theta(\\alpha)=(1-\\alpha)\\theta_{t}+\\alpha\\theta_{t+\\tau} 的损失景观。**左图（原始）**：轨迹呈现**凸谷曲线**，中点损失显著低于两端点。**右图（合并后）**：在合并策略（τ=500,n=8\\tau=500,n=8）下，景观转变为**单调下降**。 #### LAWA 框架。 当前关于预训练中模型合并的研究主要由**最新权重平均（LAWA）**框架（Kaddour，2022 (https://arxiv.org/html/2605.26484#bib.bib4)；Sanyal 等，2023 (https://arxiv.org/html/2605.26484#bib.bib7)）所涵盖。LAWA 聚合一个包含 nn 个检查点的序列，这些检查点以固定间隔 τ\\tau 采样，计算加权平均：
θavg(t)=∑i=0n−1wi⋅θt−iτ,s.t.∑i=0n−1wi=1,\\theta_{\\text{avg}}^{(t)}=\\sum_{i=0}^{n-1}w_{i}\\cdot\\theta_{t-i\\tau},\\quad\\text{s.t.}\\quad\\sum_{i=0}^{n-1}w_{i}=1, (1)
其中 wi=0n−1{w}_{i=0}^{n-1} 是标量权重，决定每个检查点的贡献。该公式允许各种加权方案，例如指数移动平均（EMA），其中 θavg(t)=γθt+(1−γ)θavg(t−τ)\\theta_{\\text{avg}}^{(t)}=\\gamma\\theta_{t}+(1-\\gamma)\\theta_{\\text{avg}}^{(t-\\tau)}。 #### 基线：预训练模型平均（PMA）。 虽然 EMA 在凸优化中是标准方法，但近期关于**预训练模型平均（PMA）**的工作（Li 等，2025 (https://arxiv.org/html/2605.26484#bib.bib1)）提供了令人信服的证据，表明简单的**均匀平均**（即 wi=1/nw_{i}=1/n）对于 LLM 预训练更为优越。PMA 认为，在 LLM 的高维非凸景观中，均匀权重能够有效近似局部损失盆地的质心，比 EMA 提供更好的稳定性。为严格建立我们的基线，我们在图 1 (https://arxiv.org/html/2605.26484#S3.F1) 中验证了这一观察。我们在 GPT-2 Small 上比较了原始轨迹与使用均匀权重（PMA）和 EMA 权重（γ∈0.1,0.2\\gamma\\in\\{0.1,0.2\\}）的 LAWA。与 Li 等（2025 (https://arxiv.org/html/2605.26484#bib.bib1)）一致，均匀平均产生了最低且最稳定的验证损失。因此，我们在本工作中采用 PMA（均匀平均）作为标准基线，记为 θavg(t)=1n∑i=0n−1θt−iτ\\theta_{\\text{avg}}^{(t)}=\\frac{1}{n}\\sum_{i=0}^{n-1}\\theta_{t-i\\tau}。 ## 4 主要发现 ### 4.1 预热：成对连通性 为揭示模型合并的几何机制，我们首先探究**相邻**检查点之间的连通性。这是一个基本测试：如果优化轨迹是平滑的，步骤之间的线性插值应产生单调改进。反之，非单调性则意味着曲率或振荡。我们分析沿着线性路径 θ(α)=(1−α)θt+αθt+τ\\theta(\\alpha)=(1-\\alpha)\\theta_{t}+\\alpha\\theta_{t+\\tau}（α∈[0,1]\\alpha\\in[0,1]）的验证损失 L(θ(α))\\mathcal{L}(\\theta(\\alpha))。图 2 (https://arxiv.org/html/2605.26484#S3.F2) 对比了原始轨迹（θrawt→θrawt+τ\\theta^{\\text{raw}}t\\to\\theta^{\\text{raw}}{t+\\tau}）与合并后轨迹（θavgt→θavgt+τ\\theta^{\\text{avg}}t\\to\\theta^{\\text{avg}}{t+\\tau}）的景观，实验条件与图 1 (https://arxiv.org/html/2605.26484#S3.F1) 相同。出现了两种不同的几何模式： - • **原始检查点呈现凸谷曲线**。如图 2 (https://arxiv.org/html/2605.26484#S3.F2)（左）所示，原始检查点之间的插值始终形成 U 形凸谷曲线。值得注意的是，中点损失（α≈0.5\\alpha\\approx 0.5）显著低于两端点。这一现象表明，连续的原始步骤并非严格沿着谷底“下坡”移动。相反，它们在损失谷的高曲率壁之间振荡，实际上是在最优路径周围“弹跳”。 - • **合并后检查点揭示单调下降**。与此形成鲜明对比的是，图 2 (https://arxiv.org/html/2605.26484#S3.F2)（右）显示合并后检查点之间的路径是严格单调的。损失从 tt 到 t+τt+\\tau 平滑下降，没有观察到原始轨迹中的障碍或盆地结构。这表明平均操作充当了几何滤波器，有效抵消了正交方向的振荡。 参考图注