Year: 2021

Abstract:
Self-supervised contrastive learning is an emerging field due to its power in providing good data representations. Such learning paradigm widely adopts the InfoNCE loss, which is closely connected with maximizing the mutual information. In this work, we propose the f -Mutual Information Contrastive Learning framework (f -MICL), which directly maximizes the f -divergence-based generalization of mutual information. We theoretically prove that, with mild assumptions, our f -MICL naturally attains the alignment for positive pairs and the uniformity for data representations, the two main factors for the success of contrastive learning. We further provide theoretical guidance on designing the similarity function and choosing the effective f -divergences for f -MICL. Using several benchmark tasks from both vision and natural language, we empirically verify that our novel method outperforms or performs on par with state-of-the-art strategies.

heading: 1 INTRODUCTION
text: 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.

heading: 2 PRELIMINARIES
text: 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.

heading: 3.1 OPTIMIZED SIMILARITY FUNCTION AND IMPLEMENTATION
text: 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).

heading: 3.2 ALIGNMENT AND UNIFORMITY
text: Notably, if we choose the f -divergence to be the KL divergence, the objective in equation 9 becomes:
sup g∈G − 1 2σ2 E(x,y)∼ppos
[ ‖xg − yg‖2 ] − µE(x,y)∼pdata⊗pdata [ exp ( −‖x g − yg‖2
2σ2
)] , (11)
which retrieves the objective of the alignment and uniformity in Wang & Isola (2020). Specifically, in equation 11 the first expectation is the same as Lalign, and the second expectation is the same as Luniform in Wang & Isola (2020) (up to the logarithmic transformation). Therefore, under our Assumption 3, we naturally decompose the objective for the KL into two: the first term leads to the alignment of positive pairs while the second term results in the uniformity of data representations. Compared to Theorem 1 of Wang & Isola (2020), our conclusion is based on the lower bound of f -MI instead of the InfoNCE loss, and does not rely on the limit of infinite negative samples. More importantly, besides the KL divergence we can extend our conclusion to various f -divergences.
Alignment Consider the objective in eq. (9). For the first term, since f is convex, the derivative f ′ is monotonic increasing. Therefore, in the ideal case, maximizing the first term would yield xg = yg for all (x, y) ∼ ppos, i.e., similar sample pairs should have aligned representations. Uniformity In Wang & Isola (2020) it is proved that for the KL divergence minimizing the second expectation in equation 11 gives the uniform distribution, if we have infinite samples. Here we discuss the general f -divergences with finite samples by considering equation 10, i.e., our empirical estimation of the objective. We list some common choices of f -divergences in Table 1, and a more detailed version can be found in Table 4 in Appendix A.1. We have the following theorem: Theorem 5 (uniformity). Suppose that the batch size N satisfies 2 ≤ N ≤ d + 1, with d the dimension of the feature space. If the real function
h(t) = f∗ ◦ f ′ ◦Gσ(t) is strictly convex on [0, 4], (12)
then all minimizers of the second term of equation 10, i.e., ∑ i 6=j f ∗ ◦ f ′ ◦Gσ(‖xgi − x g j‖2), satisfy the following condition: the feature representations of all samples are distributed uniformly on the unit hypersphere Sd−1. All f -divergences in Table 1 satisfy equation 12.
Here by “distributed uniformly” we mean that the feature vectors form a regular simplex (see Figure 2), and thus the distances between all sample pairs are the same. It reflects our intuition that the feature embeddings are evenly distributed. Although minimizing the negative term gives uniformity, the positive term is also needed for aligning similar pairs, as we observe in our experiments in §5. Therefore, there is a tradeoff between alignment and uniformity. In Theorem 5 we assumedN ≤ d+1, with d as the feature dimension. This assumption is always satisfied in our experiments in §5. For instance, for CIFAR-10 experiments we choose N = d = 512.
Not all f -divergences lead to the property of uniformity. In Appendix A.1 we also listed some f -divergences that do not satisfy equation 12 and Theorem 5, such as the Reversed Kullback–Leibler (RKL) and the Neynman χ2 divergences. Experimentally, we found that these divergences generally result in feature collapse (i.e., all feature vectors are the same) and thus poor performance in downstream applications. Compared to the existing literature (Nowozin et al., 2016; Acuna et al., 2021) where no suggestion is given for the choice of f -divergences theoretically, with Theorem 5 we provide guidance for choosing desirable f -divergences.
4 ESTIMATION OF f -MUTUAL INFORMATION
In §3.1 we have provided the empirical estimation of our objective in equation 10. However, it remains a question whether our estimation of f -mutual information is consistent. In this part we give an upper bound for the estimation error. Recall that our f -MICL objective is if (X;Y ) (equation 2) and its empirical estimation is îf (X;Y ) (equation 10), with T (x, y) = f ′ ◦Gσ(‖xg − yg‖2). Theorem 6 (estimation error). Suppose that the function T is taken from a function class T and define Tx as the function class of T (x, ·) given some x ∈ supp(pdata). Denote RPN to be the Rademacher complexity w.r.t. the distribution P with N i.i.d. drawn samples. Then for any T ∈ T , the estimation error |if (X;Y )− îf (X;Y )| is upper bounded with probability at least 1− δ:
2R ppos N (T ) + 2µ [ Ex∼pdataR pdata N (Tx) + 1
N N∑ i=1 RpdataN−1(Txj )
] + (rT + 2rf ) √ log 6/δ
2(N − 1) , (13)
with the constants rT = f ′(µ)− f ′(µe−2/σ 2 ), rf = f ∗ ◦ f ′(µ)− f∗ ◦ f ′(µe−2/σ2).
Here the constant µ is from our Gaussian kernel in equation 7. Rademacher complexity evaluates the richness of a class of real-valued functions regarding a probability distribution, and its formal definition can be found in Koltchinskii (2001) (see also Definition 11 in Appendix B).
Note that the function class T depends on the class of the feature encoder g and the f -divergence. The estimation error eq. (13) is composed of three parts: (1) the Rademacher complexity of the function class T . In general, if T is richer then its Rademacher complexity is also larger. (2) the expected Rademacher complexity of the one-side function class Tx and its empirical estimation; and (3) an error term that decreases with the number of samples. Since the feature embeddings are usually neural networks, we can use existing theory (Bartlett et al., 2019) to give more detailed bounds for the Rademacher complexities of T . Specifically, if the Vapnik–Chervonenkis (VC) dimension of T is finite, then our estimation error in eq. (13) goes to zero as N →∞ (Mohri et al., 2018). Approximation and estimation tradeoff In order to minimize the estimation error in equation 13, we should choose a simpler function class T to reduce the Rademacher complexities. However, T should also be rich enough so that equation 3 can be satisfied, since our objective if (X;Y ) should approximate the f -mutual information If (X;Y ) if we choose the optimal T . Therefore, there is a natural tradeoff between approximation and estimation errors when we change the complexity of T . It is worth mentioning that our conclusion in Theorem 6 is theoretically non-trivial since our sample pairs are non-i.i.d.: although the individual samples are assumed to be i.i.d., the negative pairs are not independently drawn (e.g., (x1, x2) and (x1, x3)), which makes the derivation challenging.

