Outer-Momentum Restarting in High-Dimensional Two-Phase Optimization

技术摘要：高维两阶段优化中的外动量重启

问题陈述
现代大规模语言模型（LM）训练日益受到分布式设置中通信成本的瓶颈制约。为缓解这一问题，DiLoCo 等两阶段优化方法逐渐受到关注，其中工作节点在执行多次局部更新（内阶段）后，通过外部优化器聚合进展（外阶段）。尽管这些方法行之有效，但它们引入了新的超参数敏感性层次。近期的理论研究表明，外部优化器在内循环诱导的“有效谱”上运行。具体而言，重球（HB）外动量可能会限制内循环已取得显著进展的方向，导致收敛缓慢或不稳定，特别是在通信周期较长或外动量值较高的情况下。本文解决的核心问题是两阶段优化器对外动量和学习率选择的脆弱性，这给预训练场景中的调优带来了复杂性。

方法论
作者提出周期性重启外动量缓冲区，作为一种控制“外记忆”的补充机制。该方法结合了理论分析与实证验证：

1. 理论框架：

   • 作者采用基于经验神经正切核（NTK）下的线性化平方损失近似的确定性、模态模型。
   • 他们分析了跨通信轮次的残差特征坐标的动态。内循环产生伪梯度 $g_t = \sigma x_t$，其中 $\sigma$ 代表有效进展。
   • 对于重球（HB）动量，转移矩阵 $T_{HB}$ 的行列式为 $\beta_{out}$，导致阻尼率独立于进展 $\sigma$。这意味着即使高进展模态也受制于相同的缓慢包络。
   • 作者推导了模态重启收缩。通过每 $K$ 轮重置外动量缓冲区（$m \leftarrow 0$），系统利用相位抵消。在复数域（高动量情况下的典型特征）中，残差坐标在包络内振荡。精心选择的重启周期 $K$ 可以将重置对齐到投影残差最小的相位，从而有效丢弃陈旧动量，同时保留内循环进展。
   • 他们推导了收缩因子 $\chi_K(\sigma)$ 的递推关系，并表明当相位依赖的投影项幅度小于 1 时，重启优于未重启的 HB。

2. 实验设置：

   • 该方法在 H200 GPU 上，使用 2 副本 DiLoCo 配置预训练 1.5 亿参数量的 LLaMA 模型进行评估。
   • 实验涵盖了重球和 Nesterov（NAG）外部优化器，涉及不同的通信周期（$S$）和超参数网格。
   • 该研究将标准 DiLoCo 与带周期性重启的 DiLoCo 进行比较，测试了全局、每模态和分块重启策略。

主要贡献

• 机制识别： 本文指出，两阶段优化中的周期性重启通过利用外动量动力学中的相位抵消起作用，这与 Nesterov 动量修改谱包络的方式不同。
• 理论推导： 提供了重启收缩因子的闭式分析，证明重置可以通过丢弃导致振荡或不对齐的过时动量状态来降低有效阻尼率。
• 鲁棒性增强： 主要贡献在于证明周期性重启显著拓宽了外超参数（学习率 $\nu$ 和动量 $\beta_{out}$）的稳定区域。
• 实用基准： 作者确立了周期性重启作为一个简单、鲁棒的基准，减少了对精细超参数调优的需求，这是大规模预训练的关键属性。

结果

• 稳定性与鲁棒性： 在 1.5 亿参数量 LLaMA 模型上的实证结果表明，若无重启，高外动量值（$\beta_{out}$）与长通信周期结合会导致性能急剧下降或发散。周期性重启在很大程度上消除了这些失效区域，允许在更广泛的 $\beta_{out}$ 和 $\nu$ 范围内进行稳定训练。
• 超参数敏感性： 研究表明，重启周期 $K$ 是比外动量系数更“宽容”的超参数。虽然未重启的运行需要随着通信周期 $S$ 的变化重新调整 $\beta_{out}$，但重启运行在固定 $K$ 的情况下，能在不同的 $S$ 值下保持稳定的性能。
• 性能： 就峰值性能而言，周期性重启在保持最佳可达验证困惑度的同时，显著降低了超参数网格上的性能方差。
• 软重启与硬重启： 对“软重启”（部分保留动量）的实验显示，其相比硬重置仅有边际改进，但引入了对训练配置敏感的额外超参数（$\alpha, \beta$），因此作者倾向于更简单的硬重启方法。

意义与主张
本文主张，周期性外动量重启是一种简单而有效的机制，可稳定高维两阶段优化。其意义在于解决了预训练中优化器调优的“精细”性质，其中两阶段方法引入了额外的复杂性层次。通过限制无益的外记忆，该方法使外部优化器对外动量和学习率的选择不那么敏感。

作者谦逊地将自己的工作定位为提供鲁棒基准而非通用解决方案。他们指出，虽然分析侧重于固定周期重置，但实证结果表明该方法对所选的具体周期相当不敏感。最后，他们提出了一个更广泛的未来工作问题：外循环是否可以自适应地决定何时刷新其记忆，这暗示了自适应重启规则、每层周期以及更大规模研究等潜在方向。