This is the implementation of our paper (accepted by ICML 2025 poster):

Hybrid Batch Normalisation: Resolving the Dilemma of Batch Normalisation in Federated Learning

(a) Original Clusters: Cluster 1 (blue squares): 300 samples centered at [5, 5] with a tight standard deviation of 0.4; Cluster 2 (coral circles): 200 samples centered at [-1, -1] with a wider standard deviation of 1.2.

(b) Local Normalistion: Each cluster is normalised independently using its own local statistics.

(c) Global Normalistion: Both clusters are normalised using the global statistics.

(d) Hybrid Normalistion: Each cluster is normalised by a specific mixture of local and global statistics. Cluster 1: 70% reliance on local statistics, 30% on global statistics; Cluster 2: 30% reliance on local statistics, 70% on global statistics

Hybrid normalisation combines global statistics with local statistics, which can standardise the size of two clusters while maintaining the global structure.

Hybrid Batch Normalisation

The specific code implementation can be found in ./FedBaseline/models/FedHBN/MyBNtool.py.

File Structure （Updating）

├── FedBaseline/
│   ├── models/
│   │   ├── FedAvg/          implementation of client and server operations for FedAvg
│   │   ├── FBN/             implementation of client and server operations for FedAvg+FBN
│   │   ├── FedFN/           implementation of client and server operations for FedAvg+FedFN
│   │   ├── FixBN/           implementation of client and server operations for FedAvg+FixBN
│   │   ├── FedHBN/          implementation of client and server operations for FedAvg+HBN
│   │   ├── Net.py           storage network architecture.
│   │   └── Test.py          used for testing.
│   ├── options/
│   │   └── options.py       experimental parameter settings.
│   └── sampling/
│       ├── dataloader.py    load data.
│       └── sampling.py      divide the data to the clients.
└── XXX_main.py              core driver program for specific algorithms.

FedAvg:Communication-Efficient Learning of Deep Networks from Decentralized Data AISTATS 2017

FBN:Overcoming the Challenges of Batch Normalization in Federated Learning arXiv

FedFN: FedFN: Feature Normalization for Alleviating Data Heterogeneity Problem in Federated Learning FL@FM-NeurIPS’23

FixBN: Making Batch Normalization Great in Federated Deep Learning FL@FM-NeurIPS’23

Runing

Cifar10 with B=4, β=0.6 by Simple-CNN

# FedAvg+HBN
python FedHBN_main.py --model 'cnn-hbn' --lr 0.01 --local_batch 4 --dataset 'cifar10' --num_classes 10 --iid False --dirichlet_alpha 0.6
# FedAvg+BN
python FedAvg_main.py --model 'cnn-bn' --lr 0.01 --local_batch 4 --dataset 'cifar10' --num_classes 10 --iid False --dirichlet_alpha 0.6
# FedAvg+GN
python FedAvg_main.py --model 'cnn-gn' --lr 0.005 --local_batch 4 --dataset 'cifar10' --num_classes 10 --iid False --dirichlet_alpha 0.6
# FedAvg+LN
python FedAvg_main.py --model 'cnn-ln' --lr 0.005 --local_batch 4 --dataset 'cifar10' --num_classes 10 --iid False --dirichlet_alpha 0.6
# FedAvg+FedFN
python FedFN_main.py --model 'cnn-fn' --lr 0.005 --local_batch 4 --dataset 'cifar10' --num_classes 10 --iid False --dirichlet_alpha 0.6
# FedAvg+FixBN
python FixBN_main.py --model 'cnn-bn' --lr 0.002 --local_batch 4 --dataset 'cifar10' --num_classes 10 --iid False --dirichlet_alpha 0.6
# FedAvg+FBN
python FBN_main.py --model 'cnn-fbn' --lr 0.01 --local_batch 4 --dataset 'cifar10' --num_classes 10 --iid False --dirichlet_alpha 0.6

Note: Do not run all methods with the same initial learning rate, as different normalisation methods may have different optimal learning rates. Some tuning experience can be found in the appendix of the paper.

Citation

@inproceedings{
chen2025hybrid,
title={Hybrid Batch Normalisation: Resolving the Dilemma of Batch Normalisation in Federated Learning},
author={Hongyao Chen and Tianyang Xu and Xiaojun Wu and Josef Kittler},
booktitle={Forty-second International Conference on Machine Learning},
year={2025},
url={https://openreview.net/forum?id=zV5pkTMHPP}
}