heading: 5 EXPERIMENTS
text: We compare our framework with various frameworks on several popular vision and language datasets, and show the wide applicability of our method. In particular, our f -MICL gives stateof-the-art performance compared to popular baseline algorithms, such as SimCLR (van den Oord et al., 2018; Chen et al., 2020), MoCo (He et al., 2020), Uniformity (Wang & Isola, 2020), and RPC (Tsai et al., 2020). Additional experimental results can be seen in Appendix C.
Specifically, our results confirm the following:
• Our f -MICL encourages alignment between positive pairs, and encourages dissimilar sample pairs to be equally far apart and thus leads to uniformity.
• By replacing the cosine similarity with the Gaussian kernels, the performance is consistently better across a variety of f -divergences in our f -MICL framework.

heading: 5.1 EXPERIMENTAL SETTINGS
text: For vision tasks, we use SGD with momentum as our optimizer, and apply the cosine learning rate schedule (Loshchilov & Hutter, 2017). Regarding the neural network architecture, we employ ResNets (He et al., 2020) as the feature encoders. For the language dataset, we use the BERT models (Devlin et al., 2019) and Adam with weight decay as the optimizer (Loshchilov & Hutter, 2018). To achieve fair comparison, we keep the network architecture and optimization the same, while only changing the objective accordingly in our comparison. See Appendix C for more detailed settings.

heading: 5.2 COMPARISON WITH BENCHMARKS
text: We compare with several state-of-the-art benchmarks in Table 2. Note that SimCLR and RPC use the cosine similarity while we use the proposed Gaussian similarity.
Vision task Our vision datasets include CIFAR-10, CIFAR-100 (Krizhevsky et al., 2009), STL-10 (Coates et al., 2011), TinyImageNet (Chrabaszcz et al., 2017) and ImageNet (Deng et al., 2009) for image classification. After learning a feature embedding, we evaluate the quality of representation using the test classification accuracies via a linear classifier. We observe from Table 2 that our proposed f -MICL consistently outperforms the benchmarks across all datasets. Specifically, we find that the JS divergence is superior in general, especially on the large datasets.
Language task To show the efficacy of our f -MICL framework, we also conduct experiments on a natural language dataset: English Wikipedia (Gao et al., 2021). We follow the experimental setting of Gao et al. (2021), which applies SimCLR with BERT models (Devlin et al., 2019; Liu et al., 2019). The application task is semantic textual similarity (STS, Agirre et al. 2013) and we report the averaged the Spearman’s correlation in Table 2 for comparison. We can see that our f -MICL performs better or on par with the benchmarks, especially for the KL divergence.

