Pith. sign in

REVIEW 2 major objections 5 minor 17 references

Federated Compositional Muon Optimizer for Matrix-Wise Models

T0 review · 2 major / 5 minor · reviewed 2026-08-16 · deepseek-v4-flash

Challenge this review Re-run · record.json Download PDF Read on arXiv ↗ Review JSON

Pith's one-line read Muon-style orthogonalized momentum can be extended to federated compositional optimization, where a variance-reduced variant reaches an $\epsilon$-stationary point with $O(\epsilon^{-3})$ per-client sample complexity.

desk verdict A solid synthesis of Muon and compositional federated optimization, but the central O(ε^{-3}) claim is proven only for an idealized exact-SVD update, not for the Newton-Schulz algorithm actually run. read the letter →

arxiv 2608.12710 v1 pith:FURIUKSJ submitted 2026-08-13 cs.LG math.OC

Wang Yan , Feihu Huang This is my paper · ORCID
classification cs.LGmath.OC
keywords federatedlearningcompositionaloptimizationMuonoptimizerorthogonalizedmomentumvariancereductionnon-convexmatrix-wisemodelsNewton-Schulziterations
verification ladder T0 review T1 audit T2 compute T3 formal

The pith

A machine-rendered reading of the paper's core claim, the machinery that carries it, and where it could break.

The reading

Federated models whose parameters are matrices—transformers, CNNs—are usually trained on single-level objectives, but many real tasks such as robust federated learning and meta learning have nested objectives: a client loss is an expected value of a function of another expected value. The paper sets out to extend the Muon optimizer, which updates matrix blocks along an orthogonalized momentum direction, to this compositional federated setting. It proposes FedCoMuon, built from compositional gradient tracking plus orthogonalized momentum, and FedCoMuon-VR, which adds momentum-based variance reduction. The central theoretical claim is that FedCoMuon-VR finds an $\epsilon$-stationary point with per-client sample complexity $O(\epsilon^{-3})$ and communication complexity $O(\epsilon^{-3})$, which the authors state is lower than existing FedMuon algorithms. If correct, this would give a Muon-type method with convergence guarantees for distributed nested objectives under non-i.i.d., non-convex conditions, and the experiments show gains on robust learning and meta-learning benchmarks.

What carries the argument

The load-bearing mechanism is the compositional matrix-momentum recursion combined with orthogonalized updates. Each client maintains $M^k_{t+1}=\beta Z^k_{t+1}+(1-\beta)M^k_t$, where $Z^k_t$ is the matrix-form stochastic compositional gradient, and then replaces the momentum matrix by its orthogonal factor pair $(U,V)$ obtained from a Newton–Schulz-approximated SVD, updating $W$ by $-\eta U V^\top$. The key identity used in the descent analysis is Lemma 2: for a compact SVD $B=U_B\Sigma_B V_B^\top$, $\langle A, U_B V_B^\top\rangle \ge \|A\|_F - 2\sqrt{n}\|A-B\|_F$, which converts progress in Frobenius norm into progress along the approximate gradient direction. FedCoMuon-VR wraps this in momentum-based variance-reduced trackers for the inner value, inner Jacobian, and outer gradient, with projections onto bounded balls, so that the estimator error contracts rather than accumulating across iterations.

What would settle it

Compute, over a run, the maximum Frobenius distance between the five-iteration Newton–Schulz output and the exact compact SVD direction; if this error does not shrink with $T$, the theorem's rate cannot apply to the implemented update, and a direct comparison run with exact compact SVD would settle whether the proof's idealization matters.

Watch

Extended reading notes

Core claim

