【论文精读】Momentum Contrast for Unsupervised Visual Representation Learning

前言

如果说MAE证实了无监督学习在CV领域能够像在NLP领域那样超过监督预训练的性能，那么MoCo就填补了CV中无监督学习和监督学习的gap。作为来自大神何恺明的团队的作品，MoCo行文充满了大佬的自信，让我们看看MoCo这篇工作在写作和方法构建上的精妙之处~

---

Paper: https://openaccess.thecvf.com/content_CVPR_2020/papers/He_Momentum_Contrast_for_Unsupervised_Visual_Representation_Learning_CVPR_2020_paper.pdf
Code: https://github.com/facebookresearch/moco
From: CVPR2020

Abstract

本文提出MoCo用于无监督视觉表征学习。它将对比学习任务看成字典查询任务，作者设计了由队列和移动平均编码器组成的动态词典。由于无需梯度反向传播，因此可以得到一个大型且一致的字典， 从而帮助无监督学习。MoCo最大的优势在于可以将学习到的表示很好迁移到下游任务。在实验中，MoCo在PASCAL VOC、COCO等七个不同任务的数据集上超过了监督预训练的结果，这表明在视觉任务上，监督学习和无监督学习的差距越来越小。

1. Introduction

无监督学习在NLP领域大为成功，但在CV领域仍然是监督学习作为主导。导致二者不同的原因可能是信号差异。语言模型有着离散的信号空间，即词向量空间，用于构建词表。相比之下，CV进一步关注字典的构建，因为图像信号处于连续的高维空间，不适合人类交互。

这里信号不同的问题很好理解。在NLP中，一种语言的单词个数几乎是固定的（当然也受到划分规则的影响），并且很容易区分，但是在CV中，一张图像可用于建模的图像块就数不胜数，其向量化得到的词表空间巨大，因此如何构建词表是重中之重。

最近在对比损失上的研究让无监督学习在CV上有了一定的进展。其核心思想是构建动态字典。字典中的keys（即token）从原始数据中采样，由编码器进行表征。无监督学习训练编码器从字典中查找，让keys对应的query尽可能接近keys，其余token远离keys。学习过程通过设计对比损失实现。
本文构建的字典有如下特点：

1. 大。
2. 训练过程中保持一致性。

直观上来说，更大的字典可以更好对底层的连续高维视觉空间进行采样，并且字典中的token应该由相同的编码器进行表征，以便查询比较保持一致。但是这样的方法会受到一定的限制。

查询比较一致这里，我的理解是采用相同的标准和规范，这样的比较才有意义，没有偏差。

本文提出MoCo作为一种构建大型且一致的字典的方法，采用对比损失的无监督学习。

作者将字典维护为数据采样的队列，将当前新的mini-batch样本进行编码入队，将旧的样本出队。队列的设计将可以将字典大小和mini-batch解耦，让字典变得很大。此外，由于字典keys来自于最前面的几个小批量，因此提出了一种可以缓慢进行的key编码器，基于动量的移动平均值， 以保持一致性。
MoCo是一种基于对比学习构建动态词典的机制，可以结合很多代理任务。本文采用的是简单的个体判别任务。即一个query是来自于key的不同视角（或者图像块），二者就可以匹配。采用该任务，MoCo在ImageNet数据集中的线性分类上展现了有力的结果。
实验部分，MoCo在七个下游任务上达到了SOTA。并且在更大的数据集上预训练还能进一步提升性能，证明MoCo可以在更真实的场景工作。这些结果都表明，MoCo在很大程度上缩小了许多CV任务中无监督和监督表示学习之间的差距，并且可以在多种应用中作为ImageNet监督预训练的替代方案。

2. Related Work

无监督/监督学习方法主要包含两个方面：代理任务和损失函数。pretext意味着进行的任务不是最终的任务，而是为了学习良好的数据表示。损失函数通常独立于代理任务。MoCo重点关注损失函数部分。

