[Paper Review] FedBN: Federated Learning on Non-IID Features via Local Batch Normalization

Xiaoxiao Li, Meirui Jiang+3 more|arXiv (Cornell University)|Feb 15, 2021

Privacy-Preserving Technologies in Data40 references77 citations

TL;DR

FedBN keeps Batch Norm parameters local in federated learning to address feature shift non-IID data, improving convergence and accuracy over FedAvg and FedProx across benchmarks and real-world datasets.

ABSTRACT

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.

Motivation & Objective

• Motivate and address feature shift non-IID in federated learning where local feature distributions differ across clients.
• Propose a lightweight method to mitigate feature shift by keeping BN parameters local (not averaged) while using FedAvg for non-BN layers.
• Provide theoretical convergence analysis showing faster convergence under feature shift compared to FedAvg.
• Empirically validate FedBN on benchmark and real-world datasets with clear improvements over FedAvg and FedProx.

Proposed method

• Introduce FedBN, which performs local updates and aggregates non-BN layers, while BN parameters remain strictly local and are not communicated.
• Model formulation with two-layer network and batch normalization as part of the architecture, and definition of the FedBN objective that omits BN parameter averaging.
• Convergence analysis using neural tangent kernel (NTK) framework showing FedBN’s convergence governed by the minimum eigenvalue of an auxiliary Gram matrix, with FedBN achieving a faster rate than FedAvg in the feature-shift setting.
• Experimental protocol across benchmark and real-world datasets comparing FedBN to FedAvg, FedProx, and SingleSet, including convergence behavior and accuracy under varying data regime and heterogeneity.

Experimental results

Research questions

• RQ1 Does feature shift in federated learning degrade FedAvg, and can local BN mitigate this degradation?
• RQ2 Can FedBN improve convergence speed and final accuracy compared to FedAvg and FedProx under feature-shift non-IID data?
• RQ3 How does FedBN perform on benchmark domains and real-world datasets with varying degrees of heterogeneity?
• RQ4 Is FedBN robust to different local update frequencies and client data sizes?

Key findings

• FedBN consistently outperforms FedAvg and FedProx on benchmark and real-world datasets under feature-shift non-IID settings.
• FedBN shows faster convergence and more stable loss curves than FedAvg in experiments.
• FedBN yields larger accuracy gains on datasets with more pronounced feature shift (e.g., SVHN) and maintains lower variance across runs.
• In real-world tasks (Office-Caltech10, DomainNet, ABIDE), FedBN achieves higher mean accuracy than alternatives across multiple clients/datasets.
• Theoretical analysis indicates FedBN has a faster convergence rate than FedAvg in the G-dominated regime under feature shift.