DeMo：解耦动量优化
机器学习
2026-02-10 v2 人工智能
摘要
同步数据并行是神经网络训练的关键趋势，但全精度梯度 all-reduce 带来了严重的通信瓶颈。我们提出一种称为 DeMo（Decoupled Momentum Optimization）的解耦动量优化方法，作为基于动量的优化器的即插即用替代方案，可显著降低通信带宽 while 保持收敛性。DeMo 的核心设计包括：(i) 解耦本地动量更新；(ii) 应用快速正交变换（如 DCT）后进行 top-k 稀疏化；(iii) 通过动量减法将动量缓冲区用作误差反馈。该设计在计算开销最小的前提下，可将每步通信量降低最高可达两个数量级。在 300M 和 1B 参数的 DeMo 语言模型实验中，DeMo 每个 GPU 的传输数据量较 AdamW-DDP 低最高可达 85 倍，同时实现了可比的损失和准确率。DeMo 具备拓扑agnostic 特性，能够支持跨多数据中心或以太网环境下的训练。代码已公开于 https://github.com/bloc97/DeMo
引用
@article{arxiv.2411.19870,
title = {DeMo: Decoupled Momentum Optimization},
author = {Bowen Peng and Lizhang Chen and Baiyu Su and Jeffrey Quesnelle and Diederik P. Kingma and Qiang Liu},
journal= {arXiv preprint arXiv:2411.19870},
year = {2026}
}