Momentum Contrast for Unsupervised Visual Representation Learning

He K., Fan H., Wu Y., Xie S., Girshick R. Momentum Contrast for Unsupervised Visual Representation Learning. arXiv preprint arXiv:1911.05722, 2019.

@Article{he2019moco,
author = {Kaiming He and Haoqi Fan and Yuxin Wu and Saining Xie and Ross Girshick},
title = {Momentum Contrast for Unsupervised Visual Representation Learning},
journal = {arXiv preprint arXiv:1911.05722},
year = {2019},
}

概

MOCO的第一个版本, 利用队列queue来更新负样本, 使其能够兼顾大样本的训练并保持负样本之间的一致性. 另外本文详述了各种训练的tricks.

主要内容

首先介绍一下contrastive loss:

$${\mathcal{L}}_q=-\log \frac{\exp (q^T{k}_{+}/T)}{{\sum }_{i=0}^K\exp (q^T{k}_i/T)},\text{\hspace{1em}(1)}$$
其中$T$是temperature, 这是和CPC不同的一点.

从算法流程讲起

首先, 如上图所示, 我们有俩个encoder, 一个是普通的用于生成pair中的第一个特征的encoder, 第二个是momentum encoder用于生成第二个特征的. 二者的初始化权重是一致的, 最大的区别是之后的参数更新方式.

1. 输入一批样本$x$;
2. 通过上述的俩个encoder获得$(q,k)$, 构成正对 { ( p i , k i ) } \{(p_i, k_i)\} {(pi​,ki​)};
3. $q$与队列中的所有的负样本特征构成负对;
4. 计算contrastive loss;
5. 梯度反向传播更新第一个encoder;
6. 第二个encoder利用momentum update.
7. 更新队列(放入$k$, 移除老的负样本特征).

queue： 这个队列里面保存了由第二个encoder提取的特征, 既然是一个队列, 那么其一定有一个更新的过程: 旧的去, 新的来.

为什么要有这个queue, 因为一个非常简单粗暴的做法就是, 把batch_size取得很大, 这样也就能够容纳很多的负样本信息了, 可以其缺点是显而易见, 这个大受限于内存受限于硬件, 而利用队列不断更新负样本信息却没有这个限制.
另外一个做法是, memory bank, 顾名思义, 其保存了整个数据的特征表示, 每一次你需要负样本信息的时候从里面采样就可以了, 但是这个更新的策略是, 所有的特征表示都是该样本最新的特征表示. 这就意味着, 经过多次迭代之后, 不同样本的之间的差异过大(对应不同的encoder), 这种负样本信息缺乏一致性(这篇文章非常非常强调这个一致性).

注: memory bank 似乎是别的方法更新负样本的, 而非直接是样本的特征表示.

momentum update

momentum update是用来更新 momentum encoder 也就是第二个encoder的, 可以从图(b)中看到, 有些机制压根就没有第二个encoder, 这个大概就是一个trick了, 作者在实验中发现如果encoder直接取同一个直接就是fail(那为啥别人可以?).

再来看(a), 它是直接用梯度更新的, 这个也是最直接的方法, 但是需要注意的是, 如果采取这个方法, 那么$x^k$的数量必然受到限制, 也就是说我们的负样本的数量没有办法很大. 但是说实话, 这个限制在使用了queue的MOCO这里是不存在的, 所以MOCO不用这种方法更新的原因大概还是为了之前的一致性吧. 设第一个encoder的参数为${\theta }_q$, 第二个encoder的参数为${\theta }_k$, 则momentum update为

$${\theta }_k\leftarrow m{\theta }_k+(1-m){\theta }_q,\text{\hspace{1em}(2)}$$
其中$m\in (0,1]$作者说取大点比较好, 比如0.999.

这便应了作者说的一致性, 大的$m$导致结果就是momentum encoder更新前后差距不大, 那么queue中特征表示的总体的"差距"(不是指范数距离)也不会太大.

tricks

1. encoder 之后的输出进行了${\ell }_2$ normalization;
2. temperature $\tau$ 设置为0.07;
3. 输入resize为224x224, 并有 random color jittering, horizontal flip, grayscale conversion的augmentation.
4. 直接使用BN似乎有碍于学习到好的特征, 故在每块GPU上单独使用BN且对于key encoder(第二个encoder), 打乱其输入的顺序 (有用吗?).

代码

原文代码