This repository contains a minimalistic (yet general and modular) JAX implementation of federated posterior averaging (FedPA) algorithm along with a variety of simulation experiments on synthetically generated problems.

The easiest way to reproduce our synthetic experiments and/or compare FedAvg and FedPA is by using the colab notebook provided in this repository (simply click the Open in Colab button at the top of README). If you would like to use our JAX implementation of FedAvg or FedPA elsewhere, the federated/ folder can be used a standalone Python package.

All the code is located under the federated folder and organized into multiple sub-modules:

• objectives: Contains implementations of synthetic objective functions. We assume that each client is represented by the corresponding objective. The objective is base on the client's data and can be (1) evaluated or (2) return a (stochastic) gradient at a given point (or parameter vector).
• inference: Contains functions for computing client updates, running SGD (and its variations), doing posterior sampling, estimating moments of a distribution from samples.
• learning: Contains implementations of the learning algorithms (FedAvg, FedPA, and their variations.)
• utils: Contains various utility functions, e.g., for timing code or plotting figures.

A note on implementation. Our minimalistic library is implemented in a functional style. The main learning function is fed_opt (defined in federated.learning.algorithms), which implements the generalized federated optimization (corresponds to Algorithm 1 in our paper, originally proposed by Reddi*, Charles*, et al. (2020)). fed_opt takes client_update_fn and server_update_fn functions as arguments, which are used for computing client and server updates, respectively. FedAvg and FedPA are implemented by providing the corresponding client_update_fn and server_update_fn arguments to fed_opt.

This mini-library does NOT support running experiments on FL benchmark tasks such as EMNIST, CIFAR100, etc. If you would like to run FedPA on these benchmarks, please use our TFF implementation.

@inproceedings{alshedivat2021federated,
  title={Federated Learning via Posterior Averaging: A New Perspective and Practical Algorithms},
  author={Al-Shedivat, Maruan and Gillenwater, Jennifer and Xing, Eric and Rostamizadeh, Afshin},
  booktitle={International Conference on Learning Representations (ICLR)},
  year={2021}
}

Apache 2.0