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: The paper proposes to use the HSIC metric in order to impose view feature dependency on the image identity during self-supervised learning (SSL). The authors made non-trivial efforts in both deriving the SSL-HSIC approach and theoretically linking SSL-HSIC with the existing mainstream SSL approaches. Empirical results show that the method offers competitive results.