PaperlayerView exact revision on arXiv

Source-linked AI summary

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

Xiaoxiao Li, Meirui Jiang, Xiaofei Zhang, Michael Kamp, Qi Dou

arXiv:2102.07623v2cs.LG

TL;DR

Federated learning must handle non-IID client data, including feature shifts that change local feature distributions. The paper proposes FedBN, which keeps batch-normalization parameters local while averaging other model parameters. FedBN improves convergence behavior and model performance across federated datasets, with convergence guarantees in the analyzed regime.

  • Problem

    Federated learning lacks methods focused on feature shift, where client feature distributions differ even when label-related distributions may be similar.

  • Method

    FedBN keeps client batch-normalization parameters local and unsynchronized while updating non-BN layers through FedAvg.

  • Results

    FedBN significantly improves convergence behavior and model performance on non-IID federated datasets, with convergence analysis conducted in the over-parameterized neural-network regime.

  • Takeaways & Limitations

    FedBN can be applied across federated datasets and is presented as a method for mitigating feature shifts without changing the standard non-BN optimization or aggregation scheme.

  • Takeaways & Limitations

    The paper leaves open which types of local-data differences benefit from FedBN and where its limits lie.

Abstract

from arXiv · show

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.

1 INTRODUCTION

Federated learning struggles with heterogeneous client data, especially feature shifts that alter local feature distributions. FedBN addresses this by keeping batch-normalization layers local while averaging non-BN parameters, improving practical performance and convergence relative to existing approaches.

  • Standard FedAvg can suffer performance degradation or diverge when client training samples are non-IID.
  • Feature shift describes client data whose feature distributions differ, including changes caused by imaging machines and protocols despite similar labels.
  • Batch normalization is motivated as a way to mitigate domain shifts and harmonize local feature distributions in federated learning.
  • Averaging both model and BN parameters can yield high generalization error, whereas averaging the model with local BN parameters performs better in the toy example.
  • FedBN keeps client BN layers local and averages non-BN layers using FedAvg, requiring no tunable parameters and minimal additional computation.
  • FedBN demonstrates significant practical improvements over FedAvg and FedProx in extensive experiments.

2 RELATED WORK

Prior federated-learning research has mainly addressed label-distribution heterogeneity and optimization or stability issues. Batch normalization has shown benefits for optimization and domain adaptation, but its role in non-IID federated training remained unexplored.

  • FedAvg variants primarily target stability, client drift, and heterogeneous label distributions.
  • FedProx addresses heterogeneity through partial information aggregation and a proximal term added to FedAvg.
  • Batch normalization has been associated with smoother optimization landscapes, robustness, generalizability, and domain-adaptation benefits.
  • The role of batch normalization in federated learning with non-IID training data remained unexplored.

3 PRELIMINARY

The paper introduces feature shift as a category of client non-IID data and distinguishes it through changes in feature-related distributions. FedAvg provides the baseline aggregation procedure in which clients send locally trained updates to a server for averaging.

  • Feature shift is introduced as a category of client non-IID data distribution.
  • Feature shift includes covariate shift, where Pi(x) varies across clients while Pi(y|x) remains the same, and concept shift, where Pi(x|y) varies while P(y) remains the same.
  • In FedAvg, each client trains a local copy of the global model, sends updates to the server, and receives a model formed by averaging client updates.

4 FEDERATED AVERAGING WITH LOCAL BATCH NORMALIZATION

This section formulates FedBN for feature-shift non-IID federated learning by keeping batch-normalization parameters local while averaging the remaining model parameters. Its simplified convergence analysis shows positive-definite training dynamics and a faster convergence rate for FedBN than FedAvg.

  • Method: FedBN averages local models like FedAvg but excludes batch-normalization parameters from server aggregation, keeping them updated locally.The non-BN layers use FedAvg without changing the optimization or aggregation scheme.
  • Problem setup: The analysis models clients with centered inputs whose covariance matrices may differ, representing feature shift rather than label-distribution heterogeneity.The setup assumes N clients, M examples per client, and a two-layer ReLU network trained by gradient descent.
  • Model formulation: FedBN assigns client-specific batch-normalization scaling parameters, whereas FedAvg uses a shared scaling parameter across clients.The analyzed network is parameterized by first-layer weights, batch-normalization scales, and top-layer parameters.
  • Convergence analysis: The convergence analysis uses neural tangent kernels and decomposes the evolution into magnitude and direction components, with convergence controlled by the smallest evolution-matrix eigenvalue.For α > 1, the analysis states that convergence is dominated by the magnitude component G(t).
  • Convergence analysis: The auxiliary Gram matrices for FedAvg and FedBN are strictly positive definite under the feature-shift assumption, with least eigenvalues denoted µ0 and µ∗0.The analysis uses these eigenvalue bounds to characterize the respective training dynamics.
  • Convergence analysis: For G-dominated convergence, FedBN has a faster convergence rate than FedAvg because its convergence factor is governed by a larger smallest Gram-matrix eigenvalue.The comparison is stated for the simplified analysis setting, including one local update per communication round.

