ICLR 2023 | MocoSFL: 低成本跨用户联邦自监督学习

2023-03-14 02:03

©PaperWeekly 原创 · 作者 | 李竞涛

单位 | 亚利桑那州立大学博士生

研究方向 | 联邦学习，机器学习隐私

十分荣幸我们的近期工作被 2023 International Conference on Learning Representations 接收为 Oral Paper（TOP 5% notable）！本文探讨了联邦自监督学习在边缘设备上的解决方案，欢迎更多感兴趣的小伙伴们关注、深挖这一有潜力的方向。

论文标题：

MocoSFL: enabling cross-client collaborative self-supervised learning

论文作者：

Jingtao Li, Lingjuan Lyu, Daisuke Iso, Chaitali Chakrabarti, Michael Spranger

作者单位：

Arizona State University (1st, 4th), **Sony AI** (2nd, 3rd, 5th，通讯单位)

文章链接：

https://openreview.net/forum?id=2QGJXyMNoPz

背景

协作学习中的数据标注难题。协作学习的深度学习方案变得越来越受欢迎，因为用户可以在不共享本地私人数据的情况下训练模型并获得性能提升，从而规避 GDPR 等隐私政策。目前，协作学习主要集中在数据有标注的监督学习应用上。然而在实践中，标记需要专业知识和时间投入，对于普通用户很难执行。

当前的解决方案：联邦自监督学习。联邦学习（FL）是当下最流行的协作学习方案。其中，最具代表性的算法是“FedAvg”，用户端（Client）将其本地模型的副本发送到服务器（Server），服务器执行加权平均操作（权重取决于数据量）以获得新的全局模型。

于是，面对标注数据的难题，zhuang2022ICLR [1] 提出了一种新的联邦自监督学习，使用发散感知聚合技术缓解了非独立同分布数据的性能下降，并取得了 SOTA 的精度。然而，上述方法很好的解决了跨企业数据孤岛（Cross-silo）的问题，但是由于它对于计算资源以及本地数据量的严苛要求，实现跨用户（Cross-client）联邦自监督几乎不可能。

我们认为，脱离原本的 FedAvg 框架并且应用合理的计算卸载（computation offloading）是解决这一问题的关键。基于上述理解和启发，在本文中我们探索了如何通过有机结合 Split Federated Learning（SFL）和 Momentum Contrast（MoCo）对比学习技术，来实现低成本的跨用户联邦自监督学习。

联邦自监督学习的两大要求

2.1 对比学习Moco简介

BYOL、SimCLR 等方案使用当前 batch 中的其他样本作为负键（negative keys）来进行对比学习。与这些方案不同，MoCo（He2020CVPR [2]）使用先前计算的正键输出作为负键，并将它们存储在特征内存中以供将来的迭代使用。这种存储负键的机制使得 MoCo 的 batch size 要求相对较小，有利于减少设备内存，也易于实现特征共享。对于损失函数，MoCo 依靠 InfoNCE 损失作为对比机制：

▲ Fig1. InfoNCE 损失函数

其中查询键 Q 和正键 K+ 是由图像的两个增广视图在目标模型（Online Model）和动量模型（Momemtum Model）的输出向量。N 表示大小为 M 的特征内存中的所有负键。重要的是，MoCo方案的成功在很大程度上取决于其负键的“难度”，这可以通过 （在第 t 次迭代）与 之间的相似度（内积）来确定；相似度越小，其“难度”越高。

我们注意到，负键的“难度”会因为模型更新而迅速降低。因此，MoCo 采用一个变化缓慢的动量模型，在每个训练步骤结束时产生一致（consistent）的负键，来维护负键的“难度”。

2.2 基于模型分割的联邦学习简介

Split Federated Learning（SFL, Thapa2022AAAI [3]）是一种注重用户端高计算效率的新型协作学习方案，结合了计算卸载和联邦学习 FedAvg。它以“割裂层”为界，将原始模型架构分为两部分：包含“割裂层”及其之前所有层的用户端模型和包含其余层的服务器端模型。

