FedBN: Federated Learning on Non-IID Features via Local Batch Normalization

By Xiaoxiao Li, Meirui Jiang, Xiaofei Zhang, Michael Kamp, Qi-Yan Dou. 2021

1,240 citations recorded by the source corpus.

The emerging paradigm of federated learning (FL) strives to enable collaborative training of deep models on the network edge without centrally aggregating raw data and hence improving data privacy. In most cases, the assumption of independent and identically distributed samples across local clients does not hold for federated learning setups. Under this setting, neural network training performance may vary significantly according to the data distribution and even hurt training convergence. Most of the previous work has focused on a difference in the distribution of labels or client shifts. Unlike those settings, we address an important problem of FL, e.g., different scanners/sensors in medical imaging, different scenery distribution in autonomous driving (highway vs. city), where local clients store examples with different distributions compared to other clients, which we denote as feature shift non-iid. In this work, we propose an effective method that uses local batch normalization to alleviate the feature shift before averaging models. The resulting scheme, called FedBN, outperforms both classical FedAvg, as well as the state-of-the-art for non-iid data (FedProx) on our extensive experiments. These empirical results are supported by a convergence analysis that shows in a simplified setting that FedBN has a faster convergence rate than FedAvg. Code is available at https://github.com/med-air/FedBN.

Subfields: Artificial Intelligence

How researchers see this paper

The researchers whose trusted networks rank it highest, strongest advocate first. Every researcher gets a different order.

In Artificial Intelligence

1. Bo Du
2. Mang Ye
3. Dacheng Tao
4. Guodong Long
5. Jing Jiang
6. Tianyi Zhou
7. Bingsheng He
8. Qinbin Li

Citations in the corpus

• Federated Class-Incremental Learning (2022, background): “Federated Learning (FL) is a decentralized learning framework that can train a global model by aggregating local model parameters [24, 44, 51, 53].”
• A Comprehensive Survey of Privacy-preserving Federated Learning (2021, background): “The focuses of FL research can be roughly divided into three aspects: (1) improving the efficiency and effectiveness of FL [24, 68, 87, 100, 120, 211, 213], (2) improving the security of FL to attacks that aim to undermine the integrity of FL models and degrade the model performance [9, 13, 185, 198”
• Model-Contrastive Federated Learning (2021, background): “We also notice that there are other related contemporary work [1, 30, 43] when preparing this paper.”
• Evaluating Gradient Inversion Attacks and Defenses in Federated Learning (2021, methodology): “In each training step, the client normalizes its private batch x∗ using the batch’s mean mean(x∗) and variance var(x∗), keeps the running estimate of mean and variance locally as in [Li et al., 2021], and shares the gradient.”
• Federated Learning in Edge Computing: A Systematic Survey (2022, methodology): “Xiaoxiao Li et al. [191] proposed an effective method, FedBN, that uses local BN to mitigate feature shift before model averaging, and instead, it keeps all BN parameters strictly local, which is the state-of-the-art solution to mitigate the statistical heterogeneity challenge in FL.”
• Federated Learning from Pre-Trained Models: A Contrastive Learning Approach (2022, methodology): “[3] Xiaoxiao Li, Meirui Jiang, Xiaofei Zhang, Michael Kamp, and Qi Dou.”
• Heterogeneous Federated Learning: State-of-the-art and Research Challenges (2023, background): “FedBN [126] solves the problem of feature skew in statistical heterogeneity by adding a batch normalization layer to the local model.”
• Federated Multi-Task Learning under a Mixture of Distributions (2021, methodology): “A possible way to ensure that Assumption 2 holds is to use the batch normalization technique from [40] to account for feature shift.”
• Learn from Others and Be Yourself in Heterogeneous Federated Learning (2022, neutral): “A myriad of methods [43,46,73,77] incorporate extra proximal terms to handle the data in label distribution skew (prior probability shift) [30], neglecting the fact that there exists domain shift (same label, different features) [60, 64, 66].”
• Federated Learning on Non-IID Data Silos: An Experimental Study (2021, background): “When preparing this paper, there are other contemporary works [2], [40], [48], [73] on federated learning under non-IID setting.”