MoCo¹¶

摘要¶

Momentum Contrast(MoCo)是一种自监督视觉表征学习方法. 它的核心思想可以从对比学习作为"字典查询"的角度来理解. MoCo在构建对比学习所需要的"正样本"和"负样本"的时候, 引入了一个动态字典以及一种动量更新策略, 从而能够在训练过程中以相对较小的计算成本构造出一个规模大且一致性较高的对比字典, 并利用对比损失来进行表征学习. 在ImageNet上进行线性分类的时候, MoCo的表现非常具有竞争力. 而在下游任务(如目标检测和分割中), MoCo学习到的特征能带来更好的迁移性能, 甚至在PASCAL VOC, COCO等多个数据集上超过了同样网络结构的有监督预训练模型. 这是自监督视觉表征学习的一个重要的里程碑.

动机¶

信号空间不同¶

自监督表征学习在NLP领域非常成功, 如GPT和BERT. 但是在视觉领域, 有监督学习仍然是主流, 无监督学习/自监督学习仍然被它甩在身后. 其根本的原因可能是由于它们信号📶空间的不同: 语言数据由离散的信号组成, 如单词, 子词单元等, 这些离散的单位可以组织成一个字典(tokenized dictionaries), 使得模型能够基于这些预先定义好的信号进行学习和预测. 然而, 视觉数据是连续的, 高维的信号, 即图像中的像素值是连续变化的, 没有像语言那样的自然离散单位, 这种连续的变化是无穷无尽♾️的, 很难构建一个字典, 模型需要先找到有效的方法来表示视觉数据, 然后再基于这种表示进行学习和预测.

前人对比学习的工作¶

最近的几项研究展示了使用和对比损失相关的方法进行自监督学习, 并取得了有潜力的结果. 虽然它们在背后都受到不同的动机驱动, 这些方法可以被认为是在构建一个动态字典. 查询(Query)以及字典中的键(Key)都是通过编码器进行编码, 生成的向量, 将query和key的在潜空间中进行对比, 通过最小化对比误差, 使得匹配的Query-Key的embedding拉近, 不匹配的Query-Key的embedding拉远.

字典的理想特征¶

从这个角度看, 作者猜测在训练的过程中创造一个大的, 一致的对比字典是非常重要的. "大"是指字典存储了大量的key, 这些key是从图像patches中采样并经过encoder得到的表示向量. 如果字典越大, 理论上就可以更加全面地覆盖或代表原始数据所在的高维连续空间. "一致"是指字典里的所有key都要由相同/高度相似的encoder生成, 否则, 在计算query和不同key相似度的时候就会产生偏差. 然而, 现存的方法中的对比字典都不能很好的达到上述的两个理想特征.

解决字典问题¶

作者提出了MoCo, 用于构建大型且一致的字典, 用于基于对比损失的自监督学, 如下图所示.

动量对比(MoCo)通过计算查询$q$和一个动态字典中的所有key$k_0, k_1, k_2, ...$的对比损失, 来训练一个视觉表示encoder. 这个动态字典的key是从一组数据样本中即时定义的(defined on-the-fly). 该字典以队列的形式构建: 将当前的小批量数据加入队列, 同时移除最旧的小批量数据, 使其和小批量数据的大小解耦. 这些key由一个逐步更新(slowly-progressing)的编码器进行编码, 该编码器(又叫做动量编码器)通过和锚点编码器之间的动量更新来驱动. 这个方法能够为学习视觉表示提供一个大且一致的字典

锚点编码器和动量编码器

在MoCo中, 训练阶段会使用到两个编码器, 锚点编码器和动量编码器, 锚点编码器通过正常的反向传播进行更新, 而动量编码器通过"动量"的方式从锚点编码器同步更新(而不是反向传播), 用于构建对比学习所需的动态字典. 然而, 在推理或者下游任务中通常只会用到锚点编码器.

他们保持字典是一个数据样本的队列. 当前的小批量样本的编码表示会入队, 最旧的会出队. 队列使得字典的大小和下批量的大小解耦, 使得字典可以被做得非常大. 并且, 随着字典中的key被后面入队的下批量样本的编码表示替代, 一个逐步更新的动量编码器, 通过和锚点编码器之间的动量更新驱动, 会被用于保持一致性.

