Federated learning enables users to train models without transmitting raw data to a central server, avoiding the acquisition of personally identifiable information. Often, this is accomplished by learning a single global model for all users, although their data distributions may differ. For instance, users of a mobile keyboard application may collaborate to train a recommendation model but have varying tastes for the choices. This heterogeneity has prompted the development of algorithms capable of customizing a global model for each user.

However, in some circumstances, privacy concerns may preclude learning a truly global model. Training a truly global federated model would entail sending updates to user embeddings to a central server, revealing the embeddings’ preferences. Even if no user-specific embeddings are used, having some parameters local to user devices reduces server-client communication and allows for responsible personalization of those parameters for each user.

Google’s New Technique

The Google researchers describe a technique for scalable partially local federated learning in “Federated Reconstruction: Partially Local Federated Learning,” presented at NeurIPS 2021. This technique eliminates the need for some model parameters to be aggregated on the server. This strategy trains a recommender model for matrix factorization. Other models benefit from this method since it trains a section of the model to be unique for each user while eliminating transmission of these parameters. The researchers successfully implemented partially local federated learning on Gboard, improving suggestion performance for hundreds of millions of keyboard users. Additionally, they are providing a TensorFlow Federated tutorial that walks through the process of Federated Reconstruction.

Federated Reconstruction

Earlier attempts with partially local federated learning relied on stateful algorithms, which require user devices to maintain a state between federated training rounds. Specifically, these systems necessitated the storage of local parameters across rounds by devices. These algorithms, however, degrade significantly in large-scale federated learning contexts. In these instances, most users will not participate in training. Those who do will likely only experience once, resulting in a rarely available state and may become stale between rounds. Additionally, non-participating users are left with untrained local settings, which precludes practical uses.

How Does It Work?

Federated Reconstruction is a stateless protocol that eliminates the requirement for user devices to retain local parameters by recreating them as needed. When users participate in training, they randomly initialize and train their local parameters using gradient descent on local data with global parameters frozen before updating any globally aggregated model parameters. They can then calculate global parameter updates while keeping local parameters locked.

This straightforward solution removes the difficulties associated with prior methods. It does not assume that users’ previous training states allow for large-scale training, and local parameters are constantly recreated, preventing staleness. Unknown users can still obtain trained models and make inferences merely by recreating local parameters using local data.

Benefits

Compared to other approaches, Federated Reconstruction develops more performant models for unseen users. The technique dramatically outperforms both centralized training and baseline Federated Averaging on a matrix factorization job with unknown users. Federated Reconstruction also enables the personalization of models for heterogeneous users while minimizing model parameter communication – even for models lacking user-specific embeddings. To demonstrate this, the researchers use Federated Reconstruction to personalize a next word prediction language model and notice a significant performance improvement, approaching the accuracy of existing customization methods despite lower communication.Federated Reconstruction also beats other personalization methods when performed at a fixed communication level.

Gboard Deployment in the Real World

To demonstrate the algorithm’s viability in large-scale settings, the researchers put it to Gboard, a popular mobile keyboard application with millions of users. Gboard users communicate via expressions (e.g., GIFs, stickers). Users’ preferences for these expressions vary, making the environment ideal for applying matrix factorization to forecast new expressions that a user would like to contribute.

Additional Investigations

The Google researchers introduced Federated Reconstruction, a partially localized federated learning technique. Federated Reconstruction enables diverse individuals to be personalized while minimizing the transmission of privacy-sensitive parameters. The researchers applied the approach to Gboard by the AI Principles, improving suggestion quality for hundreds of millions of users.

Check out the TensorFlow Federated tutorial for a technical tour of Federated Reconstruction for matrix factorization. Additionally, they published general-purpose TensorFlow Federated libraries and open-source experimentation code.

https://arxiv.org/pdf/2102.03448.pdf