无监督训练模型之------MoCo

1.导言

  参考：

• 论文：Momentum Contrast for Unsupervised Visual Representation Learning（用动量对比的方法去做无监督的表征学习）
• 李沐论文精度系列之《MoCo 论文逐段精读》、精读笔记

 1.1 前言

  MoCo于2019年11月13在 CVPR发表，并获得 CVPR2020最佳论文提名，它是用一种对比学习的方式进行无监督训练的模型。MoCo是第一个在很多主流的机器视觉领域上（比如分类、检测、分割、人体关键点检测等），都超越了有监督预训练模型的无监督模型，从某种程度上证明了无监督学习在机器视觉领域，也能达到很好的效果。

MoCo这个词，来自于论文标题的前两个单词动量对比Momentum Contrast。

1. 什么是动量。
   动量从数学上可以理解成一种加权移动平均：
   y t = m ⋅ y t − 1 + ( 1 − m ) ⋅ x t y_{t}=m\cdot y_{t-1}+(1-m)\cdot x_{t} yt​=m⋅yt−1​+(1−m)⋅xt​

• 上式中， y t y_{t} yt​和 y t − 1 y_{t-1} yt−1​是当前时刻和前一时刻的输出， x t x_{t} xt​是当前时刻的输入，m就是动量。
• 这个式子表示：当前时刻的输出，不仅依赖于当前时刻的输入，还依赖于前一时刻的输出。m越大，当前时刻的输入 x t x_{t} xt​对结果 y t y_{t} yt​就越小。
• MoCo利用了动量的这种特性，从而缓慢地更新一个编码器，让中间学习的字典中的特征尽可能地保持一致。（下文会详细讲到）

         2.然后介绍一下什么是对比学习。

• 原理 ：对比学习是无监督学习的一种，着重于学习同类实例之间的共同特征，区分非同类实例之间的不同之处。

  举个例子，从imagenet中抽出猫、猫、狗、飞机四张图，那么猫和猫的图片肯定是相似的，和狗不相似。但是和飞机比起来，猫和狗是相似的。所以对比学习就是对比着差异去学习，模型并不需要真的知道图片中代表的是什么，而只需要知道哪些图片是类似的，哪些图片是不一样的就可以了。

• 训练目的：对比学习，希望相似数据（图片）最终学到的特征是相似的，在特征空间（embedding space ）中，特征向量尽量靠近；反之还希望不同的数据学到的特征向量，尽量远离。
• pretext task（代理任务）：对比学习是不需要标签的（比如不需要知道图片是哪一类），但模型还是需要知道哪些图片是类似的，哪些是不相似的，才能训练。这就需要通过通过设计一些巧妙的代理任务，人为指定一些任务来实现。
• 应用最广的代理任务：instance discrimination 。
  • 简单说就是，从一堆图片中调出任意一张图片 x i x_i xi​，将其做一次转换（transformation ，比如随机裁剪等数据增广），得到新的图片 x i 1 x_{i1} xi1​、 x i 2 x_{i2} xi2​。那么样本 x i 1 x_{i1} xi1​叫做基准点（锚点）， x i 2 x_{i2} xi2​被认为是正样本（两者都是从 x i x_i xi​变化得到的，虽然看起来有差异，但语义信息不应该发生变化），数据集中其它所有图片都是负样本。
  • 有了正负样本的划分，就可以将数据都输入编码器进行编码提取特征了。因为所有的正负样本都是基于锚点来说的，所以 x i 1 x_{i1} xi1​会单独使用一个编码器 E 11 E_{11} E11​， x i 2 x_{i2} xi2​和其它所有负样本使用另外的编码器（可以是同一个编码器，也可以也可以使用不同的编码器。但是不同的编码器之间必须相似，这样编码的特征才有一致性，才有比较的意义）。
  • 对比学习就是要让正样本的编码特征和锚点的编码特征尽可能靠近（相似），让负样本的特征和锚点特征尽量远离。
  • instance discrimination直译过来就是个体判别，在这个任务中，只有经过这张图片转换的样本才是正样本，其它图片都是负样本，所以每张图都自成一类。对于ImageNet来说，就不是1000类，而是128万个类别。
• 目标函数：确定了代理任务，知道如何定义正负样本之后，就需要用一个目标函数，来告诉模型该如何学习，比如常见的对比学习目标函数NCE loss等。
• 特性：对比学习最大的特性，是这种方法非常的灵活，可以设置各种不同的代理任务。只要找到一种方式去定义正负样本，剩下的都是一些比较标准化的流程。

