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: This paper proposes the mean square error loss with whitening operation to project positive pairs closely to each others while projecting the different positive pairs far away from each other on a unit sphere. This way, similar to BYOL, this paper removes the construction of negative pairs while improving the MoCo-V2 slightly on not very challenging benchmarks. The authors can find my questions/comments in the list below.