（Model-Contrastive Federated Learning）模型对比联邦学习

Posted 2021-08-14 AI浩

摘要

联邦学习使多方能够在不交流本地数据的情况下协作训练机器学习模型。 联邦学习的一个关键挑战是处理各方本地数据分布的异质性。 尽管已经提出了许多研究来应对这一挑战，但我们发现它们无法在具有深度学习模型的图像数据集中实现高性能。 在本文中，我们提出了 MOON：模型对比联邦学习。 MOON 是一个简单而 有效的联邦学习框架。 MOON 的关键思想是利用模型表示之间的相似性来纠正各个方的局部训练，即在模型级别进行对比学习。 我们的大量实验表明，MOON 在各种图像分类任务上明显优于其他最先进的联邦学习算法。

1、简介

深度学习需要数据。模型训练可以从大型且具有代表性的数据集（例如 ImageNet [6] 和 COCO [31]）中受益很多。然而，数据在实践中通常分散在不同的各方（例如，移动设备和公司）之间。由于越来越多的隐私问题和数据保护法规 [40]，各方无法将他们的私人数据发送到中央服务器来训练模型。 为了解决上述挑战，联邦学习 [20, 44, 27, 26] 使多方能够在不交换本地数据的情况下共同学习机器学习模型。 一种流行的联邦学习算法是 FedAvg [34]。在每一轮 FedAvg 中，更新各方的本地模型被传送到服务器，服务器进一步聚合本地模型以更新全局模型。在学习过程中不交换原始数据。联邦学习已经成为一个重要的机器学习领域并吸引了许多研究兴趣 [34, 28, 22, 25, 41, 5, 16, 2, 11]。此外，它已应用于许多应用，例如医学成像 [21, 23]、物体检测 [32] 和地标分类 [15]。 联邦学习的一个关键挑战是不同方数据分布的异质性 [20]。在许多实际应用中，数据可能会在各方之间以不同的方式分布，这会降低联邦学习的性能 [22, 29, 24]。当各方更新其本地模型时，其本地目标可能与全局目标相距甚远。因此，平均全局模型远离全局最优。有一些研究试图解决本地训练阶段的非 IID 问题 [28, 22]。 FedProx [28] 直接通过“2-范数距离”限制局部更新，而 SCAFFOLD [22] 通过方差减少 [19] 校正局部更新。 然而，正如我们在实验中所展示的（参见第 4 节），这些方法无法在具有深度学习模型的图像数据集上实现良好的性能，这可能与 FedAvg 一样糟糕。
在这项工作中，我们基于直观观察从一个新颖的角度解决了非 IID 问题：在整个数据集上训练的全局模型能够学习比在偏斜子集上训练的局部模型更好的表示。 具体来说，我们提出了模型对比学习（MOON），它通过最大化当前局部模型学习的表示与全局模型学习的表示的一致性来纠正局部更新。 与传统的对比学习方法 [3, 4, 12, 35] 通过比较不同图像的表示在学习视觉表示方面达到最先进的结果不同，MOON 通过比较表示在模型级别进行对比学习 通过不同的模型学习。 总的来说，MOON 是一个简单有效的联邦学习框架，并且 通过基于模型的对比学习的新颖设计解决了非 IID 数据问题。
我们进行了大量实验来评估 MOON 的有效性。 MOON 在包括 CIFAR-10、CIFAR-100 和 Tiny-Imagenet 在内的各种图像分类数据集上明显优于其他最先进的联邦学习算法 [34、28、22]。 仅对 FedAvg 进行轻量级修改，MOON 在大多数情况下比现有方法的准确率至少高出 2%。 此外，在某些设置上，MOON 的改进非常显着。 例如，在 100 方的 CIFAR-100 数据集上，MOON 达到了 61.8% 的 top-1 准确率，而现有研究的最佳 top-1 准确率是 55%。

2、背景及相关工作

2.1.联邦学习