5 EXPERIMENTS

Experiments evaluate FedBN on benchmark and real-world feature-shift datasets, examining convergence, local-update epochs, client data size, heterogeneity, and comparisons with alternative methods. Across these settings, FedBN shows faster or more stable convergence and higher testing accuracy than FedAvg, while improving over alternatives on real-world datasets.

  • Convergence Rate: FedBN achieves faster and more robust convergence than FedAvg on the benchmark digits classification task.Its training loss decreases faster, more smoothly, and with more stable curves during learning.
  • Analysis of Local Updating Epochs: FedBN’s testing accuracy stably exceeds FedAvg’s across local update epochs E = 1, 4, 8, and 16.Both methods show lower accuracy as E increases, but FedBN remains ahead across the tested values.
  • Analysis of Local Dataset Size: FedBN’s improvement margin over SingleSet increases as each client’s local dataset shrinks from 100% to 1% of its original size.Testing accuracy begins to drop significantly when clients receive only 20% of their original data.
  • Effects of Statistical Heterogeneity: FedBN achieves substantially higher testing accuracy than FedAvg across all tested levels of feature-shift heterogeneity.The experiment varies client composition so that more clients from the same dataset correspond to less heterogeneity.
  • Real-World Dataset Experiments: On real-world Office-Caltech10, DomainNet, and ABIDE I datasets, FedBN improves performance over alternative federated methods, including FedProx.It improves Office-Caltech10 mean accuracy by at least 6% over all alternative methods and exceeds alternatives by over 10% on most DomainNet comparisons.

6 CONCLUSION AND DISCUSSION

FedBN keeps local Batch Normalization parameters unsynchronized to mitigate feature shifts while averaging other model parameters. The paper reports convergence and performance benefits, but leaves broader combinations and privacy effects for future study.

  • 6 CONCLUSION AND DISCUSSION: FedBN keeps local Batch Normalization parameters separate from the global model to mitigate feature shifts in non-IID data.The method averages non-BN layers using FedAvg while leaving BN layers local.
  • 6 CONCLUSION AND DISCUSSION: The paper reports convergence guarantees and empirical improvements across federated datasets, including effectiveness when a new client with an unknown domain joins.The new-client setting is evaluated separately in Appendix G.
  • 6 CONCLUSION AND DISCUSSION: FedBN can be combined with different optimization, communication, and aggregation strategies without changing its core local-BN design.The paper identifies theoretical analysis of such combinations as future work.
  • 6 CONCLUSION AND DISCUSSION: The authors identify future work on FedBN's beneficial data differences, its limits, and the privacy-preservation effect of keeping BN parameters invisible to the server.These questions remain to be quantified or analyzed.

B.3 PROOF OF COROLLARY 4.6

The proof compares FedAvg and FedBN through the smallest eigenvalues governing their exponential convergence factors. It concludes that FedBN converges faster in the analyzed setting.

  • B.3 PROOF OF COROLLARY 4.6: FedAvg and FedBN have exponential convergence factors controlled by the smallest eigenvalues of G∞ and G∗∞, respectively.The comparison uses µ0 = λmin(G∞) and µ∗0 = λmin(G∗∞).
  • B.3 PROOF OF COROLLARY 4.6: The proof constructs G∗∞ from client-specific M × M diagonal blocks of G∞ and compares their minimum eigenvalues.The resulting relation is λmin(G∗∞) = min_i∈[N]{λmin(G∞i)} ≥ λmin(G∞).
  • B.3 PROOF OF COROLLARY 4.6: FedBN has a faster convergence rate than FedAvg in the G-dominated convergence setting.The conclusion follows from the comparison of the exponential factors.

D.2 MODEL ARCHITECTURE AND TRAINING DETAILS ON BENCHMARK

