论文：Momentum Contrast for Unsupervised Visual Representation Learning
代码：https://github.com/facebookresearch/moco
出处：FAIR | 何凯明 | CVPR2020
时间：2020.03
MOCO 分类：
MOCO 下游任务：可以很好的迁移到下游任务
MOCO 的主要特点：
文章思想分析：
本文的方式可以看做为了建立大的字典来提升编码空间的信息丰富程度，也就是样本能见到足够多的负样本，从而学习到它们是我的对立样本，是和我不一样的。但建立大的字典又非常依赖 batch size（因为负样本是 batch size 内的样本，要建立大的字典就表示当前 batch 内的样本要尽可能的多），大的 batch size 又对显存要求很大，所以就陷入了如何在不依赖 batch size 的情况下来建立大的字典
那么建立大的字典就借用了队列的方式，每个 epoch 提取到的特征送入队列中，同时队列中还保存了之前很多 epoch 保留的特征，但是，不同 epoch 的编码器的权重是更新的，没法保证 query 和 key 是在一个绝对一致的编码空间的，这样的对比其实意义不大，因为标准都不一样了怎么判断相似度呢，所以作者就提出了一个缓解的方法，就是以动量的形式来更新编码器的权重，这样虽然也没法绝对的保证每次对比的编码空间是一致的，但总能拉小这种不一致性吧，所以就有了动量更新 query 编码器的形式
当预训练数据集换成 ins（1亿） 时，提升相比 imagenet 的 100w 数据有点小，所以作者觉得大规模数据集还没有很好的利用起来，这可能是源于没有一个好的代理任务而导致的
MOCO 还提出一个讨论，是不是能使用除了 instance discrimination 之外的代理任务，如 masked auto-encoding 的方式，将图像的一块挡起来，来和 MOCO 一起训练，这个想法就是后面的 MAE
无监督表达学习在自然语言处理方面已经有了很成功的应用，如 GPT 和 BERT，但当时（2020年左右）在计算机视觉中还是监督学习占主流。
其主要原因在于两者的特征信号不同，语言任务的数据是在离散数字空间（如单词），可以很方便的建立 tokenized 词典，可以很好的把一个单词对应成一个特征，一旦有了这个词典呢，无监督学习可以基于它很好的展开。因为可以简单的把字典里的所有 key 想象成一个类别，就变成了一个有监督学习的范式了，所以还是有一个类似标签一样的东西来帮助学习，所以 NLP 上无监督学习很容易建模。
但视觉任务的数据是在连续的高维空间，不像单词有很强的语义信息，没有那么容易建模，所以并不适合建立一个字典，没有字典的无监督学习很难建模
近期有一些使用对比学习的无监督方法，取得了较好的效果，都可以被归纳成一种方法，就是在构造动态的字典
当时的主流方法是怎么做的：
如何理解：
因为将对比学习转换成了字典查询的问题，所以本文中很少出现正负样本等字样，而是使用 query 表示 anchor （x_key），key 表示其他样本（x_query），对应的特征使用 k 和 q 表示。
那么从字典查询的角度来理解对比学习呢，作者就说期望构建的字典有如下两个特征：
基于上面的两个期望的特征，作者提出了 MOCO，目的就是构建一个大且连续的字典，如图 1 所示，关键的贡献在于 momentum encoder 和 queue
动量对比学习： 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
代理任务 instance discrimination （个体判别）介绍：认为每张图片都是自己的类
什么是对比学习：
对比学习是无监督学习/自监督学习，一般的研究点落在下面两个地方：
MOCO 的落脚点是在目标函数上，提出的这个大的字典主要会影响后面计算 loss，对比学习的 loss 不同于生成/判别式的目标都是固定的目标，而对比学习的目标是会变化的，在训练的过程中会变化，也就是目标是由编码器中抽出来的特征决定的，主要是在同一个特征表达空间中衡量样本之间的相似程度。
对比学习可以被看做为了字典查表任务训练一个 encoder 的任务：
那么训练这类任务的时候 loss 要有什么属性呢：
在看对比学习 loss 之前先看一下交叉熵 loss：
假设分类任务的 gt 是 one-hot 向量，那么一般会对最后一层特征的输出做 softmax 后和 gt 来做交叉熵损失。
假设第 i 个输出为 z i z_i zi，那么经过 softmax 后为： e z i ∑ i = 0 k e z i \frac{e^{z_i}}{\sum_{i=0}^ke^{z_i}} ∑i=0keziezi，交叉熵结果为：
− l o g e z i ∑ i = 0 k e z i -log\frac{e^{z_i}}{\sum_{i=0}^ke^{z_i}} −log∑i=0keziezi
这里的 k 指的是训练的时候有多少类别，比如 1000 类就是 1000，是固定的数字
但对比学习，类别数 k 将会是非常大的数，有多少图片就有多少类，计算复杂度非常高，所以就出来了 NCE loss
因为之前类别太多，没法计算，NCE 就简化成二分类问题，一个是数据样本，一个是噪声样本，将数据样本和噪声样本做对比就可以了
但如何将整个数据集中剩下的所有图片当负样本，NCE 就解决了类别多的问题，但没有解决计算复杂度高的问题
如何让 loss 计算的更快一点呢，就是取近似，与其计算整个数据集中剩下的数据的 loss，不如直接选择一些作为负样本来计算 loss 就可以了，这也就是 estimation 的含义。
如果选取样本少，近似性会越差，如果选的越多，那么近似性会越高，效果会更好，所以 MOCO 希望字典足够大，足够大的字典就能更好的近似整个空间
所以 NCE loss 就把超级多分类的问题，变成了一系列的二分类问题，从而还可以使用 softmax
这里的 InfoNCE 就是一个变体，觉得如果把数据看成二分类，可能对模型学习也不是很友好，毕竟在那么多噪声样本里，大家也不是一个类，还是看成更多的类比较好
InfoNCE[46] 公式如下
作者之所以提出动量对比是因为作者认为使用大的字典能够引入更丰富的负样本，但大的字典中不同 batch 提取特征的模型参数是一直在更新的，这样就难以使用梯度反传的方式来更新 key encoder，之前有些方法使用 query encoder 的参数来当做 key encoder 的参数，但这样就会导致特征不连续，因为 encoder 是随着训练的进行会变化的，所以会导致特征不在相同的特征空间内。
1、Dictionary as a queue：如何把一个字典看成一个队列
本文思想的核心在于将字典当做数据的队列，这样的好处就是能够对前面 batch 的编码特征进行重复使用，可以将字典大小和 batch 大小解耦开来，字典的大小可以远远大于 batch 的大小，且大小可以设置为可调节的超参数。
字典中的样本可以被逐步的替代，当前 batch 的特征入队，最老的 batch 的特征出队
这个字典中的特征其实是整个数据集中的子集，前面说过计算 loss 的时候不是在整个数据集中来计算 loss，而使用这个字典中的特征来计算，其实就可以看做一个近似
而且字典的长度基本不会影响训练时间，且队列的先进先出的特性，能让新的特征加入计算，更有时效性
2、momentum update：如何用动量的方式来更新 key encoder 的参数
使用队列可以使得字典变大，但不同 batch 提取特征的模型参数是一直在更新的，这样就难以使用梯度反传的方式来更新 key encoder，因为需要给队列中的所有 samples 传递梯度。
一个简单的做法是直接从 query encoder 来复制得到 key encoder，但效果不好
作者猜测这不好的效果来源于 encoder 剧烈的变化会降低 key 表达特征的一致性，所以提出了动量更新的方法。
所以，动量 encoder 更新 θ k \theta_k θk 的公式如下：
3、Relations to previous mechanisms：MOCO 和之前的方法到底有什么不同呢，又是如何使用动量的方式解决问题的
MoCo 和之前两种方法的对比见图 2，主要的差别就在于字典的尺寸和参数更新的一致性。
图 2a 是端到端的梯度反传方法：牺牲字典大小，保证了一致性
就是每次都通过梯度回传来更新 encoder k 的参数，这里 q 和 k 可以相同也可以不同，之所以可以相同是因为这种训练方式，每次的正负样本都是来源于当前训练的 batch 中，这些样本是高度一致的，因为是使用相同的编码器得到的，这也就意味着字典的大小和 batch 的大小是耦合的，会被 GPU 显存限制，且大的 batch size 也很难收敛
图 2b 是 memory bank 方法：牺牲了一致性，提高了字典大小，但扩展性也不好
memory bank 包括所有数据集中 sample 的特征表达，这里只有一个 query encoder，key 是不需要 encoder 的，而是从存储好特征的 memory bank 中随机采样的，如果是 ImageNet 就有 128w 个特征，每个特征 128 维，所以只需要 600M 空间来存储这些 key，查询也很高效，也能支持很大尺寸的字典。而且他的做法会更新 memory bank 中的样本特征，做法就是每次选择 m 个样本来作为负样本计算 loss，query encoder 就会更新权重，然后会使用更新后的 encoder 对这 m 个负样本生成新的特征，这样的过程会缺少一致性，每次的 encoder 都会被更新，都是不同时刻得到的，就意味着存储的特征缺乏一致性。
pretext task 也叫前置任务或代理任务，也就是该任务不是目标任务，但执行该目标可以更好的执行目标任务，本质就是迁移学习。
参考论文 [61]，作者也将来源于同一张图像的 query 和 key 当做一组 positive pair，其他都是 negative pair
query 和 key 都被其各自的 encoder 进行编码 f q f_q fq 和 f k f_k fk，编码器可以是卷积神经网络。
Algorithm 1 展示了 MoCo 的伪代码
代理任务：对 x 做两次数据增强，得到 x_q 和 x_k，作为一对正样本对
batch size 为 256，queue 长度为 65536，编码器输出特征为 128 维，也就是每个样本被编码为 128 维的特征
对 x_q 使用 f_q 提取特征得到 [256,128] 特征 q，x_k 使用 f_k 提取特征得到 [256,128] 特征 k，k 不进行梯度回传
计算 query 和正样本的相似度 logits：每个 q 对应一个 k，得到 256x1，即每个 query 样本和它对应的正样本 k+ 的相似度的
计算 query 和负样本的相似度 logits：每个 q 对应 K 个负样本（K=65536），得到 256x65536 的特征，表示每个 query 和负样本相似度
计算完 query 和正样本/负样本的相似度后，concat 起来，得到 256x65537 维特征向量，表示每个 query 和其他样本的相似度，第一列就是正样本，最好的期望就是第一列全为 1，后面所有列都为 0 ，那就学的非常好了
真值：作者使用 256 维全 0 向量表示真值 gt，放到真正的分类任务上来说，就是说每个样本学习的真正类别就是 0 类，也就是说在 0 位置上的 one-hot 值为 1，其他位置（其他 65536 个位置）上的 one-hot 值都为 0。
为什么这样来定义真值呢，因为 query 和正样本的相似度矩阵 256x1 后面 concat 了负样本相似度矩阵，也就是 query 和正样本的相似度在整个矩阵的第一列，也就是位置 0，所以对于正样本，如果找对了 key，在分类任务上找到的正确类别就是类别 0
再想想对比学习学的是啥，学的是样本的相似度，而不是真正的类别，再换句话说就是，要在 0 位置上的相似度为 1，因为 0 位置上永远是正样本。
梯度回传更新后，再使用动量方法更新 key encoder
然后更新 queue
1、ImageNet-1M(IN-1M)
约有 1.28 million 数据，共 1000 个类别
2、Instagram-1B (IG-1B)
约有 1 billion 数据，来源于 Instagram
优化器：SGD
IN-1M ：
IG-1B：
1、不同 loss 的对比
作者冻结训练好的特征（在 IN-1M 上无监督预训练），在后面接了一个 linear classification，就是一个全连接+average pooling 的操作，训练了 100 个 epoch，只训练这个分类器，也表现的较好。表明 MOCO 可以很好的迁移到下游任务，有效的弥补了有监督和无监督的鸿沟。
但作者说使用 grid search 的方式发现最优的初始化学习率是 30，所以这是很奇怪的现象，暗示无监督学习到的分布和有监督学习的分布是非常不同的。
对比不同对比学习 loss 的结果见图 3，横坐标表示负样本的数量，纵坐标是 top-one 准确率：
2、Momentum 的效果
动量从小到大带来的影响，使用 0.999 的效果最好，变小的时候下降的明显，尤其是去掉动量后整个模型都无法收敛，loss 一致是震荡。非常有利的证明了作者的论点。
K=4096
3、ImageNet 分类上的结果
所有的方法都是在 linear protocal 方式下进行的，都没有更新 backbone
上面是不用对比学习的，下面是用对比学习的
而且在无监督里边，模型的大小是很重要的，所以使用的网络结构和参数量对比也很重要。
之前的方法使用 R50 最好的效果是 58.8，MOCO 达到了 60.6，高两个点
不对模型做限制，变宽 channel 后，得到的效果也很好
4、在 PASCAL VOC 上的目标检测效果
要验证下游任务能否有好的效果才是更重要的
对比使用 ImageNet 有监督训练和 MOCO 无监督训练的模型来初始化模型做不同的下游任务
COCO 文章中有证明说，当使用较大的数据集如 COCO 的时候，随机初始化然后训练 6x~9x （1x = 12 epoch）的模型仍然可以达到很好的效果，那么也就没法证明预训练模型的效果了
但是，如果训练时间短的话呢，预训练模型还是有用的，所以作者只对比 1x~2x 的训练效果
检测器：Faster RCNN
端到端、memory bank、moco 的对比，效果都是 moco 更好，且超越了有监督训练结果
5、在 COCO 上的检测和分割
模型：Mask RCNN（with FPN）
下游任务：
总结：
训练无监督学习的方式：
# Copyright (c) Meta Platforms, Inc. and affiliates.
# This source code is licensed under the MIT license found in the
# LICENSE file in the root directory of this source tree.
import torch
import torch.nn as nn
class MoCo(nn.Module):
"""
Build a MoCo model with: a query encoder, a key encoder, and a queue
https://arxiv.org/abs/1911.05722
"""
def __init__(self, base_encoder, dim=128, K=65536, m=0.999, T=0.07, mlp=False):
"""
dim: feature dimension (default: 128)
K: queue size; number of negative keys (default: 65536)
m: moco momentum of updating key encoder (default: 0.999)
T: softmax temperature (default: 0.07)
"""
super(MoCo, self).__init__()
self.K = K
self.m = m
self.T = T
# create the encoders
# num_classes is the output fc dimension
self.encoder_q = base_encoder(num_classes=dim) # 构建 query encoder
self.encoder_k = base_encoder(num_classes=dim) # 构建 key encoder
if mlp: # hack: brute-force replacement
dim_mlp = self.encoder_q.fc.weight.shape[1]
self.encoder_q.fc = nn.Sequential(
nn.Linear(dim_mlp, dim_mlp), nn.ReLU(), self.encoder_q.fc
)
self.encoder_k.fc = nn.Sequential(
nn.Linear(dim_mlp, dim_mlp), nn.ReLU(), self.encoder_k.fc
)
for param_q, param_k in zip(
self.encoder_q.parameters(), self.encoder_k.parameters()
):
param_k.data.copy_(param_q.data) # initialize # 初始时是将 query 的参数拷贝到 key encoder 作为初始参数
param_k.requires_grad = False # not update by gradient
# create the queue
self.register_buffer("queue", torch.randn(dim, K))
self.queue = nn.functional.normalize(self.queue, dim=0)
self.register_buffer("queue_ptr", torch.zeros(1, dtype=torch.long))
@torch.no_grad()
def _momentum_update_key_encoder(self): # key encoder 的参数不进行梯度反传的更新，而是使用动量法来更新
"""
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.0 - self.m)
@torch.no_grad()
def _dequeue_and_enqueue(self, keys):
# gather keys before updating queue
keys = concat_all_gather(keys)
batch_size = keys.shape[0]
ptr = int(self.queue_ptr)
assert self.K % batch_size == 0 # for simplicity
# replace the keys at ptr (dequeue and enqueue)
self.queue[:, ptr : ptr + batch_size] = keys.T
ptr = (ptr + batch_size) % self.K # move pointer
self.queue_ptr[0] = ptr
@torch.no_grad()
def _batch_shuffle_ddp(self, x):
"""
Batch shuffle, for making use of BatchNorm.
*** Only support DistributedDataParallel (DDP) model. ***
"""
# gather from all gpus
batch_size_this = x.shape[0]
x_gather = concat_all_gather(x)
batch_size_all = x_gather.shape[0]
num_gpus = batch_size_all // batch_size_this
# random shuffle index
idx_shuffle = torch.randperm(batch_size_all).cuda()
# broadcast to all gpus
torch.distributed.broadcast(idx_shuffle, src=0)
# index for restoring
idx_unshuffle = torch.argsort(idx_shuffle)
# shuffled index for this gpu
gpu_idx = torch.distributed.get_rank()
idx_this = idx_shuffle.view(num_gpus, -1)[gpu_idx]
return x_gather[idx_this], idx_unshuffle
@torch.no_grad()
def _batch_unshuffle_ddp(self, x, idx_unshuffle):
"""
Undo batch shuffle.
*** Only support DistributedDataParallel (DDP) model. ***
"""
# gather from all gpus
batch_size_this = x.shape[0]
x_gather = concat_all_gather(x)
batch_size_all = x_gather.shape[0]
num_gpus = batch_size_all // batch_size_this
# restored index for this gpu
gpu_idx = torch.distributed.get_rank()
idx_this = idx_unshuffle.view(num_gpus, -1)[gpu_idx]
return x_gather[idx_this]
def forward(self, im_q, im_k):
"""
Input:
im_q: a batch of query images
im_k: a batch of key images
Output:
logits, targets
"""
# compute query features
q = self.encoder_q(im_q) # queries: NxC
q = nn.functional.normalize(q, dim=1)
# compute key features
with torch.no_grad(): # no gradient to keys
self._momentum_update_key_encoder() # update the key encoder
# shuffle for making use of BN
im_k, idx_unshuffle = self._batch_shuffle_ddp(im_k)
k = self.encoder_k(im_k) # keys: NxC
k = nn.functional.normalize(k, dim=1)
# undo shuffle
k = self._batch_unshuffle_ddp(k, idx_unshuffle)
# compute logits
# Einstein sum is more intuitive
# positive logits: Nx1
# 爱因斯坦求和约定 einsum：第一个参数：equation 中的箭头左边表示输入张量，以逗号分割每个输入张量，箭头右边则表示输出张量
# 第二个参数：表示实际的输入张量列表
l_pos = torch.einsum("nc,nc->n", [q, k]).unsqueeze(-1) # 这个 batch 中有 N 个样本，每个样本有一个正样本 logits 得分
# negative logits: NxK
l_neg = torch.einsum("nc,ck->nk", [q, self.queue.clone().detach()]) # 这个 batch 中有 N 个样本，有 K 个负样本，共有 NxK 个相似度得分
# logits: Nx(1+K)
logits = torch.cat([l_pos, l_neg], dim=1)
# apply temperature
logits /= self.T
# labels: positive key indicators
labels = torch.zeros(logits.shape[0], dtype=torch.long).cuda()
# 大小为 batch size 的一组全 0 向量，作为 label
# labels 中的元素实际上意味着在进行 CrossEntropyLoss 计算时，标签为 1 的 ground truth 的索引是多少，而不是 gt 为 0
# dequeue and enqueue
self._dequeue_and_enqueue(k)
return logits, labels
# utils
@torch.no_grad()
def concat_all_gather(tensor):
"""
Performs all_gather operation on the provided tensors.
*** Warning ***: torch.distributed.all_gather has no gradient.
"""
tensors_gather = [
torch.ones_like(tensor) for _ in range(torch.distributed.get_world_size())
]
torch.distributed.all_gather(tensors_gather, tensor, async_op=False)
output = torch.cat(tensors_gather, dim=0)
return output