Federated Compositional Muon Optimizer for Matrix-Wise Models

2026-08-12

AI Summary

This work addresses the challenge of combinatorial optimization for matrix-valued models in distributed settings under non-independent and identically distributed (non-IID) data and non-convex objectives. The authors propose the Federated Combinatorial Muon optimizer (FedCoMuon) and its variance-reduced variant, FedCoMuon-VR, which uniquely integrate combinatorial gradient tracking and orthogonal momentum mechanisms into federated learning, augmented with momentum-based variance reduction. Theoretical analysis demonstrates that FedCoMuon-VR achieves a sample complexity of $O(\varepsilon^{-3})$ under non-IID non-convex conditions, improving upon the existing FedMuon method. Empirical evaluations on robust federated learning and task-distributed risk-sensitive meta-learning benchmarks show substantial gains over current combinatorial optimization baselines, establishing state-of-the-art accuracy.

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(ε^{-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.

Problem

• federated optimization
• compositional optimization
• matrix-wise models
• non-iid
• non-convex

Innovation

• Federated Optimization
• Compositional Optimization
• Muon Optimizer
• Variance Reduction
• Matrix-wise Models