Bayesian Federated Learning with Stochastic Variational Inference
Mehreen Tahir, Feras M. Awaysheh, Sadi A. Alawadi, Muhammad Abdul Rehman Liaquat Ali · 2024
Federated Learning (FL) faces significant chal-lenges, such as handling non-IID (Non-Independent and Identically Distributed) data and efficiently aggregating distributed models, which can lead to slower convergence and reduced model accuracy. This paper proposes a novel framework, Bayesian Federated Learning with Stochastic Variational Inference (BayFL-SVI), to address these issues. Stochastic Variational Inference (SVI) is a scalable approximation method for Bayesian inference that optimizes the Evidence Lower Bound (ELBO) using mini-batches of data through stochastic gradient descent. By computing the ELBO for each client update, our approach quantifies the significance of these updates, effectively managing the heterogeneity of non-IID data and improving the aggregation process. This approach results in a more accurate and robust integration of client contributions, enhancing convergence rates and overall model performance. We provide theoretical analysis with convergence guarantees. Our empirical results demonstrate significant improvements in convergence rates and model accuracy, establishing a solid foundation for future studies.