FedAvg [34] 一直是联邦学习的一种事实上的方法。 FedAvg的框架如图1所示，每轮FedAvg有四个步骤。首先，服务器向各方发送一个全局模型。其次，双方执行随机梯度下降 (SGD) 以在本地更新他们的模型。第三，本地模型被发送到中央服务器。最后，服务器对模型权重进行平均以生成用于下一轮训练的全局模型。 有相当多的研究试图改进非 IID 数据的 FedAvg。这些研究可以分为两类：局部训练的改进（即图 1 的步骤 2）和聚合（即图 1 的步骤 4）。 本研究属于第一类。至于改进本地培训的研究，FedProx [28] 在本地培训期间将近端项引入目标。近端项是根据当前全局模型和当前全局模型之间的“2-范数距离”计算得出的。 本地模式。因此，局部模型更新在局部训练期间受到近端项的限制。 SCAFFOLD [22] 通过引入控制变量来纠正局部更新。与训练模型一样，控制变量也在局部训练期间由各方更新。局部控制变量和全局控制变量之间的差异用于校正局部训练中的梯度。然而，FedProx 仅展示了多项逻辑回归在 MNIST 和 EMNIST 上的实验，而 SCAFFOLD 仅展示了带有逻辑回归和 2 层全连接层的 EMNIST 实验。 FedProx 和 SCAFFOLD 在具有深度学习模型的图像数据集上的有效性尚未得到很好的探索。正如我们的实验所示，这些研究与 FedAvg 相比几乎没有优势，甚至没有优势，这促使这项研究采用一种新的方法来处理具有深度学习模型的非 IID 图像数据集。我们还注意到，在准备这篇论文时，还有其他相关的当代工作 [1, 30, 43]。我们将 MOON 与这些当代作品之间的比较留作未来研究。 对于改进聚合阶段的研究，FedMA [41] 利用贝叶斯非参数方法以分层方式匹配和平均权重。 FedAvgM [14] 在更新服务器上的全局模型时应用动量。最近的另一项研究 FedNova [42] 在平均之前对本地更新进行了标准化。我们的研究与它们是正交的，当我们在本地训练阶段工作时，可能会与这些技术相结合。另一个研究方向是个性化联邦学习 [8, 7, 10, 47, 17]，它试图为每一方学习个性化的本地模型。在本文中，我们研究了典型的联邦学习，它试图为所有各方学习一个单一的全局模型。

2.2.对比学习

自监督学习 [18, 9, 3, 4, 12, 35] 是最近的一个热门研究方向，它试图从未标记的数据中学习良好的数据表示。在这些研究中，对比学习方法 [3, 4, 12, 35] 在学习视觉表示方面取得了最先进的结果。对比学习的关键思想是减少同一图像的不同增强视图（即正对）的表示之间的距离，并增加不同图像（即，负对）的增强视图的表示之间的距离。一个典型的对比学习框架是 SimCLR [3]。给定图像 x，SimCLR 首先使用不同的数据增强算子创建该图像的两个相关视图，表示为$x_i$和$x_j$。训练基本编码器 f(·) 和投影头 g(·) 分别提取表示向量并将表示映射到潜在空间。然后，对比损失（即 NT-Xent [38]）应用于投影向量 g(f(·))，它试图最大化同一图像的不同增强视图之间的一致性。具体来说，给定 2N 个增强视图和同一图像的一对视图 $x_i$和$x_j$，这对的对比损失定义为

其中 sim(·;·) 是余弦相似度函数，$\tau$ 表示温度参数。最终的损失是通过对小批量中相同图像的所有对的对比损失求和来计算的。 除了 SimCLR，还有其他对比学习框架，如 CPC [36]、CMC [39] 和 MoCo [12]。我们选择 SimCLR 是因为它在许多计算机视觉任务中的简单性和有效性。尽管如此，对比学习的基本思想在这些研究中是相似的：从不同图像获得的表征应该彼此远离，从同一图像获得的表征应该彼此相关。这个想法很直观，并且已被证明是有效的。 最近的一项研究 [46] 将联邦学习与对比学习相结合。他们专注于无监督学习环境。与 SimCLR 一样，他们使用对比损失来比较不同图像的表示。在本文中，我们专注于监督学习设置并提出模型对比学习来比较不同模型学习的表示。

