Academic Summary: FedProto: Federated Prototype Learning Across Heterogeneous Clients

作者及机构信息

本文的作者包括 Yue Tan、Guodong Long、Lu Liu、Tianyi Zhou、Qinghua Lu、Jing Jiang 和 Chengqi Zhang，分别隶属于以下机构：Australian Artificial Intelligence Institute, FEIT, University of Technology Sydney；University of Washington；University of Maryland, College Park；和 Data61, CSIRO。研究成果发表在 “The Thirty-sixth AAAI Conference on Artificial Intelligence (AAAI-22)” 上。

研究背景与目的

研究领域： 联邦学习（Federated Learning，简称 FL）。

联邦学习是一种允许多个客户端在数据隐私保护的限制条件下协同训练全局模型的技术，适用于数据存储分散且无法共享原始数据的场景。然而，传统的联邦学习方法主要通过客户端之间共享梯度或模型参数来聚合模型。对于真实世界中存在的数据分布不一致（Non-IID）以及模型架构不一致等异构性问题，这些方法通常表现出较差的优化收敛性和泛化能力。

现有的研究中，统计异构性（Statistical Heterogeneity）和模型异构性（Model Heterogeneity）是联邦学习的两大主要挑战。尽管已有个性化联邦学习（Personalized FL）和知识蒸馏（Knowledge Distillation，简称 KD）等方法试图应对这些问题，但它们存在通信成本高或依赖公共数据集等局限。

本研究提出了一种新颖的联邦原型学习框架 FedProto，通过让客户端与服务器交换抽象类原型（prototypes）而非梯度，从而高效处理联邦学习中的异构问题。研究的目标包括： 1. 提高联邦学习对数据和模型异构性的容忍度； 2. 提供理论证明 FedProto 框架的非凸优化收敛性； 3. 在实验和应用层面验证 FedProto 的优越性能。

研究方法与实验设计

整体方法流程： FedProto 是基于原型聚合（Prototype Aggregation）的联邦学习框架，旨在通过共享原型代替传统的梯度通信，实现高效信息交换。框架中包括以下关键步骤：
1. 在客户端本地更新类别原型； 2. 将本地原型发送到中央服务器； 3. 服务器聚合收集到的本地原型，并向客户端发送全局原型； 4. 客户端通过规整化损失函数，将本地原型调整至靠近全局原型，同时优化分类任务的损失。

研究设计的主要步骤：

1. 客户端本地数据处理与建模： 每个客户端具有自己的局部数据集，其数据分布和类别可能与其他客户端不一致（统计异构性）。此外，客户端可以采用不同的模型架构或超参数（模型异构性）。局部模型中包括嵌入层（Representation Layers）和决策层（Decision Layers），分别用于特征提取和分类任务。

   • 局部原型通过计算每类嵌入特征的均值来生成。
   • 为了确保局部模型生成的原型与全局原型对齐，采用了一种新的损失函数，其中包括分类误差损失 (L_s) 和全局与局部原型距离的正则项 (L_r)。

2. 服务器端全局原型聚合： 中央服务器负责从各个客户端收集原型并进行聚合。具体地，对于同一类别，服务器通过加权平均的方式生成全局原型 ( \bar{c} )。

3. 优化与通信： FedProto 的目标函数综合考虑了所有客户端的任务损失和原型对齐的规整化项。优化通过交替最小化局部任务损失和全局原型更新来完成。

4. 支持模型和统计异构性： FedProto 的关键创新在于其完全摒弃了梯度或参数的共享，因此允许各客户端使用完全不同的模型架构和类别空间。这种灵活性显著提高了其应对实际场景中复杂异构问题的能力。

实验设计： - 数据集： 使用三个常见联邦学习数据集 MNIST、FEMNIST 和 CIFAR-10。 - 模型： MNIST 和 FEMNIST 使用两层卷积网络；CIFAR-10 使用 ResNet18。 - 异构性模拟： - 随机设置客户端各自的类别数量（n-way）及每类别样本数（k-shot），模拟统计异构性。 - 客户端间的卷积层通道数或步幅差异，用于模拟模型异构性。 - 比较方法： 使用 FedAvg、FedProx、Fesem 等传统联邦学习方法作为基线；并与个性化方法如 FedPer、FedRep 进行比对。

实验结果

主要性能比较： FedProto 在不同水平的统计异构性或模型异构性条件下，均取得了最高的平均测试准确率和最低的样本间差异。这说明 FedProto 能更好地平衡泛化性和个性化性能。

具体来说： - 在 MNIST 数据集上，FedProto 达到约 97% 的测试准确率，显著超越了 FedAvg 和 FedProx。 - 在 FEMNIST 数据集上，即使在类别和样本数量极度不平衡的情况下，FedProto 仍能保持较好的模型一致性。 - 在 CIFAR-10 数据集上，当客户端模型架构异构时，FedProto 依旧展现出优秀的韧性。

通信效率： FedProto 极大地降低了通信成本，每次通信仅需传输小规模的原型，而非完整的模型参数。在 MNIST 数据集上，每轮通信的参数量仅为传统方法的1%，收敛所需轮数也显著减少。

t-SNE 可视化： 通过对 MNIST 数据集的嵌入特征进行 t-SNE 降维，FedProto 的原型及样本点分布更加集中、可区分。这证明了 FedProto 在加强全局模型泛化性能的同时，仍然保留了每个客户端的独立特性。

可伸缩性分析： FedProto 在样本数量较少的情况下表现出更强的鲁棒性。当每类训练样本数从 120 减少到 5 时，其性能下降幅度远低于其他方法。

研究结论与意义

研究结论： FedProto 提出了一种全新的利用原型聚合的联邦学习方法，成功应对了统计异构和模型异构的双重挑战。通过规整化的优化目标和交替更新策略，该方法实现了更高效的通信与更强泛化能力。

科学价值： 1. 提供了一种无需共享参数的联邦学习方法，为高隐私场景提供了新思路。 2. 在理论上证明了 FedProto 的非凸优化收敛性，并为联邦学习体系设计贡献了全新的方法框架。

实际应用价值： FedProto 可广泛应用于异构分布场景，如个性化医疗诊断、智能交通或分布式物联网系统，并有效提升模型性能与数据隐私保护水平。

亮点与创新

• 方法创新： 提出了原型通信中替代梯度聚合的新模式。
• 适用范围广： 支持异构输入/输出空间与异构模型。
• 通信高效： 显著降低了通信成本。
• 隐私保护： 数据隐私性强，攻击者难以通过原型还原训练数据。

FedProto 为联邦学习的未来研究与实践提供了重要的理论基础与实践工具，是应对复杂异构环境中数据协同学习的有效范式。