前置任务¶

MoCo是一种用于对比学习的机制, 它通过构建"动态字典"来存储和更新大量负样本的特征表示, 并可以与多种前置任务(pretext tasks)相结合. 作者选用了一种简单的实例判别(instance discrimination task)前置任务: 一条query会和一条key配对, 当且仅当它们都是由同一张原始图像的不同视图(例如, 不同的随机裁剪)编码得到. 也就是说, 如果query和key来自同一图像的不同增强版本, 它们就被视为正样本对, 否则就是负样本对. 通过这种方式, 模型被强迫去学习区分同一张图像的不同视图与其他图像的视图之间的差别, 从而提取对下游任务有用的视觉特征. 在ImageNet数据集上, 作者先利用MoCo进行自监督预训练, 学到模型的编码器, 然后在这一编码器的输出特征上仅仅训练一个线性分类器进行图像分类. MoCo仅仅使用这项简单的实例判别前置任务就取得了和其他方法能够竞争的结果.

迁移学习¶

自监督学习的主要目的是预训练特征表示, 然后通过微调迁移到下游任务. 作者展示了在7个和检测和分割相关的下游任务中, 进行自监督预训练的MoCo可以超过在ImageNet上进行有监督预训练的对手, 在某些情况下, 差距相当显著. 在这些实验中, 他们探索了在ImageNet或十亿张Instagram图片集上训练的MoCo模型, 结果表明, MoCo在更加接近真实世界, 包含数十亿图片且相对未经整理的场景中也能良好地运行. 这些实验结果表明, MoCo在很多计算机视觉任务中很大程度上弥合了自监督和有监督学习之间的差距, 并在某些应用中可以作为在ImageNet进行有监督预训练的模型的替代方案.

流程¶

设定batch大小是$N$, 这意味着每个batch会从数据集中取出$N$个图片
为了进行对比学习, 需要为每张图片生成两份不同的增广图片, 例如, 随机裁剪, 随机水平反转, 颜色抖动, Gaussian Blur等, 每次增广都有一定的随机性, 所以同一张图像的两次增广通常不会完全相同
对于当前batch中的$N$张原图, 每张图像会有2张增广, 总共有$2N$张图像, 我们将这两份增广称为query增广(要输入$f_q$)和key增广(要输入$f_k$)
将query增广输入到$f_q$, 得到一批Query特征向量$\{q_1, q_2, ..., q_N\}$
将key增广输入到$f_k$, 得到一批Key特征向量$\{k_1, k_2, ..., k_N\}$, 其中, 每对$(q_i, k_i)$来自同一个原图的不同增广, 被视为"正对"
将字典中的所有key$\{k_1^-, k_2^-, ...k_K^-\}$作为负样本
对于每个$q_i$, 它对应的正样本Key是当前batch生成的$k_i$
对于每个$q_i$, 它对应的负样本Key是从Queue里面取出来的$k_j^-$
计算对比损失
对$f_q$进行反向传播和梯度更新
使用动量更新的方式来更新$f_k$的参数
将当前的Key特征向量$\{k_1, k_2, ..., k_N\}$放到Queue的末尾
如果队列已经满了, 把最早进入队列的那个batch移除

方法论¶

对比学习当作字典查询¶

可以把对比学习看作是在训练一个用于字典查询的编码器.

考虑一个编码查询$q$和字典中的一堆编码样本$k_0, k_1, k_2, ...$. 假设在字典中有一个唯一的key$k_+$与$q$吻合. 当$q$和$k_+$相似的时候, 对比损失函数的值较小. 点积可以用来衡量两者之间的相似度, 在本研究中, 作者选择使用的是基于点积的InfoNCE当做对比损失函数.

\[ \mathcal{L}_q = -\log \frac{\exp\left(\frac{q \cdot k_{+}}{\tau}\right)} {\sum\limits_{i=0}^{K} \exp\left(\frac{q \cdot k_{i}}{\tau}\right)} \]

