论文阅读——FedTP: 基于Transformer的个性化联邦学习

文章基本信息

作者：

• 汪婧雅团队

机构：

• 上海科技大学信息科学与技术学院

期刊：IEEE Transactions on Neural Networks and Learning Systems（中科院一区）
发表时间：

• 提交于 2022年10月22日
• 修改于 2023年2月14日
• 接收于 2023年4月8日

解决问题：提出了 FedTP 框架，将 transformer 应用于联邦学习。在聚合客户端的其他参数的同时，学习各个客户端的 self attention 中的参数。
贡献：

1. 提出了 FedTP 框架，一个使用 transformer 进行联邦学习的框架。
2. 提出了一种 learn-to-personalize 机制，用于解决客户端之间交换 self attention 的参数的问题。

代码：
GitHub - zhyczy/FedTP: This is official implementation of FedTP (TNNLS.2023) [https://ieeexplore.ieee.org/abstract/document/10130784].
文章：
FedTP: Federated Learning by Transformer Personalization

摘要

摘要部分文章主要说明了个性化联邦学习在解决客户端异构性的重要作用，但是在联邦学习中使用 transformer 模型的工作还不完善，没有讨论联邦学习的聚合算法对 self attention 的影响。
在研究 transformer 在联邦学习中的效果的过程中，本文发现 fedavg 对 self attention 有负面影响。所以提出了基于 transformer 的联邦学习框架 FedTP，该框架在学习每个客户端的自注意力的同时，聚合客户端之间的其它参数。本文还提出了一种 learn-to-personalize 机制，通过学习一个服务器端上的超网络，生成客户端的 QKV，进一步促进客户之间的协作。

结论

实验结果表明，基于 learn-to-personalize 机制的 FedTP 在 non-IID 场景下具有良好的性能，在个性化联邦学习中的性能优于现有方法。优于 FedTP 在所有客户端上学习了一个中央超网络，因此该模型更好地适应新客户端，证明其泛化能力较好。
并且，FedTP 与许多现有的个性化联邦学习方法兼容，可以轻易地将 FedTP 和这些方法结合如 FedPer、FedRod 和 KNN-Per 等。实验结果表明这些组合模型取得了较好的性能。但是 FedTP 的主要限制是训练速度较慢，优化效率是下一步的工作。

引言

引言部分先简单介绍了联邦学习，引出 non-IID 的问题以及个性化联邦学习的解决方法。
接下来介绍了数据异构性的问题以及目前大部分联邦学习框架中的模型是 CNN 等卷积模型的变体，这些基于卷积的方法对数据异构性很敏感，而现在的 transformer 模型使用自注意力机制来学习数据的全局交互，所以对 non-IID 场景更加适用。所以有一个工作 [Rethinking architecture design for tackling data heterogeneity in federated learning] 使用 transformer 代替了 CNN 等模型作为联邦学习框架的基准模型。
但是上述工作使用的是 FedAvg 算法，本文研究目的解决 FedAvg 对 transformer 中 self attention 的负面影响。作者总结了目前个性化联邦存在的问题：

1. 现有联邦学习框架大多使用 CNN，但 CNN 对 non-IID 数据非常敏感。
2. 现有的基于 transformer 的联邦学习框架没有研究聚合对 self attention 的影响。
3. 目前缺乏一个更适用于 transformer 的个性化联邦学习框架。

并且下面有一个作者比较创新的写法，在 intro 部分作者插入了一个简单的实验来表明 fedavg 对 self attention 存在负面影响，即设置了一个对比实验，使用注意力转出方法生成了注意力图如下：

可以从这个图中看出只有 FedAvg-T 的注意力图中红色部分不是关键信息。原始个性化 机制通过局部独立训练客户端的 self attention，它并没有考虑到客户端之间的关系，所以可能是次优的。因此本文设计的 FedTP 采用超网络的方式共享客户端的参数，来完善各自的 self attention。
本文贡献如下：

1. 提出 FedTP 框架
2. 提出 learn-to-personalize 机制
3. 在不同的 non-IID 数据设置下对三个 base-line 数据集进行实验

相关工作

这一部分主要从下面几个方面介绍相关工作：

1. 个性化联邦学习 PFL（personalized federated learning），作者从全局模型个性化和学习个性化模型介绍。
   1. 其中全局模型个性化目的是在给定异构数据的情况下提高单个共享全局模型的性能，一种简单的方法就是在客户端的本地数据集上微调全局模型来产生个性化参数。还有一种简单的方法是添加近端正则化项，来处理由统计异质性导致的客户端漂移问题，比如 FedProx。
   2. 学习个性化模型更适合异构客户端，一种方法是在全局模型和局部模型之间进行插值，类似于知识蒸馏的方法
2. transformer 在各个领域的应用以及和联邦学习结合的优点和作用
3. 超网络 hypernetworks，介绍了超网络的概念和相关工作。

通过 Transformer 个性化实现联邦学习

问题形式化

FedTP 框架使用传统的 ViT 用于图像任务，并使用 transformer 用于 NLP 任务，输入序列 $X$具有预定的长度$m$，图像在 ViT 的图像与处理层中被分割为一个序列。然后该序列被转换为一个对应的嵌入矩阵$H$，维度为${\mathbb{R}}^{m\times d}$。自注意力机制的 QKV 分别表示为：
$Q=H{W}^Q,K=H{W}^K,V=H{W}^V$
再将这些投影参数合并为：
$W=[W^Q,W^K,W^V]$
接下来,通过以下公式应用自注意力机制:
$\text{Attention}(Q,K,V)=\text{softmax}\left(\frac{Q{K}^T}{\sqrt{d}}\right)V$
其中 d 是 QKV 的列数.
为了模拟联邦场景,假设有$N$个客户,每个客户$i$有一个本地数据集$D_i=\{(x_i^{(j)},y_i^{(j)}){\}}_{j=1}^{m_i}$,样本数量为$m_i$，来自不同的数据分布${\mathcal{P}}_i$，目标是最小化每一个客户端个性化模型的损失函数：
$\arg {\min }_{\Theta }{\sum }_{i=1}^N\frac{m_i}{M}{\mathcal{L}}_i({\theta }_i)$
其中${\mathcal{L}}_i({\theta }_i)={\mathbb{E}}_{(x_i^{(j)},y_i^{(j)})\in D_i}l(f({\theta }_i;x_i^{(j)}),y_i^{(j)})$，$l(\cdot ,\cdot )$是损失函数，可以根据不同的任务选择不同的损失函数如交叉熵损失或均方误差。${\theta }_i$是个性化参数。

整体框架图如下：

自注意力的原始个性化

vanilla 在深度学习中就是最初的、原始的意思。所以这部分讲的是本文提出了一种个性化的自注意力机制，每个客户端在本地维护各自的自注意力层，其他参数参与联邦聚合和更新，学习公共信息。如框架图左侧所示：

左侧灰色部分即各个客户端的 MLP 层、Norm 层等参与联邦学习，而自注意力层只在客户端本地保存。
与 FedAvg 类似，个性化参数通过局部训练若干轮次后在服务器上进行聚合，个性化参数${\theta }_i$被分割为$W_i$和$\xi$前者保留在本地以存储客户特定信息，后者在客户端之间进行聚合：
${\bar{\xi }}^t={\sum }_{i=1}^N\frac{m_i}{M}{\xi }_i^{t,k}.$
推导出 FedTP 的目标函数为：
$\arg {\min }_{\Theta }{\sum }_{i=1}^N\frac{m_i}{M}{\mathcal{L}}_i(W_i,\xi )$
其中
${\mathcal{L}}_i(W_i,\xi )={\sum }_{i=1}^N\frac{m_i}{M}{\mathbb{E}}_{\left(x_i^{(j)},y_i^{(j)}\right)\in D_i}l\left(f\right(W_i,\xi ;x_i^{(j)}),y_i^{(j)}).$

这种方法忽略了客户端之间潜在的内在联系，所获得的个性化 self attention 也可能只是次优的。此外，个性化自注意力层是不可扩展的，因为它随着客户端数量增长而线性增长。并且这种方式的泛化能力也不好。

Learn-to-Personalize 自注意力（重点）

这一部分介绍本文的核心部分即 FedTP 和 learn-to-personalize 机制。我们看到框架图的右侧：

为了改进自注意力的原始个性化机制，FedTP 使用了一个超网络来生成每个客户端的自注意力层投影矩阵。超网络在服务器上学习，它是一个简单的 MLP，通过输入客户端的嵌入向量生成自注意力的投影参数$W$。超网络表示为$h(\phi ;z_i)$，其中$\phi$是参数，$z_i$可以是对应于客户端$i$的可学习嵌入向量或固定向量。对于给定的嵌入向量$z_i$，超网络输出部分权重$W_i=h(\phi ;z_i)$，然后将其分解为自注意力机制的 QKV 的投影参数：$W=[W^Q,W^K,W^V]$，通过这种方式，超网络学习一类投影参数用于个性化的自注意力。损失函数替换为：
$\begin{array}{rl}{\mathcal{L}}_i(W_i,\xi )& ={\mathcal{L}}_i(h(\phi ;z_i),\xi )\\ & =\sum _{i=1}^N\frac{m_i}{M}{\mathbb{E}}_{\left(x_i^{(j)},y_i^{(j)}\right)\in D_i}l\left(f\right(h(\phi ;z_i),\xi ;x_i^{(j)}),y_i^{(j)}).\end{array}$
算法 1 演示了 FedTP 的流程：

与原始个性化机制相比，learn-to-personalize 机制的优点如下：

1. 可以在客户端之间有效地共享 self attention 机制的参数
2. 有良好的可扩展性
3. 对新客户具有较好的泛化能力

泛化上界

这一部分主要是统计学上的证明，我还看不懂，就不介绍了，有兴趣可以看原文研究一下。

实验

实验设置

基线模型

• FedAvg
• FedProx
• FedPer
• pFedMe
• FedBN
• pFedHN
• pFedGP
• FedRoD

数据集

• CIFAR10（图像数据集）
• CIFAR100（图像数据集）
• Shakespeare（语言数据集）

对 CIFAR10 数据集和 CIFAR100 都采用了两种划分方式：

• CIFAR10
  • Pathological：每个客户端随机分配 CIFAR10 中的 2 类数据。
  • Dirichlet：直接使用狄利克雷分布进行划分，浓度参数$\alpha =0.3$。
• CIFAR100
  • Pathological：每个客户端随机分配 CIFAR100 中的 10 类数据。
  • Dirichlet：先使用浓度参数$\alpha =0.3$的狄利克雷分布进行划分，再对具体标签进行浓度参数$\beta =10.0$的狄利克雷分布划分。
• Shakespeare
  • 类似于 [Personalized federated learning through local memorization] 中的划分。

模型

• 使用具有两个卷积层和三个全连接层的 ConvNet 作为 CIFAR10 和 CIFAR100 上基线方法的局部神经网络。为了提高大规模联邦学习场景的通信效率，选择了一个参数较少的小型 ViT，由 8 个块组成，每个块中有 8 个自注意力层，对应的头数为 8，patch 大小为 4，隐藏大小为 128.
• 对于Shakespeare ，采用了堆叠两层的 LSTM 模型和一个简单的 transformer。其中有两个块作为编码器。深度都是 2。hidden 维度都是 256。

对于两种任务，采用了相同的超网络即 4 个全连接层的 MLP

实验结果

（本文实验非常详细全面，这里只给出部分实验结果）
作者在 50 个和 100 个客户端场景下测试了 FedTP 和基线方法，每个通信轮中抽样 10%客户端。对于图像任务，每个方法进行了 1500 个通信回合。

因为其中很多方法使用的模型是 CNN，作者都用 transformer 进行了替换，使得实验更加准确。
具体变化过程，FedTP 更平滑。

注意力图的可视化：


扩展性实验：


新客户端的泛化能力实验：

T-SNE可视化FedTP在CIFAR-100数据集上的学习客户端嵌入：

消融实验

客户端数据异质性影响：

噪声对实现结果的影响：

自注意力块数量的影响：

客户端采样率的影响：