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: The paper advocates Equivariance Self-Supervised Learning (E-SSL) as a more general framework than the Invariance SSL (I-SSL), which is common in SoTA vision SSL methods, such as SimCLR and Barlow Twins. The paper starts with empirical evidence (Fig-1) that adding equivariance objectives (e.g. 4-fold rotation and vertical flips) to SimCLR can improve performance. The proposed E-SSL framework boils down to adding an additional equivariance objective (mainly 4-fold rotations) to popular I-SSL methods (Eqn. 1 & 2). The empirical results show encouraging results in CIFAR-10 and ImageNet. Finally, the paper applies the proposed method in two datasets of photonic data (regression task from 2D square periodic unit cells).