[论文解读] FedBN: Federated Learning on Non-IID Features via Local Batch Normalization

Xiaoxiao Li, Meirui Jiang+3 more|arXiv (Cornell University)|Feb 15, 2021

Privacy-Preserving Technologies in Data参考文献 40被引用 77

一句话总结

FedBN 在联邦学习中将 Batch Norm 参数保持本地，以应对特征 shift 的非 IID 数据，在基准数据集和真实世界数据集上，相比 FedAvg 和 FedProx，提升了收敛性和准确性。

ABSTRACT

The emerging paradigm of federated learning (FL) strives to enable collaborative training of deep models on the network edge without centrally aggregating raw data and hence improving data privacy. In most cases, the assumption of independent and identically distributed samples across local clients does not hold for federated learning setups. Under this setting, neural network training performance may vary significantly according to the data distribution and even hurt training convergence. Most of the previous work has focused on a difference in the distribution of labels or client shifts. Unlike those settings, we address an important problem of FL, e.g., different scanners/sensors in medical imaging, different scenery distribution in autonomous driving (highway vs. city), where local clients store examples with different distributions compared to other clients, which we denote as feature shift non-iid. In this work, we propose an effective method that uses local batch normalization to alleviate the feature shift before averaging models. The resulting scheme, called FedBN, outperforms both classical FedAvg, as well as the state-of-the-art for non-iid data (FedProx) on our extensive experiments. These empirical results are supported by a convergence analysis that shows in a simplified setting that FedBN has a faster convergence rate than FedAvg. Code is available at https://github.com/med-air/FedBN.

研究动机与目标

• 在联邦学习中动机化并解决本地特征分布在各客户端不同的特征 shift 非 IID 问题。
• 提出一种轻量方法，通过保持 BN 参数本地化（不进行平均），同时对非 BN 层使用 FedAvg 来缓解特征 shift。
• 给出理论收敛分析，显示在特征 shift 下的收敛速度快于 FedAvg。
• 在基准和真实世界数据集上实证验证 FedBN，相较于 FedAvg 和 FedProx 具有明显的改进。

提出的方法

• 引入 FedBN，在局部更新并聚合非 BN 层的同时，BN 参数保持严格本地，不进行通信。
• 包含两层网络和批归一化作为架构的一部分的模型表述，以及未对 BN 参数进行平均的 FedBN 目标函数定义。
• 使用神经切线核（NTK）框架的收敛分析，表明 FedBN 的收敛受辅助 Gram 矩阵的最小特征值控制，在特征 shift 场景下，FedBN 的收敛速度快于 FedAvg。
• 在基准和真实世界数据集上的实验 protocol，将 FedBN、FedAvg、FedProx 和 SingleSet 进行对比，涵盖在不同数据 regime 和异质性下的收敛行为与准确性。

实验结果

研究问题

• RQ1 联邦学习中的特征 shift 是否会降低 FedAvg 的性能，局部 BN 是否能缓解这种降级？
• RQ2 在特征 shift 的非 IID 数据下，FedBN 是否能提升收敛速度和最终准确性，相较于 FedAvg 和 FedProx？
• RQ3 在具有不同异质程度的基准域和真实世界数据集上，FedBN 的表现如何？
• RQ4 FedBN 对不同本地更新频率和客户端数据规模是否具有鲁棒性？

主要发现

• 在特征 shift 的非 IID 设置下，FedBN 在基准和真实世界数据集上持续优于 FedAvg 和 FedProx。
• 在实验中，FedBN 展现出比 FedAvg 更快的收敛和更稳定的损失曲线。
• 在特征 shift 更明显的数据集（如 SVHN）上，FedBN 获得更大的准确性增益，且跨多次运行的方差更低。
• 在真实世界任务（Office-Caltech10、DomainNet、ABIDE）中，FedBN 在多个客户端/数据集上均实现了比替代方法更高的平均准确性。
• 理论分析表明，在特征 shift 的 G-dominated 区域，FedBN 相对于 FedAvg 具有更快的收敛速度。