ECCV 2022 | FedX: 在无监督联邦学习中进行知识蒸馏

©PaperWeekly 原创 · 作者 | 张成蹊

单位 | Freewheel机器学习工程师

研究方向 | 自然语言处理

论文的立意是在联邦学习过程中，结合对比学习，用局部与全局的两种蒸馏方式来无监督地学习样本的向量表示。FedX 不仅是一种创新的无监督学习算法，更是一种可以热插拔用于传统对比学习，使之效果提升的无监督损失函数。

关键词：联邦学习（Federated Learning），对比学习（Contrastive Learning），知识蒸馏（Knowledge Distillation），无监督学习（Unsupervised Learning）。

论文标题：

FedX: Unsupervised Federated Learning with Cross Knowledge Distillation

论文链接：

https://arxiv.org/abs/2207.09158

代码链接：

https://github.com/Sungwon-Han/FEDX

相关背景

联邦学习假设数据分散在不同的客户端中，且不能与中央服务器共享；截止目前，已经有很多有监督联邦学习用在下游任务的相关研究，如目标检测、图像分割、人员重识别等，但对于无监督联邦学习的需求也在逐步增加，以应对一些本地客户端缺失数据的场景。其中比较具有挑战性的研究点为：

1. 数据在不同客户端是离散且异构的，不满足数据独立同分布的条件，所以联邦学习算法必须考虑到独立客户端与全局数据分布不同时的有偏情况；

2. 在数据有偏到一定程度，客户端可能只拥有较少的数据，此时需要引入无监督/半监督的联邦学习方法，使得每一个客户端都从全局模型中获益，达到较好利用本地数据的目的。

联邦学习的整体流程如下图所示：

其中：

1. Local Update：每个client各自学习自己的样本表示，获得各自独立的本地参数模型；

2. Local model update：将本地参数模型上传到中心服务器；

3. Global aggregation：中心服务器对所有 client 上传的模型进行聚合，最简单的方式是（假设本地模型都是同构的）根据 client 各自拥有样本数量的占比进行加权平均；

4. Global model download：每个 client 将 global model 各自下载到本地。

论文主要聚焦讨论的是步骤 1 中的独立学习部分，也是数据偏差最大的部分。

总的来说，论文具有以下几方面的贡献：

1. 提出了无监督的联邦学习算法 FedX，使用了局部与全局两个阶段的蒸馏来学习不同客户端的数据表示；两阶段的知识蒸馏能够发现有意义的本地数据向量表示，同时通过全局的知识蒸馏消除不同客户端之间的偏差；

2. FedX 是一个可以被用于现有联邦学习扩展的算法，提升现有联邦学习算法的性能（1.58-5.52pp），同时对不同客户端的数据严格保密，不（像部分现有的无监督联邦学习算法一样）直接在不同客户端之间进行数据的共享；

3. FedX 已开源。

局部与全局蒸馏

以 SimCLR 为例，传统对比学习的损失函数如下所示：

其中， 与 分别表示原批次的训练样本，和基于原样本进行了数据增强后的训练样本，即从一个真实样本中采用窗口看到的某个切片图。 表示某种对两个向量的相似性度量方法（比如余弦相似度，值得注意的是，作者在论文中也提到了余弦相似度，不过在实际的代码实现里用的是两个正则化（torch.F.normalize）后的向量内积）， 类似知识蒸馏中用于控制程度的超参（temperature）。

可以容易地看出，对比学习损失函数的目标是：对于样本 ，已知模型最后一个隐层对于它的向量表示 ，及其增强后的样本 ，希望两者的向量更加接近，其相似度更高；而对于两个批次中其他 个负样本，希望其向量相似度越低。

2.1 局部蒸馏

作为一个独立的本地客户端，其学习的目标是基于自己拥有的本地数据，获得数据有意义的向量表示。

作者认为，局部模型可以通过提取不同样本间的结构化知识，实现局部知识传递的目标。在整个过程中，从自身学习到的知识逐步地由不同的局部样本传递给自身。在本地模型更新时，除了每个样本和自己的增强后样本进行对比学习，两者还可以从彼此的分布中进一步学习，谓之为关系损失（relational loss）。下式是关系损失的形式化定义：

该公式表示了正样本 的隐层向量 对于一个随机采样的批次 中任意一个样本 的隐层向量 的概率分布。与之对应，正样本在对比学习中的增强样本隐层向量 与后者也会有一个对应的向量分布：

这两个分布代表了一组相似的正样本关于其它负样本的相对分布情况。容易发现，如果模型真的能辨认出增强前后的样本归属于同一个样本（或者标签），那么，除了对比学习本身设计的学习目标 —— 这两个样本间的相似度高于同批次中其他负样本的相似度——以外，这两个样本对于某个随机的负样本的相似度分布也应该相似。

但是这个相似度要怎么度量呢？作者提出将这两个概率向量相加作为两者要学习的目标，即：

其中 指的是 KL 散度，也即相对熵（其实就是真实分布不是 one-hot 前提下的交叉熵）。

正如上面所述，这个损失函数最小化的目标是两个概率分布对于其分布均值之间的距离。举个具体的例子，比如正样本及其增强分别代表下图中的同一只狗的两种不同的视角：

那么，对于另一个随机采样的样本（比如猫的某个视角），正样本的向量和随机采样样本的向量的相似度，预期上应该是低于正样本的向量和另一张狗的随机采样样本的相似度的；而作为正样本的增强样本，它对于不同的随机样本，也应该具有这样相似的特性。

