AI News HubLIVE
站內改寫1 分鐘閱讀

SignMuon:通訊高效的分散式Muon最佳化

SignMuon是一種結合了signSGD的多數表決符號聚合與Muon極座標步驟的1位、矩陣感知最佳化器。每個工作節點透過牛頓-舒爾茨迭代計算動量極因子,傳輸元素級符號並由多數表決聚合。在頻譜範數平滑和有界方差隨機梯度下,實現了O(1/√T)的非凸收斂率。在330個CIFAR-10/ResNet-50配置上,SignMuon取得了最佳驗證準確率92.15%;其4-GPU多數表決變體在匹配有效批次下訓練時間減少37%。在nanoGPT上,SignMuon實現了更低的困惑度和更好的即時效能。

來源arXiv Machine Learning作者: Neel Mishra, Kushagara Trivedi, Pawan Kumar

近日,一篇題為《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。

該研究為分散式深度學習最佳化提供了新的思路,透過符號通訊和矩陣感知更新顯著降低了頻寬需求,有望在大規模模型訓練中得到廣泛應用。