FedBN-Prox: Federated BN & Proximal Regularization

• The paper introduces FedBN-P, partitioning model parameters into shared non-BN and client-specific BN to address statistical heterogeneity and client drift.
• It applies proximal regularization only on shared parameters while keeping BN parameters local, stabilizing training and improving generalization across diverse domains.
• Empirical results show FedBN-P achieves faster convergence and higher accuracy compared to FedAvg, FedBN, and FedProx, especially under non-IID and Byzantine conditions.

Searching arXiv for FedBN-Prox / FedBN-P and foundational related work. FedBN-Prox, commonly abbreviated FedBN-P, is a federated optimization scheme that combines Federated Batch Normalization with proximal regularization by partitioning model parameters into shared non-BN parameters and client-local BN parameters, regularizing only the shared subset, and excluding BN affine parameters and running statistics from server aggregation. In the OptiGradTrust framework, FedBN-P is used to address statistical heterogeneity and convergence instability, while trust-weighted aggregation addresses Byzantine behavior; conceptually, it extends the local-BN principle of FedBN to a setting in which client drift is additionally controlled by a proximal term on the shared parameters (Karami et al., 31 Jul 2025, Li et al., 2021).

1. Origins and problem setting

FedBN-P is motivated by two coupled pathologies in federated learning. First, statistical heterogeneity, especially non-IID data, induces client drift: local models trained on skewed distributions move in disparate directions, which can slow or destabilize convergence when updates are naively averaged. Second, Batch Normalization layers encode domain-specific feature statistics; aggregating those statistics across divergent client domains, such as scanners or acquisition protocols, can harm generalization. FedBN-P combines the two established responses to these problems: the FedBN insight that BN layers should remain local, and the FedProx insight that a proximal regularizer can discourage excessive deviation from the current global model (Karami et al., 31 Jul 2025).

The FedBN foundation is feature-shift non-IID federated learning. In that setting, heterogeneity is not restricted to label proportions or client sampling, but instead reflects changes in feature distributions across clients, as in different scanners or sensors in medical imaging or different scenery styles in autonomous driving. FedBN addresses that regime by keeping BN strictly client-local and aggregating only non-BN parameters, which substantially improves training behavior under feature shift and, in the simplified analysis of the FedBN paper, yields a faster convergence rate than FedAvg (Li et al., 2021).

FedBN-P preserves that locality principle and adds proximal regularization only on the shared part of the model. This design implies a specific division of labor. Local BN preserves institution-specific normalization, the proximal term reduces client drift on shared parameters, and, in OptiGradTrust, trust weighting is layered on top to down-weight malicious updates. A plausible implication is that FedBN-P is best understood not as a replacement for robust aggregation, but as the heterogeneity-control component in a larger robust FL pipeline.

2. Formal formulation

FedBN-P partitions model parameters into shared non-BN parameters $w_g$ and client-local BN parameters $w_{bn,i}$ for client $i$, where $w_{bn,i}$ includes BN affine parameters $\gamma, \beta$ and running mean/variance buffers. For client $i$ with data $D_i$, the local empirical risk is

$$F_i(w_g, w_{bn,i}) = \frac{1}{|D_i|} \sum_{(x,y)\in D_i} \ell(f(x; w_g, w_{bn,i}), y).$$

At communication round $t$, with current global shared parameters $w_g^{(t)}$, FedBN-P solves the local proximal objective

$w_{bn,i}$0

The proximal penalty is applied only to the shared parameters $w_{bn,i}$1; BN parameters remain local and are not regularized toward any server reference. In OptiGradTrust, $w_{bn,i}$2 controls how strongly clients stay near the global model, and the reported choice is $w_{bn,i}$3, selected via grid search on a validation split (Karami et al., 31 Jul 2025).

With local step index $w_{bn,i}$4, step size $w_{bn,i}$5, and Adam as optimizer, the instantaneous gradient of the augmented objective over the shared parameters is

$w_{bn,i}$6

The Adam moments are

$w_{bn,i}$7

$w_{bn,i}$8

$w_{bn,i}$9

The shared-parameter update is

$i$0

while BN parameters are updated without any proximal term:

$i$1

In the reported implementation, Adam uses learning rate $i$2, weight decay $i$3, and a cosine decay schedule. BN affine parameters and running statistics are updated on local data only (Karami et al., 31 Jul 2025).

3. One-round protocol and server-side aggregation

