Why Batch Normalization Damage Federated Learning on Non-IID Data?
It addresses a critical bottleneck in federated learning for privacy-preserving distributed training, offering a solution to performance degradation on non-i.i.d. data, though it is incremental as it builds on prior FL algorithms.
The paper tackles the problem of batch normalization (BN) impairing federated learning (FL) performance on non-i.i.d. data by providing the first convergence analysis showing that parameter mismatch causes gradient deviation, and proposes FedTAN, a new algorithm that achieves robust performance, with experimental results demonstrating superiority over existing baselines.
As a promising distributed learning paradigm, federated learning (FL) involves training deep neural network (DNN) models at the network edge while protecting the privacy of the edge clients. To train a large-scale DNN model, batch normalization (BN) has been regarded as a simple and effective means to accelerate the training and improve the generalization capability. However, recent findings indicate that BN can significantly impair the performance of FL in the presence of non-i.i.d. data. While several FL algorithms have been proposed to address this issue, their performance still falls significantly when compared to the centralized scheme. Furthermore, none of them have provided a theoretical explanation of how the BN damages the FL convergence. In this paper, we present the first convergence analysis to show that under the non-i.i.d. data, the mismatch between the local and global statistical parameters in BN causes the gradient deviation between the local and global models, which, as a result, slows down and biases the FL convergence. In view of this, we develop a new FL algorithm that is tailored to BN, called FedTAN, which is capable of achieving robust FL performance under a variety of data distributions via iterative layer-wise parameter aggregation. Comprehensive experimental results demonstrate the superiority of the proposed FedTAN over existing baselines for training BN-based DNN models.