Scaling Language Model Size in Cross-Device Federated Learning

TL;DR

This paper demonstrates that by applying various optimization techniques, larger language models can be effectively trained in cross-device federated learning, achieving comparable or better performance with reduced communication costs.

Contribution

The authors introduce a combination of methods enabling the training of significantly larger language models in federated settings, overcoming traditional bottlenecks.

Findings

Trained a 21M parameter Transformer with comparable perplexity to smaller models.

Achieved 10x reduction in client-server communication costs.

Lower perplexity than smaller LSTMs in federated learning scenarios.

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 $21$M parameter Transformer and $20.2$M parameter Conformer that achieve the same or better 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.