Re-Sampling Calibrated SNN Loss: A Robust Approach to Non-IID Data in Federated Learning

Abstract

Federated Learning (FL) represents a significant advancement in decentralised machine learning, offering a solution to the privacy concerns associated with traditional centralised approaches. However, a critical limitation of FL arises in the presence of Non-Independent and Identically Distributed (non-IID) data, which is common in real-world scenarios. Traditional FL algorithms, such as Federated Averaging (FedAvg), tend to underperform when faced with data heterogeneity across participating clients. To address this challenge, we propose CalibSNN, a method that combines calibration re-sampling with Soft Nearest Neighbour (SNN) loss to mitigate the bias and variance introduced by uneven data distributions. Calibration aligns local data distributions with global statistics, while SNN loss improves feature representations across heterogeneous clients. Through extensive experiments on diverse datasets and non-IID conditions, we demonstrate that CalibSNN significantly outperforms state-of-the-art baselines, offering a robust solution to the challenges of non-IID data in FL.