mlsanityData to 7 Oct 2026

FedJAX: Federated learning simulation with JAX

arXiv 2108.02117 ↗First posted 4 Aug 2021cs.LGcited by 32

Not on the map. Why?

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.

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

Builds on

29 of 41 references.

References · 29

  1. 2016Communication-Efficient Learning of Deep Networks from Decentralized Datacited by 7,592
  2. 2019Federated Learning: Challenges, Methods, and Future Directionscited by 1,531
  3. 2016Federated Learning: Strategies for Improving Communication Efficiencycited by 1,461
  4. 2019On the Convergence of FedAvg on Non-IID Datacited by 1,220
  5. 2020Adaptive Federated Optimizationcited by 1,036
  6. 2018LEAF: A Benchmark for Federated Settingscited by 780
  7. 2018Federated Learning for Mobile Keyboard Predictioncited by 565
  8. 2016Federated Optimization: Distributed Machine Learning for On-Device Intelligencecited by 558
  9. 2019Clustered Federated Learning: Model-Agnostic Distributed Multi-Task Optimization under Privacy Constraintscited by 499
  10. 2019Agnostic Federated Learningcited by 451
  11. 2019Fair Resource Allocation in Federated Learningcited by 389
  12. 2020FedML: A Research Library and Benchmark for Federated Machine Learningcited by 298
17 more
  1. 2018Parallel Restarted SGD with Faster Convergence and Less Communication: Demystifying Why Model Averaging Works for Deep Learningcited by 277
  2. 2019Tighter Theory for Local SGD on Identical and Heterogeneous Datacited by 266
  3. 2018Applied Federated Learning: Improving Google Keyboard Query Suggestionscited by 202
  4. 2019SCAFFOLD: Stochastic Controlled Averaging for Federated Learningcited by 174
  5. 2018Expanding the Reach of Federated Learning by Reducing Client Resource Requirementscited by 162
  6. 2018cpSGD: Communication-efficient and differentially-private distributed SGDcited by 157
  7. 2020Mime: Mimicking Centralized Stochastic Algorithms in Federated Learningcited by 152
  8. 2019On the Convergence of Local Descent Methods in Federated Learningcited by 150
  9. 2016Distributed Mean Estimation with Limited Communicationcited by 134
  10. 2018A generic framework for privacy preserving deep learningcited by 103
  11. 2020Fairness-aware Agnostic Federated Learningcited by 87
  12. 2020Ternary Compression for Communication-Efficient Federated Learningcited by 57
  13. 2020IBM Federated Learning: an Enterprise Framework White Paper V0.1cited by 57
  14. 2020Mitigating Bias in Federated Learningcited by 56
  15. 2019Private Federated Learning with Domain Adaptationcited by 36
  16. 2020Fairness and Accuracy in Federated Learningcited by 27
  17. 2021Communication-Efficient Agnostic Federated Averagingcited by 10

Built on by

Cited by year

Citing papers per year: 2021: 1; 2022: 10; 2023: 10; 2024: 9; 2025: 2.

32 papers.

Sort:

  1. 2022Improving Generalization in Federated Learning by Seeking Flat Minimacited by 98
  2. 2023Towards Federated Foundation Models: Scalable Dataset Pipelines for Group-Structured Learningcited by 29
  3. 2022Scaling Language Model Size in Cross-Device Federated Learningcited by 25
  4. 2022Towards Fair Federated Recommendation Learning: Characterizing the Inter-Dependence of System and Data Heterogeneitycited by 17
  5. 2022GradSkip: Communication-Accelerated Local Gradient Methods with Better Computational Complexitycited by 17
  6. 2023Federated Learning as Variational Inference: A Scalable Expectation Propagation Approachcited by 17
  7. 2022FedLite: A Scalable Approach for Federated Learning on Resource-constrained Clientscited by 14
  8. 2022Correlated quantization for distributed mean estimation and optimizationcited by 14
  9. 2023JaxPruner: A concise library for sparsity researchcited by 13
  10. 2021FLIX: A Simple and Communication-Efficient Alternative to Local Methods in Federated Learningcited by 12
  11. 2024Training Machine Learning models at the Edge: A Surveycited by 11
  12. 2024DrJAX: Scalable and Differentiable MapReduce Primitives in JAXcited by 9
20 more
  1. 2024Beyond Local Sharpness: Communication-Efficient Global Sharpness-aware Minimization for Federated Learningcited by 8
  2. 2023Flame: Simplifying Topology Extension in Federated Learningcited by 7
  3. 2023Window-based Model Averaging Improves Generalization in Heterogeneous Federated Learningcited by 7
  4. 2023Explicit Personalization and Local Training: Double Communication Acceleration in Federated Learningcited by 5
  5. 2024FedComLoc: Communication-Efficient Distributed Training of Sparse and Quantized Modelscited by 5
  6. 2022Protea: Client Profiling within Federated Systems using Flowercited by 4
  7. 2022Federated Hypergradient Descentcited by 2
  8. 2023Federated Automatic Differentiationcited by 2
  9. 2024Parameter-Efficient Transfer Learning under Federated Learning for Automatic Speech Recognitioncited by 2
  10. 2025The MICCAI Federated Tumor Segmentation (FeTS) Challenge 2024: Efficient and Robust Aggregation Methods for Federated Learningcited by 2
  11. 2022Rieoptax: Riemannian Optimization in JAXcited by 1
  12. 2024Efficient Language Model Architectures for Differentially Private Federated Learningcited by 1
  13. 2024Learning from straggler clients in federated learningcited by 1
  14. 2024One-Shot Federated Learning with Bayesian Pseudocoresetscited by 1
  15. 2022Encrypted machine learning of molecular quantum propertiescited by 0
  16. 2023A Comprehensive Empirical Study of Bugs in Open-Source Federated Learning Frameworkscited by 0
  17. 2023Federated Deep Equilibrium Learning: Harnessing Compact Global Representations to Enhance Personalizationcited by 0
  18. 2023Heterogeneous Federated Learning Using Knowledge Codistillationcited by 0
  19. 2024Where is the Testbed for my Federated Learning Research?cited by 0
  20. 2025Decoding FL Defenses: Systemization, Pitfalls, and Remediescited by 0