[论文解读] FedProc: Prototypical Contrastive Federated Learning on Non-IID data

Xutong Mu, Yulong Shen+5 more|arXiv (Cornell University)|Sep 25, 2021

Privacy-Preserving Technologies in Data被引用 11

一句话总结

FedProc 提出了一种新颖的联邦学习框架，利用全局类别原型作为全局知识，以纠正非独立同分布（non-IID）数据上的本地模型训练。通过引入原型对比损失和混合本地网络架构，FedProc 将本地目标与全局最优对齐，在 CIFAR-10、CIFAR-100 和 Tiny-ImageNet 上实现了 1.6% 到 7.9% 的最先进准确率提升，同时计算成本可接受。

ABSTRACT

Federated learning allows multiple clients to collaborate to train high-performance deep learning models while keeping the training data locally. However, when the local data of all clients are not independent and identically distributed (i.e., non-IID), it is challenging to implement this form of efficient collaborative learning. Although significant efforts have been dedicated to addressing this challenge, the effect on the image classification task is still not satisfactory. In this paper, we propose FedProc: prototypical contrastive federated learning, which is a simple and effective federated learning framework. The key idea is to utilize the prototypes as global knowledge to correct the local training of each client. We design a local network architecture and a global prototypical contrastive loss to regulate the training of local models, which makes local objectives consistent with the global optima. Eventually, the converged global model obtains a good performance on non-IID data. Experimental results show that, compared to state-of-the-art federated learning methods, FedProc improves the accuracy by $1.6\%\sim7.9\%$ with acceptable computation cost.

研究动机与目标

• 解决客户端间非独立同分布（non-i.i.d.）数据分布导致的联邦学习性能下降问题。
• 通过将本地训练目标与全局最优对齐，提升模型泛化能力和收敛性。
• 通过全局类别原型利用客户端特定知识，实现更优的表征学习。
• 在显著优于现有最先进联邦学习方法的非 IID 基准上，保持计算效率的同时实现显著性能提升。

提出的方法

• FedProc 引入了一种全局原型对比损失，将每个样本的表征拉近至其全局类别原型，同时推离其他类别原型。
• 该框架采用混合本地网络架构，通过对比损失与交叉熵损失的加权组合，联合优化特征学习与分类。
• 全局类别原型由服务器通过聚合客户端级别的原型计算得出，随后广播回客户端用于本地训练校正。
• 损失函数通过系数 α = 1 - t/T 动态调整表征学习与分类学习之间的权衡，早期训练阶段更侧重表征学习。
• 客户端利用全局原型作为归纳偏置，以稳定本地训练并减少与全局最优的偏离。
• 该方法兼容标准 FedAvg 通信协议，可无缝集成至现有联邦学习流程中。

Figure 1: Federated learning process.

实验结果

研究问题

• RQ1 全局类别原型能否作为有效的归纳偏置，以纠正非 IID 联邦学习中的本地模型漂移？
• RQ2 原型对比学习在数据分布偏斜的联邦设置下，如何提升泛化能力和收敛性？
• RQ3 在本地训练中动态调整表征学习与分类学习损失权重的影响是什么？
• RQ4 FedProc 在非 IID 数据上，随着客户端数量和通信轮次增加，其可扩展性如何？

主要发现

• 在 100 个客户端的 CIFAR-100 上，FedProc 实现了 70.6% 的 top-1 准确率，显著优于此前最先进方法 MOON 的 61.8%。
• 与现有 SOTA 方法相比，FedProc 在 Tiny-ImageNet 上准确率提升 7.9%，在 CIFAR-10 上提升 1.6%。
• 该方法保持了较低的计算开销，每轮训练时间与 FedAvg 相当，且在 Tiny-ImageNet 等较大数据集上优于 SCAFFOLD 和 MOON。
• FedProc 展现出优异的可扩展性，在 100 个客户端和 500 次通信轮次下，CIFAR-100 准确率达到 72.5%，较 MOON 提升 9.3%。
• 消融实验表明，动态损失权重（α = 1 - t/T）在性能上优于固定权重或两阶段训练策略。
• 全局原型对比损失有效将本地目标与全局最优对齐，减少模型漂移，并在非 IID 数据上增强泛化能力。

Figure 3: Overview of the local network architecture in FedProc. The feature extraction network (including the base encoder and the projection head) extracts the representation $z$ , which is used to calculate the global prototypical contrastive loss $\ell_{gpc}$ . By inputting the representation $z