最后，作者将对比学习的损失函数与关系损失函数加和，得到了联邦学习的步骤 1 中的最终损失函数：

2.2 全局损失

除了局部的损失函数外，作者也引入了全局的损失函数。其出发点为：局部模型不但要从自己独立的样本中来学习分布，自然也要从全局模型中学习信息，所以在全局上提出了与局部损失类似的全局损失函数。唯一的差别在于，全局损失函数的关系损失函数中，随机采样的负样本的表示向量从全局模型中取得（从作者的代码中可以发现，此时全局模型的参数是冻结的，只用于提供特征向量）。

全局模型主要目的是解决局部数据和全局数据的分布不一致所造成的的偏差，从而避免局部模型从上述过程中往局部最优收敛。

在联邦学习的过程中， 步骤 4 将全局聚合模型下载的不同客户端本地，通过该全局模型来试图对自身拥有的局部模型纠偏。值得注意的是，全局模型的结构可能与局部模型不一致，为了使得知识蒸馏能够有效进行，在全局蒸馏过程中，作者在局部模型中添加了一个额外的预测层（在代码中体现为一个两层的全连接，维度均为 ）。

基于此，全局的对比学习损失表示为：

其中上标 表示该样本的特征表示经过了局部模型的额外预测层 ，而上标 表示该样本经过了全局模型的最终预测层（softmax 的前一层输出）。在全局的对比学习损失中，旨在通过全局模型对局部样本进行纠偏。

与此对应，样本对于其它负样本的相对分布也分别有局部模型和全局模型产生：

最终的 FedX 损失函数由局部损失和全局损失加和得到，即：

FedX效果评估

作者在 CIFAR-10，SVHN，F-MNIST 三个数据集上进行了实验，使用 Dirichlet分布模拟联邦学习中各个 client 分得的样本分布，采用了 SimCLR，MoCo，BYOL，ProtoCL，FedU 一共 5 个方法作为基线模型，将各自加入关系损失后的 + FedX 方法作为提出的方法，在三个数据集上分别取得了 4.29pp，5.52pp 和1.58pp 的增长。

作者也进行了一系列的有效性实验，以证明 FedX 在不同场景，不同设定下的有效性。

作者首先通过了不同模块的消融实验来验证模块有效性（本地关系损失函数+全局关系损失函数+一个额外添加的两层全连接层变换）。下图的结果表明关系损失函数、全局蒸馏，以及为了匹配全局蒸馏而增加的额外全连接层都对模型有着积极的影响。而去除全局知识蒸馏过程会使得整个联邦学习最终准确率有着非常大的下降，从而从侧面证明了在局部样本有偏的情况下，不通过全局模型进行纠偏，会大概率使得局部模型陷入局部最优，从而进一步对全局模型产生比较大的负面影响。

在参数敏感性实验中，作者控制了不同的数据量，不同的 client 数量与不同的本地-全局模型交互次数，以探究 FedX 在不同的实验设定下是否会出现不稳定的情况 —— 实验结果表明其在不同的数据量、不同的训练轮次，以及不同的客户端数量下，都能给现有的联邦学习算法带来比较显著的提升。

除此以外，作者还进行了一个比较新颖的 Embed 空间实验：通过引入反余弦函数来量化全局模型与局部模型向量间的分布情况，预期上类内更加相似，类间更加不同。FedX 的表现如预期一样，使得类内的角度偏向更大的夹角。

值得注意的是，在论文的最后，作者也提出了将 FedX 扩展到半监督学习的场景中的设定，在这个场景下，首先用未标注数据对模型进行参数的初始化，然后再在模型上层添加一个全连接层，使用标注数据再额外训练一定的轮次做微调。发现 FedX 在半监督的算法场景下也能带来一定的提升，进一步证明了 FedX 不仅能够用于无监督场景下的联邦学习，也能用于少数据场景下的半监督学习。

更多阅读

#投 稿 通 道#

 让你的文字被更多人看到 

如何才能让更多的优质内容以更短路径到达读者群体，缩短读者寻找优质内容的成本呢？答案就是：你不认识的人。

总有一些你不认识的人，知道你想知道的东西。PaperWeekly 或许可以成为一座桥梁，促使不同背景、不同方向的学者和学术灵感相互碰撞，迸发出更多的可能性。 

PaperWeekly 鼓励高校实验室或个人，在我们的平台上分享各类优质内容，可以是最新论文解读，也可以是学术热点剖析、科研心得或竞赛经验讲解等。我们的目的只有一个，让知识真正流动起来。

📝 稿件基本要求：

• 文章确系个人原创作品，未曾在公开渠道发表，如为其他平台已发表或待发表的文章，请明确标注 

• 稿件建议以 markdown 格式撰写，文中配图以附件形式发送，要求图片清晰，无版权问题

• PaperWeekly 尊重原作者署名权，并将为每篇被采纳的原创首发稿件，提供业内具有竞争力稿酬，具体依据文章阅读量和文章质量阶梯制结算

📬 投稿通道：

• 投稿邮箱：[e-mail] 

• 来稿请备注即时联系方式（微信），以便我们在稿件选用的第一时间联系作者

• 您也可以直接添加小编微信（pwbot02）快速投稿，备注：姓名-投稿

△长按添加PaperWeekly小编

🔍

现在，在「知乎」也能找到我们了

进入知乎首页搜索「PaperWeekly」

点击「关注」订阅我们的专栏吧

·

·