SCAFFOLD: Stochastic Controlled Averaging for Federated Learning

Sai Praneeth Karimireddy; Satyen Kale; Mehryar Mohri; Sashank J. Reddi; Sebastian U. Stich; Ananda Theertha Suresh

Source-linked AI summary

SCAFFOLD: Stochastic Controlled Averaging for Federated Learning

Sai Praneeth Karimireddy, Satyen Kale, Mehryar Mohri, Sashank J. Reddi, Sebastian U. Stich, Ananda Theertha Suresh

arXiv:1910.06378v4Oct 14, 2019cs.LGcs.DCmath.OCstat.ML

TL;DR

FedAvg’s behavior under heterogeneous data is not fully understood, and client-drift can slow convergence. This paper derives tight rates and proposes SCAFFOLD, which uses control variates to correct drift; it requires fewer communication rounds, is robust to client sampling, and can benefit from client similarity for quadratics.

Problem
FedAvg’s performance on heterogeneous federated data remains an active research problem, including how client-drift affects convergence.
Method
SCAFFOLD uses server and client control variates to estimate client-drift and correct local updates.
Results
SCAFFOLD converges for arbitrarily heterogeneous data, is relatively unaffected by client sampling, and can converge faster when clients are more similar.
Takeaways & Limitations
SCAFFOLD is at least as fast as SGD and can exploit Hessian similarity to reduce communication requirements.
Takeaways & Limitations
The formal result connecting SCAFFOLD’s update to the ideal update is proved for quadratic functions; generalizing it remains an open problem.

Abstract

from arXiv · show

Federated Averaging (FedAvg) has emerged as the algorithm of choice for federated learning due to its simplicity and low communication cost. However, in spite of recent research efforts, its performance is not fully understood. We obtain tight convergence rates for FedAvg and prove that it suffers from `client-drift' when the data is heterogeneous (non-iid), resulting in unstable and slow convergence. As a solution, we propose a new algorithm (SCAFFOLD) which uses control variates (variance reduction) to correct for the `client-drift' in its local updates. We prove that SCAFFOLD requires significantly fewer communication rounds and is not affected by data heterogeneity or client sampling. Further, we show that (for quadratics) SCAFFOLD can take advantage of similarity in the client's data yielding even faster convergence. The latter is the first result to quantify the usefulness of local-steps in distributed optimization.

1. Introduction

Federated learning keeps data distributed across many clients, making unreliable networks, limited client availability, and heterogeneous data central optimization challenges. The paper analyzes FedAvg and introduces SCAFFOLD to correct client-drift, reduce communication, and exploit client similarity.

Federated learning trains a centralized server model while keeping training data distributed across clients such as phones, sensors, and hospitals.
FedAvg addresses communication bottlenecks by performing multiple local updates before sending client updates to the server.
SCAFFOLD uses server and client control variates to estimate and correct client-drift in local updates.
SCAFFOLD overcomes heterogeneity, converges in significantly fewer communication rounds, and is relatively unaffected by client sampling.
For similar clients, SCAFFOLD requires fewer communication rounds, quantifying an advantage of local steps over large-batch SGD.
The paper confirms its theoretical results on simulated and extended MNIST datasets.

2. Setup

The paper formulates federated optimization as minimizing a sum of stochastic client functions and introduces assumptions describing gradient and Hessian dissimilarity. These assumptions characterize distinct forms of client heterogeneity.

Federated optimization is formalized as minimizing a sum of stochastic client functions with access only to stochastic samples.
Each client function f_i represents that client’s loss, and the results extend to weighted aggregation.
The setup assumes a lower-bounded global objective, β-smooth client functions, unbiased stochastic gradients, bounded within-client variance, and optional strong convexity.
Bounded gradient dissimilarity (A1) and bounded Hessian dissimilarity (A2) capture different forms of client heterogeneity.
A1 and A2 are orthogonal: gradient dissimilarity can be absent when Hessian dissimilarity is large, and vice versa.

3. Convergence of FedAvg

FedAvg performs local client updates followed by server aggregation, but heterogeneous objectives make local updates drift toward client optima and slow convergence. The section derives tighter bounds and lower bounds, motivating SCAFFOLD’s correction mechanism.

