What Should Not Be Contrastive in Contrastive Learning

1 INTRODUCTION . Self-supervised learning , which uses raw image data and/or available pretext tasks as its own supervision , has become increasingly popular as the inability of supervised models to generalize beyond their training data has become apparent . Different pretext tasks have been proposed with different transformations , such as spatial patch prediction ( Doersch et al. , 2015 ; Noroozi & Favaro , 2016 ) , colorization ( Zhang et al. , 2016 ; Larsson et al. , 2016 ; Zhang et al. , 2017 ) , rotation ( Gidaris et al. , 2018 ) . Whereas pretext tasks aim to recover the transformations between different “ views ” of the same data , more recent contrastive learning methods ( Wu et al. , 2018 ; Tian et al. , 2019 ; He et al. , 2020 ; Chen et al. , 2020a ) instead try to learn to be invariant to these transformations , while remaining discriminative with respect to other data points . Here , the transformations are generated using classic data augmentation techniques which correspond to common pretext tasks , e.g. , randomizing color , texture , orientation and cropping . Yet , the inductive bias introduced through such augmentations is a double-edged sword , as each augmentation encourages invariance to a transformation which can be beneficial in some cases and harmful in others : e.g. , adding rotation may help with view-independent aerial image recognition , but significantly downgrade the capacity of a network to solve tasks such as detecting which way is up in a photograph for a display application . Current self-supervised contrastive learning methods assume implicit knowledge of downstream task invariances . In this work , we propose to learn visual representations which capture individual factors of variation in a contrastive learning framework without presuming prior knowledge of downstream invariances . Instead of mapping an image into a single embedding space which is invariant to all the handcrafted augmentations , our model learns to construct separate embedding sub-spaces , each of which is sensitive to a specific augmentation while invariant to other augmentations . We achieve this by optimizing multiple augmentation-sensitive contrastive objectives using a multi-head architecture with a shared backbone . Our model aims to preserve information with regard to each augmentation in a unified representation , as well as learn invariances to them . The general representation trained with these augmentations can then be applied to different downstream tasks , where each task is free to selectively utilize different factors of variation in our representation . We consider transfer of either the shared backbone representation , or the concatenation of all the task-specific heads ; both Color Rotation Texture Color Rotation Texture Augmentations Downstream Tasks Augmentations ? Classes ( a ) ( b ) Coarse-grained Fine-grained ( bird ) Fine-grained ( flower ) Figure 1 : Self-supervised contrastive learning relies on data augmentations as depicted in ( a ) to learn visual representations . However , current methods introduce inductive bias by encouraging neural networks to be less sensitive to information w.r.t . augmentation , which may help or may hurt . As illustrated in ( b ) , rotation invariant embeddings can help on certain flower categories , but may hurt animal recognition performance ; conversely color invariance generally seems to help coarse grained animal classification , but can hurt many flower categories and bird categories . Our method , shown in the following figure , overcomes this limitation . outperform all baselines ; the former uses same embedding dimensions as typical baselines , while the latter provides greatest overall performance in our experiments . In this paper , we experiment with three types of augmentations : rotation , color jittering , and texture randomization , as visualized in Figure 1 . We evaluate our approach across a variety of diverse tasks including large-scale classification ( Deng et al. , 2009 ) , fine-grained classification ( Wah et al. , 2011 ; Van Horn et al. , 2018 ) , few-shot classification ( Nilsback & Zisserman , 2008 ) , and classification on corrupted data ( Barbu et al. , 2019 ; Hendrycks & Dietterich , 2019 ) . Our representation shows consistent performance gains with increasing number of augmentations . Our method does not require hand-selection of data augmentation strategies , and achieves better performance against state-of-the-art MoCo baseline ( He et al. , 2020 ; Chen et al. , 2020b ) , and demonstrates superior transferability , generalizability and robustness across tasks and categories . Specifically , we obtain around 10 % improvement over MoCo in classification when applied on the iNaturalist ( Van Horn et al. , 2018 ) dataset . 2 BACKGROUND : CONTRASTIVE LEARNING FRAMEWORK . Contrastive learning learns a representation by maximizing similarity and dissimilarity over data samples which are organized into similar and dissimilar pairs , respectively . It can be formulated as a dictionary look-up problem ( He et al. , 2020 ) , where a given reference image I is augmented into two views , query and key , and the query token q should match its designated key k+ over a set of sampled negative keys { k− } from other images . In general , the framework can be summarized as the following components : ( i ) A data augmentation module T constituting n atomic augmentation operators , such as random cropping , color jittering , and random flipping . We denote a pre-defined atomic augmentation as random variable Xi . Each time the atomic augmentation is executed by sampling a specific augmentation parameter from the random variable , i.e. , xi∼Xi . One sampled data augmentation module transforms image I into a random view Ĩ , denoted as Ĩ = T [ x1 , x2 , . . . , xn ] ( I ) . Positive pair ( q , k+ ) is generated by applying two randomly sampled data augmentation on the same reference image . ( ii ) An encoder network f which extracts the feature v of an image I by mapping it into a d-dimensional space Rd . ( iii ) A projection head h which further maps extracted representations into a hyper-spherical ( normalized ) embedding space . This space is subsequently used for a specific pretext task , i.e. , contrastive loss objective for a batch of positive/negative pairs . A common choice is InfoNCE ( Oord et al. , 2018 ) : Lq = − log exp ( q·k+/τ ) exp ( q·k+/τ ) + ∑ k− exp ( q·k−/τ ) , ( 1 ) where τ is a temperature hyper-parameter scaling the distribution of distances . As a key towards learning a good feature representation ( Chen et al. , 2020a ) , a strong augmentation policy prevents the network from exploiting naı̈ve cues to match the given instances . However , in- ductive bias is introduced through the selection of augmentations , along with their hyper-parameters defining the strength of each augmentation , manifested in Equation 1 that any views by the stochastic augmentation module T of the same instance are mapped onto the same point in the embedding space . The property negatively affects the learnt representations : 1 ) Generalizability and transferability are harmed if they are applied to the tasks where the discarded information is essential , e.g. , color plays an important role in fine-grained classification of birds ; 2 ) Adding an extra augmentation is complicated as the new operator may be helpful to certain classes while harmful to others , e.g. , a rotated flower could be very similar to the original one , whereas it does not hold for a rotated car ; 3 ) The hyper-parameters which control the strength of augmentations need to be carefully tuned for each augmentation to strike a delicate balance between leaving a short-cut open and completely invalidate one source of information . 3 LOOC : LEAVE-ONE-OUT CONTRASTIVE LEARNING . We propose Leave-one-out Contrastive Learning ( LooC ) , a framework for multi-augmentation contrastive learning . Our framework can selectively prevent information loss incurred by an augmentation . Rather than projecting every view into a single embedding space which is invariant to all augmentations , in our LooC method the representations of input images are projected into several embedding spaces , each of which is not invariant to a certain augmentation while remaining invariant to others , as illustrated in Figure 2 . In this way , each embedding sub-space is specialized to a single augmentation , and the shared layers will contain both augmentation-varying and invariant information . We learn a shared representation jointly with the several embedding spaces ; we transfer either the shared representation alone , or the concatenation of all spaces , to downstream tasks . View Generation . Given a reference image and n atomic augmentations , we first augment the reference image with two sets of independently sampled augmentation parameters into the query view Iq and the first key view Ik0 , i.e. , I { q , k0 } = T [ x { q , k0 } 1 , x { q , k0 } 2 , . . . , x { q , k0 } n ] ( I ) . Additionally , we generate n views from the reference image as extra key views , denoted as Iki , ∀i ∈ { 1 , . . . , n } . For the ith additional key view , the parameter of ith atomic augmentation is copied from it of the query view , i.e. , xkii ≡ x q i , ∀i ∈ { 1 , . . . , n } ; whereas the parameter of other atomic augmentations are still independently sampled , i.e. , xkij ∼ Xj , ∀j 6= i . For instance , assume that we have a set of two atomic augmentations { random_rotation , color_jitter } , Iq and Ik1 are always augmented by the same rotation angle but different color jittering ; Iq and Ik2 are always augmented by the same color jittering but different rotation angle ; Iq and Ik0 are augmented independently , as illustrated in the left part of Figure 2 . Contrastive Embedding Space . The augmented views are encoded by a neural network encoder f ( · ) into feature vectors vq , vk0 , · · · , vkn in a joint embedding space V ∈ Rd . Subsequently , they are projected into n+1 normalized embedding spaces Z0 , Z1 , · · · , Zn ∈ Rd ′ by projection heads h : V 7→ Z , among which Z0 is invariant to all types of augmentations , whereas Zi ( ∀i ∈ { 1 , 2 , · · · , n } ) is dependent on the ith type of augmentation but invariant to other types of augmentations . In other words , in Z0 all features v should be mapped to a single point , whereas in Zi ( ∀i ∈ { 1 , 2 , · · · , n } ) only vq and vki should be mapped to a single point while vkj ∀j 6= i should be mapped to n−1 separate points , as only Iq and Iki share the same ith augmentation . We perform contrastive learning in all normalized embedding spaces based on Equation 1 , as shown in the right part of Figure 2 . For each query zq , denote zk + as the keys from the same instance , and zk − as the keys from other instances . Since all views should be mapped to the single point in Z0 , the positive pair for the query zq0 is z k+0 0 , and the negative pairs are embeddings of other instances in this embedding space { zk − 0 0 } ; for embedding spaces Z1 , · · · , Zn , the positive pair for the query zqi is z k+i i , while the negative pairs are embeddings of other instances in this embedding space { zk − i i } , and { z k+j i | ∀j ∈ { 0 , 1 , · · · , n } and j 6= i } , which are the embeddings of the same instance with different ith augmentation . The network then learns to be sensitive to one type of augmentation while insensitive to other types of augmentations in one embedding space . Denote E { + , − } i , j = exp ( z q i · z k { + , − } j i /τ ) . The overall training objective for q is : Lq = − 1 n+ 1 ( log E+0,0 E+0,0 + ∑ k− E − 0,0 + n∑ i=1 log E+i , i∑n j=0 E + i , j + ∑ k− E − i , i ) , ( 2 ) The network must preserve information w.r.t . all augmentations in the general embedding space V in order to optimize the combined learning objectives of all normalized embedding spaces . Learnt representations . The representation for downstream tasks can be from the general embedding space V ( Figure 2 , blue region ) , or the concatenation of all embedding sub-spaces ( Figure 2 , grey region ) . LooC method returns V ; we term the implementation using the concatenation of all embedding sub-spaces as LooC++ .

Summary: The paper proposes a contrastive learning approach for self-supervised learning in which multiple heads are trained to be invariant to all but one type of data augmentation. The rationale is that different downstream tasks may require different types of invariances (e.g. we may want to be rotation invariant for pictures of flowers, but not for pictures of animals), and one does not know a-priori which kind of invariances will be required. After training multiple representation heads, one can later concatenate them or use the general embedding (the input to all the variant-specific heads) for the downstream task.