新上线2天前0 投票
SignMuon：兼具通信效率与矩阵感知的分布式优化器
背景：分布式训练的通信瓶颈
训练大规模神经网络时，全精度梯度的通信开销是主要瓶颈之一。传统的分布式优化器（如 Adam、SGD）通常按坐标独立更新，忽略了权重张量的矩阵结构，导致优化效率受限。尽管 signSGD 通过 1-bit 梯度量化大幅降低了通信量，但其逐坐标处理方式仍未利用矩阵的几何信息。
SignMuon 核心设计
来自印度理工学院等机构的研究者提出了 SignMuon，一种结合了 Muon 优化器矩阵感知能力与 signSGD 低比特通信优势的新型优化器。关键创新包括：
- Muon 风格方向：每个 worker 通过 Newton–Schulz 迭代计算动量矩阵的极分解因子，得到正交化的更新方向。
- 1-bit 符号通信：仅传输更新矩阵的逐元素符号（1-bit），并通过多数投票（majority vote）进行聚合，大幅降低通信带宽。
- 可选本地极分解：在本地额外执行一步极分解，进一步强化正交性，且不增加通信成本。
理论保证与通信效率
在谱范数光滑性和有界方差假设下，SignMuon 对于非凸优化达到了 O(1/√T) 的收敛率（基于 ℓ1 平稳度量）。当噪声为单峰对称分布时，多数投票机制可将随机项降低 1/√M（M 为 worker 数），与 signSGD 一致。
在 α-β 通信模型中，分布式 SignMuon 每轮只需一次整数 sum-allreduce 操作，所有正交化都在本地完成。相比 float32，带宽降低 32 倍；即使对比 int8，也降低 4 倍。
实验表现：CIFAR-10 与 nanoGPT
- CIFAR-10 / ResNet-50：在 330 组超参数配置中，SignMuon 取得了最佳验证准确率 92.15%。其 4-GPU 多数投票变体达到 92.02%，并且在匹配有效批量时，训练时间减少 37%。
- nanoGPT：SignMuon 实现了更低的困惑度，并在任意时刻性能上优于其他基于符号的基线方法。弱扩展性测试显示，在 16 GPU 范围内性能良好。
意义与展望
SignMuon 为分布式深度学习提供了一种兼具通信效率和优化质量的实用方案。它证明了将矩阵感知优化与 1-bit 通信结合是可行的，并且在大规模训练场景中具有显著优势。未来工作可探索将其扩展到更复杂的模型架构，或与其他压缩技术（如 top-k 稀疏化）协同使用。