🤖 machine learning

Cosine-Gated Adam-Decay: Drop-In Staleness-Aware Outer Optimization for Decoupled DiLoCo

本文提出了一种即插即用且感知年龄的外层优化器 Cosine-Gated Adam-Decay（CGAD），用于异步 DiLoCo，该优化器通过衰减函数和余弦截断对伪梯度进行缩放以缓解陈旧性，相较于基于标准 Nesterov 和 Adam 的基线方法，其在多种 Llama 模型规模上均展现出更优的稳定性，并提供了一个与最大延迟无关的理论收敛界。

原作者： Vatsal Shah, Jiahao Sun

发布于 2026-05-12

📖 1 分钟阅读☕ 轻松阅读

CC BY 4.0

原作者： Vatsal Shah, Jiahao Sun

原始论文采用 CC BY 4.0 许可（http://creativecommons.org/licenses/by/4.0/）。 ✨ 这是对下方论文的AI生成解释。它不是由作者撰写或认可的。如需技术准确性，请参阅原始论文。阅读完整免责声明

想象一下，你正在试图教导一个庞大的学生团队（即计算机模型）如何写故事。在一个理想的世界里，所有人都在同一个房间里协作，彼此即时地大声交流更新。但在现实世界中，尤其是面对超大规模团队时，有些学生行动迟缓，有些相距甚远，还有些只是状态不佳。

这就是异步训练的问题。“老师”（中央计算机）接收来自学生们的建议（梯度），但当这些建议到达时，它们可能已经过时了。给出建议的学生可能早已翻到了书本的不同页面。

问题：“陈旧”建议的陷阱

目前处理这一问题的标准方法，就像是一位盲目信任每一条建议的老师，无论这些建议多么陈旧。

场景：一名学生根据 10 分钟前的位置，发送了一张纸条说：“向左转！”
现实：该学生现在已经在前方 100 米处了。如果老师根据那张旧纸条向左转，整个班级就会撞墙。
论文的发现：标准方法（称为Nesterov 动量）在建议过于陈旧时会陷入混乱。它会持续在错误的方向上累积“动量”（速度），导致训练失控，尤其是随着模型规模增大（从 2500 万参数增加到 70 亿参数）时更是如此。

解决方案：CGAD（余弦门控 Adam 衰减）

作者提出了一种名为CGAD的新方法。你可以将其视为一种智能过滤器，或是对 incoming 建议的“交通警”。

以下是其工作原理，使用一个简单的类比：

“新鲜度”折扣（指数衰减）：
想象每条建议都带有时间戳。建议越旧，其价值越低。CGAD 根据建议的年龄对其应用折扣。如果建议是 10 分钟前的，其价值减半；如果是 20 分钟前的，其价值几乎为零。这防止了老师对旧闻反应过度。
“硬性截断”（余弦门控）：
这是本文的秘诀所在。有时，某条建议极其陈旧（例如，一名学生已经一小时没说话了），以至于它不仅仅是“价值降低”，而是危险的。
- 旧方法：即使是极其陈旧的建议也会获得微小的权重。久而久之，这些微小的垃圾信息会堆积起来，混淆系统。
- CGAD 方法：它设有一条“截断线”。如果建议超过某个时间点（例如 32 轮），系统就会说：“不行，那太陈旧了。完全忽略它。”它将该值设为零。

为何这很重要：“规模保险”

论文在三种规模的模型上测试了该方法：

小型（2500 万参数）：新方法表现良好，但旧方法不会立即崩溃。这就像驾驶卡丁车；即使转向错误，你也可能只是原地打转。
中型（10 亿参数）：旧方法会严重崩溃。新方法则行驶平稳。
超大型（70 亿参数）：这是论文提出最大主张的地方。当模型如此庞大时，“陈旧建议”问题会变成一场灾难。
- 旧方法（Nesterov）失败得如此严重，以至于模型学不到任何东西（其表现甚至不如随机猜测）。
- 仅使用“折扣”的方法（Adam-Decay）表现尚可，但变得非常不可预测。一次运行可能成功，下一次可能失败。
- CGAD是唯一保持稳定的方法。这种“硬性截断”充当了规模保险。它保证即使网络混乱且缓慢，老师也绝不会听取那些会毁掉整个项目的危险陈旧建议。

