Source-linked AI summary

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

Haibo Yang, Minghong Fang, Jia Liu

arXiv:2101.11203v3cs.LGcs.DC

TL;DR

Federated learning’s linear convergence-speedup theory had been limited mainly to i.i.d. data and full worker participation, leaving non-i.i.d. partial participation open. The paper analyzes generalized FedAvg with two-sided learning rates and proves linear speedup in both full and partial participation, while examining local-step limits and practical sampling assumptions.

  • Problem

    Prior linear-speedup analyses were limited mainly to i.i.d. datasets and/or full worker participation, whereas this paper asks whether speedup remains achievable with non-i.i.d. data and partial participation.

  • Method

    The paper analyzes generalized FedAvg with two-sided learning rates for general non-convex optimization under non-i.i.d. data and full or partial worker participation.

  • Results

    The algorithm achieves linear convergence speedup under full and partial participation, with local steps allowed up to K ≤ T/m in the full-participation setting.

  • Takeaways & Limitations

    With proper learning-rate choices, local steps can help convergence while reducing communication overhead through fewer communication rounds.

  • Takeaways & Limitations

    The analysis assumes every enlisted worker is available to participate, so practical worker-state constraints and sampling design require further investigation.

Abstract

from arXiv · show

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.

1 INTRODUCTION

Federated learning must handle heterogeneous data and changing worker availability, challenging prior linear-speedup results based on i.i.d. data or full participation. This paper shows generalized FedAvg with two-sided learning rates achieves linear convergence speedup in non-i.i.d. settings with full and partial participation.

  • Motivation: Non-i.i.d. worker datasets arise from local circumstances and challenge FL algorithm design and performance analysis.
  • Motivation: Time-varying participation makes waiting for every worker infeasible, so servers may select only a subset of workers.
  • Background: FedAvg uses multiple local stochastic-gradient steps between centralized server aggregations, unlike traditional distributed learning’s single step.
  • Research question: The paper answers affirmatively whether linear convergence speedup remains achievable with non-i.i.d. datasets and full or partial participation.
  • Implications: K can be as large as T/m under full participation, while larger local-step counts can reduce communication overhead and help convergence with suitable learning rates.

2 RELATED WORK

Prior work established linear speedup mainly for i.i.d. data and full participation, while this paper analyzes generalized FedAvg for non-i.i.d. data and partial participation. Its algorithm samples active workers, performs local updates, and applies a server update with two-sided learning rates.

  • Prior results: Earlier studies commonly achieved linear speedup under i.i.d. datasets and full worker participation.
  • Assumption: The analysis allows active-worker sets of size at least n, so the fixed-size presentation is not restrictive for the stated results.
  • Algorithm: The proposed generalized FedAvg algorithm is presented with two-sided learning rates for federated optimization.
  • Algorithm: Each round samples n workers, each performs K local updates from the current model, and the server aggregates the resulting updates.
  • Research gap: Prior non-i.i.d. studies often demonstrated FedAvg variants heuristically, while convergence analyses under partial participation remained comparatively limited.

3 LINEAR SPEEDUP OF THE GENERALIZED FEDAVG WITH TWO-SIDED LEARNING RATES FOR NON-IID DATASETS

Generalized FedAvg with two-sided learning rates achieves linear convergence speedup for non-i.i.d. data under both full and partial worker participation. Proper learning-rate choices also allow local steps to help convergence while reducing communication rounds.

  • Algorithm: Two-sided FedAvg uses worker-side local SGD and server-side parameter aggregation with a separate global learning rate.The original FedAvg is recovered when the server-side learning rate equals one.
  • Convergence analysis: Choosing the local learning rate inversely with K controls variance amplification from multiple local updates.Global variability and local stochastic variance are amplified by local steps, motivating ηL = O(1/K).
  • Full worker participation: O(1/√(mKT) + 1/T) convergence is achieved with full worker participation under non-i.i.d. datasets.The result holds for generalized FedAvg under the stated non-convex assumptions and suitable learning rates.
  • Full worker participation: K ≤ T/m local steps are permitted under full participation, improving on prior bounds and potentially reducing communication overhead.More local steps imply relatively fewer communication rounds.
  • Partial worker participation: O(√K/√(nT) + 1/T) convergence is achieved under partial worker participation with either considered uniform sampling strategy.Partial participation preserves the convergence-rate structure but has a larger variance term from fewer workers and random sampling.

