[论文解读] Boosting Federated Learning Convergence with Prototype Regularization

Yu Qiao, Huy Q. Le+1 more|arXiv (Cornell University)|Jul 20, 2023

Privacy-Preserving Technologies in Data被引用 5

一句话总结

该论文提出FedPR，一种联邦学习框架，通过引入基于原型的正则化，加速了非独立同分布（non-IID）设置下的收敛速度并提升了准确率。客户端本地计算各类别的原型，服务器将这些原型聚合为全局原型；随后，利用L2距离最小化，将该全局原型用于正则化本地训练，使MNIST和Fashion-MNIST上的测试准确率分别比FedAvg高出3.3%和8.9%。

ABSTRACT

As a distributed machine learning technique, federated learning (FL) requires clients to collaboratively train a shared model with an edge server without leaking their local data. However, the heterogeneous data distribution among clients often leads to a decrease in model performance. To tackle this issue, this paper introduces a prototype-based regularization strategy to address the heterogeneity in the data distribution. Specifically, the regularization process involves the server aggregating local prototypes from distributed clients to generate a global prototype, which is then sent back to the individual clients to guide their local training. The experimental results on MNIST and Fashion-MNIST show that our proposal achieves improvements of 3.3% and 8.9% in average test accuracy, respectively, compared to the most popular baseline FedAvg. Furthermore, our approach has a fast convergence rate in heterogeneous settings.

研究动机与目标

• 解决因客户端间数据分布非独立同分布（non-IID）导致的模型性能下降问题。
• 在不损害数据隐私的前提下，提升异构数据设置下的收敛速度与测试准确率。
• 利用全局类别原型指导本地模型训练，增强客户端之间的泛化能力。
• 设计一种低复杂度、可扩展的正则化机制，将原型整合进联邦训练流程。

提出的方法

• 客户端使用每类训练样本的平均嵌入计算本地原型。
• 服务器通过简单平均将所有客户端的本地原型聚合，形成每类的全局原型。
• 将全局原型传回客户端，并通过本地特征表示与全局原型之间的L2距离损失，用于正则化本地训练。
• 本地目标函数结合标准交叉熵损失与原型正则化项：$\mathcal{L}_{i}(\omega_{i}) = \mathcal{L}_{i}(\mathcal{F}(\omega;\boldsymbol{x}_{i}),y_{i}) + \ell_{2}(f_{e}(\omega_{e};\boldsymbol{x}_{i}) - \overline{y}_{j})$。
• 采用类似FedAvg的通信机制，通过同步轮次迭代更新模型参数与全局原型。
• 框架使用4层卷积神经网络实现，并在Dirichlet分布数据偏移（$\alpha = 0.05$）下进行测试，以模拟强非IID条件。

实验结果

研究问题

• RQ1 基于原型的正则化是否能提升非IID联邦学习设置下的收敛速度与测试准确率？
• RQ2 将全局原型整合进本地训练，如何影响在异构数据下客户端之间的模型泛化能力？
• RQ3 在强数据偏移条件下，所提方法是否在准确率与收敛速率上优于FedAvg？
• RQ4 原型聚合能否在大量客户端下高效计算并实现良好扩展性，同时保持较低通信开销？

主要发现

• 在MNIST上，FedPR实现了94.62%的平均测试准确率，相较于FedAvg在$\alpha = 0.05$数据偏移下的91.57%提升了3.3%。
• 在Fashion-MNIST上，FedPR达到86.05%的平均测试准确率，相较于FedAvg的79.04%提升了8.9%。
• 与FedAvg相比，该方法表现出更快的收敛速度，表现为测试准确率在更早的通信轮次即趋于稳定。
• 原型正则化机制有效缓解了联邦设置中因类别不平衡与数据异质性导致的性能下降。
• 该框架保持了较低的算法复杂度，适合在实际联邦系统中部署。
• 将全局原型作为正则化器，通过将本地表示对齐到共享的、全局聚合的类别结构，提升了模型泛化能力。