# Federated Compositional Muon Optimizer for Matrix-Wise Models

arXiv: https://arxiv.org/abs/2608.12710 · HTML: https://mlsanity.com/paper/2608.12710/

## Key facts

- No citing papers in the map yet (data through 14 Oct 2026).
- Builds on 36 of its 47 references that are in the map.
- First posted on arXiv 13 Aug 2026.

## Builds on (most cited first, 30 of 36)

- [Attention Is All You Need](https://mlsanity.com/paper/1706.03762/) (2017)
- [Communication-Efficient Learning of Deep Networks from Decentralized Data](https://mlsanity.com/paper/1602.05629/) (2016)
- [Training data-efficient image transformers & distillation through attention](https://mlsanity.com/paper/2012.12877/) (2020)
- [Pointer Sentinel Mixture Models](https://mlsanity.com/paper/1609.07843/) (2016)
- [Advances and Open Problems in Federated Learning](https://mlsanity.com/paper/1912.04977/) (2019)
- [Federated Optimization in Heterogeneous Networks](https://mlsanity.com/paper/1812.06127/) (2018)
- [Adaptive Federated Optimization](https://mlsanity.com/paper/2003.00295/) (2020)
- [Muon is Scalable for LLM Training](https://mlsanity.com/paper/2502.16982/) (2025)
- [A Survey on Federated Learning Systems: Vision, Hype and Reality for Data Privacy and Protection](https://mlsanity.com/paper/1907.09693/) (2019)
- [Personalized Federated Learning: A Meta-Learning Approach](https://mlsanity.com/paper/2002.07948/) (2020)
- [Momentum-Based Variance Reduction in Non-Convex SGD](https://mlsanity.com/paper/1905.10018/) (2019)
- [Training Deep Learning Models with Norm-Constrained LMOs](https://mlsanity.com/paper/2502.07529/) (2025)
- [Stochastic Compositional Gradient Descent: Algorithms for Minimizing Compositions of Expected-Value Functions](https://mlsanity.com/paper/1411.3803/) (2014)
- [Understanding Gradient Orthogonalization for Deep Learning via Non-Euclidean Trust-Region Optimization](https://mlsanity.com/paper/2503.12645/) (2025)
- [A Note on the Convergence of Muon](https://mlsanity.com/paper/2502.02900/) (2025)
- [Gluon: Making Muon & Scion Great Again! (Bridging Theory and Practice of LMO-based Optimizers for LLMs)](https://mlsanity.com/paper/2505.13416/) (2025)
- [Federated Reinforcement Learning with Environment Heterogeneity](https://mlsanity.com/paper/2204.02634/) (2022)
- [FedNest: Federated Bilevel, Minimax, and Compositional Optimization](https://mlsanity.com/paper/2205.02215/) (2022)
- [A Single Time-Scale Stochastic Approximation Method for Nested Stochastic Optimization](https://mlsanity.com/paper/1812.01094/) (2018)
- [Solving Stochastic Compositional Optimization is Nearly as Easy as Solving Stochastic Optimization](https://mlsanity.com/paper/2008.10847/) (2020)
- [STEM: A Stochastic Two-Sided Momentum Algorithm Achieving Near-Optimal Sample and Communication Complexities for Federated Learning](https://mlsanity.com/paper/2106.10435/) (2021)
- [Lions and Muons: Optimization via Stochastic Frank-Wolfe under Heavy-Tailed Noise](https://mlsanity.com/paper/2506.04192/) (2025)
- [Momentum Benefits Non-IID Federated Learning Simply and Provably](https://mlsanity.com/paper/2306.16504/) (2023)
- [Convergence of Muon with Newton-Schulz](https://mlsanity.com/paper/2601.19156/) (2026)
- [On the Convergence of Muon and Beyond](https://mlsanity.com/paper/2509.15816/) (2025)
- [Multi-Level Composite Stochastic Optimization via Nested Variance Reduction](https://mlsanity.com/paper/1908.11468/) (2019)
- [LiMuon: Light and Fast Muon Optimizer for Large Models](https://mlsanity.com/paper/2509.14562/) (2025)
- [Optimal Algorithms for Stochastic Multi-Level Compositional Optimization](https://mlsanity.com/paper/2202.07530/) (2022)
- [FedMuon: Accelerating Federated Learning with Matrix Orthogonalization](https://mlsanity.com/paper/2510.27403/) (2025)
- [Compositional federated learning: Applications in distributionally robust averaging and meta learning](https://mlsanity.com/paper/2106.11264/) (2021)

## Abstract (arXiv)

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(ε-3) for finding an ε-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.