A Simple Framework for Contrastive Learning of Visual ...

1 A Simple Framework for Contrastive Learning of Visual RepresentationsTing Chen1 Simon Kornblith1 Mohammad Norouzi1 Geoffrey Hinton1 AbstractThis paper presentsSimCLR: a Simple frameworkfor Contrastive Learning of Visual simplify recently proposed Contrastive self-supervised Learning algorithms without requiringspecialized architectures or a memory bank. Inorder to understand what enables the contrastiveprediction tasks to learn useful representations,we systematically study the major components ofour Framework . We show that (1) composition ofdata augmentations plays a critical role in definingeffective predictive tasks, (2) introducing a learn-able nonlinear transformation between the repre-sentation and the Contrastive loss substantially im-proves the quality of the learned representations,and (3) Contrastive Learning benefits from largerbatch sizes and more training steps compared tosupervised Learning .

2 By combining these findings,we are able to considerably outperform previousmethods for self-supervised and semi-supervisedlearning on ImageNet. A linear classifier trainedon self-supervised representations learned by Sim-CLR achieves top-1 accuracy, which is a7% relative improvement over previous state-of-the-art, matching the performance of a supervisedResNet-50. When fine-tuned on only 1% of thelabels, we achieve top-5 accuracy, outper-forming AlexNet with 100 fewer IntroductionLearning effective Visual representations without humansupervision is a long-standing problem. Most mainstreamapproaches fall into one of two classes: generative or dis-criminative. Generative approaches learn to generate orotherwise model pixels in the input space (Hinton et al.,2006; Kingma & Welling, 2013; Goodfellow et al., 2014).1 Google Research, Brain Team.

3 Correspondence to: Ting of the37thInternational Conference on MachineLearning, Vienna, Austria, PMLR 119, 2020. Copyright 2020 bythe author(s).1 Code available at Top-1 accuracy of linear classifiers trainedon representations learned with different self-supervised meth-ods (pretrained on ImageNet). Gray cross indicates supervisedResNet-50. Our method, SimCLR, is shown in , pixel-level generation is computationally expen-sive and may not be necessary for representation approaches learn representations using objec-tive functions similar to those used for supervised Learning ,but train networks to perform pretext tasks where both the in-puts and labels are derived from an unlabeled dataset. Manysuch approaches have relied on heuristics to design pretexttasks (Doersch et al., 2015; Zhang et al., 2016; Noroozi &Favaro, 2016; Gidaris et al., 2018), which could limit thegenerality of the learned representations.

4 Discriminativeapproaches based on Contrastive Learning in the latent spacehave recently shown great promise, achieving state-of-the-art results (Hadsell et al., 2006; Dosovitskiy et al., 2014;Oord et al., 2018; Bachman et al., 2019).In this work, we introduce a Simple Framework for con-trastive Learning of Visual representations, which we callSimCLR. Not only does SimCLR outperform previous work(Figure 1), but it is also simpler, requiring neither special-ized architectures (Bachman et al., 2019; H naff et al., 2019)nor a memory bank (Wu et al., 2018; Tian et al., 2019; Heet al., 2019; Misra & van der Maaten, 2019).In order to understand what enables good Contrastive repre-sentation Learning , we systematically study the major com-ponents of our Framework and show that:A Simple Framework for Contrastive Learning of Visual Representations Composition of multiple data augmentation operationsis crucial in defining the Contrastive prediction tasks thatyield effective representations.

5 In addition, unsupervisedcontrastive Learning benefits from stronger data augmen-tation than supervised Learning . Introducing a learnable nonlinear transformation be-tween the representation and the Contrastive loss substan-tially improves the quality of the learned representations. representation Learning with Contrastive cross entropyloss benefits from normalized embeddings and an appro-priately adjusted temperature parameter. Contrastive Learning benefits from larger batch sizes andlonger training compared to its supervised supervised Learning , Contrastive Learning benefitsfrom deeper and wider combine these findings to achieve a new state-of-the-artin self-supervised and semi-supervised Learning on Ima-geNet ILSVRC-2012 (Russakovsky et al., 2015). Under thelinear evaluation protocol, SimCLR achieves top-1accuracy, which is a 7% relative improvement over previousstate-of-the-art (H naff et al.)

