FedJAX: Federated learning simulation with JAX
arXiv:2108.02117
Abstract
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.
References in corpus (19)
- Communication-Efficient Learning of Deep Networks from Decentralized Data
- Federated Learning: Strategies for Improving Communication Efficiency
- Federated Optimization: Distributed Machine Learning for On-Device Intelligence
- Federated Learning for Mobile Keyboard Prediction
- On the Convergence of FedAvg on Non-IID Data
- SCAFFOLD: Stochastic Controlled Averaging for Federated Learning
- Applied Federated Learning: Improving Google Keyboard Query Suggestions
- Federated Optimization in Heterogeneous Networks
- FedML: A Research Library and Benchmark for Federated Machine Learning
- A generic framework for privacy preserving deep learning
- LEAF: A Benchmark for Federated Settings
- On the Convergence of Local Descent Methods in Federated Learning
- Expanding the Reach of Federated Learning by Reducing Client Resource Requirements
- Distributed Mean Estimation with Limited Communication
- Adaptive Federated Optimization
- IBM Federated Learning: an Enterprise Framework White Paper V0.1
- Mime: Mimicking Centralized Stochastic Algorithms in Federated Learning
- Private Federated Learning with Domain Adaptation
- Communication-Efficient Agnostic Federated Averaging