Katharine Daly, Software Engineer, and Daniel Ramage, Research Director, Google Research

In 2017, Google introduced Federated Learning:https://research.google/blog/federated-learning-collaborative-machine-learning-without-centralized-training-data/ (FL) a machine learning technique that trains models across decentralized, private data. It has been used to power everyday helpful features, including next-word prediction:https://arxiv.org/abs/1811.03604 and Smart Compose:https://support.google.com/gboard/answer/7068415 on Gboard, reply suggestions:https://support.google.com/messages/answer/9327902?hl=en in Google Messages, and Smart Text Selection:https://research.google/blog/predicting-text-selections-with-federated-learning/ in Android.

Federated learning (FL) is a machine learning setting where multiple entities (clients) collaborate in solving a machine learning problem, under the coordination of a service provider. A complete FL system should enable clients to maintain full control over their data, the set of workloads allowed to access their data, and the anonymization properties of those workloads. FL systems should provide appropriate transparency and control to the users whose data is managed by FL clients.

Our TEE-based FL system builds on techniques developed in our earlier work on confidential federated analytics:https://research.google/blog/discovering-new-words-with-confidential-federated-analytics/ and provably private insights:https://research.google/blog/toward-provably-private-insights-into-ai-use/.

The system coordinates four core operational concepts:

For more details on the TEE-based FL system design please see our whitepaper, Toward provably private learning from federated data:https://arxiv.org/abs/2609.31494.

In earlier FL systems, device data was uploaded for the purpose of immediate aggregation, but there was no way for external observers to verify that the data was never logged or inspected. Later, Secure Aggregation:https://arxiv.org/abs/1611.04482 allowed uploads to be protected cryptographically, but was not compatible with state-of-the-art central DP guarantees:https://arxiv.org/abs/2202.08312. Our new TEE-based system represents the next milestone in our ongoing effort to completely remove the need to trust the server operator.

In our new TEE-based FL system, only metrics and differentially private model weights are visible to workload operators. Encrypted training data collected from devices can only be decrypted and processed within TEEs running Python training programs represented in the access policies, and only for a limited amount of time after upload.

Devices participating in our new TEE-based FL system know the full set of server workloads that may access data they’re uploading. The access policies representing these potential future server workloads are published to Rekor:https://docs.sigstore.dev/logging/overview/, a public transparency log, and external auditors are able to observe these logs to track the full set of server-side workloads that devices could potentially be participating in.

The KMS and data processing binaries used in our FL system can be reproducibly built from open source code published in the Confidential Federated Compute:https://github.com/google-parfait/confidential-federated-compute Github repository.

By collecting all device uploads before running the server-side training workload, we no longer have to worry about diurnal variations:https://arxiv.org/abs/1912.04977 in device availability impacting training progress. At server-side execution time, we can dynamically calculate the optimal device participation schedule within the program, and can use it to tune other DP parameters.

In the past, training these FL models could take 1-2 months each, with progress limited by device availability, on-device compute, and competition across multiple training workloads for the same set of device resources. With the new TEE-based system, bottlenecks have been moved to the server, and computation parallelization across many machines allows us to achieve significant speedups in training time, currently only limited by TEE resource availability.

In our new TEE-based FL system, computation of client gradients is shifted to the server, lifting limitations related to on-device compute resources that were present in earlier systems. This paves the way for training increasingly larger models using FL techniques. Integrating TEEs with accelerators will play an important role in such use cases.

The system we have described above is capable of executing not just FL training workloads in a verifiable manner, but also arbitrary workloads that can be expressed using Python. We are experimenting with running other types of workloads on this infrastructure, such as synthetic data generation workloads. Another area of exploration: using the data processing TEEs that execute arbitrary Python in combination with our other data processing TEEs that specialize in functionality such as LLM inference.

The authors would like to thank the collaborators who contributed to the infrastructure design and implementation: Arun Ganesh, Brendan McMahan, Brett McLarnon, Chunxiang (Jake) Zheng, Emily Glanz, Maya Spivak, Michael Reneer, Nova Fallen, Stefan Dierauf, Suxin Guo, Timon Van Overveldt, Yu Xiao, Zachary Charles, and Zachary Garrett. We also thank close partners who supported the Gboard integration: Haicheng Sun, Heng Su, Jianpeng Hou, Liyang Jiang, Noriyuki Takahashi, Wenzhi Mao, Xiaojuan Fang, Yanxiang Zhang, Yingjie Liu, Yuanbo Zhang, and Yun Wang. This work was supported by Corinna Cortes, Shumin Zhai, and Yossi Matias. We additionally thank Maysam Moussalem for feedback on the writing of this post.