arXiv:2510.15830cs.LGcs.AI2025-10被引 6

给伪梯度加Nesterov动量,让优化器提速1.5到2.5倍。

SNOO: Step-K Nesterov Outer Optimizer - The Surprising Effectiveness of Nesterov Momentum Applied to Pseudo-Gradients

  • 在快慢权重框架中对伪梯度使用Nesterov动量
  • 在1e23 FLOPs规模下实现1.5-2.5倍算力提升
  • 适合大模型训练,兼容AdamW和Muon等优化器

大语言模型的快速发展推动了对更高效优化技术的需求。其中,Lookahead类优化器采用双循环框架,维护快慢两组模型参数。在快参数上执行多步内优化,生成轨迹即伪梯度,用于更新慢参数。DiLoCo是典型代表,最初为分布式训练设计,将各工作节点的平均伪梯度引入Nesterov动量,声称在非分布式设置下甚至优于AdamW。本文通过实证发现,DiLoCo的出色表现主要源于对伪梯度应用Nesterov动量,该机制在非分布式场景中显著提升训练效率。我们称此变体为步长-K Nesterov外层优化器(SNOO)。实验表明,SNOO在不超过1e23训练FLOPs的规模下,可实现1.5–2.5×的计算因子提升,且性能随模型增大而增强。由于其极低的计算与内存开销,以及支持模型分片,SNOO是对多种内优化器(如AdamW、Muon)的实用增强。

原文摘要 · Abstract (English)

The rapid development of large language models (LLMs) has driven the demand for more efficient optimization techniques. Among these, the Lookahead family of optimizers employs a two-loop framework, maintaining fast and slow sets of model weights. Multiple inner optimizer steps on the fast weights produce a trajectory - the pseudo-gradient - that is used to update the slow weights. DiLoCo, a notable example originally designed for distributed training, applies Nesterov momentum to the averaged pseudo-gradient from multiple workers, claiming to even outperform AdamW in a non-distributed setup. In this paper, we empirically show that DiLoCo's surprising effectiveness stems primarily from applying Nesterov momentum to the pseudo-gradient, which improves training in a non-distributed setting. We call this Lookahead variant the Step-$K$ Nesterov Outer Optimizer (SNOO). We demonstrate that SNOO achieves compute factor gains of 1.5 - 2.5$\times$ in a non-distributed setting up to a scale of 1e23 training FLOPs, with improvements that increase with model size. Because of its minimal compute and memory overhead and compatibility with model sharding, SNOO is a practical enhancement for a variety of inner optimizers, including AdamW and Muon.

优化器大模型Nesterov加速

Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。