$f$-Mutual Information Contrastive Learning

1 INTRODUCTION . Contrastive learning has attracted a surge of attention recently due to its success in learning informative representation for image recognition , natural language understanding , and reinforcement learning ( Chen et al. , 2020 ; He et al. , 2020 ; Logeswaran & Lee , 2018 ; Srinivas et al. , 2020 ) . Such learning paradigm is fully unsupervised by encouraging the contrastiveness between similar and dissimilar sample pairs . Specifically , the feature embeddings of similar sample pairs are expected to be close while those of dissimilar sample pairs are expected to be far apart . To attain this goal , a softmax cross-entropy loss , a.k.a . InfoNCE , has been widely used ( Wu et al. , 2018 ; van den Oord et al. , 2018 ; Chen et al. , 2020 ; Hénaff et al. , 2020 ; He et al. , 2020 ) , which aims to maximize the probability of picking a similar sample pair among a batch of sample pairs . InfoNCE can be interpreted as a lower bound of the mutual information ( MI ) between two views of data samples ( van den Oord et al. , 2018 ; Bachman et al. , 2019 ; Tian et al. , 2020a ; Tschannen et al. , 2020 ) . This explanation is consistent with the well-known “ InfoMax principle ” ( Linsker , 1988 ) . Nevertheless , it has been shown that maximizing a tighter bound on the MI can result in worse representations ( Tschannen et al. , 2020 ) ; and reducing the MI between views while only keeping task-relevant information can improve the downstream performance ( Tian et al. , 2020b ) . These observations suggest that maximizing the MI may be insufficient in contrastive learning and thus a better objective design is required . To attain the aforementioned goal , we propose a novel contrastive learning framework , coined as f -MICL . In a nutshell , leveraging the fact that MI can be formulated as the Kullback–Leibler ( KL ) divergence between the joint distribution and the product of the marginal distributions , we replace the KL divergence with the general f -divergence family ( Ali & Silvey , 1966 ; Csiszár , 1967 ) . Doing so , we obtain a generalization of mutual information , called f -mutual information ( f -MI , Csiszár , 1967 ) . Notably , by maximizing a lower bound of f -mutual information we naturally decompose the objective function into two terms , which correspond to the properties of the alignment and the uniformity . Such characterization has been revealed in Wang & Isola ( 2020 , Theorem 1 ) for the InfoNCE loss . Compared with Wang & Isola ( 2020 ) , our result applies to a wide range of the f - divergence family such as KL , Jensen–Shannon ( JS ) , Pearson χ2 and Vincze–Le Cam . and it does not rely on the limit of an infinite number of dissimilar samples . This allows us to explore the space of f -MI and improve the performance of InfoNCE-based contrastive learning . The similarity function is crucial for the evaluation of the contrastiveness of similar and dissimilar sample pairs . Commonly used similarity functions include the cosine similarity ( Chen et al. , 2020 ; He et al. , 2020 ) , the bilinear functions ( van den Oord et al. , 2018 ; Tian et al. , 2020a ; Hénaff et al. , 2020 ) , and the neural network based scores ( Hjelm et al. , 2018 ) . While most aforementioned similarity functions for contrastive learning are heuristic and pre-designed , in this work , we provide a principled way to design the similarity function . By assuming that the joint feature distribution of two similar samples is proportional to a Gaussian kernel , we derive an optimal similarity function for practical use , which resembles the well-known radial basis functions ( Powell , 1987 ) . Figure 1 gives a high-level summary of our f -MICL framework . Given a batch of samples ( e.g. , images ) we generate positive pairs via data augmentation and negative pairs using other augmented samples in the current batch . With the optimization of our f -mutual information objective the positive pairs are aligned with each other and the data representations are uniformly distributed . Our contributions can be summarized as follows : • We propose a novel framework for contrastive learning ( called f -MICL ) by encouraging the contrastiveness of positive and negative pairs with a general f -divergence family . • With an assumption on the joint feature distribution we provide an optimal design for the similarity function with Gaussian kernels , which shows interesting connection between contrastive learning and kernel methods . • We characterize the properties of alignment and uniformity for f -MICL and provide guidance on choosing proper f -divergences . Our f -MICL objective can be estimated with finite samples and we give the corresponding error bound . • Experimentally , our framework is better or on par with popular baselines , by simply replacing , e.g. , the InfoNCE loss function with our theoretically grounded f -MICL objectives . Notations We assume a dominating measure λ ( e.g . Lebesgue ) is given and all other probability measures are represented as some density w.r.t . λ . We denote Df ( p‖q ) as the f -divergence between two densities functions p and q . Given the joint density p ( x , y ) , we denote p ( x ) = ∫ p ( x , y ) dλ ( y ) and p ( y ) = ∫ p ( x , y ) dλ ( x ) as the marginals . We use supp ( · ) to denote the support of a distribution , and f∗ to denote the Fenchel conjugate of a convex function f . Every norm presented is Euclidean . 2 PRELIMINARIES . We provide some preliminaries for our framework in the context of contrastive learning . Contrastive learning is a popular unsupervised method for learning data representations . In contrastive learning , we expect that similar sample pairs to be close to each other in the embedding space while uncorrelated pairs to be far away . Denote ppos as the distribution of positive pairs , i.e. , two samples that have similar representations . Assume that this distribution is symmetric w.r.t . the two random variables , then the resultant two marginals both follow the data distribution pdata ( Wang & Isola , 2020 ) . In this work we propose the f -mutual information framework for contrastive learning . First recall the f -mutual information ( f -MI ) between a pair of random variables X and Y : Definition 1 ( f -mutual information , Csiszár 1967 ) . Consider a pair of random variables ( X , Y ) with density function p ( x , y ) . The f -mutual information If between X and Y is defined as If ( X ; Y ) : = Df ( p ( x , y ) ‖p ( x ) p ( y ) ) = ∫ f ( p ( x , y ) p ( x ) p ( y ) ) p ( x ) p ( y ) · dλ ( x , y ) , ( 1 ) where f : R+ → R is ( closed ) convex with f ( 1 ) = 0 , and recall that p ( x ) and p ( y ) are the marginal densities of p ( x , y ) , whereas λ is a dominating measure ( e.g . Lebesgue ) . Common choices of f can be found in Table 1 and Table 4 ( Appendix A ) , which we will discuss in more details in §3 . It is well-known that f -mutual information is non-negative and symmetric , and provided that f is strictly convex , If ( X ; Y ) = 0 iff X and Y are independent ( Ali & Silvey , 1966 ) . When X and Y are of high dimension , it is quite challenging to estimate the f -divergence directly . Instead , Nguyen et al . ( 2010 ) derived a variational method by maximizing the dual problem : If ( X ; Y ) ≥ sup T∈T if ( X ; Y ) : = E ( x , y ) ∼ppos [ T ( x , y ) ] − E ( x , y ) ∼pdata⊗pdata [ f ∗ ( T ( x , y ) ) ] , ( 2 ) where f∗ ( t ) : = supx∈R+ ( xt − f ( x ) ) is the ( monotone ) Fenchel conjugate of f , and is always monotonically increasing . Here T is a class of functions T : supp ( pdata ) × supp ( pdata ) → dom f∗ . Nguyen et al . ( 2010 ) showed that the bound in eq . ( 2 ) is tight if there exists T ∗ ∈ T such that T ∗ ( x , y ) = f ′ ( p ( x , y ) / ( p ( x ) p ( y ) ) ) , for any ( x , y ) ∈ supp ( pdata ) × supp ( pdata ) , ( 3 ) and in particular , if T comprises of all ( measurable ) functions . 3 DESIGN OF f -MICL Based on the f -MI introduced in §2 , we propose a novel framework for contrastive learning . Furthermore , we characterize the properties of alignment and uniformity theoretically for general f - divergences . Following Chen et al . ( 2020 ) , we design the strcture of function T as follows : T ( x , y ) : = k ( g ( x ) , g ( y ) ) , where ‖g ( x ) ‖ = 1 for any sample x . ( 4 ) The function g produces a d-dimensional normalized feature encoding on the hypersphere Sd−1 and k is a similarity function that measures the similarity between two embeddings g ( x ) and g ( y ) . With the above interpretation , we can rewrite our objective of f -mutual information , equation 2 , as : sup g∈G , k∈K if ( X ; Y ) : = E ( x , y ) ∼ppos [ k ( g ( x ) , g ( y ) ) ] − E ( x , y ) ∼pdata⊗pdata [ f ∗ ( k ( g ( x ) , g ( y ) ) ) ] , ( 5 ) where G andK are the function classes of the feature encoder g and the similarity function k. We can treat the first term as the similarity score between positive pairs in the feature space , and the second term as the similarity score between two random samples , a.k.a . negative pairs , in the feature space . As f∗ is increasing , maximizing f -MI is equivalent to simultaneously maximizing the similarity between positive pairs and minimizing the similarity between negative pairs . 3.1 OPTIMIZED SIMILARITY FUNCTION AND IMPLEMENTATION . Let us now study how to search for the optimal similarity function k. To our best knowledge , there has been no theoretical study on the choice of similarity functions . Most existing contrastive learning methods ( e.g . Chen et al. , 2020 ; Tsai et al. , 2020 ; He et al. , 2020 ) adopt a pre-designed similarity function , such as the cosine similarity . For the ease of notation , from now on we define xg : = g ( x ) and yg : = g ( y ) . Suppose ( x , y ) ∼ ppos , then we denote pgpos as the distribution of ( xg , yg ) , and p g data as the marginal feature distribution of xg or yg . The corresponding density functions are written as pg ( x g ) , pg ( y g ) and pg ( xg , yg ) . We remind the reader of the following result : Lemma 2 ( e.g. , Nguyen et al . 2010 , Lemma 1 ) . Suppose f is differentiable , and the encoder function g is fixed . The similarity function k∗ ( xg , yg ) = f ′ ( pg ( x g , yg ) pg ( xg ) pg ( yg ) ) ( 6 ) maximizes if ( X ; Y ) in eq . ( 5 ) as long as it is contained in the function class K. Equation 6 provides an optimal similarity function , which nevertheless depends on the density functions . Comparing equation 6 with equation 5 , we realize that the optimal k∗ in fact gives the f -MI on the feature space , If ( g ( X ) , g ( Y ) ) , which is a low bound of the original f -MI , If ( X ; Y ) . To use k∗ practically we make the following assumption on the joint density : Assumption 3 . The joint feature distribution is proportional to a radial basis function ( RBF ) , i.e. , pg ( x g , yg ) ∝ ϕ ( ‖xg − yg‖2 ) for a real-valued function ϕ . Radial basis functions are widely used in kernel methods ( Powell , 1987 ; Murphy , 2012 ) , and the Gaussian kernel is perhaps the most well-known RBF . Throughout this work we mainly consider ϕ as a Gaussian kernel : ϕ ( ‖xg − yg‖2 ) = Gσ ( ‖xg − yg‖2 ) : = µ exp ( −‖x g − yg‖2 2σ2 ) , ( 7 ) with µ a constant left to be determined . Fixing yg , then pg ( · , yg ) is known as the von Mises–Fisher distribution ( von Mises , 1918 ; Fisher , 1953 ; Bingham & Mardia , 1975 ) , since xg and yg are unit vectors . With Assumption 3 on the joint density , the resultant marginal feature distribution pgdata is uniform on the hypersphere Sd−1 , where d is the dimension of the feature space ( see Prop . 8 in App . B ) . Additionally , for positive pairs the distance in the feature space , ‖xg − yg‖ , is more likely to be small . If the variance σ2 → 0 , then the Gaussian kernel becomes the Dirac delta distribution , δxg=yg . This requires that the two features xg and yg to be the same , which is desirable . In general , the radial basis function ϕ should be decreasing since a positive pair should be more likely to be adjacent in the feature space . For example , with ϕ ( t ) = 1− t/2 we obtain the cosine similarity . Based on Assumption 3 we propose the following similarity function between pairs of features : Theorem 4 ( Gaussian similarity ) . Under Assumption 3 with Gaussian kernels and the same settings as Lemma 2 , the optimal similarity function k∗ satisfies that for any xg , yg ∈ Sd−1 : k∗ ( xg , yg ) = f ′ ( CGσ ( ‖xg − yg‖2 ) ) , ( 8 ) where d is the feature dimension and C is an absolute constant . For simplicity we will rewrite k∗ ( xg , yg ) = f ′ ◦ Gσ ( ‖xg − yg‖2 ) by absorbing the constant C into Gσ , since we have left some flexibility in equation 7 . Although Assumption 3 with Gaussian kernels may not always reflect the real feature distribution , we can still use the similarity function in equation 8 , even if it might not be optimal . In our experiments in §5 , the Gaussian similarity equation 8 consistently outperforms the default cosine similarity in contrastive learning . Algorithm 1 : f -mutual information contrastive learning ( f -MICL ) Input : batch size N , function f , weighting parameter α , constant µ ( in Gσ ) , variance σ2 , optimizer 1 for each batch { zi } Ni=1 do 2 forall k ∈ [ 1 , N ] do 3 randomly sample two augmentation functions t1 , t2 4 yk ← t1 ( zk ) , xk ← t2 ( zk ) 5 compute if = 1N ∑N i=1 [ f ′ ◦Gσ ( ‖xgi − y g i ‖2 ) ] − αN ( N−1 ) ∑ i6=j f ∗ ◦ f ′ ◦Gσ ( ‖xgi − x g j‖ ) 6 update g by taking a step to maximizing if using the optimizer Bringing the optimal k∗ in equation 8 into our objective equation 5 we have the following objective : sup g∈G E ( x , y ) ∼ppos [ f ′ ◦Gσ ( ‖xg − yg‖2 ) ] − E ( x , y ) ∼pdata⊗pdata [ f∗ ◦ f ′ ◦Gσ ( ‖xg − yg‖2 ) ] , ( 9 ) where Gσ is defined in equation 7 . With a similar sampling method of positive and negative pairs in Chen et al . ( 2020 ) , given a batch of N samples we can estimate the objective in equation 9 as : îf ( X ; Y ) = 1 N N∑ i=1 f ′ ◦Gσ ( ‖xgi − y g i ‖ 2 ) − 1 N ( N−1 ) ∑ i 6=j f∗ ◦ f ′ ◦Gσ ( ‖xgi − x g j‖ 2 ) , ( 10 ) where xi and yi are two different kinds of data augmentation of the i-th sample , and xi and xj are different samples of the same kind of data augmentation . With the objective in equation 10 we propose our algorithm for contrastive learning in Algorithm 1 . Note that we treat µ and σ2 in our Gaussian kernel equation 7 as hyperparameters . To balance the two terms in our objective , we additionally include a weighting parameter α in front of the second term . We can prove that rescaling the second term with the factor α is equivalent to changing the function f to another convex function fα ( see Prop . 7 in Appendix A ) .

Summary: This paper makes two contributions to contrastive representation learning: 1. Using the more general f-mutual information rather than using Shannon mutual information for contrastive learning 2. Experimental results to compare the possible options given the new design freedom.