3. 模型对比联邦学习

3.1. 问题陈述

假设有 N 个参与方，记为$P_1$… $P_N$ 。 Party $P_i$ 有一个本地数据集$D^i$。 我们的目标是在中央服务器的帮助下，在数据集 D, Si2[N] Di 上学习机器学习模型 w，同时不交换原始数据。 目标是解决

其中 Li(w) = E(x;y)∼Di[‘i(w; (x; y))] 是 Pi 的经验损失。

3.2.动机

MOON 基于一个直观的想法：在整个数据集上训练的模型能够提取比在偏斜子集上训练的模型更好的特征表示。例如，给定一个在狗和猫图像上训练的模型，我们不能期望模型学习到的特征能够区分训练期间从未存在过的鸟类和青蛙。 为了进一步验证这种直觉，我们对 CIFAR-10 进行了一个简单的实验。具体来说，我们首先在 CIFAR-10 上训练一个 CNN 模型（详细结构见第 4.1 节）。我们使用 t-SNE [33] 来可视化来自测试数据集的图像的隐藏向量，如图 2a 所示。然后，我们以不平衡的方式将数据集划分为 10 个子集（有关划分策略，请参见第 4.1 节）并在每个子集上训练一个 CNN 模型。图 2b 显示了 t-SNE 随机选择的模型的可视化。显然，在子集上训练的模型学习的特征很差。大多数类的特征表示甚至是混合的，无法区分。然后，我们在 10 个子集上运行 FedAvg 算法，并在图 2c 中显示全局模型学习的表示和图 2d 中选择的局部模型（基于全局模型训练）学习的表示。我们可以观察到，与图 2c 相比，图 2d 中具有相同类别的点更加发散（例如，参见第 9 类）。由于偏斜的局部数据分布，局部训练阶段甚至导致模型学习到更糟糕的表示。这进一步验证了全局模型应该能够比局部模型学习更好的特征表示，并且局部更新存在漂移。因此，在非 IID 数据场景下，我们应该控制漂移并弥合局部模型和全局模型学习的表示之间的差距。

3.3.方法

基于上述直觉，我们提出了 MOON。 MOON 被设计为一种基于 FedAvg 的简单有效的方法，仅在本地训练阶段引入轻量级但新颖的修改。由于局部训练总是存在漂移，全局模型学习到的表示比局部模型更好，MOON旨在减少局部模型学习到的表示与全局模型学习的表示之间的距离，并增加与局部模型学习的表示之间的距离。本地模型学习到的表示和之前的本地模型学习到的表示。我们从对比学习的灵感中实现了这一点，对比学习现在主要用于学习视觉表示。下面，我们介绍网络架构、本地学习目标和学习过程。最后，我们讨论了对比学习的关系。

3.3.1 网络架构

该网络由三个组件组成：基本编码器、投影头和输出层。 基本编码器用于从输入中提取表示向量。 像[3]一样，引入了一个额外的投影头来将表示映射到具有固定维度的空间。 最后，当我们研究监督设置时，输出层用于为每个类生成预测值。 为了便于表示，使用模型权重 w，我们使用 Fw(·) 表示整个网络，使用 Rw(·) 表示输出层之前的网络（即，Rw(x) 是输入 x 的映射表示向量）

3.3.2 局部目标