1.2 摘要

  我们在机器视觉领域提出了一种新的无监督学习方法——MoCo。MoCo虽然是基于对比学习的，但是本文是从另外一个角度来看对比学习，即把对比学习看作是一个字典查询任务。

  比如将上面提到的 x i x_i xi​当做是query，其它包括 x i 1 x_{i1} xi1​、 x i 2 x_{i2} xi2​这些图片都是字典中的key。我们每次判断正负样本，就是看字典中的这些key和query是否相似，而这些key都是通过encoder来更新的。

  具体来说，我们构建了一个动态的字典，这个字典有两个特性：队列特性和moving-averaged encoder（这两点在下文模型结构中会具体说明，现在记住就行）。因为这样，我们的字典非常大，且特征一致性非常好，从而便于进行对比学习。

  最终，MoCo作为一个无监督的预训练模型，能够在7个下游任务（分割、检测等）上 ，超越之前的有监督的预训练模型 ，填平了CV领域中，无监督训练和有监督训练之间的坑。

1.3 导言

  GPT和BERT已经证明了无监督的表征学习在NLP领域是非常成功的，但是在视觉领域，无监督学习效果差很多，作者认为可能是二者的原始信号空间不同。

• 在NLP任务中，原始信号空间是离散的（都是一些含有不同语义的单词或者词根），信号本来就拉得比较开，容易建立tokenize（将单词映射成向量）的字典。这样无监督学习容易建模，且模型容易优化。
• CV中，视觉信号都是在一个连续且高维的空间里，不想单词那样信息和语义浓缩的那么好，不够简洁，这样就不容易建立一个这样的字典，也就不容易进行无监督学习。

  最近有一些无监督学习方法表现不错，但是都可以归结为建立动态字典。
  如果将上一节讲到的所有样本都构建到一个字典中，字典的key就是各个样本，字典的value就是编码之后的特征（后面直接以 k 0 k_0 k0​表示第一个样本的编码特征）。我们先编码好锚点的特征，当做query；其它所有样本特征当做字典中不同的key，那么那对比学习就转化成为了一个字典查询的问题了。
如下图所示，我们训练一些编码器，再根据q去字典中查找key。查找的目的，就是让已经编码好的特征q，和与它匹配的特征key（其实就是正样本 x i 2 x_{i2} xi2​的特征）最相似；与其它不匹配的特征不相似。

 

  在MoCo这篇论文当中，因为作者已经把所有的对比学习的方法归纳成为了一个动态字典的问题，所以很少使用anchor或者正负样本这些词，用的都是query和key。所以锚点 x i 1 x_{i1} xi1​用 x q u e r y x^{query} xquery表示，其编码特征用q表示。其它样本和对应特征分别用 x i k e y x_{i}^{key} xikey​和 k i k_i ki​表示。

作者认为，一个好的字典应该有两个特性：

• 字典足够大
  • 字典越大，key越多，所能表示的视觉信息、视觉特征就越丰富 ，这样拿query去做对比学习的时候，才越能学到图片的特征。
  • 反之，如果字典很小，模型很容易通过学习一些捷径来区分正负样本，这样在碰到大量的真实数据时，泛化就会特别差（我的理解是，字典中只有猫和狗，狗都是黑色，猫都是黄色。模型简单的判断图片中物体是否是黄色，来区分猫和狗，而不是真的学到了猫和狗的特征）
• 编码的特征尽量保持一致性
  字典里的key都应该用相同或者说相似的编码器去编码得到，否则模型在查找query时，可以简单的通过找到和它使用相同或者相似编码器的key，而不是真的和它含有相同语义信息的key（变相引入两一个捷径）。

  以前的对比学习，都至少被上述所说的两个方面中的一个所限制（要么一致性不好，要么字典不够大）。本文最大的贡献，就是使用队列以及动量编码器来进行对比学习，解决了这个问题。具体来说：

• key（编码特征）并不需要梯度更新，而是通过更新编码器，新的编码器使输出的key更新。
• queue ：整个队列里面的元素都是字典，队首输入当前batch的编码特征，队尾弹出最旧的batch特征。每次移除的是最老的那些key，从一致性的角度来说 ，有利于对比学习。
  • 用队列的好处是可以重复使用那些已经编码好的key，而这些key是从之前的那些mini-batch中得到的。
  • 用队列结构，就可以把的mini_batch的大小和队列的大小直接分开了，所以最后这个队列的大小，也就是字典的大小可以设的非常大，因为它大部分的元素都不是每个iteration都需要更新的。
  • 在字典里计算loss而不是整个数据集上计算loss，使用队列的数据结构，可以让维护这个字典的计算开销非常小。
