FedJAX: Federated learning simulation with JAX
It addresses the need for efficient and user-friendly tools in federated learning research, though it is incremental as it builds on existing libraries and methods.
FedJAX is a JAX-based open source library designed to simplify federated learning simulations for researchers, enabling training on datasets like EMNIST in minutes and Stack Overflow in about an hour with standard hyperparameters using TPUs.
Federated learning is a machine learning technique that enables training across decentralized data. Recently, federated learning has become an active area of research due to an increased focus on privacy and security. In light of this, a variety of open source federated learning libraries have been developed and released. We introduce FedJAX, a JAX-based open source library for federated learning simulations that emphasizes ease-of-use in research. With its simple primitives for implementing federated learning algorithms, prepackaged datasets, models and algorithms, and fast simulation speed, FedJAX aims to make developing and evaluating federated algorithms faster and easier for researchers. Our benchmark results show that FedJAX can be used to train models with federated averaging on the EMNIST dataset in a few minutes and the Stack Overflow dataset in roughly an hour with standard hyperparameters using TPUs.