# A Note on Stability for Orthogonalized Matrix Momentum with Client Sampling

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

## Key facts

- Cited by 1 papers in the map (data through 14 Oct 2026).
- Builds on 39 of its 51 references that are in the map.
- First posted on arXiv 1 Jun 2026.

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

- [Communication-Efficient Learning of Deep Networks from Decentralized Data](https://mlsanity.com/paper/1602.05629/) (2016)
- [Adaptive Federated Optimization](https://mlsanity.com/paper/2003.00295/) (2020)
- [Local SGD Converges Fast and Communicates Little](https://mlsanity.com/paper/1805.09767/) (2018)
- [Train faster, generalize better: Stability of stochastic gradient descent](https://mlsanity.com/paper/1509.01240/) (2015)
- [Muon is Scalable for LLM Training](https://mlsanity.com/paper/2502.16982/) (2025)
- [Parallel Restarted SGD with Faster Convergence and Less Communication: Demystifying Why Model Averaging Works for Deep Learning](https://mlsanity.com/paper/1807.06629/) (2018)
- [Old Optimizer, New Norm: An Anthology](https://mlsanity.com/paper/2409.20325/) (2024)
- [Training Deep Learning Models with Norm-Constrained LMOs](https://mlsanity.com/paper/2502.07529/) (2025)
- [On the Linear Speedup Analysis of Communication Efficient Momentum SGD for Distributed Non-Convex Optimization](https://mlsanity.com/paper/1905.03817/) (2019)
- [The Polar Express: Optimal Matrix Sign Methods and Their Application to the Muon Algorithm](https://mlsanity.com/paper/2505.16932/) (2025)
- [High probability generalization bounds for uniformly stable algorithms with nearly optimal rate](https://mlsanity.com/paper/1902.10710/) (2019)
- [Sharper bounds for uniformly stable algorithms](https://mlsanity.com/paper/1910.07833/) (2019)
- [Dion: Distributed Orthonormalized Updates](https://mlsanity.com/paper/2504.05295/) (2025)
- [AdaMuon: Adaptive Muon Optimizer](https://mlsanity.com/paper/2507.11005/) (2025)
- [Generalization Bounds for Uniformly Stable Algorithms](https://mlsanity.com/paper/1812.09859/) (2018)
- [Nonconvex Stochastic Optimization under Heavy-Tailed Noises: Optimal Convergence without Gradient Clipping](https://mlsanity.com/paper/2412.19529/) (2024)
- [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)
- [MuonBP: Faster Muon via Block-Periodic Orthogonalization](https://mlsanity.com/paper/2510.16981/) (2025)
- [On the Convergence of Muon and Beyond](https://mlsanity.com/paper/2509.15816/) (2025)
- [Preconditioning Benefits of Spectral Orthogonalization in Muon](https://mlsanity.com/paper/2601.13474/) (2026)
- [Beyond the Ideal: Analyzing the Inexact Muon Update](https://mlsanity.com/paper/2510.19933/) (2025)
- [FedSWA: Improving Generalization in Federated Learning with Highly Heterogeneous Data via Momentum-Based Stochastic Controlled Weight Averaging](https://mlsanity.com/paper/2507.20016/) (2025)
- [MuonEq: Balancing Before Orthogonalization with Lightweight Equilibration](https://mlsanity.com/paper/2603.28254/) (2026)
- [Muon2: Boosting Muon via Adaptive Second-Moment Preconditioning](https://mlsanity.com/paper/2604.09967/) (2026)
- [FedMuon: Accelerating Federated Learning with Matrix Orthogonalization](https://mlsanity.com/paper/2510.27403/) (2025)
- [Robust Machine Unlearning for Quantized Neural Networks via Adaptive Gradient Reweighting with Similar Labels](https://mlsanity.com/paper/2503.13917/) (2025)
- [InSPO: Unlocking Intrinsic Self-Reflection for LLM Preference Optimization](https://mlsanity.com/paper/2512.23126/) (2025)
- [FedMuon: Federated Learning with Bias-corrected LMO-based Optimization](https://mlsanity.com/paper/2509.26337/) (2025)
- [Reason in Chains, Learn in Trees: Self-Rectification and Grafting for Multi-turn Agent Policy Optimization](https://mlsanity.com/paper/2604.07165/) (2026)

## Built on by (most cited first, 1 of 1)

- [Muse: Representation Geometry of Muon Beyond Normalized Momentum](https://mlsanity.com/paper/2607.14536/) (2026)

## Abstract (arXiv)

We study finite-sample generalization for a client-sampled distributed optimization scheme with matrix-valued parameters and orthogonalized momentum updates. The central quantity is the gap between the population and empirical objectives at the returned model when only a subset of clients participates in each round. Under independent heterogeneous client data, unequal local sample counts, and fixed aggregation weights, we derive a finite-round upper-tail guarantee from a coupled-neighbor stability recursion and a weighted concentration step. The bound keeps the client-selection counts through the amplification factor (Yi(mathcal C)); in the uniform full-participation full-batch regime, it yields (widetildemathcal O(n-1+n-1/2)) scaling whenever the horizon-dependent amplification terms are controlled. The matrix-orthogonalization rule is required to be Lipschitz along paired trajectories, a condition satisfied by regularized polar-type maps and normalized finite-step Newton--Schulz orthogonalizers. For the unregularized matrix sign, the same argument requires coupled spectral separation, whereas Gaussian smoothing gives a finite-round smoothed variant. A one-dimensional counterexample shows why a gap, smoothing, or regularity condition is necessary.