论文简介

• 论文名称：Layer-wised Model Aggregation for Personalized Federated Learning
• 作者：Xiaosong Ma，香港理工大学
• 论文链接：https://arxiv.org/abs/2110.07829v1
• 论文来源：CVPR 2022

这篇论文针对的是个性化联邦，尝试将模型按层聚合，从而优化异构数据的客户端优化个性化模型聚合。提出的框架交pFedLA。
做法：给每个客户端使用一个超网络，用来识别每层的贡献因素。同时引入了一种参数化机制来更新层化的聚合权值，利用用户间模型的相似性来实现模型个性化。

现有pFL的问题：
现有的pFL方法大多用距离来度量不同不同的客户端，但这个度量无法准确反映每个本地模型的重要性，可能导致不准确的结合权重或不平衡non-IID分布式数据集，反而阻止了客户端的个性化。作者认为这个原因是神经网络不同层的作用不同，浅层关注局部特征提取，深层用于全局特征提取。因此如果直接粗暴的测量模型距离会忽略这种层级差异，导致个性化不准确，从而影响pFL的训练效率。

作者做了一个实验如下图：

作者进行了模型化方法和分层化方法的对比，实验中客户端1的FC1和FC2明显权重更大，而Conv1的权重比较小，可以为个性化模型的精度产生显著的性能增益。

作者的方法：
在server端，为每个client引入一个专用的超网络，以学习在pFL训练过程中跨客户端层的权重。

论文方法

论文的方法框架如下：

客户端数据集大小为$M= \sum^N_{i=1} m_i$。$\theta_i$表示客户端$i$的模型参数，pFL的训练目标为：



损失函数$L$为：



损失函数用的CE。

作者的思路是：具有相似数据分布的客户端应当具有较高的聚合权重，来加强彼此之间的相互贡献。因此在server端建立了一个相似度的矩阵：

$\alpha_i^{ln}$表示客户端$i$中第$n$层的聚合权向量。$\alpha_i^{ln, N}$表示客户端$N$第$n$层的聚合权。对于所有的$n$层，$\sum^N_{j=1} \alpha_i^{ln,j} =1 $。

聚合权重矩阵$\alpha$的更新过程

每个超网络都是由一个全连接的层构成的，输入是一个embedding，由模型参数自动更新，输出是权重矩阵$\alpha$。将客户端$i$上的超网络定义为：

{\theta^{l1},\theta^{l2},...,\theta^{ln}}是训练过程中的中间参数，$\theta^{ln} = {\theta^{ln}_1, \theta^{ln}_2, ...,\theta^{ln}_N, }$是各客户端第$n$层的集合，$\theta_N^{ln}$是客户端N第$n$层的参数。

$\psi_i$是第$i$个客户端的超网络的参数，超网络的输入为嵌入向量$v_i$，输出为聚合权值矩阵$\alpha_i$。

在与中间参数$\theta^{l1},...,\theta^{ln}$进行加权组合以后，客户端$i$将会对私有数据进行本地训练。
$v_i$和$\psi_i$都在训练过程中进行更新。

--模型聚合
pFedLA的框架中，客户端$i$的模型参数是根据$\alpha_i$通过加权聚合得到的：

。

因此，客户端参数$\theta_i$的优化问题，就转化成了超网络的嵌入向量$v_i$和参数$\psi_i$。

$V$和$\Psi$的更新过程

更新过程通过链式法则推导而来：
通过链式法则求得$v_i$和$\psi_i$的梯度：

$\nabla_\bar{\theta_i} L_i$可以从客户端$i$的本地训练中计算得到，$\nabla_{v_i / \psi_i} H N_i (v_i;\psi_i)$是$\alpha_i$在$v_i / \psi_i$上的梯度。因此可以用更简单通用的方法来更新$v_i$和$\psi_i$。

$\delta_i$表示客户端$i$通过本地培训后参数模型的变化。pFedLA在每轮通信过程中更新客户端$i$的超网络嵌入向量的参数，然后更新聚合权重矩阵$\alpha_i$。

伪码如图：