如图 3 所示，我们的局部损失由两部分组成。 第一部分是监督学习中的典型损失项（例如交叉熵损失），表示为“sup”。 第二部分是我们提出的模型对比损失项，表示为“con”。 假设派派正在进行本地培训。 它从服务器接收全局模型 wt 并在本地训练阶段将模型更新为 wt i。 对于每个输入 x，我们从全局模型 wt（即 zglob = Rwt(x)）中提取 x 的表示，从上一轮 wt−1 i 的局部模型中提取 x 的表示（即 zprev = Rwt−1 i (x))，以及来自正在更新的本地模型的 x 表示 w i（即 z = Rwt i (x)）。 由于全局模型应该能够提取更好的表示，我们的目标是减少 z 和 zglob 之间的距离，并增加 z 和 z prev 之间的距离。 与 NT-Xent 损失 [38] 类似，我们将模型对比损失定义为

其中 τ 表示温度参数。 aninput (x; y) 的损失由下式计算

其中 µ 是控制模型对比损失权重的超参数。 局部目标是最小化

整体的联邦学习算法如算法1所示。在每一轮中，服务器向各方发送全局模型，从各方接收本地模型，并使用加权平均更新全局模型。 在局部训练中，每一方使用随机梯度下降用其局部数据更新全局模型，而目标在等式中定义。 (5).

为简单起见，我们在算法 1 中不应用采样技术来描述 MOON。当每轮只有一组样本方参与联邦学习时，MOON 仍然适用。 与 FedAvg 一样，每一方都维护其本地模型，该模型将被全局模型取代，并且只有在该方被选中参与一轮时才会更新。 MOON 只需要该方拥有的最新本地模型，即使它可能不会在第 (t-1) 轮中更新（例如，wt-1 i = wit-2）。
值得注意的是，考虑到理想情况下局部模型足够好并且学习（几乎）与全局模型相同的表示（即 zglob = zprev），模型对比损失将是一个常数（即 - log 1 2)。 因此，MOON 将产生与 FedAvg 相同的结果，因为不存在异质性问题。 从这个意义上说，无论漂移量如何，我们的方法都是稳健的。

3.4. 与对比学习的比较

MOON 和 SimCLR 之间的比较如图 4 所示。模型对比损失比较了不同模型学习的表示，而对比损失比较了不同图像的表示。 我们还强调了 MOON 与传统对比学习之间的主要区别：MOON 目前用于联邦环境中的监督学习，而对比学习用于集中环境中的无监督学习。 从对比学习中汲取灵感，MOON 是一种新的学习方法，用于处理联邦学习中各方之间的非 IID 数据分布。

4. 实验

4.1. 实验装置

我们将 MOON 与三种最先进的方法进行比较，包括 (1) FedAvg [34]、(2) FedProx [28] 和 (3) SCAFFOLD [22]。我们还比较了一种名为 SOLO 的基线方法，其中每一方都使用其本地数据训练模型，而无需联合学习。我们在三个数据集上进行了实验，包括 CIFAR-10、CIFAR-100 和 Tiny-Imagenet1（100,000 张图像，200 个类别）。此外，我们尝试了两种不同的网络架构。对于 CIFAR-10，我们使用 CNN 网络作为基本编码器，它有两个 5x5 卷积层，然后是 2x2 最大池化（第一个具有 6 个通道，第二个具有 16 个通道）和两个具有 ReLU 激活的全连接层（第一个120 个单位，第二个单位为 84 个单位）。对于 CIFAR-100 和 Tiny-Imagenet，我们使用 ResNet-50 [13] 作为基本编码器。对于所有数据集，如 [3]，我们使用 2 层 MLP 作为投影头。投影头的输出尺寸默认设置为256。请注意，为了公平比较，所有基线都使用与 MOON 相同的网络架构（包括投影头）。

