从MoCo入门自监督学习与对比学习

2022-11-18 15:01:25 浏览数 (3)

  • 谈到CV领域的自监督学习,离不开的就是恺明大神的MoCo和Hinton团队的SimCLR,之所以选择MoCo来进行解读,主要是MoCo相较SimCLR对GPU友好太多,可以让普通学生也有机会能去跑一下对比学习的实验。本文以MoCo论文原文自顶向下的方式进行阐述(主要follow朱毅老师解读的角度),希望让读者阅读完除了对MoCO算法理解,对自监督和对比学习领域也能入门。

Introduction

最近有一些基于对比学习的无监督学习的工作取得了不错的效果,这类对比学习方法的本质上是构造一个动态的字典。我们先解释一下对比学习。

举个例子,比如现在我们有3张图片,分别是x1、x2、x3。(1)定义正负样本:先对x1做2次随机裁剪,得到x11和x12两张新的图片,我们会先选择一张图片作为anchor基准点,比如我们选择x11为anchor,那么x12就为正样本positive,即相对于x1来说,x2是正样本(因为他们都是从同一张图片裁剪出来的)。剩下的x2和x3就为负样本negetive(因为他们根本就不是一张图片)。(2)提取特征:定义好正负样本之后,接下来就是将样本通过编码器得到对应的特征,基准样本x11会通过基准编码器E11得到一个基准特征f11,正样本x12会通过正样本编码器E12得到特征f12,负样本x2和x3也会通过正样本编码器得到特征f2和f3.之所以负样本和正样本用一个编码器是因为正负样本都是相对于基准样本来说的。(3)对比学习:对比学习就是让(基准样本特征f11,正样本特征f12)这样的特征对在特征空间里尽可能相近,让(基准样本f11,负样本f2 f3)在特征空间里尽可能远离,这样我们就理解了什么叫对比学习。

那为什么对比学习可以理解为一个动态的字典,是因为我们可以将基准样本特征f11看做是一个query,然后正负样本的特征当做是key,那么对比学习就转换为一个字典查询的问题,具体来说对比学习就是去训练一些编码器,从而进行字典查找的任务,查找的目的就是将已经编码好的query(f11),尽可能和他匹配的那个key的特征相似(f12),和其他负样本的key远离(f2 f3),因此就变成了一个最小化对比学习损失的问题(具体loss后面会介绍)。

从将对比学习当做动态字典的角度来看,作者认为要有好的结果这个字典需要有2个特性,第一个就是这个字典必须要大,第二个就是在训练的过程中要保持一致性。(1)字典越大,就能更好的从这个连续高维的视觉空间抽样,因为字典里的key越多,所能表示的视觉信息视觉特征就越多,那么当你拿query去和key做比对的时候,你就更能学到将物体分开,更本质的特征。(如果特征太少,可能模型就走了捷径,通过某个不是本质的特征来进行分类)。(2)一致性指的是字典里的key,都应该用相同或相似的编码器去产生得到的,这样你跟query去对比时,才能保证对比时尽量一致,否则如果不同的key使用不同的编码器得到的,那么query可能会选择和它本身使用相同或相似编码器得到的key,这样也就是变相走了捷径。但是目前所有已有的对比学习方法,都至少不满足其中一个特性。(后面具体展开)

MoCo本质上是希望给无监督学习构造一个又大又有一致性的字典,具体架构如上图所示。首先我们有一个query的图片,还有一些key的图片,这些图片分别通过各自的编码器得到一系列特征,然后query的特征去跟所有key的特征做对比,最后用对比学习的loss去训练整个模型。和前面的不同主要在于queue和momentum encoder。

为什么我们要用一个队列去存储key的特征而不是用字典呢?主要是受限于显卡的内存。因为如果字典很大,那我们对应就要输入很多图片,那显卡的内存肯定是吃不消的,所以我们需要想一个办法,让这个字典的大小和每次模型去做前向过程时的batch size大小剥离开,于是作者就想到一个巧妙的方法,即用队列这种数据结构来完成这种任务。具体来说就是这个队列可以很大,但是我们每次更新这个队列都是一点一点进行的,即我们现在读入一个mini-batch的图片,那么最早的那个mini-batch就移出队列,这样就把batch size和队列大小分开了,这样队列的大小可以设为非常大,因为每次更新只更新这个队列里面mini-batch大小的特征梯度。

说完了MoCo的架构,接下来该说应该选用什么样的代理任务来充当自监督信号进行模型的训练了。因为MoCo只是建立中间模型的一个方式,它只是为对比学习提供了一个动态的字典,本文选择了一个比较简单的instance discrimination任务作为代理任务。instance discrimination即一个query和一个key,是同一个图片不同的视角(如裁剪),那么这个query就能和key配上对。

