[Paper Review] Federated Learning for Emoji Prediction in a Mobile Keyboard

Swaroop Ramaswamy Pillai, Rajiv Mathews+2 more|arXiv (Cornell University)|Jun 11, 2019

Digital Communication and LanguageComputer Science164 citations

TL;DR

The paper trains a CIFG-LSTM emoji predictor for Gboard using federated learning, showing FL can outperform server-trained models while keeping user data on-device, and introduces triggering and diversification mechanisms.

ABSTRACT

We show that a word-level recurrent neural network can predict emoji from text typed on a mobile keyboard. We demonstrate the usefulness of transfer learning for predicting emoji by pretraining the model using a language modeling task. We also propose mechanisms to trigger emoji and tune the diversity of candidates. The model is trained using a distributed on-device learning framework called federated learning. The federated model is shown to achieve better performance than a server-trained model. This work demonstrates the feasibility of using federated learning to train production-quality models for natural language understanding tasks while keeping users' data on their devices.

Motivation & Objective

• Demonstrate emoji prediction from text on mobile keyboards using a word-level RNN.
• Show that pretraining on language modeling speeds up emoji prediction learning.
• Propose triggering and diversification mechanisms to balance usefulness and variety of emoji suggestions.

Proposed method

• Use a two-layer CIFG-LSTM with 256 units per layer and 10,000 word vocabulary to predict 100 emoji.
• Pretrain layers (except output projection) on a language modeling task to improve convergence.
• Employ Federated Averaging to aggregate client updates from on-device training rounds.
• Introduce an on-device triggering mechanism by adding an <UNK> class to control when emoji predictions are shown.
• Apply diversification by scaling predicted emoji probabilities with empirical emoji frequencies (S_i = P_hat(emoji=i|text) / P(emoji=i)^alpha).
• Train and evaluate with server-based and federated setups; compare Accuracy@1 and AUC metrics.

Experimental results

Research questions

• RQ1 Can a CIFG-LSTM emoji predictor trained with federated learning outperform a server-trained model for on-device emoji suggestion?
• RQ2 Does language-model pretraining improve convergence and performance of the emoji predictor in a federated setting?
• RQ3 How do triggering and diversification strategies affect emoji prediction usefulness and variety?

Key findings

• Federated training achieved higher Accuracy@1 than the best server-trained model (0.256 vs 0.239).
• Federated model achieved similar or better prediction performance but had lower AUC than server-trained model when evaluated on server data.
• Large client batch sizes (B) and more devices per round (K) improve model quality, with diminishing returns beyond certain values.
• Momentum-based server updates (momentum=0.9 with Nesterov) outperform SGD variants in federated experiments.
• Live traffic show higher CTR, emoji shares, and DAU for federated model compared to server-trained model.
• Low latency inference (~1 ms) with TensorFlow Lite on mobile.