On the paper's own terms, the discovery is that the Muon update does not have to be confined to single-level federated objectives: by maintaining a moving-average tracker $u^k_t$ of the inner mapping $g_k(W)$, forming the matrix product estimator $Z^k_t = \nabla g_k(W^k_t;\xi)(\nabla_y f_k(u^k_t;\zeta)\otimes I_n)$, and running a momentum recursion on the matrix $M^k_t$, each client obtains an orthogonalized descent direction that still points downhill on the compositional objective. The variance-reduced variant replaces the raw estimators with momentum-based trackers $u$, $v$, and $H$ evaluated on shared fresh samples, and the analysis shows the compositional bias and client drift can be controlled simultaneously. Theorems 1 and 2 quantify this: FedCoMuon needs $O(\epsilon^{-4})$ per-client samples and $O(\epsilon^{-3})$ communication rounds for an $\epsilon$-stationary point, while FedCoMuon-VR lowers the sample complexity to $O(\epsilon^{-3})$ with the same communication complexity, under Assumptions 1–3, 6, and 7. The experiments on robust federated learning (MNIST, WikiText-2) and task-distributed meta learning (CIFAR-10 with CNN and ViT-Tiny) report that the proposed methods are competitive with compositional baselines and reach the best accuracy in several configurations.

Load-bearing premise

Theorems 1 and 2 rely on Lemma 2, which requires the exact compact SVD $M=U\Sigma V^\top$, while Algorithms 1 and 2 compute the update direction with Newton–Schulz iterations and no term in the theorems bounds the distance between the Newton–Schulz output and the exact $U V^\top$.

Editorial extensions

If this is right

  • FedCoMuon finds an $\epsilon$-stationary point with $O(\epsilon^{-4})$ per-client samples and $O(\epsilon^{-3})$ communication rounds; FedCoMuon-VR improves the sample bound to $O(\epsilon^{-3})$ while keeping $O(\epsilon^{-3})$ communication.
  • Muon-style orthogonalized momentum remains convergent when the objective is a nested expectation and client data are non-i.i.d., provided the compositional bias is tracked and variance-reduced.
  • For matrix-wise models with many parameter blocks, the same update rule works for robust federated learning and task-distributed meta learning, where standard compositional federated methods are worse in the reported experiments.
  • The claimed improvement over existing FedMuon algorithms applies specifically to the compositional setting: the FedMuon baselines are designed for single-level objectives and do not carry these guarantees.

Reading between the lines

Editorial extensions of the paper, not claims the author makes directly.

  • If a bound on the Newton–Schulz approximation error can be added to the analysis, the same descent identity would likely extend the rate guarantees to the exact implemented update rather than to the idealized exact-SVD direction.
  • The $O(\epsilon^{-3})$ sample complexity suggests the compositional bias, not the orthogonalization, is the bottleneck; a similar rate may hold for other structured optimizers that correct nested gradients with momentum.
  • A direct experiment comparing five Newton–Schulz iterations against an exact compact SVD on the same tasks would isolate whether the reported empirical gains come from the orthogonalization mechanism or from the variance-reduced compositional tracker.
  • The framework could plausibly be adapted to federated bilevel or multi-level objectives, where the same tracker-plus-orthogonalization pattern applies, though the paper does not analyze those settings.
Share X Bluesky LinkedIn Reddit HN

Signed reviews

No signed human review yet.

Editorial analysis

A structured set of objections, weighed in public.

Desk editor's note, referee report, and a circularity audit.

Referee Report

2 major / 5 minor

Summary. The paper proposes two federated compositional Muon optimizers, FedCoMuon and its variance-reduced variant FedCoMuon-VR, for matrix-wise objectives of the form (1/K) sum_k E[ f_k(E[g_k(W)]) ]. It claims non-convex, non-i.i.d. convergence guarantees: FedCoMuon finds an epsilon-stationary solution with O(epsilon^-4) per-client sample complexity and O(epsilon^-3) communication complexity, while FedCoMuon-VR achieves O(epsilon^-3) sample and communication complexity. The proof appendix supplies detailed algebraic derivations and explicit parameter choices. The experiments cover robust federated learning on MNIST and WikiText-2, and task-distributed meta-learning on CIFAR-10 with CNN and ViT-Tiny models.

Significance. If the advertised rates hold for the algorithms as implemented, the O(epsilon^-3) per-client sample complexity for federated compositional optimization with an orthogonalized matrix update would be a useful step beyond existing FedMuon and compositional federated baselines. The appendix is thorough: constants are explicit, the parameter tuning is algebraic, and no fitted constants enter the theorems. The main obstacle is that the analysis applies to an exact compact-SVD update, while the algorithms and experiments use five fixed Newton-Schulz iterations; until that gap is closed, the central claim is not established for the implemented method.

