Can contrastive learning avoid shortcut solutions?

The generalization of representations learned via contrastive learning depends crucially on what features of the data are extracted . However , we observe that the contrastive loss does not always sufficiently guide which features are extracted , a behavior that can negatively impact the performance on downstream tasks via “ shortcuts ” , i.e. , by inadvertently suppressing important predictive features . We find that feature extraction is influenced by the difficulty of the so-called instance discrimination task ( i.e. , the task of discriminating pairs of similar points from pairs of dissimilar ones ) . Although harder pairs improve the representation of some features , the improvement comes at the cost of suppressing previously well represented features . In response , we propose implicit feature modification ( IFM ) , a method for altering positive and negative samples in order to guide contrastive models towards capturing a wider variety of predictive features . Empirically , we observe that IFM reduces feature suppression , and as a result improves performance on vision and medical imaging tasks . The code is available at : https : //github . com/joshr17/IFM . 1 Introduction . Representations trained with contrastive learning are adept at solving various vision tasks including classification , object detection , instance segmentation , and more [ 5 , 15 , 44 ] . In contrastive learning , encoders are trained to discriminate pairs of positive ( similar ) inputs from a selection of negative ( dissimilar ) pairs . This task is called instance discrimination : It is often framed using the InfoNCE loss [ 14 , 33 ] , whose minimization forces encoders to extract input features that are sufficient to discriminate similar and dissimilar pairs . However , learning features that are discriminative during training does not guarantee a model will generalize . Many studies find inductive biases in supervised learning toward simple “ shortcut ” features and decision rules [ 16 , 21 , 32 ] which result in unpredictable model behavior under perturbations [ 22 , 43 ] and failure outside the training distribution [ 2 , 37 ] . Simplicity bias has various potential sources [ 11 ] including training methods [ 8 , 29 , 41 ] and architecture design [ 10 , 17 ] . Bias towards shortcut decision rules also hampers transferability in contrastive learning [ 4 ] , where it is in addition influenced by the instance discrimination task . These difficulties lead us to ask : can the contrastive instance discrimination task itself be modified to avoid learning shortcut solutions ? We approach this question by studying the relation between contrastive instance discrimination and feature learning . First , we theoretically explain why optimizing the InfoNCE loss alone does not guarantee avoidance of shortcut solutions that suppress ( i.e. , discard ) certain input features [ 4 , 11 ] . Second , despite this negative result , we show that it is still possible to trade off representation of Correspondence to Joshua Robinson ( joshrob @ mit.edu ) . 35th Conference on Neural Information Processing Systems ( NeurIPS 2021 ) . one feature for another using simple methods for adjusting the difficulty of instance discrimination . However , these methods have an important drawback : improved learning of one feature often comes at the cost of harming another . That is , feature suppression is still prevalent . In response , we propose implicit feature modification , a technique that encourages encoders to discriminate instances using multiple input features . Our method introduces no computational overhead , reduces feature suppression ( without trade-offs ) , and improves generalization on various downstream tasks . Contributions . In summary , this paper makes the following main contributions : 1 . It analyzes feature suppression in contrastive learning , and explains why feature suppression can occur when optimizing the InfoNCE loss . 2 . It studies the relation between instance discrimination tasks and feature learning ; concretely , adjustments to instance discrimination difficulty leads to different features being learned . 3 . It proposes implicit feature modification , a simple and efficient method that reduces the tendency to use feature suppressing shortcut solutions and improves generalization . 1.1 Related work . Unsupervised representation learning is enjoying a renaissance driven by steady advances in effective frameworks [ 3 , 5 , 15 , 18 , 33 , 44 , 45 , 51 ] . As well as many effective contrastive methods , Siamese approaches that avoid representation collapse without explicitly use of negatives have also been proposed [ 6 , 13 , 51 ] . Pretext task design has been at the core of progress in self-supervised learning . Previously popular tasks include image colorization [ 54 ] and inpainting [ 35 ] , and theoretical work shows pre-trained encoders can provably generalize if a pretext task necessitates the learning of features that solve downstream tasks [ 27 , 39 ] . In contrastive learning , augmentation strategies are a key design component [ 5 , 48 , 50 ] , as are negative mining techniques [ 9 , 15 , 25 , 40 ] . While feature learning in contrastive learning has received less attention , recent work finds that low- and mid-level features are more important for transfer learning [ 55 ] , and feature suppression can occur [ 4 ] just as with supervised learning [ 10 , 16 ] . Combining contrastive learning with an auto-encoder has also been considered [ 28 ] , but was found to harm representation of some features in order to avoid suppression of others . Our work is distinguished from prior work through our focus on how the design of the instance discrimination task itself affects which features are learned . 2 Feature suppression in contrastive learning . Feature suppression refers to the phenomenon where , in the presence of multiple predictive input features , a model uses only a subset of them and ignores the others . The selected subset often corresponds to intuitively “ simpler ” features , e.g. , color as opposed to shape . Such features lead to “ shortcut ” decision rules that might perform well on training data , but can harm generalization and lead to poor robustness to data shifts . Feature suppression has been identified as a common problem in deep learning [ 11 ] , and both supervised and contrastive learning suffer from biases induced by the choice of optimizer and architecture . However , contrastive learning bears an additional potential source of bias : the choice of instance discrimination task . Which positive and negative pairs are presented critically affects which features are discriminative , and hence which features are learned . In this work we study the relation between feature suppression and instance discrimination . First , we explain why optimizing the InfoNCE loss is insufficient in general to avoid feature suppression , and show how it can lead to counter-intuitive generalization ( Sec . 2.2 ) . Given this negative result , we then ask if it is at least possible to control which features a contrastive encoder learns ? We find that this is indeed the case , and that adjustments to the instance discrimination task lead to different features being learned ( Sec . 2.3 ) . However , the primary drawback of these adjustments is that improving one feature often comes at the cost of harming representation of another . That is , feature suppression is still prevalent . Addressing this drawback is the focus of Sec . 3 . 2.1 Setup and definition of feature suppression . Formally , we assume that the data has underlying feature spaces Z1 , . . . , Zn with a distribution pj on each Zj . Each j 2 [ n ] , corresponding to a latent space Zj , models a distinct feature . We write the product as ZS = Q j2S Zj , and simply write Z instead of Z [ n ] where [ n ] = { 1 , . . . , n } . A set of features z = ( zj ) j2 [ n ] 2 Z is generated by sampling each coordinate zj 2 Zj independently , and we denote the measure on Z induced by z by . Further , let ( ·|zS ) denote the conditional measure on Z for fixed zS . For S ✓ [ n ] we use zS to denote the projection of z onto ZS . Finally , an injective map g : Z ! X produces observations x = g ( z ) . Our aim is to train an encoder f : X ! Sd 1 to map input data x to the surface of the unit sphere Sd 1 = { u 2 Rd : kuk2 = 1 } in such a way that f extracts useful information . To formally define feature suppression , we need the pushforward h # ⌫ ( V ) = ⌫ ( h 1 ( V ) ) of a measure ⌫ on a space U for a measurable map h : U ! V and measurable V ✓ V , where h 1 ( V ) denotes the preimage . Consider an encoder f : X ! Sd 1 and features S ✓ [ n ] . For each zS 2 ZS , let µ ( ·|zS ) = ( f g ) # ( ·|zS ) be the pushforward measure on Sd 1 by f g of the conditional ( ·|zS ) . 1. f suppresses S if for any pair zS , z̄S 2 ZS , we have µ ( ·|zS ) = µ ( ·|z̄S ) . 2. f distinguishes S if for any pair of distinct zS , z̄S 2 ZS , measures µ ( ·|zS ) , µ ( ·|z̄S ) have disjoint support . Feature suppression is thus captured in a distributional manner , stating that S is suppressed if the encoder distributes inputs in a way that is invariant to the value zS . Distinguishing features , meanwhile , asks that the encoder f separates points with different features zS into disjoint regions . We consider training an encoder f : X ! Sd 1 to optimize the InfoNCE loss [ 33 , 14 ] , Lm ( f ) = Ex , x+ , { x i } mi=1 log e f ( x ) > f ( x+ ) /⌧ ef ( x ) > f ( x+ ) /⌧ + Pm i=1 e f ( x ) > f ( x i ) /⌧ , ( 1 ) where ⌧ is known as the temperature . Positive pairs x , x+ are generated by first sampling z ⇠ , then independently sampling two random augmentations a , a+ ⇠ A , a : X ! X from a distribution A , and setting x = a ( g ( z ) ) and x+ = a+ ( g ( z ) ) . We assume A samples the identity function a ( x ) = x with non-zero probability ( “ x is similar to itself ” ) , and that there are no collisions : a ( x ) 6= a0 ( x0 ) for all a , a0 , and all x 6= x0 . Each negative example x i is generated as x i = ai ( g ( zi ) ) , by independently sampling features zi ⇠ and an augmentation ai ⇠ A . 2.2 Why optimizing the InfoNCE loss can still lead to feature suppression . Do optimal solutions to the InfoNCE loss automatically avoid shortcut solutions ? Unfortunately , as we show in this section , this is not the case in general ; there exist both optimal solutions of the InfoNCE loss that do and solutions that do not suppress a given feature . Following previous work [ 40 , 49 , 56 ] , we analyze the loss as the number of negatives goes to infinity , L = lim m ! 1 Lm ( f ) logm 2⌧ = 12⌧ Ex , x+kf ( x ) f ( x + ) k2 + Ex+ log ⇥ Ex ef ( x + ) > f ( x ) /⌧ ⇤ . We subtract logm to ensure the limit is finite , and use x to denote a random sample with the same distribution as x i . Prop . 2.2 ( proved in App . A ) shows that , assuming the marginals pj are uniform , the InfoNCE loss is optimized both by encoders that suppress feature j , and by encoders that distinguish j . Suppose that pj is uniform on Zj = Sd 1 for all j 2 [ n ] . Then for any feature j 2 [ n ] there exists an encoder fsupp that suppresses feature j and encoder fdisc that discriminates j but both attain minf : measurable L ( f ) . The condition that pj is uniformly distributed on Zj = Sd 1 is similar to conditions used in previous work [ 56 ] . Prop . 2.2 shows that empirical observations of feature suppression [ 4 ] ( see also Fig . 3 ) are not simply due to a failure to sufficiently optimize the loss , but that the possibility of feature suppression is built into the loss . What does Prop . 2.2 imply for the generalization behavior of encoders ? Besides explaining why feature suppression can occur , Prop . 2.2 also suggests another counter-intuitive possibility : lower InfoNCE loss may actually lead to worse performance on some tasks . To empirically study whether this possibility manifests in practice , we use two datasets with known semantic features : ( 1 ) In the Trifeature data , [ 16 ] each image is 128⇥128 and has three features : color , shape , and texture , each taking possible 10 values . See Fig . 10 , App . C for sample images . ( 2 ) In the STL-digits data , samples combine MNIST digits and STL10 objects by placing copies of a randomly selected MNIST digit on top of an STL10 image . See Fig . 11 App . C for sample images . We train encoders with ResNet-18 backbone using SimCLR [ 5 ] . To study correlations between the loss value and error on downstream tasks , we train 33 encoders on Trifea- ture and 7 encoders on STL-digits with different hyperparameter settings ( see App . C.2 for full details on training and hyperparameters ) . For Trifeature , we compute the Pearson correlation between InfoNCE loss and linear readout error when predicting { color , shape , texture } . Likewise , for STL-digits we compute correlations between the InfoNCE loss and MNIST and STL10 prediction error . Fig . 2 shows that performance on different downstream tasks is not always positively correlated . For Trifeature , color error is negatively correlated with shape and texture , while for STL-digits there is a strong negative correlation between MNIST digit error and STL10 error . Importantly , lower InfoNCE loss is correlated with lower prediction error for color and MNIST-digit , but with larger error for shape , texture and STL10 . Hence , lower InfoNCE loss can improve representation of some features ( color , MNIST digit ) , but may actually hurt others . This conflict is likely due to the simpler color and MNIST digit features being used as shortcuts . Our observation is an important addition to the statement of Wang and Isola [ 49 ] that lower InfoNCE loss improves generalization : the situation is more subtle – whether lower InfoNCE helps generalization on a task depends on the use of shortcuts .

Summary: This paper investigates feature learning in contrastive learning. It argues that feature suppression occurs in contrastive learning, where the model ignores a subset of the features. It argues that by varying the difficulty of the contrastive discrimination task, it is possible to incentivize the learning of particular features, but again at the expense of other features. To deal with this feature suppression problem, the authors propose a new method called implicit feature modification, where the basic idea is to continuously remove features from the embedding representation that are being used to discriminate negatives from positives. This encourages the model to learn new, additional features to perform the discrimination task. The authors show that this scheme makes the model learn more robust features (in the sense of Ilyas et al. (2019) --ref. 22--).