FedBN-P uses a communication pattern in which the server broadcasts only shared parameters, and each client preserves its own BN state across rounds. In round $i$4, the server selects participating clients $i$5 and broadcasts $i$6. Each client initializes its shared parameters from that broadcast, retains its local $i$7 from the previous round, performs local training for $i$8 epochs, and returns only the final shared parameters $i$9 or the corresponding shared-parameter delta

$w_{bn,i}$0

BN parameters and statistics are not uploaded. BN affine parameters $w_{bn,i}$1 and running statistics (mean/variance) are excluded from aggregation and remain local to each client throughout training (Karami et al., 31 Jul 2025).

In OptiGradTrust, server aggregation is trust-weighted. A trust weight $w_{bn,i}$2 is computed for each client from a six-dimensional gradient fingerprint and a dual-attention + RL policy module. With normalized weights satisfying $w_{bn,i}$3, the server update on the shared parameters is

$w_{bn,i}$4

which is equivalently a weighted average of client models for the shared parameters. The BN subset is neither aggregated nor broadcast from the server; clients keep and evolve their own BN layers (Karami et al., 31 Jul 2025).

This protocol encodes the central structural claim of FedBN-P: only the non-BN subset participates in both proximal regularization and inter-client consensus. A common misconception is to interpret FedBN-P as a method that uses BN during local training but still averages BN statistics globally. The described algorithm does not do that. The BN affine parameters and running statistics are never uploaded or averaged.

4. Relation to neighboring methods and theoretical status

FedBN-P is most naturally situated relative to FedAvg, FedProx, and FedBN. FedAvg provides no explicit heterogeneity control and aggregates BN; under non-IID conditions, this can degrade accuracy and produce oscillatory convergence. FedProx adds a proximal term but still aggregates BN, so it improves stability yet can remain vulnerable when BN statistics differ widely across clients. FedBN keeps BN local but uses plain averaging for shared parameters; it boosts accuracy under non-IID feature shift but can converge more slowly or be less stable on highly skewed data. FedBN-P jointly applies BN localization and proximal regularization on the shared part of the model, then, in OptiGradTrust, layers trust-weighted aggregation on top. The resulting three-part synergy is local BN for distribution shift, proximal regularization for convergence stability, and trust weights for Byzantine robustness (Karami et al., 31 Jul 2025).

The broader comparison class includes SCAFFOLD, FedNova, and FedAdam/FedDWA, which are alternative heterogeneity handlers with different control or normalization ideas but no BN localization and typically no adversarial robustness. FedBN-P therefore differs from purely optimization-based or purely aggregation-based defenses. Its heterogeneity mechanism is not only a correction term or a server optimizer; it is an architectural partition of the parameter space.

The theoretical status of FedBN-P is limited. The OptiGradTrust paper does not introduce new theorems for FedBN-P. Instead, it appeals to two established intuitions: proximal regularization mitigates client drift under heterogeneity by penalizing large deviations from the global model, and keeping BN local preserves per-client feature normalization that reflects local data statistics, reducing cross-client distribution shift induced by averaging normalization layers. This suggests that the method’s justification is compositional rather than theorem-driven: it inherits plausibility from FedProx-style stabilization and FedBN-style feature-shift handling, but no dedicated convergence proof for their combination is presented (Karami et al., 31 Jul 2025, Li et al., 2021).

FedBN-P should also be distinguished from other proximal federated methods that are unrelated to BN localization. For example, the federated binary matrix factorization method Felb uses proximal operators for relaxed factor matrices and server-side proximal aggregation, but it does not use the acronyms “FedBN-Prox” or “FedBN-P” (Dalleiger et al., 2024).

5. Empirical evidence in OptiGradTrust

The OptiGradTrust evaluation reports FedBN-P on three model-dataset pairs: a 3-layer CNN on MNIST, ResNet-18 on CIFAR-10, and ResNet-18 at $w_{bn,i}$5 on a synthetic Alzheimer’s MRI dataset with 6,983 scans and 4 classes. The optimizer is Adam with learning rate $w_{bn,i}$6 and cosine decay, weight decay $w_{bn,i}$7, batch size 64 except 16 for MRI, 25–30 communication rounds, 5–8 local epochs, 10 clients, and 100% participation. Non-IID partitions include Dirichlet $w_{bn,i}$8, label skew at 70% and 90%, and quantity skew via log-normal (Karami et al., 31 Jul 2025).

