Improving Batch Normalization in Federated Learning with Non-IID Features

ICLR 2026 Conference Submission17440 Authors

Keywords: federated learning, non-IID features

Abstract:

Batch normalization (BN) has become a standard practice in deep neural networks, offering significant advantages in convergence speed and training stability. Recent studies indicate that applying BN in non-IID federated learning (FL) scenarios may result in performance degradation. However, these studies either replace it with alternative normalization methods or incur significant communication overhead.
Moreover, they focus primarily on the non-IID labels scenario, neglecting the impact of non-IID features on BN in FL. In this paper, we carefully examine the challenges of BN under feature-shift FL and aim to retain BN in the model, which is crucial for leveraging abundant pretrained backbones. We revisit the FL procedure and show that BN statistics can vary significantly across clients in the presence of non-IID features, resulting in a notable train–test inconsistency. To address this issue, we propose Local–Global Consistency Regularization (GReg), a simple method that achieves alignment with global statistics during local training through an additional KL-based regularization term. Extensive experiments on three natural image benchmarks and a medical image benchmark demonstrate that GReg consistently improves FedAvg and several state-of-the-art FL methods. In addition, GReg exhibits strong generalization to unseen clients, works across diverse CNN and ViT-style architectures, and is suitable for both cross-silo and cross-device FL settings.

Primary Area: other topics in machine learning (i.e., none of the above)

Submission Number: 17440