Multi-Stage Robust Federated Learning: Addressing Label Noise under Data Heterogeneity and Imbalance

Federated learning (FL) enables collaborative model training while preserving data privacy, but the presence of noisy labels in local datasets remains a significant challenge, particularly under heterogeneous noise conditions and class imbalance. In this work, we introduce a novel Multi-stage Robust Federated Learning (MRFL) framework to address these issues. In the warm-up noise detection stage, MRFL computes per-class average losses on each client and employs a Gaussian mixture model to accurately identify clients with substantial label noise. In the subsequent noise-robust training stage, a robust loss function and noise solver are designed to distinguish clean from noisy samples, while semi-supervised learning is used to recover valuable information from tail classes. Moreover, a robust weighted aggregation strategy is adopted to mitigate the adverse effects of noisy clients. Extensive experiments on CIFAR-10/100-LT and ICH datasets demonstrate that MRFL outperforms state-of-the-art methods in federated noisy label learning scenarios characterized by data heterogeneity and imbalance.