Stochastic Controlled Averaging for Federated Learning with Communication Compression

Xinmeng Huang; Ping Li; Xiaoyun Li

Source-linked AI summary

Stochastic Controlled Averaging for Federated Learning with Communication Compression

Xinmeng Huang, Ping Li, Xiaoyun Li

arXiv:2308.08165v2Aug 16, 2023math.OCcs.DCcs.LGstat.ML

TL;DR

Compressed federated learning faces challenges from compression distortion, partial participation, and data heterogeneity. The paper proposes SCALLION and SCAFCOM for unbiased and biased compression, respectively, and reports performance comparable to full-precision approaches at substantially reduced communication cost.

Problem
Compressed FL must address compression distortion alongside partial participation and data heterogeneity, while existing approaches have limited accommodation or stringent compression requirements.
Method
The paper proposes SCALLION and SCAFCOM, building on a new stochastic controlled averaging implementation to support unbiased and biased compression.
Results
SCALLION and SCAFCOM perform comparably to full-precision FL approaches at substantially reduced communication cost and match or improve compressed-FL convergence rates and complexities.
Takeaways & Limitations
The proposed methods outperform prior compressed FL methods while supporting both unbiased and biased compression.
Takeaways & Limitations
One compared communication-complexity result requires homogeneous (iid) clients.

Abstract

from arXiv · show

Communication compression, a technique aiming to reduce the information volume to be transmitted over the air, has gained great interests in Federated Learning (FL) for the potential of alleviating its communication overhead. However, communication compression brings forth new challenges in FL due to the interplay of compression-incurred information distortion and inherent characteristics of FL such as partial participation and data heterogeneity. Despite the recent development, the performance of compressed FL approaches has not been fully exploited. The existing approaches either cannot accommodate arbitrary data heterogeneity or partial participation, or require stringent conditions on compression. In this paper, we revisit the seminal stochastic controlled averaging method by proposing an equivalent but more efficient/simplified formulation with halved uplink communication costs. Building upon this implementation, we propose two compressed FL algorithms, SCALLION and SCAFCOM, to support unbiased and biased compression, respectively. Both the proposed methods outperform the existing compressed FL methods in terms of communication and computation complexities. Moreover, SCALLION and SCAFCOM accommodates arbitrary data heterogeneity and do not make any additional assumptions on compression errors. Experiments show that SCALLION and SCAFCOM can match the performance of corresponding full-precision FL approaches with substantially reduced uplink communication, and outperform recent compressed FL methods under the same communication budget.

1 Introduction

Federated learning must balance heterogeneous data, partial participation, and heavy communication costs; compression can worsen client-drift and aggregation errors. This paper develops SCALLION and SCAFCOM to support unbiased and biased compression with reduced communication and broad robustness.

Challenges: Severe data heterogeneity makes local models favor local losses over the global loss, hindering global convergence quality.
Challenges: Partial client participation slows global convergence because only a fraction of clients contribute data and information in each round.
Motivation: Communication compression reduces transmitted information but introduces distortion that interacts with client drift and inaccurate aggregation.
Contributions: The proposed SCAFFOLD reformulation halves uplink communication by requiring each participating client to transmit one increment variable instead of two.
Contributions: SCALLION supports unbiased compression with arbitrary data heterogeneity and partial participation, while achieving state-of-the-art communication and computation complexities.
Contributions: SCAFCOM supports biased compression, and experiments show both methods match full-precision performance with substantially reduced communication while outperforming recent compressed FL methods under equal budgets.

2 Related Work

Communication compression in federated learning uses quantization or sparsification, but compression interacts with biasedness, data heterogeneity, and partial participation. Existing methods may require stringent compressor conditions, whereas the proposed algorithms operate under minimal assumptions and support arbitrary client heterogeneity and partial participation.

Compression: Quantization maps vectors to grid values, while sparsification transmits only a subset of entries; either approach may be unbiased or biased.
Compression and error feedback: Directly using biased compressors may slow convergence or lead to divergence, motivating error-feedback methods for reducing compression distortion.
Federated learning: SCAFFOLD uses control variables to mitigate data heterogeneity and partial participation in federated learning.
Compressed federated learning: Existing compressed FL methods may lack robustness to arbitrary client heterogeneity and partial participation or require assumptions beyond standard unbiased or contractive compressibility.
Proposed methods: The proposed algorithms work under minimal assumptions, accommodate arbitrary client heterogeneity and partial participation, and outperform previous methods theoretically and empirically.
Federated learning with momentum: Momentum has been incorporated into federated learning and empirically shown to enhance performance, while related work studies its role in mitigating client drift.

3 Problem Setup

The problem setup defines federated learning as minimizing an aggregate objective using client-specific data distributions and seeks a stationary point for a generally non-convex objective. Client data may be heterogeneous, so a globally stationary model need not be stationary for individual clients; the analysis uses standard smoothness and stochastic-gradient assumptions.

