Enhancing Federated Learning Robustness in Non-IID Data Environments via MMD-Based Distribution Alignment

Federated learning(FL), due to its distributed nature, is highly susceptible to malicious attacks. Although various Byzantine-robust FL methods exist, they often fail to maintain robustness in practical scenarios due to the non-independent and identically distributed (Non-IID) nature of client data. Moreover, existing FL methods often suffer from weight divergence caused by heterogeneous data distributions across clients. To address these issues, we propose a novel federated learning framework that aligns local data distributions across different clients to enhance robustness for Non-IID data in adversarial environments. It contains a feature transformation layer that incorporates Maximum Mean Discrepancy (MMD) as a regularization term to avoid weight divergence through aligning local and global data distributions without sharing raw data. Our approach dynamically updates the statistical information of both local and global data, including the mean and variance, ensuring that local models are closely aligned with the global model throughout training. Experimental results on MNIST and CIFAR-10 datasets demonstrate that our proposed framework significantly improves robustness both in the absence of attacks and against untargeted attacks such as sign-flipping and additive noise.