heading: 5.3 ABLATION STUDY
text: To develop a thorough understanding of f -MICL we also perform the ablation study regarding the choice of the similarity function, the design of the f -divergences, and the batch size. More ablation study can be found in Appendix C.
Cosine similarity vs. Gaussian In Table 3 We compare the cosine and Gaussian similarities for different f -divergences on CIFAR-10. It can be seen that under our f -MICL framework the Gaussian similarity consistently outperforms the cosine similarity for various f -divergences. This agrees with our Theorem 2 and equation 8, and also implies the validity of Assumption 3.
Non-satisfying f -divergences We have shown analytically that not all f -divergences are proper in our f -MICL framework (see the proof of Theorem 5, Appendix B). For example, in our experiments on CIFAR-10 and CIFAR-100, the RKL and Neyman χ2 do not perform well: the test classification accuracies on these two datasets are 10.00% and 1.00%, respectively, which means that the classifier simply outputs the random guess. Moreover, we observe from Figure 3 that for the non-satisfying f -divergences such as the RKL, the features collapse to a constant. These observations emphasize the importance of choosing proper f -divergences and confirms our conclusion in Theorem 5.
Sensitivity to batch size We study the sensitivity to the batch size of our f -MICL framework on CIFAR-10. On the right panel of Figure 3, we evaluate the classification accuracy by varying the batch size for different f -divergences and SimCLR. We can see that for all different batch sizes and with proper choice of f -divergences, our performance is always better than SimCLR. In other words, we require less negative samples to achieve the same performance.

heading: 5.4 UNIFORMITY TEST
text: To check the uniformity of feature vectors (Theorem 5) we plot the pairwise distance ‖xgi − x g j‖ of the feature representations within the same batch on CIFAR-10 and CIFAR-100. We compute the distances between the normalized features of every pair from a random batch, and then sort the pairs with the increasing order. From Figure 3 we can see that f -MICL gives nearly uniform distances for dissimilar pairs (orange regions) on both datasets with various proper f -divergences. In contrast, a random initialized model gives a less uniform distribution for dissimilar pairs. Besides, for f -MICL we observe small pairwise distances for similar pairs (green regions). On CIFAR-100 there are less similar pairs compared to CIFAR-10 as there are more classes.

heading: 6 RELATED WORKS
text: Contrastive learning Self-supervised contrastive learning learns representations by encouraging the contrastiveness between positive pairs and negative pairs. Recently it has been shown analytically that improving the contrastiveness can benefit the downstream applications (Saunshi et al., 2019; Tosh et al., 2021). For popular contrastive learning methods such as Contrastive Predictive Coding (CPC) (van den Oord et al., 2018), SimCLR (Chen et al., 2020), and MoCo (He et al., 2020), their loss functions can be interpreted as a lower bound of mutual information, which is essentially the KL divergence between the joint distribution ppos and the product of margin distributions pdata ⊗ pdata. Besides the KL divergence, other statistical divergences or distances have been individually studied under the context of contrastive learning, e.g., the Wasserstein distance (Ozair et al., 2019), Pearson χ2 divergence (Tsai et al., 2021), and Jensen–Shannon divergence (Hjelm et al., 2018). Our paper differs from Hjelm et al. (2018); Tsai et al. (2021) in the following ways:
• Instead of treating f -divergences as estimators of mutual information, we directly consider f -mutual information, which builds solid foundation for our objective.
• Instead of using the default cosine similarity, we challenge this routine and construct the optimal design with the Gaussian kernel, which consistently performs better (Table 3).
• We show that optimizing the f -MICL objective of the contrastive learning yields alignment and uniformity for proper f -divergences with finite samples.
Uniformity Our work is inspired by Wang & Isola (2020), which shows the alignment and uniformity in contrastive learning. However, Wang & Isola (2020) focuses on the InfoNCE loss and the loss function Luniform they consider is somewhat ad hoc. In contrast, our f -MICL objective is directly built from the f -mutual information theory. Another difference is that Wang & Isola (2020) proves uniformity in the limit of infinite samples, while our Theorem 5 only needs finite samples.
f -divergences have been widely used in, e.g., generative models (Nowozin et al., 2016), variational inference (Wan et al., 2020), and domain adaptation (Acuna et al., 2021), for measuring the discrepancy of two distributions. Compared to the existing work such as f -GAN (Nowozin et al., 2016) which evaluates the similarity of two marginal distributions (i.e., the real and the modeled data distributions), our f -MICL objective measures the similarity between the joint distribution and the product of data distributions. Furthermore, while f -GAN aims to learn a generative model and thus minimizes the f -divergence, our f -MICL intends to encourage the contrastiveness and maximizes the f -divergence. Finally, we provide a theoretical criterion for choosing proper f -divergences.
Metric learning Our work is closely related to metric learning (Kaya & Bilge, 2019; Suárez-Dı́az et al., 2018), which aims to learn a distance metric bringing similar objects closer and distancing dissimilar objects further. Distance metrics that are designed heuristically or independently of the problem may not lead to a satisfactory performance in applications. In contrastive learning, a predefined similarity metric, e.g., the cosine similarity (Chen et al., 2020; He et al., 2020) or a bilinear function (van den Oord et al., 2018; Tian et al., 2020a; Hénaff et al., 2020) is commonly employed to measure the similarity of sample pairs. Comparably, we have derived an optimal similarity function assuming that the joint positive distribution is proportional to a Gaussian kernel.