Objective: Federated learning minimizes an objective built from local client objectives and seeks a stationary point because non-convex optima are generally intractable.
Data heterogeneity: Client data distributions may differ, producing unequal local objectives and making a globally stationary model non-stationary for individual clients.
Data heterogeneity: Homogeneous clients share a common data distribution, yielding identical local objectives across clients.
Assumptions: The convergence analysis assumes L-Lipschitz local gradients and bounded stochastic-gradient variance with iid samples per client.
Assumptions: Smoothness and gradient stochasticity are the only two assumptions required for the paper’s theoretical analysis.

4 SCALLION: Single-round Compressed Communication

SCALLION reformulates SCAFFOLD so each participating client sends one variable, then compresses local-model increments with unbiased compressors. It preserves robustness to heterogeneity and partial participation while matching state-of-the-art complexity without additional compression-error assumptions.

Communication reformulation: SCAFFOLD's equivalent implementation reduces client-to-server communication from two model-sized variables to one increment variable per participating client.The local updates remain unchanged; only the communication procedure is modified.
Supported FL settings: SCALLION supports arbitrary data heterogeneity, local updates, and partial client participation under standard unbiased compressibility.Its control-variable formulation is designed to avoid client drift and maintain nearly synchronized local updates under heterogeneous data.
Relation to SCAFFOLD: With identity compressors and α = 1, SCALLION reduces to SCAFFOLD.The scaling factor α stabilizes control-variable updates and acts as their learning rate.
Theory and complexity: SCALLION matches state-of-the-art asymptotic communication and computation complexities under non-iid clients without relying on a large uniform compression-error bound.The paper reports a superior convergence rate relative to FedCOMGATE despite matching its asymptotic complexities.

5 SCAFCOM: Biased Compression with Momentum

SCAFCOM extends controlled averaging with local momentum to handle biased, contractive compression. It provides convergence and complexity guarantees while supporting heterogeneous data and partial participation more broadly than prior biased-compression methods.

Motivation: SCALLION's analysis does not yield fast convergence rates for biased compressors, motivating the complementary SCAFCOM method.SCAFCOM is introduced specifically to accommodate biased communication compression in federated learning.
Compression model: SCAFCOM uses q2-contractive compressors, which may have smaller squared errors than unbiased compressors but do not preserve unbiasedness.Examples include Top-r and grouped-sign operators.
Momentum-enhanced compression: SCAFCOM adds a client-side momentum variable whose accumulated message has smaller variance under biased compression.The momentum-enhanced update is communicated instead of the corresponding uncorrected quantity.
Algorithmic relations: SCAFCOM transmits one compressed variable in the uplink and recovers SCAFFOLD when β = 1 and the compressors are identity mappings.Setting β = 1 and Ci = I is also stated to recover SCAFFOLD in the algorithmic description.
Convergence and complexity: For biased compression, SCAFCOM attains asymptotic communication complexity N S(1−q)ϵ and computation complexity 1 Sϵ2 for E[∥∇f(ˆx)∥2] ≤ ϵ.The theorem assumes the paper's stated assumptions and learning-rate conditions.
Comparison with prior methods: SCAFCOM improves prior biased-compression communication complexity by at least a factor 1/(1 −q) and supports arbitrary heterogeneity and partial participation.The cited comparison states that several prior methods require homogeneous or bounded-heterogeneity settings, while some converge only under full participation.

6 Experiments

Experiments evaluate SCALLION and SCAFCOM against full-precision and compressed FL baselines on heterogeneous, partially participating MNIST and FMNIST tasks. The proposed methods closely match or outperform baselines while substantially reducing communication.

Datasets, Algorithms and Training Setup: Experiments use MNIST and FMNIST with 200 clients, 20 participants per round, and highly heterogeneous data limited to at most two classes per client.Each participating client performs 10 local steps with mini-batches of 32, and results average five runs.
Results: SCAFCOM outperforms Fed-EF in both training loss and test accuracy at the same compression degree.The paper attributes this comparison to controlled variables and local momentum in SCAFCOM.
Results: SCAFCOM with Top-0.01 achieves test accuracy very close to full-precision SCAFFOLD, while Top-0.05 essentially matches it on both datasets.These results correspond to biased compression experiments.
Results: 20 - 100x uplink communication savings reach the same performance as full-precision SCAFFOLD.The reported savings accompany SCAFCOM’s close or matching accuracy under Top-r compression.
Results: SCALLION outperforms FedCOMGATE at the same compression degree, while its 2-bit and 4-bit curves essentially overlap full-precision SCAFFOLD.The 4-bit compressor provides around 100x communication compression, and 2-bit compression saves more communication.
Results: SCAFCOM converges across β ∈ [0.05, 1], with β = 0.2 performing best on both datasets, whereas SCALLION becomes unstable for α > 0.5 under 2-bit quantization.The paper reports α = 0.1 as SCALLION’s best overall setting in the corresponding experiments.

7 Conclusion

The paper proposes SCALLION and SCAFCOM as compressed FL methods built on a more communication-efficient SCAFFOLD formulation. The methods support unbiased and biased compression, achieve favorable theoretical and empirical results, and match full-precision performance at lower communication cost.

