Cautious Optimizers: Improving Training with One Line of Code

该论文提出了一种仅需在 PyTorch 中修改一行代码即可应用于任意动量优化器的“谨慎优化器”（C-Optim），该方法在理论保证了收敛性的同时，显著提升了 LLM 预训练和图像分类任务的效率且无需额外超参数调整。

要点

这篇论文提出了一种非常巧妙的优化方法，叫做**“谨慎优化器”（Cautious Optimizers）**。它的核心思想可以用一句话概括：给现有的优化算法加了一行代码，让它在“走路”时变得更加小心，结果反而走得更快、更稳。

为了让你轻松理解，我们可以把训练人工智能模型想象成**“在迷雾中下山”**。

1. 背景：下山的故事

想象你是一位盲人探险家（AI 模型），你的目标是从山顶（高错误率）走到山脚（低错误率，即最佳状态）。

• 梯度（Gradient）：是你脚下的坡度，告诉你哪里是下坡。
• 动量（Momentum）：是你跑起来的惯性。如果你跑得太快，惯性可能会让你冲过头，甚至冲上对面的山坡，然后再滚回来。这就是为什么传统的优化器（如 AdamW）有时候会**“震荡”**，在谷底来回晃悠，浪费了很多时间。

2. 问题：惯性带来的“鲁莽”

过去十年，大家一直在寻找跑得更快、更稳的下山方法。虽然有很多新发明，但AdamW 依然是大家最信赖的“老大哥”。
但是，AdamW 有个缺点：它太依赖“惯性”了。有时候，它冲得太猛，明明前面是上坡（损失在增加），它却因为惯性还在往前冲。这就好比开车下坡时，司机没看路，只顾着踩油门，结果车子在谷底左右乱撞，很久才能停稳。

3. 解决方案：加一个“刹车检查员”

这篇论文的作者提出了一种极其简单的方法，只需要在代码里加一行，就能让任何基于“动量”的优化器变得“谨慎”。

这个“谨慎”是怎么工作的呢？
想象你在下山时，身边多了一个**“谨慎检查员”**。

• 当你准备迈出一大步（优化器计算的更新方向）时，检查员会先看看：“这一步的方向，和脚下的坡度方向一致吗？”
• 如果一致（都是下坡）：检查员说“好，冲！”（允许更新）。
• 如果不一致（你想往前冲，但脚下其实是上坡）：检查员立刻大喊“停！”（把这次更新直接抹零，不让走这一步）。

这就好比：

• 普通优化器：像是一个鲁莽的赛车手，不管前面是不是墙，只要油门踩下去就冲，撞墙了再倒车。
• 谨慎优化器：像是一个经验丰富的老司机，每次踩油门前都会看一眼路况。如果方向不对，他干脆就不踩油门，避免浪费时间和燃料。

4. 为什么这行代码这么厉害？

作者发现，加上这个“检查员”后，发生了两件神奇的事：

1. 不再乱撞（更稳定）：因为阻止了那些“往山上冲”的错误步伐，模型在训练过程中不会忽高忽低，损失函数（下山的高度）会单调下降，非常平稳。
2. 跑得更快（更高效）：听起来“不踩油门”会变慢？恰恰相反！因为省去了“冲过头再倒车”的时间，模型能更直接、更有效地到达山脚。在训练大语言模型（LLM）时，这意味着在同样的时间内，模型能“吃”进更多的数据，变得更聪明。

5. 实验结果：真的有效吗？

作者在几个大任务上测试了这种方法：

• 大语言模型（LLM）：就像训练 ChatGPT 这样的模型。使用“谨慎版”的 AdamW（叫 C-AdamW）和 Lion（叫 C-Lion），在同样的训练时间下，模型的效果（困惑度）更好，而且不需要重新调整复杂的参数。
• 图像分类：在识别图片的任务中，准确率也提高了。
• 通用性：这个方法不仅对 AdamW 有效，对 Lion、MARS 等其他优化器也有效。就像给所有车都装上了同一个“智能刹车系统”，大家都跑得更好了。

