Population Normalization for Federated Learning

Batch normalization (BN) is a standard technique for training deep neural networks. However, its effectiveness in Federated Learning (FL) applications, which typically involve heterogeneous training data and resource-limited clients, is compromised for two reasons. One is the population statistics (i.e., mean and variance) of the heterogeneous datasets differ greatly, which leads to inconsistent BN layers among the client models and finally causes these models to drift further away from each other. The other is estimating statistics from a mini-batch can be inaccurate since the batch size has to be small in resource-limited clients. In this paper, we propose a technique named Population Normalization for FL, in which the statistics are learned as trainable parameters during training instead of calculated from mini-batches directly as BN does. Thus, our normalization layers are homogeneous among the clients and the impact of small batch size is eliminated as the model can be well-trained even when the batch size equals to one. To make our method more flexible in real applications, we investigate the role of the stochastic uncertainty existing in the statistic estimation of BN and show that when a large batch size is available we can inject simple artificial noise into our method to mimic this stochastic uncertainty and improve the generalization ability. The experimental results demonstrate the effectiveness of our method in various FL tasks.