DeMuon: A Decentralized Muon for Matrix Optimization over Graphs
This work addresses the problem of efficient decentralized optimization for distributed systems, offering a provable extension to a known method, but it appears incremental as it builds directly on an existing centralized approach.
The paper tackles decentralized matrix optimization over communication graphs by proposing DeMuon, which extends a centralized method with gradient tracking to handle heterogeneity and noise, achieving iteration complexity matching centralized algorithms and showing improved performance in experiments on transformer pretraining.
In this paper, we propose DeMuon, a method for decentralized matrix optimization over a given communication topology. DeMuon incorporates matrix orthogonalization via Newton-Schulz iterations-a technique inherited from its centralized predecessor, Muon-and employs gradient tracking to mitigate heterogeneity among local functions. Under heavy-tailed noise conditions and additional mild assumptions, we establish the iteration complexity of DeMuon for reaching an approximate stochastic stationary point. This complexity result matches the best-known complexity bounds of centralized algorithms in terms of dependence on the target tolerance. To the best of our knowledge, DeMuon is the first direct extension of Muon to decentralized optimization over graphs with provable complexity guarantees. We conduct preliminary numerical experiments on decentralized transformer pretraining over graphs with varying degrees of connectivity. Our numerical results demonstrate a clear margin of improvement of DeMuon over other popular decentralized algorithms across different network topologies.