heading: 7 CONCLUSION AND FUTURE WORK
text: In this work we proposed f -MICL, a novel contrastive learning framework with a generalization of mutual information, called f -mutual information (f -MI). Our objective corresponds to the variational bound of f -divergences. By making the assumption that the joint feature distribution is proportional to a Gaussian kernel, we naturally characterized the properties of the alignment and uniformity. Our f -MI based objective can be well estimated with finite samples. Experimentally, we showed the efficacy of f -MICL across a wide array of datasets and the advantage of using the Gaussian similarity. Our results imply the following: (1) even though mutual information is widely used in contrastive learning, generalization to f -MI can bring us better performance; (2) the cosine similarity, the de facto option in contrastive learning, can be replaced with more effective similarity functions such as the Gaussian similarity. As our work shows interesting connection between contrastive learning and kernel methods, it would be promising to explore the usage of more RBF kernels in contrastive learning.
REPRODUCIBILITY STATEMENT
Our novel f -MICL objectives are implemented in anonymous downloadable source code as in the supplementary. For theoretical results, clear explanations of any assumptions and complete proofs are included in Appendix A and Appendix B. For detailed experimental settings and complete experimental results, the complete descriptions can be found in Appendix C.

heading: A ADDITIONAL THEORETICAL RESULTS
text: In this appendix we provide additional theoretical results, including additional f -divergences and the theory for weighting parameters.
A.1 ADDITIONAL f -DIVERGENCES
We expand Table 1 and give more examples of f -divergences in Table 4. As we will see in the proof of Theorem 5, Table 1 gives a special class of f -divergences that guarantees uniformity. A detailed description of f -divergences can be found in e.g. Sason & Verdú (2016).
A.2 WEIGHTING PARAMETERS
In Algorithm 1 we added a weighting parameter α to balance the alignment and uniformity. We prove that even after adding this parameter we are still maximizing the f -mutual information, although with respect to a different f .
Proposition 7 (weighting parameter). Given α > 0 and a closed convex function f : R+ → R such that f(1) = 0, define fα : α dom f → R with fα(x) = αf(x/α) − αf(1/α) for any x ∈ dom f . Then Ifα is still a valid f -mutual information (see Definition 1). Besides, by replacing f with fα in equation 9 we have the following optimization problem:
sup g∈G
E(x,y)∼ppos [ f ′ ( Gσ(‖xg−yg‖2)/α )] − αE(x,y)∼pdata⊗pdata [ f∗ ◦ f ′ ( Gσ(‖xg−yg‖2)/α )] ,
where Gσ(‖xg − yg‖2) = µ exp ( −‖x g−yg‖2 2σ2 ) is the Gaussian kernel.
Note that α dom f means the scalar multiplication of a set which is applied element-wisely. According to Definition 1, fα is also a valid f -divergence. This proposition tells us that rescaling the second term with factor α is equivalent to changing the function f to another convex function fα. The transformation from f to αf(x/α) is also known as right scalar multiplication (e.g. Chapter X, Urruty & Lemaréchal, 1993). Let us now move on to our proof:
Proof. By definition we know that fα is convex and closed with fα(1) = 0, and thus Ifα is a valid f -mutual information according to Definition 1. Moreover, we have f ′α(x) = f ′(x/α) for any
x ∈ α dom f and f∗α(t) = sup
x∈dom fα xt− fα(x)
= sup x∈αdom f
xt− αf(x/α) + αf(1/α)
= sup x/α∈dom f
(x/α) · (αt)− αf(x/α) + αf(1/α)
= α sup x/α∈dom f
((x/α) · t− f(x/α)) + αf(1/α)
= αf∗(t) + αf(1/α), (A.1) where in the last line we used the definition of f∗(t). Plugging f ′α and f ∗ α into equation 9 yields the desired result.
A.3 UNIFORM DISTRIBUTIONS
In the following we prove that under Assumption 3 the marginal feature distribution pgdata is uniform. Proposition 8. Under Assumption 3, the marginal feature distribution pgdata is uniform on the hypersphere Sd−1, with d the dimension of the feature space.
Proof. Under Assumption 3 we have:
pg(x g, yg) = C0 exp ( −‖x g − yg‖2
2σ2
) , (A.2)
where C0 is a normalizing constant. Integrating yg we have the marginal distribution for xg:
pg(x g) = ∫ Sd−1 C0 exp ( −‖x g − yg‖2 2σ2 ) dyg. (A.3)
It suffices to show that pg(x g 1) = pg(x g 2) for any x g 1, x g 2 ∈ Sd−1. SupposeQ is the orthogonal matrix such that: Qxg1 = x g 2. (A.4) Such a matrixQ always exists and constructingQ is not difficult. For example, assume that {xg1, x g 2} span a plane with an orthonormal basis e1, e2, and xg1 = cos θ1 · e1 + sin θ1 · e2, xg2 = cos θ2 · e1 + sin θ2 · e2, (A.5) then Q can take the following form:
Q = [e1 e2] [ cos(θ2 − θ1) − sin(θ2 − θ1) sin(θ2 − θ1) cos(θ2 − θ1) ] [ e>1 e>2 ] (A.6)
such that Q is orthogonal and satisfies Qxg1 = x g 2. Hence we have:
pg(x g 1) = ∫ Sd−1 C0 exp ( −‖x g 1 − yg‖2 2σ2 ) dyg
= ∫ Sd−1 C0 exp ( −‖Q >xg2 − yg‖2 2σ2 ) dyg
= ∫ Sd−1 C0 exp ( −‖Q >xg2 −Q>zg‖2 2σ2 ) d(Q>zg)
= ∫ Sd−1 C0 exp ( −‖x g 2 − zg‖2 2σ2 ) dzg = pg(x g 2), (A.7)
where in the second line we used equation A.4; in the third line we made the transformation yg = Q>zg with zg a unit vector; in the fourth line we used the fact that applying a orthogonal matrix does not change the norm and that the corresponding Jacobian determinant is one.
In fact, the proof above can be generalized from the Gaussian kernel to any radial basis functions, by replacing the Gaussian kernel with ϕ(‖xg − yg‖2), and repeating the same proof. Here ϕ can be any function such that the integral ∫ Sd−1 ϕ(‖x g − yg‖2)dyg is finite.

