Equivariant Self-Supervised Learning: Encouraging Equivariance in Representations

1 INTRODUCTION . Human knowledge about what makes a good representation and the abundance of unlabeled data has enabled the learning of useful representations via self-supervised learning ( SSL ) pretext tasks . State-of-the-art SSL methods encourage the representations not to contain information about the way the inputs are transformed , i.e . to be invariant to a set of manually chosen transformations . One such method is contrastive learning , which sets up a binary classification problem to learn invariant features . Given a set of datapoints ( say images ) , different transformations of the same data point constitute positive examples , whereas transformations of other data points constitute the negatives ( He et al. , 2020 ; Chen et al. , 2020 ) . Beyond contrastive learning , many SSL methods also rely on learning representations by encouraging invariance ( Grill et al. , 2020 ; Chen & He , 2021 ; Caron et al. , 2021 ; Zbontar et al. , 2021 ) . Here , we refer to such methods as Invariant-SSL ( I-SSL ) . The natural question in I-SSL is to what transformations should the representations be insensitive ( Chen et al. , 2020 ; Tian et al. , 2020 ; Xiao et al. , 2020 ) . Chen et al . ( 2020 ) highlighted the importance of transformations and empirically evaluated which transformations are useful for contrastive learning ( e.g. , see Figure 5 in their paper ) . Some transformations , such as four-fold rotations , despite preserving semantic information , were shown to be harmful for contrastive learning . This does not mean that four-fold rotations are not useful for I-SSL at all . In fact , predicting four-fold rotations is a good proxy task for evaluating the representations produced with contrastive learning ( Reed et al. , 2021 ) . Furthermore , instead of being insensitive to rotations ( invariance ) , training a neural network to predict them , i.e . to be sensitive to four-fold rotations , results in good image representations ( Gidaris et al. , 2018 ; 2019 ) . These results indicate that the choice of making features sensitive or insensitive to a particular group of transformations can have a substantial effect on the performance of downstream tasks . However , the prior work in SSL has exclusively focused on being either entirely insensitive ( Grill et al. , 2020 ; Chen & He , 2021 ; Caron et al. , 2021 ; Zbontar et al. , 2021 ) or sensitive ( Agrawal et al. , 2015 ; Doersch et al. , 2015 ; Zhang et al. , 2016 ; Noroozi & Favaro , 2016 ; Gidaris et al. , 2018 ) to a set of transformations . In particular , the I-SSL literature has proposed to simply remove transformations that hurt performance when applied as invariance . To understand how sensitivity/ insensitivity to a particular transformation affects the resulting features , we ran a series of experiments summarized in Figure 1 . We trained and tested a simple I-SSL baseline , SimCLR ( Chen et al. , 2020 ) , on CIFAR-10 using only the random resized cropping transformation ( solid yellow line ) . The test accuracy is calculated as the retrieval accuracy of a k-nearest neighbors ( kNN ) classifier with a memory bank consisting of the representations on the training set obtained after pre-training for 800 epochs . Next , in addition to being invariant to resized cropping , we additionally encouraged the model to be either sensitive ( shown in pink ) or insensitive ( shown in blue ) to a second transformation . We encourage insensitivity by adding the transformation to the SimCLR data augmentation , and sensitivity by predicting it ( see Section 4 ) . We varied the choice of this second transformation . We found that for some transformations , such as horizontal flips and grayscale , insenstivity results in better features , but is detrimental for transformations , such as fourfold rotations , vertical flips , 2x2 jigsaws ( 4 ! = 24 classes ) , four-fold Gaussian blurs ( 4 levels of blurring ) and color inversions . When we encourage sensitivity to these transformations , the trend is reversed . In summary , we observe that if invariance to a particular transformation hurts feature learning , then imposing sensitivity to the same transformation may improve performance . This leads us to conjecture that instead of choosing the features to be only invariant or only sensitive as done in prior work , it may be possible to learn better features by imposing invariance to certain transformations ( e.g. , cropping ) and sensitivity to other transformations ( e.g. , four-fold transformations ) . The concepts of sensitivity and insensitivity are both captured by the mathematical idea of equivariance ( Agrawal et al. , 2015 ; Jayaraman & Grauman , 2015 ; Cohen & Welling , 2016 ) . Let G be a group of transformations . For any g ∈ G let Tg ( x ) denote the function with which g transforms an input image x . For instance , if G is the group of four-fold rotations then Tg ( x ) rotates the image x by a multiple of π/2 . Let f be the encoder network that computes feature representation , f ( x ) . I-SSL encourages the property of “ invariance to G , ” which states f ( Tg ( x ) ) = f ( x ) , i.e . the output representation , f ( x ) , does not vary with Tg . Equivariance , a generalization of invariance , is defined as , ∀x : f ( Tg ( x ) ) = T ′g ( f ( x ) ) , where T ′g is a fixed transformation ( i.e. , without any parameters ) . Intuitively , equivariance encourages the feature representation to change in a well defined manner to the transformation applied to the input . Thus , invariance is a trivial instance of equivariance , where T ′g is the identity function , i.e . T ′ g ( f ( x ) ) = f ( x ) .While there are many possible choices for T ′ g ( Cohen & Welling , 2016 ; Bronstein et al. , 2021 ) , I-SSL uses only the trivial choice that encourages f to be insensitive to G. In contrast , if T ′g is not the identity , then f will be sensitive to G and we say that the “ equivariance to G ” will be non-trivial . Therefore , in order to encourage potentially more useful equivariance properties , we generalize SSL to an Equivariant Self-Supervised Learning ( E-SSL ) framework . In our experiments on standard computer vision data , such as the small-scale CIFAR-10 ( Torralba et al. , 2008 ; Krizhevsky , 2009 ) and the large-scale ImageNet ( Deng et al. , 2009 ) , we show that extending I-SSL to E-SSL by also predicting four-fold rotations improves the semantic quality of the representations . We show that this approach works for other transformations too , such as vertical flips , 2x2 jigsaws , four-fold Gaussian blurs and color inversions , but focus on four-fold rotations as the most promising improvement we obtain with initial E-SSL experiments in Figure 1 . We also note that the applications of E-SSL in this paper are task specific , meaning that the representations from E-SSL may work best for a particular downstream task that benefits from equivariance dictated by the available data . E-SSL can be further extended to applications in science ; in particular , we focus on predictive tasks using ( unlabelled and labelled ) data collected via experiments or simulations . The downstream tasks in prediction problems in science are often fixed and can be aided by incorporating scientific insights . Here , we also explore the generality of E-SSL beyond computer vision , on a different application : regression problems in photonics science and demonstrate examples where E-SSL is effective over I-SSL . Our contributions can be summarized as follows : • We introduce E-SSL , a generalization of popular SSL methods that highlights the complementary nature of invariance and equivariance . To our knowledge , we are the first to create a method that benefits from such complementarity . • We improve state-of-the-art SSL methods on CIFAR-10 and ImageNet by encouraging equivariance to four-fold rotations . We also show that E-SSL is more general and works for many other transformations , previously unexplored in related works . • We demonstrate the usefulness of E-SSL beyond computer vision with experiments on regression problems in photonics science . We also show that our method works both for finite and infinite groups . The rest of the paper is organized as follows . In Section 3 we introduce our experimental method for E-SSL . In Section 4 we present our main experiments in computer vision . In Section 2 we elaborate on related work . In Section 5 provide a discussion around our work that extends our study beond computer vision . In Section 6 we conclude and point to future work . 2 RELATED WORK . To encourage non-trivial equivariance , we observe that a simple task that predicts the synthetic transformation applied to the input , works well and improves I-SSL already ; some prediction tasks create representations that can be transferred to other tasks of interest , such as classification , object detection and segmentation . While prediction tasks alone have been realized successfully before in SSL ( Agrawal et al. , 2015 ; Doersch et al. , 2015 ; Zhang et al. , 2016 ; Misra et al. , 2016 ; Noroozi & Favaro , 2016 ; Zamir et al. , 2016 ; Lee et al. , 2017 ; Mundhenk et al. , 2018 ; Gidaris et al. , 2018 ; Zhang et al. , 2019 ; Zhang , 2020 ) , to our knowledge we are the first to combine simple predictive objectives of synthetic transformations with I-SSL , and successfully improve the semantic quality of representations . We found that the notion of equivariance captures the generality of our method . To improve representations with pretext tasks , Gidaris et al . ( 2018 ) use four-fold rotations prediction as a pretext task for learning useful visual representations via a new model named RotNet . Feng et al . ( 2019 ) learn decoupled representations : one part trained with four-fold rotations prediction and another with non-parametric instance discrimination ( Wu et al. , 2018 ) and invariance to four-fold rotations . Yamaguchi et al . ( 2021 ) use a joint training objective between four-fold rotations prediction and image enhancement prediction . Xiao et al . ( 2020 ) propose to learn representations as follows : for each atomic augmentation from the contrastive learning ’ s augmentation policy , they leave it out and project to a new space on which I-SSL encourages invariance to all augmentations , but the leftout one . The resulting representation could either be a concatenation of all projected left-out views ’ representations , or the representation in the shared space , before the individual projections . Our method differs from the above contributions in that E-SSL is the only hybrid framework that encourages both insensitive representations for some transformations and sensitive representations for others and does not require representations to be sensitive and insensitive to a particular transformation at the same time . Thus , what distinguished our work is the complementary nature of invariance and equivariance for multiple transformations , including finite and infinite groups . To obtain performance gains from transformations , Tian et al . ( 2020 ) study which transformations are the best for contrastive learning through the lens of mutual information . Reed et al . ( 2021 ) use four-fold rotations prediction as an evaluation measure to tune optimal augmentations for contrastive learning . Wang & Qi ( 2021 ) use strong augmentations to improve contrastive learning by matching the distributions of strongly and weakly augmented views ’ representation similarities to a memory bank . A growing body of work encourages invariance to domain agnostic transformations ( Tamkin et al. , 2021 ; Lee et al. , 2021 ; Verma et al. , 2021 ) or strengthens invariance with regularization ( Foster et al. , 2021 ) . Our framework is different from the above works , because we work with transformations that encourage equivariance beyond invariance . To understand and improve equivariant properties of neural networks , Lenc & Vedaldi ( 2015 ) study emerging equivariant properties of neural networks and ( Cohen & Welling , 2016 ; Bronstein et al. , 2021 ) construct equivariant neural networks . In contrast , our work does not enforce strict equivariance , but only encourages equivariant properties for the encoder network through the choice of the loss function . While strict equivariance is concerned with groups , some of the transformations , such as random resized cropping and Gaussian blurs , may not even be groups , but they could still be analyzed in the E-SSL framework . Thus , ours is a flexible framework , which allows us to consider a variety of transformations and how the encoder might exhibit equivariant properties to them .

Summary: This work presents a framework for learning representations with invariances (insensitivity) to some transformations and sensitivities to others.  Previous work has only considered representations with insensitivity or sensitivity but not a combination.  The authors use the concept of invariance or equivariance to symmetry group transformations to learn such representations.  The advantage of the proposed method is demonstrated relative to SimCLR and SimSiam on Cifar-10, ImageNet, and a new scientific application domain: learning frequency responses of photonic crystals.  The authors also prove a theoretical characterization of when their methods works.