Applying client-specific Batch Normalization to unseen federated clients

Determine how to apply client-specific Batch Normalization (BN) layer parameters and running statistics, learned locally under the FedBN strategy during federated learning across heterogeneous brain MRI datasets with different diseases and modality combinations, to unseen target clients during inference.

Background

The paper trains a single federated segmentation model across multiple brain MRI datasets that differ in diseases and available MRI modalities. To handle data heterogeneity, the authors consider keeping BN parameters and statistics client-specific (as in FedBN), which improves performance on source clients.

However, when deploying the model to an unseen client (a target site not participating in training), it is unclear how to utilize these source-specific BN parameters and statistics. One option is to average BN parameters and estimate running statistics using target data, but this requires access to target data, which may be impractical or undesirable. This unresolved issue motivates exploring normalization-free methods and highlights a need for principled strategies for BN at inference on unseen clients.

References

However, it is not clear how to apply the source-specific BN layer parameters/statistics to unseen (target) clients during inference.

Feasibility of Federated Learning from Client Databases with Different Brain Diseases and MRI Modalities  (2406.11636 - Wagner et al., 2024) in Section 2.2 (Training a Unified Model for Multiple Segmentation Tasks) — Feature Normalization for Heterogeneous Clients