Why Batch Normalization Damage Federated Learning on Non-IID Data?

Area of Science:

• Artificial Intelligence
• Machine Learning
• Distributed Systems

Background:

• Federated learning (FL) trains models at the edge, preserving client privacy.
• Batch normalization (BN) accelerates deep neural network (DNN) training but degrades FL performance with non-i.i.d. data.
• Existing FL algorithms show limited improvement over centralized methods and lack theoretical convergence analysis for BN issues.

Purpose of the Study:

• To provide the first theoretical convergence analysis of batch normalization's impact on federated learning.
• To identify the cause of performance degradation in FL with non-i.i.d. data when using BN.
• To develop a novel FL algorithm robust to data heterogeneity for BN-based DNNs.

Main Methods:

• Developed a convergence analysis demonstrating how BN's statistical parameter mismatch under non-i.i.d. data causes gradient deviation.
• Introduced FedTAN, a new FL algorithm employing iterative layer-wise parameter aggregation.
• Conducted comprehensive experiments to evaluate FedTAN against existing methods.

Main Results:

• The theoretical analysis confirmed that BN's local-global statistical mismatch slows and biases FL convergence.
• FedTAN achieved robust performance across various data distributions, outperforming baseline algorithms.
• Experimental results validated the superiority of FedTAN for training BN-based DNNs in federated settings.

Conclusions:

• Batch normalization's inherent properties create convergence challenges in federated learning with non-i.i.d. data.
• The proposed FedTAN algorithm effectively mitigates these challenges through tailored aggregation strategies.
• FedTAN offers a promising solution for efficient and robust federated learning of DNNs utilizing batch normalization.