6 , 2019). When fine-tuned withonly 1% of the ImageNet labels, SimCLR achieves accuracy, a relative improvement of 10% (H naff et al.,2019). When fine-tuned on other natural image classifica-tion datasets, SimCLR performs on par with or better thana strong supervised baseline (Kornblith et al., 2019) on 10out of 12 The Contrastive Learning FrameworkInspired by recent Contrastive Learning algorithms (see Sec-tion 7 for an overview), SimCLR learns representationsby maximizing agreement between differently augmentedviews of the same data example via a Contrastive loss inthe latent space. As illustrated in Figure 2, this frameworkcomprises the following four major components. A stochasticdata augmentationmodule that transformsany given data example randomly resulting in two cor-related views of the same example, denoted xiand xj,which we consider as a positive pair.

7 In this work, wesequentially apply three Simple augmentations:randomcroppingfollowed by resize back to the original size,ran-dom color distortions, andrandom Gaussian blur. Asshown in Section 3, the combination of random crop andcolor distortion is crucial to achieve a good performance. A neural networkbase encoderf( )that extracts repre-sentation vectors from augmented data examples. Ourframework allows various choices of the network archi-tecture without any constraints. We opt for simplicityand adopt the commonly used ResNet (He et al., 2016) representation x xi xjhihjzizjt Tt Tf( )f( )g( )g( )Maximize agreementFigure Simple Framework for Contrastive Learning of visualrepresentations. Two separate data augmentation operators aresampled from the same family of augmentations (t Tandt T) and applied to each data example to obtain two correlatedviews.

8 A base encoder networkf( )and a projection headg( )are trained to maximize agreement using a Contrastive loss. Aftertraining is completed, we throw away the projection headg( )anduse encoderf( )and representationhfor downstream obtainhi=f( xi) = ResNet( xi)wherehi Rdisthe output after the average pooling layer. A small neural networkprojection headg( )that mapsrepresentations to the space where Contrastive loss isapplied. We use a MLP with one hidden layer to obtainzi=g(hi) =W(2) (W(1)hi)where is a ReLU non-linearity. As shown in section 4, we find it beneficial todefine the Contrastive loss onzi s rather thanhi s. Acontrastive loss functiondefined for a Contrastive pre-diction task. Given a set{ xk}including a positive pairof examples xiand xj, thecontrastive prediction taskaims to identify xjin{ xk}k6=ifor a given randomly sample a minibatch ofNexamples and definethe Contrastive prediction task on pairs of augmented exam-ples derived from the minibatch, resulting in2 Ndata do not sample negative examples explicitly.

9 Instead,given a positive pair, similar to (Chen et al., 2017), we treatthe other2(N 1)augmented examples within a minibatchas negative examples. Letsim(u,v) =u>v/ u v de-note the dot product between`2normalizeduandv( similarity). Then the loss function for a positive pairof examples(i,j)is defined as`i,j= logexp(sim(zi,zj)/ ) 2Nk=11[k6=i]exp(sim(zi,zk)/ ),(1)where1[k6=i] {0,1}is an indicator function evaluating to1iffk6=iand denotes a temperature parameter. The fi-nal loss is computed across all positive pairs, both(i,j)and(j,i), in a mini-batch. This loss has been used inprevious work (Sohn, 2016; Wu et al., 2018; Oord et al.,2018); for convenience, we term itNT-Xent(the normalizedtemperature-scaled cross entropy loss).A Simple Framework for Contrastive Learning of Visual RepresentationsAlgorithm 1 SimCLR s main Learning :batch sizeN, constant , structure off,g, minibatch{xk}Nk=1dofor allk {1.}

10 ,N}dodraw two augmentation functionst T,t T# the first augmentation x2k 1=t(xk)h2k 1=f( x2k 1)# representationz2k 1=g(h2k 1)# projection# the second augmentation x2k=t (xk)h2k=f( x2k)# representationz2k=g(h2k)# projectionend forfor alli {1,..,2N}andj {1,..,2N}dosi,j=z>izj/( zi zj )# pairwise similarityend fordefine`(i,j)as`(i,j)= logexp(si,j/ ) 2Nk=11[k6=i]exp(si,k/ )L=12N Nk=1[`(2k 1,2k) +`(2k,2k 1)]update networksfandgto minimizeLend forreturnencoder networkf( ), and throw awayg( )Algorithm 1 summarizes the proposed Training with Large Batch SizeTo keep it Simple , we do not train the model with a memorybank (Wu et al., 2018; He et al., 2019). Instead, we varythe training batch sizeNfrom 256 to 8192. A batch sizeof 8192 gives us 16382 negative examples per positive pairfrom both augmentation views. Training with large batchsize may be unstable when using standard SGD/Momentumwith linear Learning rate scaling (Goyal et al.)