Swift Convergence: Federated Learning Enhanced with GMMs for Image Classification

In Federated learning (FL), FederatedAveraging(FedAvg) is widely used to compute the weighted mean of local models in the parametric space over time on a central server by exchanging intermediate updates over multiple rounds of communication. However, this approach requires many communication rounds for the central model to learn data generalization. Each local model updates the central model parameters in different directions of the multi-dimensional parameter space. To address this challenge, we propose FedGMMinit: Federated Initialization with Gaussian Mixture Model, which adjusts initial central model gradients by pre-training the model on synthetic data generated from a Gaussian Mixture Model (GMM). For each label in the client’s dataset, a GMM is built. The pre-trained weights are then communicated to the selected clients to initialize FedAvg. To maintain data privacy, only the client’s representation of the Gaussian is passed to the server. Our proposed approach is tested on MNIST digit datasets for image classification. It shows a reduction of 10–15 communication rounds required by the central model to achieve target accuracy for both IID and non-IID distributions. In the scope of the study, we also discovered that clustering clients and training them with global models also contributed to the overall improvement of convergence. We call this clustering method as FedGMMCluster.