Papers
Topics
Authors
Recent
Search
2000 character limit reached

FedMuon: Matrix-Aware Federated Optimization

Updated 4 July 2026
  • 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 XtRm×nX_t \in \mathbb{R}^{m\times n} has a momentum matrix MtM_t, which is converted into the closest orthonormal matrix

Ot=argminORm×nOMtF2s.t.OO=In.O_t = \arg\min_{O \in \mathbb{R}^{m\times n}} \|O - M_t\|_F^2 \quad \text{s.t.} \quad O^\top O = I_n.

If Mt=UtStVtM_t = U_t S_t V_t^\top is the SVD, then Ot=UtVtO_t = U_t V_t^\top, and the update is

Xt+1=XtηUtVt.X_{t+1} = X_t - \eta\, U_t V_t^\top.

This orthonormalization fixes the scale of the update direction: OtO_t has orthonormal columns, Ot2=1\|O_t\|_2 = 1, and OtFn\|O_t\|_F \le \sqrt{n}. As a consequence, the per-step movement obeys Xt+1XtFηn\|X_{t+1}-X_t\|_F \le \eta \sqrt{n}, 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 MtM_t0, the exact Muon LMO is

MtM_t1

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 MtM_t2 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: MtM_t3 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 MtM_t4 and orthonormalized ascent on the dual matrix MtM_t5, 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 MtM_t6 and MtM_t7 every MtM_t8 iterations, and the server averages both: MtM_t9 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 Ot=argminORm×nOMtF2s.t.OO=In.O_t = \arg\min_{O \in \mathbb{R}^{m\times n}} \|O - M_t\|_F^2 \quad \text{s.t.} \quad O^\top O = I_n.0, a client control variate Ot=argminORm×nOMtF2s.t.OO=In.O_t = \arg\min_{O \in \mathbb{R}^{m\times n}} \|O - M_t\|_F^2 \quad \text{s.t.} \quad O^\top O = I_n.1, and receives a global control variate Ot=argminORm×nOMtF2s.t.OO=In.O_t = \arg\min_{O \in \mathbb{R}^{m\times n}} \|O - M_t\|_F^2 \quad \text{s.t.} \quad O^\top O = I_n.2. The local direction is formed from

Ot=argminORm×nOMtF2s.t.OO=In.O_t = \arg\min_{O \in \mathbb{R}^{m\times n}} \|O - M_t\|_F^2 \quad \text{s.t.} \quad O^\top O = I_n.3

after which the client applies Ot=argminORm×nOMtF2s.t.OO=In.O_t = \arg\min_{O \in \mathbb{R}^{m\times n}} \|O - M_t\|_F^2 \quad \text{s.t.} \quad O^\top O = I_n.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

Ot=argminORm×nOMtF2s.t.OO=In.O_t = \arg\min_{O \in \mathbb{R}^{m\times n}} \|O - M_t\|_F^2 \quad \text{s.t.} \quad O^\top O = I_n.5

where Ot=argminORm×nOMtF2s.t.OO=In.O_t = \arg\min_{O \in \mathbb{R}^{m\times n}} \|O - M_t\|_F^2 \quad \text{s.t.} \quad O^\top O = I_n.6 is the local Muon direction and

Ot=argminORm×nOMtF2s.t.OO=In.O_t = \arg\min_{O \in \mathbb{R}^{m\times n}} \|O - M_t\|_F^2 \quad \text{s.t.} \quad O^\top O = I_n.7

is a server-side estimate of the global direction. In the same method, momentum aggregation replaces per-round momentum reinitialization: Ot=argminORm×nOMtF2s.t.OO=In.O_t = \arg\min_{O \in \mathbb{R}^{m\times n}} \|O - M_t\|_F^2 \quad \text{s.t.} \quad O^\top O = I_n.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 Ot=argminORm×nOMtF2s.t.OO=In.O_t = \arg\min_{O \in \mathbb{R}^{m\times n}} \|O - M_t\|_F^2 \quad \text{s.t.} \quad O^\top O = I_n.9, global momenta Mt=UtStVtM_t = U_t S_t V_t^\top0, and control variates Mt=UtStVtM_t = U_t S_t V_t^\top1. Clients perform local Muon-style updates

Mt=UtStVtM_t = U_t S_t V_t^\top2

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 Mt=UtStVtM_t = U_t S_t V_t^\top3, Mt=UtStVtM_t = U_t S_t V_t^\top4, and Mt=UtStVtM_t = U_t S_t V_t^\top5 that depend only on the number of workers and the time horizon, not on problem-specific constants such as the Lipschitz constant Mt=UtStVtM_t = U_t S_t V_t^\top6, the stochastic noise level Mt=UtStVtM_t = U_t S_t V_t^\top7, or the heterogeneity parameter Mt=UtStVtM_t = U_t S_t V_t^\top8. Under bounded-variance noise, the dominant nonconvex rate is

Mt=UtStVtM_t = U_t S_t V_t^\top9

