Improving Transformation Invariance in Contrastive Representation Learning

1 INTRODUCTION . Learning meaningful representations of data is a central endeavour in artificial intelligence . Such representations should retain important information about the original input whilst using fewer bits to store it ( van der Maaten et al. , 2009 ; Gregor et al. , 2016 ) . Semantically meaningful representations may discard a great deal of information about the input , whilst capturing what is relevant . Knowing what to discard , as well as what to keep , is key to obtaining powerful representations . By defining transformations that are believed a priori to distort the original without altering semantic features of interest , we can learn representations that are ( approximately ) invariant to these transformations ( Hadsell et al. , 2006 ) . Such representations may be more efficient and more generalizable than lossless encodings . Whilst less effective for reconstruction , these representations are useful in many downstream tasks that relate only to the semantic features of the input . Representation invariance is also a critically important task in of itself : it can lead to improved robustness and remove noise ( Du et al. , 2020 ) , afford fairness in downstream predictions ( Jaiswal et al. , 2020 ) , and enhance interpretability ( Xu et al. , 2018 ) . Contrastive learning is a recent and highly successful self-supervized approach to representation learning that has achieved state-of-the-art performance in tasks that rely on semantic features , rather than exact reconstruction ( van den Oord et al. , 2018 ; Hjelm et al. , 2018 ; Bachman et al. , 2019 ; He et al. , 2019 ) . These methods learn to match two different transformations of the same object in representation space , distinguishing them from contrasts that are representations of other objects . The objective functions used for contrastive learning encourage representations to remain similar under transformation , whilst simultaneously requiring different inputs to be well spread out in representation space ( Wang & Isola , 2020 ) . As such , the choice of transformations is key to their success ( Chen et al. , 2020a ) . Typical choices include random cropping and colour distortion . However , representations are compared using a similarity function that can be maximized even for representations that are far apart , meaning that the invariance learned is relatively weak . Unfor∗Equal contribution tunately , directly changing the similarity measure hampers the algorithm ( Wu et al. , 2018 ; Chen et al. , 2020a ) . We therefore investigate methods to improve contrastive representations by explicitly encouraging stronger invariance to the set of transformations , without changing the core selfsupervized objective ; we look to extract more information about how representations are changing with respect to transformation , and use this to direct the encoder towards greater invariance . To this end , we first develop a gradient regularization term that , when included in the training loss , forces the encoder to learn a representation function that varies slowly with continuous transformations . This can be seen as constraining the encoder to be approximately transformation invariant . We demonstrate empirically that while the parameters of the transformation can be recovered from standard contrastive learning representations using just linear regression , this is no longer the case when our regularization is used . Moreover , our representations perform better on downstream tasks and are robust to the introduction of nuisance transformations at test time . Test representations are conventionally produced using untransformed inputs ( Hjelm et al. , 2018 ; Kolesnikov et al. , 2019 ) , but this fails to combine information from different transformations and views of the object , or to emulate settings in which transformation noise can not simply be removed at test time . Our second key proposal is to instead create test time representations by feature averaging over multiple , differently transformed , inputs to address these concerns and to more directly impose invariance . We show theoretically that this leads to improved performance under linear evaluation protocols , further confirming this result empirically . We evaluate our approaches first on CIFAR-10 and CIFAR-100 ( Krizhevsky et al. , 2009 ) , using transformations appropriate to natural images and evaluating on a downstream classification task . To validate that our ideas transfer to other settings , and to use our gradient regularizer within a fully differentiable generative process , we further introduce a new synthetic dataset called Spirograph . This provides a greater variety of downstream regression tasks , and allows us to explore the interplay between nuisance transformations and generative factors of interest . We confirm that using our regularizer during training and our feature averaging at test time both improve performance in terms of transformation invariance , downstream tasks , and robustness to train–test distributional shift . In summary , the contributions of this paper are as follows : • We derive a novel contrastive learning objective that leads to more invariant representations . • We propose test time feature averaging to enforce further invariance . • We introduce the Spirograph dataset . • We show empirically that our approaches lead to more invariant representations and achieve state-of-the-art performance for existing downstream task benchmarks . 2 PROBABILISTIC FORMULATION OF CONTRASTIVE LEARNING . The goal of unsupervized representation learning is to encode high-dimensional data , such as images , retaining information that may be pertinent to downstream tasks and discarding information that is not . To formalize this , we consider a data distribution p ( x ) on X and an encoder fθ : X → Z which is a parametrized function mapping from data space to representation space . Contrastive learning is a self-supervized approach to representation learning that learns to make representations of differently transformed versions of the same input more similar than representations of other inputs . Of central importance is the set of transformations , also called augmentations ( Chen et al. , 2020a ) or views ( Tian et al. , 2019 ) , used to distort the data input x . In the common application of computer vision , it is typical to include resized cropping ; brightness , contrast , saturation and hue distortion ; greyscale conversion ; and horizontal flipping . We will later introduce the Spirograph dataset which uses quite different transformations . In general , transformations are assumed to change the input only cosmetically , so all semantic features such as the class label are preserved ; the set of transformations indicates changes which can be safely ignored by the encoder . Formally , we consider a transformation set T ⊆ { t : X → X } and a probability distribution p ( t ) on this set . A representation z of x is obtained by applying a random transformation t to x and then encoding the result using fθ . Therefore , we do not have one representation of x , but an implicit distribution p ( z|x ) . A sample of p ( z|x ) is obtained by sampling t ∼ p ( t ) and setting z = fθ ( t ( x ) ) . If the encoder is to discard irrelevant information , we would expect different encodings of x formed with different transformations t to be close in representation space . Altering the transformation should not lead to big changes in the representations of the same input . In other words , the distribution p ( z|x ) should place most probability mass in a small region . However , this does not provide a sufficient training signal for the encoder fθ as it fails to penalize trivial solutions in which all x are mapped to the same z . To preserve meaningful information about the input x whilst discarding purely cosmetic features , we should require p ( z|x ) to be focused around a single z whilst simultaneously requiring the representations of different inputs not to be close . That is , the marginal p ( z ) = Ep ( x ) [ p ( z|x ) ] should distribute probability mass over representation space . This intuition is directly reflected in contrastive learning . Most state-of-the-art contrastive learning methods utilize the InfoNCE objective ( van den Oord et al. , 2018 ) , or close variants of it ( Chen et al. , 2020a ) . InfoNCE uses a batch x1 , ... , xK of inputs , from which we form pairs of representations ( z1 , z′1 ) , ... , ( zK , z ′ K ) by applying two random transformations to each input followed by the encoder fθ . In probabilistic language xi ∼ p ( x ) for i = 1 , ... , K ( 1 ) zi , z ′ i ∼ p ( z|x = xi ) conditionally independently given xi , for i = 1 , ... , K , ( 2 ) such that zi , z′i = fθ ( t ( x ) ) , fθ ( t ′ ( x ) ) for i.i.d . transformations t , t′ ∼ p ( t ) . Given a learnable similarity score sφ : Z × Z → R , contrastive learning methods minimize the following loss L ( θ , φ ) = − 1 K K∑ i=1 sφ ( zi , z ′ i ) + 1 K K∑ i=1 log K∑ j=1 exp [ sφ ( zi , z ′ j ) ] . ( 3 ) Written in this way , we see that the loss will be minimized when sφ ( zi , z′i ) is large , but sφ ( zi , z ′ j ) is small for i 6= j . In other words , InfoNCE makes the two samples zi , z′i of p ( z|x = xi ) similar , whilst making samples zi , z′j of p ( z ) dissimilar . This can also be understood through the lens of mutual information , for more details see Appendix A . In practice , the similarity measure used generally takes the form ( Chen et al. , 2020a ) sφ ( z , z ′ ) = gφ ( z ) > gφ ( z ′ ) τ‖gφ ( z ) ‖2‖gφ ( z′ ) ‖2 ( 4 ) where gφ is a small neural network and τ is a temperature hyperparameter . If the encoder fθ is perfectly invariant to the transformations , then zi = z′i and sφ ( zi , z ′ i ) will be maximal . However , there are many ways to maximize the InfoNCE objective without encouraging strong invariance in the encoder.1 In this paper , we show how we can learn stronger invariances , above and beyond what is learned through the above approach , and that this benefits downstream task performance . 3 INVARIANCE BY GRADIENT REGULARIZATION . Contrastive learning with InfoNCE can gently encourage invariance by maximizing sφ ( z , z′ ) , but does not provide a strong signal to ensure this invariance . Our first core contribution is to show how we can use gradient methods to directly regulate how the representation changes with the transformation and thus ensure the desired invariance . The key underlying idea is to differentiate the representation with respect to the transformation , and then encourage this gradient to be small so that the representation changes slowly as the transformation is varied . To formalize this , we begin by looking more closely at the transformations T which are used to define the distribution p ( z|x ) . Many transformations , such as brightness adjustment , are controlled by a transformation parameter . We can include these parameters in our set-up by writing the transformation t as a map from both input space X and transformation parameter space U , i.e . t : X ×U → X . In this formulation , we sample a random transformation parameter from u ∼ p ( u ) which is a distribution on U . A sample from p ( z|x ) is then obtained by taking z = fθ ( t ( x , u ) ) , with t now regarded as a fixed function . 1This is because the function gφ is not an injection , so we may have gφ ( z ) = gφ ( z′ ) but z 6= z′ . Johnson & Lindenstrauss ( 1984 ) gives conditions under which a projection of this form will preserve approximate distances , in particular , the required projection dimension is much larger than the typical value 128 . The advantage of this change of perspective is that it opens up additional ways to learn stronger invariance of the encoder . In particular , it may make sense to consider the gradient ∇uz , which describes the rate of change of z with respect to the transformation . This only makes sense for some transformation parameters—we can differentiate with respect to the brightness scaling but not with respect to a horizontal flip . To separate out differentiable and non-differentiable parameters we write u = α , β where α are the parameters for which it makes sense to consider the derivative ∇αz . Intuitively , this gradient should be small to ensure that representations change only slowly as the transformation parameterα is varied . For clarity of exposition , and for implementation practicalities , it is important to consider gradients of a scalar function , so we introduce an arbitrary direction vector e ∈ Z and define F ( α , β , x , e ) = e · fθ ( t ( x , α , β ) ) ‖fθ ( t ( x , α , β ) ) ‖2 ( 5 ) so that F : A× B × X × Z → R calculates the scalar projection of the normalized representation z/‖z‖2 in the e direction . To encourage an encoder that is invariant to changes in α , we would like to minimize the expected conditional variance of F with respect to α : V = Ep ( x ) p ( β ) p ( e ) [ Varp ( α ) [ F ( α , β , x , e ) | x , β , e ] ] , ( 6 ) where we have exploited independence to write p ( x , β , e ) = p ( x ) p ( β ) p ( e ) . Defining V requires a distribution for e to be specified . For this , we make components of e independent Rademacher random variables , justification for which is included in Appendix B . A naive estimator of V can be formed using a direct nested Monte Carlo estimator ( Rainforth et al. , 2018 ) of sample variances , which , including Bessel ’ s correction , is given by V ≈ 1 K K∑ i=1 1 L− 1 L∑ j=1 F ( αij , βi , xi , ei ) 2 − 1 L ( L− 1 ) [ L∑ k=1 F ( αik , βi , xi , ei ) ] 2 ( 7 ) where xi , βi , ei ∼ p ( x ) p ( β ) p ( e ) and αij ∼ p ( α ) . However , this estimator requires LK forward passes through the encoder fθ to evaluate . As an alternative to this computationally prohibitive approach , we consider a first-order approximation2 to F F ( α′ , β , x , e ) − F ( α , β , x , e ) = ∇αF ( α , β , x , e ) · ( α′ −α ) + o ( ‖α′ −α‖ ) ( 8 ) and the following alternative form for the conditional variance ( see Appendix B for a derivation ) Varp ( α ) [ F ( α , β , x , e ) | x , β , e ] = 12Ep ( α ) p ( α′ ) [ ( F ( α , β , x , e ) − F ( α′ , β , x , e ) ) 2 | x , β , e ] ( 9 ) Combining these two ideas , we have V = Ep ( x ) p ( β ) p ( e ) [ 1 2Ep ( α ) p ( α′ ) [ ( F ( α , β , x , e ) − F ( α′ , β , x , e ) ) 2 | x , β , e ] ] ( 10 ) ≈ Ep ( x ) p ( β ) p ( e ) [ 1 2Ep ( α ) p ( α′ ) [ ( ∇αF ( α , β , x , e ) · ( α′ −α ) ) 2 | x , β , e ] ] . ( 11 ) Here we have an approximation of the conditional variance V that uses gradient information . Including this as a regularizer within contrastive learning will encourage the encoder to reduce the magnitude of the conditional variance V , forcing the representation to change slowly as the transformation is varied and thus inducing approximate invariance to the transformations . An unbiased estimator of equation 11 using a batch x1 , ... , xK is V̂regularizer = 1 K K∑ i=1 1 2L L∑ j=1 [ ∇αF ( αi , βi , xi , ei ) · ( α′ij −αi ) ] 2 ( 12 ) where xi , αi , βi , ei , ∼ p ( x ) p ( α ) p ( β ) p ( e ) , α′ij ∼ p ( α ) . We can cheaply use a large number of samples for α′ without having to take any additional forward passes through the encoder : we only require K evaluations of F . Our final loss function is L ( θ , φ ) =− 1 K K∑ i=1 sφ ( zi , z ′ i ) + 1 K K∑ i=1 log K∑ j=1 exp [ sφ ( zi , z ′ j ) ] + λ LK K∑ i=1 L∑ j=1 [ ∇αF ( αi , βi , xi , ei ) · ( α′ij −αi ) ] 2 ( 13 ) 2We use the notation a ( x ) = o ( b ( x ) ) to mean a ( x ) /b ( x ) → 0 as x → ∞ . where λ is a hyperparameter controlling the regularization strength . This loss does not require us to encode a larger number of differently transformed inputs . Instead , it uses the gradient at ( x , α , β , e ) to control properties of the encoder in a neighbourhood ofα . This can effectively reduce the representation gradient along the directions corresponding to many different transformations . This , in turn , creates an encoder that is approximately invariant to the transformations .

Summary: Given one image, the paper first generates different views which are controlled by differentiable parameter \alpha, and then minimizes the additional "conditional variance" term~(expectation of these views' squared differences). Therefore, the paper encourages representations of the same image remain similar under the augmentation. A testing strategy is further proposed by voting features with different augmentations. Results demonstrate the effectiveness.