Skip to content

Federated Compositional Muon Optimizer for Matrix-Wise Models

Aug 2026 · 0 citations · 47 references
Computer Science Mathematics

TL;DR

This work proposes an effective federated compositional Muon (FedCoMuon) optimizer to solve distributed matrix-wise compositional optimization problems and proposes a variance reduced variant of FedCoMuon (FedCoMuon-VR) based on a momentum-based variance reduced technique.

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.

View source

Similar papers

Preprint Aug 2026

A Momentum-Based Variance-Reduced Algorithm for Federated Multiobjective Optimization

Federated learning has traditionally been formulated as a single-objective optimization problem, primarily focused on maximizing model utility. In real-world applications, however, machine learning models often need to optimize multiple and potentially conflicting objectives simultaneously. This motivates federated multiobjective optimization (FMOO), which provides a natural framework for jointly handling multiple task-specific objectives in federated learning. In this paper, we propose a momentum-based variance-reduced algorithm for federated multiobjective optimization. The method incorporates a momentum-driven gradient estimator into the local updates to reduce the variance of stochastic updates, leading to an improved convergence rate. We establish theoretical guarantees showing that the expected Pareto stationarity measure of a randomly selected output iterate decays at a rate of $\mathcal{O}(T^{-2/3})$, improving upon the $\mathcal{O}(T^{-1/2})$ rates established for existing methods such as FSMGDA and FedCMOO. Numerical experiments on federated multiobjective optimization benchmarks demonstrate the effectiveness and competitive performance of the proposed algorithm.

Yong Zhao, Chunlin You, M. N. Dao et al. · 0 citations
Preprint Aug 2026

Scaling Muon for Diffusion Transformers

Periodic Row-wise Muon is introduced, which performs a full NS5 spectral update once every \(K\) steps and applies a low compute and communication cost row-wise constrained update based on the current momentum at the remaining steps to preserve Muon's generative quality advantage while translating it into end-to-end training efficiency for large DiTs.

Chenghao Li, Xiao Han, Xinxin Huang et al. · 0 citations
Preprint Jul 2026

Reassessing Muon for Matrix Factorization

It is found that Muon does not consistently outperform AdamW in this setting and that several previously reported advantages are sensitive to hyperparameter choices, providing a more nuanced picture of when spectrum-aware orthogonalization is beneficial and arguing for evaluating modern optimizers on controlled problems in addition to end-to-end benchmarks.

Alipanah Parviz, Gal Mishne, Alex Cloninger · 0 citations
Preprint Jul 2026

Federated Lightweight Fine-Tuning

Federated fine-tuning is bottlenecked by communication: FedAvg and pseudo-gradient schemes transmit a payload that scales with the model, and gradient compression shrinks it by only a constant factor. We take a different lever. Mapping networks generate a network's weights from a small trainable latent through a frozen affine projection; because the map is shared and affine, averaging latents is exactly averaging the generated weights. We turn this into a practical low-bandwidth federated channel with two changes: a low-rank, seed-regenerable factorisation of the projection (cutting generator memory from ~80 GB to ~10 MB), and a delta formulation $\theta = \theta^{\mathrm{pre}} + U V^{\top} z$ that learns an additive correction around a shared centrally-pretrained base -- federated fine-tuning, which is what makes the method work at scale. A frozen orthogonal classifier head further removes the head from the payload while improving accuracy. On CIFAR-100 with ResNet-18+GroupNorm, our method (FLITE, Federated Low-rank Iterative Training Engine) communicates 1,280 floats (~5 KB) per client per round -- an 8718x reduction -- and reaches 74.67%, within ~0.5 pp of full-weight FedAvg. The averaging identity holds to floating-point precision ($6 \times 10^{-8}$); the method sits one to two orders of magnitude below PowerSGD and top-k on the bandwidth-accuracy Pareto; it matches or exceeds full-weight FedAvg under strong non-IID skew. int4 latents reach 648 bytes per round at unchanged accuracy, whereas int4 full-weight FedAvg collapses to chance.

R. Achanta, Will Reed · 0 citations
Review Open access Aug 2026

Federated Learning with SVMs: Dynamic SGD for Efficient Hyperparameter Optimization

In the current digital landscape, organizations are actively looking for user data to make informed decisions at the edge. This raises a need for a solution that prioritizes user privacy while leveraging user's data. Federated learning (FL) emerges as a viable solution to address this issue. However, these methods are dominated by complex, resource-intensive neural networks (NN). This necessitates the development of an FL technique that protects user privacy while employing lightweight models. Support vector machines (SVMs) provide a lightweight alternative, but their traditional training techniques, such as quadratic programming and sequential minimal optimization, are inefficient and incompatible with FL. Existing stochastic gradient descent (SGD)-based SVM variants also depend on fixed or heuristic learning rates, which restrict convergence under non-IID data and client heterogeneity. In this paper, we propose a novel meta-learned dynamic learning rate controller for SGD-trained SVMs in federated settings. Unlike standard adaptive optimizers, the controller adapts learning rates over each epoch and clients. This process enables faster convergence and lower communication overhead to heterogeneous data. Empirical evaluations on benchmark datasets show that federated SVM framework provides effective results comparable to NN-based FL approaches while significantly less computational and communication overhead. Received: 29 September 2025 | Revised: 30 March 2026 | Accepted: 9 July 2026Conflicts of Interest The authors declare that they have no conflicts of interest to this work. Data Availability Statement The data that support the findings of this study are openly available in Kaggle at https://git-disl.github.io/GTDLBench/datasets/mnist_datasets/ (MNIST), the NIST repository at https://www.nist.gov/itl/products-and-services/emnist-dataset (FEMNIST), and the University of Toronto's repository at https://www.cs.toronto.edu/~kriz/cifar.html (CIFAR-10). Author Contribution Statement Deebakkarthi Chinnasame Rani: Methodology, Software, Validation, Formal analysis, Investigation, Resources, Writing – original draft, Writing – review & editing. Gowtham Ramesh: Conceptualization, Methodology, Software, Validation, Formal analysis, Investigation, Resources, Writing – original draft, Writing – review & editing, Supervision, Project administration. Sountharrajan Sehar: Validation, Formal analysis, Investigation, Resources, Data curation, Writing – review & editing, Project administration. Varun Aiyaswamy Kannan: Formal analysis, Resources, Data curation. Elambharathi Padmavathi Thangavel: Formal analysis, Writing – original draft, Visualization. Bharath Kumar Kanapareddy: Investigation, Resources, Data curation.

Deebakkarthi Chinnasame Rani, Gowtham Ramesh, Sountharrajan Sehar et al. · 0 citations
2025

FedRAM: Federated Reweighting and Aggregation for Multi-Task Learning

FedRAM is proposed, a three-step framework that progressively updates two scalar hyperparameters: the task importance weight and the client aggregation coefficient, where the proxy model serves as an intermediate between the local reference model and the global agent model.

Fan Wu, Xinyu Yan, Jiabei Liu et al. · 0 citations

Related blog posts

MIT News · Artificial Intelligence Aug 27, 2026

Looking beyond natural sequences

A new machine-learning framework aims to improve the success rate of computational protein design while moving away from results that reproduce sequences found in nature.