Loss functions

定义损失函数最寻常的方式是衡量模型预测与目标的不同。例如通过L1、L2损失重建输入的像素，通过交叉熵或基于边际的损失将输入定义为分类任务等等。
但是对比损失不同，它用于表示空间中样本对的相似性，其中目标在训练过程中动态变化，而不是将输入与固定目标相匹配。
此外还有对抗损失，用于衡量概率分布之间的差异。这是一种广泛的生成无监督数据的方法。GAN和噪声对比估计（NCE）之间存在关系。

Pretext tasks

代理任务种类有很多，包括对破坏的输入进行修复，如去噪自编码器、上下文自编码器等生成任务，还有一些生成伪标签的代理任务，如对单个图像进行转换、块排序，还有聚类等。

Contrastive learning vs. pretext tasks

各种代理任务都可以基于某种形式的对比损失函数。如个体判别任务与基于示例的任务以及NCE有关。对比预测编码（CPC）中的代理任务是上下文自编码的一种形式，而在对比多视图编码（CMC）中，与着色相关。

3. Method

3.1. Contrastive Learning as Dictionary Look-up

对比学习最新的工作可以认为是字典查找的编码器，考虑一个查询$q$和一组编码的字典key集合 { k 0 , k 1 , k 2 . . . } \{k_0,k_1,k_2...\} {k0​,k1​,k2​...}。假设集合中只有一个key和$q$匹配。对比损失是一个函数，当$q$与其正确的key相似并与其他key远离时，其值较低。本文采用了名为InfoNCE的对比损失函数如下：
${\mathcal{L}}_q=-\log \frac{\exp \left(q\cdot k_{+}/\tau \right)}{{\sum }_{i=0}^K\exp \left(q\cdot k_i/\tau \right)}$
其中$\tau$是温度超参数。样本由一个正样本和$K$个负样本组成。直观上，这个损失函数就是尝试将$q$分类为$K+1$类的softmax对数损失。当然对比损失也可以基于边界和其他NCE变体。
对比学习作为无监督目标函数，用于训练表征查询和keys的编码器网络。一般来说，查询表示为$q=f_{\mathrm{q}}(x^q)$，其中$f_{\mathrm{q}}$是编码器网络，同理$k=f_{\mathrm{k}}(x^k)$。其实例化取决于具体的代理任务，两个编码器网络可以是相同的，部分共享或者不同。

3.2. Momentum Contrast

根据上述描述，对比学习是一种在图像等高维连续输入上构建离散字典的方法。字典是动态的，因为编码器在训练中不断变化。作者假设好的特征可以通过涵盖丰富的负样本的大型字典来学习，因此作者理想的字典是大型且尽可能保持一致的。为此，作者提出了动量对比。

这里尽可能保持一致需要进一步说明，如果字典中编码的key不一致，即每个样本都是不同的或者相差很大的编码器进行编码的，这会导致模型学习到一些捷径，而不是真正学习到图像的特征。

Dictionary as a queue.

方法的核心在于将字典维护为数据样本队列。这使得作者可以复用之前mini-batch的编码keys。这种方式的引入将字典大小和mini-batch解耦，字典的大小可以比mini-batch大得多，并且可以灵活设置为超参数。
字典中的样本会逐步被取代，遵循FIFO方式，当前mini-batch入队，旧的mini-batch出队列。这样的方式可以将旧的过时的编码keys删除，防止带来不一致性。

Momentum update.