major comments (2)
  1. [Section 5 and Appendix B, Lemma 2 (Eq. (12)) and Lemma 14 (Eq. (83)); Algorithms 1 and 2, line 5; Appendix C] The convergence theorems are proven only for an idealized exact compact-SVD update, not for the Newton-Schulz update actually used in the algorithms. Lemma 2 and its use in Lemma 14 require the identity <M, U V^T> = ||M||_* for the exact polar factor of a compact SVD. However, Algorithms 1 and 2 compute the update direction by 'SVD ... with the Newton-Schulz approach', and Appendix C fixes five Newton-Schulz iterations. For the implemented update Q_t^k = NS_5(M_t^k), the proofs would need a lower bound on <grad F, Q_t^k> that controls the approximation error <grad F, Q_t^k - U_t^k(V_t^k)^T>. No such bound appears in Theorem 1 or Theorem 2, and with only five iterations the error ||Q - U V^T||_F can be O(1) for matrices with a wide singular-value spread, e.g. M = diag(1, epsilon, ..., epsilon) with small epsilon. After telescoping in Lemma 14, this would leave a non-vanishing additive term in the stationarity bound. Consequently the claimed O(epsilon^-3) sample complexity of FedCoMuon-VR, and similarly the FedCoMuon rate, is not established for the algorithm whose numerical results are reported. To fix this, the authors should either analyze the actual Newton-Schulz update with a formal error bound, or change the algorithm to use an exact SVD/polar factor and state this explicitly in the theorems.
  2. [Abstract and Remark 2] The paper claims that FedCoMuon-VR has a lower sample complexity than existing FedMuon algorithms, but it never states the sample or communication complexities of the cited FedMuon baselines [Takezawa et al., 2026, Liu et al., 2025b, Zhang and Gao, 2025, Qian et al., 2026]. As written, this central comparative claim is not verifiable from the manuscript. I request a quantitative comparison table or explicit cited rates for the baselines.
minor comments (5)
  1. [Appendix B, after Lemma 14] The theorem immediately following Lemma 14 is labeled 'Theorem 1 (Convergence of FedCoMuon-VR with first-order product tracker)' but it is Theorem 2 of the paper; the numbering should be corrected.
  2. [Lemma 13 proof] The proof of Lemma 13 refers to 'Lemma C.4' when bounding the client product disagreement; the correct reference appears to be Lemma 11. The appendix contains no Lemma C.4.
  3. [Algorithms 1 and 2, line 5] The pseudocode returns factors (U_t^k, Sigma_t^k, V_t^k) from an operation described as 'SVD ... with the Newton-Schulz approach', which is notationally inconsistent. The authors should either define the approximate polar factor Q_t^k explicitly in the pseudocode or give the Newton-Schulz iteration as a separate subroutine.
  4. [Section 5, synchronization convention] Lemmas 5 and 13 rely on the convention that after a communication round all clients set W and M to the server averages, but this convention is only stated informally in Section 4. I suggest stating it once before the lemmas to make the proofs self-contained.
  5. [Section 6 and Appendix C] The experiments report averages over only two random seeds and do not show error bars or standard deviations. Given the modest number of runs, statements such as 'achieve the best overall performance' should be interpreted with caution.

Circularity Check

0 steps flagged · score 0.0 of 10

No significant circularity: the convergence rates are derived from explicit assumptions; the Newton-Schulz versus exact-SVD gap is a proof gap, not a circular reduction.

full rationale

The paper's central claims, the O(epsilon^{-4}) and O(epsilon^{-3}) sample complexities of FedCoMuon and FedCoMuon-VR, are obtained by telescoping/summation arguments (Lemmas 3, 5, 7, 8-14, and Theorem 1/2 in the appendix) from Assumptions 1-7. No fitted constant enters the bounds, and no parameter is tuned to a subset of data and then relabeled as a prediction. The self-citations, e.g., ComFedL [Huang and Li, 2021] and Huang et al. [2026], appear as background and baseline comparisons, not as premises that force the convergence results. The claimed superiority over existing FedMuon algorithms is a comparison with externally reported complexities, not a restatement of this paper's own assumptions. The one substantive concern is the discrepancy between Algorithm line 5, which says 'SVD ... with the Newton-Schulz approach', and Lemma 2, which requires an exact compact SVD M = U Sigma V^T for the key inequality <A, U_B V_B^T> >= ||A||_F - 2 sqrt(n) ||A-B||_F. If the Newton-Schulz iterate is not close to the polar factor, the theorems describe an idealized algorithm rather than the implemented one. That is a correctness/completeness gap in the proof, not a circularity: the argument does not assume its conclusion, and the missing approximation error is not a fitted input being renamed as a prediction. Since no circular step can be quoted or exhibited, the honest finding is no significant circularity.

