🤖 machine learning

SignMuon: Communication-Efficient Distributed Muon Optimization

SignMuon 是一种通信高效的 1 比特分布式优化器，它将 Muon 的矩阵感知极步框架与 signSGD 的多数投票聚合相结合，从而在大规模神经网络训练中实现最先进的精度并显著降低带宽需求。

原作者： Neel Mishra, Kushagara Trivedi, Pawan Kumar

发布于 2026-05-19

📖 1 分钟阅读☕ 轻松阅读

原作者： Neel Mishra, Kushagara Trivedi, Pawan Kumar

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

想象一下，你正在训练一支庞大的机器人团队去玩一个复杂的游戏，比如解谜或在照片中识别猫。为了变得更好，它们需要在每一轮之后分享所学到的内容。

在人工智能领域，这种“学习”是通过一种称为分布式训练的过程实现的。你拥有许多协同工作的计算机（工作节点）。问题在于，它们必须来回发送海量数据以保持同步。这就像试图协调一个合唱团，其中每位歌手在唱出每一个音符后，都必须向其他人吼叫一份长达 10 页的剧本。吼叫所花费的时间挤占了唱歌的时间，导致一切变慢。

本文介绍了一种名为Sign-Muon的新方法，旨在解决这一瓶颈。其工作原理可分解为以下简单概念：

1. 问题：过多的 chatter（ chatter 指无意义的闲聊或数据交换）

通常情况下，当这些计算机分享它们的进展时，它们会发送完整的、高精度的数字（例如"0.004321"）。这就像发送一份详尽的 100 页报告。这需要大量的带宽（网络速度）和时间。

2. 第一个想法：“是/否”的吼叫（SignSGD）

一些研究人员此前曾提出一种捷径：与其发送完整的数字，不如只发送符号。误差是上升了还是下降了？只需发送“正”（+）或“负”（-）。

类比：与其吼叫整份报告，不如只吼叫“向上！”或“向下！”
优势：这将消息大小减少了 32 倍（从 32 位数字减少到 1 位符号）。这就像发送一封信而不是整本书。
局限：如果你只吼叫“向上”或“向下”，你就会丢失信息的形状。这就像试图仅凭“左”和“右”来导航城市，却不知道要走多远或街道的布局。

3. 第二个想法：“完美转向”（Muon）

另一组研究人员开发了一种名为Muon的优化器。它将数据视为三维对象（矩阵），而不是扁平的列表。

类比：想象数据是一个旋转的陀螺。Muon 不仅关注速度，还关注旋转的轴。它使用一种数学技巧（称为极分解）来确保团队以最高效、最“正交”的方向移动，就像舞者与音乐完美同步地移动一样。
局限：进行这种数学运算是繁重的，而且如果你尝试用“是/否”的吼叫方法来做，你就会失去让舞蹈看起来优美所需的精度。

4. 解决方案：Sign-Muon（集两者之长）

作者将这两个想法结合成了Sign-Muon。以下是他们使用的逐步过程：

本地思考：每台计算机在本地进行繁重的数学运算。它为自己计算出完美的“舞步”（Muon 方向）。
吼叫：它不发送复杂的数学结果，而只发送该动作的符号（上/下，左/右）。
投票：所有计算机聚集在一起并进行多数投票。如果 10 台计算机中有 6 台说“向上”，团队就“向上”移动。这抵消了单个计算机的噪声和误差。
结果：团队移动的方向既高效（得益于本地 Muon 数学），又通信成本低廉（因为它们只发送了 1 位符号）。

为什么这很重要？

该论文声称这种方法是一个“双赢”：

速度：因为它们只发送 1 位符号，通信速度比发送完整数字快 32 倍。
质量：因为他们在缩小消息之前使用了 Muon 数学，所以没有丢失数据的“形状”。他们仍然朝着最明智的方向移动。
证据：他们在图像识别（CIFAR-10）和语言模型（nanoGPT）上测试了这一点。
- 在图像方面，与其他方法相比，他们获得了最高的准确率（92.15%）。
- 在使用多台计算机时，他们的训练速度快了 37%。
- 在语言模型方面，他们取得了比其他基于符号的方法更好的结果（更低的“困惑度”）。

总结

Sign-Muon 就像一群探险家，他们同意只互相发送简单的指南针方向（北/南）以节省能量，但在发送方向之前，他们每个人都使用高科技地图来找出完美的北/南路径。结果是一个移动速度极快、通信极少，但到达目的地时比那些吼叫完整报告或只是猜测方向的团队更聪明的团队。

技术摘要：SignMuon——通信高效的分布式 Muon 优化

1. 问题陈述

大型神经网络（特别是大语言模型，LLMs）的分布式训练在通信带宽和延迟方面面临关键瓶颈。在同步数据并行训练中，工作节点必须持续同步大量数据（梯度、更新量或优化器状态）。现有方法存在两个主要局限：

全精度通信：标准优化器传输全精度（例如 float32）梯度或更新量，这在带宽方面代价高昂。
对矩阵结构的逐坐标忽视：许多流行的通信高效方法（如 signSGD）将权重张量视为一维向量。它们忽略了参数固有的矩阵结构（例如 MLP 权重、注意力投影），从而可能错失源自矩阵几何的优化收益。

现有解决方案通常仅解决上述问题之一：基于符号的方法实现了极致的压缩，但缺乏矩阵感知能力；而像 Muon 这样具有矩阵感知能力的优化器改善了收敛几何，却保留了全精度同步的高通信成本。

2. 方法：Sign-Muon

作者提出了Sign-Muon，这是一种混合优化器，它结合了基于符号方法的通信效率与 Muon 的矩阵感知几何特性。

