自监督模型---MoCo

论文地址：Momentum Contrast for Unsupervised Visual Representation Learning

开源代码：PyTorch implementation of MoCo

摘要

该论文提出了动量对比度(MoCo)用于无监督的视觉表示学习。构建了一个带有一个队列和一个移动平均编码器的动态字典。这使得建立一个大的和一致的字典的动态，促进对比无监督学习。MoCo在ImageNet分类的通用线性协议下提供了有竞争的结果。更重要的是，MoCo学习到的表示可以很好地转移到下游任务中。MoCo可以在PASCAL VOC、COCO和其他数据集上的7个检测/分割任务中优于其有监督的训练，有时会大大超过它。这表明，在许多视觉任务中，无监督和监督表征学习之间的差距在很大程度上被缩小。

介绍

无监督表示学习在自然语言处理中非常成功，例如，如GPT和BERT所示。但是有监督的预训练在计算机视觉中仍然占主导地位，因为无监督的方法通常会落后。其原因可能是由于它们各自的信号空间的不同。语言任务具有离散的信号空间（单词、子词单位等）。用于构建标记化字典，它可以基于无监督学习。相比之下，计算机视觉进一步关注字典构建，因为原始信号是在一个连续的高维空间中，并不是为人类交流而构建的结构（例如，不像单词）。

最近的几项研究在使用与对比损失相关的方法进行无监督视觉表征学习方面提供了有希望的结果。尽管受到各种动机的驱动，但这些方法可以被认为是构建动态字典。字典中的“键”（tokens）从数据（例如，图像或patches）中采样，并由编码器网络表示。无监督学习训练编码器来执行字典查找：一个编码的“查询”应该与它的匹配键相似，而与其他键不同。学习被表述为最小化对比损失。

从这个角度来看，我们假设构建字典：(i)大的，(ii)在训练期间随着改变。直观地说，一个更大的字典可以更好地采样底层的连续的、高维的视觉空间，而字典中的键应该由相同或相似的编码器表示，以便它们与查询的比较是一致的。 然而，使用对比损失的现有方法在这两个方面之一被限制。

我们提出了动量对比(MoCo)，作为一种构建大型和一致的学习字典的方法，用于具有对比损失的无监督学习（图1）。我们将字典维护为数据样本队列：当前小批量的编码表示入队，最旧的出队。 队列将字典大小与小批量大小解耦，允许它很大。 此外，由于字典键来自前面的几个小批量，因此提出了一种缓慢进展的键编码器，实现为查询编码器的基于动量的移动平均值，以保持一致性。

动量对比度(MoCo)通过使用对比度损失将编码的查询q与编码键的字典匹配来训练视觉表示编码器。字典键{k0、k1、k2、……}是由一组数据样本动态定义的。字典被构建为一个队列，当前的小批进入队列，最古老的小批移除队列，将其与小批量大小解耦。keys由一个缓慢发展的编码器编码，由查询编码器的动量更新驱动。这种方法可以使用一个大而一致的字典来学习视觉表示。

MoCo 是一种为对比学习构建动态词典的机制，可用于各种借口任务(pretext tasks)。在本文中，我们遵循一个简单的实例判别任务 ：如果它们是同一图像的编码视图（例如，不同的裁剪），则查询匹配一个键。 使用这个借口任务，MoCo在ImageNet 数据集中的线性分类通用协议下显示了有竞争力的结果。 

无监督学习的一个主要目的是预训练可以通过微调转移到下游任务的表示（即特征）。 我们表明，在与检测或分割相关的 7 个下游任务中，MoCo 无监督预训练可以超越其 ImageNet 监督对应物，在某些情况下，差距不小。 在这些实验中，我们探索了在 ImageNet 或 10 亿张 Instagram 图像集上预训练的 MoCo，证明 MoCo 可以在更真实的、10 亿张图像规模和相对未经处理的场景中很好地工作。 这些结果表明，MoCo 在很大程度上缩小了许多计算机视觉任务中无监督和有监督表示学习之间的差距，并且可以在多个应用中作为 ImageNet 监督预训练的替代方案。

方法

对比学习作为字典查找

对比学习及其最近的发展，可以被认为是为字典查找任务训练编码器，如下所述。考虑一个编码查询 q 和一组编码样本 {k0, k1,k2, ...}（即字典的键）。

考虑一个编码的查询q和一组编码的样本{k0k1k2…}，它们是字典的键。假设在字典中有一个q匹配的键(记为k+)。对比损失是一个函数，当q与正键k+相似，而与所有其他键不同(被认为是q的负键)时，它的值较低。利用点积度量相似性，本文考虑了一种对比损失函数的形式，称为InfoNCE：