• momentum encoder：
  • 如果只有当前batch的key是从当前的编码器得到特征，其它的key都是另外时刻的编码器输出的特征，这样就无法保证字典中key的一致性。所以作者又提出了动量编码器
  • 动量编码器，即编码器参数的更新方式就是 y t = m ⋅ y t − 1 + ( 1 − m ) ⋅ x t y_{t}=m\cdot y_{t-1}+(1-m)\cdot x_{t} yt​=m⋅yt−1​+(1−m)⋅xt​（MoCo中m=0.999）。
  • 初始化的编码器来自于query的编码器，之后每次更新，只有1‰的参数会从query的编码器参数里拿过来更新，所以这个编码器参数更新的非常缓慢。从而保证了字典中所有的key都是由相似的编码器抽取得到的，尽最大可能地保持了他们的一致性。（直接更新编码器k的所有参数，会导致编码器更新过快，降低了这个队列中所有key的特征的一致性）
• 动态字典：字典中的key都是随机取样的，而且key的编码器在训练的过程中也是在不停的改变。

 2. 相关工作

2.1 SimCLR：端到端的学习方式（Inva Spread也是）

端到端学习，顾名思义就是编码器都是可以通过梯度回传来更新模型参数的，优缺点都很明显：

• 缺点：字典大小和mini_batch大小一致，但是现在一般是存不了太大的batch的，而且太大的batch难以优化，处理不好的话，不容易收敛，所以最终模型效果没那么好。
• 优点：因为进行梯度回传，所以编码器可以实时更新，字典中的key的特征一致性非常高
• SimCLR最终使用batch_size=8192来做训练（google有TPU，内存大，可以无脑上batch-size），可以支持模型做对比学习

 

 

2.2 memory bank （InstDisc模型）

• 在memory bank中，q的编码器是梯度更新的，但是字典中的k，是没有单独的编码器。
• memory bank把整个数据集的特征都存到了一起。每次训练时，只需要从memory bank中采样一些key来作为字典（比如 k 1 k_1 k1​、 k 2 k_2 k2​、 k 3 k_3 k3​），然后正常计算q和k的loss，进行梯度回传更新编码器。
• 编码器更新后，重新编码 k 1 k_1 k1​、 k 2 k_2 k2​、 k 3 k_3 k3​得到新的值，替换原来对应的值，这样就完成了一次memory bank的更新，依此类推。

  ImageNet虽然有128万张图片，即128w的key，但是特征维度为dim=128，用memory bank存下来只需要600M，所以这样做是没问题的。但是对于一个拥有亿级图片规模的数据，存储所有的特征就需要几十G甚至上百G的内存了，所以memory bank的扩展性不如MoCo好。

但是这样做有一个明显的问题，就是特征的一致性非常差。表现在：

• 编码器q是梯度回传更新的，所以更新的很快，这样key都是在不同时刻编码器编码的特征，所以特征一致性很差
• memory bank存储了所有的图片，也就意味着模型训练了整整一个epoch才能把整个memory bank更新一遍，那也就意味着，当开始下一个epoch训练的时候，假如选了三个key，那这三个key的特征都是上一个epoch不知道哪个时间点算出来的特征了，这也就导致query的特征和key的特征差的特别远。
• memory bank 的作者也意识到了这一点，所以使用另外一个loss（proximal optimization），目的就是为了让训练变得更加平滑，而且也提到了动量更新，只不过它的动量更新的是特征。

  由此，作者才提出了MoCo，采用队列的形式去实现字典，使其不必受限于字典的大小；使用动量编码器进行缓慢更新，使特征保持一致性。

3.算法

3.1 损失函数

在本文中，采取了一个叫做InfoNCE的对比学习函数来训练整个模型。

  式子中，τ是一个超参数。如果去掉τ，整个式子其实就是交叉熵损失函数（cross entropy loss ），在后面的伪代码中，也是基于cross entropy loss实现。

• 分子表示q和正样本做计算，分母其实是k个负样本上做累加和，因为是从0到k，所以是k+1个样本，也就指的是字典里所有的key。
• 直接计算复杂度太大： MoCo使用 instance discrimination作为代理任务，那么光是ImageNet数据集，就有128万个类别，直接计算，复杂度会非常高，难以训练（128万类的softmax）。
• NCE loss（noise contrastive estimation ）：将超级多分类转为二分类——数据类别data sample和噪声类别noisy sample。这样解决了类别多的问题。

  estimation：近似的意思。为了降低计算复杂度，不是在每次迭代时遍历整个数据集128万张负样本，而是只从数据集中选一些负样本来计算loss（也就是选队列字典中的6万多个负样本），相当于一种近似。所以这也是MoCo一直强调的希望字典足够大，因为越大的字典，越能够提供更好的近似。

