Making Batch Normalization Great in Federated Deep Learning
This addresses a practical challenge for federated learning practitioners by clarifying when BN is effective and providing an easy-to-implement solution, though it is incremental as it builds on prior empirical studies.
The paper tackled the problem of Batch Normalization (BN) underperforming in federated learning (FL) compared to Group Normalization (GN), finding that BN actually outperforms GN in many FL settings except in high-frequency communication and extreme non-IID regimes, and proposed FIXBN, a simple practice that reduces BN's issues without extra costs, performing favorably across FL settings.
Batch Normalization (BN) is widely used in {centralized} deep learning to improve convergence and generalization. However, in {federated} learning (FL) with decentralized data, prior work has observed that training with BN could hinder performance and suggested replacing it with Group Normalization (GN). In this paper, we revisit this substitution by expanding the empirical study conducted in prior work. Surprisingly, we find that BN outperforms GN in many FL settings. The exceptions are high-frequency communication and extreme non-IID regimes. We reinvestigate factors that are believed to cause this problem, including the mismatch of BN statistics across clients and the deviation of gradients during local training. We empirically identify a simple practice that could reduce the impacts of these factors while maintaining the strength of BN. Our approach, which we named FIXBN, is fairly easy to implement, without any additional training or communication costs, and performs favorably across a wide range of FL settings. We hope that our study could serve as a valuable reference for future practical usage and theoretical analysis in FL.