6. 总结：一行代码的魔法

这篇论文最迷人的地方在于它的简单。

• 以前：大家为了改进优化器，发明各种复杂的数学公式，需要调参调到头秃。
• 现在：作者发现，只要加一行代码（判断更新方向和梯度方向是否一致，不一致就跳过），就能让现有的优化器性能提升。

打个比方：
这就好比你在玩一个复杂的电子游戏，大家都在研究怎么升级装备、怎么加点才能通关。结果有人发现，只要在游戏里加一个“如果走错路就原地踏步”的简单规则，大家反而通关得更快了。

一句话总结：
“谨慎优化器”通过让 AI 在训练时“三思而后行”，避免了盲目冲动带来的浪费，用最少的一行代码，换来了更快的训练速度和更稳的模型表现。

技术摘要

这是一篇发表于 ICLR 2026 的会议论文《CAUTIOUS OPTIMIZERS: IMPROVING TRAINING WITH ONE LINE OF CODE》（谨慎优化器：用一行代码改进训练）的技术总结。

1. 研究背景与问题 (Problem)

• 现状：AdamW 长期以来是大语言模型（LLM）预训练和其他深度学习任务中的默认优化器。尽管社区多年来一直在寻找更快、更稳定的优化器（如 Lion, SHAMPOO, SOAP 等），但 AdamW 的主导地位从未受到实质性挑战。
• 痛点：
  1. 现有替代方案的局限性：许多新的优化器虽然声称性能提升，但通常需要大量的超参数调整（Hyperparameter Tuning）才能获得最佳效果，这限制了它们的广泛采用。
  2. 动量优化器的不稳定性：基于动量（Momentum）的优化器（如 Adam, Lion）在更新方向 $u_t$ 与当前梯度 $g_t$ 不一致时，可能会导致损失函数暂时增加，从而减缓收敛速度或引起震荡。
  3. 训练效率需求：随着模型规模扩大，更快的优化器意味着在相同时间内可以消耗更多的训练 Token，从而训练出能力更强的模型。

2. 核心方法论 (Methodology)

作者提出了一种名为**“谨慎优化器”（Cautious Optimizers, C-Opt）**的通用框架。其核心思想极其简单：只有当建议的更新方向与当前梯度方向一致（符号相同）时，才执行参数更新。

2.1 核心算法逻辑

对于任意基于动量的优化器，其标准更新公式为 $w_{t+1} \leftarrow w_t - \epsilon_t u_t$。
谨慎优化器将其修改为：
$$w_{t+1} \leftarrow w_t - \epsilon_t u_t \circ \phi(u_t \circ g_t)$$
其中：

• $\circ$ 表示逐元素乘积。
• $g_t = \nabla L(w_t)$ 是当前梯度。
• $\phi$ 是一个重加权函数，基于 $u_t$ 和 $g_t$ 的乘积进行掩码（Masking）。

具体实现（一行代码）：
在 PyTorch 中，对于参数 $p$ 和更新量 $u$，梯度 $g$：

m = (u * g > 0).to(g.dtype)  # 如果更新方向和梯度方向一致，m 为 1，否则为 0
p.add(u * m / (m.mean() + eps), alpha=-lr) # 更新并重新缩放学习率

• 掩码机制：如果某个维度的更新方向与梯度方向相反（符号不一致），则该维度的更新被置零。
• 重缩放（Rescaling）：由于掩码会减少有效更新的维度，导致更新幅度减小，因此引入了缩放因子 $\alpha(x) = \frac{dim}{nnz(x>0) + \xi}$ 来补偿，保持更新幅度的期望值不变。

2.2 理论分析

