⚡ electrical engineering

DeMuon: A Decentralized Muon for Matrix Optimization over Graphs

This paper introduces DeMuon, the first decentralized extension of the Muon optimizer that combines Newton-Schulz orthogonalization with gradient tracking to achieve provable convergence and superior performance in matrix optimization over communication graphs, particularly under heavy-tailed noise conditions.

Original authors: Chuan He, Shuyi Ren, Jingwei Mao, Erik G. Larsson

Published 2026-06-03

📖 5 min read🧠 Deep dive

CC BY 4.0

Original authors: Chuan He, Shuyi Ren, Jingwei Mao, Erik G. Larsson

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Imagine a group of friends trying to solve a giant, complex puzzle together. They are all in different rooms (decentralized), and they can only talk to their immediate neighbors. They don't have a boss or a central leader telling them what to do; they have to figure it out by sharing what they see and adjusting their own pieces based on what their neighbors tell them.

This paper introduces a new way for these friends to solve the puzzle faster and more accurately. They call their new method DeMuon.

Here is the breakdown of how it works, using simple analogies:

The Problem: The "Matrix" Puzzle

In the world of artificial intelligence (specifically deep learning), the "puzzle pieces" aren't just single numbers; they are giant grids of numbers called matrices.

Old Way (Vectorizing): Traditionally, computers treated these big grids like long, flat lists of numbers (vectors). It's like trying to solve a 3D jigsaw puzzle by flattening it into a 2D sheet first. It works, but it's clumsy and misses the shape of the pieces.
The New Way (Muon): A recent method called Muon realized that treating the pieces as their natural 3D shape (matrices) is much better. It uses a special "compass" (called the spectral norm) to decide which direction to move the pieces. This works incredibly well when everyone is in the same room (centralized).

The Challenge: Going Decentralized

The authors asked: Can we use this smart "Muon" compass when our friends are in different rooms and can't talk to a central boss?
This is hard because:

Different Views: Each friend sees a slightly different part of the puzzle (local data).
No Boss: They can't just ask a leader, "What's the best move?" They have to guess the "global" best move by listening to neighbors.
Confusion: If they aren't careful, they might all start moving in different directions, and the puzzle will never get solved.

The Solution: DeMuon

The paper proposes DeMuon, a method that lets the friends solve the puzzle together without a boss. It combines two main tricks:

1. The "Shared Compass" (Gradient Tracking)
Imagine each friend has a compass. Because they are in different rooms, their compasses point in slightly different directions.

Old Decentralized Methods: Friends would just point their compasses at their neighbors and hope they align.
DeMuon's Trick: They use a technique called gradient tracking. It's like a relay race where each friend doesn't just pass their current direction, but also passes a "correction note" about how their direction changed since the last step. This helps the whole group agree on the true global direction, even if they are far apart.

2. The "Matrix Orthogonalization" (The Muon Magic)
When a friend decides to move their puzzle piece, they don't just push it randomly. They use the Muon technique, which is like a specialized "shape-shifter."

Instead of just pushing the piece forward, Muon checks the "shape" of the piece (using the spectral norm) and rotates it to be perfectly aligned before moving.
Think of it like a dancer who doesn't just walk forward, but first strikes a perfect pose to ensure they are balanced. This prevents the puzzle pieces from getting "stuck" or moving inefficiently.

The Super-Charged Version: DeMuon-A

The authors also created an even faster version called DeMuon-A.

The Analogy: If DeMuon is a runner who looks at the ground and takes a step, DeMuon-A is a runner who looks at the ground, predicts where they will be in two steps, and then takes a giant leap based on that prediction.
How it works: It uses a technique called multi-extrapolation. It asks, "If I keep moving this way, where will I be?" and uses that prediction to take a bigger, smarter step. This requires the puzzle to be "smooth" (predictable), but when it works, it converges to the solution much faster.

What Did They Prove?

The authors did two main things:

Math Proof: They used advanced math to prove that if the friends follow these rules, they will eventually agree on the solution (consensus) and find the best possible arrangement of the puzzle pieces (stationarity). They proved this works even if the friends are in a messy network (some can talk to many, some to few).
Real-World Test: They tested this on training a language model (a type of AI that writes text).
- They set up 8 computers (nodes) connected in different ways (like a perfect circle, a ring, or a messy web).
- Result: DeMuon and DeMuon-A learned the language task much better and faster than the standard methods (like DSGD). They reached a lower "error" score, meaning the AI was smarter.

Summary

DeMuon is a new way for a group of computers to train AI models together without a central boss.
It keeps the "smart shape-shifting" (matrix optimization) of the original Muon method.
It adds a "relay race" system (gradient tracking) so everyone agrees on the direction.
DeMuon-A adds a "prediction leap" to make it even faster.
The paper proves mathematically that it works and shows through experiments that it beats current methods in speed and accuracy.

The paper does not claim this is for medical use or specific future applications; it is strictly about improving the mathematical efficiency of training AI models in a decentralized network.

Technical Summary: DeMuon - A Decentralized Muon for Matrix Optimization over Graphs

Problem Statement