with Ot=UtVtO_t = U_t V_t^\top0 iterations and communication complexity Ot=UtVtO_t = U_t V_t^\top1 to reach an Ot=UtVtO_t = U_t V_t^\top2-stationary point. Under heavy-tailed noise with moment index Ot=UtVtO_t = U_t V_t^\top3, the dominant term becomes

Ot=UtVtO_t = U_t V_t^\top4

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 Ot=UtVtO_t = U_t V_t^\top5 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 Ot=UtVtO_t = U_t V_t^\top6, the result controls a Frobenius-norm quantity; as the number of iterations increases, the bound interpolates toward the trace norm through a Schatten-Ot=UtVtO_t = U_t V_t^\top7 parameter that decreases from Ot=UtVtO_t = U_t V_t^\top8 toward Ot=UtVtO_t = U_t V_t^\top9 (Takezawa et al., 30 Sep 2025).

The alignment-based FedMuon develops a different guarantee. Its final rate is

Xt+1=XtηUtVt.X_{t+1} = X_t - \eta\, U_t V_t^\top.0

where Xt+1=XtηUtVt.X_{t+1} = X_t - \eta\, U_t V_t^\top.1 is the number of participating clients, Xt+1=XtηUtVt.X_{t+1} = X_t - \eta\, U_t V_t^\top.2 is the number of local steps, Xt+1=XtηUtVt.X_{t+1} = X_t - \eta\, U_t V_t^\top.3 is the number of communication rounds, and Xt+1=XtηUtVt.X_{t+1} = X_t - \eta\, U_t V_t^\top.4. The paper explicitly contrasts this with local-SGD-type bounds that contain a heterogeneity term Xt+1=XtηUtVt.X_{t+1} = X_t - \eta\, U_t V_t^\top.5; its claim is that local-global alignment removes dependence on Xt+1=XtηUtVt.X_{t+1} = X_t - \eta\, U_t V_t^\top.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 Xt+1=XtηUtVt.X_{t+1} = X_t - \eta\, U_t V_t^\top.7 clients, Xt+1=XtηUtVt.X_{t+1} = X_t - \eta\, U_t V_t^\top.8 local steps, and Xt+1=XtηUtVt.X_{t+1} = X_t - \eta\, U_t V_t^\top.9 communication rounds, the reported rate is

OtO_t0

where OtO_t1 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 OtO_t2 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(OtO_t3) 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 OtO_t4 to OtO_t5, 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 OtO_t6 total clients, OtO_t7 sampled clients, and OtO_t8 local steps per round. Data heterogeneity is controlled by Dirichlet partitions with OtO_t9 and Ot2=1\|O_t\|_2 = 10. 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 Ot2=1\|O_t\|_2 = 11 trains reasonably well, that increasing Ot2=1\|O_t\|_2 = 12 from Ot2=1\|O_t\|_2 = 13 to Ot2=1\|O_t\|_2 = 14 yields a substantial boost, and that the best settings in the reported experiment are Ot2=1\|O_t\|_2 = 15 in the homogeneous case and Ot2=1\|O_t\|_2 = 16 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-Ot2=1\|O_t\|_2 = 17, the best baseline accuracies reported are Ot2=1\|O_t\|_2 = 18 for FedCM and Ot2=1\|O_t\|_2 = 19 for Local Muon, whereas FedMuon reaches OtFn\|O_t\|_F \le \sqrt{n}0 test accuracy with loss OtFn\|O_t\|_F \le \sqrt{n}1. Under Dirichlet-OtFn\|O_t\|_F \le \sqrt{n}2, the best baselines are OtFn\|O_t\|_F \le \sqrt{n}3 for FedCM and OtFn\|O_t\|_F \le \sqrt{n}4 for Local Muon, while FedMuon reaches OtFn\|O_t\|_F \le \sqrt{n}5 with loss OtFn\|O_t\|_F \le \sqrt{n}6. For ViT-Tiny on the same task, FedMuon reaches OtFn\|O_t\|_F \le \sqrt{n}7 under Dirichlet-OtFn\|O_t\|_F \le \sqrt{n}8 and OtFn\|O_t\|_F \le \sqrt{n}9 under Dirichlet-Xt+1XtFηn\|X_{t+1}-X_t\|_F \le \eta \sqrt{n}0, exceeding the cited Local Muon and Local AdamW baselines. On Swin Transformer fine-tuning, the paper reports Xt+1XtFηn\|X_{t+1}-X_t\|_F \le \eta \sqrt{n}1 on CIFAR-100 and Xt+1XtFηn\|X_{t+1}-X_t\|_F \le \eta \sqrt{n}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 Xt+1XtFηn\|X_{t+1}-X_t\|_F \le \eta \sqrt{n}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 Xt+1XtFηn\|X_{t+1}-X_t\|_F \le \eta \sqrt{n}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 Xt+1XtFηn\|X_{t+1}-X_t\|_F \le \eta \sqrt{n}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 Xt+1XtFηn\|X_{t+1}-X_t\|_F \le \eta \sqrt{n}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.

Topic to Video (Beta)

No one has generated a video about this topic yet.

Whiteboard

No one has generated a whiteboard explanation for this topic yet.

Follow Topic

Get notified by email when new papers are published related to FedMuon.