Assumptions & free parameters 3 free parameters · 8 assumptions · 0 invented entities

The central rate claims rest on a standard set of smoothness, bounded-variance, and heterogeneity assumptions from compositional federated optimization, plus an unstated identification of the implemented Newton-Schulz update with the exact SVD used in the proofs. No constants are fitted to data in the theory; experimental hyperparameters are grid-selected.

free parameters (3)
  • Experimental learning rates and momentum hyperparameters for FedCoMuon and FedCoMuon-VR = Per task, e.g., MNIST FedCoMuon-VR uses alpha=0.2, beta=0.8, gamma=0.9, rho=0.2; learning rates and other…
    Selected by grid search on validation; they support the empirical accuracy claims but are not part of the convergence theorem.
  • Regularization parameter lambda = 0.5
    Grid-searched in meta-learning experiments; not load-bearing for the theoretical claims.
  • Initial batch size b for FedCoMuon-VR = Set as T^{2/3} in theory; not specified as a finite number in the experiments
    Theorem 2 requires b = T^{2/3}, but the experiments do not state the initial batch size, so the empirical runs may not operate in the analyzed regime.
assumptions (8)
  • domain assumption Assumption 1: the global objective F(W) is lower bounded
    Used for telescoping the descent inequality over T iterations.
  • domain assumption Assumption 2: unbiased stochastic oracles and bounded variances for g, Jacobian of g, and outer gradient of f
    Controls the noise of the compositional gradient estimators.
  • domain assumption Assumption 3: bounded second moments of stochastic Jacobians and outer gradients
    Gives bounded constants C_g and C_f used throughout the proofs.
  • domain assumption Assumption 4: L_g-smoothness of g_k and L_f-smoothness of f_k
    Implies L_F-smoothness of the composed objective, used in every descent lemma.
  • domain assumption Assumption 5: bounded gradient heterogeneity for FedCoMuon
    Controls client drift in Theorem 1 through the delta term.
  • domain assumption Assumption 6: mean-square sample smoothness for FedCoMuon-VR
    Required for the STORM-style recursive estimators to have controlled second moments.
  • domain assumption Assumption 7: bounded client heterogeneity in g, its Jacobian, and outer gradients
    Produces the non-i.i.d. heterogeneity term in Theorem 2.
  • ad hoc to paper Exact compact SVD is used in place of Newton-Schulz approximation
    Lemma 2 requires an exact compact SVD M = U Sigma V^T, while Algorithms 1 and 2 state Newton-Schulz orthogonalization; no approximation error bound is provided.

how reviews work

0 comments
Cite this review

Pith. "Pith review of Federated Compositional Muon Optimizer for Matrix-Wise Models." pith.science (2026). https://pith.science/paper/FURIUKSJ