InfoNCE:NCE的一个简单的变体.

• 作者认为如果只把问题看作是一个二分类（只有数据样本和噪声样本）的话，可能对模型学习不是很友好，毕竟在那么多的噪声样本中，大家很有可能不是一个类，所以还是把它看成一个多分类的问题比较合理。
• 公式中的q * k，其实就相当于是logit，也可以类比为softmax中的z。
• τ：一个超参数，用来控制分布的形状 。τ越大，分布中的数值越小，经过exp之后就更小了，分布就会变得更平滑，相当于对比损失对所有的负样本都一视同仁，导致学习的模型没有轻重
• τ越小，分布更集中，模型只关注那些特别困难的样本，其实那些负样本很有可能是潜在的正样本，如果模型过度地关注这些特别困难的负样本，会导致模型很难收敛，或者学好的特征不好去泛化。

 3.2 伪代码

  对于整个模型来说，在代理任务不一样的时候，输入 x q x^q xq和 x k x^k xk既可以是图片，也可以是图片块（CPC），或者是含有上下文的一系列的图片块。
  query的编码器和key的编码器既可以是相同的（模型的架构一样，参数完全共享，比如Inva Spread），或者说它们的参数是部分共享的，也可以是彻底不一样的两个网络（CMC，多视角多编码器）。

上面提到的CPC、CMC、Inva Spread、SimCLR、InstDisc在后面对比学习综述中都会简单介绍。

下面是论文中作者给出的伪代码，其中：

• fq、fk分别是query和key的编码器
• queue这个队列指的是字典，里面一共有k个key，所以它的维度是c*k，c指的是每个特征的维度（c=128）
• m是动量，t是InfoNCE里面的超参数τ
• aug表示数据增强

1. 初始化编码器fq，并将其参数赋值给编码器f_k
2. 从data loader里拿一个batch的数据（n=bacth_size=256，n是采样数）
3. 通过数据增强得到正样本对x_q和x_k，然后通过各自的编码器得到特征q和特征k（大小都是N*C）。key不需要梯度回传，所以用.detach() 去掉梯度信息。
4. 计算N张图片的自己与自己的增强图的特征的匹配度
   q 、k之间计算logit（正样本），也就是之前公式1中算InfoNCE loss的时候的分子 q ∗ k + q * k+ q∗k+，其特征维度就变成了n * 1（256，1）。
5. 计算N张图片与队列中的K张图的特征的匹配度
   q、queue拿出来计算，得到InfoNCE的分母，也就得到了负样本的logit，维度是n*k（256*65536，MoCo中，字典大小为65536）
6. 将正负样本logit进行cat拼接
7. 通过交叉熵损失函数实现loss计算。具体的，设置一个全0向量作为ground truth来进行计算。
   因为按照作者的这种实现方式，所有的正样本永远都是在logit的第一个位置上，也就是位置0，所以对于正样本来说，如果找对了那个key，在分类任务中得到的正确的类别就是类别0，所以巧妙地使用了这种方式创建了一个ground truth，从而计算出了对比学习的loss
8. 根据loss进行梯度回传，更新编码器fq
9. 动量更新编码器f_k
10. 更新队列（队首压入新的batch编码的key，队尾弹出最旧的key）

f_k.params = f_q.params # 初始化
for x in loader: # 输入一个图像序列x，包含N张图，没有标签
    x_q = aug(x) # 用于查询的图（数据增强得到）
    x_k = aug(x) # 模板图（数据增强得到），自监督就体现在这里，只有图x和x的数据增强才被归为一类
    q = f_q.forward(x_q) # 提取查询特征，输出NxC
    k = f_k.forward(x_k) # 提取模板特征，输出NxC
    # 不使用梯度更新f_k的参数，这是因为文章假设用于提取模板的表示应该是稳定的，不应立即更新
    k = k.detach() 
    # 这里bmm是分批矩阵乘法
    l_pos = bmm(q.view(N,1,C), k.view(N,C,1)) # 输出Nx1，也就是自己与自己的增强图的特征的匹配度
    l_neg = mm(q.view(N,C), queue.view(C,K)) # 输出Nxk，自己与上一批次所有图的匹配度（全不匹配）
    logits = cat([l_pos, l_neg], dim=1) # 输出Nx(1+k)
    labels = zeros(N)
    # NCE损失函数，就是为了保证自己与自己衍生的匹配度输出越大越好，否则越小越好
    loss = CrossEntropyLoss(logits/t, labels) 
    loss.backward()
    update(f_q.params) # f_q使用梯度立即更新
    # 由于假设模板特征的表示方法是稳定的，因此它更新得更慢，这里使用动量法更新，相当于做了个滤波。
    f_k.params = m*f_k.params+(1-m)*f_q.params 
    enqueue(queue, k) # 为了生成反例，所以引入了队列
    dequeue(queue)

 4.实验

