[Paper Review] Federated Learning for Mobile Keyboard Prediction

Andrew Hard, K. Praveen Kumar Rao+7 more|arXiv (Cornell University)|Nov 8, 2018

Privacy-Preserving Technologies in Data27 references1,105 citations

TL;DR

The paper trains a CIFG LSTM language model for next-word prediction using federated learning on mobile devices and shows it can outperform a server-trained model and a baseline n-gram model in recall metrics. It demonstrates feasibility and privacy benefits of on-device training for mobile keyboards.

ABSTRACT

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.

Motivation & Objective

• Motivate and enable next-word prediction on mobile keyboards without sending raw user data to servers.
• Develop a CIFG-based LSTM language model suitable for on-device training and inference.
• Evaluate federated learning (FederatedAveraging) against server-based SGD in a commercial setting.
• Demonstrate privacy advantages and practical feasibility of federated training for mobile NLP tasks.

Proposed method

• Use Coupled Input-Forget Gates (CIFG) LSTM with 670 units and a 10,000-word vocabulary for on-device next-word prediction.
• Train from scratch on server and on-device data using federated averaging (FedAvg) with client updates aggregated by a server.
• Keep updates as model weights only, with on-device data never uploaded.
• Evaluate recall (top-1 and top-3) and production metrics via server-hosted logs, client caches, and live experiments.
• Quantize weights and tie input embedding with output projection to reduce size to ~1.4 MB on-device.

Experimental results

Research questions

• RQ1 Can federated learning match or exceed server-based training for neural language models on mobile keyboards?
• RQ2 Does on-device federated training improve next-word prediction recall compared with a strong n-gram baseline?
• RQ3 What are the privacy and deployment implications of deploying federated learning for mobile keyboard prediction?
• RQ4 How does federated CIFG perform in real-world live user experiments compared with server-trained models?

Key findings

• Federated CIFG achieves higher recall than the baseline n-gram model on both server logs and client caches.
• Top-1 recall: Federated CIFG 16.4% vs Server CIFG 16.5% on server data; Top-3 recall: Federated CIFG 27.0% vs Server CIFG 27.1% on server data.
• On client-owned caches, Federated CIFG shows Top-1 recall 15.8% vs Server CIFG 15.0% (higher by 0.8 percentage points).
• Live production experiments show Federated CIFG top-1 recall 5.82% vs Server CIFG 5.76% (top-1), and Top-3 recall 13.75% vs 13.63% (top-3).
• Federated CIFG yields a 1% relative improvement in top-1 and top-3 impression recall over server CIFG in live experiments.
• Federated training processed 600 million sentences across 1.5 million clients over ~4–5 days, converging after 3000 rounds.