Intrinsic Gradient Compression for Scalable and Efficient Federated Learning (2022)

Abstract: Federated learning is a rapidly growing area of research, holding the promise of privacy-preserving 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 bandwidth-constrained. 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 significantly 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 @xcite )) * equal contribution and (2) by training with non-standard methods that naturally use less bandwidth (e.g. prototypical networks @xcite ). 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 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 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: 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. 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, section 6 concludes. ## 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: as the approximate gradient. this inspires us to consider the more general family of optimization algorithms given by 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. 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. ## 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 @xcite , we simulate the federated learning in order to scale to large numbers of clients (upwards of 10, 000). we simulate on 8 commercial-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 × 32px. we use the same experimental setup as @xcite : 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 communi-cation between all clients (i.e. 24 training epochs). we use a resnet-9 architecture with 6,570,880 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 @xcite , 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. we perform 30 rounds (i.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 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, . . . , 25600. our results are given in figure 2 . first, along similar lines to @xcite , 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 per-sonachat @xcite dataset of dialogues between amazon mechanical turk workers as-signed to act out specific personalities. 1 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. 2 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 @xcite and the recent fetchsgd @xcite . 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). ## 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. intrinsic dimensionality as mentioned in the main paper, the concept of measuring the intrinsic dimensional of loss landscapes was introduced by @xcite . @xcite 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. @xcite 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. @xcite introduced a maximum likelihood approach to estimating intrinsic dimensionality based on nearest-neighbors, while @xcite 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. @xcite finds that the representations produced by popular image and face representation learning models (resnet-50 and sphereface) have quite low intrinsic dimensionalities (16 and 19, respectively). along similar lines, @xcite 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 @xcite , how to deal with heterogeneous data and networks @xcite , how to reconcile weights/gradients from multiple clients @xcite @xcite , how to manage clients in a fault-tolerant manner, deployment on mobile/iot devices @xcite , and fairness @xcite . numerous works focus on making federated training more efficient, with the ultimate goal of reducing communication cost and training time. the classic fedavg @xcite algorithm tries to do this by communicating weights rather than gradients. fedprox @xcite generalizes and re-parametrizes fedavg. fedma @xcite continues to improve this approach by matching and averaging hidden layers of networks with similar activations at each communication round. fedaws @xcite considers federated averaging in the case where each client has data from only a single class. @xcite 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. @xcite describes an approach that is theoretically optimal but not practical for large non-linear models. @xcite proposes adaptive gradient sparsification for federated learning, in which a subset of the full gradient is communicated at each round. fetchsgd @xcite compresses gradients by sketching; it is the current state-of-the-art algorithm 2 time-varying intrinsic gradient compression input: learning rate η, timesteps t , local batch size ℓ, clients per round w for @xmath0, 2, • • • e do create matrix ae i.i.d. ∼ a where @xmath1. current, final vector: σ current @xmath2, σ final @xmath3, 2 • • • t do randomly select w clients c1, . . . cw . loop {in parallel on clients {ci} w @xmath4, σ final e-1 , calculate current θ c i @xmath5. update σ @xmath6. compute stochastic gradient g t i on batch bi of size ℓ: g t @xmath7, zj). sketch g t i : s (e)t @xmath8. end loop aggregate sketches s (e)@xmath9: ∆ (e)@xmath10: θ @xmath11, σ current @xmath12. end for let σ final @xmath13. end for in gradient compression for federated learning. we describe it in further depth in the main paper. finally, @xcite and @xcite ) accelerate training by bringing adaptive optimizers built for centralized learning into the federated setting.