Intrinsic Gradient Compression for Scalable and Efficient Federated Learning
Abstract
Federated learning is a rapidly growing area of research, holding the promise of privacypreserving distributed training on edge devices.
The largest barrier to wider adoption of federated learning is the communication cost of model updates, which is accentuated by the fact that many edge devices are bandwidthconstrained.
At the same time, within the machine learning theory community, a separate line of research has emerged around optimizing networks within a subspace of the full space of all parameters.
The dimension of the smallest subspace for which these methods still yield strong results is called the intrinsic dimension .
In this work, we prove a general correspondence between the notions of intrinsic dimension and gradient compressibility, and we show that a family of low-bandwidth federated learning algorithms, which we call intrinsic gradient compression algorithms , naturally emerges from this correspondence.
Finally, we conduct large-scale NLP experiments using transformer models with over 100M parameters (GPT-2 and
BERT), and show that our method outperforms the state-of-the-art in gradient compression.
Introduction
Federated learning is a nascent area of study which seeks to perform machine learning in a privacypreserving way.
However, federated learning with deep neural networks suffers from a problem with communication bandwidth: it is very costly to send gradient/model updates over a network, especially when communicating with mobile phones and edge devices.
To reduce bandwidth for federated learning, it is natural to utilize various forms of compression.
Previous works have tried to achieve compression in two ways: (1) by compressing the information communicated in standard gradient descent algorithms (e.g. quantizing gradients (Wen et al., 2017)) and (2) by training with non-standard methods that naturally use less bandwidth (e.g. prototypical networks (Tan et al., 2021)).
At the same time, in the machine learning theory community, researchers have been working to understand what at first seems like an entirely different question: why do hugely overparametrized models generalize so well?
One promising approach to this answering this question has utilized the concept of intrinsic dimension , defined for a given optimization problem as the smallest dimension d for which we can solve the problem when the weights are restricted to a a d -dimensional manifold.
To be precise, it is the smallest d for which the standard loss minimization problem
[Equation]
has a satisfactory solution, where the image of g is a d -dimensional manifold.
If the intrinsic dimension of a problem is low, then even if a model is vastly overparameterized, only a small number of parameters need to be tuned in order to obtain a good solution, which is often enough to imply certain generalization guarantees.
We begin this paper by observing that the two problems above are naturally related.
If one can find a solution to the problem by only tuning d parameters, as in Equation (1), then a corresponding low bandwidth algorithm can be found by simply running stochastic gradient descent in the reduced parameter space (in this case, R d ).
However, simply optimizing a subset of a model’s parameters is often insufficient for training models (especially when training from scratch, rather than finetuning).
Thus, we are inspired to seek a more general characterization of algorithms that use a low amount of bandwidth.
In order to do this, we rewrite the optimization problem in Equation (1) in the original parameter space.
When g ( θ ′ )
=
Aθ ′ for some matrix A (so the lowdimensional manifold is a low-dimensional sub- space), stochastic gradient descent can be rewritten as
[Equation]
We call this method static intrinsic gradient compression , because our gradients are projected into a static (“intrinsic”) subspace.
Now, Equation (2) admits a natural generalization, which allows us to explore more of the parameter space while still preserving a low level of upload bandwidth usage:
[Equation]
where A t may vary with time.
We call the set of all such algorithms intrinsic gradient compression algorithms , and consider three particular instantiations: static, time-varying, and k -varying, each of which perform in different use cases.
Our approach is model-agnostic and highly scalable.
In experiments across multiple federated learning benchmarks (language modeling, text classification, and image classification), we vastly outperform prior gradient compression methods, and show strong performance even at very high compression rates (e.g. up to 1000 × ).
Our contributions are as follows.
• We find a general class of optimization algorithms based on the notion of intrinsic dimension that use low amounts of upload bandwidth, which we denote intrinsic gradient compression algorithms .
• We specify three such algorithms: static compression, time-varying compression and Kvarying compression, with different levels of upload and download bandwidth for use in various federated settings.
• In a set of experiments, we show that these methods significantly outperform prior approaches to federated learning with gradient compression, obtaining large reductions in bandwidth at the same level of performance.
In Section 2, we describe the preliminaries needed to contextualize our work, namely ideas from intrinsic dimension, federated learning, and gradient compression.
In Section 3, we show how the algorithm used by intrinsic dimension naturally generalizes to algorithms which use little upload bandwidth.
In Section 4 we consider special instantiations of these algorithms in federated learning settings which attain low upload and download bandwidth, and in Section 5 show that they achieve state of the art results.
Finally, Sectionconcludes.
Preliminaries 2. Intrinsic Dimension
The concept of intrinsic dimension was introduced in the work of (Li et al., 2018), as a way of evaluating the true difficulty of an optimization problem.
While this can usually be done by counting the number of parameters, some optimization problems are easier than others in that solutions may be far more plentiful..
One can write
[Equation]
where g : R d → R D and thus we’ve transformed the problem into an optimization problem over θ 2 .
If we can still find good solutions to the original problem where θ 2 ∈ Θ 2 , then the problem’s intrinsic dimension may be lower, and thus the question may be easier than previously expected.
Throughout this paper we will always take g ( θ ′ )
=
Aθ ′
+ θfor a D × d matrix A , and take Θ 2 = R d , and Θ 1 = R D , where D > d , where θ 0 is the original value of the expression.
The intrinisic dimension g ( ℓ, L ) with respect to a task ℓ and performance threshold L is equal to the smallest integer d so that optimizing Equation (4) on task ℓ could lead to a solution of performance at least equal to T .
The intrinsic dimension is not exactly knowable, because we cannot find the “best performing model” exactly.
However, if say, training with some optimization algorithm gives us a solution to Equation (4) with loss ≤ L and with d dimensions, we can say with certainty that g ( ℓ, T ) ≤ d .
2. Federated Learning
Federated learning is a paradigm built around protecting the privacy of user data.
The standard model involves a server and many clients, where the raw data must remain on the client’s device but the server learns a model.
Generally, this is implemented by only the gradients of the model on the data being sent to the central server, which then runs a standard algorithm.
A common example of this is the FedAvg algorithm (McMahan et al., 2017), where models are trained to near-completion on a each client’s data, and the data is then averaged.
In what follows, we define an epoch to be a single pass over every client.
2. Gradient Compression
Sending full gradients in standard uncompressed form uses far more bandwidth than we are afforded in certain settings.
For example, in a 1 billion parameter model (hardly particularly large by current standards) one gradient update would take 4 gigabytes of bandwidth uncompressed.
Thus, there has been substantial amounts of work in compressing the gradient, like (Albasyoni et al., 2020), which finds an optimal gradient compression algorithm, albeit one which is computationally infeasible.
2. Related Work: Model Pruning and Model Compression Related Work: Model Pruning
There has been great interest in compressing models by using fewer weights, starting with the work of (Hinton et al., 2015; Han et al., 2015).
One related work is
Diff Pruning (Guo et al., 2021), which constrains the number of weights that can be changed from a pretrained model.
In essence, diff pruning attempts to solve an L 0 minimization problem on the weights of the model, and approaches this by means of a relaxation.
A number of other works have explored the idea of finetuning by only modifying a subset of a model’s parameters.
(Jiang et al., 2019)
and (Bibikar et al., 2021) utilize sparsity to reduce communication costs during training.
(Ravfogel et al., 2021) finetunes only the layer biases of large models.
Similarly, (Houlsby et al., 2019) finetunes lowparameter adapters between each layer.
Compared to (Ravfogel et al., 2021)
our method is far more flexible, allowing any number of parameters to be changed.
Compared to (Houlsby et al., 2019) our methods are architecture-independent, and can be applied to any model.
Related Work: Federated Learning Federated
learning is a machine learning paradigm in which a model is trained by a collection of clients, each with their own private local data.
From the introduction of federated learning (McMahan et al., 2017), it was clear that communication costs represented a significant challenge: sending gradients or weights over a network is costly due to the large size of modern machine learning models.
(McMahan et al., 2017) introduced the FedAvg algorithm, which aims to reduce communication costs by sending and averaging weights, rather than gradients.
Specifically, clients train their model locally for a given number of epochs, send it to the server, and received an averaged copy of the model weights.
However, sending the full set of model weights often remains very costly (especially when clients only have a small amount of local data, such that many rounds of communication are necessary); as a result, FedAvg performs poorly in heavilybandwidth-constrained settings.
Recently, FetchSGD (Rothchild et al., 2020) aimed to address this issue differently by utilizing the concept of sketching.
Rather than transmitting full gradients from the client to the server, they send a sketch of the gradient.
This approach performs well, but only yields moderate compression rates.
We compare to FetchSGD in Section 5.
A Family of Low-Bandwidth Algorithms
In this section, we characterize a family of lowbandwidth optimization algorithms based on the notion of intrinsic dimension.
We start from the optimization problem induced by intrinsic dimension (Equation (4)).
If we directly run gradient descent on Equation (4) with respect to the intrinsic weights θ ′ , we obtain an equation of the following form:
[Equation]
Then, left-multiplying both sides by A we obtain
[Equation]
Note that here, we can interpret
A ⊤ ∇ θ ( ℓ ( f ( θ )))
| θ = θ t as a compressed gradient with dimension d , and AA ⊤ ∇ θ ( ℓ ( f ( θ )))
| θ = θ t as the approximate gradient.
This inspires us to consider the more general family of optimization algorithms given by
[Equation]
where v t is a D dimensional vector computed from data available at timestep t that plays a similar role to a gradient, but may not be an exact gradient, and the A t are all D
× d matrices known ahead of time (say, generated with random seeds).
One intuitive way of interpreting this algorithm is that θ t +1 − θ t is constrained to lie in a low-dimensional subspace, input: learning rate η , timesteps T , local batch size ℓ , clients per round W
Create matrix A ∈ R D × d with E [ AA ⊤ ] = I D .
Current Vector: Σ 0 = 0 for t = 1 , 2 · · · T do
Randomly select W clients c 1 , . . .
c W . loop {In parallel on clients { c i } Wi =1 } Download Σ t − 1 , calculate current θ t −
= θ 0 +
A (Σ t − 1 ) .
Compute stochastic gradient g ti on batch B i of size ℓ : g ti = 1 ℓ (cid:80) ℓj =1 ∇ θ L ( θ t − 1 , z j ) .
Sketch g ti to S ti =
A ⊤ g ti and upload it to the aggregator.
end loop
Aggregate sketches S t = 1 W (cid:80)
Wi =1 S ti
Unsketch: ∆ t =
AS t
Update: θ t
= θ t − 1 − η ∆ t , Σ t =
Σ t − 1 − ηS t .
end for namely that given by the span of A t . This family of algorithms can be made to use only d upload bandwidth, as only the vector A ⊤ t ( v t ) must be uploaded. Furthermore, note that Equation (6) has no references to the intrinsic weights θ ′ , meaning that it represents a general optimization algorithm in the original space. Formally,
Theorem 3.1.
All algorithms of the form
[Equation]
can be simulated with d upload bandwidth in a standard federated learning setting, where v t is a function that can be calculated by the client at time t combined with all data from the server.
We call all such algorithms intrinsic gradient compression algorithms .
Note that this theorem only bounds the upload bandwidth capacity needed to run gradient descent, and does not bound the download bandwidth.
In the particular instantiations we consider, we will demonstrate that one can also bound the download bandwidth.
Intrinsic Gradient Compression Algorithms
While Theorem 3.1 shows that any algorithm of the form Equation (6) can be implemented with low levels of upload bandwidth, not every algorithm of the form Equation (6) can be implemented with low levels of download bandwidth as well.
Theorem 3.1 gives rise to a family of algorithms we denote intrinsic gradient compression algorithms .
In this section, we describe three particular intrinsic gradient compression algorithms which use low amounts of both upload and download bandwidth.
These federated learning algorithms can be decomposed into three main phases
Reconciliation:
The client reconciles its model with the server’s copy of the model.
• Compression: The local model calculates, compresses, and sends its local gradient to the server.
• Decompression: The server model updates its own copy of the model using the estimated gradient from the local model.
In general, reconciliation will be by far the most complex part of each algorithm, and the other steps are essentially shared across algorithms.
We show how to implement SGD for each variant, and note that this choice of optimization algorithm is quite necessary – other optimization algorithms like SGD with momentum cause the parameters to not move in the low-dimensional subspace, which makes the compression impossible.
While one can implement a variant which resets the momentum every epoch, momentum is rarely a useful optimization in federated learning due to the non-i.i.d.
nature of the batches) so we do not consider this.
Static Intrinsic Gradient Compression
In this subsection, we seek to implement the static intrinsic gradient compression algorithm
[Equation]
in a federated learning setting.
In the reconciliation phase, since we know that the parameters θ c (which denotes the current parameters of the server) will always be equal to θ 0 + A Σ for some Σ ∈ R d , the server can just send Σ to the client, which will take d download bandwidth.
For compression, the client compresses the gradient by multiplying by A ⊤ , and for decompression the server multiplies this by A .
The full algorithm is given in Algorithm 1.
Time-Varying Intrinsic Gradient Compression
In this subsection, we implement the time-varying intrinsic gradient compression algorithm
[Equation]
in a federated learning setting, where e is the epoch.
In this case, we show that our algorithm can be implemented with at most 2 d bandwidth used per
Intrinsic Gradient Compression Method
Upload
Download
Dimensions Explored
Static dE dE d
Time-Varying dE 2 dE dE
K -Varying dE 2 dEK dEK
No Compression
DE DE
D client per timestep, so over E epochs there is 2 dE bandwidth used total on downloading. Since this bandwidth is twice that of static subspace compression, but we search E times more directions in the space, this algorithm is particularly useful when we have many epochs.
Letting θ c e be the client parameters at epoch e , note that we have the value of θ ce − 1 when performing reconciliation.
Now we can write
[Equation]
and we can see that ( θ ce − θ final e − 1 ) lies in the column space of A e and ( θ final e − 1 − θ ce − 1 ) lies in the column space of A e − 1 , which is enough to find the full algorithm, given in Algorithm 2.
K -Varying Intrinsic Gradient Compression
In this subsection, we describe how to implement the K -varying intrinsic gradient compression algorithm
[Equation]
where { A ( i ) e } Ki =1 is the set of K compression matrices used at epoch e , and i is a randomly chosen integer between 1 and K inclusive.
This method is motivated from the fact that in many cases, the upload speed is much slower than the download speed, so we may only want to project the gradient into part of the subspace currently being explored, as opposed to the complete subspace.
This allows each client to explore d directions at a time, but for dK directions to be explored across the entire epoch.
As such, the algorithm identical time-varying compression, and is given in Algorithm 3.
Choice of Compression Matrix
Finally, we we discuss the choice of compression matrix for A .
We note that our methods are agnostic to the specific choice of A , and depend only on the existence of efficient subroutines for calculating the matrix-vector products Ax and A ⊤ y .
Nonetheless, the choice of
A has significant implications for the resulting accuracy of the algorithms.
In order to maintain the most proximity to the original stochastic gradient descent algorithm, we will choose normalized A so that E [ AA ⊤ ] = I D .
The naive choice is to let A be a D × d random dense matrix, but such a choice is impossible due to memory constraints.
For example, if we aim to train even a small version of BERT (100M parameters) with an intrinsic dimension of 1000 , we would need to store a matrix with 10 11 entries.
The approach taken by (Aghajanyan et al., 2021;
Li et al., 2018) for large-scale experiments, which we follow, utilizes the Fastfood transform (Le et al., 2013), in which A can be expressed as the D × d matrix A i =
Unpad D B
i H Π
i G i H Pad 2 ℓ where 2 ℓ is the smallest power of two larger than D , H is a standard Hadamard matrix,
B i is a random diagonal matrix with independent Rademacher entries (random signs), Π is a random permutation matrix, G is a random diagonal matrix with independent standard normal entries, Pad 2 ℓ to be a linear operator which simply pads a d -dimensional vector v with zeroes until it has size 2 ℓ , and Unpad D is a linear operator which takes the first D elements from a 2 ℓ -dimensional vector.
Since we can quickly compute a matrix-vector product by
H with a fast Walsh-Hadamard transform, we can perform a matrix multiplication by A
i
A ⊤ i in
O ( ℓ 2 ℓ ) =
O ( D log D ) time
and O ( D ) space.
Performance Comparison
We show the theoretical tradeoffs between each of these algorithms in
Table 1.
[Table]
Experiments
We evaluate our method across a range of benchmarks to showcase the potential of our three algorithms.
These include two natural language processing tasks (language modeling and text classification), as well as a computer vision task (image classification).
As with previous works (Rothchild et al., 2020;
McMahan et al., 2017), we simulate the federated learning in order to scale to large numbers of clients (upwards of 10 , 000 ).
We simulate oncommercial-grade GPUs for the language modeling experiments and 1 GPU for the other experiments.
We perform experiments in both non-IID (language modeling, image classification) and IID (text classification) settings, because both scenarios are common in real-world federated learning.
Image Classification (ResNet-9 on CIFAR-10)
First, we consider image classification on the
CIFAR-10 dataset, a collection of 50,000 images with resolution 32 × 32 px.
We use the same experimental setup as (Rothchild et al., 2020): we split the data between 10,000 clients in a non-IID fashion, such that each client only has data from a single class.
At each step, we sample 100 clients at random, such that each gradient step corresponds to 500 images.
We perform 24 rounds of communication between all clients (i.e. 24 training epochs).
We use a ResNet-9 architecture with 6,570, trainable parameters for our fair comparison to previous work.
Note that the model does not have batch normalization, as batch normalization would not make sense in a setting where each client has so few examples.
Due to the substantial number of epochs performed here, we experiment with both static and time-varying gradient compression ( kvarying compression is better suited to settings involving fewer rounds of communication).
We perform experiments across intrinsic dimensions of 4000, 8000, 16000, 32000, 64000, 128000, and 256000.
Our results are shown in Figure 1.
Whereas
FedAvg and Top-K struggle at even modest compression rates (e.g. 3 × ), the intrinsic gradient compression methods deliver strong performance at much larger compression rates.
The intrinsic methods outperform the current state-of-the-art gradient compression method, FetchSGD (Rothchild et al., 2020), by a large margin, and easily scale to high compression rates (e.g. 100 × ).
Finally, we see that time-varying intrinsic compression generally outperforms static compression for the same communication cost.
Text Classification (BERT on SST-2)
Next, we consider text classification on the Stanford Senti- [Figure] ment Treebank-v2 (SST-2) dataset (Socher et al., 2013), a common sentiment analysis dataset.
For this experiment, we consider an IID data split into 50 and 500 clients, respectively.
We employ the popular BERT (Devlin et al., 2019) transformer architecture with 109M parameters.
The purpose of this experiment is to push the limits of gradient compression; we project the 109M-dimension
BERT gradients into as few as 200 dimensions.
We perform 30 rounds (i [Figure] e. 30 epochs) of training for all compressed runs, while we perform 6 for the uncompressed baseline (as it converges more quickly).
Federated learning experiments has previously been criticized for being challenging to reproduce; as a result, we perform each run five [Figure] times over different random seeds.
We report the mean, min, max, and standard deviation of the runs in Appendix D.
Due to the substantial number of epochs performed here, it is natural to apply static and timevarying intrinsic gradient compression.
We use intrinsic dimensions of 200, 400, 800, . . .
Our results are given in Figure 2.
First, along similar lines to (Aghajanyan et al., 2021), we find that it is possible to achieve remarkably high compression ratios for text classification: we achieve close to full performance even when compressing the 109M-dimension parameter vector into an intrinsic space of dimension 16,384.
Furthermore, we find that time-varying intrinsic gradient compression consistently outperforms static intrinsic gradient compression at the same compression rate.
Language Modeling (GPT-2 on PersonaChat)
Lastly, we consider language modeling on the PersonaChat (Zhang et al., 2018) dataset of dialogues between Amazon Mechanical Turk workers as- signed to act out specific personalities.
The dataset has a non-IID split into 17,568 clients in which each client is assigned all data corresponding to given personality; as a result, it is widely used in federated learning simulations.
We perform language modeling using the GPT-2 transformer architecture (124M parameters).
For fair comparison to previous work, we conduct only two rounds of training across the clients (i.e. two epochs).
Due to the low number of training rounds, it is natural to apply static and K-varying gradient compression.
Specifically, we apply both of these algorithms to train GPT-2 using intrinsic dimensions of 16384, 65536, 262144, 1048576, and 4194304.
Our results are shown in Figure 2.
Overall, intrinsic dimension-based gradient compression vastly outperforms a wide range of prior approaches to reducing communication in federated learning.
On the low-compression end of the spectrum, we obtain nearly full performance with superior compression rates to FedAvg (McMahan et al., 2017) and the recent FetchSGD (Rothchild et al., 2020).
On the high-compression end of the spectrum, we scale better than previous approaches.
For example, we obtain a perplexity of around 20 even with an extremely high compression rate of 1898.
Finally, we see that K -varying intrinsic compression performs similarly to (or slightly worse) than static compression at the same level of overall compression.
However, if it is more important to conserve upload bandwidth than download bandwidth, then K -varying intrinsic gradient compression significantly outperforms static intrinsic gradient compression (see Section 4).
5. Gradient Compression Results
One of the primary motivations of federated learning is the desire for individual clients to be able to retain data privacy while still participating in model training.
However, a number of works have shown that if the client does not have a large amount of data and the client sends back their full local gradient, it is possible to approximately reconstruct their local data from the model.
This is a significant problem, because their data would then effectively be visible to the central server and any attackers that intercept their communications.
Here, we show that compressing gradients with our approach can mitigate this problem.
Specifically, we check if our compressed gradients can be reconstructed with the procedure proposed by (Zhu et al., 2019).
As in (Zhu et al., 2019), we use a
ResNet-152 model a randomly selected image from
ImageNet and run for 24,000 iterations (by which time the method has converged).
We reconstruct the image both from the full gradient (the center image) and from a the intrinsically-compressed image (the right image) with intrinsic dimension 65,536.
As seen in Figure 3, given the full gradient it is possible to obtain a fairly good reconstruction of the image.
By contrast, with our method, the reconstruction is visually much less similar from original image.
Of course, our method does not solve the problem entirely; an outline of the dog in the image is still visible because the compressed gradient still contains some information about the local data.
To solve the issue entirely, it would be necessary to use a method such as differential privacy.
Conclusion
Federated learning holds the promise of large-scale model training while simultaneously letting users retain control over their data.
In this paper, we preset a set of novel algorithms for scalable and efficient federated learning.
These algorithms are particularly helpful for NLP training, where models often have hundreds of millions of parameters.
Our experiments finetuning BERT and GPT-2 that our proposed method significantly improves upon the state-of-the-art in gradient compression for federated learning.
In future work, we hope to deploy our system in a real-world federated learning setting with a large number of physical devices, rather than solely in simulation.  Acknowledgments
We would like to thank Ankur Moitra, Yang Liu, and Demi Guo for helpful discussions.
L. M. K. is supported by the Rhodes Trust.
Appendix A Algorithms
In Algorithm 2 and Algorithm 3 below, we provide the full time-varying and K -varying intrinsic gradient compression algorithms, which were omitted from the main text.
B Proofs B. Proof of Theorem 3.
First, note that the server knows the value of A t .
Then, for any local vector v t , the client can send
A ⊤ t ( v t ) to the server, and the server can calculate
A t A ⊤ t , enabling it to continue executing the algorithm.
C Additional Related work
In the main paper, we described the prior work in federated learning and machine learning theory that was directly relevant to our paper’s method.
Here, we describe a number of less-related works that could not be included in the main paper due to space constraints.
Intrinsic Dimensionality
As mentioned in the main paper, the concept of measuring the intrinsic dimensional of loss landscapes was introduced by (Li et al., 2018).
(Li et al., 2018) consider optimizing a D -parameter model in a random ddimensional subspace of the full parameter space.
They define the intrinsic dimension of the optimization problem as the minimum dimension d for which a solution to the problem can be found, where a “solution” refers attaining a certain percentage of the maximum possible validation accuracy (i.e. the validation accuracy obtained by optimizing in all D dimensions).
They use a fixed cut-off of 90 % accuracy for their experiments. (Aghajanyan et al., 2021) followed up on this work by considering the setting of finetuning models in natural language processing.
They show that the intrinsic dimension of some of these tasks (e.g. text classification on MRPC) is surprisingly low.
A number of works have tried to measure the intrinsic dimension of datasets, rather than objective landscapes.
(Levina and Bickel, 2005) introduced a maximum likelihood approach to estimating intrinsic dimensionality based on nearest-neighbors, while (Ceruti et al., 2014) employed angle and norm-based similarity.
More recently, () further extended this line of work to use minimal neighborhood information.
Finally, some works have tried to measure the intrinsic dimensionality of image representations and datasets.
(Gong et al., 2019) finds that the representations produced by popular image and face representation learning models (ResNetand SphereFace) have quite low intrinsic dimensionalities (16 and 19, respectively).
Along similar lines, (Pope et al., 2021) showed that popular image datasets (MNIST, CIFAR 10, ImageNet) also have low intrinsic dimensionality.
Federated Learning
Federated learning is generally concerned with the distributed training of machine learning models across many devices, each of which holds private data.
Many aspects of this federated setup are separate subfields of research, including how to ensure the privacy of client-held data (Xie et al., 2020;
Bhagoji et al., 2019), how to deal with heterogeneous data and networks (Li et al., 2020a,b; Yu et al., 2020), how to reconcile weights/gradients from multiple clients (Li et al., 2020a; Wang et al., 2020; Li et al., 2020c), how to manage clients in a fault-tolerant manner, deployment on mobile/iot devices (He et al., 2020), and fairness (Mohri et al., 2019).
Numerous works focus on making federated training more efficient, with the ultimate goal of reducing communication cost and training time.
The classic FedAvg (McMahan et al., 2017) algorithm tries to do this by communicating weights rather than gradients.
FedProx (Li et al., 2020a) generalizes and re-parametrizes FedAvg.
FedMA (Wang et al., 2020) continues to improve this approach by matching and averaging hidden layers of networks with similar activations at each communication round.
FedAwS
(Yu et al., 2020) considers federated averaging in the case where each client has data from only a single class.
(Malinovsky et al., 2020) analyzes a generalization of these weightaveraging approaches from a theoretical viewpoint.
Relative to the weight averaging approach, the approach of compressing and sending gradients is relatively understudied.
(Albasyoni et al., 2020) describes an approach that is theoretically optimal but not practical for large non-linear models.
(Han et al., 2020) proposes adaptive gradient sparsification for federated learning, in which a subset of the full gradient is communicated at each round.
FetchSGD (Rothchild et al., 2020) compresses gradients by sketching; it is the current state-of-the-art input: learning rate η , timesteps T , local batch size ℓ , clients per round W for e = 1 , 2 , · · · E do
Create matrix A e i.i.d.
∼ A where A ∈ R D × d with E [ AA ⊤ ] = I D .
Current, Final Vector: Σ current e = 0 , Σ final e =for t = 1 , 2 · · · T do
Randomly select W clients c 1 , . . .
c W . loop {In parallel on clients { c i } Wi =1 }
Download Σ current e , Σ final e − 1 , calculate current θ c i e
=
θ
c
i e − 1 + A e − 1 (Σ final e − 1 − Σ last ) + A e (Σ current e ) .
Update Σ last =
Σ current e .
Compute stochastic gradient g ti on batch B i of size ℓ : g ti = 1 ℓ (cid:80) ℓj =1 ∇ θ L ( θ c i e , z j ) .
Sketch g ti : S ( e ) t i =
A ⊤ e g ti and upload it to the aggregator.
end loop
Aggregate sketches S ( e ) t = 1 W (cid:80)
Wi =1 S ( e ) t
i
Unsketch: ∆ ( e ) t =
A e S ( e ) t
Update: θ current = θ current − η ∆ ( e ) t , Σ current e =
Σ current e − ηS ( e ) t .
end for
Let Σ final e =
Σ current e .
end for in gradient compression for federated learning.
We describe it in further depth in the main paper.
Finally, (Reddi et al., 2021) and (Li et al., 2020c) accelerate training by bringing adaptive optimizers built for centralized learning into the federated setting.
D Further Experimental Analysis
In the main paper, we included a number of figures demonstrating our performance in comparison to prior work.
Here, we include tables with our precise results for clarity and in order to facilitate future comparison with our work.
D. Further PersonaChat Analysis
Section 4 shows full results on PersonaChat, complete with upload and download compression.
Overall compression is calculated as average compression over both upload and download.
We compare with FedAvg (McMahan et al., 2017), Top-K, and FetchSGD (Rothchild et al., 2020).
FedAvg is the baseline federated learning approach involving sending and averaging weights.
Top-K refers to sending the top gradients, sorted by magnitude.
FetchSGD compresses the weights with sketching.
Our method significantly outperforms competing approaches across the board.
We obtain an accuracy close to that of uncompressed optimization using INSERTx overall compression; FedAvg and Top-K both fail to achieve such strong results, while FetchSGD does so at a significantly lower compression rate.
Next we compare static and K-varying intrinsic gradient compression.
When comparing overall compression rates, static compression is slightly better than K-varying compression.
However, Kvarying compression is optimized for low upload bandwidth; it obtains much better upload compression rates than static compression at the same accuracy.
For example, K-varying compression with k = 8 and d = 65536 yields perplexity 17 . 6 at upload compression 1900 × , whereas static compression with d = 262144 yields perplexity 17 .at upload compression 475 × .
D. Further SST-2 Analysis
In Table 3, we show full results for the SSTdataset with static and time-varying gradient compression for a range of intrinsic dimensions.
We include in this experiment an demonstration of the robustness of our method to variation in random seeds; we run each experiment five times using separate random seeds (i.e. different intrinsic subspaces and model initializations).
We report standard errors in Table 3; variability is quite low.
We also see that time-varying intrinsic gradient compression outperforms static intrinsic compression, especially for low intrinsic dimensions.
For example, time-varying compression at d =outperforms static compression with d = 400 , and time-varying compression with d = 400 outperforms static compression with d = 800 .
input: distinct subspaces K , learning rate η , timesteps T , local batch size ℓ , clients per round W for e = 1 , 2 , . . .
E do
Create matrices A (1) e , A (2) e , . . .
A ( K ) e i.i.d.
∼ A where A ∈ R D × d with E [ AA ⊤ ] = I D .
Current, Final Vector: Σ current( k ) e = 0 , Σ final( k ) e = 0 for k = 1 , 2 , . . .
K .
for t = 1 , 2 · · · T do
Randomly select W clients c 1 , . . .
c W . loop {In parallel on clients { c i } Wi =1 }
Download Σ current( k ) e , Σ final( k ) e −for k = 1 , . . .
K , and calculate: θ c i e =
θ c i e − 1 + (cid:80)
Kk =1 (cid:16) A e − 1 (Σ final( k ) e −− Σ last( k ) )
+ A e (Σ current( k ) e ) .
(cid:17) Σ last( k ) =
Σ c ( k ) e for k = 1 , 2 , . . .
K .
Choose a random k 1 ∼ DUnif ( { 1 , 2 , . . .
K } )
Compute stochastic gradient g ti on batch B i of size ℓ : g ti = 1 ℓ (cid:80) ℓj =1 ∇ θ L ( θ c i e , z j ) .
Sketch g ti : S ( e ) t i = ( k 1 , A ( k 1 )
⊤ e g ti ) and upload it to the aggregator.
end loop
Write sketches received as { S ( e ) t w } Ww =
= { ( j w , C ( e ) t w ) }
Ww =1 .
Unsketch S ( e ) t to get ∆ ( e ) t = 1 W (cid:80) Ww =1 A ( j w ) e
C ( e ) t w
Update: θ current = θ current − η ∆ ( e ) t , for k = 1 , 2 . . .
K do
Update: Σ current( k )
[Equation]
end for end for
Let Σ final ( k ) e
= Σ c ( k ) e for k = 1 , 2 , . . .
K .
end for
[Figure] [Table]