$\tau$是一个温度超参数, 这个思想和蒸馏中生成软目标中使用温度的思想是一致的. $K$是队列的大小, 字典的大小也是抽取的负样本的个数. 分母是$1$个正样本和$K$个负样本相似度的和. 在这种设置下, 可以将对比损失函数看作是一个具有$K+1$类别的Softmax分类器, 损失函数视图最大化正样本(即$k_+$)的概率, 同时最小化负样本(即$k_1, k_2, ..., k_K$)的概率. 对比损失函数也可以基于其他形式, 如基于margin(规定相似样本之间的margin), 或者是NCE损失的变体.

通常, 我们将查询表示为$q=f_q(x^q)$, $f_q$是编码器网络, $x^q$是一个查询样本, 类似的, 我们有$k=f_k(x^k)$. 查询的样本$x^k_q$的形式根据前置任务的不同而不同, 如可以是完整的图像, 图像patches等等. 并且, $f_q$和$f_k$分别表示的编码器网络在不同的对比学习方法中可能不同, 它们可以是相同的网络, 如SimCLR, 部分共享的网络, 如MoCo(一个是锚点编码器, 一个是动量编码器), 或者是不同的网络.

动量对比¶

从以上视角来看, 对比学习是在一个高维度连续输入(例如图像)上构建离散字典的方法, 请见信号空间不同. 这个词典是动态的, 因为用于对比的key是从字典中随机采样的, 并且key编码器在训练的过程中不断演化. 正如字典的理想特征中提到的, 作者猜想一个大的和一个在演化过程中尽量能够保持一致的解码器是学到优良表示的必要条件. 基于这样的猜想, 作者提出了动量对比.

队列=字典¶

MoCo最核心的思想是字典被维护为一个数据样本的队列. 我们可以重复利用历史仍保留在队列中的mini-batches中已经编码过的key, 并且字典的大小可以设置得比典型的mini-batch大很多, 可以作为一个独立的超参数, 实现了字典大小和mini-batch size的解耦.

这个字典中的样本被渐进式地替换. 当前的mini-batch入队, 最旧的mini-batch出队. 这意味着字典中的样本始终是整个数据集的一个采样子集. 维护一个固定大小的队列所需的额外计算资源是有限的, 计算复杂度较低. 并且, 随着训练的进行, key编码器的参数会不断更新, 意味着新编码的key会更准确地放映负样本的特征表示.

动量更新¶

通过使用队列, 可以维护一个包含大量样本的字典. 但是当字典变大的时候, 更新编码器的过程变得不可行, 这是因为梯度需要对队列中所有的样本进行回传, 而一个队列通常包含成千上万个mini-batch. 一个非常天真的想法是将$f_q$复制到$f_k$上, 忽视那个梯度. 但是这种做法在实验中产生了较差的结果. 作者猜测这可能是因为频繁的变动$f_k$导致了key表示的不一致性. 因此, 他们提出了动量更新的方法来解决这个问题.

使用$\theta_k$表示$f_k$的参数, 使用$\theta_q$表示$f_q$的参数, 通过下列公式更新$\theta_k$:

\[\theta_k\leftarrow m\theta_k+(1-m)\theta_q\]

其中, $m\in [0, 1)$是动量系数. 反向传播只会更新$\theta_q$. 动量更新使得$\theta_k$比$\theta_q$的变化更加平滑. 因此, Queue中不同mini-batch的key是由不同的差距很小的动量编码器编码的. 在实验中, 相对较大的动量系数, 如$m=0.999$比较小的值, 如$m=0.9$效果要更好, 这意味着缓慢渐进式更新是使用queue的核心要素.

和之前工作的关系¶

MoCo是一个通用的使用对比损失的机制. 作者将其和显存的通用机制比较, 在字典大小和一致性上, 展现了不同的设置. 如下图所示.

三种对比损失机制在理念上的比较. 这里只展示了一对query和key. 三种机制的主要区别是在key如何维护和key编码器如何更新上. (a) query和key的编码器是端到端反向传播更新的 (b) key是从一个记忆池中采样的 (c) MoCo使用动量编码器即时对key进行编码, 维护key的一个队列

