[论文解读] Making Batch Normalization Great in Federated Deep Learning

Jike Zhong, Hong-You Chen+1 more|arXiv (Cornell University)|Mar 12, 2023

Neural Networks and Applications被引用 4

一句话总结

该论文提出FixBN，一种简单而有效的方法，通过解决非独立同分布（non-IID）数据和高通信频率引起的梯度偏差与统计不匹配问题，恢复联邦学习（FL）中批量归一化（BN）的性能。FixBN在不增加额外训练或通信成本的前提下，保留了BN的优势，且在各种联邦学习设置中，尤其是在低通信和中等非IIDs场景下，表现优于BN和组归一化（GN）。

ABSTRACT

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.

研究动机与目标

• 重新评估联邦学习中批量归一化（BN）的使用，挑战因感知到不稳定而普遍用组归一化（GN）替代BN的常见做法。
• 识别并解决BN在联邦学习中性能退化的主要原因，特别是在高通信频率和极端非IIDs设置下。
• 提出一种实用且轻量级的解决方案，保留BN在收敛性和泛化能力方面的优势，同时消除其在联邦学习中的缺陷。
• 通过使BN在联邦学习中可稳定使用，弥合集中式与联邦深度学习之间的差距，且无需改变模型架构或训练成本。

提出的方法

• FixBN通过在每个通信轮次将客户端特定的移动平均与全局模型同步，修改了本地训练期间BN层的运行统计量更新规则。
• 确保BN层中的运行均值和方差在各客户端间一致更新，从而减少本地与全局数据分布之间的统计不匹配。
• 采用延迟更新策略：仅在全局模型聚合后才更新本地运行统计量，防止因非IIDs小批量统计量导致的发散。
• FixBN不改变前向传播过程，也无需额外参数，因此与现有联邦学习框架和训练流程完全兼容。
• 该方法轻量化，除标准FedAvg外不增加任何额外通信或计算开销。
• 其设计支持与预训练模型及广泛使用的基于BN的架构（如ResNet）向后兼容，实现无缝集成。

实验结果

研究问题

• RQ1 在何种联邦学习设置中，批量归一化（BN）优于组归一化（GN），又在何时会失效？
• RQ2 BN在联邦学习中性能退化的主要原因是什么，特别是在高通信频率和极端非IIDs场景下？
• RQ3 是否可以在不改变模型架构或训练过程的前提下，缓解BN中的梯度偏差与统计不匹配问题？
• RQ4 对BN运行统计量更新规则进行简单修改，能否在保留其收敛优势的同时恢复其在联邦学习中的性能？

主要发现

• FixBN显著提升了高通信频率联邦学习设置下BN的性能，在E=1本地步数时仍能恢复集中式学习的准确率。
• 在CIFAR-10数据集上，使用非IIDs分片且E=100时，FixBN达到76.56%的测试准确率，优于GN（59.69%）和标准BN（53.97%）。
• 在低通信场景（如10K轮次）中，FixBN保持强劲性能（非IIDs CIFAR-10上达87.71%），优于FedAvg+BN（45.96%）和FedAvg+GN（82.66%）。
• FixBN使BN能从增加的本地步数（E）中受益，随着E增大准确率持续提升，而GN在相同条件下性能反而下降。
• FixBN优于其他归一化方法，如Fixup（70.66%）和带权重归一化的GN（66.90%），展现出其鲁棒性与有效性。
• 该方法在不增加通信或训练成本的前提下解决了梯度偏差与统计不匹配问题，使其在真实联邦学习部署中具有实际可行性。