Scaling Language Model Size in Cross-Device Federated Learning (2022)

Abstract: Most studies in cross-device federated learning focus on small models, due to the server-client communication and on-device computation bottlenecks. In this work, we leverage various techniques for mitigating these bottlenecks to train larger language models in cross-device federated learning. With systematic applications of partial model training, quantization, efficient transfer learning, and communication-efficient optimizers, we are able to train a 21M parameter Transformer that achieves the same perplexity as that of a similarly sized LSTM with \sim10\times smaller client-to-server communication cost and 11% lower perplexity than smaller LSTMs commonly studied in literature.

## introduction federated learning is a distributed training technique, where a model is trained on data distributed across clients or edge devices without usergenerated data ever leaving the device, providing an additional layer of privacy and security @xcite . we refer readers to @xcite for a detailed literature survey on federated learning. federated learning has been used in several applications including virtual keyboard applications @xcite , keyword spotting @xcite , and healthcare @xcite . language models (lm) have many uses in language-based applications including virtual keyboard @xcite and automatic speech recognition @xcite @xcite . recently, there has been increased interest in training progressively larger and deeper lms with impressive quality improvements in downstream tasks, including question answering, text classification, and text summarization @xcite @xcite @xcite . these models tend to be variants of the transformer @xcite . federated learning is typically studied in two scenarios: cross-silo, where the number of clients is small, and cross-device, where the number of clients can be in the order of millions @xcite . in this work we focus on cross-device, where devices are typically edge devices such as cell phones, with limited computation and communication capabilities. hence, the major benchmark lms tend to be very limited in size @xcite @xcite @xcite because memory, computation, and communication are critical bottlenecks @xcite . in particular, previous works that train federated lms in production settings have used coupled input forget gate (cifg) long shortterm memory (lstm) models with fewer than 4 million parameters @xcite @xcite . these resource constraints have motivated research into various efficient algorithms for training larger models with federated learning @xcite . however, most of these techniques are still evaluated on relatively small models compared to their server-based counterparts. in this work, we systematically evaluate multiple strategies for mitigating communication and computation costs of training larger lms to determine if the impressive quality gains from larger models can also be achieved in cross-device federated learning. while there are previous works on efficient transformers @xcite , we forgo these efficient variants as they may actually be more inefficient when sequences are short @xcite . additionally, @xcite ; @xcite ; @xcite trained large transformer models in the cross-silo setting, where devices have more resources, whereas we focus on the resource-constrained cross-device setting. recent large lms, such as @xcite , contain hundreds of billions of parameters, which is substantially bigger than the memory limits of edge devices. therefore in this work, we consider large models to be at most 25 million parameters, which is still considerably larger than existing models trained on-device. the rest of the paper is organized as follows. in section 2, we overview our contributions. in section 3, we detail the dataset and models. we then analyze techniques to reduce the per-round cost in section 4, and the number of communication rounds in section 5. finally in section 6, we combine techniques and demonstrate that large transformers can be trained using many fewer rounds and significantly lower communication and computation cost. ## our contributions we explore two regimes: small models typically studied in cross-device federated learning with fewer than 5m parameters and new larger models with at most 25m parameters. we study two architectures: cifg-lstm @xcite , or lstm for simplicity, @xcite and transformer @xcite . our contributions are the following: ## dataset and models in this section, we describe the models and dataset used in the rest of the paper. we train on the stack overflow federated dataset from tff (2018), which contains posts from the public forum grouped by username. following trends in training transformers, we use sentence-piece @xcite for sub-word tokenization with a vocabulary size of 4k. the sentence-piece model is computed based on the entire stack overflow training corpus in an offline process on server. during federated learning, this fixed sentence-piece model is transmitted to each client to encode the local text data. doing so provides greater coverage for cross-dataset applications as well as potential downstream speech applications such as asr @xcite . we measure performance on next-subword prediction using test perplexity. see appendix a for descriptive dataset statistics. all experiments were implemented using jax @xcite and fedjax (ro et al., 2021) federated simulation libraries. we first did a hyperparameter search for each model and size (≤ 5m @xmath0, with fedadam (reddi et al., 2020), or fedavg for simplicity, with 200 clients per round for 3k rounds, resulting in four models: small lstm (4.7m), large lstm (18.8m), small transformer (4.1m), and large transformer (21m). we then trained the chosen architectures with 800 clients per round for 10k rounds in figure 1 . as expected, the larger variants significantly outperform their smaller counterparts with the large transformer achieving the best perplexity. however, the larger models are more expensive to train per round and although the large transformer achieves the best perplexity, it only surpasses the large lstm after 4k rounds. next, we focus on techniques to reduce this cost per round and number of rounds. for more details about the architecture search, the selected models, and their performance, see appendix a. ## cost per round the larger models have 18.8m and 21m parameters (150mb and 168mb, at 32 bits per parameter) which need to be downloaded, trained, and uploaded at each round, a strain on both communication and computation on device. there are often strict time or transfer byte limits for each round of training, which can prohibit some devices from training these models due to slower transfer/processing speeds @xcite . we show that we can significantly reduce these costs by partial model training and quantization techniques. partial model training: training only a subset of the model can reduce the computational cost of training and has been examined in both federated @xcite and nonfederated @xcite settings. additionally, reducing the number of trainable parameters can also decrease communication cost since only the trainable parameters need to be uploaded. we follow the partial variable training (pvt) per client per round strategy @xcite as it only freezes a subset of the original model and can be applied generally to multiple model architecture types. for more experiment details, see appendix b. we report test perplexity as a function of number of trainable variables in figure 2 . large lstm seems to be able to handle more aggressive parameter freezing compared to large transformer in terms of quality regression. however, training only 40% of variables for the large transformer (6.3m) achieves better performance than the full large lstm (18.8m). quantization: to reduce communication costs, various quantization strategies can decrease the number of bits required to represent model parameters @xcite @xcite . we examine stochastic k-level uniform quantization @xcite as it can be applied to model parameters on download (server-to-client) and model updates on upload (client-to-server) communication with adjustable levels of compression, and compare with terngrad, an upload technique @xcite . we focus analysis on larger models which are more affected by quantization. the lstm appears more "quantizable" during download than the transformer, with less regression in figure 3 . the perplexity of the transformer with 16 download bits matches that of the baseline transformer and with 12 bits its perplexity is close to that of the lstm. for both the models, 8 bit upload matches the corresponding baselines, or even 6 bits for the lstm in figure 4 . terngrad, requiring log 2 (3) bits, outperforms the 4 bit in the transformer but not for the lstm in figure 5 . more details are in appendix c. ## number of communication rounds transfer learning: transfer learning leverages pretrained models to improve model quality @xcite . by pretraining, the number of communication rounds required for model convergence can be significantly reduced @xcite . we use two datasets for pretraining: a large corpus of digitized books @xcite and the one billion word benchmark (lm1b) @xcite . after pretraining using synchronous sgd for 30m steps, we finetune on stack overflow using fedavg. for additional details, see appendix d. we report results for each of the pretraining datasets and random initialization in figure 6 . books consistently outperforms lm1b for both the lstm and transformer. pretraining greatly benefits the large transformer compared to the large lstm, reducing the number of rounds needed to reach the final 10k without pretraining by 4k rounds. furthermore, at round 2k, the large transformer already outperforms the large lstm, making the number of rounds needed for training similar to that of smaller models used in mobile keyboard prediction @xcite . different optimizers: since the introduction of fedavg, several variations continue to be developed @xcite ; reddi et al., 2020). specifically, we examine mimelite (karimireddy et al., 2020) and fedprox @xcite as they have been shown to reduce the total amount of rounds required for provable convergence. however, in figure 7 , fedprox and mimelite do not improve convergence speed over fedavg. more details can be found in appendix e. ## combination of techniques we experiment with combining partial model training, quantization, and transfer learning to train efficient larger models. for these experiments, we train on just 40% of trainable parameters with pvt and warm start after pretraining on the books corpus. combining download quantization with these techniques did not perform as well, so we only apply 8 bit uniform quantization on upload, which is the tightest communication bottleneck @xcite reports that mobile upload speeds worldwide are over 4× slower than download as of may 2021). for the full experiment details, refer to appendix f. we report the test perplexity in terms of total upload communication cost in figure 8 . restricting for small upload costs (< 200gb), the efficient models outperform all others with the efficient large transformer yielding the best perplexity. furthermore, the efficient large transformer also achieves the same perplexity as the large lstm with no efficient techniques. ## conclusion we systematically studied several techniques for addressing the communication and computation bottlenecks of federated learning. we further demonstrated that these techniques, individually or in combination, can scale to larger models in crossdevice federated learning. extending this study to other architectures and efficient strategies remains an interesting open question. also search over client learning rate, client batch size, client max number of examples (with client number of epochs fixed to 1), client ℓ 2 norm for clipping, and server learning rate. the search ranges as well as selected values for each model are detailed in as an additional baseline comparison, we also train each model using synchronous sgd to observe model quality in terms of number of gradient computations. these centralized baselines provide a rough estimate of an upper bound on model quality for federated learning. to produce a reasonable comparison between the federated and centralized experiments, we compare by number of gradient computations. we approximate the number of gradient steps taken for federated learning with 200 clients per round for 10k communication rounds. we train the centralized models using the adam optimizer and run periodic evaluation on the test set at the same frequency as the federated experiments. we report and compare final metrics between centralized training and federated averaging on the test set in figure 10 . observing the test perplexity over gradient steps, it is evident that the relative rankings of the models remain consistent between centralized and federated baselines. additionally, by 10k rounds, the large federated models seem to approach somewhat close in perplexity to their centralized counterparts.