使用反向传播进行end-to-end更新是一种很自然的机制. 它使用当前的mini-batch作为负样本字典, 所以负样本key之间的编码具有一致性, 但是字典的大小和mini-batch的大小是相等的, 它的大小受限于GPU所能容纳的batch size. 最近的一些同类的方法通过使用局部位置信息来构造更多的样本, 从而增加字典大小, 这些方法需要对网络结构进行特殊的设计, 可能在迁移到下游任务的时候带来额外的困难.

另外一种机制叫做记忆池. 记忆池中包含了数据集中所有样本的表示. 每个mini-batch中, 会随机从记忆池中采样出字典(负样本), 其余样本的表示不参与反向传播, 这可以支持非常大的字典规模. 但是, 记忆池中的每个样本的表示是"它上一次被模型采样"时才更新的, 这意味着同一批采样的key可能来自网络不同训练阶段的表示, 不够一致, 会产生偏差. 某些研究针对记忆池引入了一种动量更新的方式, 但是它更新的是"同一个样本在记忆池中的表示"(保证同一个样本的表示不会因为一次迭代变化过大), 而不是更新编码器参数, 所以这个动量更新方式和MoCo无关. MoCo不需要为每个样本都存储一个特征表示, 因此更加节省内存. 即使数据规模达到数十亿样本, MoCo也能训练, 而记忆池在这种个规模下往往难以维护(需要在池中存储所有样本的特征表示).

前置任务¶

许多前置任务都基于对比学习. 由于本文的重点不在于设计新的前置任务, 所以作者遵循了instance discrimination这个任务. 在这个任务中, 我们考虑一个query和一个key为一个正对, 如果它们来源于同一张图片. 我们对同一张图片进行随机数据增强, 生成两个随机的views作为正样本对. query和key分别由$f_q$和$f_k$编码, 编码器可以是任何卷积神经网络. 下列是MoCo的伪代码. 对于当前的mini-batch, 我们对query和对应的key进行编码, 它们形成了正对, 同时, 负样本来自于队列.

# 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

技术细节¶

他们采用的是ResNet作为encoder. 最后的全连接层有一个固定维度的输出(128-D). 这个输出向量使用L2归一化. 这个是query或者key的表示. 温度$\tau$被设置为$0.07$. 数据增广方法如下: 首先, 从一张随机缩放的图像中裁剪得到224*224的图像块, 然后对图像进行一系列随机变换, 包括随机颜色抖动(color jittering), 随机水平反转(random horizontal fliop)以及随机灰度转换(random grayscale conversion), 这些都可以通过PyTorch的torchvision工具包轻松实现.

洗牌式BN¶

在深度学习任务中, BN通过计算当前mini-batch的每一列的均值和方差, 对激活值进行归一化, 以加速收敛并稳定训练. 但是在多GPU并行训练的时候, 通常每张卡只在自己那部分mini-batch数据上计算均值和方差, 如果各卡的数据分布差异较大, BN在单卡上估计得到的统计量就可能出现较大偏差, 从而影响模型的收敛和性能.

洗牌式BN的核心思路为跨卡随机打乱, 在进行BN的统计量计算之前, 将不同卡上的数据进行某种随机方式的重组, 例如, 卡1上的一部分样本会和卡2上的一部分样本交换, 从而使得每个卡上所拥有的数据分布都相对接近全局分布. 在各卡完成数据洗牌之后, 再各自计算新的mini-batch的均值和方差, 然后进行标准化运算. 完成BN操作后, 需要把刚才打乱过的数据还原到原本的顺序, 再继续后续的正向和方向传播训练.

但是作者提出的Shuffling BN和上述的BN有一些些小小的差异:

作者在做实验的时候发现, 模型似乎在前置任务上"作弊", 找到了一个低损失解. 他们认为这可能是因为batch内的样本之间有一种由BN引起的统计信息的交互而导致信息泄露. 解决方法如下.