无监督的表征学习,最主要的目的就是在一个很大没有标注的数据集做完预训练后,这个预训练好的特征是可以直接迁移到下游任务上的。MoCo在7个下游任务中都能超过用ImageNet做预训练的有监督学习方式!另外作者为了实验完整性,又做了一个新的实验,还把MoCo在facebook自己的一个十亿数据上做了预训练,结果也还能提升,证明MoCo可以在真实图片,并且有亿级规模图片的数据集上做的很好。因此这些结果证实了MoCo可以在很多CV任务上把无监督学习和有监督学习的坑填平,甚至可以在真实应用中取代大家之前一直使用的用ImageNet预训练的模型。

Related Work

无监督学习或者自监督学习,一般就是在代理任务(pretext task)和目标函数上做文章。因为无监督学习区别于有监督学习最大的区别就是他的数据没有人工标注的label,那这个时候我们就可以使用代理任务去生成“label”这种自监督的信号,然后再用一些特别的目标函数去梯度反传训练模型,因此related work从这两块展开。

代理任务一般指的是那些大家不太感兴趣的任务,即不是分类 分割 检测这种有实际应用的任务,这类任务的提出主要是为了学习一个好的特征。而MoCo主要就是从目标函数上下手,他提出的又大又一致的字典主要影响的是后面InfoNCE这个目标函数的计算。

目标函数

一般无监督学习分为生成式和判别式。生成式的无监督学习的损失函数就是去衡量模型的输出和应该得到的那个固定的目标之间的差异,比如用auto-encoders,输入一张低分辨的图,通过编码器解码器还原出一张高分辨率的图,这里你可以使用L1 loss或L2 loss衡量原图和新建图之间的差异;

判别式无监督学习有一些不同的做法。比如上图Unsupervised Visual Representation Learning by Context Prediction这篇文章的做法,如果有一张图片打成九宫格,现在将中间这个格子给你,再随机从剩下的格子中挑一格给你,代理任务就是预测随机挑选的格位于给你的格子的哪个方位,即把代理任务转换为8分类的任务。再比如下图Colorful Image Colorization这篇文章,他的代理任务就是输入图像的灰度图,预测图片的色彩

除了判别式和生成式这种常见的目标函数,还有对比型的目标函数和对抗性的目标函数。对比学习的目标函数主要是去一个特征空间里,衡量各个样本对之间的这个相似性,达到的目标就是让相似物体特征尽量拉近,不相似物体之间的特征推开尽量远。对比学习和前两种目标函数最大的不同是生成式(重建整张图)和判别式(预测8个位置)的目标都是一个固定的目标,但是对比学习的目标在训练过程中是不停的改变的,目标是由一个编码器抽出来的这个数据特征而决定的,也就是MoCo这篇文章的字典。

还有一种方法就是对抗型的目标函数,他主要衡量的是两个概率分布之间的差异,一开始主要是用来做无监督的数据生成的。但是后来也有一些对抗性的方法拿来去做特征学习了,因为大家觉得要是能生成很好的图片,那应该是学到了这个数据的底层分布,那这样模型学出来的特征也是不错的。

代理任务

代理任务的类型就比较多了,比如有的文章的代理任务是重建整张图,有的任务是重建重建整个patch,还有很多代理任务是去生成一些伪标签,比如exemplar image,即给一张图片做不同的数据增广,这些生成的图片都是同个类等等。

对比学习和代理任务

不同的代理任务是可以和某种形式的对比学习目标函数配对使用,比如MoCo使用的instance discrimination的方式,就和examplar-based代理任务很相关

Method

对比学习和字典查找

我们从softmax开始讲起。softmax表达式大家都很了解,即如下式所示

softmax(x)_i = frac {e^{x_i}}{sum _j e^{x_j}}

如果在有监督学习范式下,即我们有一个one-hot向量,那我们就可以在softmax前加上一个-log,就成为了cross entropy loss

crossentropy(x)_i = -logfrac {e^{x_i}}{sum _j e^{x_j}}

这里的k指的是这个数据集一共有多少类别,比如ImageNet的话就是1000类,是一个固定的数字。对比学习理论上也是可以使用交叉熵去学习,但是实际上行不通。比如我们使用instance discrimination这个代理任务去当自监督信号的话,那这里的类别数k将会是一个巨大的数字,比如在ImageNet下这个k就不是1000了,而是128万了,即有多少图片就有多少类。但是softmax操作在有这么多类别时是没法计算的,而且这里还有exp的操作,当这个向量维度是几百万的时候,计算复杂度是相当高的,如果每个前向都要算一次这样的loss,训练时间会被大大延长。因此NCE(noise contrastive estimation) loss就被提出来了。

前面提到因为类别太多造成无法计算,那么NCE就简化成了一个二分类问题。一个类别就是数据类别data sample,一个类别就是噪声类别noise sample,每次只要拿数据样本和噪声样本对比即可。但是如果还是把整个数据集剩下的图片都当做是负样本,但是计算复杂度还是没有降下来,那还有什么办法可以加速计算呢?只能取近似了。即与其在整个数据集上算loss,不如从这个数据集里去选一些负样本算loss就可以了。但是我们也知道,按照一般的规律如果这里选择的样本很少,效果就没那么近似了,所以MoCo一直在强调字典要足够大,这样做近似才能足够准。