@misc{pith2026260812710,
  author       = {Pith},
  title        = {Pith review of: Federated Compositional Muon Optimizer for Matrix-Wise Models},
  year         = {2026},
  howpublished = {\url{https://pith.science/paper/FURIUKSJ}},
  note         = {Machine review of arXiv:2608.12710}
}
abstract

Muon, a more recently developed optimizer, is useful for matrix-wise models in AI areas. Although many works have studied Muon and its variants, these methods are still not particularly well-suited for hierarchical structured problems. To fill this gap, we propose an effective federated compositional Muon (FedCoMuon) optimizer to solve distributed matrix-wise compositional optimization problems. Specifically, our FedCoMuon optimizer builds on compositional gradient tracking and orthogonalized momentum. Moreover, we propose a variance reduced variant of FedCoMuon (FedCoMuon-VR) based on a momentum-based variance reduced technique. In theory, we analyze the convergence properties of our algorithms under the non-i.i.d. and non-convex settings. In particular, we prove that our FedCoMuon-VR obtains a lower sample complexity of $O(\epsilon^{-3})$ for finding an $\epsilon$-stationary solution than the existing FedMuon algorithms. Extensive numerical experiments on robust federated learning and task-distributed risk-sensitive meta learning show that our proposed methods are competitive with existing compositional baselines and achieve the best reported accuracy in several settings.

Figures

Figures reproduced from arXiv: 2608.12710 by the authors.

Figure 1
Figure 1. Training and test performance on the imbalanced MNIST robust federated learning task. [PITH_FULL_IMAGE:figures/full_fig_p009_1.png] view at source ↗
Figure 2
Figure 2. Effect of the synchronization gap τ on our algorithm for the robust federated learning task. 6 Numerical Experiments In the section, we evaluate FedCoMuon and FedCoMuon-VR on robust federated learning and task-distributed meta learning. In the experiment, we compare our methods with task-specific standard federated baselines: FedAvg [McMahan et al., 2017] for robust federated learning and FedMAML [Fallah et al., 202… view at source ↗
Figure 3
Figure 3. Test loss and perplexity of each method on the WikiText-2 language-modeling task. [PITH_FULL_IMAGE:figures/full_fig_p010_3.png] view at source ↗
Figures from the paper (3 more)
Figure 4
Figure 4. Figure 4: Test accuracy and loss of each method on the task-distributed CNN meta-learning task [PITH_FULL_IMAGE:figures/full_fig_p011_4.png]
Figure 5
Figure 5. Figure 5: Effect of the synchronization gap τ on task-distributed CNN meta learning with hetero￾geneous CIFAR-10 data. following risk-sensitive compositional MAML objective: min W∈Rm×n 1 K X K k=1 exp ℓ k [PITH_FULL_IMAGE:figures/full_fig_p012_5.png]
Figure 6
Figure 6. Figure 6: Test accuracy and loss of each method on the task-distributed ViT-Tiny meta learning [PITH_FULL_IMAGE:figures/full_fig_p013_6.png]

Discussion (0). Continue with ORCID to comment.

Reference graph

Works this paper leans on

17 extracted references · 1 canonical work pages

  1. [1]

    On the convergence of muon and beyond.arXiv preprint arXiv:2509.15816,

    Da Chang, Yongxiang Liu, and Ganzhao Yuan. On the convergence of muon and beyond.arXiv preprint arXiv:2509.15816,

    arXiv

  2. [7]

    Understanding gradient orthogonalization for deep learning via non-euclidean trust-region optimization.arXiv preprint arXiv:2503.12645,

    Dmitry Kovalev. Understanding gradient orthogonalization for deep learning via non-euclidean trust-region optimization.arXiv preprint arXiv:2503.12645,

    arXiv

  3. [9]

    Muon is scalable for LLM training.arXiv preprint arXiv:2502.16982, 2025a

    Jingyuan Liu, Jianlin Su, Xingcheng Yao, Zhejun Jiang, Guokun Lai, Yulun Du, Yidao Qin, Weixin Xu, Enzhe Lu, Junjie Yan, et al. Muon is scalable for LLM training.arXiv preprint arXiv:2502.16982, 2025a. Junkang Liu, Fanhua Shang, Junchao Zhou, Hongying Liu, Yuanyuan Liu, and Jin Liu. FedMuon: Accel- erating federated learning with matrix orthogonalization....

    arXiv

  4. [12]

    Communication-efficient gluon in federated learning.arXiv preprint arXiv:2604.10689,

    Xun Qian, Alexander Gaponov, Grigory Malinovsky, and Peter Richtárik. Communication-efficient gluon in federated learning.arXiv preprint arXiv:2604.10689,

    arXiv

  5. [13]

    Reddi, Zachary Charles, Manzil Zaheer, Zachary Garrett, Keith Rush, Jakub Konečný, Sanjiv Kumar, and H

    Sashank J. Reddi, Zachary Charles, Manzil Zaheer, Zachary Garrett, Keith Rush, Jakub Konečný, Sanjiv Kumar, and H. Brendan McMahan. Adaptive federated optimization.arXiv preprint arXiv:2003.00295,

    arXiv 2003

  6. [14]

    Gluon: Making muon & scion great again!(bridging theory and practice of lmo-based optimizers for llms).arXiv preprint arXiv:2505.13416,

    15 Artem Riabinin, Egor Shulgin, Kaja Gruntkowska, and Peter Richtárik. Gluon: Making muon & scion great again!(bridging theory and practice of lmo-based optimizers for llms).arXiv preprint arXiv:2505.13416,

    arXiv

  7. [15]

    Lions and muons: Optimization via stochastic frank-wolfe.arXiv preprint arXiv:2506.04192,

    Maria-Eleni Sfyraki and Jun-Kun Wang. Lions and muons: Optimization via stochastic frank-wolfe.arXiv preprint arXiv:2506.04192,

    arXiv

  8. [16]

    On the convergence analysis of muon.arXiv preprint arXiv:2505.23737,

    Wei Shen, Ruichuan Huang, Minhui Huang, Cong Shen, and Jiawei Zhang. On the convergence analysis of muon.arXiv preprint arXiv:2505.23737,

    arXiv

Show all 17 references
  1. [17]

    On provable benefits of muon in federated learning.arXiv preprint arXiv:2510.03866,

    Xinwen Zhang and Hongchang Gao. On provable benefits of muon in federated learning.arXiv preprint arXiv:2510.03866,

  2. [2015]

    A note on the convergence of muon and further.arXiv preprint arXiv:2502.02900,

    Jiaxiang Li and Mingyi Hong. A note on the convergence of muon and further.arXiv preprint arXiv:2502.02900,

    arXiv

  3. [2017]

    Pointer sentinel mixture models

    Stephen Merity, Caiming Xiong, James Bradbury, and Richard Socher. Pointer sentinel mixture models. arXiv preprint arXiv:1609.07843,

    arXiv

  4. [2019]

    Personalized federated learning: A meta-learning approach.arXiv preprint arXiv:2002.07948,

    Alireza Fallah, Aryan Mokhtari, and Asuman Ozdaglar. Personalized federated learning: A meta-learning approach.arXiv preprint arXiv:2002.07948,

    2002 arXiv

  5. [2020]

    Compositional federated learning: Applications in distributionally robust aver- aging and meta learning.arXiv preprint arXiv:2106.11264,

    Feihu Huang and Jian Li. Compositional federated learning: Applications in distributionally robust aver- aging and meta learning.arXiv preprint arXiv:2106.11264,

    arXiv

  6. [2021]

    LiMuon: Light and fast muon optimizer for large models

    Feihu Huang, Yuning Luo, and Songcan Chen. LiMuon: Light and fast muon optimizer for large models. arXiv preprint arXiv:2509.14562,

    arXiv

  7. [2024]

    Brendan McMahan, Brendan Avent, Aurélien Bellet, Mehdi Bennis, Arjun Nitin Bhagoji, Keith Bonawitz, Zachary Charles, Graham Cormode, Rachel Cummings, et al

    Peter Kairouz, H. Brendan McMahan, Brendan Avent, Aurélien Bellet, Mehdi Bennis, Arjun Nitin Bhagoji, Keith Bonawitz, Zachary Charles, Graham Cormode, Rachel Cummings, et al. Advances and open problems in federated learning.arXiv preprint arXiv:1912.04977,

    1912 arXiv

  8. [2025]

    Muon is provably faster with momentum variance reduction.arXiv preprint arXiv:2512.16598,

    Xun Qian, Hussein Rammal, Dmitry Kovalev, and Peter Richtarik. Muon is provably faster with momentum variance reduction.arXiv preprint arXiv:2512.16598,

  9. [2026]

    URLhttps://openreview.net/forum?id= 4uRlbSNevR

    ISSN 2835-8856. URLhttps://openreview.net/forum?id= 4uRlbSNevR. Gyu Yeol Kim and Min-hwan Oh. Convergence of muon with newton-schulz.arXiv preprint arXiv:2601.19156,

Pith tools

Reviewed August 16, 2026 · model on record in the stance chip above.