对通信效率的启发式改进

这里作者降低通信效率的思考是，并不是模型的每一层都需要重新训练，因为如果有的层在本地的权重很大，那就可以用本地模型的该层，不用再从服务器端那里重新下载再个性化了。

通过对本地模型所有层的权重排一个序，选出前$k$个权重最大的层保留，剩余的层再用server端的全局模型进行个性化，以此来节约通信开销，而且让权重高的层不执行聚合过程，还能保证全局端得模型不会被Non-iid情况所污染。

$AT_k$是top-K选择函数。

HeurpFedLA的伪码如下：其实就是在pFedLA的基础上添加了一个是否让该层参与全局训练的判断。

实验环节

数据集：EMNIST、FashionMNIST、CIFAR-10、CIFAR-100。
数据分布：non-iid

• non-IID_1：每个客户端被随机分配4个类（CIFAR100中每个客户端12个类），每个类上的数据量相同
• non-IID_2：每个客户端包含所有类，且每个类的数据不是均匀分布的（还是没具体说怎么分的）
  Baseline：
• FedAvg+Local Training
• Per-Fedavg
• pFedMe
• pFedHN（也是一个用超网络优化pFL的方法，找机会看一看）
• FedBN
• FedRep
• FedFomo

两种框架跟诸多baseline的实验对比：

在两种不同non-iid设置下表现都为最佳，在non-iid_1的情况下，HeurpFedLA的方法甚至有时候比pFedLA还要好，挺神奇的。

权重变化过程
为了证明该方法能够使得数据分布相似的客户端之间产生更高的权重，作者进行了如下实验：挑选了8个客户单，其中4个具有随机数据类，作为目标客户端；4个具有相同的类，作为对比客户端，还有1个与目标客户端有3个相同类的相似客户端。训练过程中记录目标客户端上每一层的权值，包括训练阶段的前期、中期、后期的局和权重的演变。权重演变如下图：

图中可以观察到：随着训练过程的推进，other clients的权重越来越少，因为其数据分布跟目标客户端有很大差异。且具有相似数据分布的客户端（same、same labels clients）比similar client具有更高的权重，说明超网络能够学习到不同client之间数据分布的相似性。

客户端内部权重热图

四张图分别是4个不同数据集的情况。而四张图的对角线都很明显，说明每个客户端的自身权重是最大的，说明pFedLA可以学习到异构客户端之间的相似性。

通信效率分析
这部分实验是为了演示HeurpFedLA在通信效率上的改进。
HeurpFedLA的改进策略是选出top-K个权重最大的模型层保留，不参与全局模型的个性化改进，以此节约通信成本。

从这张图中可以看到，对于（a）的EMNIST数据集，FC3的权重最大，而且当拒绝上传FC3的时候，模型的精度反而还升上去了，作者的解释是，聚合过程中移除分类器层可以提高Non-iid数据集的模型性能，因为这样可以避免在聚合过程中收到其他客户端无关知识的干扰。对于（b）中的CIFAR-10数据集，FC1的权重最大，移除该层后，模型准确度有所降低（降低了1.4%左右）。

超参$k$的选取
$k$指的是HeurpFedLA的改进策略中，保持的权重最高的层数不受server端模型改变的个数。如上面的Table 3所示。k取0~3的时候，模型的精度变化并不明显，因此可以根据通信带宽的限制选取不同的$k$值。

总结

这篇文章比较创新的一点就是把FL从按本地模型聚合，细化到了按层聚合，按层聚合的过程中针对每层权重的问题，引入了超网络的思想来给每个客户端进行参数调节，很像是元学习的思想。特别是第二个框架HeurpFedLA，虽然并没有降低多少通信效率，但是的确一定程度上做到了pFL问题的“各客户端的特殊性”和“全局模型的普适性”之间的权衡。

不过这篇论文一直没有论证FL按层聚合的可行性，按理说模型每层的作用不同（浅层一边做些边缘检测之类的工作，深层会做一些语义特征的识别），每一层都混起来求聚合，真的会让每一层的效果都变好吗？