我们使用 PyTorch [37] 来实现 MOON 和其他基线。该代码是公开可用的2。我们对所有方法使用学习率为 0.01 的 SGD 优化器。 SGD 权重衰减设置为 0.00001，SGD 动量设置为 0.9。批量大小设置为 64。 SOLO 的本地 epoch 数设置为 300。除非明确指定，否则所有联合学习方法的本地 epoch 数都设置为 10。 CIFAR-10/100 的通信轮数设置为 100，Tiny-ImageNet 的通信轮数设置为 20，其中所有联合学习方法在更多通信时几乎没有或没有准确度增益。对于 MOON，我们将温度参数默认设置为 0.5，如 [3]。
与之前的研究 [45, 41] 一样，我们使用 Dirichlet 分布来生成各方之间的非 IID 数据分区。具体来说，我们对 pk ∼ DirN(β) 进行采样，并将 k 类实例的 pk;j 比例分配给第 j 方，其中 Dir(β) 是具有浓度参数 β（默认为 0.5）的狄利克雷分布。通过上述划分策略，每一方在某些类别中可以拥有相对较少（甚至没有）的数据样本。默认情况下，我们将参与方的数量设置为 10。默认设置下各方的数据分布如图5所示，更多实验结果请参考附录。

4.2. 精度比较

对于 MOON，我们从 f0:1 调整 µ； 1; 5; 10g并报告最佳结果。 CIFAR-10、CIFAR-100 和 Tiny-Imagenet 的最佳 MOON µ 分别为 5、1 和 1。 请注意，FedProx 还有一个超参数 µ 来控制其近端项的权重（即 LF edP rox = ‘F edAvg + µ‘prox）。 对于 FedProx，我们从 f0:001 调整 µ； 0:01; 0:1; 1g（范围也在之前的论文[28]中使用）并报告最佳结果。 用于 CIFAR-10、CIFAR-100 和 Tiny-Imagenet 的 FedProx 的最佳 µ 分别为 0:01、0:001 和 0:001。 除非明确指定，否则我们将这些 µ 设置用于所有剩余的实验。

表 1 显示了具有上述默认设置的所有方法的 top-1 测试精度。 在非 IID 设置下，SOLO 表现出比其他联邦学习方法差得多的准确性。 这证明了联邦学习的好处。 比较不同的联邦学习方法，我们可以观察到 MOON 始终是所有任务中最好的方法。 它可以在所有任务的平均准确率上超过 FedAvg 2.6%。 对于 FedProx，其准确性非常接近 FedAvg。 FedProx 中的近端项对训练几乎没有影响，因为 µ 很小。 但是，当 µ 没有设置为非常小的值时，FedProx 的收敛速度非常慢（参见第 4.3 节），并且 FedProx 的准确性很差。 对于 SCAFFOLD，它在 CIFAR-100 和 Tiny-Imagenet 上的准确度比其他联邦学习方法要差得多。

4.3.沟通效率

图 6 显示了训练期间每一轮的准确率。正如我们所看到的，模型对比损失项对最佳 μ 的收敛速度几乎没有影响。 MOON 的准确率提升速度与开始时的 FedAvg 几乎相同，但后期可以从模型对比损失中获得更好的准确率。由于 FedProx 中的最佳 µ 值通常较小，因此具有最佳 µ 的 FedProx 非常接近 FedAvg，尤其是在 CIFAR-10 和 CIFAR-100 上。然而，当设置 µ = 1 时，由于附加的近端项，FedProx 变得非常慢。这意味着限制局部模型和全局模型之间的“2-范数距离”不是有效的解决方案。我们的模型对比损失可以在不减慢收敛速度的情况下有效提高准确性。

我们在表 2 中显示了与在 CIFAR-10/100 上运行 FedAvg 100 轮或在 Tiny-Imagenet 上运行 20 轮相同精度的通信轮数。我们可以观察到，在 MOON 中通信轮数显着减少。月亮
与 FedAvg 相比，在 CIFAR-100 和 Tiny-Imagenet 上需要大约一半的通信轮数。在 CIFAR-10 上，MOON 的加速比甚至接近 4。MOON 的通信效率比其他方法高得多。

4.4.本地时代的数量

我们研究了局部时期数对最终模型准确性的影响。结果如图7所示，当local epochs为1时，local update很小。因此，在通信轮数相同的情况下，训练速度较慢，准确度相对较低。所有方法都具有接近的准确性（MOON 仍然是最好的）。当局部 epoch 数变得太大时，所有方法的准确度都会下降，这是由于局部更新的漂移，即局部最优与全局最优不一致。尽管如此，MOON 显然优于其他方法。这进一步验证了 MOON 可以有效减轻过多局部更新带来的漂移的负面影响。

