# Federated Automatic Differentiation

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

## Key facts

- Cited by 2 papers in the map (data through 14 Oct 2026).
- Builds on 69 of its 98 references that are in the map.
- First posted on arXiv 18 Jan 2023.

## Builds on (most cited first, 30 of 69)

- [Adam: A Method for Stochastic Optimization](https://mlsanity.com/paper/1412.6980/) (2014)
- [PyTorch: An Imperative Style, High-Performance Deep Learning Library](https://mlsanity.com/paper/1912.01703/) (2019)
- [Communication-Efficient Learning of Deep Networks from Decentralized Data](https://mlsanity.com/paper/1602.05629/) (2016)
- [Model-Agnostic Meta-Learning for Fast Adaptation of Deep Networks](https://mlsanity.com/paper/1703.03400/) (2017)
- [Outrageously Large Neural Networks: The Sparsely-Gated Mixture-of-Experts Layer](https://mlsanity.com/paper/1701.06538/) (2017)
- [Advances and Open Problems in Federated Learning](https://mlsanity.com/paper/1912.04977/) (2019)
- [Federated Optimization in Heterogeneous Networks](https://mlsanity.com/paper/1812.06127/) (2018)
- [Neural Architecture Search with Reinforcement Learning](https://mlsanity.com/paper/1611.01578/) (2016)
- [Adaptive Federated Optimization](https://mlsanity.com/paper/2003.00295/) (2020)
- [Tackling the Objective Inconsistency Problem in Heterogeneous Federated Optimization](https://mlsanity.com/paper/2007.07481/) (2020)
- [Automatic differentiation in machine learning: a survey](https://mlsanity.com/paper/1502.05767/) (2015)
- [Measuring the Effects of Non-Identical Data Distribution for Federated Visual Classification](https://mlsanity.com/paper/1909.06335/) (2019)
- [Efficient Neural Architecture Search via Parameter Sharing](https://mlsanity.com/paper/1802.03268/) (2018)
- [Towards Federated Learning at Scale: System Design](https://mlsanity.com/paper/1902.01046/) (2019)
- [Byzantine-Robust Distributed Learning: Towards Optimal Statistical Rates](https://mlsanity.com/paper/1803.01498/) (2018)
- [Learning to learn by gradient descent by gradient descent](https://mlsanity.com/paper/1606.04474/) (2016)
- [Federated Learning for Mobile Keyboard Prediction](https://mlsanity.com/paper/1811.03604/) (2018)
- [Meta-Learning with Latent Embedding Optimization](https://mlsanity.com/paper/1807.05960/) (2018)
- [Gradient-based Hyperparameter Optimization through Reversible Learning](https://mlsanity.com/paper/1502.03492/) (2015)
- [Robust Aggregation for Federated Learning](https://mlsanity.com/paper/1912.13445/) (2019)
- [The Non-IID Data Quagmire of Decentralized Machine Learning](https://mlsanity.com/paper/1910.00189/) (2019)
- [Differentially Private Learning with Adaptive Clipping](https://mlsanity.com/paper/1905.03871/) (2019)
- [Applied Federated Learning: Improving Google Keyboard Query Suggestions](https://mlsanity.com/paper/1812.02903/) (2018)
- [Optimal rates for zero-order convex optimization: the power of two function evaluations](https://mlsanity.com/paper/1312.2139/) (2013)
- [Learning to Optimize](https://mlsanity.com/paper/1606.01885/) (2016)
- [Federated Learning with Buffered Asynchronous Aggregation](https://mlsanity.com/paper/2106.06639/) (2021)
- [Federated Learning with Compression: Unified Analysis and Sharp Guarantees](https://mlsanity.com/paper/2007.01154/) (2020)
- [ProxSkip: Yes! Local Gradient Steps Provably Lead to Communication Acceleration! Finally!](https://mlsanity.com/paper/2202.09357/) (2022)
- [Linear Convergence in Federated Learning: Tackling Client Heterogeneity and Sparse Gradients](https://mlsanity.com/paper/2102.07053/) (2021)
- [Online Learning Rate Adaptation with Hypergradient Descent](https://mlsanity.com/paper/1703.04782/) (2017)

## Built on by (most cited first, 2 of 2)

- [DrJAX: Scalable and Differentiable MapReduce Primitives in JAX](https://mlsanity.com/paper/2403.07128/) (2024)
- [A Typed Tensor Language for Shared-State Federated Computation](https://mlsanity.com/paper/2605.21103/) (2026)

## Abstract (arXiv)

Federated learning (FL) is a general framework for learning across an axis of group partitioned data (heterogeneous clients) while preserving data privacy, under the orchestration of a central server. FL methods often compute gradients of loss functions purely locally (ie. entirely at each client, or entirely at the server), typically using automatic differentiation (AD) techniques. We propose a federated automatic differentiation (FAD) framework that 1) enables computing derivatives of functions involving client and server computation as well as communication between them and 2) operates in a manner compatible with existing federated technology. In other words, FAD computes derivatives across communication boundaries. We show, in analogy with traditional AD, that FAD may be implemented using various accumulation modes, which introduce distinct computation-communication trade-offs and systems requirements. Further, we show that a broad class of federated computations is closed under these various modes of FAD, implying in particular that if the original computation can be implemented using privacy-preserving primitives, its derivative may be computed using only these same primitives. We then show how FAD can be used to create algorithms that dynamically learn components of the algorithm itself. In particular, we show that FedAvg-style algorithms can exhibit significantly improved performance by using FAD to adjust the server optimization step automatically, or by using FAD to learn weighting schemes for computing weighted averages across clients.