This paper addresses the problem of fully decentralized matrix optimization over communication graphs. The objective is to solve the multi-agent optimization problem:
$\min_{X \in \mathbb{R}^{m \times n}} f(X) := \frac{1}{N} \sum_{i=1}^N f_i(X)$
where $N$ agents are connected via a directed graph $\mathcal{G}=(\mathcal{V}, \mathcal{E})$ , and each node $i$ holds a local objective function $f_i$ . Unlike traditional decentralized optimization methods that often vectorize matrix variables, this work focuses on preserving the matrix structure, specifically leveraging techniques from the centralized Muon algorithm (Jordan et al., 2024). The goal is to achieve global convergence and consensus among agents without a central server, while mitigating heterogeneity in local objective functions.

Methodology

The paper proposes two algorithms: DeMuon (a decentralized variant of Muon) and DeMuon-A (an accelerated variant).

1. DeMuon: Decentralized Muon

DeMuon integrates spectral-norm-based updates into a gradient-tracking framework. The algorithm maintains three sequences: local gradient estimators $\{M^{(k)}\}$ , global gradient estimators $\{V^{(k)}\}$ , and local iterates $\{X^{(k)}\}$ .

Local Momentum: Each node updates its local estimator $M^{(k)}_i$ using an exponentially weighted moving average of stochastic gradients.
Gradient Tracking: Nodes update $V^{(k)}_i$ by aggregating information from neighbors to track the global gradient, ensuring consensus on the gradient direction despite local heterogeneity.
Matrix Orthogonalization: The local iterate $X^{(k+1)}_i$ is updated by aggregating neighbors' updates, specifically applying a matrix sign function ( $\text{msgn}(\cdot)$ ) to the tracked gradient. This step solves a norm-constrained linear subproblem using the spectral norm ( $\|\cdot\|$ ), which is central to the Muon methodology.

Theoretical Guarantees for DeMuon:
Under assumptions of bounded noise and gradient Lipschitz continuity with respect to the spectral and nuclear norms, the authors establish:

Consensus Error: The error $\|X^{(k)} - \mathbf{1} \otimes \bar{X}^{(k)}\|$ is bounded by $O(\frac{\sqrt{N}\lambda \eta}{1-\lambda})$ , where $\lambda$ characterizes the mixing rate of the graph.
Stationarity Error: The average stationarity error converges at a rate of $O(K^{-1/4})$ with properly tuned step size $\eta$ and momentum $\theta$ . This matches the optimal convergence rate of centralized normalized SGD with momentum.

2. DeMuon-A: Accelerated Decentralized Muon

DeMuon-A extends the multi-extrapolation technique (He et al., 2026b) to the decentralized matrix setting.

Mechanism: It performs $q = p-1$ local extrapolation steps per iteration to generate points $Z^{(k,s)}_i$ , exploiting higher-order smoothness of the objective function.
Assumptions: It requires a higher-order Lipschitz condition on the $p$ -th derivative of the functions $f_i$ , induced by the spectral and nuclear norms.
Convergence: Under these conditions, DeMuon-A achieves an improved stationarity error rate of $O(K^{-p/(3p+1)})$ and a consensus error rate of $O(K^{-(2p+1)/(3p+1)})$ .

Key Contributions

Algorithmic Proposal: The authors introduce DeMuon, the first fully decentralized matrix optimization method that incorporates Muon's spectral-norm-based orthogonalization and gradient tracking.
Convergence Analysis: They provide rigorous global convergence guarantees for both DeMuon and DeMuon-A in terms of stationarity and consensus errors. Notably, the analysis relies on matrix martingale moment inequalities to establish dimension-free convergence (independent of matrix dimensions $m, n$ ).
Acceleration: The paper proposes DeMuon-A, which integrates multi-extrapolation to leverage higher-order smoothness, offering improved convergence rates over the base DeMuon algorithm.
Empirical Validation: Numerical experiments on distributed deep learning (Transformer training) and distributed matrix regression demonstrate the practical advantages of the proposed methods over existing decentralized baselines.

Experimental Results

The authors evaluated DeMuon against three baselines: DSGD (vanilla decentralized SGD), DSGD-C (clipped SGD), and DSGD-N (normalized SGD with gradient tracking).

Setup: Experiments were conducted on a 7.4M-parameter GPT model trained on the Multi30k dataset across three graph topologies: Complete, Directed Exponential, and Ring graphs.
Performance: DeMuon consistently outperformed all baselines in terms of training loss, validation loss, and perplexity across all topologies.
Consensus: The consensus error for DeMuon remained competitive with or lower than the baselines, indicating effective agreement among agents.
Specifics: In the Complete Graph topology, DeMuon achieved a final validation loss of 0.240 (perplexity 1.27), significantly lower than DSGD-N (0.302) and DSGD (7.031).

Significance and Claims

The paper claims that DeMuon successfully bridges the gap between the empirical success of Muon in centralized deep learning and the requirements of decentralized optimization.

Dimension-Free Convergence: Unlike some recent works (e.g., Zhang et al., 2026a) that may depend explicitly on problem dimensionality, the authors claim their use of matrix martingale moment inequalities allows for convergence guarantees that do not explicitly depend on the matrix dimensions $(m, n)$ .
Practical Advantage: The results suggest that the specific structure of Muon (spectral norm constraints and matrix orthogonalization) provides clear practical advantages in decentralized settings, particularly for training large language models where matrix structures are inherent.
Theoretical Rigor: The work establishes that the benefits of momentum and normalization in matrix optimization can be preserved in a fully decentralized environment without a central coordinator, provided that gradient tracking and appropriate norm-based updates are utilized.

The authors position this work as a step toward bringing advanced matrix optimization techniques to the decentralized world, motivated by applications in distributed deep learning and classical problems like distributed PCA and matrix completion.