Source-linked AI summary
Federated Learning for Mobile Keyboard Prediction
Andrew Hard, Kanishka Rao, Rajiv Mathews, Swaroop Ramaswamy, Françoise Beaufays, Sean Augenstein, Hubert Eichner, Chloé Kiddon, Daniel Ramage
TL;DR
The paper asks whether a mobile-keyboard language model can be trained without exporting sensitive user text while addressing the mismatch between public or logged data and users’ typing distribution. It trains a CIFG recurrent model centrally and with Federated Averaging on client devices, finding that federated training improves recall and supports a privacy-oriented deployment model. The authors conclude that federated learning is feasible for commercial mobile-keyboard language modeling and can improve model quality while retaining data on devices.
Problem
Training data must represent users’ typing, but public datasets may mismatch that distribution and centrally collecting personal text requires logging, storage, security, and remote data retention.
Method
The paper trains a CIFG recurrent neural language model from scratch with server-based SGD and client-device training using Federated Averaging.
Results
The federated CIFG improves client-cache top-1 recall by 5% relative, or 0.8% absolute, over the server-trained CIFG, while both CIFG models improve recall over the n-gram FST baseline.
Takeaways & Limitations
Federated learning can train a commercial mobile-keyboard language model across distributed devices while improving language-model quality and retaining user data on devices.
Takeaways & Limitations
Federated training requires users to trust the aggregation server not to scrutinize individual weight uploads, and participating devices must meet charging, network, idle, memory, and location requirements.
Abstract
from arXiv · showhide
We train a recurrent neural network language model using a distributed, on-device learning framework called federated learning for the purpose of next-word prediction in a virtual keyboard for smartphones. Server-based training using stochastic gradient descent is compared with training on client devices using the Federated Averaging algorithm. The federated algorithm, which enables training on a higher-quality dataset for this use case, is shown to achieve better prediction recall. This work demonstrates the feasibility and benefit of training language models on client devices without exporting sensitive user data to servers. The federated learning environment gives users greater control over the use of their data and simplifies the task of incorporating privacy by default with distributed training and aggregation across a population of client devices.
1. INTRODUCTION
Mobile keyboards need fast, compact next-word prediction, but training data must represent users without requiring remote collection of personal text. This paper evaluates recurrent neural models trained centrally and federatedly against Gboard’s n-gram baseline.
- Gboard provides next-word prediction from preceding user text, such as suggesting “and,” “too,” or “so much” after “I love you”.
- Before this work, Gboard generated predictions with a Katz-smoothed interpolated 5-gram finite-state-transducer language model containing 1.25 million n-grams.Personalized user history, contacts, and email n-gram models augmented the primary model.
- Mobile keyboard models must remain small, provide inference responses within 20 milliseconds, and limit CPU use to preserve battery life.These constraints typically restrict models to tens of megabytes and vocabularies of hundreds of thousands of words.
- RNNs can use dynamically sized context windows, while LSTM variants address exploding and vanishing gradients in back-propagation through time.
- Representative user text is difficult to obtain centrally because public data may mismatch users’ distribution, while collecting personal text requires logging, storage, infrastructure, and security.Users may be uncomfortable with remote collection and storage of personal data.
- Federated learning offers an alternative to server-based data collection and training, and the paper reports recall improvements over the FST decoder baseline.
2. RELATED WORK
Prior work covers mobile decoding, recurrent language modeling, and federated training, but published work had not yet applied federated learning to a neural mobile-keyboard language model.
- Prior mobile-keyboard research explored FST decoding, LSTM-based gestured-input decoding, and RNN language models optimized under latency and memory constraints.
- Federated learning extends server-based distributed training to client devices using locally stored data, and has been combined with differential privacy for language models.
- Gboard had previously used federated learning for search-query suggestions, but those results had not been published.
- The authors state that no existing publications had trained a neural language model for a mobile keyboard with federated learning.
3. MODEL ARCHITECTURE
The next-word predictor uses a compact CIFG recurrent architecture designed for mobile deployment, with tied embeddings and output projections to reduce model size and training cost.
- CIFG reduces computations and parameter-set size without affecting model performance, supporting the mobile-device environment.The model is trained with TensorFlow, and on-device inference is supported by TensorFlow Lite.
- Tied input-embedding and output-projection matrices reduce model size and accelerate training.A softmax converts output logits into normalized probabilities for cross-entropy training.
- The model uses a 10,000-word vocabulary, embedding dimension D = 96, and a single CIFG layer with 670 units.
- Federated learning repeatedly computes local updates on client devices, aggregates them on a server, and returns a new model to clients.
4. FEDERATED LEARNING
Federated learning trains the model by keeping user data on client devices, sending ephemeral model updates to a server, and aggregating those updates into successive global models.
- Clients process locally stored personal data and share model updates instead of uploading training data for centralized training.Aggregated updates from many clients produce an improved global model.
- Federated Averaging sends global model w_t to participating clients, which begin from a common initialization and have device-specific datasets of n_k examples.In Gboard studies, n_k varies with users’ typing volume.
- Each client computes an average gradient with local SGD, after which the server performs weighted aggregation to obtain the next global model.
- Client batch size, client epochs, and the number of clients per round are tuned to improve performance.
- Updates are ephemeral, processed in memory, immediately discarded after aggregation, and limited to model weights rather than user data.The procedure still requires users to trust that the aggregation server will not scrutinize individual uploads.
5. EXPERIMENTS
The experiments train CIFG next-word prediction models with server-based SGD and federated averaging, using logged text or client-held caches. Federated training aggregates updates from eligible mobile devices across repeated training rounds.
- Both server-based and federated training start CIFG language models from random weight initializations.
- Server-based training with logs data: Server training uses anonymized, opt-in Gboard logs from U.S. English-speaking users, comprising approximately 7.5 billion training sentences.The logs contain short phrases and are filtered to begin with a start-of-sentence token.
- Federated training with client caches: Federated training uses text and decoder-generated prediction candidates stored in local caches on Gboard client devices.
- Federated training with client caches: Eligible federated clients require at least 2 GB of available memory and must be charging, idle, and connected to an un-metered network.These participation criteria are specific to this Gboard implementation rather than inherent to federated learning.
6. RESULTS
The results compare CIFG and n-gram FST prediction recall across server logs, client caches, and live traffic. Federated CIFG generally outperforms both the baseline and server-trained CIFG, although logs-data recall is comparable between the two CIFG models.
- Recall is defined as correct predictions divided by total tokens, with top-1 recall especially relevant to the center suggestion position.Top-3 recall is also evaluated because Gboard displays three candidates.
- Both server-trained and federated CIFG models improve top-1 and top-3 recall over the n-gram FST baseline.The comparison favors CIFG despite the n-gram model’s larger vocabulary and personalized components.
- Evaluation across data sources: 5% relative (0.8% absolute): federated CIFG’s top-1 recall improvement over server-trained CIFG on client cache data.On server-hosted logs data, the two CIFG models have comparable recall.
7. CONCLUSION
The paper concludes that federated learning can train a CIFG language model from scratch for mobile keyboard prediction. The federated model outperforms the server-trained CIFG and baseline n-gram model on the task.
- Federated learning trains a CIFG language model from scratch that outperforms both an identical server-trained CIFG and a baseline n-gram model.The paper presents this as an early commercial application of federated language modeling.