联邦学习简介之 FedProto: Federated Prototype Learning across Heterogeneous Clients 阅读笔记

一. 前言

《FedProto: Federated Prototype Learning across Heterogeneous Clients》发表在 2022 年的人工智能顶级会议 AAAI 上，此篇论文旨在解决联邦学习客户端模型和数据异构的问题，并设计了一个简洁且优美的原型学习方法。此篇论文是我个人非常喜欢的一篇联邦学习论文，此方法简洁优美且有效！此篇博文对其核心思想进行阐述，希望可以帮助对此篇论文感兴趣的小伙伴～～～

注：笔者水平有限，若有描述不当或错误的地方，敬请大家批评指正！

二. 问题

联邦学习各个客户端不仅数据分布异构，而且模型架构异构。对于数据分布异构的情况，各个客户端数据的类别数可能并不相同，例如，客户端 1 可能具有类别 3 和 4 ，而客户端 2 可能具有类别 2 和 3。此外，由于模型架构异构，服务器侧无法像联邦平均算法一样执行局部模型参数的聚合。

三. 方法

3.1 核心思想

一方面，FeProto 针对每个类别，通过对齐全局原型（Prototype，即类中心）与局部原型的方式，以解决各个客户端数据分布异构的问题；另一方面，FeProto 通过使用训练完成的类原型执行推理任务，以解决各个客户端模型架构异构的问题。FedProto 算法概览图如图 1 所示。

图 1. FedProto 算法概览 [1]

3.2 核心公式

通过上述核心思想，可以看出，FedProto 最核心的计算在于计算所有类别的全局原型，每个客户端中所有类别的局部原型，局部模型优化目标，基于原型的模型推理。

1. 全局原型聚合：全局原型聚合是将各个客户端上传的相同类别的原型进行加权平均，便可以得到该类别的全局原型，具体公式如下：
   �¯(�)=1|��|∑�∈��|��,�|����(�) （1）
   其中， ��(�) 表示来自于第 � 个客户端的属于类别 � 的原型， |��| 表示拥有类别 � 的客户端集合数量， �� 表示所有客户端中属于类别 � 的数据总量，|��,�| 表示第 � 个客户端中属于类别 � 的数据总量。注：如上所述，在数据异构的场景中，所有的客户端数据并不能保证涵盖所有的类别，因此，这里使用 �� 指定具有类别 � 的客户端集合。
    
2. 局部原型计算：对于每个类别，局部原型计算只需将对应类别所有的数据平均即可得到，具体公式如下：
   ��(�)=1|��,�|∑(�,�)∈��,���(�;�) （2）
   其中， ��(�;⋅) 表示特征提取器，其参数为 � ， ��,� 表示第 � 个客户端中属于类别 � 的数据集合。
    
3. 局部模型优化目标：在计算得到每个类别的局部原型之后，通过最小化每个类别的局部原型与全局原型的距离，以最小化不同客户端数据分布之间的差异，具体公式如下：
   ��(�¯�,��)=∑��(��(�),�¯�(�)) （3）
   其中， �� 表示客户端 � 具有的类别局部原型集合， �¯� 表示与客户端 � 具有类别相对应的全局原型集合， �(⋅,⋅) 表示一个距离函数，可以是 L1 距离和 L2 距离等。
   
   基于上述公式，在结合类别分类损失，可得第 � 个客户端的局部模型优化目标：
   �(��,��)=��(��(��;�),�)+���(�¯�,��) （4）
   其中， �(��;�) 表示局部模型，其参数为 ��（注： �� 包含了特征提取器和类别分类器的参数）， ��(⋅,⋅) 表示监督分类损失，如交叉熵损失。
    
4. 基于原型的模型推理：由于模型架构异构，因此无法通过训练全局的模型对测试样本进行推理。因此，
   FedProto 提出使用全局类别原型基于 1-NN（最近邻分类） 思想对测试样本进行推理，具体公式如下： �^=arg⁡min�‖�(�;�)−�(�)‖22 （5）

3.3 具体步骤

1. 随机初始化所有类别的全局原型集合 {�¯(�)}�=1|�| ，并下发至各个客户端。其中， |�| 表示类别总数。
    
2. 对于每一个客户端 � ，计算新一轮的局部原型集合 �� ，重复执行以下步骤 � 次：

• 采用 mini-batch 方法对客户端 � 的数据进行随机采样，然后利用公式（2）计算局部原型；
• 利用局部原型最小化客户端的最优化目标，即公式（4）；

3. 对于每一个客户端 � ，将步骤 2 得到的最终的局部原型上传至服务器侧。

4. 服务器侧基于公式（1）进行局部原型的聚合，得到所有类别的全局原型，然后将其下发至所有客户端。重复执行步骤 2-4 � 次。

四. 总结

FedProto 思想非常简单，也取得了很好的效果，难得的一篇联邦学习方向的佳作。原文文笔也很不错，非常值得一读！！！

五. 参考文献

[1] Tan Y, Long G, Liu L, et al. Fedproto: Federated prototype learning across heterogeneous clients[C] // AAAI. 2022, 36(8): 8432-8440.