Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

FedMuon: Federated Learning with Bias-corrected LMO-based Optimization

Yuki Takezawa, Anastasia Koloskova, Xiaowen Jiang, Sebastian U. Stich

TL;DR

We address the challenge of training neural networks in federated settings using Muon, an optimizer built on a linear minimization oracle (LMO). Naïve integration (LocalMuon) can fail to converge due to LMO bias, so we propose FedMuon with a bias-correction mechanism and provide convergence guarantees, including for inexact LMO solved via Newton-Schulz iterations. Theoretical results show FedMuon converges to a stationary point, with rates close to FedAvg/SCAFFOLD and potential gains when LMO accuracy increases; the dependence on norm choice and Hessian spectrum is analyzed. Empirically, FedMuon outperforms state-of-the-art adaptive federated learning methods on FashionMNIST and CIFAR-10, including under data heterogeneity, validating the approach's practical impact for scalable distributed training with LMO-based optimizers.

Abstract

Recently, a new optimization method based on the linear minimization oracle (LMO), called Muon, has been attracting increasing attention since it can train neural networks faster than existing adaptive optimization methods, such as Adam. In this paper, we study how Muon can be utilized in federated learning. We first show that straightforwardly using Muon as the local optimizer of FedAvg does not converge to the stationary point since the LMO is a biased operator. We then propose FedMuon which can mitigate this issue. We also analyze how solving the LMO approximately affects the convergence rate and find that, surprisingly, FedMuon can converge for any number of Newton-Schulz iterations, while it can converge faster as we solve the LMO more accurately. Through experiments, we demonstrated that FedMuon can outperform the state-of-the-art federated learning methods.

FedMuon: Federated Learning with Bias-corrected LMO-based Optimization

TL;DR

We address the challenge of training neural networks in federated settings using Muon, an optimizer built on a linear minimization oracle (LMO). Naïve integration (LocalMuon) can fail to converge due to LMO bias, so we propose FedMuon with a bias-correction mechanism and provide convergence guarantees, including for inexact LMO solved via Newton-Schulz iterations. Theoretical results show FedMuon converges to a stationary point, with rates close to FedAvg/SCAFFOLD and potential gains when LMO accuracy increases; the dependence on norm choice and Hessian spectrum is analyzed. Empirically, FedMuon outperforms state-of-the-art adaptive federated learning methods on FashionMNIST and CIFAR-10, including under data heterogeneity, validating the approach's practical impact for scalable distributed training with LMO-based optimizers.

Abstract

Recently, a new optimization method based on the linear minimization oracle (LMO), called Muon, has been attracting increasing attention since it can train neural networks faster than existing adaptive optimization methods, such as Adam. In this paper, we study how Muon can be utilized in federated learning. We first show that straightforwardly using Muon as the local optimizer of FedAvg does not converge to the stationary point since the LMO is a biased operator. We then propose FedMuon which can mitigate this issue. We also analyze how solving the LMO approximately affects the convergence rate and find that, surprisingly, FedMuon can converge for any number of Newton-Schulz iterations, while it can converge faster as we solve the LMO more accurately. Through experiments, we demonstrated that FedMuon can outperform the state-of-the-art federated learning methods.

Paper Structure

This paper contains 34 sections, 17 theorems, 108 equations, 3 figures, 3 tables, 4 algorithms.

Table of Contents

  1. Introduction
  2. Notation:
  3. Preliminary
  4. Federated Learning:
  5. Optimizer with Linear Minimization Oracle:
  6. LocalMuon does not Always Converge
  7. FedMuon
  8. Convergence Analysis
  9. Assumptions
  10. Convergence Result
  11. Discussion:
  12. FedMuon with Inexact LMO
  13. Discussion:
  14. Comparison with Existing Analysis with Inexact LMO:
  15. Proof Sketch:
  16. ...and 19 more sections

Key Result

Theorem 1

For simplicity, we consider the initialization ${\bm{M}}_i^{(0)} = 0$. There exist convex functions $\{ f_i \}_{i=1}^n$ such that for any $r \geq 1$ rounds, the output of LocalMuon (eq:simple_rule_1eq:simple_rule_2eq:simple_rule_3) is the same as the initial parameter and does not converge to the op where $\zeta^2_\star \coloneqq \frac{1}{n} \sum_{i=1}^n \left\| \nabla f_i ({\bm{X}}^\star) \right\

Figures (3)

  • Figure 1: Homogeneous Case ($\beta=10.0$)
  • Figure 2: Heterogeneous Case ($\beta=0.1$)
  • Figure 4: Training curves of FedMuon with various number of Newton-Schulz iterations. We used FashionMNIST and LeNet.

Theorems & Definitions (37)

  • Theorem 1
  • Remark 1: conway2019course
  • Remark 2
  • Example 1
  • Theorem 2
  • Theorem 3
  • Remark 3
  • Remark 4
  • proof
  • Lemma 1
  • ...and 27 more