Self-Supervised Learning with Kernel Dependence Maximization

1 Introduction . Learning general-purpose visual representations without human supervision is a long-standing goal of machine learning . Specifically , we wish to find a feature extractor that captures the image semantics of a large unlabeled collection of images , so that e.g . various image understanding tasks can be achieved with simple linear models . One approach takes the latent representation of a likelihood-based generative model [ 2–8 ] ; such models , though , solve a harder problem than necessary since semantic features need not capture low-level details of the input . Another option is to train a self-supervised model for a “ pretext task , ” such as predicting the position of image patches , identifying rotations , or image inpainting [ 9–14 ] . Designing good pretext tasks , however , is a subtle art , with little theoretical guidance available . Recently , a class of models based on contrastive learning [ 15–22 ] has seen substantial success : dataset images are cropped , rotated , color shifted , etc . into several views , and features are then trained to pull together representations of the “ positive ” pairs of views of the same ∗These authors contributed equally . †Work done in part while at the Toyota Technological Institute at Chicago . 35th Conference on Neural Information Processing Systems ( NeurIPS 2021 ) . source image , and push apart those of “ negative ” pairs ( from different images ) . These methods are either understood from an information theoretic perspective as estimating the mutual information between the “ positives ” [ 15 ] , or explained as aligning features subject to a uniformity constraint [ 23 ] . Another line of research [ 24 , 25 ] attempts to learn representation without the “ negative ” pairs , but requires either a target network or stop-gradient operation to avoid collapsing . We examine the contrastive framework from a statistical dependence point of view : feature representations for a given transformed image should be highly dependent on the image identity ( Figure 2 ) . To measure dependence , we turn to the Hilbert-Schmidt Independence Criterion ( HSIC ) [ 26 ] , and propose a new loss for self-supervised learning which we call SSL-HSIC . Our loss is inspired by HSIC Bottleneck [ 27 , 28 ] , an alternative to Information Bottleneck [ 29 ] , where we use the image identity as the label , but change the regularization term . Through the dependence maximization perspective , we present a unified view of various selfsupervised losses . Previous work [ 30 ] has shown that the success of InfoNCE can not be solely attributed to properties of mutual information , in particular because mutual information ( unlike kernel measures of dependence ) has no notion of geometry in feature space : for instance , all invertible encoders achieve maximal mutual information , but they can output dramatically different representations with very different downstream performance [ 30 ] . Variational bounds on mutual information do impart notions of locality that allow them to succeed in practice , departing from the mutual information quantity that they try to estimate . We prove that InfoNCE , a popular such bound , in fact approximates SSL-HSIC with a variance-based regularization . Thus , InfoNCE can be thought of as working because it implicitly estimates a kernel-based notion of dependence . We additionally show SSL-HSIC is related to metric learning , where the features learn to align to the structure induced by the self-supervised labels . This perspective is closely related to the objective of BYOL [ 25 ] , and can explain properties such as alignment and uniformity [ 23 ] observed in contrastive learning . Our perspective brings additional advantages , in computation and in simplicity of the algorithm , compared with existing approaches . Unlike the indirect variational bounds on mutual information [ 15 , 31 , 32 ] , SSL-HSIC can be directly estimated from mini-batches of data . Unlike “ negative-free ” methods , the SSL-HSIC loss itself penalizes trivial solutions , so techniques such as target networks are not needed for reasonable outcomes . Using a target network does improve the performance of our method , however , suggesting target networks have other advantages that are not yet well understood . Finally , we employ random Fourier features [ 33 ] in our implementation , resulting in cost linear in batch size . Our main contributions are as follows : • We introduce SSL-HSIC , a principled self-supervised loss using kernel dependence maximization . • We present a unified view of contrastive learning through dependence maximization , by establishing relationships between SSL-HSIC , InfoNCE , and metric learning . • Our method achieves top-1 accuracy of 74.8 % and top-5 accuracy of 92.2 % with linear evaluations ( see Figure 1 for a comparison with other methods ) , top-1 accuracy of 80.2 % and Top-5 accuracy of 94.7 % with fine-tuning , and competitive performance on a diverse set of downstream tasks . 2 Background . 2.1 Self-supervised learning . Recent developments in self-supervised learning , such as contrastive learning , try to ensure that features of two random views of an image are more associated with each other than with random views of other images . Typically , this is done through some variant of a classification loss , with one “ positive ” pair and many “ negatives. ” Other methods can learn solely from “ positive ” pairs , however . There have been many variations of this general framework in the past few years . Van den Oord et al . [ 15 ] first formulated the InfoNCE loss , which estimates a lower bound of the mutual information between the feature and the context . SimCLR [ 19 , 34 ] carefully investigates the contribution of different data augmentations , and scales up the training batch size to include more negative examples . MoCo [ 17 ] increases the number of negative examples by using a memory bank . BYOL [ 25 ] learns solely on positive image pairs , training so that representations of one view match that of the other under a moving average of the featurizer . Instead of the moving average , SimSiam [ 24 ] suggests a stop-gradient on one of the encoders is enough to prevent BYOL from finding trivial solutions . SwAV [ 18 ] clusters the representation online , and uses distance from the cluster centers rather than computing pairwise distances of the data . Barlow Twins [ 22 ] uses an objective related to the cross-correlation matrix of the two views , motivated by redundancy reduction . It is perhaps the most related to our work in the literature ( and their covariance matrix can be connected to HSIC [ 35 ] ) , but our method measures dependency more directly . While Barlow Twins decorrelates components of final representations , we maximize the dependence between the image ’ s abstract identity and its transformations . On the theory side , InfoNCE is proposed as a variational bound on Mutual Information between the representation of two views of the same image [ 15 , 32 ] . Tschannen et al . [ 30 ] observe that InfoNCE performance can not be explained solely by the properties of the mutual information , however , but is influenced more by other factors , such as the formulation of the estimator and the architecture of the feature extractor . Essentially , representations with the same MI can have drastically different representational qualities . To see this , consider a problem with two inputs , A andB ( Figure 3 , green and purple ) , and a one-dimensional featurizer , parameterized by the integer M , which maps A to Uniform ( { 0 , 2 , . . . , 2M } ) and B to Uniform ( { 1 , 3 , . . . , 2M + 1 } ) . When M = 0 , the inputs are encoded into linearly separable features A = 0 and B = 1 ( Figure 3 , bottom ) . Otherwise when M > 0 , they are interspersed like ABABABAB – a representation which is much harder to work with for downstream learners . Nevertheless , the mutual information between the features of any two augmentations of the same input ( a positive pair ) is independent of M , that is H [ Z1 ] −H [ Z1|Z2 ] = log 2 for any M . Note that InfoNCE would strongly prefer M = 0 , indeed behaving very differently from MI . Later theories suggest that contrastive losses balance alignment of individual features and uniformity of the feature distribution [ 23 ] , or in general alignment and some loss-defined distribution [ 36 ] . We propose to interpret the contrastive loss through the lens of statistical dependence , and relate it to metric learning , which naturally leads to alignment and uniformity . 2.2 Hilbert-Schmidt Independence Criterion ( HSIC ) . The Hilbert-Schmidt Independence Criterion ( HSIC ) [ 26 ] is a kernel-based measure of dependence between probability distributions . Like mutual information , for a wide range of kernels HSIC ( X , Y ) = 0 if and only if X and Y are independent [ 37 ] , and large values of the measure correspond to “ more dependence. ” Unlike mutual information , HSIC incorporates a notion of ge- ometry ( via the kernel choice ) , and is both statistically and computationally easy to estimate . It has been used in a variety of applications , particularly for independence testing [ 38 ] , but it has also been maximized in applications such as feature selection [ 39 ] , clustering [ 40 , 41 ] , active learning [ 42 ] , and as a classification loss called HSIC Bottleneck [ 27 , 28 ] ( similar ideas were expressed in [ 43 , 44 ] ) . HSIC measures the dependence between two random variables by first taking a nonlinear feature transformation of each , say φ : X → F and ψ : Y → G ( with F and G reproducing kernel Hilbert spaces , RKHSes1 ) , and then evaluating the norm of the cross-covariance between those features : HSIC ( X , Y ) = ∥∥E [ φ ( X ) ψ ( Y ) > ] − E [ φ ( X ) ] E [ ψ ( Y ) ] > ∥∥2 HS . ( 1 ) Here ‖·‖HS is the Hilbert-Schmidt norm , which in finite dimensions is the usual Frobenius norm . HSIC measures the scale of the correlation in these nonlinear features , which allows it to identify nonlinear dependencies between X and Y with appropriate features φ and ψ . Inner products in an RKHS are by definition kernel functions : k ( x , x′ ) = 〈φ ( x ) , φ ( x′ ) 〉F and l ( y , y′ ) = 〈ψ ( y ) , ψ ( y′ ) 〉G . Let ( X ′ , Y ′ ) , ( X ′′ , Y ′′ ) be independent copies of ( X , Y ) ; this gives HSIC ( X , Y ) = E [ k ( X , X ′ ) l ( Y , Y ′ ) ] − 2E [ k ( X , X ′ ) l ( Y , Y ′′ ) ] + E [ k ( X , X ′ ) ] E [ l ( Y , Y ′ ) ] . ( 2 ) HSIC is also straightforward to estimate : given i.i.d . samples { ( x1 , y1 ) , . . . , ( xN , yN ) } drawn i.i.d . from the joint distribution of ( X , Y ) , Gretton et al . [ 26 ] propose an estimator ĤSIC ( X , Y ) = 1 ( N − 1 ) 2 Tr ( KHLH ) , ( 3 ) where Kij = k ( xi , xj ) and Lij = l ( yi , yj ) are the kernel matrices , and H = I − 1N 11 > is called the centering matrix . This estimator has an O ( 1/N ) bias , which is not a concern for our uses ; however , an unbiased estimator with the same computational cost is available [ 39 ] .

Summary:   This paper proposed a new loss function for self-supervised learning using the Hilbert-Schmidt Independence Criterion (HSIC) to measure the kernel space dependence between the learned latent representation of images and their augmented versions. The main contribution is that authors demonstrated the connections between the proposed framework and InfoNCE, metric learning and Maximum Mean Discrepancy, respectively. An algorithm with random fourier feature kernel estimation is proposed and the method is shown to have comparable performance with state-of-the-art self-supervised learning (SSL) paradigms. The sample complexity and bias analysis for kernel estimation are provided.