结论

论文声称，只需添加一个“新鲜度过滤器”来忽略过于陈旧的建议，你就可以在混乱、缓慢或不可靠的网络上训练大规模 AI 模型，而无需担心训练崩溃。它将一个脆弱的系统转变为一个稳健的系统，确保当你最终部署模型时，它确实能正常工作。

简而言之：不要听取过于陈旧的建议。如果它真的太旧，就彻底扔掉它。这条简单的规则拯救了巨型 AI 模型的训练，使其免于崩溃。

技术摘要：用于解耦 DiLoCo 的余弦门控 Adam 衰减（CGAD）

1. 问题陈述

在异步解耦 DiLoCo（DDiLoCo）系统中，工作节点在执行多次局部优化步骤后，才与全局参数服务器交换伪梯度。该架构面临的一个关键挑战是陈旧性（staleness）：由于网络延迟、落后工作节点或异步预算限制，在第 $t$ 个外层轮次应用的伪梯度可能是在 $\tau$ 个轮次之前计算的。

DiLoCo 的标准外层优化器是带有 Nesterov 动量的 SGD。该方法对每个传入的伪梯度应用统一的更新步长，而不考虑其年龄（ $\tau$ ）。随着 $\tau$ 增加，动量缓冲区会累积对损失景观的陈旧视图。因此，优化器会采取过冲的步骤，指向那些不再下降当前损失的方向，导致不稳定性或发散，特别是在模型规模增大时。现有的异步动量-SGD 理论分析表明，收敛速率会以与 $\tau_{max}^2$ （最大延迟的平方）成正比的因子退化。

2. 方法：余弦门控 Adam 衰减（CGAD）

作者提出了余弦门控 Adam 衰减（CGAD），作为 Nesterov 外层优化器的即插即用替代品，它明确地将更新条件设定为伪梯度的年龄。

核心机制

CGAD 修改了标准 Adam 更新，在伪梯度 $g_f$ 进入 Adam 的一阶和二阶矩缓冲区之前，将其按因子 $\sigma(\tau)$ 进行缩放：
$\sigma(\tau) = \gamma(\tau) e^{-\alpha \tau}$

其中：

指数衰减（ $e^{-\alpha \tau}$ ）： 模拟陈旧梯度的信息衰减。随着 $\tau$ 增大，梯度对动量缓冲区的贡献呈指数级减少。
余弦门（ $\gamma(\tau)$ ）： 定义为平滑截止函数：
$\gamma(\tau) = \frac{1}{2} \left(1 + \cos\left(\frac{\pi \tau}{\tau_{cut}}\right)\right) \cdot \mathbb{I}[\tau \le \tau_{cut}]$
该门对于新鲜梯度等于 1，在选定的截止值 $\tau_{cut}$ 处平滑衰减至 0，并严格将 $\tau > \tau_{cut}$ 的贡献归零。

关键设计选择

缓冲区前门控（Pre-Buffer Gating）： 缩放 $\sigma(\tau)$ 在梯度更新一阶矩（ $m$ ）和二阶矩（ $v$ ）缓冲区之前应用。这防止了陈旧梯度膨胀二阶矩（ $v$ ），否则这会削弱后续的新鲜更新。
分片变体（PA-CGAD）： 对于部分同步调度器（如 Streaming DiLoCo），该方法跟踪各个分片（fragment）的年龄（ $a_f$ ），并基于 $\max(\tau, a_f)$ 对更新进行门控，从而为网络延迟和调度延迟提供统一的陈旧性模型。
超参数： 该方法引入了两个超参数， $\alpha$ （衰减率）和 $\tau_{cut}$ （截止值），其默认值（ $\alpha=0.2, \tau_{cut}=32$ ）可跨模型规模迁移。

3. 理论贡献

本文针对平滑非凸目标上的理想化门控自适应更新，提供了非渐近收敛界。

