FedMuon: Matrix-Aware Federated Optimization
- FedMuon is a federated learning framework that employs matrix-aware, orthonormalized updates derived from the Muon optimizer to handle client drift and noisy gradients.
- It introduces diverse mechanisms such as periodic aggregation, bias correction through control variates, and local-global alignment to mitigate heterogeneity across clients.
- Rigorous convergence analyses show that decoupling step size from gradient magnitude enables competitive performance under nonconvex objectives and heavy-tailed noise.
FedMuon is the name used in several 2025 arXiv works for federated optimization algorithms derived from Muon, a matrix-aware optimizer that replaces raw gradient or momentum directions with orthonormalized, or equivalently linear-minimization-oracle-based, update directions. Across these papers, the common objective is to import Muon’s matrix geometry into federated learning while controlling client drift, communication cost, and stochastic noise; the resulting methods differ materially in how they correct heterogeneity, how they aggregate optimizer state, and what guarantees they prove for nonconvex or minimax objectives (Zhang et al., 4 Oct 2025).
1. Muon foundations and the matrix-orthogonalization viewpoint
Muon is formulated for matrix-shaped parameters. In one line of analysis, each parameter matrix has a momentum matrix , which is converted into the closest orthonormal matrix
If is the SVD, then , and the update is
This orthonormalization fixes the scale of the update direction: has orthonormal columns, , and . As a consequence, the per-step movement obeys , independently of gradient magnitude (Zhang et al., 4 Oct 2025).
A second line of analysis expresses Muon as a linear minimization oracle (LMO) method over the spectral-norm unit ball. For a matrix momentum 0, the exact Muon LMO is
1
This makes Muon a steepest-descent method in a non-Euclidean geometry: the update direction is determined by singular vectors rather than by raw element-wise scaling. The exact SVD is typically replaced in practice by Newton–Schulz iterations, which preserve the singular vectors while progressively transforming the singular values toward the polar factor (Takezawa et al., 30 Sep 2025).
The matrix-orthogonalization interpretation and the LMO interpretation are mathematically aligned. Both imply that Muon uses bounded-norm, matrix-structured directions, and both motivate its federated variants: the optimizer is attractive because deep networks are dominated by matrix-valued layers, yet federated optimization has historically relied mainly on element-wise local optimizers such as SGD, Adam, and AdamW (Liu et al., 31 Oct 2025).
2. The term “FedMuon” and its distinct algorithmic meanings
The 2025 literature uses “FedMuon” for several non-identical federated algorithms. The overlap lies in the use of Muon-style local updates; the divergence lies in how each method addresses the interaction between matrix orthogonalization and federated heterogeneity.
| Variant | Defining mechanism | Primary setting |
|---|---|---|
| FedMuon (Zhang et al., 4 Oct 2025) | Local Muon with periodic averaging of parameters and momentum matrices | Nonconvex FL; problem-parameter-free hyperparameters; heavy-tailed noise |
| FedMuon (Takezawa et al., 30 Sep 2025) | Bias-corrected LMO updates with SCAFFOLD-like control variates | Nonconvex FL with exact or inexact LMO |
| FedMuon (Liu et al., 31 Oct 2025) | Momentum aggregation and local-global alignment | Communication-efficient FL under IID and non-IID data |
| FedMuon-DA (Zhang et al., 6 Nov 2025) | Muon-style primal descent and dual ascent with control variates and distributed averaging | Federated nonconvex-PL minimax optimization under heavy-tailed noise |
One formulation treats FedMuon as the direct federated analogue of LocalSGD or FedAvg, but with Muon replacing the local SGDM step. Each client performs several local Muon updates, and every 2 iterations the server averages both the model and the momentum matrices before broadcasting them back to all workers (Zhang et al., 4 Oct 2025).
A second formulation argues that this direct substitution is insufficient because the LMO is a biased operator in the federated setting: 3 That paper therefore defines FedMuon through bias correction using control variates, so that each client applies the LMO not to its raw local momentum but to a corrected direction intended to approximate the global gradient (Takezawa et al., 30 Sep 2025).
A third formulation emphasizes client drift caused by client-specific orthogonalization in non-IID FL. It introduces two specific mechanisms: momentum aggregation, where clients initialize local momentum from an aggregated global momentum, and local-global alignment, where each local Muon direction is blended with a server-side estimate of the global update direction (Liu et al., 31 Oct 2025).
A fourth formulation, FedMuon-DA, extends the idea beyond minimization to federated minimax optimization. It uses Muon-style orthonormalized descent on the primal matrix 4 and orthonormalized ascent on the dual matrix 5, combined with momentum and control variates under a nonconvex-PL structure (Zhang et al., 6 Nov 2025).
3. Core algorithmic mechanisms
Despite the naming collision, several recurring mechanisms define the FedMuon literature. The first is local matrix-valued momentum. Clients do not maintain scalar or vector preconditioners alone; they maintain momentum matrices whose singular structure determines the update direction. This is central both in the orthonormalized-update formulation and in the LMO formulation.
The second is periodic server aggregation, but the aggregated state differs by paper. In the periodic-averaging FedMuon, each worker sends both 6 and 7 every 8 iterations, and the server averages both: 9 This preserves a shared matrix momentum across communication rounds without introducing extra control variates (Zhang et al., 4 Oct 2025).
The third is explicit bias correction. In the LMO-based FedMuon, each client maintains a local momentum 0, a client control variate 1, and receives a global control variate 2. The local direction is formed from
3
after which the client applies 4. The stated purpose is to ensure that, when local drift is small, each client’s LMO sees a proxy for the global gradient rather than a purely local direction (Takezawa et al., 30 Sep 2025).
The fourth is local-global alignment. In the alignment-based FedMuon, the client update is
5
where 6 is the local Muon direction and
7
is a server-side estimate of the global direction. In the same method, momentum aggregation replaces per-round momentum reinitialization: 8 The resulting server state couples clients through both a shared direction estimate and a shared momentum (Liu et al., 31 Oct 2025).
The fifth is minimax distributed averaging. FedMuon-DA maintains global variables 9, global momenta 0, and control variates 1. Clients perform local Muon-style updates
2
and the server updates the global model by scaled averaging of local endpoints. This is the only variant in the supplied literature that is explicitly formulated for federated minimax optimization rather than plain minimization (Zhang et al., 6 Nov 2025).
4. Convergence theory and the role of orthonormalized directions
The strongest shared theoretical theme is that orthonormalized update directions decouple step length from raw gradient magnitude. In the periodic-averaging FedMuon, this yields hyperparameter schedules for 3, 4, and 5 that depend only on the number of workers and the time horizon, not on problem-specific constants such as the Lipschitz constant 6, the stochastic noise level 7, or the heterogeneity parameter 8. Under bounded-variance noise, the dominant nonconvex rate is
9
with 0 iterations and communication complexity 1 to reach an 2-stationary point. Under heavy-tailed noise with moment index 3, the dominant term becomes
4
and this guarantee is obtained without gradient clipping (Zhang et al., 4 Oct 2025).
The bias-corrected LMO paper begins with a negative result. It proves that naive LocalMuon can fail to move at all: there exist convex losses for which the global parameter remains at initialization for every round, and the gradient norm remains bounded away from zero. The mechanism is cancellation of per-client LMO directions under aggregation. FedMuon in that paper repairs the issue with control variates and proves nonconvex convergence for exact LMOs, with a dominant term matching the 5 order familiar from FedAvg and SCAFFOLD analyses, up to norm equivalence factors. It also analyzes inexact Muon via Newton–Schulz iterations and shows that convergence holds for any number of Newton–Schulz steps. When the number of iterations is 6, the result controls a Frobenius-norm quantity; as the number of iterations increases, the bound interpolates toward the trace norm through a Schatten-7 parameter that decreases from 8 toward 9 (Takezawa et al., 30 Sep 2025).
The alignment-based FedMuon develops a different guarantee. Its final rate is
0
where 1 is the number of participating clients, 2 is the number of local steps, 3 is the number of communication rounds, and 4. The paper explicitly contrasts this with local-SGD-type bounds that contain a heterogeneity term 5; its claim is that local-global alignment removes dependence on 6 from the final rate (Liu et al., 31 Oct 2025).
FedMuon-DA extends heavy-tail analysis to federated minimax optimization under a PL condition in the maximization variable. With 7 clients, 8 local steps, and 9 communication rounds, the reported rate is
0
where 1 is the heavy-tail moment index. The analysis does not assume bounded heterogeneity and frames Muon’s orthonormalization as a matrix-valued analogue of normalized-gradient methods (Zhang et al., 6 Nov 2025).
5. Empirical behavior across vision, language, and minimax tasks
The periodic-averaging FedMuon is evaluated with 2 workers under synchronous full participation on image classification and text classification. The tasks are CIFAR-10 and CIFAR-100 with ResNet-18 and ViT, and Sentiment140 with an RNN. Homogeneous splits are used for main experiments, and a heterogeneous Dirichlet(3) split is used for non-IID CIFAR-10. The reported findings are that FedMuon shows faster training loss decrease and higher test accuracy than all baselines on ResNet-18, that the advantage becomes more pronounced when the communication period increases from 4 to 5, that it remains competitive or superior under non-IID data without control variates, and that the gains are even larger on ViT. On Sentiment140, where the paper characterizes the regime as heavy-tailed, FedMuon is reported to achieve consistently better training and test performance than both vanilla and adaptive baselines, despite not using clipping (Zhang et al., 4 Oct 2025).
The bias-corrected LMO-based FedMuon is evaluated on FashionMNIST with LeNet and CIFAR-10 with ResNet-18 plus GroupNorm, with 6 total clients, 7 sampled clients, and 8 local steps per round. Data heterogeneity is controlled by Dirichlet partitions with 9 and 0. The paper reports that FedMuon achieves the highest test accuracy across all settings. It also reports a targeted comparison with LocalMuon: under nearly homogeneous data LocalMuon is competitive but slightly worse, whereas under heterogeneous data it significantly underperforms FedMuon, matching the theoretical claim that naive local LMOs are problematic under heterogeneity. An ablation on Newton–Schulz accuracy finds that even 1 trains reasonably well, that increasing 2 from 3 to 4 yields a substantial boost, and that the best settings in the reported experiment are 5 in the homogeneous case and 6 in the heterogeneous case (Takezawa et al., 30 Sep 2025).
The alignment-based FedMuon reports detailed quantitative gains on larger non-IID vision and language benchmarks. On CIFAR-100 with ResNet-18 under Dirichlet-7, the best baseline accuracies reported are 8 for FedCM and 9 for Local Muon, whereas FedMuon reaches 0 test accuracy with loss 1. Under Dirichlet-2, the best baselines are 3 for FedCM and 4 for Local Muon, while FedMuon reaches 5 with loss 6. For ViT-Tiny on the same task, FedMuon reaches 7 under Dirichlet-8 and 9 under Dirichlet-0, exceeding the cited Local Muon and Local AdamW baselines. On Swin Transformer fine-tuning, the paper reports 1 on CIFAR-100 and 2 on Tiny ImageNet, both above the listed strongest baselines. On GLUE with RoBERTa-Base and LoRA, it reports the best scores on CoLA, RTE, SST-2, QQP, MRPC, QNLI, and MNLI, with especially pronounced gains on CoLA and RTE. An ablation shows that removing momentum aggregation or global alignment hurts performance, that 3 is best among the tested local-global mixing values, and that SVD compression of momentum nearly preserves accuracy while reducing communication overhead (Liu et al., 31 Oct 2025).
FedMuon-DA is evaluated on deep AUC maximization for Sentiment140, SST2, and Yelp Review with a two-layer RNN. The data are made homogeneous or heterogeneous through class-imbalance constructions, and the baselines include LocalSGDAM, FedSGDA-M, F-SAGDA, FESS-GDA, and FedSGDA-Clip with tuned threshold 4. The paper reports that FedMuon-DA and Fed-NSGDA-M converge faster and achieve higher test AUC than the baselines in both homogeneous and heterogeneous settings, while the clipping-based baseline shows marked sensitivity to the clipping threshold (Zhang et al., 6 Nov 2025).
6. Conceptual issues, limitations, and open directions
A central conceptual issue is terminological rather than mathematical: “FedMuon” is not a single canonical algorithm in the supplied literature. One paper uses the term for periodic averaging of model and momentum with no control variates; another uses it for SCAFFOLD-like bias correction of LMOs; another uses it for momentum aggregation plus local-global alignment; and a fourth paper uses “FedMuon-DA” for a minimax extension. This suggests that references to “FedMuon” require disambiguation by paper and mechanism, especially when comparing theory or implementation.
A second issue is the status of naive LocalMuon. The literature does not support a uniform conclusion. One paper proves non-convergence examples for naive Muon inside FedAvg under heterogeneity because averaged LMOs can cancel and fail to approximate the global direction (Takezawa et al., 30 Sep 2025). Another paper reports that Local Muon significantly accelerates convergence and reduces communication rounds in IID settings, but also states that non-IID FL creates severe client drift through independent client preconditioners and moment reinitialization (Liu et al., 31 Oct 2025). The consistent reading is therefore conditional: Muon can be highly effective as a local optimizer, but federated stability depends on how heterogeneity is corrected.
The main practical limitation is orthonormalization cost. Multiple papers note that exact SVD is too expensive per step, so practical implementations use Newton–Schulz or related approximations. One paper states that five Newton–Schulz iterations add about 5 computation versus AdamW; another emphasizes that per-iteration SVD or Newton–Schulz overhead can be substantial for very large layers or resource-constrained clients (Liu et al., 31 Oct 2025). A related communication issue appears when momentum matrices are transmitted: without compression this can double per-round communication, which motivates low-rank SVD compression of momentum (Liu et al., 31 Oct 2025).
The theoretical limitations are paper-specific. The periodic-averaging FedMuon still assumes global smoothness and a standard heterogeneity bound in its minimization theory, even though its recommended hyperparameters are independent of those constants; its experiments use 6 with full participation, and it does not analyze privacy mechanisms (Zhang et al., 4 Oct 2025). The bias-corrected LMO FedMuon assumes smoothness, unbiased gradients, and bounded variance, and its emphasized theorem is stated for full participation in the main text even though partial participation is treated in the appendices (Takezawa et al., 30 Sep 2025). The alignment-based FedMuon presents broad empirical evidence but remains more complex to implement than FedAvg or server-side adaptive methods, and its experiments are concentrated on moderate-scale CNN, Transformer, and LoRA-tuned language setups (Liu et al., 31 Oct 2025). FedMuon-DA, finally, is restricted to nonconvex-PL minimax problems under heavy-tailed noise, so its guarantees do not directly cover fully general non-PL minimax dynamics (Zhang et al., 6 Nov 2025).
The open directions identified across the literature are relatively consistent: more efficient orthonormalization, combinations with adaptive preconditioning or layer-wise strategies, better handling of partial participation and asynchronous FL, privacy-aware variants, and extension to larger-scale language-model or cross-device deployments. A plausible implication is that the lasting contribution of FedMuon may be less the stabilization of one specific algorithm than the broader demonstration that matrix-aware update geometry can be integrated into federated optimization with rigorous convergence analyses in nonconvex, heterogeneous, and heavy-tailed regimes.