Achieving Linear Speedup with Partial Worker Participation in Non-IID Federated Learning

TL;DR

This work extends federated learning theory to non-i.i.d. data with partial worker participation by analyzing a generalized FedAvg with two-sided learning rates. It proves that linear convergence speedup is achievable in this challenging setting, with convergence rates $\mathcal{O}\left(\frac{1}{\sqrt{mKT}} + \frac{1}{T}\right)$ for full participation and $\mathcal{O}\left(\frac{\sqrt{K}}{\sqrt{nT}} + \frac{1}{T}\right)$ for partial participation, and shows that local updates can aid convergence when properly tuned, with a maximum effective $K$ of $\frac{T}{m}$. The approach avoids bounded-gradient assumptions and is validated by MNIST and CIFAR-10 experiments. These results offer practical guidance for scaling FL systems with heterogeneous data and intermittent client participation, reducing communication without sacrificing convergence quality.

Abstract

Federated learning (FL) is a distributed machine learning architecture that leverages a large number of workers to jointly learn a model with decentralized data. FL has received increasing attention in recent years thanks to its data privacy protection, communication efficiency and a linear speedup for convergence in training (i.e., convergence performance increases linearly with respect to the number of workers). However, existing studies on linear speedup for convergence are only limited to the assumptions of i.i.d. datasets across workers and/or full worker participation, both of which rarely hold in practice. So far, it remains an open question whether or not the linear speedup for convergence is achievable under non-i.i.d. datasets with partial worker participation in FL. In this paper, we show that the answer is affirmative. Specifically, we show that the federated averaging (FedAvg) algorithm (with two-sided learning rates) on non-i.i.d. datasets in non-convex settings achieves a convergence rate $\mathcal{O}(\frac{1}{\sqrt{mKT}} + \frac{1}{T})$ for full worker participation and a convergence rate $\mathcal{O}(\frac{\sqrt{K}}{\sqrt{nT}} + \frac{1}{T})$ for partial worker participation, where $K$ is the number of local steps, $T$ is the number of total communication rounds, $m$ is the total worker number and $n$ is the worker number in one communication round if for partial worker participation. Our results also reveal that the local steps in FL could help the convergence and show that the maximum number of local steps can be improved to $T/m$ in full worker participation. We conduct extensive experiments on MNIST and CIFAR-10 to verify our theoretical results.

Figures (7)

• Figure 1: Training loss (top) and test accuracy (bottom) for the 2NN model with hyper-parameters setting: local learning rate 0.1, global learning rate 1.0: (a) worker number 100, local steps 5 epochs; (b) local steps 5 epochs; (c) 5 digits in each worker's dataset.
• Figure 2: Training loss (top) and test accuracy (bottom) for three models on MNIST with hyperparameters setting: local learning rate 0.1, global learning rate 1.0, local steps 5 epochs.
• Figure 3: Training loss (top) and test accuracy (bottom) for three models on MNIST with hyperparameters setting: local learning rate 0.1, global learning rate 1.0, worker number 100.
• Figure 4: Training loss (top) and test accuracy (bottom) for three models on MNIST with hyperparameters setting: local learning rate 0.1, global learning rate 1.0, worker number 10.
• Figure 5: Test accuracy with respect to worker number on CIFAR-10 dataset.

Theorems & Definitions (5)

• proof
• proof
• Lemma 1: Unbiased Sampling
• Lemma 2: Lemma 4 in reddi2020adaptive
• proof