Can Model Merging Improve Aggregation in DiLoCo?¶
会议: NeurIPS 2026
arXiv: 2607.03011
代码: https://github.com/shoroi/isoloco
领域: 模型压缩
关键词: 低通信分布式训练, 模型合并, 任务算术, 各向同性聚合, 伪梯度正交化
一句话总结¶
本文首次从数学上揭示了分布式低通信训练（DiLoCo）的伪梯度聚合与任务算术模型合并在结构上的严格等价性，并提出基于谱均衡各向同性矫正与 Nesterov 动量的聚合算法 IsoLoCo，显著缓解了多节点与大局部步数下的负干扰性能退化。
研究背景与动机¶
随着超大语言模型预训练规模的快速扩张，传统数据并行（Data-Parallel, DP）对加速器集群间高带宽网络与紧密互联的强依赖，日益成为分布式计算的主要瓶颈。在跨数据中心或异构弱互联算力岛场景下，通信开销往往主导了整体训练时长。低通信分布式优化方法（如 DiLoCo）通过让各个计算节点在本地执行多轮（如几十到上百步）局部优化，随后仅周期性通信并平均各节点的参数位移（伪梯度）来驱动全局更新，极大压缩了节点间的同步频次。然而，这类方法普遍面临严峻的可扩展性瓶颈：在总计算量（FLOPs）严格对齐的前提下，一旦计算节点数量增加或本地训练步长拉长，本地模型副本在参数空间中显著发散，导致简单算术平均得到的全局更新方向质量剧烈恶化，与紧密同步的数据并行标准相比产生严重的性能塌陷。
与分布式优化并行发展的模型合并（Model Merging）研究中，基于任务算术（Task Arithmetic）合并多个独立微调模型时也面临几乎相同的困境——所谓的“参数负干扰”（Negative Parameter Interference）。已有研究表明，当合并的模型数量增多或单个模型微调步数延长时，不同模型在参数空间中的更新方向会发生严重谱冲突与干涉，导致合并后的多任务模型性能大幅落后于单个专家。然而，过去这两条研究脉络长期割裂：分布式训练聚焦于跨算力孤岛预训练相同目标，而模型合并主要被视为训练后的静态检查点融合策略。
本文的切入角度是：如果在形式上将 DiLoCo 每一轮外层同步的前序检查点视作预训练基座模型，各个工作节点在本地多步迭代后的参数位移视作“任务向量”，那么 DiLoCo 的外层聚合步在数学结构上完全等价于一轮单步的任务算术模型合并。核心 idea：将模型合并中消除参数负干扰的各向同性谱均衡技术（Iso-C）引入 DiLoCo 外层伪梯度聚合，并结合 Nesterov 动量提出 IsoLoCo，通过压平高能量奇异值方向从根本上化解多节点发散干扰，打破低通信分布式训练随节点数扩展的性能瓶颈。
方法详解¶
整体框架¶
IsoLoCo 针对多算力岛弱互联场景构建。在每个全局外层步 \(t\)，中央调度器向 \(R\) 个本地工作节点分发全局参数 \(\theta^{(t-1)}\)。各节点在各自的不相交数据分片上独立运行 \(H\) 步内部优化器（如 AdamW），生成各自的本地模型副本 \(\theta_i^{(t)}\)。节点将本地累积位移（即伪梯度 \(\Delta_i^{(t)} = \theta^{(t-1)} - \theta_i^{(t)}\)）汇总到中央调度端后，IsoLoCo 不再直接使用加权平均更新全局模型，而是依次执行三阶段聚合流程：首先在算术均值上计算二维权重矩阵的奇异值分解并进行谱均衡矫正；随后将矫正后的伪梯度注入 Nesterov 动量缓存；最后完成全局模型的前瞻更新。
%%{init: {'flowchart': {'rankSpacing': 24, 'nodeSpacing': 28, 'padding': 6, 'wrappingWidth': 400}}}%%
flowchart TD
A["全局参数分发与本地多步优化<br/>各节点独立迭代 H 步生成伪梯度"] --> B["伪梯度均值计算<br/>跨 R 个节点收集并计算局部参数差值均值"]
B --> C["各向同性谱均衡矫正<br/>2D参数矩阵SVD分解并以均值替换奇异值"]
C --> D["Nesterov 动量前瞻更新<br/>将谱矫正梯度注入动量缓冲完成全局步推进"]
D --> E["更新后的全局参数<br/>广播至下一轮外层优化迭代"]关键设计¶
1. DiLoCo 与任务算术的结构等价性映射：揭示伪梯度聚合的本质退化根源
传统 DiLoCo 在每一轮外层优化步 \(t\) 中，先由各工作节点计算其局部伪梯度 \(\Delta_i^{(t)} = \theta_i^{(t)} - \theta^{(t-1)}\)，并在外层使用学习率为 \(\eta\) 的 SGD 进行全局更新： $\(\theta^{(t)} = \theta^{(t-1)} + \eta \Delta^{(t)} = \theta^{(t-1)} + \eta \frac{1}{R} \sum_{i=1}^{R} \left(\theta_i^{(t)} - \theta^{(t-1)}\right)\)$ 这一形式与任务算术模型合并公式 \(\operatorname{Merge}(\theta_0, \{\theta_t\}) = \theta_0 + \lambda \sum \tau_t\) 完全一致，其中全局同步点 \(\theta^{(t-1)}\) 对应预训练基模 \(\theta_0\)，各节点副本 \(\theta_i^{(t)}\) 对应微调专家，伪梯度对应任务向量 \(\tau_i\)，节点数 \(R\) 对应任务数。由此清晰揭示出：DiLoCo 在大节点数 \(R\) 与大局部步数 \(H\) 下出现的性能骤降，本质上正是模型合并领域熟知的参数负干扰问题。因此，能够抑制任务向量谱冲突的合并算子可以直接迁移为外层优化器的新型聚合法则。
2. 各向同性谱均衡矫正：压缩高能奇异值方向以抑制参数负干扰
在对主流合并算法（TIES、DARE、TSV-M、Iso-C、Iso-CTS）的横向基准测试中，基于稀疏剪枝的算法收效甚微，而基于白化与正交补空间的算法计算复杂度与节点数成正比（需要 \(O(R)\) 次 SVD），在每轮外层同步中开销极其昂贵。IsoLoCo 选取并改造了各向同性公共子空间矫正机制（Iso-C）。对于所有二维权重矩阵（如注意力投影层与前馈网络权重），算法对节点平均后的伪梯度矩阵 \(\Delta_\ell^{(t)}\) 进行奇异值分解 \(U \Sigma V^\top\)，提取出奇异值向量 \(\boldsymbol{\sigma} = \operatorname{Diag}(\Sigma)\)，并计算奇异值均值 \(\bar{\sigma} = \frac{1}{r} \sum_{k=1}^r \boldsymbol{\sigma}_k\)。随后，算法使用该均值强制替换对角谱矩阵，将梯度矩阵投影为： $\(\Delta_\ell^{(t)} \leftarrow \bar{\sigma} U V^\top\)$ 这一操作在保留参数更新主要子空间几何方向的前提下，强行将各方向的能量拉平。方向敏感度分析证实，性能提升的核心源于压平了少数主导更新的高能奇异值方向（clip_high），防止了个别节点因局部过度探索而沿特定主方向形成破坏性的参数漂移。对于一维参数（如 LayerNorm 与偏置项），则保持标准算术平均。
3. Nesterov 动量注入与外层前瞻推进：协同弥补单步谱矫正的时序惯性
模型合并原本针对训练完成后的静态单次操作，缺乏针对动态序列优化的时序关联机制。直接将无动量的 Iso-C 替代 DiLoCo SGD 虽然表现优异，但未充分利用外层轨迹的历史动量信息。IsoLoCo 将谱矫正后的伪梯度 \(\Delta^{(t)}\) 融入外层 Nesterov 动量更新架构： $\(B^{(t)} = \mu B^{(t-1)} + \Delta^{(t)}\)$ $\(P^{(t)} = \Delta^{(t)} + \mu B^{(t)}\)$ $\(\theta^{(t)} = \theta^{(t-1)} - \eta_{\text{out}} P^{(t)}\)$ 通过先对当前伪梯度实施谱均衡、再在历史动量缓冲上叠加具有前瞻特性的速度向量，IsoLoCo 既消除了当期多节点更新的发散冲突，又维持了沿损失曲面平滑谷底快速滑移的外层优化动力。
4. 基于 Newton-Schulz 迭代与 RMS 代理的显存与计算加速
标准 SVD 运算在 GPU 架构上缺乏硬件级 Tensor Core 加速，尤其在频繁外层同步中会拖慢端到端吞吐。IsoLoCo 进一步提出了面向工业部署的快速实现方案，利用五阶多项式 Newton-Schulz 迭代以纯通用矩阵乘法（GEMM）近似提取极性正交矩阵 \(\operatorname{NS}(M) \approx U V^\top\)。针对 Newton-Schulz 无法直接输出奇异值均值的问题，作者利用矩阵 Frobenius 范数构建了均方根（RMS）极速代理指标： $\(\operatorname{RMS}(\boldsymbol{\sigma}) = \frac{\|M\|_F}{\sqrt{r}} = \sqrt{\frac{1}{r}\sum_{i=1}^r \sigma_i^2}\)$ 由于柯西-施瓦茨不等式保证 \(\operatorname{RMS}(\boldsymbol{\sigma}) \ge \operatorname{mean}(\boldsymbol{\sigma})\)，该标量可直接在前向极性化前由未分解矩阵一步算得。实验证明该快速变体相比完整 SVD 实现了近两个数量级的正交化加速，带来整网端到端约 15% 的训练墙钟时间缩短，而验证损失性能差距仅约 0.45%。
实验关键数据¶
主实验¶
在 DCLM 数据集上，按照 Chinchilla 最优计算预算（训练 token 数为参数量的 20 倍）训练 LLaMA-2 架构语言模型。基准对照包括全同步 AdamW 数据并行（DP）、DiLoCo（SGD + Nesterov 动量）、DiLoCo SGD 以及无动量的合并算法外层基线。
表 1：不同外层聚合机制在 DiLoCo SGD 架构下的验证损失对比（178M 模型，节点数 \(R=8\)，局部步数 \(H=30\)）
| 方法 | 外层聚合机制类别 | 验证损失 (Valid Loss) | 相对 DP 性能退化 |
|---|---|---|---|
| AdamW DP | 同步数据并行基准 | 2.881 | 0.0% (黄金标准) |
| DiLoCo SGD | 简单伪梯度算术平均 | 2.990 | +3.78% |
| TIES | 稀疏剪枝 + 符号投票合并 | 2.997 | +4.03% |
| DARE | 随机丢弃与重缩放合并 | 2.993 | +3.89% |
| TSV-M | 低秩截断 SVD + 白化解相关 | 3.160 | +9.68% |
| Iso-CTS | 各向同性 + 任务特定正交补 | 2.986 | +3.64% |
| Iso-C | 各向同性公共子空间谱均衡 | 2.929 | +1.67% |
表 2：跨节点规模 \(R\)（\(1 \to 128\)）与模型容量的最终验证损失对比（固定 \(H=30\)）
| 模型规模 | 方法 | \(R=1\) | \(R=2\) | \(R=4\) | \(R=8\) | \(R=16\) | \(R=32\) | \(R=64\) | \(R=128\) |
|---|---|---|---|---|---|---|---|---|---|
| 178M | DiLoCo SGD | 2.887 | 2.905 | 2.937 | 2.990 | 3.064 | 3.178 | 3.357 | 3.601 |
| 178M | DiLoCo | 2.869 | 2.882 | 2.911 | 2.950 | 3.005 | 3.077 | 3.195 | 3.355 |
| 178M | DiLoCo SGD + Iso-C | 2.855 | 2.872 | 2.891 | 2.929 | 2.992 | 3.070 | 3.162 | 3.280 |
| 178M | IsoLoCo (本文) | 2.842 | 2.852 | 2.864 | 2.889 | 2.930 | 2.992 | 3.062 | 3.160 |
| 512M | DiLoCo | 2.630 | 2.639 | 2.653 | 2.675 | 2.706 | 2.767 | 2.863 | — (OOM) |
| 512M | IsoLoCo (本文) | 2.589 | 2.593 | 2.599 | 2.610 | 2.628 | 2.676 | 2.716 | — (OOM) |
| 1B (\(R=8\)) | DiLoCo: 2.450 | IsoLoCo: 2.421 | (相对增益: 1.2%) | — | — | — | — | — | — |
消融实验¶
表 3：不同同步间隔步数 \(H\) 下的性能容忍度消融（178M 模型，\(R=8\)）
| 方法 | \(H=30\) | \(H=60\) | \(H=120\) | \(H=240\) |
|---|---|---|---|---|
| DiLoCo | 2.950 | 2.957 | 2.976 | 3.008 |
| IsoLoCo (本文) | 2.889 | 2.907 | 2.944 | 2.991 |
表 4：关键设计与各向同性矫正机制的消融分析（178M 模型，\(R=8\), \(H=30\)）
| 配置 / 变体 | 验证损失 (Valid Loss) | 说明 |
|---|---|---|
| IsoLoCo (完整模型) | 2.889 | SVD 分解 + 奇异值均值缩放 + Nesterov 动量 |
| 快速实现 (Newton-Schulz + RMS) | 2.902 | 用 GEMM 迭代与 RMS 代替 SVD，端到端加速 15% |
| 非 Nesterov 动量 (标准动量) | 2.895 | 损失上升约 0.2%，需更小的动量系数 |
| 对最终更新量 \(P^{(t)}\) 做正交矫正 | 2.894 | 劣于直接对伪梯度 \(\Delta^{(t)}\) 进行各向同性矫正 |
仅削减高能奇异值 (clip_high) |
2.894 | 接近完整模型表现，证实压制支配性能量方向为主因 |
仅提升低能奇异值 (clip_low) |
2.931 | 仅收窄少量差距，表明低能量补足并非关键增益源 |
| 谱缩放采用 Muon 式 \(\sqrt{\text{fan\_out}/\text{fan\_in}}\) | 3.005 | 性能严重下滑，证明数据驱动的均值缩放必不可少 |
关键发现¶
- 节点扩展优势急剧放大：在 178M 规模下，当节点数从 \(R=1\) 增长至 \(R=128\) 时，DiLoCo 相对数据并行的性能恶化从负增益暴增至 16.5%，而 IsoLoCo 仅退化 9.7%；与 DiLoCo 相比，性能差距从 0.027 显著拉大到 0.195，证实谱均衡从根本上遏制了多节点参数干涉。
- 高能谱削减是核心机制：奇异值消融明确证实，单纯将高于均值的奇异值截断为均值（
clip_high），即可贡献 IsoLoCo 绝大部分增益；而提升低能方向（clip_low）收效有限。 - 与内部正交优化器正交互补：当将内部优化器由 AdamW 替换为 Muon 时（MuLoCo），IsoLoCo 的外层谱矫正能够与其内层权重正交化形成良好配合。在 \(R \ge 32\) 的超大规模节点配置下，IsoLoCo + Muon 取得了全场最优验证损失，并在相同外部谱矫正下大幅超越单一维度的内层正交化策略。
亮点与洞察¶
- 跨领域等价性洞察：敏锐捕捉到了分布式低通信训练中的伪梯度外层平均与模型微调后任务算术合并在数学结构上的完全一致性，打破了优化算法与模型融合之间的学术壁垒。
- 极简而精准的谱各向同性修正：没有盲目照搬依赖繁琐多步正交补的复杂合并方案（如 TSV-M 或 Iso-CTS），而是选定计算开销低且匹配分布式任务同质特性的各向同性公共子空间操作，仅需 \(O(1)\) 次 SVD 即可取得显著效果。
- 工程落地的硬件友好化：通过引入五阶 Newton-Schulz 多项式结合 RMS 标量代理，把原本难以在 GPU 集群上高效并行化的 SVD 转化为纯 Tensor Core 友好的 GEMM 运算，实现了算法先进性与系统级吞吐的双赢。
局限与展望¶
- 理论收敛性证明尚待补全：论文以实证分析和几何谱直觉为主，尚未给出带有各向同性非线性投影算子的非凸随机优化严格收敛阶证明。
- 超大规模与异构网络验证缺失：实验最大模型停留在 1B 规模，尚未测试千亿级密集模型或混合专家（MoE）结构，且实验基于同构集群仿真，在真实跨广域网（WAN）的极端非均匀延迟环境下表现仍需实测检验。
- 未来方向：探索 IsoLoCo 与极端梯度量化/有损压缩技术的正交组合，以及在单节点超大批次预训练和持续学习场景下的泛化应用。
相关工作与启发¶
- vs DiLoCo (Douillard et al., 2024): DiLoCo 外层直接采用伪梯度的简单算术平均配合 Nesterov 动量，在工作节点数增多或通信间隔变长时因参数负干涉性能恶化严重；IsoLoCo 针对二维矩阵做各向同性谱均衡矫正，显著收窄了多节点下与数据并行的性能差距。
- vs MuLoCo (Thérien et al., 2026): MuLoCo 仅在节点内部训练步采用 Muon 进行内层正交化，外层聚合仍是简单平均；IsoLoCo 作用于外层伪梯度聚合步，单次外层正交开销仅为内层 Muon 的 \(1/H\)，并在超大节点规模下与 Muon 具备互补叠加效应。
- vs Iso-C (Marczak et al., 2025): Iso-C 最初设计用于异构下游多任务检查点的后训练单次合并；本文将其泛化为预训练动态外层迭代步，结合 Nesterov 动量与 Newton-Schulz 硬件加速机制，成功改造为高效分布式优化器。
评分¶
- 新颖性: ⭐⭐⭐⭐⭐ 首次确立 DiLoCo 伪梯度聚合与任务算术模型合并的深层数学等价性，并提出谱矫正外层优化方案。
- 实验充分度: ⭐⭐⭐⭐⭐ 涵盖 178M/512M/1B 模型、节点数从 1 扩展至 128、同步间隔从 30 步至 240 步，并配有深入的谱方向消融与计算时间评测。
- 写作质量: ⭐⭐⭐⭐⭐ 叙述逻辑严丝合缝，公式克制且直击要害，算法对比与图表呈现极为清晰。
- 价值: ⭐⭐⭐⭐⭐ 为异构弱连接算力集群下大模型高效预训练提供了极具实用价值的优化工具，工程与学术价值俱佳。