In optimizer comparison under standard settings, final test accuracy (%) averaged over seeds is reported as follows. On MNIST, FedBN-P reaches 99.32, compared with FedBN 99.35, FedProx 99.25, and FedAvg 99.20. On CIFAR-10, FedBN-P attains 83.67, compared with FedBN 83.10, FedProx 82.80, and FedAvg 82.50. On Alzheimer’s MRI, FedBN-P reaches 96.10, compared with FedBN 96.25, FedProx 95.47, and FedAvg 94.68. For convergence on Alzheimer’s MRI, FedBN-P reaches target accuracy in 26 rounds versus 30 for FedBN and 24 for FedProx, which is described as a Pareto-like balance between speed and accuracy. The accompanying ablation states that FedBN-P trails FedBN by only 0.15 percentage points in IID accuracy on Alzheimer’s MRI while converging faster, and outperforms FedProx by 5.8 percentage points under heterogeneity (Karami et al., 31 Jul 2025).

Under IID attacks with 30% Byzantine clients, the reported averages across attack types are approximately 99.21 on MNIST, 82.44 on CIFAR-10, and 96.61 on Alzheimer’s MRI. Under non-IID attacks, the reported averages are 98.12, 81.36, and 95.00 for Dirichlet $w_{bn,i}$9; 97.42, 79.42, and 93.80 for Dirichlet $\gamma, \beta$0; and 97.21, 79.20, and 93.26 under label-skew 90%, respectively for MNIST, CIFAR-10, and Alzheimer’s MRI (Karami et al., 31 Jul 2025).

When compared with FLGuard under Dirichlet $\gamma, \beta$1, FedBN-P within OptiGradTrust achieves 98.12 versus 96.53 on MNIST, 81.36 versus 79.62 on CIFAR-10, and 95.00 versus 93.41 on Alzheimer’s MRI, corresponding to gains of +1.59, +1.74, and +1.59 percentage points. Under IID, the corresponding figures are 99.21 versus 98.98, 82.44 versus 81.72, and 96.61 versus 95.89. The paper summarizes the combined result as state-of-the-art performance up to +1.6 percentage points over FLGuard under non-IID conditions (Karami et al., 31 Jul 2025).

These results are used to support three robustness claims internal to the framework: local BN prevents degradation from mixing incompatible normalization statistics, the proximal term stabilizes local training trajectories and reduces oscillation, and trust weighting down-weights malicious updates. A plausible implication is that the empirical gains depend on complementarity rather than redundancy among these components.

6. Practical use, edge cases, and limitations

The practical guidance reported for FedBN-P centers on the proximal coefficient $\gamma, \beta$2, the optimizer schedule, and BN behavior. For $\gamma, \beta$3, the recommendation is to start with $\gamma, \beta$4, with OptiGradTrust using $\gamma, \beta$5. The guidance is to increase $\gamma, \beta$6 as non-IID severity grows, such as smaller Dirichlet $\gamma, \beta$7 or higher label skew, and to reduce $\gamma, \beta$8 if convergence slows excessively or if local steps become too conservative. Lower learning rates pair well with higher $\gamma, \beta$9, and a cosine decay helps settle to a robust consensus. Weight decay is kept modest at $i$0 in the experiments (Karami et al., 31 Jul 2025).

The BN-specific prescription is strict locality. All BN affine parameters and running statistics remain local. Typical BN momentum of approximately 0.1 and $i$1 of approximately $i$2 are used unless there are strong reasons to tune them. Very small batches can destabilize BN; the reported guidance is to consider ghost BN or, if necessary, GroupNorm or LayerNorm for extremely small local batches. With extreme skew, such as Dirichlet $i$3, the recommendation is to increase $i$4 and consider slightly more local epochs $i$5 with caution while monitoring validation stability (Karami et al., 31 Jul 2025).

The behavior in edge cases is explicitly specified. If a model contains some BN layers and some layers without BN, only the non-BN subset is aggregated with proximal regularization, while all BN layers remain per-client. If there are no BN layers, FedBN-P reduces to FedProx plus trust-weighted aggregation. In that case, the stability benefit of proximal regularization and the attack robustness of trust weighting remain, but the BN component is inert. Client dropout is handled by aggregating over the active set $i$6 each round and normalizing trust weights over active clients (Karami et al., 31 Jul 2025).

The principal limitations are likewise explicit. No new convergence theorem is provided for FedBN-P; its arguments are empirical and by analogy to FedProx plus FedBN. BN dependence on batch size makes extremely small local batches problematic. Hyperparameter sensitivity remains, because $i$7 interacts with learning rate and the number of local epochs, so a brief validation sweep is recommended per domain. When there are no BN layers, the benefit over FedProx alone comes only from trust weighting if that mechanism is enabled. Finally, trust-weight computation, including components such as Shapley sampling, adds server cost, although that overhead is orthogonal to FedBN-P’s local computation (Karami et al., 31 Jul 2025).