4.1 对比其他模型

下图是端到端学习、memory bank和MoCo三种流派的模型，在只做特征提取时候的精度对比：

 

• 横坐标用k表示，指的是用了多少个负样本，也可以粗略地理解为字典的大小
• 纵坐标指的是在ImageNet数据集上的top 1的准确率
• 端到端学习，受限于显卡内存，实验结果只有三个点（字典最大1024）
• MoCo性能最好，对硬件要求最低，而且扩展性也比较好

 4.2 imagenet数据集结果对比

 

 

• 表格中上半部分都不是使用的对比学习，下半部分都是使用的对比学习，可以看到对比学习效果明显更好
• 精度结果后面的特殊标记表示用fast auto augment做了数据增强（ImageNet有监督训练的数据增强策略）

 4.3 迁移学习效果

归一化
  预训练好的MoCo做微调，其学习率需要设为30，远大于以前模型的微调时的一些学习率（比如lr=0.03）说明MoCo学到的特征跟有监督学到的特征的分布是非常不一样的，但是不能每次微调时都去grid search找一下它最佳的学习率是多少，这样失去了微调的意义。
  当分布不一致的时候，最常想到的方法就是归一化，所以作者这里使用了特征归一化的方法（整个模型都做BN，包括检测时用到的FPN结构，也使用BN）。做完归一化之后，就可以拿这些有监督训练用的超参数来做微调了。

 在keypoint detection人体关键点检测、pose estimation姿态检测、实例分割、语义分割四个任务中做测试：

• 第一行使用的是随机初始化的模型再做微调，所以它是一个基线网络，分数比较低
• 第二行使用的是有监督的ImageNet的预训练的模型做初始化然后再做微调，也就是一个比较强的极限结果
• 最后两行分别是MoCo在ImageNet上和在Instagram 1Billion上做无监督预训练当作模型的初始化，然后再做微调

  结论：MoCo预训练的模型在大部分时候都比ImageNet的有监督预训练模型要好，在实例分割和语义分割的任务上有时候会稍差一些。

  所以大家怀疑对比学习可能不太适合做这种每个像素的都要预测的任务，基于这一点，后续发展出dence contrast或者是pixel contrast。

 5.总结

  MoCo 的主要贡献就是把之前对比学习的一些方法都归纳总结成了一个字典查询的问题，并提出了队列存储和动量编码器。前者解决字典太大不好存储和训练的问题，后者解决了字典特征 不一致的问题；从而形成一个又大又一致的字典，能帮助模型更好的进行对比学习。

  MoCo跟Inst Disc是非常相似的，比如它用队列取代了原来的memory bank作为一个额外的数据结构去存储负样本，用动量编码器去取代了原来loss里的约束项，这样就可以动量的更新编码器，而不是动量的去更新特征，从而能得到更好的结果。其整体的出发点以及一些实现的细节（比如backbone和lr、batch_size，dim、τ等等超参数都是一样的）和Inst Disc都是非常类似的，所以可以说MoCo是Inst Disc的改进工作。但是MoCo真正出色的地方其实有两点 ：

• 使用动量编码器。这个改进简单有效，并在后面一系列工作中被一直沿用（比如SimCLR、BYOL），所以也非常深刻。
• 写作高人一等。直接把之前所有的方法都总结成了一个字典查找的问题，所以直接把问题给归纳升华了。而且提出CV和NLP的对比学习大一统框架，论文的泛用性彻底扩大了。

  MoCo还有一个优点，就是训练比较便宜。在一张8卡V100 16G GPUs上，训练200个epoch只需要53小时（batch_size=256，GPU memory=5.3G），完全就是大佬给我们送福利。MoCo这篇论文以及它高效的实现，能让大多数人有机会用普通的GPU就能跑对比学习实验，做自己的研究。

  最后，因为MoCo在各个视觉任务上取得了更好的性能，也激发了很多后续分析性的工作，去研究MoCo学出来的特征到底和有监督学出来的特征有什么不同，还能从别的什么方向去提高对比学习。