使用队列的方式虽然能使字典变大，但也使得梯度反向传播困难。一个简单的方法是采用查询的编码器，但是效果不好，可能是因为编码器快速变化降低了一致性。为此作者设计了动量更新。
形式上，将$f_{\mathrm{k}}$的参数表示为${\theta }_{\mathrm{k}}$,$f_{\mathrm{q}}$的参数表示为${\theta }_{\mathrm{q}}$，更新公式如下：
$${\theta }_{\mathrm{k}}\leftarrow m{\theta }_{\mathrm{k}}+(1-m){\theta }_{\mathrm{q}}$$
$m\in [0,1)$是动量系数。这种更新方式更为平滑，虽然keys仍然是不同的编码器编码得到的，但是编码器之间的差距很小。实验中$m=0.999$比$m=0.9$效果更好，说明缓慢更新的编码器是利用队列的核心。

Relations to previous mechanisms.

上图是MoCo与两种现有通用机制的对比，它们在大小和一致性上表现不同的特性。
端到端的更新方式如a所示，它使用当前小批量数据作为字典大小，保证了keys的一致性，但是字典大小受到GPU显存的限制，并且大批量优化也是一个挑战。现有的方法设计特殊的代理任务，但是需要特殊的网络设计，让模型处理下游任务变得复杂。
另一种机制是记忆库方法，如b所示，它存储了所有样本的表征，每个mini-batch都是随机从记忆库中采样，没有反向传播，可以支持大的字典。
第四部分将对这三种机制进行对比。

3.3. Pretext Task

对比学习可以驱动各种代理任务。本文的重点不在代理任务设计上，因此采用简单的个体判别任务。
同一张图片的查询和key作为正样本对，其他作为负样本。查询和key属于同一张图片的不同视图。查询和keys通过不同的卷积神经网络进行编码。

算法1提供了MoCo执行代理任务的伪代码。对于当前的mini-batch，对查询和相应的key进行编码，形成正样本，负样本来自队列。

Technical details

作者采用ResNet作为编码器。最后一个全连接层具有固定维度的输出。$\tau$设置为0.07，此外作者还采用随机颜色抖动、水平翻转、灰度转换等数据增强方式。

Shuffling BN

编码器都采用了批归一化方法BN，但是BN会通过捷径阻碍模型学习，可能的原因是样本之间因为批处理导致了信息泄露。
作者设计了shuffling BN解决这个问题。作者采用多个GPU进行训练，并为每个GPU单独执行BN。对于key编码器，作者在分配给GPU之前打乱了当前mini-batch中样本的顺序，并在编码后还原顺序。这确保了查询和其对应的key来自不同的子集。
这个方法应用在MoCo和端到端的方法中，记忆库方法无需这种方法，因为其正样本key来自于过去不同的mini-batch。

4. Experiments

作者分别在ImageNet-1M (IN-1M)和Instagram-1B (IG-1B)两个数据集上进行无监督预训练。前者数据集类别分布均衡且特征明显，后者更接近现实数据。训练超参数设置见原文。

4.1. Linear Classification Protocol

作者首先冻结原始模型，并训练一个线性层执行分类任务，在ImageNet数据集上对MoCo进行验证。作者通过网格搜索发现，最优的学习率为30，这说明无监督学习学到的特征分布和有监督学习学到的特征分布有很大不同。

Ablation: contrastive loss mechanisms

作者对比了三种不同的机制，为了关注于对比损失机制，作者在相同的代理任务下执行了这三种机制。结果如下图所示：

可以看到，所有机制都受益于样本个数，这支持了作者构建大字典的动机。
端到端的方法在K很小的时候和MoCo性能相近，但是受限于mini-batch的大小（显存不支持大批量，并且大批量优化困难），最多只能支持1024个样本。并且即使显存足够，折线图的趋势是否可以外推到更大的K也值得怀疑。
记忆库的机制可以支持大的字典大小。但是在任何规模的K下都比不如MoCo。这印证了作者假设：keys来自不同的编码器导致不一致。

Ablation: momentum

下表展示了不同动量m下MoCo训练的accuracy。

可以看到当m在范围0.99~0.9999性能表现最好，m过小时，准确率会大幅下降，为0时将无法收敛。

Comparison with previous results

以前的无监督学习方法在模型大小上存在明显的差异。为了公平比较，作者加入了准确率和参数量的权衡。下表是比较结果：

