Whitening for Self-Supervised Representation Learning

1 INTRODUCTION . One of the current main bottlenecks in deep network training is the dependence on large annotated training datasets , and this motivates the recent surge of interest in unsupervised methods . Specifically , in self-supervised representation learning , a network is ( pre- ) trained without any form of manual annotation , thus providing a means to extract information from unlabeled-data sources ( e.g. , text corpora , videos , images from the Internet , etc. ) . In self-supervision , label information is replaced by a prediction problem using some form of context or using a pretext task . Pioneering work in this direction was done in Natural Language Processing ( NLP ) , in which the co-occurrence of words in a sentence is used to learn a language model ( Mikolov et al. , 2013a ; b ; Devlin et al. , 2019 ) . In Computer Vision , typical contexts or pretext tasks are based on : ( 1 ) the temporal consistency in videos ( Wang & Gupta , 2015 ; Misra et al. , 2016 ; Dwibedi et al. , 2019 ) , ( 2 ) the spatial order of patches in still images ( Noroozi & Favaro , 2016 ; Misra & van der Maaten , 2019 ; Hénaff et al. , 2019 ) or ( 3 ) simple image transformation techniques ( Ji et al. , 2019 ; He et al. , 2019 ; Wu et al. , 2018 ) . The intuitive idea behind most of these methods is to collect pairs of positive and negative samples : two positive samples should share the same semantics , while negatives should be perceptually different . A triplet loss ( Sohn , 2016 ; Schroff et al. , 2015 ; Hermans et al. , 2017 ; Wang & Gupta , 2015 ; Misra et al. , 2016 ) can then be used to learn a metric space which should represent the human perceptual similarity . However , most of the recent studies use a contrastive loss ( Hadsell et al. , 2006 ) or one of its variants ( Gutmann & Hyvärinen , 2010 ; van den Oord et al. , 2018 ; Hjelm et al. , 2019 ) , while Tschannen et al . ( 2019 ) show the relation between the triplet loss and the contrastive loss . It is worth noticing that the success of both kinds of losses is strongly affected by the number and the quality of the negative samples . For instance , in the case of the triplet loss , a common practice is to select hard/semi-hard negatives ( Schroff et al. , 2015 ; Hermans et al. , 2017 ) . On the other hand , Hjelm et al . ( 2019 ) have shown that the contrastive loss needs a large number of negatives to be competitive . This implies using batches with a large size , which is computationally demanding , especially with high-resolution images . In order to alleviate this problem , Wu et al . ( 2018 ) use a memory bank of negatives , which is composed of feature-vector representations of all the training samples . He et al . ( 2019 ) conjecture that the use of large and fixed-representation vocabularies is one of the keys to the success of self-supervision in NLP . The solution proposed by He et al . ( 2019 ) extends Wu et al . ( 2018 ) using a memory-efficient queue of the last visited negatives , together with a momentum encoder which preserves the intra-queue representation consistency . Chen et al . ( 2020 ) have performed large-scale experiments confirming that a large number of negatives ( and therefore a large batch size ) is required for the contrastive loss to be efficient . Concurrently with our work , Grill et al . ( 2020 ) have suggested that it is not necessary to rely on the contrastive scheme , introducing a high-performing alternative based on bootstrapping . In this paper we propose a new self-supervised loss function which first scatters all the sample representations in a spherical distribution1 and then penalizes the positive pairs which are far from each other . In more detail , given a set of samples V = { vi } , corresponding to the current minibatch of images B = { xi } , we first project the elements of V onto a spherical distribution using a whitening transform ( Siarohin et al. , 2019 ) . The whitened representations { zi } , corresponding to V , are normalized and then used to compute a Mean Squared Error ( MSE ) loss which accumulates the error taking into account only positive pairs ( zi , zj ) . We do not need to contrast positives against negatives as in the contrastive loss or in the triplet loss because the optimization process leads to shrinking the distance between positive pairs and , indirectly , scatters the other samples to satisfy the overall spherical-distribution constraint . In summary , our contributions are the following : • We propose a new loss function , Whitening MSE ( W-MSE ) , for self-supervised training . W-MSE constrains the batch samples to lie in a spherical distribution and it is an alternative to positive-negative instance contrasting methods . • Our loss does not rely on negatives , thus including more positive samples in the batch can be beneficial ; we indeed demonstrate that multiple positive pairs extracted from one image improve the performance . • We empirically show that our W-MSE loss outperforms the commonly adopted contrastive loss when measured using different standard classification protocols . We show that W-MSE is competitive with respect to state-of-the-art self-supervised methods . 2 BACKGROUND AND RELATED WORK . A typical self-supervised method is composed of two main components : a pretext task , which exploits some a-priori knowledge about the domain to automatically extract supervision from data , and a loss function . In this section we briefly review both aspects , and we additionally analyse the recent literature concerning feature whitening . Pretext Tasks . The temporal consistency in a video provides an intuitive form of self-supervision : temporally-close frames usually contain a similar semantic content ( Wang & Gupta , 2015 ; van den Oord et al. , 2018 ) . Misra et al . ( 2016 ) extended this idea using the relative temporal order of 3 frames , while Dwibedi et al . ( 2019 ) used a temporal cycle consistency for self-supervision , which is based on comparing two videos sharing the same semantics and computing inter-video frame-toframe nearest neighbour assignments . When dealing with still images , the most common pretext task is instance discrimination ( Wu et al . ( 2018 ) ) : from a training image x , a composition of data-augmentation techniques are used to extract two different views of x ( xi and xj ) . Commonly adopted transformations are : image cropping , rotation , color jittering , Sobel filtering , etc .. The learner is then required to discriminate ( xi , xj ) from other views extracted from other samples ( Wu et al. , 2018 ; Ji et al. , 2019 ; He et al. , 2019 ; Chen et al. , 2020 ) . Denoising auto-encoders ( Vincent et al. , 2008 ) add random noise to the input image and try to recover the original image . More sophisticated pretext tasks consist in predicting the spatial order of image patches ( Noroozi & Favaro , 2016 ; Misra & van der Maaten , 2019 ) or in reconstructing large masked regions of the image ( Pathak et al. , 2016 ) . Hjelm et al . ( 2019 ) ; Bachman et al . ( 2019 ) compare the holistic representation of an input image with a patch of the same image . Hénaff et al . 1Here and in the following , with “ spherical distribution ” we mean a distribution with a zero-mean and an identity-matrix covariance . ( 2019 ) use a similar idea , where the comparison depends on the patch order : the appearance of a given patch should be predicted given the appearance of the patches which lie above it in the image . In this paper we use standard data augmentation techniques on still images to obtain positive pairs , which is a simple method to get self-supervision ( Chen et al. , 2020 ) and does not require a pretexttask specific network architecture ( Hjelm et al. , 2019 ; Bachman et al. , 2019 ; Hénaff et al. , 2019 ) . Loss functions . Denoising auto-encoders use a reconstruction loss which compares the generated image with the input image before adding noise . Other generative methods use an adversarial loss in which a discriminator provides supervisory information to the generator ( Donahue et al. , 2017 ; Donahue & Simonyan , 2019 ) . Early self-supervised ( deep ) discriminative methods used a triplet loss ( Wang & Gupta , 2015 ; Misra et al. , 2016 ) : given two positive images xi , xj and a negative xk ( Sec . 1 ) , together with their corresponding latent-space representations zi , zj , zk , this loss penalizes those cases in which zi and zk are closer to each other than zi and zj plus a margin m : LTriplet = −max ( zTi zk − zTi zj +m , 0 ) . ( 1 ) Most of the recent self-supervised discriminative methods are based on some contrastive loss ( Hadsell et al. , 2006 ) variant , in which zi and zj are contrasted against a set of negative pairs . Following the common formulation proposed by van den Oord et al . ( 2018 ) : LContrastive = − log exp ( zTi zj/τ ) ∑K k=1 , k 6=i exp ( z T i zk/τ ) , ( 2 ) where τ is a temperature hyperparameter which should be manually set and the sum in the denominator is over a set of K − 1 negative samples . Usually K is the size of the current batch , i.e. , K = 2N , being N the number of the positive pairs . However , as shown by Hjelm et al . ( 2019 ) , the contrastive loss ( 2 ) requires a large number of negative samples to be competitive . Wu et al . ( 2018 ) ; He et al . ( 2019 ) use a set of negatives much larger than the current batch , by pre-computing latent-space representations of old samples . SimCLR ( Chen et al . ( 2020 ) ) uses a simpler , but computationally very demanding , solution based on large batches . While recent works ( van den Oord et al. , 2018 ; Hénaff et al. , 2019 ; Hjelm et al. , 2019 ; Bachman et al. , 2019 ; Ravanelli & Bengio , 2018 ) draw a relation between the contrastive loss and an estimate of the mutual information between the latent-space image representations , Tschannen et al . ( 2019 ) showed that the success of this loss is likely related to learning a metric space , similarly to what happens with a triplet loss . On the other hand , Wang & Isola ( 2020 ) showed that the L2 normalized contrastive loss asymptotically converges to the minimization of two desirable characteristics of the latent-space representations on the surface of the unit hypersphere : uniformity and semantic alignment . In the same paper , the authors propose two new losses ( Luniform and Lalign ) which explicitly deal with these characteristics . Concurrently with our work , BYOL ( Grill et al . ( 2020 ) ) proposes a “ bootstrapping ” scheme which is alternative to the positive-negative contrastive learning . In BYOL , an “ online ” network is optimised to predict the output of a “ target ” network , whose parameters are a running average of the online network . The predictions of the two networks are compared using an additional prediction network and an MSE loss . However , very recently , Fetterman & Albrecht ( 2020 ) and Tian et al . ( 2020 ) have empirically shown that BYOL can avoid a collapsed solution through the use of the Batch Norm ( BN ) ( Ioffe & Szegedy , 2015 ) which avoids constant representations . Our work can be seen as a generalization of this finding with a much simpler network architecture ( more details in Sec . 3.1 ) . In this paper we propose a different loss which is competitive with respect to other alternatives . Our loss formulation is simpler because it does not require a proper setting of the τ hyperparameter in equation 2 , m in equation 1 , or additional networks with a specific weight update schemes as in BYOL . Feature Whitening . We adopt the efficient and stable Cholesky decomposition ( Dereniowski & Marek , 2004 ) based whitening transform proposed by Siarohin et al . ( 2019 ) to project our latentspace vectors into a spherical distribution ( more details in Sec . 3 ) . Note that Huang et al . ( 2018 ) ; Siarohin et al . ( 2019 ) use whitening transforms in the intermediate layers of the network for a completely different task : extending BN to a multivariate batch normalization .

Summary: The paper proposes to first do representation "whitening", so that the representations are scattered in the space and not collapsing to a single data point; then compute distance metric on top of that (e.g. Euclidean, cosine similarity). A nice thing about explicit scattering is that it does not require large numbers of negative examples to pull the features apart. Experiments are done on several toy datasets like CIFAR.