Simple and Scalable Federated Learning with Uncertainty via Improved Variational Online Newton

The standard setup for federated learning (FL) is as a distributed optimization problem where each client learns a local model using its own private data, and these local models are aggregated at a central server. Recent works have also focused on settings where, not just the predictive accuracy but the model and predictive uncertainty estimates are also of interest. Bayesian FL methods have recently emerged as a promising way to achieve this by solving a distributed posterior inference problem. However, computing as well as aggregating the local client posteriors is usually much more expensive (both in terms of local computation as well as the client-server communication) than optimization based FL approaches, such as FedAvg. We present a simple and scalable Bayesian FL method in which, in each round, each client approximates its local posterior using the improved variational online Newton method, which has almost the same cost as simply running an Adam optimizer, making distributed inference mimic distributed optimization. We also present an efficient aggregation method for the client posteriors to learn the global model at the server. Our method achieves improved predictive accuracies as well as better uncertainty estimates as compared to the baselines which include both optimization based FL as well as Bayesian FL methods.