核心机制

Sign-Muon 通过结合两个不同的组件来运行：

Muon 风格的极化方向：每个工作节点不使用原始梯度，而是计算动量矩阵，并通过极分解（polar decomposition）推导出更新方向。这一步骤在本地通过牛顿 - 舒尔茨（Newton–Schulz, NS）迭代进行近似，使动量矩阵正交化。该步骤确保更新方向尊重权重张量的谱范数几何特性。
1 比特符号聚合：工作节点不传输全精度的极化方向，而是仅传输计算方向的逐元素符号（ $\{-1, +1\}$ ）。
多数投票聚合：在分布式设置中，工作节点通过多数投票聚合这些符号矩阵。这被高效地实现为整数SUM all-reduce（跨工作节点对符号求和），随后进行本地阈值操作（ $\text{sign}(\sum S^{(m)})$ ）。

算法流程（分布式）

本地计算：每个工作节点 $m$ 计算随机梯度 $G^{(m)}_t$ ，更新本地动量 $M^{(m)}_{t+1}$ ，并通过牛顿 - 舒尔茨迭代计算极化因子 $U^{(m)}_t$ 。
符号提取：工作节点提取逐元素符号 $S^{(m)}_t = \text{sign}(U^{(m)}_t)$ 。
通信：工作节点执行单次集体通信步骤：
- All-Reduce 变体：整数 SUM all-reduce 聚合符号。
- All-Gather 变体（1 比特）：符号被打包为比特（每个条目 1 比特），并通过 All-Gather 分发打包后的缓冲区。
本地更新：每个工作节点在本地计算多数投票 $\bar{S}_t$ 并更新参数： $W_{t+1} = W_t - \eta_t \bar{S}_t$ 。
可选本地正交化：工作节点可选择对聚合后的符号矩阵应用本地极化步骤，以在不产生额外通信的情况下进一步增强正交性。

通信效率

负载大小：该方法将通信量减少至每个参数 1 比特（若为实施简便使用 int8，则为 8 比特）。
缩减因子：与 float32 all-reduce 相比，Sign-Muon 实现了32 倍的带宽缩减（使用比特打包）或4 倍缩减（使用 int8）。
本地计算：所有谱范数归一化和正交化（牛顿 - 舒尔茨）均在本地执行，不产生额外的网络通信成本。

3. 理论贡献

本文在谱范数平滑性和有界方差随机梯度的假设下提供了收敛性分析。

收敛速率：对于基于 $\ell_1$ 的平稳性度量，Sign-Muon 实现了 $O(1/\sqrt{T})$ 的非凸收敛速率。
噪声降低：在单峰对称噪声的假设下，跨 $M$ 个工作节点的多数投票机制将随机噪声项降低了 $1/\sqrt{M}$ 倍，这与分布式 signSGD 的理论收益相匹配。
几何特性：分析利用了谱范数下降引理，这使其区别于依赖逐坐标平滑性的标准 signSGD 分析。谱归一化确保更新方向的谱范数至多为 1。

4. 实证结果

作者在两个工作负载上评估了 Sign-Muon：CIFAR-10（使用 ResNet 架构的图像分类）和nanoGPT（语言建模）。

CIFAR-10 (ResNet-50)

准确率：在 330 种超参数配置下，Sign-Muon 达到了92.15% 的最佳验证准确率。
分布式效率：4-GPU 多数投票变体达到了92.02%的准确率（与单 GPU 性能持平），同时在匹配有效批大小的情况下将训练时间减少了37%。
扩展性：与 SignAdam 和 SignSGD 相比，Sign-Muon 表现出更优越的弱扩展行为。在 16 个 GPU 上，其训练时间开销显著更低（+77.8%，而 All-Gather 设置下的基线为 +150–163%）。

nanoGPT

性能：与其他基于符号的基线相比，Sign-Muon 实现了更低的困惑度和更好的“任意时刻”性能（随挂钟时间运行的最小困惑度）。
扩展性：观察到在高达 16 个 GPU 上的有利弱扩展性。
超参数：该方法对归一化方案（Frobenius 与谱缩放）表现出鲁棒性，且仅需微调（例如，通常 1 次牛顿 - 舒尔茨迭代就足够了）。

内存与开销

内存：Sign-Muon 保持与 SGD 相似的持久内存占用（一个动量缓冲区），显著低于维护两个动量向量的 Adam/AdamW。
计算：本地牛顿 - 舒尔茨迭代增加了每一步的计算开销，但这被通信时间的巨大减少所抵消，特别是在带宽受限的环境中。

5. 意义与主张

本文主张 Sign-Muon 成功弥合了此前两个截然不同的研究路线之间的差距：

通信效率：它保留了基于符号方法的极致压缩（1 比特）和鲁棒性（多数投票）。
矩阵感知：它融入了 Muon 的几何优势（极分解），同时不牺牲通信效率。

关键主张：

弥合差距：Sign-Muon 是首个将多数投票符号聚合直接嵌入 Muon 极化步骤框架的优化器。
可扩展性：它为带宽成为主要瓶颈的大规模分布式训练提供了实用解决方案，与 float32 相比将通信量减少了高达 32 倍。
理论严谨性：该方法在谱范数平滑性下提供了严格的收敛保证，这是一种比逐坐标平滑性更适用于矩阵结构化更新的度量。
实用性：在标准基准（CIFAR-10, nanoGPT）上的实证结果表明，Sign-Muon 不仅匹配而且往往超越了全精度及其他基于符号的基线的性能，同时在分布式设置中显著减少了训练时间。

作者强调，当网络带宽成为瓶颈时，该方法尤为有效，因为正交化的本地计算在通信成本方面是“免费”的。