MoCo在ResNet50下取得了有竞争的60.6%的准确率，比其他相同规模的模型都要好。随着模型增大，MoCo的性能也变得更好。

4.2. Transferring Features

无监督学习的主要目标是学习可迁移的特征。作者将MoCo和ImageNet监督预训练在各种下游任务上进行比较。作为先决条件，作者首先讨论两个重要问题：归一化和调度。

Normalization

根据4.1的描述，监督学习和无监督学习在学习的特征分布上有很大的不同，但是下游任务上的超参数会优先选择监督预训练的设置。为了缓解这个问题，作者在微调期间采用特征归一化，并且不冻结参数。

这里需要对统一超参数进行解释，虽然无监督学习会在类似于30的学习率上表现更好，但是下游任务不可能为了不同的无监督学习方法都单独寻找最优的学习率等超参数，因此需要将无监督学习的超参数和监督学习的超参数对齐。

Schedules

如果微调时间足够长，随机初始化的模型重新训练是可以达到MoCo性能的。作者的目标是研究MoCo的迁移能力。所以作者的实验在受控的时间表上进行。
总而言之，虽然设置上与监督学习对齐让MoCo处于不利的地位，但是MoCo仍具有竞争力，这样做还能在多个数据集和任务上进行比较。

4.2.1 PASCAL VOC Object Detection

在目标检测任务上，作者基于VOC数据集，由上表可知，IN-1M上的监督预训练和IN-1M上的MoCo性能相当，但是IG-1B上的MoCo则超过了它。有意思的是，迁移的准确率还和检测架构有关。预训练和模型架构的关系应该成为考虑的因素。

此外还进行了三种不同机制在目标检测任务下的比较。另外两种方法在所有指标上都不如MoCo。

先前的对比学习方法都无法赶上其各自的监督学习的结果。如上表所示。

4.2.2 COCO Object Detection and Segmentation

目标检测和语义分割任务上，作者在train2017上进行端到端的微调，在val2017上进行评估。MoCo在所有指标上都优于其ImageNet监督学习对手。

4.2.3 More Downstream Tasks

上表显示了更多下游任务的结果。总的来说，MoCo都比ImageNet监督训练的方法要好。但是在VOC语义分割上，MoCo至少差了0.8个点。总而言之，MoCo在很大程度上缩小了CV多任务中监督学习和无监督学习之间的差距。值得注意的是，IG-1B上预训练的MoCo始终优于IN-1M上预训练的MoCo。这表明MoCo在大规模的数据集上可以表现良好。

5. Discussion and Conclusion

MoCo展示了无监督学习在CV领域的优势。但是有一些问题值得讨论，MoCo的性能从IN-1M到IG-1B的改进虽然明显，但是较小，这说明大规模数据没能充分利用，可能更高级的代理任务可以改善这个问题，比如可以采用masked auto-encoding（即后续的MAE）。

阅读总结

MoCo作为CV领域对比学习的成功范例，填补了监督学习和无监督学习在视觉领域的gap，也应证了大规模的无监督预训练方式是可行的， 虽然在更大规模的数据上性能提升较小，这需要进一步的研究。
MoCo的核心自然在于队列字典和动量编码器的设计。队列保证可以构建更大的字典，FIFO的方式也让最新的keys留在队列中，而动量编码器尽可能保证队列keys的一致性。解决了一致性和大二者不可得兼的问题，方法简单，设计精巧。
此外，对于无监督预训练提取特征的方式，我也有了更深的认识，本文的对比学习方式，本质上就是分类任务，通过构建自监督的分类标签来对模型进行更新，这样的预训练较为简单，远不如生成任务复杂，尤其是接近像素级别的生成任务，这也就是分类预训练在下游任务的表现不如生成式预训练的原因，比如作者两年后的工作MAE，当然模型架构也有一定的影响因素。