FedAvg procedure: FedAvg samples clients, performs K local updates from the current server model, and aggregates their updates using local and global step-sizes.
Rate comparisons: Table 2 compares communication rounds for strongly convex and non-convex objectives under gradient and Hessian dissimilarity assumptions.
FedAvg convergence: FedAvg’s heterogeneous-client rates are tighter than prior results, but SGD remains faster in the summarized comparisons.
FedAvg convergence: FedAvg’s improved bounds use separate local and global step-sizes, reducing client-drift through smaller local steps while maintaining progress with a larger global step-size.
Client-drift: Client-drift occurs because local updates move toward individual client optima rather than the global optimum, slowing convergence as gradient dissimilarity increases.
SCAFFOLD: SCAFFOLD uses control variates in its local update to correct client-drift and mitigate its convergence effects.
Lower bounds: Matching lower bounds show that FedAvg can be slower than SGD even with full-batch gradients and participation from every client.

4. SCAFFOLD algorithm

SCAFFOLD corrects client-drift by maintaining server and client control variates, using their difference to adjust local updates before aggregation. Stateful client variates enable synchronization despite heterogeneous data.

State and initialization: SCAFFOLD maintains a server control variate c and a client control variate c_i for each client.These states accompany the server model and are initialized consistently across clients.
Local computation: Each participating client copies the server model, performs K corrected local updates, and then updates its client control variate.The server communicates its parameters to sampled clients at each communication round.
Control-variate options: Option I computes a fresh gradient at the server model, whereas Option II reuses previously computed gradients and is cheaper to compute.The experiments use Option II, although Option I can be more stable depending on the application.
Drift correction: The correction c − c_i estimates client-drift and adjusts each local update toward the server’s update direction.This control-variate strategy is intended to correct differences between local and global gradients.
Effect of the correction: The corrected local updates remain synchronized and converge for arbitrarily heterogeneous clients.If every client control variate is set to zero, SCAFFOLD becomes equivalent to FEDAVG.

5. Convergence of SCAFFOLD

SCAFFOLD’s convergence guarantees do not depend on client similarity and match the performance of SGD with K times larger batch size in the no-sampling strongly convex case. Its correction removes the dependence on function dissimilarity from the update error.

General convergence: SCAFFOLD’s convergence rate holds for arbitrarily heterogeneous client functions without assuming similarity between them.The theorem is stated for any β-smooth functions {f_i}.
Strongly convex rates: In the strongly convex no-sampling case, SCAFFOLD matches SGD with K times larger batch size and is therefore at least as fast as SGD.The stated rates also match known distributed-optimization lower bounds up to acceleration.
Non-convex rates: For non-convex objectives, SCAFFOLD’s rates match those of SAGA, while client-sampled standard SGD converges at a sub-linear rate even when σ = 0.With σ = 0, K = 1, and S = 1, Option I reduces to SAGA.
Mechanism: FEDAVG’s local update differs from the ideal global-gradient update by gradient dissimilarity, whereas SCAFFOLD uses ∇f_i(y) − c_i + c.The correction changes the discrepancy so it depends on the quality of c_i approximating the client gradient.
Mechanism: SCAFFOLD’s update error is independent of client-function similarity and instead depends on how well c_i approximates ∇f_i(y).Smoothness supports approximating the client gradient because it does not change too quickly.

6. Usefulness of local steps

For quadratic objectives, SCAFFOLD can exploit similarity between clients to reduce communication complexity, with Hessian similarity—not proximity of client optima—determining the benefit. The generalization beyond quadratics remains open.

Quadratic convergence: For β-smooth quadratic functions with bounded Hessian dissimilarity δ, Theorem IV gives a convergence guarantee for SCAFFOLD without client sampling.The theorem assumes S = N and targets expected error below ϵ.
Similarity benefits: When σ = 0 and K is large, SCAFFOLD’s complexity becomes δ/µ, showing that greater client similarity can reduce communication requirements.The comparison concerns quadratic functions and large numbers of local steps.
Limiting cases: For identical functions, SCAFFOLD and FEDAVG improve linearly as the number of local steps increases; for arbitrarily different functions, their complexity matches heterogeneous lower bounds.The identical-function case has δ = 0, while the maximally dissimilar case may have δ = 2β.
Choosing local steps: The optimal number of local steps is K = β/δ; smaller K increases communication, while larger K wastes computation.This trade-off follows from the complexity bound for somewhat similar functions.
Role of similarity: Hessian similarity determines the benefit of local steps even when client optima are arbitrarily far apart and gradient dissimilarity is unbounded.Thus, similar curvature rather than similar minimizers is the relevant structural condition in this result.
Scope: Generalizing the similarity-based result beyond quadratic functions is identified as a challenging open problem.The formal approximation result is proved for quadratic functions.