• 哈密顿量保持（Hamiltonian Preservation）：论文证明了该修改保留了原始优化器（如 Adam, Lion）的哈密顿函数（Lyapunov 函数）结构。
• 单调下降保证：在连续时间动力学分析中，证明了谨慎优化器不仅保证总能量（哈密顿量）单调下降，还能保证损失函数 $L(w)$ 本身单调下降（在步长足够小的情况下）。相比之下，传统动量优化器由于惯性效应，损失函数可能会出现震荡。
• 收敛性：在温和条件下，证明了谨慎优化器不会陷入非平稳点，且收敛到局部最优解的速率优于或等于基线优化器。
• 离散时间分析：证明了在单步更新中，谨慎优化器带来的损失下降量至少不小于原始优化器。

3. 主要贡献 (Key Contributions)

1. 极简实现：提出了一种通用的性能提升方案，只需在现有动量优化器（如 AdamW, Lion）的代码中添加一行代码即可实现。
2. 理论突破：
   • 揭示了基于哈密顿动力学的一类新优化器家族。
   • 从理论上证明了该修改在保留基线优化器收敛保证的同时，加速了损失函数的下降。
   • 证明了该方法能保证损失函数的单调下降，解决了传统动量方法中常见的震荡问题。
3. 广泛的实证效果：在从 LLM 预训练到图像分类的多种任务中，展示了 consistent 的性能提升，且无需重新调整超参数。

4. 实验结果 (Results)

4.1 2D 优化玩具实验

• 在 2D 二次函数优化中，对比了带 Polyak 动量的梯度下降（GDM）与谨慎版本（C-GDM）。
• 结果：C-GDM 消除了震荡，实现了损失函数的单调下降，收敛轨迹更平滑，且收敛速率更快。

4.2 大语言模型（LLM）预训练

• 设置：在 C4 数据集上预训练 1 亿参数量的 LLaMA 模型，以及在不同规模（1.2B）下使用 FineWeb-Edu 数据集。
• 对比：C-AdamW vs AdamW, C-Lion vs Lion。
• 结果：
  • 困惑度（Perplexity）降低：在所有模型规模下，C-AdamW 和 C-Lion 均取得了更低的验证集困惑度（例如 1.2B 模型上 C-AdamW 降低了 0.32%）。
  • 鲁棒性：C-Opt 对超参数（特别是学习率）更不敏感。在 Lion 实验中，C-Lion 甚至能在基线 Lion 发散的学习率下保持稳定训练。
  • 下游任务：在 1.2B 模型的 7 个下游任务（如 MMLU, ARC 等）中，C-AdamW 训练的模型在 5 个任务上表现更好。

4.3 图像分类

• 设置：在 Mini-ImageNet 上使用 ViT 架构。
• 结果：C-AdamW, C-LaProp, C-MARS 均显著优于其对应的基线优化器（Top-1 准确率提升 1% - 3% 不等）。

4.4 效率分析

• 吞吐量：在分布式训练（DDP）设置下，由于掩码和缩放操作是逐元素的，通信开销极小。实验显示 C-AdamW 的 Token 吞吐量仅比 AdamW 低约 3%，考虑到性能提升，性价比极高。

5. 意义与结论 (Significance & Conclusion)

• 通用性与易用性：该工作打破了“新优化器需要复杂调参”的魔咒。它提供了一种“即插即用”的增强方案，适用于任何基于动量的优化器。
• 理论指导实践：通过哈密顿动力学视角，解释了为什么“拒绝反向更新”能带来性能提升，为优化器设计提供了新的理论依据。
• 实际影响：对于大规模模型训练，微小的收敛加速和稳定性提升都能转化为巨大的算力节省和模型性能提升。
• 未来方向：作者建议将此方法应用于强化学习（RLHF）、连续学习，以及探索在特征空间（Eigenspace）而非参数空间进行掩码的可能性。

总结：这篇论文提出了一种极其简单但理论扎实且效果显著的优化器改进方案。通过引入“谨慎”机制（仅在更新方向与梯度一致时更新），它在不增加额外超参数调整成本的前提下，显著提升了 AdamW 和 Lion 等主流优化器在 LLM 预训练和图像分类任务中的性能与稳定性。