Table of Contents
Fetching ...

Momentum Benefits Non-IID Federated Learning Simply and Provably

Ziheng Cheng, Xinmeng Huang, Pengfei Wu, Kun Yuan

TL;DR

This paper demonstrates that incorporating momentum allows FedAvg to converge without relying on the assumption of bounded data heterogeneity even using a constant local learning rate, and develops new variance-reduced extensions of FedAvg and SCAFFOLD, which exhibit state-of-the-art convergence rates.

Abstract

Federated learning is a powerful paradigm for large-scale machine learning, but it faces significant challenges due to unreliable network connections, slow communication, and substantial data heterogeneity across clients. FedAvg and SCAFFOLD are two prominent algorithms to address these challenges. In particular, FedAvg employs multiple local updates before communicating with a central server, while SCAFFOLD maintains a control variable on each client to compensate for ``client drift'' in its local updates. Various methods have been proposed to enhance the convergence of these two algorithms, but they either make impractical adjustments to the algorithmic structure or rely on the assumption of bounded data heterogeneity. This paper explores the utilization of momentum to enhance the performance of FedAvg and SCAFFOLD. When all clients participate in the training process, we demonstrate that incorporating momentum allows FedAvg to converge without relying on the assumption of bounded data heterogeneity even using a constant local learning rate. This is novel and fairly surprising as existing analyses for FedAvg require bounded data heterogeneity even with diminishing local learning rates. In partial client participation, we show that momentum enables SCAFFOLD to converge provably faster without imposing any additional assumptions. Furthermore, we use momentum to develop new variance-reduced extensions of FedAvg and SCAFFOLD, which exhibit state-of-the-art convergence rates. Our experimental results support all theoretical findings.

Momentum Benefits Non-IID Federated Learning Simply and Provably

TL;DR

This paper demonstrates that incorporating momentum allows FedAvg to converge without relying on the assumption of bounded data heterogeneity even using a constant local learning rate, and develops new variance-reduced extensions of FedAvg and SCAFFOLD, which exhibit state-of-the-art convergence rates.

Abstract

Federated learning is a powerful paradigm for large-scale machine learning, but it faces significant challenges due to unreliable network connections, slow communication, and substantial data heterogeneity across clients. FedAvg and SCAFFOLD are two prominent algorithms to address these challenges. In particular, FedAvg employs multiple local updates before communicating with a central server, while SCAFFOLD maintains a control variable on each client to compensate for ``client drift'' in its local updates. Various methods have been proposed to enhance the convergence of these two algorithms, but they either make impractical adjustments to the algorithmic structure or rely on the assumption of bounded data heterogeneity. This paper explores the utilization of momentum to enhance the performance of FedAvg and SCAFFOLD. When all clients participate in the training process, we demonstrate that incorporating momentum allows FedAvg to converge without relying on the assumption of bounded data heterogeneity even using a constant local learning rate. This is novel and fairly surprising as existing analyses for FedAvg require bounded data heterogeneity even with diminishing local learning rates. In partial client participation, we show that momentum enables SCAFFOLD to converge provably faster without imposing any additional assumptions. Furthermore, we use momentum to develop new variance-reduced extensions of FedAvg and SCAFFOLD, which exhibit state-of-the-art convergence rates. Our experimental results support all theoretical findings.
Paper Structure (42 sections, 22 theorems, 120 equations, 8 figures, 2 tables, 4 algorithms)

This paper contains 42 sections, 22 theorems, 120 equations, 8 figures, 2 tables, 4 algorithms.

Key Result

Theorem 1

Under Assumption asp:smooth and asp:sgd_var, if we set $g^0=0$, $\beta$, $\gamma$, and $\eta$ as in equation eqn:fedavg-m-para, FedAvg-M enjoys where $\Delta \triangleq f(x^0)-\min_x f(x)$ and $\lesssim$ absorbs numeric numbers. See proof in Appendix app:fedavgm.

Figures (8)

  • Figure 1: Test loss of three-layer MLP versus the number of communication rounds
  • Figure 2: Test accuracy of ResNet18 versus the number of communication rounds
  • Figure 3: Test loss of ResNet18 versus the number of communication rounds
  • Figure 4: Comparing the test accuracy of VR methods with ResNet-18
  • Figure 5: Test loss of MNIST versus the number of communication rounds
  • ...and 3 more figures

Theorems & Definitions (39)

  • Theorem 1
  • Theorem 2
  • Theorem 3
  • Theorem 4
  • Lemma 5
  • proof
  • Lemma 6: karimireddy2020scaffold
  • Lemma 7
  • proof
  • Lemma 8
  • ...and 29 more