Swag Technique and Dirichlet Distribution to Address Non-Iid Data in Federated Learning

Abstract

Federated learning deals with the challenge of accessing data from different information sources while preserving their privacy in centralized learning. We can use this paradigm to learn a common global model for multiple clients using model aggregation cycles, without sharing data. Here aggregating the local models is a crucial part of the training. However, the model may experience accuracy and performance loss while aggregating heterogeneous data. We propose a new aggregation method with sampling, FEDSBME, using the Bayesian inference. We sample the local models of the participating clients and build a Bayesian ensemble model to create a powerful aggregation. The sampling of local models is performed using two approaches, SWAG and Dirichlet distribution sampling. Our experimental results prove that our suggested approach can preserve the accuracy and performance of the model when clients’ data are heterogeneous (non-iid) and with deeper neural networks.