InfoNCE就是NCE的一个简单变种,他认为如果只把问题看作是二分类,可能对模型学习没那么友好,毕竟在那么多噪声样本里大家很可能不是一个类,所以还是看做一个多分类的问题比较合适。因此InfoNCE就被提出了:

Momentum Contrast

把字典看作是队列

在intruduction我们提到,McMo把字典看作是一个队列,每次有一批新的key进来,就有一批老的key移出去。所以作者说用队列的好处,可以让我们重复用那些已经编码好的key,而这些key是从之前的那些mini-batch里得到的,这样就可以把batch的大小和字典的大小剥离开了,就可以在模型的训练过程中使用一个比较标准的mini-batch size,一般就是128或者256,但是字典大小变得非常大,可以当做一个超参数一样单独设置。这个字典是所有数据的一个子集,我们前面提过InfoNCE算loss的时候,只是取一个近似,而不是在整个数据集上算一个loss。另外使用队列这个数据结构,可以让维护这个字典的计算开销非常小。

最后作者又重复了一下为什么要使用队列这个数据结构的原因,使用队列是因为队列里有这个先进先出的特性,这样每次移出队列,都是最早的那些mini-batch,这样对对比学习来说是很有利的,因为从一致性的角度来说,最早计算的那些mini-batch的key是最过时的,即跟最新算的mini-batch是最不一致的(因为最早产生的mini-batch的更新 是用最老的编码器产生的)。

动量更新

这里introduction介绍过了就不重复了。熟悉RL的小伙伴估计已经看出了这和我们double DQN的思想非常一致。。。因为用了动量更新的方式,所以虽然队列里的key都是由不同的编码器得到的,但是因为编码器之间的区别太小,所以key的一致性还是很强的。在作者的实验里,用了相对比较大的动量,就是0.999,效果比0.9好得多。因此说明了动量更新对队列一致性的重要性。

和之前方法的对比

如上图a所示,之前的方法第一种类型为端对端学习。端对端学习顾名思义即query和key的编码器都可以通过梯度回传来更新模型参数,这里两个编码器是同一个编码器,用的数据也都是一个mini-batch的数据,做一次forward就能得到所有样本的特征,而且这些样本特征因为是通过同样编码器得到,所以也是高度一致的。但是这种方法的局限性就在于字典的大小,因为在端对端的学习框架中,字典的大小和mini-batch size的大小是等价的,那如果我们想要一个很大的字典,里面有成千上万个key的话,也就意味着mini batch的大小也是成千上万的,这样难度就很高的,因为现在的GPU一般塞不下这么大的batch size的。就算你的GPU能塞下这么大的batch size,大batch size的优化也是一个难点,如果处理不得当,模型是很难收敛的。因此端对端学习的方式就是受限于字典的大小。另一篇很出名的SimCLR就是这种端对端训练的方式,当然他用了更多的这个数据增强,而且提出了在编码器之后再用一个projector会让学到的特征大大变好,他们之所以这么做,因为google有tpu,tpu内存大,所以可以用大的batch size,他们就是用了8192为batch size,这样最后负样本的规模就可以做端到端的学习了。

看完端对端的学习后,我们发现他的优点在于编码器是可以实时更新的,所以他字典里的一致性是非常高的,但他的缺点在于字典的大小就是mini-batch的大小,导致这个字典不能设的过大,否则硬件吃不消。因此第二种方式应运而生:关注字典的大,而不太关注一致性。在memory bank里,只有一个query编码器,是可以通过梯度回传更新的。memory bank就是把整个数据集的特征都存在一起,比如imagenet的memory bank就有128万个特征,听起来虽然很大,但是每个特征只有128维,所以这样算下来整个memory bank也只需要600兆的空间就能都存下来。有了memory bank后,每次训练只需要从这个memory bank里去随机抽样很多key当字典即可。但这样的话由于memory bank里不同的key是由不同时刻不同的编码器计算出来的,所以每次抽出来若干个key进行计算时,无法保证这些key的特征具有一致性。

伪代码

最后作者提供了一个伪代码:首先我们初始化一个query编码器,然后把参数复制给key编码器。接下来从dataloader拿一个batch的数据,接下来就是得到正样本对。所以从x开始,先做一次数据增强,得到了一个query图片,然后再做一次数据增强,得到了一个key图片。然后将query图片通过query编码器得到特征,key图片通过key编码器得到特征,当然这里key的特征需要用detach做梯度截断来保证key编码器不更新。接下来就是算NCE loss,对query编码器更新,对key编码器进行动量更新,最后将新的key入队,老的key退出队列。

  • 实验部分我就不做一一解读了,总之就是MoCo潜力很大,在很多CV任务超过了有监督的方法,而且学到的特征可以很好的迁移到下游任务中,至此本文介绍结束,如果读者一直阅读到这关于MoCo的细节和无监督,对比学习应该也已经了解了。

参考

  1. 论文https://arxiv.org/pdf/1911.05722.pdf
  2. 朱毅老师的解读https://www.bilibili.com/video/BV1C3411s7t9/?spm_id_from=333.999.0.0&vd_source=d8673d961b825117f97b7ea28c540f9d

0 人点赞