paper:Momentum Contrast for Unsupervised Visual Representation Learning
Code:https://github.com/facebookresearch/moco
Abstrcat
MoCo从对比学习的角度出发,构建了一个包含队列和移动平均编码器的动态字典。这使得能够即时构建一个大型的、一致的字典,从而促进对比无监督学习。贡献如下:
- 一个队列:队列中的样本无需梯度回传,可以放很多负样本,让字典变得很大
- 一个移动平均的编码器:让字典的特征尽可能的保持一致
- 一个大的、一致的字典,有利于 无监督的对比学习 训练
亮点:在分类、检测和语义分割等cv任务证明无监督学习也不比有监督学习差,其中的分类还是做的是Linear protocol测试,就是冻结主干网络的参数,然后只微调最后的分类的全连接层,相当于把backbone当作是一个特征提取器,可以证明backbone学习图片特征的好坏。
1. Introduction
paper将一系列的对比学习方法归纳为一个字典查询的问题building dynamic dictionaries。将负样本图片通过编码器后所得的输出看成是一个特征key,将正样本图片通过另外一个编码器所得到的输出看成是一个query。对比学习本质上,就是希望在字典中找到与query最匹配的那个key,而这个key是正样本通过一些列的数据增强变化获得,所以语义信息应该相同,在特征空间上也应该类似,而与其他的负样本的特征key应该尽可能的远离。整个学习的过程变成了一个对比学习的框架,从而只需要最小化一个对比学习的目标函数就可以。
从动态字典的角度看对比学习,字典应该满足大且保持一致性两个特征。
- Large
让字典更大,可以从连续高维空间做更多的采样。字典 key 越多,表示的视觉信息越丰富,匹配时更容易找到具有区分性的本质特征。 如果字典小、key少,模型可能学到 shortcut 捷径,不能泛化
- Consistent
字典里的 key (k0, k1, k2, …, kN) 应该由相同的 or 相似的编码器生成。如果字典的 key 是由不同的编码器得到的,query q 做字典查询时,很有可能 找到和 query 使用同一个 or 相似编码器生成的 key,而不是语义相似的 key,这是另一种形式的 shortcut solution
现有的对比学习方法某有上面两个问题的限制,所以MoCo提出了动量编码器Momentum Encoder来解决这个问题。
- queue 数据结构
- 剥离 字典的大小 和 显卡内存的限制,让字典的大小 和 模型每次做前向传播的 batch size 的大小 分开。这是因为字典很大(成千上万),意味着要输入很多很多的图片,显卡内存吃不消
- current mini-batch enqueued and the oldest mini-batch dequeued 当前 mini-batch 入队,最早进入队列的 mini-batch 出队,将batch_size的大小与队列的大小进行解耦
- 队列的大小 == 字典的大小,但是每次做 iteration 更新,并不需要更新字典中所有 key 元素的值。
- Momentum Encoder
- 问题:为了保持特征的一致性,字典中的key应该使用一个或者是相似的编码器产生的,但是如果只有一小部分也就是当前的batch是从当前的编码器得到,而之前的key都是用不同时刻的编码去抽取的特征,这就引起不一致的问题。
MoCo中代理任务的选择是instance discrimination,因为其比较简单而且效果也比较好。但是MoCo建立模型比较灵活,可以和很多个代理任务一起使用。
- instance discrimination
在instance discrimination(个体判别)的代理任务中,把一张图片xi随机裁剪出两个子图片,再分别通过一系列的数据增强,那么现在就就得到了两张看起来应该很不一样的图片。但是,这两张子图又是来自于同一张大图xi经过某些变化而得到的,他们的语义信息应该不发生变化,这两张图片就被称作是正样本。然后,instance discrimination这个代理任务认为这个数据集剩下的图片都可以被当做是不相似,每张图片自成一类。那么此时的ImageNet-1k中就不是1000个类别,而是100w多个类别,每个图片都是自己的正样本,其余都是负样本。
无监督学习本来的目的就是在一个很大的无标注的数据集上进行训练,使得模型学习到的特征可以很好的迁移到下游任务上,MoCo做到了。MoCo用 instance discrimination在ImageNet对下游任务进行无监督预训练可以与用其进行有监督预训练打个平手,或者有更好的表现。其在七个计算机视觉任务中表现得不错,部分甚至超越了ImageNet有监督训练的结果。而且,MoCo在更大数量级上进行训练,在性能上还有提升。这满足了无监督学习的期待,更多的数据,更大的模型来实现更强的性能。
MoCo的效果也说明,之前使用ImageNet来做有监督预训练的模型可以考虑用其的无监督预训练模型来替换掉。所用的应用,都可以用MoCo无监督的训练方式训练好的模型,说不定效果更好。
2. Related Work
无监督或者是自监督的学习方法一般包括两个方面:代理任务pretext task、损失函数loss function。其中,pretext task一般指大家不太感兴趣的有实际应用场景的任务(分类,分割,检测),代理任务的提出主要是为了学习一个比较好的特征。损失函数通常可以独立于代理任务进行研究,MoCo侧重于损失函数方面。
2.1 Loss functions
- 生成式损失函数
生成式损失函数的定义是衡量一个模型的输出,和他应该得到的那个固定的目标之间的差距。
例如Auto-encoder一个编码器解码器结构,输入一张原图或者是一张被干扰的图,经过编码器、解码器重构输入的图,可以用L1loss or L2loss,总之衡量的是原图和重构图之间的差异。这就是生成式网络的做法。
- 判别式损失函数
例如一张图片打成有序号的9宫格,给中间的第5格和剩下随机挑一格,预测随机挑的这一格是中间第5格的方位(8个方位可选)。这类似于一个分类任务,因为每一个方格都自带序号,输出分到 8 个方位的哪一类。
- 对比学习损失函数
对比学习损失函数的想法是在一个特征空间里,衡量各个样本对之间的相似性。所要达到的目标就是让相似物体的特征尽量进,不相似物体之间的特征推开得尽量远。
对比学习目标函数的不同是,无论是判别式还是生成式,预测八个位置还是重建整张图,他们的目标都是一个固定的目标。而对比学习的目标在训练的过程中是不停的改变的,其目标其实是由一个编码器抽出来的数据特征所决定的。
- 对抗学习损失函数
对比学习损失函数的想法是衡量两个概率分布之间的差异。
对抗学习的目标函数主要是作无监督数据生成的,后来也有一些对抗性方法用来去做特征学习了。因为,如果你能生成很好很真实的图片,按道理来说,你是已经学习到了这个数据的底层分布,这样模型学习出来的特征应该也是不错的。这个想法感觉挺有趣的,paper引用了两篇文章:
1)Jeff Donahue, Philipp Krähenbühl, and Trevor Darrell. Adversarial feature learning. InICLR, 2017
2)Jeff Donahue and Karen Simonyan. Large scale adversarial representation learning.arXiv:1907.02544, 2019
2.2 Pretext tasks
代理任务发展得多种多样:
- denoising auto-encoders 重建整张图
- context auto-encoders 重建某个 patch
- cross-channel auto-encoders (colorization) 给图片上色当自监督信号
- pseudo-labels 图片生成伪标签
- exemplar image 给同一张图片做不同的数据增广,它们都属于同一个类。
- patch ordering 九宫格方法:打乱了以后预测 patch 的顺序, or 随机选一个 patch 预测方位 eight positions
- 利用视频的顺序做 tracking
- 做聚类的方法 clustering features
2.3 Contrastive learning & pretext task
不同的代理任务是可以和某种形式的对比学习的目标函数配对使用的。
- MoCo 论文里 instance discrimination 个体判别方法 —— examplar based 代理任务很相关
- CPC contrastive predictive coding 用上下文信息预测未来 —— context auto-encoding 上下文自编码
- CMC contrastive multiview coding 利用一个物体的不同视角做对比 —— colorization 图片上色(同一个图片的 2 个视角:黑白 和 彩色)
2.4 Unsupervised or Supervised
与有监督学习工作的对比,有监督学习的整个流程是,现在有一个输入x,然后通过一个模型得到一个y,也就是输出,然后将输出与ground truth进行比较,然后需要一个目标函数去衡量这个比较的结果,这个就是有监督学习的流程。
而对于无监督或者是自监督来说,缺少的就是这个标签ground truth,代理任务的用去就是去生成一个自监督的信号,从而去充当这个ground truth这个标签的信息。那么现在有输出也有一个伪标签信息,就需要一个目标函数衡量差异。
所以,整个流程下来,有监督与无监督的区别,就是使用了特别的代理任务已经去衡量输出与伪标签差异的损失函数。
3. MoCo Method
3.1 Contrastive Learning as Dictionary Look-up
3.2 Momentum Contrast
对比学习是一种在图像等高维连续输入上构建离散字典的方法。字典是动态的,因为key是随机采样的,而且键编码器f k 在训练过程中也会变化。
paper里面提出将字典构建成一个队列,队列里面的元素就是一系列的特征key。在模型训练的过程中,每一个mini-batch就会有新的一批key进入队列中,同时也会有一批老的key移出队列。如此可以将字典的大小与mini-batch大小剥离开,就可以在模型训练过程中使用一个比较标准的mini-batch size(128/256),但是字典的大小既可以变得比较灵活。字典总是表示所有数据的一个采样子集,而维护这个字典的额外计算是可管理的。此外,删除最老的mini-batch size可能是有益的,因为这些key是最过时的,因此与最新的key最不一致。
但是,通过队列的方式没有办法通过反向传播的方式去更新其参数,因为梯度需要传播到队列中的所有样本,但是这些样本又不是同时处理的。一个简单的方式是一次训练结束时,将编码器f q 的参数值赋予编码器f k ,但是这样做的效果不好,可能是降低了这个队列里所有key的特征一致性,因为这些key是由不同时刻的编码器所提取的特征,也就是相当于使用了不同的编码器了,所以不一致。
所以提出了动量更新编码器f k ,公式为:
- end-to-end方式
端对端的意思就是f k , f q 都是可以通过方向传播来更新参数,所以其可以是不同的网络。在这种方式中,字典的大小与mini-batch size的大小是等价的,需要GPU比较大,而且大mini-batch size的优化也比较麻烦。
SimCLR就是这种end-to-end的方式,之所以可以这样坐是因为有TPU,硬件设备好,可以无脑上batch size,SimCLR就选用了8192当作batch size
这种方式的优点是编码器是可以实时更新的,所以导致字典里面的key一致性非常高。缺点就是字典大小与batch size大小挂钩,所以字典大小不能设置过大,否则硬件资源吃不消。
- memory bank方式
在memory bank中,只有一个编码器可以梯度回传更新,对于字典这边是没有单独的编码器的。memory bank就是把整个数据集的特征都存在一起,但是每一个特征只有128维,所以对于ImageNet数据集的100w个特征来说,存储量还是比较小的。所以,每次模型做训练的时候从memory bank去随机抽取很多个key出来当作字典就可以了。
但是memory bank的特征是在不同时刻的编码器得到的,而且这些编码器也是通过梯度回传会快的更新,再更新memory bank,所以得到的特征都缺乏一致性。
而且由于是随机取样,没有选择特征没有新旧的差别,都需要遍历一整便memory bank才进行下一轮训练,使得query的特征与key的特征相差非常远
所以,之前的方法都受限于字典大小与特征一致性这两个方面中的至少一个。所以,MoCo采用了队列的形式去实现了一个字典,从而使得他不像end-to-end方式一样受限于batch size的大小;同时为了提高字典中特征的一致性,MoCo使用了动量编码器。
伪代码如下所示:
# Algorithm 1 Pseudocode of MoCo in a PyTorch-like style # 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 (256x128) k = f_k.forward(x_k) # keys: NxC (256x128) k = k.detach() # no gradient to keys # positive logits: Nx1 (256x1) l_pos = bmm(q.view(N,1,C), k.view(N,C,1)) # q·k+ # negative logits: NxK (256x65536) l_neg = mm(q.view(N,C), queue.view(C,K)) # sum q·ki # logits: Nx(1+K) (256x65537) 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 # bmm: batch matrix multiplication; mm: matrix multiplication; cat: concatenation.
3.3 Shuffling BN
用了BN之后,很有可能会造成batch里的这些样本中间信息会泄漏,因为BN要算这些样本的running mean和running variance,也就是说他能通过这些泄漏的信息很容易的去找到那个正样本,而不需要去学一个真正好的模型。
而BN这个操作大部分是在当前的这个GPU上面计算的,所以在做多卡训练之前先把样本的顺序打乱,再送到所以的GPU上去,算完了特征再把顺序回复去算最后的loss,这样的操作loss是不会改变的,但是每隔GPU上的bn计算就变了,就不会再存在这种信息泄露的问题。
4. Experiments
paper在测试分类效果的时候,是将backbone冻结住的,只训练一个分类头,也就是相当于是一个特征提取器。其中,使用grid search搜索出最好的learning rate是30,这个比较诡异,这说明无监督学习学到的特征分布与有监督学习学到的特征分步是非常不一样的。
- 三种对比学习流派的性能对比:
MoCo的硬件资源要求较低,但性能最好,可拓展性也最好。
- 动量更新带来的提升:
所以一个变化缓慢的编码器对对比学习是有好处的,因为保留了特征一致性。
- ImageNet上使用linear classification protocol 的对比结果:
- MoCo在下游任务的迁移效果:
无监督学习最主要的目标,就是去学习一个可以迁移的特征。用ImageNet去做有监督的预训练,最有用的时候就是在下游任务上做微调,可以用预训练好的模型去做这个模型的初始化。从而当下游任务只有很少的数据时,也能获得比较好的效果。
- pascal voc效果:
- coco效果:
ps,这里补充两点:
- Normalization
当用学习好的参数做一个特征提取器时,最好的学习率是30,但是对每一个下游任务上都做一次Grid Search来找出其最佳的学习率的大小,这样就失去了无监督预训练的意义,所以需要对整个模型就行微调归一化操作。
这里使用了sync BN,其会把多机训练上,所以GPU的Batch Norm的统计量都合起来,算完running mean和running variance再去做BN层的更新,这样会让特征归一化做得更加彻底一点,模型也会更稳定。这样就可以拿有监督的参数来做微调。
- Schedule
当下游任务的数据集足够大的时候(比如COCO),可以不需要预训练,直接从随机初始化开始从头训练,最后的效果一样可以很好,只是需要的训练时长可能要比预训练的时长要多(in contrast to 6×∼9×)。
总结:
MoCo在一系列的任务和数据集上效果很好,但是还是存在一点瑕疵。比如使用了1000倍数据规模的提升, 但MoCo 性能的提升不高,可能大规模数据集可能没有完全被利用,可以尝试开发其它的代理任务 pretext task。而且除了 instance discrimination 代理任务,类似 NLP 的代理任务 masked auto-encoding可以在cv领域进行(有了后来的MAE)。
MoCo 设计的初衷:去构造一个大的字典,从而让正负样本能够更有效地去对比,提供一个稳定的自监督信号,最后去训练这个模型。
ps:MoCo系列工作的泛化会比其他的对比学习模型好
参考资料:
1)视频讲解:https://www.bilibili.com/video/BV1C3411s7t9?spm_id_from=333.999.0.0
2)视频笔记:https://www.bilibili.com/read/cv14463867?from=note