Federated Batch Normalization (FedBN)

In FedBN, each client retains its own local batch normalization (BN) parameters—specifically the running mean and running variance—and does not share them with the central server during aggregation. This is a fundamental departure from standard federated averaging, where all model parameters, including BN statistics, are averaged. By keeping these statistics local, the model preserves the client-specific feature distribution normalization, which is crucial when data is statistically heterogeneous.

The core problem FedBN solves is feature shift, a type of data heterogeneity where the distribution of intermediate layer activations differs significantly across clients. When BN statistics from different clients are averaged, the resulting global statistics are a poor representation of any individual client's data, leading to degraded performance. FedBN prevents this by allowing each client's model to maintain normalization tuned to its unique data domain, effectively decoupling the feature alignment problem from the weight optimization problem.

FedBN employs a masked aggregation strategy. During the federated averaging step, only a subset of model parameters is aggregated by the server. Typically, this includes all weight and bias parameters of convolutional and linear layers. The parameters of batch normalization layers (scale γ and shift β, along with the running statistics) are explicitly excluded from aggregation and remain fixed on the client device. This selective process is what enables personalization at the feature level.

FedBN is most effective and commonly applied in architectures that contain batch normalization layers, such as CNNs (e.g., ResNet, VGG) for vision tasks. Its efficacy is tied to the presence of these normalization modules. The technique assumes that the statistical heterogeneity is primarily captured in the feature distributions normalized by these layers. It is less directly applicable to models that use alternative normalization schemes (e.g., LayerNorm common in Transformers) without architectural adaptation.

FedBN maintains the same communication cost per round as standard Federated Averaging (FedAvg) for the parameters it does share. It does not add extra uplink communication overhead because the BN parameters are simply not transmitted. However, it may slightly reduce the downlink cost, as the global model broadcast to clients lacks meaningful BN statistics, which are immediately overwritten by the client's local cached values. The primary efficiency gain is in improved convergence, reducing the total number of communication rounds needed to achieve a target accuracy under Non-IID conditions.

FedBN is often considered a lightweight personalization method. By preserving client-specific normalization, it creates a form of inherent model specialization without requiring multiple model instances or complex meta-learning loops. It serves as a strong baseline and is frequently combined with other personalized federated learning techniques. For example, a global model trained with FedBN can be used as a better initialization for further fine-tuning on each client, as the feature representations are already aligned with the local data.

Non-Independent and Identically Distributed (Non-IID) data is the core statistical challenge in federated learning. It means data samples across different clients are not drawn from the same underlying probability distribution. This violates a standard assumption of centralized machine learning and causes:

• Client Drift: Local models diverge from the global objective.
• Convergence Instability: The global model may not converge to a good solution.
• Performance Degradation: The final model underperforms on clients with unique data distributions. Real-world examples include medical imaging data from different hospitals or user typing patterns on different smartphones.

Statistical heterogeneity is the umbrella term for the variations in data distributions across federated clients. It is the practical manifestation of Non-IID data and encompasses several dimensions:

• Feature Distribution Shift (Covariate Shift): The distribution of input features P(X) varies.
• Label Distribution Shift (Prior Probability Shift): The distribution of output labels P(Y) varies (e.g., label skew).
• Concept Shift: The relationship between features and labels P(Y|X) differs.
• Quantity Skew: The amount of data per client varies dramatically. FedBN directly addresses feature distribution shift by localizing the normalization of activations.

Personalized Federated Learning is a family of techniques designed to produce models tailored to individual clients' local data, rather than a single global model. It is a primary solution to Non-IID data. Key approaches include:

• Local Fine-Tuning: Clients personalize the global model with a few local steps post-aggregation.
• Multi-Task Learning: Framing each client's problem as a related but distinct task.
• Model Mixture: Using a weighted combination of global and local models.
• Meta-Learning: Learning a model initialization that can be quickly adapted. FedBN is a form of architectural personalization, where the batch normalization layers become client-specific.

FedProx is a foundational federated optimization algorithm that mitigates the effects of statistical and system heterogeneity. It modifies the local client training objective by adding a proximal term. This term penalizes the local model for straying too far from the current global model.

• Mechanism: The local loss function becomes L_local(w) + (μ/2) * ||w - w_global||^2.
• Effect: It limits client drift, allowing for variable amounts of local work (partial participation) and more stable convergence under Non-IID conditions. While FedProx regularizes the model weights, FedBN personalizes the normalization statistics—they are complementary techniques.

SCAFFOLD (Stochastic Controlled Averaging) is an algorithm that uses control variates—correction terms stored on the server and each client—to reduce the variance between local updates. It directly tackles the client drift caused by Non-IID data.

• Client Control Variate: Tracks the difference between the client's gradient and the global gradient direction.
• Server Control Variate: The average of all client control variates.
• Update Correction: Local training steps are corrected using these variates, aligning client updates closer to the global optimum. This leads to significantly faster and more stable convergence compared to standard Federated Averaging in heterogeneous settings.

Clustered Federated Learning is an approach that groups clients into clusters based on the similarity of their data distributions. Instead of one global model, a separate model is trained for each cluster.

• Process: Clients are partitioned (explicitly or via learned representations) into groups with more homogeneous data.
• Benefit: It transforms a highly Non-IID problem across all clients into several more IID-like problems within clusters.
• Relation to FedBN: FedBN personalizes at the individual client level. Clustered FL personalizes at the group level. They represent different points on the spectrum from one global model to many fully local models.