我们将用户端模型的副本分发到用户端的本地设备上，并在云服务器上实例化服务器端模型。为完成每个训练步骤，用户端需要将潜向量（latent vector），即用户端模型的输出，发送到服务器，然后由服务器端模型处理潜向量，计算损失并执行反向传播，并将相应的梯度向量返回给用户端。经过一定数量的训练步骤后，所有用户端模型将会由 Server 进行类似 FedAvg 的同步。

2.3 对于计算资源的要求

本地训练自监督模型需要庞大的计算资源。传统的监督学习可以采用精简小模型或剪枝方法来减少计算资源，但这种方法不适合自监督学习，因为它对模型容量（capacity）的要求更高。

先前的一些研究表明，像 Mobilenet-V3、EfficientNet 等紧凑的架构与较大的 ResNet-18 架构相比存在超过 10% 的精度下降，而使用更大的 ResNet-50 模型相比于 ResNet-18 在 ImageNet 数据集上的精度提高了超过 15%。这意味着要训练一个高精度的自监督模型，需要模型很大，对内存要求非常高。

另一方面，为了减小内存，使用较小的批大小通常是一个很好的解决方案。但是，如图 Fig2(a) 所示，在自监督学习中，使用更小的批大小会大大降低准确度。因此，这并不是一个好的选择。

▲ Fig2. 联邦自监督学习的严苛要求

2.4 对于本地数据量的要求

联邦自监督学习在跨用户应用中面临的另一个主要困难是本地数据的要求。对于跨用户应用，每个用户可用的数据量可能非常有限。例如，在跨企业数据孤岛的应用场景中，一个用户端可以是拥有大量数据的医院。相比之下，在跨用户应用中，一个用户端可能只是一位数据量有限的患者。

问题在于，当用户端没有足够的本地数据时，很难找到有“难度”的负键。而当数据量更大时，出现有“难度”的负键的机会会大得多。因此，现有的联邦自监督学习只能在跨企业数据孤岛应用中取得成功，因为用户端拥有大量数据并且可以在本地执行有效的对比学习。如 Fig2 (b) 所示，当用户端拥有 10K 个数据样本时，我们观察到高精度；而当用户端只有 500 个样本时，精度会迅速降至约 30%。

本文提出的框架（MocoSFL）

▲ Fig3. MocoSFL框架

本文提出的 MocoSFL 是 SFL 和 MoCo-V2 的有机结合，如 Fig.3 所示。MocoSFL 有三个关键组件。

第一，在每个训练步骤中，我们将所有用户端发送的 latent vector 合并起来，视作一个大 batch，由服务器端模型一并处理，以支持用户端小批量训练。第二，我们使用一个共享的特征存储器，在每个训练步骤中由所有用户端提供的正键来更新，以达成负键的共享。第三，我们通过使用更高的同步频率来减小用户端模型的差异，以提高模型在非独立同分布用户数据上的性能。以下是对 MocoSF L特性的总结：

1. 较小的本地计算，内存要求（SFL）

2. 等效大 Batch 训练，进一步节省内存（SFL）

3. 降低同步的通信成本，从而可以提高同步频率（SFL）

4. 存储负键，便于用户间共享（MoCo）

5. 对 Batch Size 要求较低（MoCo）

（详情论证请见文章第四章节）

实验结果

首先，我们在 CIFAR-10 和 CIFAR-100 上验证了本文在非独立同分布跨孤岛应用下的精度，并超越了SOTA：

▲ Fig4. MocoSFL在跨孤岛应用场景下的表现以及和SOTA的对比

我们也在这两个数据集以及 ImageNet 子集上验证了在跨用户应用下的精度（SOTA 精度很低, 详见 Fig 2 (b)，因此不提供）：

▲ Fig5. MocoSFL在跨用户应用场景下的表现

鉴于分割联邦学习的 latent vector 的通信负担和数据安全隐患（这里主要讨论 Model Inversion Attack），我们通过提出 TAResSFL 模组，对 MocoSFL 进行了优化。优化后攻击的 MSE 相对于优化前有了大概 10 倍的提升，Inversion 攻击提取的原始图像也变得相当模糊：