4 DISCUSSION

The analysis attributes the convergence behavior to two-sided learning rates, worker sampling, and the interaction between local steps and data heterogeneity. It also identifies worker availability as an important practical boundary.

  • Convergence Rate: Two-sided learning rates control the constant term involving stochastic variance, data heterogeneity, and local-step amplification.Setting ηL inversely with K makes the constant term controllable.
  • Number of Local Steps: K ≤ T/m is the improved maximum local-step bound, and properly tuned local steps can help convergence.The paper notes that other results found local steps could slow convergence, so their effect remains context-dependent.
  • Number of Workers: Increasing the number of workers per communication round improves convergence by making the sampled worker distribution better approximate the overall worker distribution.The paper uses uniform sampling with or without replacement.
  • Practical boundary: The analysis assumes every enlisted worker is available to participate, which may not hold when devices must be charging or idle.The paper leaves joint design of sampling schemes and generalized FedAvg for future investigation.

5 NUMERICAL RESULTS

Experiments on non-i.i.d. MNIST and CIFAR-10 evaluate generalized FedAvg across models, worker participation levels, local steps, and comparison with SCAFFOLD. The reported results broadly support the theoretical analysis.

  • Experimental setup: Experiments use logistic regression, a two-hidden-layer neural network, a CNN, and ResNet on non-i.i.d. MNIST and CIFAR-10.Some experimental results are placed in the supplementary material.
  • Impact of non-i.i.d. datasets: For 2NN on MNIST, generalized FedAvg converges under non-i.i.d. data with suitable learning rates.With five digits per worker, full and partial participation achieve convergence comparable to the i.i.d. ten-digit case.
  • Impact of worker number: Partial participation with n = 10 produces zigzagging, slower convergence than full participation with n = 100 under the same hyper-parameters.The effect is more prominent for highly non-i.i.d. datasets.
  • Impact of local steps: Local steps help convergence for both full and partial worker participation in the reported Figure 1(c) experiments.The paper contrasts this with prior results under unbalanced non-i.i.d. MNIST datasets.
  • Comparison with SCAFFOLD: The SCAFFOLD comparison measures communication rounds, total communication load, and estimated wall-clock time to reach target test accuracy.The comparison uses the same GPU and reports settings for non-i.i.d. digits_2 and i.i.d. digits_10 data.

6 CONCLUSIONS AND FUTURE WORK

The paper proves linear convergence speedup for generalized FedAvg under non-i.i.d. data with full and partial participation. It also shows benefits from local steps while motivating further work on worker sampling and changing participant sets.

  • Conclusions: Generalized FedAvg with two-sided learning rates achieves linear convergence speedup for general non-convex optimization on non-i.i.d. datasets.The result covers both full and partial worker participation.
  • Conclusions: Local steps can help convergence, and the maximum number of local steps under full participation is improved to T/m.The conclusion frames these results as part of the paper’s theoretical understanding of FL.
  • Future work: Future work includes optimal sampling under partial participation and active worker sets whose size and membership vary across communication rounds.The paper identifies these as open research directions.

A APPENDIX I: PROOFS

The appendix presents detailed proofs for both full and partial worker participation.

  • The appendix contains detailed proof material for full worker participation.
  • The appendix also contains detailed proof material for partial worker participation.
  • The proofs are organized into Sections A.1 and A.2.

A.1 PROOF OF THEOREM 1

Theorem 1 is proved for FedAvg with full worker participation by combining smoothness, variance, and bounded-step-size arguments over communication rounds.

  • Theorem 1 setup: Theorem 1 analyzes Algorithm 1 with constant local and global learning rates under full worker participation.The rates satisfy ηL ≤ 1/(8LK) and ηηL ≤ 1/(KL).
  • Descent analysis: The proof takes expectation over communication-round randomness using the smoothness assumption.
  • Descent analysis: The cross term A1 is bounded using a norm identity, while A2 is bounded separately.
  • Variance control: The bounds use the smoothness and bounded-variance assumptions together with auxiliary lemmas and independence of zero-mean terms.
  • Final aggregation: The proof concludes by rearranging the one-round inequality and summing it from t = 0 to T − 1.