Contributions: SCALLION and SCAFCOM support unbiased and biased compression, respectively, using a new implementation of stochastic controlled averaging with local momentum.The formulation uses a single variable for uplink communication per client.
Theory: The methods match or improve state-of-the-art convergence rates and complexities for compressed FL under minimal assumptions.The paper reports improvements in communication and computation complexities over prior results.
Scope: SCALLION and SCAFCOM are presented as robust to arbitrary data heterogeneity, partial participation, and local updates while accommodating communication compression.The paper describes them as relying solely on standard compressibilities.
Experiments: Experiments show performance comparable to full-precision FL with substantially reduced communication and better performance than recent compressed FL methods under the same budget.These conclusions summarize the empirical evaluation.
Future Work: Future extensions include adaptive optimization, privacy, and fairness in federated learning.These are identified as possible integrations or extensions rather than evaluated capabilities.

A Detailed Implementations of SCAFFOLD

The original SCAFFOLD implementation requires clients to communicate both local-model and control-variable increments. An equivalent implementation reduces this to one uplink communication round per participating client while preserving the training trajectory.

Original Implementation: Original SCAFFOLD requires each participating client to transmit both its local model increment and control-variable increment.These are sent to the server at the end of local updates.
Communication Cost: The original implementation therefore uses two rounds of uplink communication per training iteration.The communication requirement follows the separate transmission of the two increments.
Equivalent Implementation: An equivalent SCAFFOLD implementation uses only a single uplink communication round for each participating client.This formulation is described in Algorithm 4.

B Preliminaries of Proofs

The paper formulates SCALLION and SCAFCOM through stochastic recursions and establishes preliminary lemmas for client sampling, conditional variance, and auxiliary-variable analysis. These ingredients support the methods’ convergence analysis.

Method Recursions: SCALLION and SCAFCOM share the server recursion xt+1 = xt − γ d̃t+1, with γ defined as ηgηlK.The recursion is the common update form for both compressed methods.
Auxiliary Variables: The analysis introduces auxiliary variables and exclusive recursions for SCALLION and SCAFCOM to handle their respective updates.The construction includes Ut and method-specific recursion relations.
Probabilistic Tools: The preliminaries define expectation over all randomness and state lemmas for sampling without replacement and dependent random variables.These lemmas provide mean and variance bounds used in the analysis.
Probabilistic Tools: A conditional version of the dependent-variable bound gives a tighter variance estimate when variables are Markov-correlated.The bound uses conditional means and variances.
Auxiliary Bounds: Lemma 3 supplies bounds under Assumption 1 for convex combinations of vectors and is proved using Lemma 2, Sedrakyan’s inequality, and norm inequalities.The proof derives two upper bounds for the relevant expressions.

C Proof of SCALLION

The section establishes SCALLION’s convergence under unbiased compression through a sequence of lemmas, relaxations, and a final theorem. The proof uses Assumptions 1 and 2 and specific parameter choices to derive the convergence result.

Descent analysis: The descent analysis bounds the deviation between d_t+1 and ∇f(x_t) using prior lemmas, the control term c_t, and Assumption 1.The proof also invokes Young’s inequality to relax intermediate terms.
Proof assembly: The SCALLION proof repeatedly substitutes intermediate inequalities into earlier bounds to control compressed-update errors and complete the auxiliary results.The derivation uses mutual independence and Definition 2 for the compressed quantities.
Auxiliary bounds: The proof begins by establishing auxiliary lemmas under Assumptions 1 and 2 for all t ≥ 0.These lemmas support the subsequent descent and error-control bounds.
Convergence theorem: Theorem 3 assumes mutually independent ω-unbiased compressors and an initialization of c0, then states that SCALLION converges.The theorem’s bound is expressed using ∆ ≜ f(x0) − min f(x), with parameter choices involving γ and α.
Parameter selection: The final derivation uses γ = ηgηlK, α ≤ 1/[4(ω+1)], α = O((1 + ω)^−1), and the initialization conditions to obtain the theorem’s result.Substituting the selected α completes the proof.

D Proof of SCAFCOM

The section proves SCAFCOM’s convergence for biased compression through auxiliary bounds, descent analysis, and parameter selection. The final theorem applies to q^2-contractive compressors under Assumptions 1 and 2.

Descent analysis: The descent analysis controls SCAFCOM’s error terms through Lemma 2, Definition 2, Young’s inequality, and Assumption 1.Several intermediate inequalities are combined to obtain the needed recursion.
Proof assembly: The proof combines intermediate inequalities while using N^−1 ≤ S^−1 ≤ 1 and q^2 ≤ 1 to complete one auxiliary result.The argument also uses the relation 1 ≤ 1/(1 − q) in a later bound.
Convergence theorem: Theorem 4 assumes clients use q^2-contractive compressors and states that SCAFCOM converges.The theorem’s associated objective-gap notation is ∆ ≜ f(x0) − min f(x).
Parameter selection: The final argument averages the derived recursion, defines a Lyapunov function, and selects β after choosing γ and initializing the auxiliary variables.The proof concludes by substituting the chosen β into the final bound.