Bayesian Federated Learning with Hamiltonian Monte Carlo: Algorithm and Theory

TL;DR

This paper presents FA-HMC, a new Bayesian federated learning algorithm using Hamiltonian Monte Carlo, with proven convergence guarantees and superior empirical performance over existing methods.

Contribution

Introduces FA-HMC, a novel federated learning algorithm combining Hamiltonian Monte Carlo with rigorous convergence analysis and empirical validation.

Findings

FA-HMC outperforms FA-LD in experiments.

Convergence guarantees established under strong convexity.

Analysis of effects of dimension, noise, and communication frequency.

Abstract

This work introduces a novel and efficient Bayesian federated learning algorithm, namely, the Federated Averaging stochastic Hamiltonian Monte Carlo (FA-HMC), for parameter estimation and uncertainty quantification. We establish rigorous convergence guarantees of FA-HMC on non-iid distributed data sets, under the strong convexity and Hessian smoothness assumptions. Our analysis investigates the effects of parameter space dimension, noise on gradients and momentum, and the frequency of communication (between the central node and local nodes) on the convergence and communication costs of FA-HMC. Beyond that, we establish the tightness of our analysis by showing that the convergence rate cannot be improved even for continuous FA-HMC process. Moreover, extensive empirical studies demonstrate that FA-HMC outperforms the existing Federated Averaging-Langevin Monte Carlo (FA-LD) algorithm.