# FedJAX: Federated learning simulation with JAX

arXiv: https://arxiv.org/abs/2108.02117 · HTML: https://mlsanity.com/paper/2108.02117/

## Key facts

- Cited by 32 papers in the map (data through 14 Oct 2026).
- Builds on 29 of its 41 references that are in the map.
- First posted on arXiv 4 Aug 2021.

## Builds on (most cited first, 29 of 29)

- [Communication-Efficient Learning of Deep Networks from Decentralized Data](https://mlsanity.com/paper/1602.05629/) (2016)
- [Federated Learning: Challenges, Methods, and Future Directions](https://mlsanity.com/paper/1908.07873/) (2019)
- [Federated Learning: Strategies for Improving Communication Efficiency](https://mlsanity.com/paper/1610.05492/) (2016)
- [On the Convergence of FedAvg on Non-IID Data](https://mlsanity.com/paper/1907.02189/) (2019)
- [Adaptive Federated Optimization](https://mlsanity.com/paper/2003.00295/) (2020)
- [LEAF: A Benchmark for Federated Settings](https://mlsanity.com/paper/1812.01097/) (2018)
- [Federated Learning for Mobile Keyboard Prediction](https://mlsanity.com/paper/1811.03604/) (2018)
- [Federated Optimization: Distributed Machine Learning for On-Device Intelligence](https://mlsanity.com/paper/1610.02527/) (2016)
- [Clustered Federated Learning: Model-Agnostic Distributed Multi-Task Optimization under Privacy Constraints](https://mlsanity.com/paper/1910.01991/) (2019)
- [Agnostic Federated Learning](https://mlsanity.com/paper/1902.00146/) (2019)
- [Fair Resource Allocation in Federated Learning](https://mlsanity.com/paper/1905.10497/) (2019)
- [FedML: A Research Library and Benchmark for Federated Machine Learning](https://mlsanity.com/paper/2007.13518/) (2020)
- [Parallel Restarted SGD with Faster Convergence and Less Communication: Demystifying Why Model Averaging Works for Deep Learning](https://mlsanity.com/paper/1807.06629/) (2018)
- [Tighter Theory for Local SGD on Identical and Heterogeneous Data](https://mlsanity.com/paper/1909.04746/) (2019)
- [Applied Federated Learning: Improving Google Keyboard Query Suggestions](https://mlsanity.com/paper/1812.02903/) (2018)
- [SCAFFOLD: Stochastic Controlled Averaging for Federated Learning](https://mlsanity.com/paper/1910.06378/) (2019)
- [Expanding the Reach of Federated Learning by Reducing Client Resource Requirements](https://mlsanity.com/paper/1812.07210/) (2018)
- [cpSGD: Communication-efficient and differentially-private distributed SGD](https://mlsanity.com/paper/1805.10559/) (2018)
- [Mime: Mimicking Centralized Stochastic Algorithms in Federated Learning](https://mlsanity.com/paper/2008.03606/) (2020)
- [On the Convergence of Local Descent Methods in Federated Learning](https://mlsanity.com/paper/1910.14425/) (2019)
- [Distributed Mean Estimation with Limited Communication](https://mlsanity.com/paper/1611.00429/) (2016)
- [A generic framework for privacy preserving deep learning](https://mlsanity.com/paper/1811.04017/) (2018)
- [Fairness-aware Agnostic Federated Learning](https://mlsanity.com/paper/2010.05057/) (2020)
- [Ternary Compression for Communication-Efficient Federated Learning](https://mlsanity.com/paper/2003.03564/) (2020)
- [IBM Federated Learning: an Enterprise Framework White Paper V0.1](https://mlsanity.com/paper/2007.10987/) (2020)
- [Mitigating Bias in Federated Learning](https://mlsanity.com/paper/2012.02447/) (2020)
- [Private Federated Learning with Domain Adaptation](https://mlsanity.com/paper/1912.06733/) (2019)
- [Fairness and Accuracy in Federated Learning](https://mlsanity.com/paper/2012.10069/) (2020)
- [Communication-Efficient Agnostic Federated Averaging](https://mlsanity.com/paper/2104.02748/) (2021)

## Built on by (most cited first, 30 of 32)

- [Improving Generalization in Federated Learning by Seeking Flat Minima](https://mlsanity.com/paper/2203.11834/) (2022)
- [Towards Federated Foundation Models: Scalable Dataset Pipelines for Group-Structured Learning](https://mlsanity.com/paper/2307.09619/) (2023)
- [Scaling Language Model Size in Cross-Device Federated Learning](https://mlsanity.com/paper/2204.09715/) (2022)
- [GradSkip: Communication-Accelerated Local Gradient Methods with Better Computational Complexity](https://mlsanity.com/paper/2210.16402/) (2022)
- [Towards Fair Federated Recommendation Learning: Characterizing the Inter-Dependence of System and Data Heterogeneity](https://mlsanity.com/paper/2206.02633/) (2022)
- [Federated Learning as Variational Inference: A Scalable Expectation Propagation Approach](https://mlsanity.com/paper/2302.04228/) (2023)
- [FedLite: A Scalable Approach for Federated Learning on Resource-constrained Clients](https://mlsanity.com/paper/2201.11865/) (2022)
- [Correlated quantization for distributed mean estimation and optimization](https://mlsanity.com/paper/2203.04925/) (2022)
- [JaxPruner: A concise library for sparsity research](https://mlsanity.com/paper/2304.14082/) (2023)
- [FLIX: A Simple and Communication-Efficient Alternative to Local Methods in Federated Learning](https://mlsanity.com/paper/2111.11556/) (2021)
- [Training Machine Learning models at the Edge: A Survey](https://mlsanity.com/paper/2403.02619/) (2024)
- [DrJAX: Scalable and Differentiable MapReduce Primitives in JAX](https://mlsanity.com/paper/2403.07128/) (2024)
- [Beyond Local Sharpness: Communication-Efficient Global Sharpness-aware Minimization for Federated Learning](https://mlsanity.com/paper/2412.03752/) (2024)
- [Flame: Simplifying Topology Extension in Federated Learning](https://mlsanity.com/paper/2305.05118/) (2023)
- [Window-based Model Averaging Improves Generalization in Heterogeneous Federated Learning](https://mlsanity.com/paper/2310.01366/) (2023)
- [Explicit Personalization and Local Training: Double Communication Acceleration in Federated Learning](https://mlsanity.com/paper/2305.13170/) (2023)
- [FedComLoc: Communication-Efficient Distributed Training of Sparse and Quantized Models](https://mlsanity.com/paper/2403.09904/) (2024)
- [Protea: Client Profiling within Federated Systems using Flower](https://mlsanity.com/paper/2207.01053/) (2022)
- [The MICCAI Federated Tumor Segmentation (FeTS) Challenge 2024: Efficient and Robust Aggregation Methods for Federated Learning](https://mlsanity.com/paper/2512.06206/) (2025)
- [Federated Hypergradient Descent](https://mlsanity.com/paper/2211.02106/) (2022)
- [Federated Automatic Differentiation](https://mlsanity.com/paper/2301.07806/) (2023)
- [Efficient Language Model Architectures for Differentially Private Federated Learning](https://mlsanity.com/paper/2403.08100/) (2024)
- [Parameter-Efficient Transfer Learning under Federated Learning for Automatic Speech Recognition](https://mlsanity.com/paper/2408.11873/) (2024)
- [Rieoptax: Riemannian Optimization in JAX](https://mlsanity.com/paper/2210.04840/) (2022)
- [Learning from straggler clients in federated learning](https://mlsanity.com/paper/2403.09086/) (2024)
- [One-Shot Federated Learning with Bayesian Pseudocoresets](https://mlsanity.com/paper/2406.02177/) (2024)
- [Encrypted machine learning of molecular quantum properties](https://mlsanity.com/paper/2212.04322/) (2022)
- [A Comprehensive Empirical Study of Bugs in Open-Source Federated Learning Frameworks](https://mlsanity.com/paper/2308.05014/) (2023)
- [Federated Deep Equilibrium Learning: Harnessing Compact Global Representations to Enhance Personalization](https://mlsanity.com/paper/2309.15659/) (2023)
- [Heterogeneous Federated Learning Using Knowledge Codistillation](https://mlsanity.com/paper/2310.02549/) (2023)

## Abstract (arXiv)

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.