[论文解读] FedBE: Making Bayesian Model Ensemble Applicable to Federated Learning

Hong-You Chen, Wei‐Lun Chao|arXiv (Cornell University)|Sep 4, 2020

Privacy-Preserving Technologies in Data参考文献 55被引用 28

一句话总结

FedBE 提出了一种基于贝叶斯模型集成的联邦学习聚合方法，可在非独立同分布（non-i.i.d.）数据和深度神经网络下提升模型鲁棒性。通过从拟合于本地模型的分布中采样高质量全局模型，并利用随机权重平均（SWA）进行知识蒸馏，FedBE 在深度网络和异构设置下显著优于 FedAvg，准确率更高。

ABSTRACT

Federated learning aims to collaboratively train a strong global model by accessing users' locally trained models but not their own data. A crucial step is therefore to aggregate local models into a global model, which has been shown challenging when users have non-i.i.d. data. In this paper, we propose a novel aggregation algorithm named FedBE, which takes a Bayesian inference perspective by sampling higher-quality global models and combining them via Bayesian model Ensemble, leading to much robust aggregation. We show that an effective model distribution can be constructed by simply fitting a Gaussian or Dirichlet distribution to the local models. Our empirical studies validate FedBE's superior performance, especially when users' data are not i.i.d. and when the neural networks go deeper. Moreover, FedBE is compatible with recent efforts in regularizing users' model training, making it an easily applicable module: you only need to replace the aggregation method but leave other parts of your federated learning algorithm intact. Our code is publicly available at https://github.com/hongyouc/FedBE.

研究动机与目标

• 解决 FedAvg 在非独立同分布（non-i.i.d.）数据设置下因模型漂移和泛化能力差导致的性能下降问题。
• 通过引入贝叶斯模型集成，克服简单模型平均的局限性，以捕捉不确定性并提升预测鲁棒性。
• 通过将集成预测蒸馏为单个全局模型，实现客户端可复用的有效多轮联邦学习。
• 通过仅替换聚合步骤而不修改客户端训练过程，确保与现有联邦学习框架的兼容性。
• 展示对服务器无标签数据与测试数据之间分布偏移的鲁棒性，提升实际适用性。

提出的方法

• 通过将高斯或狄利克雷分布拟合到本地客户端模型，构建模型分布，从而采样多样化的全局模型候选。
• 通过聚合多个采样得到的全局模型的预测结果，实现贝叶斯模型集成，以提升鲁棒性和准确率。
• 利用无标签服务器数据生成伪标签，用于将集成预测蒸馏到单个全局模型中。
• 在蒸馏过程中应用随机权重平均（SWA），以防止对噪声集成预测的过拟合，提升泛化能力。
• 使用蒸馏损失训练一个学生型全局模型，其中集成预测作为教师信号。
• 通过仅替换聚合步骤，将 FedBE 集成到现有联邦学习框架中，同时保留客户端训练和动量机制。

实验结果

研究问题

• RQ1 与 FedAvg 相比，贝叶斯模型集成是否能在非独立同分布（non-i.i.i.d.）数据条件下提升联邦学习中的模型聚合性能？
• RQ2 在深层神经网络（如 ResNet）中，FedBE 的表现如何，尤其是在 FedAvg 通常性能下降的场景下？
• RQ3 在服务器无标签数据分布与测试数据分布不同的情况下，使用无标签服务器数据进行蒸馏的影响是什么？
• RQ4 在部分客户端参与和系统异构性等实际系统约束下，FedBE 是否能保持良好性能？
• RQ5 将 FedBE 与现有正则化技术（如 FedProx）结合使用时，整体性能如何变化？

主要发现

• 在非独立同分布（non-i.i.d.）的 Tiny-ImageNet 数据集上，FedBE 在 ResNet20 上比 FedAvg 提高 5.9% 准确率（32.4% vs. 35.4%），在 MobileNetV2 上提高 3.4%（25.5% vs. 27.8%）。
• 在非独立同分布（non-i.i.d.）的 CIFAR-10 上，FedBE 在 ResNet20 上达到 77.1% 准确率，显著优于 FedAvg（69.9%）和 FedProx（69.4%）。
• 即使服务器无标签数据的分布与测试数据不同，FedBE 仍能保持高准确率，表现出对领域偏移的鲁棒性。
• 在部分客户端参与（每轮仅 10/100 个客户端参与）的设置下，FedBE 在 Tiny-ImageNet（ResNet20）上达到 35.4% 准确率，优于 FedAvg（32.4%）。
• 在系统异构性（本地训练轮数不一致）条件下，FedBE 在 CIFAR-10（ResNet20）上达到 77.1% 准确率，优于 FedAvg（69.9%）和 FedProx（69.4%）。
• 与 FedProx 结合使用时，FedBE 进一步提升性能，在 CIFAR-10（ResNet20）上达到 77.5% 准确率，表明其与现有正则化方法具有兼容性和协同效应。