本文提出了SimCLR(Simple Framework for Contrastive Learning of Visual Representations):一种用于对视觉表示进行对比学习的简单框架。我们简化了最近提出的对比自监督学习算法,不需要专门的架构或存储库。为了理解对比预测任务如何学习有用的表示,我们系统地研究了我们框架的主要组成部分。我们发现(1)数据增强的组合在定义有效的预测任务中起着关键作用,(2)在表示和对比损失之间引入可学习的非线性变换显著提高了学习表示的质量,(3)与监督学习相比,对比学习受益于更大的批量大小和更多的训练步骤。通过结合这些发现,我们能够在ImageNet上显著超越以前的自监督和半监督学习方法。一个在SimCLR学习的自监督表示上训练的线性分类器达到了76.5%的top-1准确率,相对于先前最先进方法提高了7%,与有监督的ResNet-50的性能相当。当仅在1%的标签上进行微调时,我们实现了85.8%的top-5准确率,在使用100倍更少的标签的情况下超过了AlexNet。
在没有人类监督的情况下学习有效的视觉表示是一个长期存在的问题。大多数主流方法可以分为两类:生成式和判别式。生成式方法学习在输入空间中生成或建模像素(Hinton等,2006;Kingma和Welling,2013;Goodfellow等,2014)。然而,像素级的生成计算成本高昂,并且在表示学习中可能并非必需。判别式方法使用类似于有监督学习的目标函数来学习表示,但训练网络执行预设任务,其中输入和标签都来自无标签数据集。许多这样的方法依赖于启发式方法来设计预设任务,这可能限制了学习表示的泛化性。基于潜空间对比学习的判别式方法最近取得了很大的进展,实现了最先进的结果(Hadsell等,2006;Dosovitskiy等,2014;Oord等,2018;Bachman等,2019)。
图1:使用不同的自监督方法(在ImageNet上预训练)学习的表示,通过线性分类器在ImageNet上的Top-1准确率。灰色叉表示有监督的ResNet-50模型。我们的方法SimCLR以粗体显示。
在这项工作中,我们引入了一个称为SimCLR的简单框架,用于对视觉表示进行对比学习。SimCLR不仅在性能上超越了之前的工作(图1),而且更简单,不需要特定的架构(Bachman等,2019;Hénaff等,2019)或者存储器(Wu等,2018;Tian等,2019;He等,2019;Misra和van der Maaten,2019)。
为了理解好的对比表示学习是如何实现的,我们系统地研究了我们框架的主要组成部分,并展示了以下结果:
• 多个数据增强操作的组合对于定义有效的对比预测任务至关重要。此外,无监督的对比学习比有监督学习受益于更强的数据增强。
• 引入可学习的非线性转换在表示和对比损失之间显著提高了学习表示的质量。
• 使用对比交叉熵损失的表示学习受益于归一化的嵌入和适当调整的温度参数。
• 对比学习受益于较大的批量大小和更长的训练时间,相比之下,对比学习的监督对照组。与有监督学习类似,对比学习受益于更深和更宽的网络。
我们将这些发现结合起来,在ImageNet ILSVRC-2012数据集上实现了自我监督学习和半监督学习的新的最先进结果(Russakovsky等,2015)。在线性评估协议下,SimCLR达到了76.5%的top-1准确率,相对于之前的最先进方法(Hénaff等,2019)提升了7%。当只使用ImageNet标签的1%进行微调时,SimCLR达到了85.8%的top-5准确率,相对提升了10%(Hénaff等,2019)。在其他自然图像分类数据集上进行微调时,SimCLR在12个数据集中有10个数据集与或优于强大的有监督基线(Kornblith等,2019)的表现相当。
受最近的对比学习算法的启发(请参见第7节概述),SimCLR通过在潜空间中使用对比损失来最大化同一数据示例的不同增强视图之间的一致性来学习表示。如图2所示,该框架包括以下四个主要组成部分。
• 一个随机数据增强模块,随机地对给定的数据样本进行变换,从而产生相同样本的两个相关视图,分别表示为˜xi和˜xj,我们将其视为正样本对。在本工作中,我们依次应用三种简单的增强方式:随机裁剪并重新调整到原始大小、随机颜色扭曲和随机高斯模糊。正如第3节所示,随机裁剪和颜色扭曲的组合对于获得良好的性能至关重要。
• 一个神经网络基础编码器 f(·),从增强的数据样本中提取表示向量。我们的框架允许使用各种网络架构而不受任何约束。为了简单起见,我们采用了常用的ResNet (He et al., 2016) 以获取 hi = f(˜xi) = ResNet(˜xi),其中 hi ∈ Rd 是平均池化层之后的输出。
• 一个小型神经网络投影头 g(·),将表示映射到应用对比损失的空间中。我们使用一个具有一个隐藏层的MLP来获得 zi = g(hi) = W(2)σ(W(1)hi),其中 σ 是ReLU非线性函数。如第4节所示,我们发现在 zi 上定义对比损失比在 hi 上定义对比损失更有益。
• 为对比预测任务定义的对比损失函数。给定一个包含正例对˜xi和˜xj的样本集合{˜xk},对比预测任务旨在从{˜xk}k=i中为给定的˜xi找到˜xj。
图2. 用于对比学习视觉表示的简单框架。从相同的数据增强集合中分别采样两个不同的数据增强操作符(t ∼ T 和 t‘~ T),并将它们应用于每个数据样本,得到两个相关的视图。通过对比损失函数,对基础编码器网络 f(·) 和投影头 g(·) 进行训练,以最大程度地提高一致性。训练完成后,我们舍弃投影头 g(·),并使用编码器 f(·) 和表示 h 进行下游任务。
我们随机抽取一个大小为N的小批量样本,并在从小批量中派生的增强样本对上定义对比预测任务,得到2N个数据点。我们不显式地抽样负样本。相反,对于一个正例对,类似于(Chen et al., 2017),我们将小批量中的其他2(N−1)个增强样本视为负样本。记作sim(u, v) = uTv/||u|| ||v||,表示L2归一化后的u和v的点积(即余弦相似度)。那么对于正例对(i, j)的损失函数定义如下:
其中,1[k!=i] ∈ {0, 1} 是一个指示函数,当 k != i 时取值为 1;τ 是一个温度参数。最终的损失函数在一个小批量的所有正样本对(包括 (i, j) 和 (j, i))上进行计算。这个损失函数在先前的工作中已经被使用过(Sohn, 2016; Wu et al., 2018; Oord et al., 2018),为了方便起见,我们将其称为 NT-Xent(归一化温度缩放的交叉熵损失)。
算法1总结了所提出的方法。
为了保持简单,我们不使用记忆库(Wu et al., 2018; He et al., 2019)来训练模型。相反,我们将训练批次大小N从256增加到8192。一个批次大小为8192给我们提供了来自两种增强视图的每个正样本对的16382个负样本。使用大批次大小进行训练时,使用标准的SGD/Momentum和线性学习率缩放(Goyal et al., 2017)可能不稳定。为了稳定训练,我们使用LARS优化器(You et al., 2017)适用于所有批次大小。我们使用Cloud TPUs进行模型训练,根据批次大小使用32到128个核心(cores)。
全局批归一化(Global BN)。标准的ResNet使用批归一化(Batch Normalization)(Ioffe&Szegedy,2015)。在使用数据并行训练的分布式训练中,批归一化的均值和方差通常在每个设备上局部聚合。在我们的对比学习中,由于正样本是在同一个设备上计算的,模型可以利用本地信息泄漏来提高预测准确性而不会改善表示。我们通过在训练期间在所有设备上聚合批归一化的均值和方差来解决这个问题。其他方法包括在设备之间对数据示例进行洗牌(He et al.,2019),或者用层归一化(Layer Normalization)替换批归一化(Hénaff et al.,2019)。
图3. 实线矩形表示图像,虚线矩形表示随机裁剪。通过随机裁剪图像,我们采样包括全局到局部视图(B → A)或相邻视图(D → C)的对比预测任务。
在这里,我们列出了我们实证研究的协议,旨在了解我们框架中的不同设计选择。
数据集和评估指标。我们大部分的无监督预训练研究(学习没有标签的编码器网络f)是在ImageNet ILSVRC-2012数据集(Rus-sakovsky等人,2015)上进行的。附录B.9中还包含了一些关于CIFAR-10数据集(Krizhevsky和Hinton,2009)的额外预训练实验。我们还在一系列用于迁移学习的数据集上测试了预训练结果。为了评估学到的表示,我们采用了广泛使用的线性评估协议(Zhang等人,2016;Oord等人,2018;Bachman等人,2019;Kolesnikov等人,2019),即在冻结的基本网络之上训练一个线性分类器,并将测试准确率作为表示质量的代理指标。除了线性评估,我们还与半监督学习和迁移学习的最新方法进行比较。
默认设置。除非另有说明,我们使用随机裁剪和调整大小(带有随机翻转),颜色扭曲和高斯模糊进行数据增强(详见附录A)。我们使用ResNet-50作为基本编码器网络,并使用2层MLP投影头将表示投影到128维的潜在空间。作为损失函数,我们使用NT-Xent,并使用LARS进行优化,学习率为4.8(等于0.3乘以BatchSize/256),权重衰减为10的负6次方。我们以批量大小4096进行100个周期的训练。此外,我们在前10个周期使用线性预热,并使用余弦衰减调度(无重启)来衰减学习率(Loshchilov和Hutter,2016)。
数据增强定义了预测任务。虽然数据增强在监督和无监督表示学习中被广泛使用(Krizhevsky等,2012; Hénaff等,2019; Bachman等,2019),但尚未被视为定义对比预测任务的系统方法。许多现有方法通过改变架构来定义对比预测任务。例如,Hjelm等人(2018);Bachman等人(2019)通过限制网络架构中的感受野实现了从全局到局部视图的预测,而Oord等人(2018);Hénaff等人(2019)通过固定的图像分割过程和上下文聚合网络实现了相邻视图的预测。我们展示了通过对目标图像进行简单的随机裁剪(带有调整大小)可以避免这种复杂性,这创建了一个包含上述两种任务的预测任务族,如图3所示。这个简单的设计选择方便地将预测任务与其他组件(如神经网络架构)解耦。可以通过扩展增强的族群并随机组合它们来定义更广泛的对比预测任务。
为了系统研究数据增强的影响,我们在这里考虑了几种常见的增强方法。一种增强方法涉及数据的空间/几何变换,例如裁剪和调整大小(包括水平翻转)、旋转(Gidaris等,2018)和遮挡(DeVries和Taylor,2017)。另一种增强方法涉及外观变换,例如颜色扭曲(包括颜色丢失、亮度、对比度、饱和度和色调)(Howard,2013; Szegedy等,2015)、高斯模糊和Sobel滤波。图4可视化了我们在这项工作中研究的增强方法。
图4. 所研究的数据增强操作的示意图。每个增强操作可以随机地以一些内部参数(例如旋转角度、噪声水平)对数据进行变换。请注意,我们在消融实验中只测试了这些操作符,用于训练我们的模型的增强策略仅包括随机裁剪(带有翻转和调整大小)、颜色扭曲和高斯模糊。(原始图片 cc-by: Von.grzanka)
图5. 在单个或多个数据增强操作下的线性评估(ImageNet top-1准确率),仅应用于一个分支。除最后一列外,对角线上的条目对应于单个变换,非对角线上的条目对应于两个变换的组合(按顺序应用)。最后一列反映了行的平均值。
为了理解单个数据增强操作的效果以及增强组合的重要性,我们研究了在单个或两个增强操作下应用我们框架的性能。由于ImageNet图像的尺寸不同,我们总是先进行裁剪和调整图像大小(Krizhevsky等,2012;Szegedy等,2015),这使得在没有裁剪的情况下研究其他增强操作变得困难。为了消除这个混淆,我们考虑了这个消融实验中的非对称数据转换设置。具体而言,我们总是先随机裁剪图像并调整它们到相同的分辨率,然后仅将目标变换应用于图2中的一个分支,同时将另一个分支保持为恒等变换(即t(xi) = xi)。注意,这种非对称的数据增强会降低性能。尽管如此,这个设置不会实质性地改变单个数据增强操作及其组合的影响。
图5展示了在单个和组合变换下的线性评估结果。我们观察到,没有单个变换足以学习到良好的表示,即使模型在对比任务中几乎可以完美地识别出正样本对。当组合增强操作时,对比预测任务变得更加困难,但表示的质量显著提高。附录B.2对更广泛的增强操作组合进行了进一步的研究。
其中一种增强操作组合显得特别突出:随机裁剪和随机颜色扭曲。我们推测,当只使用随机裁剪作为数据增强时,一个严重的问题是图像中的大多数补丁共享相似的颜色分布。图6显示,仅凭颜色直方图就足以区分图像。神经网络可能会利用这种捷径来解决预测任务。因此,将裁剪与颜色扭曲结合起来学习具有泛化能力的特征至关重要。
为了进一步证明颜色增强的重要性,我们调整了颜色增强的强度,如表1所示。更强的颜色增强显著提高了学习的无监督模型的线性评估结果。在这种情况下,AutoAugment(Cubuk等,2019)是使用监督学习找到的复杂增强策略,但其表现并不比简单的裁剪+(更强的)颜色扭曲更好。当使用相同的增强策略训练监督模型时,我们观察到更强的颜色增强并不能提高甚至会损害它们的性能。因此,我们的实验证明,相比于监督学习,无监督对比学习受益于更强的(颜色)数据增强。尽管之前的研究已经报道了数据增强对自监督学习的有用性(Doersch等,2015;Bachman等,2019;Hénaff等,2019;Asano等,2019),但我们的研究表明,对于对比学习来说,即使对于监督学习没有准确性提升的数据增强仍然可以极大地帮助。
图7展示了在不同深度和宽度下的线性评估结果。蓝色圆点表示我们的模型经过100个epoch的训练,红色星星表示我们的模型经过1000个epoch的训练,绿色十字表示经过90个epoch的监督式ResNet训练的模型(He等,2016)。
图7显示,或许并不令人意外的是,增加深度和宽度都会改善性能。虽然这些发现对于监督学习也成立(He等,2016),但我们发现,随着模型规模的增加,无监督模型上训练的线性分类器与监督模型之间的差距变小,这表明无监督学习比监督学习更加受益于更大的模型。
表1. 使用线性评估的无监督ResNet-50和监督ResNet-50的Top-1准确率,在不同颜色失真强度(参见附录A)和其他数据变换下。Strength 1 (+Blur) 是我们的默认数据增强策略。
图7. 对具有不同深度和宽度的模型进行线性评估。蓝色圆点代表我们的模型,训练了100个epochs;红色星号代表我们的模型,训练了1000个epochs;绿色十字代表经过90个epochs训练的监督ResNet模型(He等人,2016)。
表2. 负损失函数及其梯度。所有输入向量,即u、v+、v-,都进行了L2归一化。NT-Xent是“归一化温度缩放交叉熵”的缩写。不同的损失函数对正负样本施加不同的加权。
图8. 使用不同的投影头g(·)和不同维度的z = g(h)进行线性评估的表示学习结果。这里表示h(投影前)是2048维。
我们随后研究了包含投影头(即g(h))的重要性。图8展示了使用三种不同投影头架构的线性评估结果:(1)恒等映射;(2)线性投影,这是一些先前方法使用的方式(Wu等,2018);(3)默认的非线性投影,包含一个额外的隐藏层(使用ReLU激活函数),类似于Bachman等(2019)的方法。我们观察到非线性投影比线性投影更好(增加3%),而且比没有投影头的情况要好得多(超过10%)。当使用投影头时,无论输出维度如何,都会观察到类似的结果。此外,即使使用非线性投影,投影头之前的隐藏层h仍然比投影头之后的层z=g(h)更好(超过10%),这表明投影头之前的隐藏层是比投影头之后的层更好的表示。
我们推测使用非线性投影之前的表示的重要性是由于对比损失引起的信息丢失。特别是,z = g(h)被训练成对数据变换具有不变性。因此,g可能会删除对下游任务有用的信息,例如物体的颜色或方向。通过利用非线性转换g(·),可以在h中形成和保持更多的信息。为了验证这个假设,我们进行了实验,使用h或g(h)来学习预训练期间应用的变换。在这里,我们设置g(h) = W(2)σ(W(1)h),具有相同的输入和输出维度(即2048)。表3显示h包含关于应用的变换的更多信息,而g(h)失去了信息。更详细的分析可以在附录B.4中找到。
表3. 在不同表示上训练额外的多层感知器(MLP)来预测应用的变换的准确率。除了裁剪和颜色增强外,在预训练期间我们还额外且独立地添加了旋转({0°, 90°, 180°, 270°}之一)、高斯噪声和Sobel滤波器变换,对应表中的最后三行。h和g(h)的维度均为2048。
我们将NT-Xent损失与其他常用的对比损失函数进行了比较,如逻辑损失(Mikolov等,2013)和边际损失(Schroff等,2015)。表2展示了目标函数以及对损失函数输入的梯度。通过观察梯度,我们可以得出以下结论:1)通过温度参数的2归一化(即余弦相似度)有效地对不同的样本进行加权,合适的温度参数可以帮助模型从困难的负例中学习;2)与交叉熵不同,其他目标函数不通过负例的相对难度对其进行加权。因此,对于这些损失函数,我们必须进行半硬负例挖掘(Schroff等,2015):即不使用所有损失项来计算梯度,而是使用半硬负例(即在损失边界内,与正例距离最近但比正例更远的样本)来计算梯度。
为了公平比较,我们对所有损失函数使用相同的2归一化,并调整超参数,并报告它们的最佳结果。表4显示,尽管(半硬)负例挖掘有所帮助,但最佳结果仍远远不及我们默认的NT-Xent损失。
接下来,我们测试了在我们默认的NT-Xent损失中2归一化(即余弦相似度与点积)和温度τ的重要性。表5显示,在没有归一化和适当的温度缩放的情况下,性能明显较差。没有2归一化,对比任务的准确率更高,但线性评估下得到的表示更差。
表4. 使用不同损失函数训练的模型的线性评估结果(top-1)。"sh"表示使用半硬负样本挖掘。
表5. 使用不同的2范数和温度τ对NT-Xent损失进行训练的模型的线性评估结果。对于对比分布,采样了4096个样本。
图9. 使用不同的批次大小和训练时期训练的线性评估模型(ResNet-50)。每个柱状图代表一次从头开始的单次运行。
Figure 9展示了在不同的训练时长下,批次大小对模型训练的影响。我们发现,当训练时长较短(例如100个epochs)时,较大的批次大小明显优于较小的批次大小。随着训练步数/epochs的增加,不同批次大小之间的差距减小甚至消失,前提是批次是随机重新采样的。与监督学习不同(Goyal等,2017),在对比学习中,较大的批次大小提供了更多的负例样本,有助于收敛(即在给定准确性下需要较少的epochs和步数)。延长训练时长也提供了更多的负例样本,提高了结果。在附录B.1中,我们提供了更长训练步数的结果。
在本小节中,与Kolesnikov等人(2019)和He等人(2019)类似,我们使用ResNet-50的三种不同的隐藏层宽度(宽度乘数为1×、2×和4×)。为了更好地收敛,我们在这里训练的模型为1000个epochs。
线性评估。表6将我们的结果与先前的方法(Zhuang等,2019; He等,2019; Misra&van der Maaten,2019; Hénaff等,2019; Kolesnikov等,2019; Donahue&Simonyan,2019; Bachman等,2019; Tian等,2019)在线性评估设置中进行了比较(详见附录B.6)。表1展示了不同方法之间的更多数值比较。我们能够使用标准的网络获得比以前需要专门设计的架构的方法更好的结果。我们使用的ResNet-50(4×)获得的最佳结果可以与经过监督预训练的ResNet-50相媲美。
表6. 使用不同自监督方法学习的表示进行训练的线性分类器在ImageNet上的准确率。
表7.ImageNet的准确性,用很少的标签训练的模型。
表8. 我们自监督方法与监督基准在12个自然图像分类数据集上的迁移学习性能进行了比较,使用在ImageNet上预训练的ResNet-50(4×)模型。未显著差异于最佳结果(p > 0.05,置换检验)的结果以粗体显示。有关实验细节和使用标准ResNet-50的结果,请参阅附录B.8。
半监督学习。我们按照Zhai等人的方法(2019)从标记的ILSVRC-12训练数据集中以类平衡的方式采样了1%或10%的数据(每个类别大约12.8和128张图像)。我们简单地对标记数据进行整个基础网络的微调,没有使用正则化(详见附录B.5)。表7显示了我们的结果与最近方法(Zhai等人,2019;Xie等人,2019;Sohn等人,2020;Wu等人,2018;Donahue和Simonyan,2019;Misra和van der Maaten,2019;Hénaff等人,2019)的比较。由于对超参数(包括数据增强)进行了深入的搜索,因此(Zhai等人,2019)中的监督基线较强。再次说明,我们的方法在使用1%和10%标签时显著优于现有技术。有趣的是,我们对预训练的ResNet-50(2×、4×)在完整的ImageNet上进行微调时,与从头开始训练相比也显著更好(最多达到2%,详见附录B.2)。
迁移学习。我们在12个自然图像数据集上评估了迁移学习的性能,包括线性评估(固定特征提取器)和微调设置。我们按照Kornblith等人的方法(2019)为每个模型-数据集组合进行超参数调优,并在验证集上选择最佳超参数。表8显示了ResNet-50(4×)模型的结果。在进行微调时,我们的自监督模型在5个数据集上显著优于监督基线,而监督基线只在2个数据集上优于我们的模型(即Pets和Flowers)。在其余的5个数据集上,两个模型的性能相当。附录B.8中提供了完整的实验细节以及使用标准ResNet-50架构的结果。
让图像的表示在小的变换下保持一致的思想可以追溯到Becker和Hinton的工作(1992年)。我们通过利用最近在数据增强、网络架构和对比损失方面的进展对其进行了扩展。类似的一致性思想在其他背景下也得到了探索,例如在半监督学习中,目标是确保类别标签预测的一致性(Xie等人,2019年;Berthelot等人,2019年)。
手工设计的先验任务。自监督学习的最近复兴始于人为设计的先验任务,例如相对补丁预测(Doersch等人,2015年),解决拼图难题(Noroozi&Favaro,2016年),彩色化(Zhang等人,2016年)和旋转预测(Gidaris等人,2018年;Chen等人,2019年)。尽管可以通过使用更大的网络和更长的训练时间获得良好的结果(Kolesnikov等人,2019年),但这些先验任务依赖于一定程度上的特定启发式方法,限制了学习表示的通用性。
对比式视觉表示学习。从Hadsell等人(2006年)开始,这些方法通过对比正样本和负样本来学习表示。沿着这个思路,Dosovitskiy等人(2014年)提出将每个实例视为由特征向量表示的类别(以参数形式)。Wu等人(2018年)提出使用存储实例类别表示向量的内存库,这种方法在几篇最近的论文中得到采用和扩展(Zhuang等人,2019年;Tian等人,2019年;He等人,2019年;Misra&van der Maaten,2019年)。其他工作探索使用批内样本进行负样本采样,而不是使用内存库(Doersch&Zisserman,2017年;Ye等人,2019年;Ji等人,2019年)。
最近的文献试图将他们的方法的成功与潜在表示之间的互信息最大化联系起来(Oord等人,2018年;Hénaff等人,2019年;Hjelm等人,2018年;Bachman等人,2019年)。然而,目前还不清楚对比式方法的成功是由互信息还是由对比损失的特定形式决定的(Tschannen等人,2019年)。我们注意到,我们框架中的几乎所有个别组成部分都出现在先前的工作中,尽管具体的实例可能不同。相对于先前的工作,我们框架的优越性并不是由任何单一的设计选择来解释的,而是由它们的组合。在附录C中,我们对我们的设计选择与先前工作的对比进行了全面的比较。
在本文中,我们提出了一个简单的框架及其在对比式视觉表示学习中的实例化。我们仔细研究了其组成部分,并展示了不同设计选择的影响。通过结合我们的发现,我们在自监督学习、半监督学习和迁移学习方面取得了显著的改进。我们的方法与仅使用ImageNet上的标准监督学习不同之处在于数据增强的选择、网络末端使用非线性投影层以及损失函数的选择。这个简单框架的优势表明,尽管最近对自监督学习产生了极大的兴趣,但它仍然被低估。