4.5.可扩展性

为了展示 MOON 的可扩展性，我们在 CIFAR-100 上尝试了更多的参与方。具体来说，我们尝试了两种设置：（1）我们将数据集划分为 50 方，每一方都参与联邦学习。 (2) 我们将数据集划分为 100 方，每轮随机抽取 20 方参与联邦学习（FedAvg [34] 中引入的客户端采样技术）。这
结果显示在表 3 和图 8 中。对于 MOON，我们显示了 µ = 1（第 4.2 节中的最佳 µ）和 µ = 10 的结果。对于 MOON (µ = 1)，它的性能优于 FedAvg 和 FedProx 超过 2% 50 方的 200 轮准确率和 100 方的 500 轮准确率 3%。此外，对于 MOON (μ = 10)，虽然大的模型对比损失在开始时减慢了训练速度，如图 8 所示，但 MOON 可以通过更多的通信轮次远远优于其他方法。与 FedAvg 和 FedProx 相比，MOON 在 50 方的 200 轮和 100 方的 500 轮时的准确率提高了约 7%。 SCAFFOLD 的准确性较低，参与方相对较多。

4.6. 异质性

我们通过改变 Dirichlet 分布的浓度参数 β 来研究数据异质性对 CIFAR-100 的影响。 对于较小的β，分区将更加不平衡。 结果如表 4 所示。 MOON 在三个不平衡级别中始终达到最佳精度。当不平衡级别降低时（即 β = 5），FedProx 比 FedAvg 差，而 MOON 仍然以 2% 以上的精度优于 FedAvg。 实验证明了 MOON 的有效性和鲁棒性。

4.7.损失函数

为了最大化全局模型学习到的表示与局部模型学习到的表示之间的一致性，我们的模型对比损失’con 是受 NT-Xent 损失 [3] 的启发而提出的。另一个直观的选择是使用‘2正则化，局部损失为
‘ = ‘
sup + µ kz − zglobk2 (6)
在这里，我们比较了使用不同类型的损失函数来限制表示的方法：没有附加项（即 FedAvg：L = ‘sup）、‘2 范数和我们的模型对比损失。结果如表 5 所示。我们可以观察到，与 CIFAR-10 上的 FedAvg 相比，仅使用 2 范数甚至无法提高准确性。使用‘2范数可以提高在 CIFAR-100 和
Tiny-Imagenet，准确率还是低于MOON。我们的模型对比损失是约束表示的有效方法。
我们的模型对比损失从两个方面影响局部模型。首先，局部模型学习到全局模型的近似表示。其次，局部模型也学习到比前一个更好的表示，直到局部模型足够好（即 z = zglob 并且“con”成为一个常数）。

5. 结论

联邦学习已成为解决医学成像、对象检测和地标分类等许多领域中数据孤岛难题的一种很有前途的方法。非 IID 是联邦学习有效性的关键挑战。为了提高联合深度学习模型在非 IID 数据集上的性能，我们提出了模型对比学习 (MOON)，这是一种简单有效的联合学习方法。 MOON 引入了一个新的学习概念，即模型级的对比学习。我们广泛的实验表明，在各种图像分类任务上，MOON 比最先进的方法取得了显着的改进。由于 MOON 不要求输入为图像，因此它可能适用于非视觉问题。

致谢

这项研究由新加坡国家研究基金会根据其人工智能新加坡计划（AISG 奖号：AISG2-RP-2020-018）提供支持。本材料中表达的任何意见、发现和结论或建议均为作者的观点，并不反映新加坡国家研究基金会的观点。感谢Jianxin Wu、Chaoyang He、Shixuan Sun、Yaqi Xie 和 Yuhang Chen 的反馈。作者还要感谢赵玉志、王伟和莫莎对计算资源的支持。