import torch.nn as nn

criterion = nn.CrossEntropyLoss()


output, target = model(im_q=images[0], im_k=images[1])
loss = criterion(output, target)

 其中，τ是每个的温度超参数。这个总和超过一个正样本和K个负样本。直观地说，这种损失是一个(k+1)方式的基于softmax的分类器的对数损失，该分类器试图将q分类为k+。对比损失函数也可以基于其他形式的，如基于保证margin-based的损失和NCE损失的变体。

对比损失作为一个无监督的目标函数来训练表示查询和关键[29]的编码器网络。一般来说，查询表示是，其中fq是一个编码器网络，xq是一个查询样本(同样，)。它们的实例化取决于具体的借口任务。输入的xq和xk可以是图像、补丁，或包含一组补丁的上下文。网络fq和fk可以是相同的，部分共享的，或不同的。

 动量对比

 从上述角度来看，对比学习是一种在图像等高维连续输入上构建离散字典的一种方法。字典是动态的，因为keys是随机采样的，并且keys编码器在训练过程中优化。我们的假设是，好的特征可以通过一个包含大量负样本的大型字典来学习，而字典键的编码器尽管还在优化，但仍尽可能地保持一致。基于这个动机，我们将呈现出下面所描述的动量对比。

字典作为队列。我们方法的核心是将字典作为数据样本队列。这允许我们重用前面的小批中的编码keys。队列的引入可以将字典大小与小批处理大小解耦。我们的字典大小可以比一个典型的小批量大小大得多，并且可以灵活地、独立地设置为一个超参数。

字典中的样本逐渐被替换。当前的小批将被排队到字典中，队列中最老的小批将被删除。字典总是表示所有数据的一个采样子集，而维护此字典的额外计算是可管理的。此外，删除最古老的小批处理可能是有益的，因为它的key编码器是最过时的，因此与最新的key编码器最不一致。(最新的momentum encoder 参数通过学习已经优化了，而最开始得到的keys其momentum encoder 是没有优化或者说次优化的)。

动量更新。使用队列可以使字典变大，但它也使通过反向传播更新键编码器变得困难（梯度应该传播到队列中的所有样本）。一个 “na¨ıve”解决方案是从查询编码器fq复制key编码器fk(即两个网络参数共享)，忽略这个梯度。但是这个解决方案在实验中产生的结果很差。我们假设，这种失败是由快速变化的编码器，减少了key 表示的一致性。我们建议动量更新来解决这个问题。形式上，fk参数为θk，fq参数为θq，我们更新θk：

 这里的m∈[0,1)是一个动量系数。只有参数θq才会通过反向传播进行更新。源码里面使用的是两个相同的编码器:

@torch.no_grad()
 def _momentum_update_key_encoder(self):
     """
     Momentum update of the key encoder
     """
     for param_q, param_k in zip(self.encoder_q.parameters(), self.encoder_k.parameters()):
         param_k.data = param_k.data * self.m + param_q.data * (1. - self.m)

动量更新使θk比θq进化得更顺利。因此，尽管队列中的key由不同的编码器编码，但这些编码器之间的差异可以很小。在实验中，一个相对较大的动量(例如，m=0.999，我们的默认值)比一个较小的值(例如，m=0.9)要好得多，这表明一个缓慢优化的key编码器(slowly evolving key encoder)是利用队列的核心。

与以前的机制之间的关系。MoCo是使用对比损失的一般机制。我们将其与图2中两种现有的通用机制进行了比较。它们在字典的大小和一致性上表现出不同的属性。

三种对比损失机制的概念性比较。 这里我们将举例说明一对查询和密钥。这三种机制在如何维护密钥和如何更新密钥编码器方面有所不同。 (a): 用于计算查询和密钥表示的编码器通过反向传播进行端到端进行更新（这两个编码器可以有所不同）。 (b): 密钥表示法从内存库中采样。 (c): MoCo通过一个动量更新的编码器动态地对新键进行编码，并维护一个键的队列（图中没有说明）。

通过反向传播进行的端到端更新是一种自然的机制(图2a)。它使用当前小批处理中的示例作为字典，因此keys被一致地编码（由相同的一组编码器参数编码）。但是字典的大小与小批量的大小相结合，受到GPU内存大小的限制。它也受到了大型小批量优化的挑战。最近的一些方法是基于由局部位置驱动的借口任务，其中多个位置可以使字典的大小变大。但是这些借口任务可能需要特殊的网络设计，如细化输入或定制接受场大小，这可能会使这些网络向下游任务的转移复杂化。