Query的mini-batch不打乱, Key的mini-batch打乱. 比如原始的Query mini-batch为$[x_{q_1}, x_{q_2}, x_{q_3}, x_{q_4}]$. $[x_{q_1}, x_{q_2}]$, $[x_{q_3}, x_{q_4}]$分别被发送到GPU1, GPU2上进行BN. 编码完成之后, 得到表示$q_1, q_2, q_3, q_4$; 原始的Key mini-batch为$[x_{k_1}, x_{k_2}, x_{k_3}, x_{k_4}]$, 打乱顺序后, $[x_{k_3}, x_{k_2}]$, $[x_{k_1}, x_{k_4}]$会被分别发送到GPU1, GPU2上进行BN. 编码完成之后, 得到的表示应该是$k_3, k_2, k_1, k_4$, 然后恢复顺序, 得到$k_1, k_2, k_3, k_4$. 最终计算正样本对$(q_i, k_i)$的对比损失. 下面是一个简单的流程图:

Query branch (no shuffle):
    Original batch: [x_{q_1}, x_{q_2}, x_{q_3}, x_{q_4}]
            |
            | Split to multi-GPU
            |   GPU1: [x_{q_1}, x_{q_2}]
            |   GPU2: [x_{q_3}, x_{q_4}]
            v
    BN -> Encoder -> outputs: q_1, q_2, q_3, q_4


Key branch (with shuffle):
    Original batch: [x_{k_1}, x_{k_2}, x_{k_3}, x_{k_4}]
            |
            | Shuffle (example):
            |   -> [x_{k_3}, x_{k_2}] , [x_{k_1}, x_{k_4}]
            | Split to multi-GPU
            |   GPU1: [x_{k_3}, x_{k_2}]
            |   GPU2: [x_{k_1}, x_{k_4}]
            v
    BN -> Encoder -> outputs in the order: k_3, k_2, k_1, k_4
            |
            | Reorder to match original indices
            v
    Final keys: [k_1, k_2, k_3, k_4]


Compute contrastive loss:
    Positive pairs: (q_1, k_1), (q_2, k_2), (q_3, k_3), (q_4, k_4)

注意, GPU1上有$x_{q_1}, x_{q_2}, x_{k_3}, x_{k_2}$, GPU2上有$x_{q_3}, x_{q_4}, x_{k_1}, x_{k_4}$. 使用这种方法, 会使得计算query编码和正样本key编码的batch统计信息来源于两个不同的子集, 这会显著改善"作弊"现象并使得训练从BN中获益.

关于为什么泄露的问题

既然是不同的encoder, 不同的BN层, 为什么统计信息会产生泄露.

在作者的实验中, 他们在自己提出的方法和与之对应的end-to-end方法消融实验中都使用了Shuffled BN. 但是对于memory bank的实验不需要Shuffled-BN. 因为memory-bank中的正样本来自于记忆池中先前的不同mini-batches, 所以不需要Shuffled BN.

He, K., Fan, H., Wu, Y., Xie, S., & Girshick, R. (2020). Momentum contrast for unsupervised visual representation learning (No. arXiv:1911.05722). arXiv. https://doi.org/10.48550/arXiv.1911.05722 ↩

MoCo¹¶

摘要¶

动机¶

信号空间不同¶

前人对比学习的工作¶

字典的理想特征¶

解决字典问题¶

前置任务¶

迁移学习¶

相关工作¶

损失函数¶

前置任务¶

对比学习大框架¶

流程¶

方法论¶

对比学习当作字典查询¶

动量对比¶

队列=字典¶

动量更新¶

和之前工作的关系¶

前置任务¶

技术细节¶

洗牌式BN¶

评论

MoCo1¶

摘要¶

动机¶

信号空间不同¶

前人对比学习的工作¶

字典的理想特征¶

解决字典问题¶

前置任务¶

迁移学习¶

相关工作¶

损失函数¶

前置任务¶

对比学习大框架¶

流程¶

方法论¶

对比学习当作字典查询¶

动量对比¶

队列=字典¶

动量更新¶

和之前工作的关系¶

前置任务¶

技术细节¶

洗牌式BN¶

评论

MoCo¹¶