[Paper Review] Achieving Linear Speedup with Partial Worker Participation in Non-IID Federated Learning

Haibo Yang, Minghong Fang+1 more|arXiv (Cornell University)|Jan 26, 2021

Privacy-Preserving Technologies in Data29 references73 citations

TL;DR

This paper proves that FedAvg with two-sided learning rates achieves linear convergence speedup under non-IID data and partial worker participation, with local steps up to T/m helping convergence when properly tuned.

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.

Motivation & Objective

• Motivate federated learning where data across workers are non-identically distributed (non-IID) and workers may partially participate.
• Analyze convergence behavior of FedAvg with two-sided learning rates under non-IID settings.
• Determine whether linear speedup with respect to the number of workers is achievable under partial participation.
• Quantify how local update steps K affect convergence and communication efficiency.

Proposed method

• Study a generalized FedAvg where each participating worker performs K local SGD steps with a local learning rate, and the server uses a separate global learning rate to update the model.
• Assume Lipschitz continuous gradients, unbiased local gradient estimators, and bounded local/global variance to model data heterogeneity.
• Derive convergence bounds for both full and partial worker participation in non-convex settings.
• Show that appropriate two-sided learning rates yield convergence rates of O(1/√(mKT) + 1/T) for full participation and O(√K/√(nT) + 1/T) for partial participation.
• Demonstrate that the maximum number of local steps can be as large as K ≤ T/m to achieve linear speedup, and that local steps can be beneficial with proper parameter choices.
• Validate theoretical results with experiments on MNIST and CIFAR-10 across various models.

Experimental results

Research questions

• RQ1 Can linear speedup be achieved in federated learning with non-IID data when only a subset of workers participate in each round?
• RQ2 How do two-sided learning rates (distinct local and server learning rates) impact convergence under non-IID data and partial participation?
• RQ3 What is the relationship between the number of local updates K, total rounds T, and the achieved convergence rate in non-IID FL?
• RQ4 Do local update steps help or hinder convergence under non-IID data, and under what conditions is K up to T/m beneficial?

Key findings

• For non-convex problems with non-IID data, FedAvg with two-sided learning rates achieves a convergence rate of O(1/√(mKT) + 1/T) under full participation and O(√K/√(nT) + 1/T) under partial participation, implying linear speedup for large T.
• The maximum permissible local updates can be increased to K ≤ T/m to maintain linear speedup, improving communication efficiency by reducing rounds.
• Without requiring a bounded gradient assumption, the analysis shows the two-sided learning rate approach controls variance terms and achieves desirable convergence behavior.
• Local updates (larger K) can aid convergence with proper learning-rate scheduling, contrasting some prior work that suggested local steps hurt performance in non-IID FL.
• Experiments on MNIST and CIFAR-10 corroborate the theoretical rates and illustrate effects of non-IIDness, number of participating workers, and local steps on convergence and accuracy.