7. Experiments

Experiments show that SCAFFOLD consistently outperforms SGD and FedAvg, while the value of local steps depends on both client similarity and algorithm. SCAFFOLD remains robust to heterogeneity and client sampling and benefits from increased similarity.

Overall results: SCAFFOLD consistently outperforms SGD, FedAvg, and FedProx across the tested parameter regimes.The experiments cover simulated and real datasets, including logistic regression and neural networks.
Simulated results: SCAFFOLD is unaffected by gradient heterogeneity in simulation and improves with more local steps, with K = 2 twice as fast as SGD and K = 10 about 5 times faster.FedAvg instead slows as local steps and gradient dissimilarity increase.
Local steps and heterogeneity: For heterogeneous clients, SCAFFOLD remains faster than SGD while FedAvg can become slower as local steps increase.With more than five epochs at zero similarity, FedAvg becomes slower than SGD, whereas SCAFFOLD remains better than SGD.
Client similarity: As client similarity increases, SCAFFOLD and FedAvg become faster and benefit more from additional local steps, while SGD remains relatively constant.The observed improvement is associated with decreasing gradient and Hessian dissimilarity.
Client sampling: SCAFFOLD is resilient to client sampling, showing only a sub-linear slowdown as the sampled-client fraction decreases, especially for more similar clients.Table 4 reports communication rounds to reach 0.45 test accuracy with five epochs fixed.
Non-convex results: On non-convex experiments, SCAFFOLD reaches higher final test accuracy than FedAvg, although broader experiments are needed before drawing conclusions.The neural-network comparison is described as preliminary and beyond the current experimental scope for definitive conclusions.

8. Conclusion

The paper analyzes how heterogeneity affects federated optimization, showing that FedAvg can be slowed by gradient dissimilarity and proposing SCAFFOLD to address it. Experiments and theory support SCAFFOLD’s convergence advantages, including benefits from client similarity.

Conclusion: FedAvg can be severely hampered by gradient dissimilarity and may be slower than SGD.The paper attributes this slowdown to client-drift in heterogeneous federated optimization.
Conclusion: SCAFFOLD uses control variates to overcome gradient dissimilarity.Its control-variates correction targets the client-drift created by heterogeneous local updates.
Conclusion: On EMNIST with a non-convex neural network, SCAFFOLD achieved the best test accuracy after 1k rounds, while SGD achieved the least.The experiment used 5 epochs per round and sampled 20% of clients each round.
Conclusion: SCAFFOLD is at least as fast as SGD and can be much faster with favorable Hessian dissimilarity.The paper connects this speedup to similarity in client data for the studied quadratic setting.
Related work: The paper relates SCAFFOLD to distributed variance-reduction methods while distinguishing its applicability to the stochastic setting.It describes SCAFFOLD as an improved version of DANE using stochastic gradient steps instead of a proximal update.

D.2. Rates of convergence (Theorem I)

The section states FedAvg convergence guarantees under smoothness and bounded-gradient assumptions across strongly convex, general convex, and non-convex cases.

D.2. Rates of convergence (Theorem I): FedAvg’s expected error is bounded under β-smooth functions satisfying assumption A1.The theorem provides parameter conditions involving local and global step-sizes and communication rounds.
D.2. Rates of convergence (Theorem I): Theorem V covers strongly convex, general convex, and non-convex client objectives with separate step-size and round conditions.The strongly convex case assumes µ > 0; the general convex case sets µ = 0.

D.3. Proof of convergence

The proof decomposes FedAvg’s progress into convergence and drift terms, then bounds local-update drift under smoothness and related assumptions to derive the stated rates.

