On Convergence of Federated Averaging Langevin Dynamics

TL;DR

This paper introduces a federated averaging Langevin algorithm (FA-LD) for uncertainty quantification in distributed settings, providing theoretical guarantees and analyzing the impact of noise, heterogeneity, and update schemes on convergence.

Contribution

The paper generalizes federated Langevin dynamics beyond normal distributions and offers convergence analysis under non-i.i.d. data and various noise and update schemes.

Findings

Convergence guarantees for FA-LD with non-i.i.d. data.

Trade-offs between communication, accuracy, and privacy.

Bias introduced by partial device updates in inactive networks.

Abstract

We propose a federated averaging Langevin algorithm (FA-LD) for uncertainty quantification and mean predictions with distributed clients. In particular, we generalize beyond normal posterior distributions and consider a general class of models. We develop theoretical guarantees for FA-LD for strongly log-concave distributions with non-i.i.d data and study how the injected noise and the stochastic-gradient noise, the heterogeneity of data, and the varying learning rates affect the convergence. Such an analysis sheds light on the optimal choice of local updates to minimize communication costs. Important to our approach is that the communication efficiency does not deteriorate with the injected noise in the Langevin algorithms. In addition, we examine in our FA-LD algorithm both independent and correlated noise used over different clients. We observe there is a trade-off between the pairs among communication, accuracy, and data privacy. As local devices may become inactive in federated networks, we also show convergence results based on different averaging schemes where only partial device updates are available. In such a case, we discover an additional bias that does not decay to zero.