Robust One Round Federated Learning with Predictive Space Bayesian Inference
Making predictions robust is an important challenge. A separate challenge in
federated learning (FL) is to reduce the number of communication rounds,
particularly since doing so reduces performance in heterogeneous data settings.
To tackle both issues, we take a Bayesian perspective on the problem of
learning a global model. We show how the global predictive posterior can be
approximated using client predictive posteriors. This is unlike other works
which aggregate the local model space posteriors into the global model space
posterior, and are susceptible to high approximation errors due to the
posterior's high dimensional multimodal nature. In contrast, our method
performs the aggregation on the predictive posteriors, which are typically
easier to approximate owing to the low-dimensionality of the output space. We
present an algorithm based on this idea, which performs MCMC sampling at each
client to obtain an estimate of the local posterior, and then aggregates these
in one round to obtain a global ensemble model. Through empirical evaluation on
several classification and regression tasks, we show that despite using one
round of communication, the method is competitive with other FL techniques, and
outperforms them on heterogeneous settings. The code is publicly available at
https://github.com/hasanmohsin/FedPredSpace_1Round.