A.2 PROOF OF THEOREM 2

Theorem 2 extends the convergence proof to partial worker participation, analyzing two uniform sampling strategies and showing that at least n participating workers suffice for the same rate.

  • Theorem 2 setup: Theorem 2 analyzes Algorithm 1 with partial worker participation under constant global and local learning rates.
  • Randomness and sampling: The proof treats random worker sampling and stochastic gradients as the two sources of randomness.
  • Two sampling strategies: The analysis bounds the partial-participation terms separately for sampling with replacement and without replacement.
  • Final aggregation: The proof obtains the partial-participation convergence results by substituting intermediate bounds, rearranging, and summing over communication rounds.
  • Rate condition: The same convergence rate is guaranteed when each communication round includes at least n workers, rather than exactly n.
  • Randomness and sampling: Both sampling strategies use an unbiased participation estimator.

B APPENDIX II: EXPERIMENTS

The experiments use non-i.i.d. versions of MNIST and CIFAR-10 to examine the paper’s theoretical setting.

  • The experiments use non-i.i.d. datasets across workers.
  • MNIST is included in the experimental datasets.
  • CIFAR-10 is included in the experimental datasets.

B.1 MNIST

The MNIST experiments evaluate non-i.i.d. data partitioning, model architectures, and default training settings across three models.

  • Data and setup: MNIST data are distributed across 100 workers according to digit classes to create controllable statistical heterogeneity.Each worker receives samples from a specified number of digits; p = 1 is highly heterogeneous, whereas p = 10 is essentially i.i.d.
  • Experimental procedure: Hyperparameters are tuned through grid-search experiments, with the tuning results reported in Table 3.The CNN architecture used for MNIST is reported separately in Table 4.
  • Models: The experiments run multinomial logistic regression, a two-hidden-layer fully connected network, and a convolutional neural network.The fully connected network has two hidden layers with 200 ReLU neurons each.
  • Training configuration: The default configuration uses server learning rate η = 1.0, client learning rate ηL = 0.1, K = 10 local epochs, 100 clients, and n = 10.The data are evenly partitioned across clients in the i.i.d. setting and otherwise distributed using the stated non-i.i.d. strategy.

B.3 DISCUSSION

The discussion finds that FedAvg can converge on non-i.i.d. MNIST data, while partial participation increases randomness and local steps provide only modest convergence benefits. Comparisons with SCAFFOLD expose a computation–communication trade-off, and sampling design remains an important practical boundary.

  • Impact of non-i.i.d. datasets: FedAvg converges under non-i.i.d. MNIST datasets with an appropriate learning-rate choice, and p = 5 can approach the i.i.d. convergence speed.This comparable speed is reported for both full and partial worker participation across the examined models.
  • Impact of worker number: Partial worker participation produces zigzagging, slower convergence because sampling gives the server only a coarse estimate of system heterogeneity.The problem is more prominent for highly non-i.i.d. datasets.
  • Impact of worker number: With p = 1 and 10 sampled workers, active datasets may cover only 4 or 5 digits, whereas p = 5 is likely to cover all 10 digits.Broader digit coverage lets the server better mitigate system heterogeneity in each communication round.
  • Practical considerations: The experiments use random equal-probability sampling without replacement, while practical worker states require joint design of sampling, worker count, and FedAvg.The authors identify this joint design as requiring further investigation.
  • Impact of local steps: Local steps help convergence under both full and partial participation, but their effect is slight relative to non-i.i.d. degree and worker number.The discussion reports this pattern for MNIST and CIFAR-10 experiments.
  • Comparison with SCAFFOLD: At ϵ = 75%, SCAFFOLD uses fewer communication rounds and less computation time but approximately twice FedAvg’s communication, prolonging communication time.The comparison uses the same NVIDIA V100 and a fixed 20MB/s upload and download bandwidth.
Loading 2101.11203v3…