Hybrid-Regularized Magnitude Pruning for Robust Federated Learning under Covariate Shift

TL;DR

This paper introduces a hybrid regularization and pruning approach to improve the robustness and generalization of federated learning models under data heterogeneity, especially in sensitive applications like medical imaging.

Contribution

It proposes a novel federated learning framework combining pruning and regularization to enhance model robustness against client data distribution shifts and introduces a new benchmark dataset, CelebA-Gender.

Findings

Improves model robustness and generalization in heterogeneous federated learning.

Outperforms standard federated learning baselines on multiple datasets.

Demonstrates effectiveness of the proposed method on a new dataset, CelebA-Gender.

Abstract

Federated Learning offers a solution for decentralised model training, addressing the difficulties associated with distributed data and privacy in machine learning. However, the fact of data heterogeneity in federated learning frequently hinders the global model's generalisation, leading to low performance and adaptability to unseen data. This problem is particularly critical for specialised applications such as medical imaging, where both the data and the number of clients are limited. In this paper, we empirically demonstrate that inconsistencies in client-side training distributions substantially degrade the performance of federated learning models across multiple benchmark datasets. We propose a novel FL framework using a combination of pruning and regularisation of clients' training to improve the sparsity, redundancy, and robustness of neural connections, and thereby the resilience to model aggregation. To address a relatively unexplored dimension of data heterogeneity, we further introduce a novel benchmark dataset, CelebA-Gender, specifically designed to control for within-class distributional shifts across clients based on attribute variations, thereby complementing the predominant focus on inter-class imbalance in prior federated learning research. Comprehensive experiments on many datasets like CIFAR-10, MNIST, and the newly introduced CelebA-Gender dataset demonstrate that our method consistently outperforms standard FL baselines, yielding more robust and generalizable models in heterogeneous settings.