The benchmark experiments use heterogeneous digit and domain datasets, standardized neural-network training settings, and multiple analyses of update frequency, client data size, heterogeneity, and competing methods.

  • D.2 MODEL ARCHITECTURE AND TRAINING DETAILS ON BENCHMARK: The experiments analyze convergence rate, local update epochs, local dataset size, statistical heterogeneity, and comparisons with state-of-the-art methods.These analyses are organized in Tables 4–8.
  • D.2 MODEL ARCHITECTURE AND TRAINING DETAILS ON BENCHMARK: Training uses SGD with learning rate 10^-2, cross-entropy loss, batch size 32, and 300 training epochs.FedProx uses µ = 10^-2 selected by grid search from its default settings.
  • D.2 MODEL ARCHITECTURE AND TRAINING DETAILS ON BENCHMARK: Office-Caltech10 and DomainNet experiments use resized 256×256×3 images, batch size 32, 300 epochs, and matched per-dataset sample counts.Office-Caltech10 uses 62 training samples and DomainNet uses 105 training samples per dataset.

D.4 ABIDE DATASET AND TRAINING DETAILS

The ABIDE experiments use multi-site medical data, while additional analyses test FedBN across local update frequencies, centralized training, batch sizes, and client dataset sizes. Results report faster convergence and consistent gains over FedAvg in the examined settings.

  • D.4 ABIDE DATASET AND TRAINING DETAILS: The medical experiment uses resting-state fMRI data from four ABIDE sites treated as clients: UM, NYU, USM, and UCLA.The data are preprocessed ABIDE I ASD and matched-control data.
  • D.4 ABIDE DATASET AND TRAINING DETAILS: FedBN converges faster than FedAvg for local update frequencies E = 1, 4, 8, and 16.This experiment is presented as support for the theoretical convergence analysis.
  • D.4 ABIDE DATASET AND TRAINING DETAILS: FedBN achieves performance comparable to centralized training with data-specific BN layers on the digit test sets.The comparison uses testing accuracy across individual digit datasets.
  • D.4 ABIDE DATASET AND TRAINING DETAILS: Across combinations of local epochs and batch sizes, FedBN consistently outperforms FedAvg and is robust to batch-size selection.The tested settings include E ∈ {1, 4, 16} and B ∈ {10, 50, ∞}.
  • D.4 ABIDE DATASET AND TRAINING DETAILS: FedBN's improvement margin increases as local client dataset sizes decrease, while accuracy drops substantially when clients retain only 20% of original data.The analysis compares FedBN with SingleSet.

E.6 TRAINING ON UNEQUAL DATASET SIZE

The benchmark controls for unequal client dataset sizes by equalizing samples, then evaluates settings that retain unequal data through repeated sampling. FedBN continues to outperform other methods in these unequal-size settings.

  • The benchmark truncates each dataset to the smallest client sample size to control non-related factors and isolate local BN effects.This preprocessing equalizes the number of data points across clients.
  • With unequal original dataset sizes, clients repeat sampling so they use the same batch size and local iterations per epoch.
  • FedBN still consistently outperforms other methods when clients use 10% or the full original dataset sizes.These settings are reported in Tables 15 and 16.

F SYNTHETIC DATA EXPERIMENT

The synthetic experiment uses paired Gaussian classification distributions with different covariance structures to study FedBN under controlled feature variation. Both methods reach perfect accuracy, while FedBN converges faster than FedAvg.

  • The experiment samples binary-class data from two Gaussian pairs whose covariance matrices differ in their off-diagonal structure.Σ1 is an identity diagonal matrix, whereas Σ2 has non-zero off-diagonal entries.
  • 100% accuracy is achieved by both FedAvg and FedBN on the synthetic binary-classification task.
  • FedBN obtains significantly faster convergence than FedAvg in the training-loss comparison over steps.The comparison is shown in Figure 9.
  • Client 1 uses Diagonal Gaussian data, while client 2 uses a combination of Diagonal Gaussian and Full Gaussian data.

G TRANSFER LEARNING AND TESTING ON UNKNOWN DOMAIN CLIENT

The paper examines transferring FedBN to new domain clients during training and testing on clients with unknown statistics. These procedures retain shared non-BN parameters while adapting batch-normalization statistics locally.

  • The analysis covers both adding a new unknown-domain client during training and testing an unknown-domain client afterward.
  • A new domain client receives the global model’s non-BN parameters and learns its own local BN parameters from its mean and variance statistics.
  • Testing on an unseen client requires access to local BN parameters and computation of mean and variance on that client’s data.The client uses averaged trainable BN parameters learned at existing FL clients.
  • Table 17 reports results for generalizing the global model to unseen-domain clients.
Loading 2102.07623v2…