另一种机制是提出的内存库方法(图2b)。内存库由数据集中所有样本的表示组成。每个小批处理的字典都是从内存库中随机抽样的，没有反向传播，因此它可以支持一个大的字典大小。但是，存储库中样本的表示在上次看到时已更新，因此采样的keys本质上是关于过去 epoch 中多个不同步骤的编码器，因此不太一致。 有些方法的记忆库上采用了动量更新。 它的动量更新是在同一样本的表示上，而不是在编码器上。 这种动量更新与我们的方法无关，因为 MoCo 不会跟踪每个样本。 此外，我们的方法更节省内存，并且可以在十亿规模的数据上进行训练，这对于内存库来说是难以处理的。

# f_q, f_k: encoder networks for query and key
# queue: dictionary as a queue of K keys (CxK) # m: momentum
# t: temperature
f_k.params = f_q.params # initialize
for x in loader: # load a minibatch x with N samples
    x_q = aug(x) # a randomly augmented version
    x_k = aug(x) # another randomly augmented version

    q = f_q.forward(x_q) # queries: NxC
    k = f_k.forward(x_k) # keys: NxC
    k = k.detach() # no gradient to keys

    # positive logits: Nx1
    l_pos = bmm(q.view(N,1,C), k.view(N,C,1))

    # negative logits: NxK
    l_neg = mm(q.view(N,C), queue.view(C,K))

    # logits: Nx(1+K)
    logits = cat([l_pos, l_neg], dim=1)
    # contrastive loss, Eqn.(1)
    labels = zeros(N) # positives are the 0-th
    loss = CrossEntropyLoss(logits/t, labels)

    # SGD update: query network
    loss.backward()
    update(f_q.params)

    # momentum update: key network
    f_k.params = m*f_k.params+(1-m)*f_q.params

    # update dictionary
    enqueue(queue, k) # enqueue the current minibatch
    dequeue(queue) # dequeue the earliest minibatch

借口任务(Pretext Task)

对比学习可以驱动各种借口任务。由于本文的重点不是设计一个新的借口任务，我们主要在中的实例识别任务之后使用一个简单的借口任务。

在之后，如果一个查询和一个键来自同一图像，则我们将它们视为正对，否则我们将它们视为负样本对。在之后，我们在随机数据增强下对同一图像进行两个随机的“视图”，形成一个正对。查询和键分别由它们的编码器fq和fk进行编码。该编码器可以是任何卷积神经网络。

对于当前的小批处理，我们对查询及其相应的键进行编码，它们形成了正样本对。阴性样本来自于队列。

技术细节。我们采用ResNet作为编码器，其最后一个全连接层（在全局平均池化之后）具有固定维的输出(128-D)。这个输出向量由它的l2-范数进行归一化。这是查询或键的表示形式。InfoNCE等式中的温度τ被设置为0.07。数据增强设置遵循：从一个随机调整大小的图像中提取224×224像素的裁剪，然后进行随机颜色抖动、随机水平翻转和随机灰度转换。

Shufflfling BN。我们的编码器fq和fk都有批处理归一化(BN)，就像在标准的ResNet中一样。在实验中，我们发现使用BN可以阻止模型学习良好的表示，就像在中所报道的那样(它避免了使用BN)。该模型似乎“欺骗”借口任务，容易找到低损失的解决方案。这可能是因为样品之间的批内通信(由BN引起)会泄漏信息。

我们通过变换BN来解决这个问题。我们使用多个GPU进行训练，并为每个GPU独立地对样本执行BN（如通常做法那样）。对于关键编码器fk，我们在将其分配到gpu之间之前，在当前的小批中打乱样本顺序（并在编码后打乱）；查询编码器fq的小批量的样本顺序没有改变。这确保了用于计算查询及其正键的批处理统计信息来自两个不同的子集。这有效地解决了作弊问题，并允许训练受益于BN。

实验略。

总结

我们的方法在各种计算机视觉任务和数据集中显示了无监督学习的积极结果。有一些开放式的问题值得讨论。MoCo从IN-1M到IG-1B的改进始终是明显的，但相对较小，这表明更大规模的数据可能没有得到充分利用。我们希望一项先进的借口任务能改善这一点。除了简单的实例识别任务之外，还可以采用MoCo来进行掩码自动编码等借口任务。我们希望MoCo能够用于处理其他涉及对比学习的借口任务。