SignMuon：通信高效的分佈式Muon優化
SignMuon是一種結合了signSGD的多數表決符號聚合與Muon極座標步驟的1位、矩陣感知優化器。每個工作節點通過牛頓-舒爾茨迭代計算動量極因子，傳輸元素級符號並由多數表決聚合。在頻譜範數平滑和有界方差隨機梯度下，實現了O(1/√T)的非凸收斂率。在330個CIFAR-10/ResNet-50配置上，SignMuon取得了最佳驗證準確率92.15%；其4-GPU多數表決變體在匹配有效批次下訓練時間減少37%。在nanoGPT上，SignMuon實現了更低的困惑度和更好的即時性能。
近日，一篇題為《SignMuon: Communication-Efficient Distributed Muon Optimization》的論文在arXiv上發佈，提出了一種名為SignMuon的新型優化器，旨在解決大規模神經網絡分佈式訓練中的通信瓶頸問題。該優化器融合了signSGD的多數表決符號聚合機制和Muon優化器的極座標更新框架，實現了1比特的通信開銷，同時保持了對權重張量矩陣結構的感知能力。
在分佈式訓練中，每個工作節點首先利用牛頓-舒爾茨迭代計算動量的極因子，從而獲得類Muon的更新方向；然後僅傳輸每個元素的符號，並通過多數表決進行聚合。此外，還可以選擇在本地執行額外的極座標步驟，以進一步增強正交性，且不增加通信成本。
理論分析表明，在頻譜範數平滑和有界方差隨機梯度的條件下，該算法對於基於ℓ_1的平穩性度量能夠達到O(1/√T)的非凸收斂速率。當噪聲滿足單峯對稱分佈時，M個工作節點的多數表決可將隨機項降低1/√M，與signSGD的理論優勢一致。在α-β模型中，分佈式SignMuon每次迭代僅需一次整數和規約操作，所有正交化步驟均在本地完成，相比於float32實現了32倍的帶寬減少（相對於int8為4倍）。
實驗部分，研究者對330個CIFAR-10/ResNet-50配置進行了測試，SignMuon取得了最佳驗證準確率92.15%；其4-GPU多數表決變體在匹配有效批次的情況下，達到92.02%的準確率，同時訓練時間減少了37%。在nanoGPT語言模型上，SignMuon相比其他基於符號的基線方法取得了更低的困惑度和更優的即時性能，並在弱擴展實驗中表現出色，最高可支持16個GPU。
該研究為分佈式深度學習優化提供了新的思路，通過符號通信和矩陣感知更新顯著降低了帶寬需求，有望在大規模模型訓練中得到廣泛應用。