D.3. Proof of convergence: FedAvg’s server update is analyzed by separating one-round progress from drift caused by local updates.The drift term is explicitly defined and the resulting inequalities are combined to establish convergence.
D.3. Proof of convergence: A bounded-drift lemma controls FedAvg updates when the local step-size satisfies ηl ≤ 1/((1+B^2)8βKηg).The bound relies on assumptions A1 and A3–A5.
D.3. Proof of convergence: The resulting recursive bound is converted into convex and non-convex convergence rates using weighted averaging and the sub-linear convergence lemma.The non-convex proof follows the same general technique after applying the sub-linear rate result.

D.4. Lower bound for FEDAVG (Theorem II)

The lower-bound analysis constructs heterogeneous strongly convex examples showing that FedAvg’s client-drift cannot generally be removed by step-size selection. The section also states SCAFFOLD’s convergence results and implementation structure.

D.4. Lower bound for FEDAVG (Theorem II): The lower-bound theorem applies to FedAvg with ηg = 1, K > 1, and arbitrary possibly adaptive positive step-sizes.The stated setup also assumes non-adaptive convex-combination aggregation of client updates.
D.4. Lower bound for FEDAVG (Theorem II): FedAvg’s lower bound uses heterogeneous strongly convex functions whose averaged objective is strongly convex.The construction uses simple one-dimensional client functions parameterized by µ and G.
D.4. Lower bound for FEDAVG (Theorem II): Client-drift creates an error term that increases with the step-size, while ordinary gradient-step convergence improves with it.The proof shows that balancing these opposing effects cannot eliminate the contribution of gradient dissimilarity G.
E. Convergence of SCAFFOLD: SCAFFOLD’s general convex rate incurs only an additive term of order 1/R, while its strongly convex dependence appears only in logarithmic terms.The section presents these consequences as part of SCAFFOLD’s convergence analysis.
E. Convergence of SCAFFOLD: SCAFFOLD maintains server and client control variates while sampling a subset of clients each round.Client updates are identical to those with all clients participating, although only sampled clients implement updates.

E.1. Convergence of SCAFFOLD for convex functions (Theorem III)

SCAFFOLD’s convex analysis bounds server-update variance, control-lag, and client-drift before combining them into one-round progress and convergence rates.

Proof strategy: The analysis decomposes convergence into server-update variance, control-lag from client sampling, and client-drift from local updates.These components are bounded separately and then combined to establish progress in a single communication round.
Convergence analysis: Under suitable step-size conditions, the analysis bounds client-drift and derives one-round progress for convex and non-convex objectives.The convex drift bound appears in Lemma 14, while Lemma 15 combines the error bounds to characterize progress.
Convergence rates: The resulting proof framework yields final strongly convex and general convex convergence rates by unrolling the recursive progress bound.The strongly convex and general convex rates are obtained from Lemma 15 using the stated step-size choices.

F. Usefulness of local steps (Theorem IV)

For quadratic client objectives, SCAFFOLD’s convergence analysis quantifies how similarity between clients can make local steps useful, including under heterogeneous curvature.

Theorem IV: Theorem VIII analyzes SCAFFOLD on quadratic functions without client sampling and gives rates for strongly convex and general convex cases.The theorem uses ηg = 1 and specifies local-step-size and communication-round conditions for each case.
Similarity and heterogeneity: δ = 0 matches the Woodworth et al. lower bound up to acceleration, even when client functions can remain highly heterogeneous.The paper notes that objective perturbation can produce high gradient dissimilarity without Hessian dissimilarity.
Non-convex scope: For non-convex quadratics, the paper reports first-of-kind similarity-based rates but assumes f⋆ lower-bounds the returned iterate’s function value.This assumption implicitly requires the iterates to remain bounded because non-convex quadratics lack a global lower bound.
Proof strategy: The quadratic analysis relies on specialized local updates, averaging arguments, and bounds for client drift and output variance.The proof develops these bounds in Lemmas 20–24 and uses weak convexity when averaging parameters.
Quadratic update: The simplified quadratic update uses gradients at the previous global iterate as control variates and averages client iterates at round ends.This update differs from the usual variants in its control-variate evaluation point.