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

TL;DR

This paper introduces Hybrid Batch Normalisation (HBN), a novel normalisation method designed to improve federated learning by effectively handling non-i.i.d. data and small batch sizes through a hybrid approach to statistical parameter updates.

Contribution

The paper proposes HBN, a new normalisation technique that separates statistical and learnable parameter updates, enhancing federated learning performance with global statistics and adaptive mixing.

Findings

HBN improves federated learning accuracy with small batch sizes.

HBN effectively handles heterogeneous, non-i.i.d. data distributions.

HBN outperforms existing normalisation methods in federated settings.

Abstract

Batch Normalisation (BN) is widely used in conventional deep neural network training to harmonise the input-output distributions for each batch of data. However, federated learning, a distributed learning paradigm, faces the challenge of dealing with non-independent and identically distributed data among the client nodes. Due to the lack of a coherent methodology for updating BN statistical parameters, standard BN degrades the federated learning performance. To this end, it is urgent to explore an alternative normalisation solution for federated learning. In this work, we resolve the dilemma of the BN layer in federated learning by developing a customised normalisation approach, Hybrid Batch Normalisation (HBN). HBN separates the update of statistical parameters (i.e. , means and variances used for evaluation) from that of learnable parameters (i.e. , parameters that require gradient updates), obtaining unbiased estimates of global statistical parameters in distributed scenarios. In contrast with the existing solutions, we emphasise the supportive power of global statistics for federated learning. The HBN layer introduces a learnable hybrid distribution factor, allowing each computing node to adaptively mix the statistical parameters of the current batch with the global statistics. Our HBN can serve as a powerful plugin to advance federated learning performance. It reflects promising merits across a wide range of federated learning settings, especially for small batch sizes and heterogeneous data.