Stochastic Controlled Averaging for Federated Learning with Communication Compression
Xinmeng Huang, Ping Li, Xiaoyun Li
Introduction
Federated learning (FL) is a powerful paradigm for large-scale machine learning (Konečnỳ et al., 2016; McMahan et al., 2017). In situations where data and computational resources are dispersed among diverse clients such as phones, tablets, sensors, banks, hospitals, and other devices and agents, federated learning facilitates local data processing and collaboration among these clients (Kairouz et al., 2021). FL enjoys the advantage of distributed optimization on the efficiency of computational resources as the local clients conduct computations simultaneously. Moreover, since the centralized model is trained without transmitting decentralized data from clients directly to servers, FL provides the first layer of protection of data privacy as the local data never leaves the local device.
Due to its nature and application scenarios, federated learning encounters several significant challenges in algorithmic development and theory (Yang et al., 2020; Li and Li, 2023):
Severe data heterogeneity. Unlike in classic distributed training, the local data distribution in FL can vary significantly (i.e., non-iid clients), reflecting practical scenarios where local data held by clients is highly personalized (Zhao et al., 2018; Kairouz et al., 2021; Yuan et al., 2021a; Li et al., 2022a). When multiple local training steps are taken, the local models become “biased” toward minimizing the local losses instead of the global loss, hindering the convergence quality of the global model (Mohri et al., 2019; Li et al., 2020c, a).
Partial client participation. Another practical issue in FL systems is partial participation, where not all clients can always join the training, e.g., due to unstable connections or active selection (Li et al., 2020a). Consequently, only a fraction of clients are involved in each FL training round to interact with the central server. This slows down the convergence of the global model because of less accessible data/information per round (Charles et al., 2021; Chen et al., 2022; Li and Li, 2023).
Heavy communication workload. The cost of model transmission can be a major challenge in FL systems with limited bandwidth (e.g., portable wireless devices), especially when training large models with millions or billions of model parameters. Therefore, communication compression, a technique that aims to reduce the volume of information transmitted, has gained growing research interests in FL (Basu et al., 2019; Reisizadeh et al., 2020; Haddadpour et al., 2021; Li and Li, 2023).
The classic FL approach, FedAvg (Konečnỳ et al., 2016; McMahan et al., 2017; Stich, 2019; Yu et al., 2019a; Lin et al., 2020; Wang and Joshi, 2021), performs multiple gradient-descent steps within each accessible client before communicating with the central server. While showing success in certain scenarios, FedAvg is notably hampered by data heterogeneity and partial client participation (Karimireddy et al., 2020b; Li et al., 2020c; Yang et al., 2021) due to the “client drift” effect. Furthermore, when communication compression is employed, the adverse effect of data heterogeneity can be amplified due to the interplay of client drift and inaccurate message aggregation caused by compression (Basu et al., 2019; Reisizadeh et al., 2020; Haddadpour et al., 2021; Gao et al., 2021; Malekijoo et al., 2021; Li and Li, 2023); see Figure 1 for illustration. The inaccurate aggregation incurred by compression imposes more obstacles to obtaining stable and robust performance in FL systems, particularly when data heterogeneity is severe, and compressors are biased (i.e., the compressed output is a biased estimate of input) (Li and Li, 2023; Gao et al., 2021; Basu et al., 2019).
While having the potential to reduce communication costs in FL, communication compression brings forth new challenges in addition to FL’s inherent characteristics like local updates and partial participation. This naturally raises the question regarding the utility of compressed FL approaches:
Can we design FL approaches that accommodate arbitrary data heterogeneity, local updates, and partial participation, as well as support communication compression?
Despite several attempts, none of the existing algorithms have successfully achieved this goal, to the best of our knowledge. For instance, FedPAQ (Reisizadeh et al., 2020), FedCOM (Haddadpour et al., 2021), QSPARSE-SGD (Basu et al., 2019), Local-SGD-C (Gao et al., 2021) consider compressed FL algorithms under homogeneous data (i.e., iid clients). FedCOMGATE (Haddadpour et al., 2021), designed for unbiased compressors, does not support biased compressors, and their analysis does not validate the utility under partial client participation. Fed-EF (Li and Li, 2023) focuses on biased communication compression in FL with error feedback (Seide et al., 2014; Karimireddy et al., 2019) and partial client participation. However, the convergence analysis requires the assumption of bounded gradient dissimilarity on the data heterogeneity and shows an extra slow-down factor in the convergence rate under partial client participation, suggesting a theoretical limitation of error feedback in FL. Moreover, both Haddadpour et al. (2021) and Li and Li (2023) impose stringent conditions on compression errors (see Remark 2 for more details).
Given these limitations, the motivation of this work is to develop new compressed FL approaches that are practical to implement, robust to data heterogeneity and partial participation, support both biased and unbiased compressors, and exhibit superior theoretical convergence.
In this paper, we propose two algorithms, SCALLION and SCAFCOM, which cover unbiased and biased compression and offer enhanced communication efficiency, faster convergence rates, and robustness to arbitrary data heterogeneity and partial participation. Table 1 presents a comprehensive comparison of communication and computation complexities and associated restrictions of existing algorithms, as well as our newly proposed approaches.
It is worth emphasizing that, the theoretical analysis in our paper only requires the smoothness of local objectives and bounded variance of stochastic gradients, without any additional assumptions on data heterogeneity or compression errors (see Remark 2 for more details), as opposed to all prior related works. The keys to this significant improvement are our new formulation of stochastic controlled averaging and the introduction of momentum. The main contributions are:
We revisit the SCAFFOLD method (Karimireddy et al., 2020b) by proposing a simplified and more communication-efficient formulation. The new implementation reduces the uplink communication cost by half, requiring each client to transmit only one increment variable (of the same size as the model) when participating in a training round, instead of two variables in the original implementation (Karimireddy et al., 2020b).
Building upon our new formulation of SCAFFOLD, we propose the SCALLION method that employs unbiased compressors for the communication of increment variables. We establish its convergence result for non-convex objectives. SCALLION obtains the state-of-the-art communication and computation complexities for FL under unbiased compressors and supports partial client participation.
We further develop SCAFCOM which enables biased compressors for broader applications. Local momentum is applied to guarantee fast convergence and improve empirical performance. The communication and computation complexities of SCAFCOM improve prior results by significant margins, particularly when compression is aggressive.
We conduct experiments to illustrate the effectiveness of SCALLION and SCAFCOM and support our theories. Our empirical results show that the proposed methods achieve comparable performance to full-precision FL methods with substantially reduced communication costs, and outperform recent compressed FL methods under the same communication budget.
Related Work
Two popular approaches are commonly employed to compress communication in distributed systems: quantization and sparsification. Quantization involves mapping input vectors to a set of grid values, and the output can be either unbiased (random dithering) or biased (deterministic dithering) of the input value. Notable examples include Sign-SGD (Seide et al., 2014; Bernstein et al., 2018), low-bit fixed rounding (Dettmers, 2016), Q-SGD (Alistarh et al., 2017), TurnGrad (Wen et al., 2017), and natural compression (Horvóth et al., 2022). On the other hand, sparsification operators only transmit a small subset of entries from the input vector, which can also be unbiased or biased (Wangni et al., 2018; Stich et al., 2018). Theoretical analyses of biased compressors often impose stringent assumptions, such as bounded gradients (Karimireddy et al., 2019; Zhao et al., 2019; Beznosikov et al., 2020) due to the challenges incurred by biasedness. A more detailed summary of unbiased and biased compressors can be found in Huang et al. (2022); Safaryan et al. (2022); He et al. (2023b), among others.
Usually, in distributed training, unbiased compressors can be applied in place of the full-precision gradients to get reasonable theoretical rates and empirical performance. However, directly using biased compressors may slow down convergence or even lead to divergence (Beznosikov et al., 2020; Li and Li, 2023). To alleviate the information distortion caused by compression, the technique of error feedback (EF) was first proposed in Seide et al. (2014). Error feedback has proven particularly effective in addressing biased compressors (Stich et al., 2018; Karimireddy et al., 2019), and it has inspired numerous subsequent distributed approaches (e.g., Wu et al., 2018; Alistarh et al., 2018; Li et al., 2022b). Moreover, a variant scheme of error feedback called EF21 was introduced recently (Richtárik et al., 2021). EF21 compresses increments of deterministic gradients and offers superior theoretical guarantees compared to vanilla error feedback.
Federated learning with compression.
Federated learning has gained great prominence since the introduction of FedAvg, proposed by McMahan et al. (2017) to improve the communication efficiency of classic distributed training. Subsequent studies have explored its theoretical convergence and empirical performance, revealing its susceptibility to data heterogeneity (i.e., non-iid clients) due to the “client-drift” effect, particularly when not all clients participate in training (Stich, 2019; Yu et al., 2019b; Wang and Joshi, 2021; Lin et al., 2020; Wang et al., 2020b; Li et al., 2020c; Yang et al., 2021). Substantial efforts have been made to address client heterogeneity in FL (Liang et al., 2019; Li et al., 2020b, a; Wang et al., 2020a; Zhang et al., 2021; Haddadpour et al., 2021; Yuan and Li, 2022; Alghunaim, 2023; Cheng et al., 2023), and develop other FL protocols involving variance reduction techniques or adaptive optimizers (Karimireddy et al., 2020b; Reddi et al., 2021; Chen et al., 2020; Karimi et al., 2023). Notably, SCAFFOLD introduced by Karimireddy et al. (2020b) leverages control variables to mitigate the impact of data herogeneity and partial client participation.
To further reduce communication costs, communication compression has been integrated into federated learning algorithms, leading to methods such as FedPAQ (Reisizadeh et al., 2020), FedCOMGATE (Haddadpour et al., 2021), Fed-EF (Li and Li, 2023), etc. However, due to the information distortion incurred by compression, the existing communication-compressed FL methods either lack the robustness to arbitrary client heterogeneity and partial participation or rely on stringent conditions of compressors used by clients, going beyond standard unbiased/contractive compressibility. In contrast, our proposed algorithms work under minimal assumptions, which accommodate arbitrary client heterogeneity, partial participation and standard compressibilities while outperforming previous methods theoretically and empirically.
Federated learning with momentum.
The utilization of momentum in optimization traces back to Nesterov’s acceleration (Yurri, 2004) and the heavy-ball method (Polyak, 1964) in deterministic settings, which has been extended to the stochastic scenario (Yan et al., 2018; Yu et al., 2019a; Liu et al., 2020) and other domains (Yuan et al., 2021b; He et al., 2023b, a; Chen et al., 2023). A recent work (Fatkhullin et al., 2023) suggests the benefits of momentum in error feedback for distributed optimization. In the context of federated learning, momentum has been widely incorporated and empirically shown to enhance performance (Wang et al., 2020b; Karimireddy et al., 2020a; Khanduri et al., 2021; Das et al., 2022). For FL, Cheng et al. (2023) demonstrates that momentum can mitigate the client drift phenomenon in FedAvg under full client participation. It is important to note that the algorithms and analysis in this paper are different from these prior works because of the unique challenges posed by the interplay of local updates, partial client participation, and communication compression.
Problem Setup
Formally, in federated learning, we aim to minimize the following objective:
where represents a local data sample of client , represents the loss function evaluated at model and sample , and is the local objective w.r.t. data distribution at client . Since finding the optimum of non-convex objectives is generally intractable, we devote to finding a stationary point of .
In practice, the data distributions across clients may vary significantly, resulting in the inequality for different clients and . Consequently, a globally stationary model with may not be a stationary point of the local objectives, leading to large values of . This phenomenon is widely referred to as data heterogeneity. If all local clients were homogeneous, meaning that the local data samples of different clients follow a common distribution , we would have and each globally stationary model would also be stationary for each client.
where are iid random samples for each client .
Assumptions 1-2 are standard in the analysis of FL algorithms. It is worth highlighting that, these are the only two assumptions required for all the theoretical analysis in this paper.
SCALLION: Single-round Compressed Communication
In this section, we first revisit the seminal SCAFFOLD algorithm (Karimireddy et al., 2020b), which requires communicating two variables (of the same size as the model) from client to server per communication round. We present a new formulation with only a single variable for uplink communication for each client participating in a training round. We then propose SCALLION, which employs unbiased compressors to further reduce the communication workload of SCAFFOLD. SCALLION embraces arbitrary data heterogeneity, local updates, and partial client participation. Theoretical analysis is provided, showing that SCALLION converges at a state-of-the-art rate under standard unbiased compressibility.
The SCAFFOLD approach (Karimireddy et al., 2020b) maintains local control variables on clients and a global control variable on the server. Let (with ) be the set of accessible (active) clients to interact with the server in the -th round. In each training round, SCAFFOLD conducts local updates within each accessible client by
The increments of local model and control variable , of each participating client , are then sent to the central server and aggregated to update the global model parameters:
where is the global learning rate. The detailed description of SCAFFOLD can be found in Appendix A. Notably, the control variables of SCAFFOLD track local gradients such that and , thereby mimicking the ideal update through given and . Consequently, the local updates are nearly synchronized in the presence of data heterogeneity without suffering from client drift.
While the introduction of control variables enables SCAFFOLD to converge robustly with arbitrarily heterogeneous clients and partial client participation, the original implementation of SCAFFOLD described above requires clients to communicate both updates of local models and control variables (also see (Karimireddy et al., 2020b, Alg. 1, line 13)). This results in a doubled client-to-server communication cost and more obstacles to employing communication compression, compared to its counterparts without control variables such as FedAvg.
2 Development of SCALLION
Now we present an equivalent implementation of SCAFFOLD which only requires a single variable for uplink communication and is readily employable for communication compression. Expanding the updates of local models and control variables used by exploiting (1) and (2), we have
In (4) and (5), we see that the updates of local models and control variables share a common component, the increment variables . Since the global control variable is inherently maintained by the server, updates and thus can be recovered by the server upon receiving the increment variables . Therefore, the server model and control variable can be equivalently updated as
Based on our above formulation, by communicating the increment variables and applying the server-side updates (6) accordingly, SCAFFOLD can be implemented equivalently with a halved uplink communication cost, compared to the original one (Karimireddy et al., 2020b). A detailed description of the new implementation can be found in Algorithm 4 in Appendix A. It is also worth noting that the new implementation only modifies the communication procedure, and the same local updates as in Karimireddy et al. (2020b) remain in our implementation.
Importantly, the new implementation of SCAFFOLD provides a simpler and more natural backbone for communication compression as only the transmission of is to be compressed. Moreover, unlike compressing local gradients as adopted in Reisizadeh et al. (2020); Haddadpour et al. (2021); Basu et al. (2019); Gao et al. (2021); Li and Li (2023), compressing asymptotically eliminates compression errors even in the presence of client heterogeneity. Consider the case of deterministic gradients for simplicity. Based on the update rules of SCAFFOLD (Algorithm 4), if hypothetically the training approached a steady stage where is close to a stationary point , we expect to have and . Consequently, the directions for local updates satisfy so that . Therefore, the definition of in (4) implies
Namely, the increment variable gradually vanishes as the algorithm iterates. Therefore, taking an -unbiased compressor as an example (see Definition 1), compressing results in a vanishing compression error
regardless of data heterogeneity. In contrast, if one considers compressing local gradients directly, a constantly large compression error is introduced in each communication round
The constant can be extremely large when the data heterogeneity is severe, resulting in the susceptibility of algorithms to data heterogeneity.
Comparison with FedPAQ (Reisizadeh et al., 2020), FedCOM (Haddadpour et al., 2021), Fed-EF (Li and Li, 2023).
All of them boil down to the FedAvg algorithm (McMahan et al., 2017) when no compression is conducted. As such, their convergence is significantly hampered by data heterogeneity across clients due to client drift. The former two works do not consider partial participation, and Fed-EF suffers from an extra slow-down factor in the convergence rate under partial participation. In opposition, SCALLION roots from SCAFFOLD, and is robust to arbitrary data heterogeneity and partial participation.
Comparison with FedCOMGATE (Haddadpour et al., 2021).
3 Convergence of SCALLION
To study the convergence of SCALLION under communication compression, we consider compressors satisfying the following standard unbiased compressibility.
where the expectation is taken over the randomness of the compressor .
Examples that satisfy Definition 1 include random sparsification and dithering as stated below.
For any , the random- sparsification is defined as where denotes the entry-wise product and is a uniformly random binary vector with non-zero entries. This random- sparsification is an -unbiased compressor with .
where and are the floor and ceiling functions, respectively. This random dithering with -bits per entry is an -unbiased compressor with .
When communication compression with -unbiased compressors is employed, the convergence of the proposed SCALLION (Algorithm 1) is justified as follows.
Under Assumptions 1 and 2, supposing clients apply mutually independent -unbiased compressors, if we initialize and , and set learning rates , as well as scaling factor properly, then SCALLION converges as
where . A detailed version and the proof are in Appendix C.
The initialization of and does not affect the convergence rate and the asymptotic complexities. The one in Theorem 1 is conducted for neatness. In practice, we can simply set .
Due to comprehensive challenges in compressed FL, to facilitate convergence analysis, most existing approaches require additional stringent conditions or assumptions that are not necessarily valid in practice, including but not restricted to:
As a result, their convergence rates inevitably depend on the large constants , , , . In contrast, the results presented in our work do not rely on any such condition.
Comparison with prior compressed FL methods.
Table 1 provides a summary of non-convex FL methods employing unbiased compressors under full client participation. We observe that SCALLION matches the state-of-the-art asymptotic communication and computation complexities under non-iid clients. In particular, while having the same asymptotic complexities as FedCOMGATE (Haddadpour et al., 2021), SCALLION does not incur the dependence on a large uniform bound of compression errors (see Remark 2) in convergence and thus has a superior convergence rate.
To sum up, based on the above discussion, we demonstrate that SCALLION theoretically improves existing FL methods with unbiased compression. In particular, SCALLION is the first stochastic FL method, to the best of our knowledge, that accommodates arbitrary data heterogeneity, partial client participation, and local updates, without any additional assumptions on compression errors.
SCAFCOM: Biased Compression with Momentum
While SCALLION achieves superior convergence speed under unbiased compression, its analysis cannot be adapted to biased compressors (also known as contractive compressors) to attain fast convergence rates. In this section, we propose an algorithm called SCAFCOM as a complement of SCALLION to accommodate biased communication compression in FL.
In the literature, biased compressors are commonly modeled by the following contractive compressibility.
where the expectation is taken over the randomness of the compressor .
Notably, compared to unbiased compressors satisfying Definition 1, contractive compressors, though potentially having smaller squared compression errors, no longer enjoy the unbiasedness. Common examples of contractive compressors include (Li and Li, 2023):
For any , the Top- operator is defined as where is the set of the largest entries of in absolute values. Top- operator is a -contractive compressor with .
Due to the lack of unbiasedness, compared to their counterparts with unbiased compressors, approaches employing with biased compressors in the literature typically (i) require stringent assumptions, e.g., bounded gradients (Seide et al., 2014; Koloskova et al., 2019; Basu et al., 2019; Li et al., 2022b) or bounded gradient dissimilarity (Huang et al., 2022; Li and Li, 2023), (ii) rely on impractical algorithmic structure, e.g., large data batches in gradient computation (Huang et al., 2022), (iii) have weak convergence guarantees, e.g., no improvement in the scaling of the number of clients (i.e., linear speedup) (Fatkhullin et al., 2021) or worse dependence on compression parameter (Zhao et al., 2022).
Recently, Fatkhullin et al. (2023) shows that tactfully incorporating momentum into communication compression can effectively mitigate the influence of biased compression. Inspired by their findings, we introduce an extra momentum variable on each client to overcome the adverse effect of biased compression. This leads to the SCAFCOM method, as presented in Algorithm 2. When client participates in the -th round, an additional momentum variable is updated as
where are the intermediate local models and is the momentum factor. We then set as the message to be communicated, as opposed to in SCAFFOLD and in SCALLION. Compared to the gradient yielded by single local loop, the momentum variable has a smaller variance due to its accumulation nature, thereby refining the convergence behavior under biased compression. Finally, note that similar to SCALLION, SCAFCOM only transmits one compressed variable in the uplink communication, and recovers SCAFFOLD when and are the identity mapping (i.e., no compression).
Notably, the difference between SCAFCOM and SCALLION lies in the utilization of momentum; see the colored highlights in Algorithm 1 and 2. Specifically, if we replace line 10 of SCAFCOM with the following formula:
then SCAFCOM recovers SCALLION (Algorithm 1) with . Note that in this case, the memorization of is no longer needed to be retained, which is consistent with the design of SCALLION. We also remark that the roles of the scaling factor and momentum vary in SCALLION and SCAFCOM. In SCALLION, stabilizes the updates of control variables while SCAFCOM sets to mainly address the biasedness issue of contractive compressors.
Connection with error feedback.
While SCAFCOM does not directly pertain to the vanilla error feedback (Seide et al., 2014; Stich, 2019), a technique widely used to tackle biased compression, SCAFCOM relates to the newly proposed EF21 mechanism (Richtárik et al., 2021). If one sets in SCAFCOM, then the message would be compressed and the control variable would be updated as . Under the simplification where (i.e., full-batch gradients), (i.e., no local updates), (i.e., full client participation), it becomes and the global model is updated through with , recovering the recursion of EF21.
2 Convergence of SCAFCOM
With the help of local momentum, the convergence of SCAFCOM under -contractive compression can be established as follows.
Under Assumption 1 and 2, supposing clients apply -contractive compressors , if we initialize and , and set learning rates , as well as momentum properly, then SCAFCOM converges as
where . A detailed version and the proof are in Appendix D.
The initialization of , , does not affect the convergence rate and the asymptotic complexities. The one in Theorem 2 is conducted for neatness. In practice, we can simply set .
Furthermore, it is known that one can convert any -unbiased compressor into a -contractive compressors with through scaling (see, e.g., (Safaryan et al., 2022, Lemma 1) and (Huang et al., 2022, Lemma 1)). Consequently, SCAFCOM can also employ unbiased compressors after the scaling with convergence guaranteed as:
When employing unbiased compressors (after scaling) in communication compression, then SCAFCOM converges as
Corollary 1 is obtained by directly plugging in the relation into Theorem 2 without exploiting the unbiasedness property of the compressors. However, it is feasible to refine the and terms in (11) by taking advantage of unbiasedness. We omit the proof here for conciseness.
Comparison with prior compressed FL methods.
In Table 1, we compare SCAFCOM with existing FL algorithms with biased compression under full client participation. We observe that SCAFCOM outperforms prior results with biased compression (QSPARSE-SGD (Basu et al., 2019), Local-SGD-C (Gao et al., 2021), and Fed-EF (Li and Li, 2023), etc) in the asymptotic communication complexity by at least a factor . Moreover, inferior to SCAFCOM, the existing FL methods with biased compression cannot tolerate unbounded data heterogeneity or even require homogeneous data. In addition, QSPARSE-SGD and Local-SGD-C only converge under full client participation. Notably, when employing unbiased compression, SCAFCOM enhances the asymptotic computation complexity by a factor of compared to SCALLION, surpassing all prior FL methods with unbiased compression. Furthermore, under partial client participation, our rate is better than that of Fed-EF (Li and Li, 2023) by a factor of thanks to control variables, overcoming the drawback of the standard error feedback under partial participation in distributed/federated learning.
Based on discussions in Section 5, we demonstrate that SCAFCOM, as a unified approach, outperforms existing compressed FL methods under both unbiased and biased compression. In particular, SCAFCOM is the first stochastic FL method, to the best of our knowledge, that accommodates arbitrary client heterogeneity, partial client participation, and local updates, as well as support communication compression relying only on standard contractive compressibility.
Experiments
We present a set of experiments on FL benchmark datasets to demonstrate the efficacy of our proposed algorithms. Since the (substantial) saving in communication overhead of various compressors is straightforward and has been well demonstrated in prior compressed FL works (e.g., via communication vs. test accuracy plots in Haddadpour et al. (2021); Li and Li (2023)), in this section, our empirical results mainly focus on:
Validating that SCALLION and SCAFCOM can empirically match the full-precision SCAFFOLD with considerably reduced communication costs.
Showing the advantages of SCALLION and SCAFCOM over prior methods with the same communication budget and training rounds.
We test our algorithms on two standard FL datasets: MNIST dataset (LeCun, 1998) and Fashion MNIST dataset (Xiao et al., 2017). The MNIST dataset contains 60,000 training images and 10,000 test images. Each image is a gray-scale handwritten digit from 0 to 9 (10 classes in total) with 784 pixels. The FMNIST dataset has the same training and test dataset sizes and the number of pixels per image whereas each image falls into 10 categories of fashion products (e.g., bag, dress), making the learning task more challenging. Following (Karimireddy et al., 2020b), we train a (non-convex) fully-connected neural network with 2 hidden layers with 256 and 128 neurons, respectively. We use ReLU as the activation function and the cross-entropy loss as the training objective.
Algorithms.
We implement our two proposed methods and two recent compressed FL algorithms, with biased and unbiased compression, respectively:
(Biased) Fed-EF (Li and Li, 2023): Federated learning with biased compression and standard error feedback. Since our proposed algorithms conduct SGD-type updates in the server, we compare them with its Fed-EF-SGD variant.
(Biased) SCAFCOM (our Algorithm 2): Biased compression for FL with stochastic controlled averaging and local momentum. The momentum in Algorithm 2 is tuned over a fine grid on .
(Unbiased) FedCOMGATE (Haddadpour et al., 2021): Federated learning with unbiased compression. This algorithm uses the gradient-tracking technique to alleviate data heterogeneity.
(Unbiased) SCALLION (our Algorithm 1): Unbiased compression for FL with stochastic controlled averaging. The local scaling factor in Algorithm 1 is tuned over a fine grid on .
Besides the compressed FL algorithms, we also test the corresponding full-precision baselines: Fed-SGD (also known as FedAvg (Yang et al., 2021)) and SCAFFOLD (Karimireddy et al., 2020b), both with two-sided (global and local) learning rates. For a fair comparison, we execute SCAFFOLD with our new implementation in experiments, corresponding to the special cases of SCAFCOM (, ) and of SCALLION (, ). Notably, under a fixed random seed, our implementation yields the same training trajectory as Karimireddy et al. (2020b) at a halved uplink communication cost (by only sending one variable per participating client).
In the experiments, biased compression is simulated with Top- operators (our Example 3). Specifically, we experiment with Top-0.01 and Top-0.05, where only the largest and entries in absolute values are transmitted in communication. For unbiased compression, we utilize random dithering (our Example 2), with 2 bits and 4 bits per entry, respectively. We tune the combination of the global learning rate and the local learning rate over the 2D grid . The combination of learning rates with the highest test accuracy is reported for each algorithm and hyper-parameter choice (e.g., , , and degree of compression).
Federated learning setting.
In our experiments, the training data are distributed across clients, in a highly heterogeneous setting following (Li and Li, 2023). The training data samples are split into 400 shards each containing samples from only one class. Then, each client is randomly assigned two shards of data. Therefore, every client only possesses training samples from at most two classes. All the clients share the same initial model at . In each round of client-server interaction, we uniformly randomly pick clients to participate in FL training, i.e., the partial participation rate is . Each participating client performs local training steps using the local data, with a mini-batch size 32. All the presented results are averaged over 5 independent runs with the same model initialization for all the algorithms.
2 Results
Since all the compressed FL methods in our experiments require transmitting one variable in the uplink communication, their communication costs are essentially the same when the same compressor is applied. Therefore, for clarity of comparisons, we will plot the metrics versus the number of training rounds.
In Figure 2, we first present the train loss and test accuracy of our proposed SCAFCOM (Algorithm 2) with momentum and Fed-EF (Li and Li, 2023), both using biased Top- compressors. We observe:
In general, under the same degree of compression (i.e., the value of in the case), SCAFCOM outperforms Fed-EF in terms of both training loss and test accuracy, thanks to controlled variables and the local momentum in SCAFCOM.
On both datasets, SCAFCOM with Top-0.01 can achieve very close test accuracy as the full-precision SCAFFOLD, and SCAFCOM with Top-0.05 essentially match those of full-precision SCAFFOLD. Hence, we can reach the same performance while saving 20 - 100x uplink communication costs.
For both SCAFCOM and Fed-EF, as the degree of compression decreases (i.e., increases), their performance approaches that of the corresponding FL methods under full-precision communication (i.e., SCAFFOLD and Fed-SGD).
SCALLION with unbiased compression.
In Figure 3, we plot the same set of experimental results and compare SCALLION () with FedCOMGATE (Haddadpour et al., 2021), both applying unbiased random dithering (Alistarh et al., 2017) with and bits per entry. Similarly, we see that SCALLION outperforms FedCOMGATE under the same degree of compression (number of bits per entry). The SCALLION curves of both 2-bit and 4-bit compression basically overlap that of SCAFFOLD, and 4-bit compression slightly performs better than 2-bit compression in later training rounds. Since random dithering also introduces sparsity in compressed variables, the 4-bit compressor already provides around 100x communication compression, and the 2-bit compressor saves more communication costs.
Impact of β𝛽\beta and α𝛼\alpha.
The momentum factor in SCAFCOM and the scaling factor in SCALLION are two important tuning parameters of our proposed methods. As an example, in Figure 4, we report the test accuracy of SCAFCOM with Top-0.01 (left column) and SCALLION (right column) 2-bit random dithering, for various and values, respectively. From the results, we see that SCAFCOM can converge with a wide range of , and performs the best on both datasets (so we presented the results with in Figure 2). For SCALLION, we report three -values, . When , the training of SCALLION becomes unstable for 2-bit quantization. As we use more bits, larger could be allowed. This is because, random dithering may hugely scale up the transmitted (compressed) entries, especially for low-bit quantization. When the scaling factor is too large in this case, the updates of local control variables become unstable, which further incapacitates the proper dynamic the local/global training. Thus, for SCALLION with low-bit random dithering, we typically need a relatively small . As presented in Figure 3, yields the best overall performance. In general, we should tune parameter and in SCAFCOM and SCALLION practically to reach the best performance.
Conclusion
This paper proposes two compressed federated learning (FL) algorithms, SCALLION and SCAFCOM, to support unbiased and biased compression in FL. The proposed methods are built upon our new implementation of the stochastic controlled averaging approach (SCAFFOLD), along with local momentum, and communication compression. Theoretically, under minimal assumptions, SCALLION and SCAFCOM match or improve the state-of-the-art convergence rates and complexities of compressed FL algorithms. Specifically, SCALLION and SCAFCOM are the first stochastic FL methods, to the best of our knowledge, that exhibit robustness to arbitrary data heterogeneity, partial participation, local updates, and also accommodate communication compression relying solely on standard compressibilities. Empirically, experiments show that SCALLION and SCAFCOM outperform prior compressed FL methods and perform comparably to full-precision FL approaches at a substantially reduced communication cost. In the future, our algorithms and techniques might be integrated with or extended to, for example, adaptive optimization, privacy, and fairness in federated learning.
References
Appendix A Detailed Implementations of SCAFFOLD
The original implementation of SCAFFOLD (Karimireddy et al., 2020b) is stated in Algorithm 3 where no compression is employed in communication. In this implementation, each participating client needs to transmit the increments of both local model and control variable to the server at the end of local updates, resulting to two rounds of uplink communication for per training iteration.
By communicating the increment variable , we can implement SCAFFOLD equivalently with only a single round of uplink communication for each participating client, as described in Algorithm 4.
Appendix B Preliminaries of Proofs
and
For SCAFCOM.
We additionally let . Then, due to client sampling, it holds that
and . Similarly, we let and .
Under Assumption 1, for any and , it holds that
We mainly focus on proving (14) as (16) can be established similarly. Using Lemma 2, we have
By further applying Sedrakyan’s inequality and Assumption 1, we have
The other upper bound of (14) follows . ∎
Appendix C Proof of SCALLION
In this subsection, we prove the convergence result of SCALLION with unbiased compression, where we additionally define .
Under Assumptions 1 and 2, it holds for all and that
For , using Lemma 2 and the fact that and , we have
Similarly, using (16) and Assumption 1, we have
Plugging (35) and (38) into (32), we obtain
Then applying the same relaxation in (38), we obtain
Plugging (41) and (44) into (30) and noting , we completes the proof. ∎
Given Lemma 4, the rest is to bound , .
Under Assumptions 1 and 2, it holds for all that
Using (14), Young’s inequality, and Assumption 1, we further have
Using Young’s inequality and Assumption 1, we can obtain
Combining (50), (53), (59) together and using completes the proof. ∎
Under Assumptions 1 and 2, suppose , then it holds for all that
Plugging (66) and (68) into (64), we obtain
where we use in the last inequality. By further using Young’s inequality and Assumption 1, we obtain
Under Assumptions 1 and 2, it holds for any and that
Under Assumptions 1 and 2, if we initialize , with and ( as ), set , , and
where .
Adding to (28), we have
Adding to (81), we have
Defining the Lyapunov function ()
Using and Lemma 7, we have
Due to the choice of and , it holds that
Averaging (98) over and noting , , we obtain
By the definition of , it holds that
where we use the choice of , , and the initialization of and in the second inequality. Due to the choice of , we have and thus
Plugging the choice of completes the proof. ∎
Appendix D Proof of SCAFCOM
In this subsection, we prove the convergence result of SCAFCOM with biased compression, where we additionally define .
Under Assumptions 1 and 2, it holds for all and that
For , using Lemma 2 and the fact that and , we have
Similarly, using (16) and Assumption 1, we have
Plugging (116) and (120) into (113), we obtain
Then applying the same relaxation in (120), we obtain
Plugging (122) and (127) into (109) and noting , , we completes the proof. ∎
Given Lemma 8, the rest is to bound , , and .
Under Assumptions 1 and 2, it holds for all that
Using (14), Young’s inequality, and Assumption 1, we further have
Using Young’s inequality and Assumption 1, we can obtain
Combining (132), (135), (139) together and using , , completes the proof. ∎
Under Assumptions 1 and 2, it holds for all that
by applying (16). By further using Young’s inequality and Assumption 1, we obtain
Under Assumptions 1 and 2, it holds for all that
where . Using Lemma 2 and Assumption 1, we have
By further using Sedrakyan’s inequality and Assumption 1, we obtain
By combinining (165) with (152) and using , we finish the proof. ∎
Under Assumptions 1 and 2, it holds for any and that
When , trivially for all so we consider below. Using Young’s inequality, we have
By further using Young’s inequality and Assumption 1, we obtain
By combinining the above inequalities together, we have
where we use so that in the last inequality. Iterating and averaging (182) over , we obtain
where we use the fact in the last inequality. ∎
Under Assumptions 1 and 2, if we initialize , with and ( as ), set , ,
where .
Adding to (107), we have
Using and to simplify coefficients, we obtain
Now adding to (199) and defining the Lyapunov function ()
Using and to simplify coefficients, we obtain
Using and Lemma 12, we have
Due to the choice of , it holds that
Averaging (215) over and noting , we obtain
Note that, by the definition of , it holds that
where we use the choice of and the initialization of , , and in the second inequality. Due to the choice of , we have
Plugging the choice of , we complete the proof.