陈旧性偏差独立性： 收敛界中的陈旧性偏差项取决于门控超参数 $\alpha$ （具体为 $O(1/e^\alpha)$ ），而不是实际的最大延迟 $\tau_{max}$ 。
与标准分析的对比： 异步动量-SGD 的标准分析得出的复杂度项随 $\tau_{max}^2$ 缩放。通过将该项替换为依赖于 $\alpha$ 的项，只要 $\tau_{cut} \ge 1/\alpha$ ，门控机制在理论上就将收敛惩罚与最坏情况延迟分布解耦。
范围： 作者指出，这主要是对门控机制的动机说明，而非针对经验 Adam 的完整收敛定理，因为它假设了有界的归一化步长，而不是从经验统计中推导这些步长。

4. 实验结果

实验在三个规模（25 M、1 B 和 7 B 参数）的 Llama 风格语言模型预训练上进行，使用了高达 $\tau=16$ 的受控延迟。

稳定性与收敛性

1 B 规模： 在 $\tau=16$ 时，标准 Nesterov 配方发散（损失 $\approx 319$ ），而 CGAD 实现了 7.44 的损失，与其 $\tau=0$ 的性能（7.48）相当。Nesterov 方法在 $\tau \ge 8$ 时完全失去收敛性。
25 M 规模： CGAD 和其他陈旧性感知基线（Adam-Decay、SDM、Delayed Nesterov）均保持稳定，而 Nesterov 在 $\tau \ge 4$ 时发散。在此规模下，陈旧性感知方法之间的差异微乎其微。

规模保险与鲁棒性

在7 B 规模下，受限于内存条件（bf16 权重、8 位 AdamW），出现了最显著的发现：

种子间方差： 最接近的基线 Adam-Decay（不带余弦截止的 CGAD）表现出种子间标准差（ $\sigma$ ）随模型从 25 M 扩展到 7 B 而增加了 27 倍（在 $\tau=8$ 时从 0.07 增加到 1.89）。相比之下，CGAD 的 $\sigma$ 保持平稳（约 0.38）。
单次运行风险： 使用“均值 + $\sigma$ "作为单次部署风险的代理指标，Adam-Decay 在 7 B / $\tau=8$ 时的风险（13.19）超出了随机水平损失（10.83）2.4 nats，意味着单次生产运行失败的概率很高。CGAD 的风险（11.28）远低于此阈值。
失效机制： 作者将 Adam-Decay 在 7 B 规模下的失效归因于陈旧梯度在二阶矩缓冲区（ $v$ ）中累积的平方舍入噪声，这些梯度从未被完全归零。CGAD 中的余弦截止通过严格将 $\tau_{cut}$ 之外的贡献归零来防止这种情况。

调参敏感性

Nesterov 与 CGAD： 即使重新调整 Nesterov（完全移除动量），它在高延迟下也无法匹配经过调整的 CGAD 配置的性能。调整后的 Nesterov 需要针对不同延迟范围使用不同的超参数，而 CGAD 在整个范围内保持了鲁棒性。
超参数敏感性： 最优 $\alpha$ 随规模略有变化（25 M 时为 0.2，1 B 在 $\tau=8$ 时为 0.4），但默认值 $\alpha=0.2$ 在整个延迟谱系中仍保持鲁棒。

5. 意义与主张

本文主张 CGAD 是异步 DiLoCo 训练中最具鲁棒性的扩展选择，特别解决了 Nesterov 动量在大受控延迟下的脆弱性问题。

规模保险： 主要贡献是余弦截止，它充当了“规模保险”。虽然仅靠指数衰减（Adam-Decay）足以应对较小模型或全精度设置，但对于在压缩优化器状态（8 位/bf16）下运行的前沿规模模型（7 B+），硬截止至关重要。它防止了陈旧方差的累积，从而避免了灾难性的种子间方差和单次运行失败。
即插即用兼容性： CGAD 仅需对现有 Adam 实现进行一行修改，并增加两个超参数，其默认值可跨规模迁移。
实际影响： 对于前沿规模的异步训练，该方法决定了单次训练运行是成功发布还是失败（损失高于随机水平）。

作者在理论范围上保持谦逊，承认收敛界适用于理想化更新，且经验优越性是在 DiLoCo 家族和 Llama 风格预训练背景下具体展示的。他们指出，实验中使用的内部步数（ $H=8$ ）低于生产环境 DiLoCo 的设置（ $H=100-500$ ），尽管敏感性研究表明该排名在更高的 $H$ 下依然存在。