heading: B PROOFS
text: 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.
Proof. From Definition 1, we are computing the following supremum:
sup g,k
∫ ( pg(x g, yg)
pg(xg)pg(yg) k(xg, yg)− f∗ ◦ k(xg, yg)
) dpgdata ⊗ p g data. (B.1)
Suppose k is unconstrained and we fix g. The optimal solution should satisfy:
pg(x g, yg)
pg(xg)pg(yg) ∈ (∂f∗)(k∗(xg, yg)), (B.2)
almost surely for (x, y) ∼ pdata ⊗ pdata. From (3.11) of Rockafellar (1966) this is equivalent to:
k∗(xg, yg) ∈ ∂f ( pg(x g, yg)
pg(xg)pg(yg)
) . (B.3)
If f is differentiable, then for any u ∈ dom f , ∂f(u) = {f ′(u)} is a singleton.
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.
Proof. Simply combine Proposition 8 with Lemma 2.
Theorem 5 (uniformity). Suppose that the batch size N satisfies 2 ≤ N ≤ d + 1, with d the dimension of the feature space. If the real function
h(t) = f∗ ◦ f ′ ◦Gσ(t) is strictly convex on [0, 4], (12)
then all minimizers of the second term of equation 10, i.e., ∑ i 6=j f ∗ ◦ f ′ ◦Gσ(‖xgi − x g j‖2), satisfy the following condition: the feature representations of all samples are distributed uniformly on the unit hypersphere Sd−1. All f -divergences in Table 1 satisfy equation 12.
Proof. From the definition of h it is clear that h is decreasing since f∗ and f ′ are both monotonically increasing white Gσ is decreasing. Using h we rewrite the second term of equation 10 as
min xg1 ,...,x g N∈Sd−1 ∑ i,j h(‖xgi − x g j‖ 2). (B.4)
When N ∈ [2, d + 1], there exists a neat characterization of the minimizers, see e.g. Borodachov et al. (2019, Theorem 2.4.1). We include the proof below for completeness.
Apply Jensen’s inequality, we have:
1
N2 ∑ i,j h(‖xi − xj‖2) ≥ h  1 N2 ∑ i,j ‖xi − xj‖2 
= h  1 N2 ∑ i,j ‖xi − xj‖2 
= h  1 N2 ∑ i,j (2− 2xi · xj)  = h 2 1− ∥∥∥∥∥ 1N N∑ i=1 xi ∥∥∥∥∥ 2 
≥ h (2) , (B.5)
where in the first line we used Jensen’s inequality; in the third line we used ‖xi‖ = ‖xj‖ = 1 for any i, j ∈ [N ]; in the last line we note that ‖ ∑N i=1 xi‖ ≥ 0 and h is a decreasing function. When h is strictly convex and decreasing, it is in fact strictly decreasing, and hence the two inequalities above can be attained iff
x̄ := 1
N ∑ i xi = 0, and ‖xi − xj‖2 ≡ c for all i 6= j, (B.6)
namely that {x1, . . . , xN} form a regular simplex with its center at the origin. We remark that when h is merely convex, points forming a centered regular simplex may form a strict subset of the minimizers.
To see the necessity of N ≤ d+ 1, let us note that
x>i xj =
{ 1, i = j
− 1N−1 , i 6= j , (B.7)
since ∑ ij ‖xi − xj‖2 = 2N2 = N(N − 1)c =⇒ c = 2N N − 1 = 2 + 2 N − 1 . (B.8)
Performing simple Gaussian elimination we note that the matrix X>X has rank N − 1 where X = [x1, . . . , xN ] ∈ Rd×N . Therefore, we must have N − 1 ≤ d. Lastly, we need to show when h is a (strictly) convex function, which may not always be true depending on the f -divergences. We give the following characterization (we ignore the constants µ and 2σ2 in equation 7 as they do not affect convexity):
• h strictly convex: hKL(t) = e−t, hJS(t) = log(1 + e−t)− log 2, hPearson(t) = e−2t − 1, hSH(t) = e −t/2 − 1, hTsallis(t) = e−αt, hVLC = 3− 41+e−t ;
• h convex but not strictly convex: hRKL(t) = −t− 1 (RKL stands for Reversed Kullback– Leibler, see Appendix A.1);
• h concave: hNeyman(t) = 2− 2et (Neyman stands for Neyman χ2, see Appendix A.1).
Only for the last case we do not have the guarantee that the minimizing configurations could form a regular simplex. For RKL, in fact, any configuration that centers at the origin suffices since h is a linear function.
Theorem 6 (estimation error). Suppose that the function T is taken from a function class T and define Tx as the function class of T (x, ·) given some x ∈ supp(pdata). Denote RPN to be the
Rademacher complexity w.r.t. the distribution P with N i.i.d. drawn samples. Then for any T ∈ T , the estimation error |if (X;Y )− îf (X;Y )| is upper bounded with probability at least 1− δ:
2R ppos N (T ) + 2µ [ Ex∼pdataR pdata N (Tx) + 1
N N∑ i=1 RpdataN−1(Txj )
] + (rT + 2rf ) √ log 6/δ
2(N − 1) , (13)
with the constants rT = f ′(µ)− f ′(µe−2/σ 2 ), rf = f ∗ ◦ f ′(µ)− f∗ ◦ f ′(µe−2/σ2).
Proof. Our proof uses the following notion of Rademacher complexity:
Definition 11 (Rademacher complexity, Koltchinskii (2001)). Let F be a family of functions from a subset of Euclidean spaceW to [a, b] and S = (w1, . . . , wm) a fixed sample of sizem with elements inW . The empirical Rademacher complexity of F w.r.t. the sample set S is defined as:
R̂S(F) = Eσi [ sup f∈F 1 m m∑ i=1 σif(wi) ] , (B.9)
where σi’s are independent uniform random variables taking values {−1,+1}. Let P be the underlying distribution of samples. The Rademacher complexity Rm(F) of function class F is the expectation over sample sets:
RPm(F) = ES∼Pm [R̂S(F)]. (B.10)
With the definition of Rademacher complexity we have the following lemma:
Lemma 12. Let F be a family of functions from a subset of Euclidean spaceW to [a, b] with b > a. Then for any 0 < δ < 1 with probability at least 1− δ, the following holds for any f ∈ F:
|Ew∼Pw [f(w)]− 1
m m∑ i=1 f(wi)| ≤ 2RPwm (F) + (b− a) √ log(2/δ) 2m , (B.11)
with wi i.i.d. examples from any distribution Pw.
Proof. Our proof follows from Lemma 11 of Zhang et al. (2021). For any function f ∈ F we can construct f̃ = (f − a)/(b − a) such that f̃ : W → [0, 1]. The rest follows from the linearity of Rademacher complexity.
Let us first recall our definition equation 4: T (x, y) := k(g(x), g(y)). From the definitions of if (X;Y ) and îf (X;Y ) we write the following:
|if (X;Y )− îf (X;Y )| = ∣∣∣∣∣E(x,z)∼ppos [T (x, z)]− Ex∼pdataEy∼pdata [f∗(T (x, y))]− −  1 N N∑ i=1 T (xi, zi)− 1 N(N − 1) ∑ i 6=j f∗(T (xi, xj))
∣∣∣∣∣∣ ≤
∣∣∣∣∣E(x,z)∼ppos [T (x, z)]− 1N N∑ i=1 T (xi, zi) ∣∣∣∣∣+ +
∣∣∣∣∣∣Ex∼pdataEy∼pdata [f∗(T (x, y))]− 1N(N − 1) ∑ i 6=j f∗(T (xi, xj)) ∣∣∣∣∣∣ , (B.12)
where we used the triangle inequality. With Lemma 12, we can upper bound the first term of equation B.12 as:
2R ppos N (T ) + (bT − aT )
√ log(2/δ)
2N , w.p. 1− δ, (B.13)
where [aT , bT ] is the range of T . Now we upper bound the second term of equation B.12. Note that we cannot apply Lemma 12 to equation B.12 directly since (xi, xj)’s are not i.i.d. samples. For instance, the pairs (x1, x2) and (x1, x3) are not independent since they share the same x1. However, all xi’s are i.i.d. sampled. We can upper bound the second term of equation B.12 as:∣∣∣∣∣∣Ex∼pdataEy∼pdataf∗(T (x, y))− 1N(N − 1) ∑ i 6=j f∗(T (xi, xj))
∣∣∣∣∣∣ =∣∣∣∣∣∣Ex∼pdataEy∼pdataf∗(T (x, y))− Ex∼pdata 1N N∑ j=1 f∗(T (x, xj))+
+Ex∼pdata 1
N N∑ j=1 f∗(T (x, xj))− 1 N(N − 1) ∑ i 6=j f∗(T (xi, xj)) ∣∣∣∣∣∣ ≤∣∣∣∣∣∣Ex∼pdataEy∼pdataf∗(T (x, y))− Ex∼pdata 1N N∑ j=1 f∗(T (x, xj))
∣∣∣∣∣∣+ +
∣∣∣∣∣∣Ex∼pdata 1N N∑ j=1 f∗(T (x, xj))− 1 N(N − 1) ∑ i 6=j f∗(T (xi, xj)) ∣∣∣∣∣∣ . (B.14) The first term of equation B.14 can be upper bounded as:∣∣∣∣∣∣Ex∼pdata Ey∼pdataf∗(T (x, y))− N∑ j=1 f∗(T (x, xj))
∣∣∣∣∣∣ ≤ 2Ex∼pdataR pdata N (f ∗ ◦ Tx) + (bf − af ) √ log(2/δ)
2N
≤ 2LfEx∼pdataR pdata N (Tx) + (bf − af )
√ log(2/δ)
2N (B.15)
with probability at least 1−δ, where we used Lemma 12 in the first line and Talagrand’s lemma (e.g., Lemma 5.7 of Mohri et al. (2018)) in the second line. We denote Lf as the Lipschitz constant of f∗ and [af , bf ] as the range of f∗ ◦ T . Similarly, we can upper bound the second term of eq. (B.14) as:
Lf N N∑ j=1 RpdataN−1(Txj ) + (bf − af )
√ log(2/δ)
2(N − 1) , (B.16)
with probability at least 1− δ.
Absolute constants Let us compute the constants af , bf , aT , bT , Lf . First note that our feature embeddings are normalized. For any xg, yg ∈ Sd−1 we have:
0 ≤ ‖xg − yg‖2 ≤ 2(‖xg‖2 + ‖yg‖2) = 4. (B.17)
From the expression of Gaussian similarity equation 7 and the monotonicity of f ′ and f∗, we have:
bT = f ′(µ), aT = f ′(µe−2/σ 2 ), bf = f ∗ ◦ f ′(µ), af = f∗ ◦ f ′(µe−2/σ 2 ). (B.18)
For simplicity we will write rT = bT − aT and rf = bf − af . We can also compute the Lipschitz constant of f∗, which is bounded by the norm of the gradient:
Lf ≤ sup xg,yg∈Sd−1 (f∗)′ ◦ f ′ ◦Gσ(‖xg − yg‖2) = sup xg,yg∈Sd−1 Gσ(‖xg − yg‖2) = µ, (B.19)
where we used the equality (f∗)′ ◦ f ′ = id (Rockafellar, 1966), with id the identity function. Combining equation B.13, equation B.15 and equation B.16 we can use the union bound to finish the proof.

heading: C ADDITIONAL EXPERIMENTAL RESULTS
text: We present additional experiment details in this appendix, to further support our experiments in the main paper.
C.1 IMPLEMENTATION DETAILS
In this paper, we follow the implementations in SimCLR (https://github.com/sthalles/ SimCLR). For vision tasks, we use ResNet (He et al., 2016) as the feature encoder, and we adopt the similar procedure of SimCLR for sampling. For the language dataset, we follow the exact experimental setting of Gao et al. (2021) and only change the objective. Our experimental settings are detailed below:
• Hardware and package: We train on a GPU cluster with NVIDIA T4 and P100. The platform we use is pytorch. Specifically, the pairwise summation can be easily implemented using torch.nn.functional.pdist from pytorch.
• Datasets: the datasets we consider include CIFAR-10, CIFAR-100 (Krizhevsky et al., 2009), STL-10 (Coates et al., 2011), TinyImageNet (Chrabaszcz et al., 2017), ImageNet (Deng et al., 2009) and English Wikipedia (Gao et al., 2021).
• Augmentation method: For each sample in a dataset we create a sample pair, a.k.a. positive pair, using two different augmentation functions. For image samples, we choose the augmentation functions to be the standard ones in contrastive learning, e.g., in Chen et al. (2020) and He et al. (2020). The augmentation is a composition of random flipping, cropping, color jittering and gray scaling. For text samples, following the augmentation method of Gao et al. (2021) we use dropout masks.
• Neural architecture: For CIFAR-10 and CIFAR-100 we use ResNet-18 (He et al., 2016); for STL-10, TinyImageNet and ImageNet we use ResNet-50 (He et al., 2016); for the Wikipedia dataset we use BERTbase (Devlin et al., 2019).
• Batch size and embedding dimension: for experiments in CIFAR-10 and CIFAR-100 we choose batch size 512, and for STL-10 we choose batch size 64 to accommodate one GPU training. Finally, for TinyImageNet and ImageNet, we choose batch size 256. For all the vision datasets, we choose the embedding dimension to be 512. Regarding the language dataset, the batch size is 64 with the feature dimension 768. In all of these cases, our assumption N ≤ d+ 1 in Theorem 5 is satisfied.
• Hyperparameters: in all our experiments we fix the constant factor µ = 1. We find that in practice the weight parameter α often needs to be large (e.g., in the Wikipedia dataset), which requires moderate tuning. Note that we also implement RPC (Tsai et al., 2021) in our paper. For all the datasets, we follow Tsai et al. (2021) and choose the relative parameters α = 1.0, β = 0.005 and γ = 1.0 for all datasets.
• Optimizer and learning rate scheduler: For vision tasks, we use SGD with momentum for optimization and the cosine learning rate scheduler (Loshchilov & Hutter, 2017). For the natural language task, we use Adam with weight decay (Loshchilov & Hutter, 2018) and the linear decay scheduler.
• Evaluation metric: for vision tasks, we use k-nearest-neighbor (kNN) and linear evaluation to evaluate the performance, based on the learned embeddings. For the NLP task, we use the Spearman’s correlation to evaluate the averaged semantic textual similarity score (Gao et al., 2021).
• Baseline methods: for the four baseline methods, we follow the implementations in: – MoCo: https://github.com/facebookresearch/moco – SimCLR: https://github.com/sthalles/SimCLR – Uniformity: https://github.com/SsnL/align_uniform – RPC: https://github.com/martinmamql/relative_predictive_ coding.
For fair comparison we use the experimental settings in Table 5 for all the baseline methods, which might differ from the original settings.
Table 5 gives common choices of hyperparameters for different datasets. Note that we may need to further finetune α and σ for different f -divergences. See our supplementary code for more details.
C.2 ADDITIONAL ABLATION STUDY ON WEIGHTING PARAMETER
We provide additional ablation study on the weighting parameter α. We perform experiments using a vision dataset (CIFAR-10) and a language dataset (Wikipedia). For CIFAR-10, we vary α from 0.1 to 50 for KL and JS divergences and run for 200 epochs. Table 6 justifies our choice of α in Table 5, where the downstream test classification accuracy indicates the optimal performance when choosing α = 40. For the Wikipedia dataset, we observe that a much bigger α is desirable for maximum performance. We vary α from 1 to 106 for KL and Pearson χ2 divergences and run for 1 epoch, as there is a large number of samples (106) in the language dataset. Table 7 justifies our choice of α in Table 5, where the best performance is reached at α = 409600. Such an α is found by starting from α = 100 and doubling iteratively.
C.3 ADDITIONAL EXPERIMENTS
Our final experiments show that f -MICL is stable in terms of training and the variation of performance is well controlled.
Training stability We depict the training loss curves of different divergences on CIFAR-10 in Figure 4. This figure shows that our methods exhibit stable training dynamics with fast convergence.
kNN evaluation and additional f -divergences We show more detailed results of Table 2 in Table 8, including experiments using k-nearest neighbour (kNN) evaluation. Additionally, we have added experiments on other f -divergences such as Squared Hellinger and Tsallis-α divergences.
Verification of Assumption 3 Throughout our paper we made an assumption (Assumption 3) that the joint feature distribution is a Gaussian kernel. However, is it a valid assumption? In this experiment we try to show some empirical evidence that this assumption approximately holds in practice. Recall that Assumption 3 says that the joint feature distribution of positive pairs is:
pg(x g, yg) ∝ exp ( −‖x g − yg‖2
2σ2
) (C.1)
if the RBF kernel is Gaussian. In order to estimate the joint density of positive pairs, we use normalizing flows (Dinh et al., 2017), which is a popular method for the density estimation. Popular normalizing flow models include RealNVP (Dinh et al., 2017), NICE (Dinh et al., 2014) and Glow (Kingma & Dhariwal, 2018). Equation C.1 is equivalent to the following:
log pg(x g, yg) = −‖x g − yg‖2
2σ2 + const, (C.2)
and thus it suffices to show that the log likelihood is linear w.r.t. the distances between each positive pair. In Figure 5, we plot the relation between log pg , estimated by RealNVP1, and the squared distances ‖xg − yg‖2. The representations are learned by SimCLR, and f -MICL with the KL divergence on the CIFAR-10 dataset. To alleviate the estimation error in the flow model, we divide the distances into small intervals and compute the average log likelihood within each interval. From this figure we can see that the log likelihood is roughly linear w.r.t. the squared distance, and thus verifying our Assumption 3.
1Code available at https://github.com/ikostrikov/pytorch-flows.

Review summaries:
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. 

This paper proposes the f-mutual information objective for contrastive learning, which generalizes the existing mutual information framework. Theoretically, it makes a connection with alignment and uniformity in Wang & Isola (2020) when the joint feature distribution can be represented by a radial basis function, and derives its estimation error. Numerically, it demonstrates the efficacy of f-mutual information with Gaussian similarity.

This paper proposed to combine contrastive learning with $f$-divergence, which naturally attains the alignment for positive pairs and the uniformity for data representations. The authors provided some theoretical results on choosing the correct function f and the upper bound of the estimation error. Finally, the authors carried out extensive experimental results to validate the effectiveness of the method.