Contrastive Learning Can Find An Optimal Basis For Approximately Invariant Functions
Abstract
Contrastive learning is a powerful framework for learning self-supervised representations that generalize well to downstream supervised tasks.
We show that multiple existing contrastive learning methods can be reinterpeted as learning a positivedefinite kernel that approximates a particular contrastive kernel defined by the positive pairs.
The principal components of the data under this kernel exactly correspond to the eigenfunctions of a positive-pair Markov chain, and these eigenfunctions can be used to build a representation that provably minimizes the worst-case approximation error of linear predictors under the assumption that positive pairs have similar labels.
We give generalization bounds for downstream linear prediction using this optimal representation, and show how to approximate this representation using kernel PCA.
We also explore kernel-based representations on a noisy MNIST task for which the positive pair distribution has a closed form, and compare the properties of the true eigenfunctions with their learned approximations.
1. Introduction
Contrastive learning models, such as SimCLR (Chen et al., 2020), have shown great success at learning self-supervised representations for downstream supervised learning tasks, by learning to map “positive pairs” (generally augmentations of the same image) close together in an embedding space.
Inspired by their success, a variety of theoretical analyses have been proposed, including uniformity and alignment of hyperspherical embeddings (Wang and Isola, 2020), conditional independence structure with landmark embeddings (Tosh et al., 2021) and spectral analysis of an author(s). augmentation graph (HaoChen et al., 2021).
These analyses are generally quite specific to particular loss functions, and primarily focus on classification tasks.
In this work, we generalize and extend these analyses using techniques from kernel methods, showing that multiple contrastive learning methods can be interpreted as approximating a particular “contrastive kernel”, and that this kernel can be used to build good representations for downstream classification or regression.
We assume dataset examples are drawn from a distribution x ∼ p
( X ) over a (discrete) set X , and that from each example we sample multiple augmented views a 1 , a 2 ∼ p ( A |
X = x ) over another (discrete) set A .
At pretraining time, we sample positive pairs a 1 , a 2 from the distribution p + ( a 1 , a 2 ) =
(cid:80) x p ( x ) p ( a 1 | x ) p ( a 2 | x ) , and negative pairs independently from the distribution p ( a ) =
(cid:80) x p ( x ) p ( a | x ) , which we use to train a contrastive learning model;
for convenience we assume population access to p ( X ) .
Afterward, we seek to approximate a particular target function g : A → R using a linear predictor, given a small number of labeled (augmented) views ( a i , y i ) ∼ p
( A ) p ( Y | A ) with
E [ y i | a i ] =
g ( a i ) .
Our main assumption is that the target function g has similar values for positive pairs:
Assumption 1.1 (Approximate invariance / small positive– pair discrepancy) .
The target function g satisfies
[Equation]
where the expectation is over paired views p + ( a 1 , a 2 | x ) of examples x drawn from the true data distribution p ( x ) .
We prove that a particular representation based on Kernel
PCA (Schölkopf et al., 1997) is optimal under this assumption, in the sense that it minimizes least-squares approximation error of the worst-case target function, and give generalization bounds for its performance.
We then show that existing contrastive learning methods can be used to approximately recover this representation.
2. Contrastive Learning Models Are Kernels
We start by identifying a common structure between common contrastive learning models and objectives (see Table
[Figure]
[Table]
[Figure]
[Table]
[Figure]
[Table]
[Figure]
[Table]
1 and Appendix A): many of them are a combination of a parameterized positive-definite kernel and an objective whose minimum (in the infinite data and capacity limit) is exactly the probability ratio p + ( a 1 , a 2 ) /p
( a 1 ) p ( a 2 ) (where all probabilities are with respect to the data distribution).
Furthermore, this probability ratio is also a positive-definite kernel (i.e. it is an inner product in a transformed space):
Definition 2.1.
The contrastive kernel associated with distributions p ( x ) and p ( a | x ) is the probability ratio
[Equation]
where ϕ Ctr ( a ) ∈ R |X| is the vector
[Equation]
Data points a 1 and a 2 that are more likely to be a positive pair than a negative pair are mapped to vectors ϕ Ctr ( a 1 ) and ϕ Ctr ( a 2 ) with a high inner product.
Conversely, if a 1 and a 2 have zero probability of being a positive pair, ϕ Ctr ( a 1 ) and ϕ Ctr ( a 2 ) are orthogonal.
Since multiple contrastive learning approaches have this kernel as their optimum, we can try to understand the behavior of contrastive learning by investigating properties of this kernel.
See Figure 1 for a comparison of the true and approximate kernels.
3. Kernel Principal Components Are Positive-Pair Eigenfunctions
To obtain a lower-dimensional representation from the kernel K Ctr or an approximation (cid:98) K θ , we can use Kernel Principal Components Analysis (Schölkopf et al., 1997), which implicitly computes the directions of maximum variance of a dataset within the kernel’s implicit inner product space.
The output of kernel PCA is a set of principal component projection functions f 1 , f 2 , . . .
and corresponding eigenvalues λ 1 , λ 2 , . . .
, such that f i ( a ) measures the projection of ϕ Ctr ( a ) onto the i th eigenvector of an implicit covariance matrix, and λ i measures the variance in that direction.
(We will assume the f i are scaled so that E [ f i ( a ) 2 ]
= 1 .)
Given a dataset of augmentations and a learned kernel (cid:98) K θ , we can run Kernel PCA and extract the first d projection functions [ ˆ f 1 ( a ) , ˆ f 2 ( a ) , . . .
, ˆ f d ( a )] to construct a “kernel
PCA representation” of any new augmented data point a .
As our dataset grows, and as our learned kernel (cid:98) K θ approaches the true contrastive kernel K Ctr , we would expect this representation to approximate the first d projection functions [ f 1 ( a ) , f 2 ( a ) , . . .
, f d ( a )] of the population covariance of the augmentation distribution under K Ctr (i.e. the covariance of ϕ Ctr ( A ) ).
It turn out that this representation is a particularly good choice for downstream supervised linear prediction under Assumption 1.1, because these populationlevel PCA projection functions are closely connected to a particular positive-pair Markov chain described below.
Starting with an example a 1 , we could sample another example a 2 proportional to how likely it ( a 1 , a 2 ) would be a positive pair, e.g. according to p + ( a 2 | a 1 ) =
[Equation]
x p ( a 2 | x ) p ( x | a 1 ) .
We can then repeat the process, sampling a 3 according to p
+ ( a 3 | a 2 ) .
This defines a Markov chain over the space A of (augmented) examples, which we visualize in Figure 2.
To the extent that some function g is invariant to the chosen augmentations, we would also expect the value of g to change slowly (or not at all) along this random walk: we would expect g ( a 1 )
≈ g
( a 2 )
≈ g ( a 3 )
≈ · · · .
Let P : R |A|×|A| be the Markov chain transition matrix, such that P ij = p ( a
=
i | a 1 = j ) .
We can then compute the left eigenvectors f i and eigenvalues λ
′ i of this matrix, such that f i P =
λ
i f i .
Each element [ f i ] a of an eigenvector f
i is a value associated with augmentation a ∈ A , and the eigenvector equation implies that
[Equation]
In other words, for a fixed a 1 ∈ A , taking the expected value of entries of f i over positive-pair neighbors a 2 of ais equivalent to scaling the entry for a 1 by λ i .
Note that we are free to choose each f i so that E (cid:2)
[ f i ] 2 a (cid:3) = 1 .
Our key insight is that the values of the principal component projection functions f i are the same as the entries of the
Markov chain eigenvectors f i , and that expressing other functions in terms of them allows you to measure invariance:
Theorem 3.1.
f i ( a ) =
[ f i ] a and λ i = λ
′ i for all i and all a ∈ A .
Furthermore, given any function g : A → R , if we let c i =
E [ g ( a ) f
i ( a )] , then the following are true:
[Equation]
This theorem has a few important consequences.
First, if we let g = f i for some i , we find that [Equation] Thus, there is a linear relationship between the variance λ i captured by the PCA projection functions and the positivepair discrepancy of those functions. Second, by substituting into Equation 2, we find that E p + ( a 2 | a 1 )
[ f i ( a 2 )]
= λ
i f i ( a 1 ) ; the population-level PCA projections
f i are thus also eigenfunctions of the Markov chain, and the eigenvalues of the kernel’s implicit covariance matrix E [ ϕ Ctr ( a ) ϕ Ctr ( a ) T ] are also eigenvalues of the Markov transition matrix P .
Third, Equation 4 implies that any target function g satisfying Assumption 1.1 must satsify 2 (cid:80) i (1 − λ i ) c 2 i ≤ ε , and thus must have coefficients c i concentrated on eigenfunctions with λ i close to 1 .
Indeed, if ε = 0 (i.e. if g is perfectly invariant to augmentations) then the only eigenfunctions with nonzero weights must be those with λ i = 1 .
Such eigenfunctions identify the subsets of augmentations for which positive pairs always remain within one subset; these are the communicating classes of the Markov chain, or equivalently, the connected components of the augmentation graph described by HaoChen et al.
4. Eigenfunction Representations are Optimal
We now give a more precise analysis of the quality of these eigenfunctions for downstream supervised fine-tuning.
We focus on the class of linear predictors on top of a particular k -dimensional representation r : A → R d , e.g. we will approximate g with a parameterized function ˆ g β ( a ) =
β ⊤ r ( a ) .
It turns out that the representation consisting of the d eigenfunctions with largest eigenvalues λ 1 ≥ λ 2 ≥ . . .
is the best choice under two simultaneous criteria.
Proposition 4.1.
Let F r = { a (cid:55)→ β ⊤ r ( a ) : β ∈ R d } be the family of linear predictors from representation r , and S ϵ be the set of functions satisfying Assumption 1.1.
Let r d ∗ ( a ) =
[ f 1 ( a ) , f 2 ( a ) , . . . ,
f d ( a )] be the representation consisting of the d largest eigenfunctions of the positive pair
Markov chain.
Then F r d ∗ maximizes the invariance of the least-invariant unit-norm predictor in F r d ∗ :
[Equation]
Simultaneously, F r d ∗ minimizes the least-squares approximation error for the worst-case target function in S ϵ :
[Equation]
Equation 6 states that the function class F r has an implicit regularization effect: it contains the functions that change as little as possible over positive pairs, relative to their norm.
Equation 7 reveals that this function class is also the optimal choice for least-squares approximation of a function satisfying Assumption 1.1.
Together, these findings suggest that this representation should give good generalization performance with only a few labeled examples, as long as
Assumption 1.1 holds.
Consider the loss function ℓ (ˆ y, y ) :
= | ˆ y − y | , and the associated risk R ( g ) =
E [ ℓ ( g ( A ) , Y )] of a predictor g : A → R .
Let g ∗ ∈ argmin R ( g ) be the optimal predictor with risk
R ∗ =
R ( g ∗ ) , and assume it satisfies Assumption 1.1.
Expand g ∗ as a linear combination of the basis of eigenfunctions ( f i )
Ki
=1 as g ∗ =
(cid:80)
Ki
=1 β ∗ i f i , and define the vector β ∗ := ( β ∗ 1 , . . .
, β ∗ d ) .
Then the following generalization bound holds for downstream supervised learning:
Proposition 4.2.
Let ( A i , Y i ) n i =1 be i.i.d.
samples, choose
R ≥ 0 , and consider the constrained empirical risk minimizer
[Equation]
Then the expected excess risk of ˆ β R is bounded by:
[Equation]
where E ( β ) :
= R ( β ) − R ∗ is the excess risk and ( x ) + := max { x, 0 } . Note that ∥ β ∗ ∥ 22 = E [ g ∗ ( a ) 2 ] by Equation 3, so if
E [ g ∗ ( a ) 2 ] ≤ R then the second term vanishes.
The third term bounds the error incurred by using only the first d eigenfunctions, since β ∗ 2 i ≤ ε 2(1 − λ i ) by Equation 4.
5. Related work
Our work is closely connected to the spectral graph theory analysis by HaoChen et al.
Their analysis focuses on the eigenvectors of the normalized adjacency matrix of the augmentation graph, and they introduce the spectral contrastive loss and show that its optimum recovers the top eigenvectors up to an invertible transformation.
Our work extends theirs by unifying their loss with other losses under a kernel framework, and showing that the kernel principal components are optimal under Assumption 1.1.
(Their eigenvectors turn out to be adjusted versions of our eigenfunctions; see Appendices A.3 and B.1.)
We note that our assumption is closely related to the Laplacian matrix of the augmentation graph; similar assumptions have been used before for label propagation (Bengio et al., 2006) and Laplacian filtering (Zhou and Srebro, 2011).
This assumption is also used as a “consistency regularizer” for semi-supervised learning (Sajjadi et al., 2016; Laine and Aila, 2016).
Tian (2022) gave a different unification of contrastive losses using a game-theoretic perspective, and showed that linear networks on contrastive losses are related to PCA in the input space (whereas our analysis applies to the unconstrained minimizer of those losses, with PCA in the contrastive kernel’s implicit inner product space).
There has been a variety of work on learning kernels with neural networks (Wilson et al., 2016; Sun et al., 2018) and extracting kernels from neural net architectures (Jacot et al., 2018; Shankar et al., 2020; Amid et al., 2022), which may be relevant to contrastive learning in light of our analysis.
6. Experiments
Our theoretical analysis suggests that, to the extent that the learned kernel (cid:98) K θ approximates the data-dependent contrastive kernel K Ctr , kernel PCA on a learned kernel (cid:98) K θ may yield a good approximation of the optimal eigenfunction basis for downstream prediction.
To explore the accuracy of this approximation in practice, and evaluate the quality of the extracted representations, we constructed a toy problem for which the exact contrastive kernel is tractable.
We selected a subset of MNIST digits (LeCun et al., 1998) as our set of base images X , then carefully chose the augmentation distribution p ( a | x ) so that we could compute the exact contrastive kernel K Ctr .
Specifically, we defined p ( a | x ) for each x ∈ X by first selecting one of 64 fixed perturbations of x , then sampling a subset of k filled pixels from this perturbed copy.
We can then tractably compute p ( a | x ) for any a and x , enabling us to evaluate K Ctr and also to sample from the positive pair Markov chain. As shown in Figure 2, there is often uncertainty about which original image generated a given augmentation due to the sparsity of pixels.
We next trained three neural network models with varying kernel heads and loss functions: a linear kernel head with the spectral loss (based on HaoChen et al.
; a temperature-normalized exponential head exp( h θ ( a 1 ) T h θ ( a 2 ) /τ ) with a weighted combination of the NT-XEnt and NT-Logistic losses (constraining h θ ( a ) to be unit-norm, based on Chen et al.
(2020)); and a scaled temperature-normalized exponential head exp( h θ ( a 1 ) T h θ ( a 2 ) /τ )
· s θ ( a 1 ) s θ ( a 2 ) , which includes a learned adjustment s θ : A → R + , again trained with a weighted combination of the NT-XEnt and NT-Logistic losses.
All models used a simplified ResNet architecture followed by a fully-connected “projection head”, and were trained on 512 MNIST digits from each of the 10 classes. For visualization purposes, we also trained a model with a two-dimensional embedding space and a scaled rational quadratic kernel head on MNIST digits 0, 1, and 2 only; this model is shown in Figures 1 and 3. Properties of extracted principal components. We used
Kernel PCA over a random set of augmentations to to estimate the first 256 principal component projection functions for K Ctr and for each of our learned kernels. We then investigated (a) whether the learned principal components were aligned with the principal components of K Ctr , and (b) whether their eigenvalues are related to their positive-pair discrepancies (estimated from a random sample of positive pairs) as predicted by Equation 5. The results of these comparisons are shown in Figure 4, for p ( a | x ) using k = 20 sampled pixels per digit. Since
Kernel PCA is performed on a random subset, there is some measurement noise; thus even two runs of KPCA on K Ctr produce slightly different results.
The linear kernel with spectral loss is able to recover the top principal components quite accurately, and also shows the expected relationship between eigenvalues and positive-pair discrepancies.
The temperature-normalized exponential kernel is less well-behaved, but this appears to be due to the constrained form of the kernel head; adding the scale adjustment s θ ( a 1 ) s
θ ( a 2 ) leads to better alignment and to the expected linear eigenvalue-discrepancy relationship.
We also found that increasing augmentation strength (smaller k ) leads to
[Table]
a better alignment for all kernels, whereas decreasing it (larger k ) makes alignment worse; see Appendix D.6. Downstream prediction ability. We next compare the quality of various representations for two downstream prediction tasks: classification with a linear layer [Table]
and linear least-squares regression on the one-hot indicator vectors for each digit class.
We consider three types of representation: the PCA projection functions for K Ctr , the PCA projection functions for each learned kernel (cid:98) K θ , and the intermediate layer embedding vector between the ResNet layers and the projection head for each model as proposed by Chen et al.
For each, we fit a regularized linear predictor on 160 labeled training examples (16 augmented samples from each class), using 160 additional validation examples to tune the regularization strength. For PCA representations, we additionally tune the representation dimension d . The results are shown in 2.
Performance is fairly similar across representations, suggesting that the contrastive kernel
K Ctr captures much of the variability between augmentation strengths, although some learned representations achieve better accuracy. Interestingly, adding the scale parameter to the NT-Exp model improves the kernel PCA representation but degrades performance for the ResNet embedding.
7. Discussion
We have shown that multiple existing contrastive objectives approximate a particular kernel, and that its principal component projection functions give an optimal representation for supervised linear prediction under the assumption that positive pairs have similar labels.
This analysis is based on the “worst” target function that satisfies our assumption; if we have additional knowledge about the target function, incorporating it may improve the representation.
Empirically, we find that Kernel PCA with learned models can produce good approximations of the largest-eigenvalue principal components, but that the quality of the approximation depends on the constraints of the kernel head as well as the strength of the augmentations.
Overall, we hope that our kernel and Markov chain perspective is a useful lens for further study of contrastive learning representations.
A. Existing objectives are minimized by the contrastive kernel
In this section, we show that the contrastive kernel is the optimum of the objectives shown in Table 1 (although this optimum is not unique for cross-entropy loss).
Throughout this section, we use p ( X ) to denote the true data distribution over unperturbed examples, p ( A | X ) to denote the distribution of augmented views conditioned on a particular unperturbed example, and p ( A ) to denote the marginal distribution of augmented views, e.g.
[Equation]
We use p + ( A 1 , A 2 ) to denote the positive pair distribution induced by p ( X ) and p ( A | X ) , defined by p
+ ( A 1 = a 1 , A 2 = a 2 ) =
(cid:88)
[Equation] p [Equation] p
[Equation] A = a [Equation] | X = x
[Equation] p ( A = a 2 | X = x ) . For notational convenience, we will use shorthand p ( x ) for p ( X = x ) , p ( a ) for p ( A = a ) , p + ( a 1 , a 2 ) for p + ( A 1 = a 1 , A 2 = a 2 ) , and so on.
A.1. NT-XEnt and the InfoNCE objective
The NT-XEnt objective described by Chen et al.
(2020) (and previously used by Sohn (2016);
Van den Oord et al.
Wu et al.
(2018)) has the form
[Equation]
For simplicity, we assume that all of the negative samples a − i are drawn independently from the marginal distribution when computing the loss for a +1 and a +2 .
(In practice, implementations often generate negative samples by taking elements of other positive pairs, e.g. ( a − 1 , a − 2 ) ∼ p + ( a − 1 , a − 2 ) , ( a − 3 , a − 4 ) ∼ p + ( a − 3 , a − 4 ) , and so on.)
We can decompose this objective into two parts: an InfoNCE-like loss (Van den Oord et al., 2018)
[Equation]
combined with a particular parameterized function
[Equation]
where h θ : A → R n +1 maps inputs to points on the n -dimensional hypersphere.
We first observe that (cid:98) K
θ ( a 1 , a 2 ) defines a positive definite kernel, within the family of “dot product kernels”.
Indeed, when h θ is restricted to the unit hypersphere, this parameterization is equivalent to the squared-exponential kernel (also called radial-basis-function kernel)
[Equation]
using the identity ∥ h θ ( a
[Equation]
[Equation]
Although this kernel is positive definite, it has “infinite rank” and cannot be expressed as an inner product of finite-dimensional embedding vectors; nevertheless, we can still run algorithms such as Kernel PCA on a finite dataset.
(See Bach (2021, Chapter 7) for some additional background on positive-definite kernels, and Scetbon and Harchaoui (2021) for discussion of other dot product kernels on the unit hypersphere.)
We next discuss the minimum of the NT-XEnt objective, under the unconstrained setting where we allow (cid:98) K θ ( a 1 , a 2 ) to be an arbitrary symmetric function.
The InfoNCE loss, in its more general form, is not necessarily symmetric: it is based on a distribution of contexts p ( C ) , positive samples p ( A | C ) , and negative samples drawn from the marginal distribution P ( A ) , and is given by
[Equation]
Van den Oord et al.
(2018) show that every minimizer of this objective is of the form
[Equation]
for some function z ( c ) that does not depend on a .
In other words, holding c fixed, f ∗ ( c, a ) ∝ p ( a,c ) p ( a ) p ( c ) .
Intuitively, this is because the exact probability of ( c, a + ) being the positive pair given c and the set { a + , a − 1 , . .
. , a − K } is also proportional to this density ratio, and the InfoNCE objective is a cross-entropy objective for identifying the positive pair.
(See also Poole et al.
for a different proof.)
In the case of the NT-XEnt contrastive objective, we choose the context C to be one of the augmentations A +1 , and the positive sample to be the other augmentation A +2 drawn according to p
+ ( A +2 |
A +1 ) .
We furthermore restrict our attention to symmetric functions (cid:98) K θ , e.g. functions for which (cid:98) K
θ ( a 1 , a 2 ) =
(cid:98) K
θ ( a 2 , a 1 ) .
In this case, the minimizer is
[Equation]
so we must have z ( a 1 ) =
z ( a 2 ) for any pair ( a 1 , a 2 ) with positive probability under p + .
If we assume that the positive pair Markov chain is irreducible, e.g. that there is a single communicating class and it is possible to reach any augmentation in A from any other augmentation over a long enough trajectory, then the function z ( a ) must be constant everywhere, and thus
[Equation]
for some Z ∈ R + .
In this case, (cid:98) K ∗ is equivalent to K Ctr up to a scaling constant.
If the Markov chain has multiple communicating classes (e.g. positive pairs are only sampled within a single communicating class), the function z ( a ) may assign a different value to different equivalence classes.
Nevertheless, any such minimizer is still a kernel, since we can write it as
[Equation]
Indeed, such a minimizer is equivalent to K Ctr except that it scales the inner product by the value of z ( a ) for each communicating class.
It is still possible to extract the set of Markov chain eigenfunctions from the set of principal components of this kernel, although one must correct for the scaling factor when computing the eigenvalues; see Appendix
B.2 for details.
Alternatively, one can ensure a unique minimum by combining the NT-Xent/InfoNCE loss with either the spectral or logistic losses (discussed below).
A.2. Logistic losses and NT-Logistic
Logistic losses have also been proposed for contrastive learning, including the NT-Logistic objective as described in Chen et al.
and other versions described by Mikolov et al.
and Tosh et al.
Such losses take the form
[Equation]
where negative samples are drawn independently from the marginal distribution p ( A ) .
Tosh et al.
(2021) motivates this loss based on a binary classification: choose a label Y to be 0 or 1 with probability 1 / 2 each, sample a positive pair ( a 1 , a 2 ) ∼ p + ( A 1 , A 2 ) if Y = 1 and a negative pair a 1 ∼ p ( A ) , a 2 ∼ p
( A )
if Y = 0 , then use a learned model to predict
Y given the pair.
The minimizer of this loss is then the conditional log-odds-ratio f
[Equation]
For the particular case of the NT-Logistic objective, we parameterize f θ as
[Equation]
where we again define
[Equation]
The optimum (if we ignore the constraints of this particular form of ˆ K and minimize over all functions of two variables) is then
[Equation]
Note that in this case there is no proportionality constant.
A.3. The Spectral Contrastive Loss
HaoChen et al.
(2021) propose the Spectral Contrastive Loss as an alternative to other contrastive losses with provable performance guarantees.
The loss is defined as
[Equation]
where they choose
[Equation]
for a learned embedding function h θ : A → R d .
We note that this directly satisfies the definition of a kernel
[Equation]
in that it is an inner product in a transformed space.
One interesting property of this kernel approximation is that it can be negative, whereas the exponential-based kernel approximations in the previous sections are always nonnegative.
HaoChen et al.
show that this loss can be rewritten as
[Equation]
where C = (cid:80) a 1 ,a 2 p + ( a 1 ,a 2 ) 2 p ( a 1 ) p ( a 2 ) is a constant independent of the model.
If we again ignore the constraints on ˆ K θ , the minimum of the spectral loss must occur when
[Equation]
for all a 1 and a 2 , or in other words, when
[Equation]
HaoChen et al. continue by expanding their definition of (cid:98) K θ for a fixed representation d , and showing that it relates to the spectral decomposition of a particular augmentation graph.
It turns out that this decomposition is equivalent to our decomposition in terms of eigenfunctions except for a scaling factor of p ( a ) 1 / 2 ; we discuss this connection more in
Appendix B.1.
B. Relationship between contrastive kernel and Markov chain eigenfunctions
In this section, we describe the relationship between the contrastive kernel principal components and the Markov chain eigenfunctions in more detail.
We start by introducing some notation that will be useful.
Throughout, we will identify functions f : A → R as vectors f : R A , which will allow us to use matrix notation for many of the relevant quantities.
We will also use e i to represent the vector that has a one at the i th position and zeros in all other positions.
We will let
[Equation]
be diagonal matrices containing the marginal probabilities of each element in X and A , respectively, under the true data distribution.
We will assume that the distribution has full support, and thus both D X and D A are invertible.
We also define the matrices
[Equation]
Equivalently
[Equation]
← A ) ⊤ .
From these, we can construct the positive pair probability matrix
[Equation]
The Contrastive Kernel.
Writing the contrastive kernel K Ctr in matrix form, such that [ K Ctr ] i,j = K Ctr ( i, j ) , our definition
K Ctr ( a 1 , a 2 ) =
p + ( a 1 ,a 2 ) p ( a 1 ) p ( a 2 ) becomes the matrix equation
[Equation]
One way to expand K Ctr is as a product
[Equation]
where Φ Ctr ∈ R
X×A is a matrix whose columns are given by ϕ Ctr : [Equation] [Equation] [Table]
[Equation]
[Table] [Equation] [Table] [Equation]
[Table] [Equation] [Table] [Equation]
  [Equation]  .
Equivalently, we have ϕ Ctr ( a ) =
Φ Ctr e a .
Since K Ctr ( a 1 , a 2 ) =
ϕ
Ctr ( a 1 ) ⊤ ϕ Ctr ( a 2 ) , ϕ Ctr is called a feature map for K Ctr .
Note that there are multiple possible feature maps for K Ctr : given any orthonormal matrix Q , the function ϕ Q ( a ) =
Qϕ Ctr ( a ) is also a feature map for the kernel, since
[Equation]
Performing kernel PCA under K Ctr is equivalent to performing ordinary PCA over one of its feature maps (since the principal component projection functions are independent of the particular feature map chosen).
We thus focus on analyzing the principal component projection functions for the feature map ϕ Ctr .
The population level principal components are the eigenvectors of the (uncentered) covariance matrix
[Equation]
Note that we are working with the uncentered principal components, as is commmon for kernel PCA: we do not subtract the mean before computing the covariance.
Since Σ is positive semidefinite, it can be diagonalized as
[Equation]
where U =
[ u 1 , u 2 , . . .
, u k ] is orthonormal and Λ =
diag
( λ 1 , λ 2 , . . .
, λ k ) is a diagonal matrix of eigenvalues (here k = |X| is the dimension of the feature map).
Each of the vectors u
i is one of the population principal components of the transformed distribution ϕ Ctr ( A ) , giving the directions of maximum variance, and the λ i measure the variance in that direction.
Given a new augmentation a ∈ A , we can then compute the projection of ϕ Ctr ( a ) into each of these principal component directions as ˜ f
i ( a ) =
u ⊤ i ϕ Ctr ( a ) .
Since the λ i measure the variance in each direction, we can optionally rescale these projection functions as
[Equation]
so that E [ f i ( a ) 2 ] = 1 .
(These are only well defined for λ i > 0 , but we are free to extend our set with additional functions to span the space of all functions
A → R .)
The Markov Chain.
We now redirect our attention to the positive pair Markov chain.
The Markov chain transition matrix is defined by [ P A ← A ] a 1 ,a 2 = p
+ ( a 1 | a 2 ) , or in matrix form
[Equation]
We are interested in the left eigenvectors f ⊤ i P
A ← A = λ ′ i
f ⊤ i of this matrix P
A ← A , or equivalently the right eigenvectors of its transpose P ⊤
A ← A = P A → A , given by P A → A f i = λ ′ i f i .
Observe that then
[Equation]
so equivalently
[Equation]
It follows that
D 1 / 2 A
f
i is an eigenvector of the symmetric matrix
[Equation]
with the same eigenvalue λ
i , and so we can diagonalize M as M = V Λ ′ V ⊤ , where
[Equation]
where
[Equation]
(We note that the matrix M is exactly the symmetrized adjacency matrix described by HaoChen et al.
Putting them together.
We can now prove our main result from Section 3.
Theorem 3.1.
f i ( a ) =
[ f i ] a and λ i = λ
′ i for all i and all a ∈ A .
Furthermore, given any function g : A → R , if we let c i =
E [ g ( a ) f
i ( a )] , then the following are true:
[Equation]
Proof.
For the first claim, consider the matrix B =
Φ Ctr D 1 / 2 A , where Φ Ctr = D 1 / 2 X P X
→ A D − 1 A described above.
Take the singular value decomposition B =
U Λ 1 / 2 V ⊤ , where U and V are orthonormal and Λ 1 / 2 is diagonal.
Now observe that
[Equation]
and
[Equation]
Thus, Σ and M must have the same eigenvalues.
For any i for which λ i > 0 , we then have
[Equation]
(When λ i = 0 , the normalized principal component projection functions are not uniquely defined [Equation] so we are free to simply set them equal to [ f i ] a [Equation]
Next, consider an arbitrary function g : A → R [Equation] and let g ∈ R
A be its vector form
[Equation] so that g ( a ) = g a = g ⊤ e a . Also define c i =
E [ g ( a ) f
i ( a )] and c = (cid:2) cc [Equation]
c k (cid:3) ⊤ .
Then
[Equation]
so we must have
[Equation]
and thus g ( a ) = (cid:80) i c
i f i ( a ) .
Additionally, we see that
[Equation]
[Table]
Note that this also implies that the functions f i are orthonormal under the base measure p ( a ) , e.g. E [ f i ( a ) 2 ] = 1 and E [ f i ( a ) f j ( a )]
= 0 for i ̸
= j .
Finally, consider the quantity E p + (cid:104)(cid:0) g
( a 1 ) − g ( a 2 ) (cid:1) 2 (cid:105) .
Observe that
[Equation]
Expanding the second term, we have
[Equation]
We conclude that E p +
[Equation]
λ [Equation]
B.1. Relationship to the eigenvectors of the symmetrized adjacency matrix
Interestingly, the matrix M described above is exactly the symmetrized adjacency matrix discussed by HaoChen et al.
HaoChen et al.
motivate their loss as estimating the eigenvectors of M up to a scaling term by p ( a ) 1 / 2 , due to prior work showing that eigenvalues give information about clustering structure in graphs.
The connection between the symmetrized adjacency matrix M and the positive-pair Markov chainis well known; indeed,
HaoChen et al.
briefly discuss the positive-pair Markov chain in their Section 2, and Levin and Peres (2017, Chapter 12) introduce the matrix M when discussing the spectral decomposition of a general symmetric Markov chain.
Another way of thinking about this reweighting is as a change of measure .
The eigenvectors of M are orthonormal with respect to the counting measure over A , e.g. if you sum squared values over all of A , you obtain 1, and the dot product of different eigenvectors is zero.
On the other hand, the eigenvectors f i (or, equivalently, the eigenfunctions f i ) of the Markov chain are orthonormal with respect to the measure p ( A ) , e.g. if you take the expectation of squared values over random augmentations, you obtain 1, and the uncentered covariance of different eigenfunctions is zero.
We believe that including the reweighting terms is a more natural perspective which gives additional insight into what contrastive losses are fundamentally doing.
We also note that our Assumption 1.1 is related to the probability-weighted Laplacian matrix of the augmentation graph, given by L
= D
A − P A,A .
Indeed, we have
[Equation]
B.2. Recovering proportionality constants
As described in Appendix A.1, minimizing the NT-XEnt / InfoNCE loss may not exactly produce the contrastive kernel
K Ctr , but may instead learn a scaled version
[Equation]
where z : A → R + is some function which is constant on each communicating class on the Markov chain.
(Since
K Ctr ( a 1 , a 2 ) =
whenever a 1 and a 2 are in separate communicating classes, we could equivalently say (cid:98) K ∗ ( a 1 , a 2 ) =
z ( a 1 ) K Ctr ( a 1 , a 2 ) =
z ( a 2 ) K Ctr ( a 1 , a 2 ) .)
When the Markov chain has one communicating class, (cid:98) K ∗ is simply a scaled version of K Ctr .
In this case, all of the principal component projection functions for (cid:98) K ∗ are still the eigenfunctions of the Markov chain, but the eigenvalues may be scaled by that constant.
The true eigenvalues of the eigenfunctions can then be estimated using equation Equation 5, which states that E [( f i ( a 1 ) − f i ( a 2 )) 2 ]
=
2(1 − λ i ) .
When the Markov chain has multiple communicating classes, we can partition the eigenfunctions so that each eigenfunction is nonzero on a single communicating class.
Since the scaling function z acts as a scaling factor for each communicating class, the principal component functions will then be scaled copies of these partitioned eigenfunctions.
We can then similarly estimate the true eigenvalues for each of these eigenfunctions using Equation 5.
C. Generalization properties of the eigenfunction representation C.1. Min-max optimality of eigenfunctions
We now prove the min-max optimality of the eigenfunctions with respect to their L2 norm.
Proposition 4.1.
Let F r = { a (cid:55)→ β ⊤ r ( a ) : β ∈ R d } be the family of linear predictors from representation r , and S ϵ be the set of functions satisfying Assumption 1.1.
Let r d ∗ ( a ) =
[ f 1 ( a ) , f 2 ( a ) , . . . ,
f d ( a )] be the representation consisting of the d largest eigenfunctions of the positive pair Markov chain.
Then F r d ∗ maximizes the invariance of the least-invariant unit-norm predictor in F r d ∗ : [Equation]
[Equation]
Simultaneously, F r d ∗ minimizes the least-squares approximation error for the worst-case target function in S ϵ :
[Equation]
Proof.
We will start by deriving Equation 7, and derive Equation 6 afterward.
We can think of Equation 7 as equivalent to the following adversarial game:
Player chooses a dimension- d subspace F ⊂ A → R of functions.
Adversary chooses a function
g ∈
A → R with a fixed level of invariance E [( g ( a 1 ) − g ( a 2 )) 2 ] = 2 g T ( D A − P A,A ) g
= ϵ .
Without loss of generality, we let ϵ = 2 so that g T ( D A − P A,A ) g = 1 ; other values of ϵ will just lead to scaling the function g . 3.
Player chooses the best ˆ g ∈ F to minimize E [(ˆ g ( a ) − g ( a )) 2 ]
We can analyze this game by working backward from the innermost step, step 3.
Given the function class F and adversarially chosen target function g , choosing ˆ g to minimize the expected squared error is equivalent to finding the orthogonal projection of g into F with respect to the measure p
( A ) , e.g. with respect to the weighted L2 norm L (2; p ( A )) .
More precisely, we want
[Equation]
But this is just finding the ˆ g ∈ F which minimizes ∥ D 1 / 2 A ˆ g − D 1 / 2 A g ∥ 2 .
This is given by the orthogonal projection of the vector D
1 / 2 A g into D 1 / 2 A F , under the ordinary L2 norm.
We can now define R as the orthogonal projection operator on D 1 / 2 A F , such that R h ∈
D 1 / 2 A F (e.g. D − 1 / 2 A
R h ∈ F ), and for D 1 / 2 A f ∈ D 1 / 2 A F (e.g. f ∈ F ), we have D 1 / 2 A f = RD 1 /
A f
(e.g. f = D − 1 / 2 A
RD 1 / 2 A f ).
Observe that R is real and symmetric, and has eigenvalue 1 with multiplicity d and all other eigenvalues are 0.
Since R characterizes the subset, we will find it convenient to redefine our objective for the initial player as choosing R , and then letting
F = { D − 1 / 2 A
RD 1 / 2 A g : g ∈
A → R } .
We then have
[Equation]
and the cost is
[Equation]
We next consider step 2.
Given R , what g should the adversary pick?
Letting L =
D
A − P A,A , the adversary is constrained to pick g such that g T L g =
( L 1 / 2 g ) T ( L 1 / 2 g )
= 1 .
We note that L is not full rank: in particular, any eigenvector of M with eigenvalue 1 is an eigenvector of L of eigenvector zero.
Any function g chosen by the adversary must then be the sum of two parts:
• a component in in the range of L , of the form (cid:0) L † (cid:1) 1 / 2 u where ∥ u ∥ 2 = 1 and † represents the Moore-Penrose pseudoinverse, • and a component in the null space of L .
Overall, we can thus write g = (cid:0) L † (cid:1) 1 / 2 u + h where ∥ u ∥ 2 = 1 and h ⊤ L h = 0 .
Similarly, the response ˆ g must also have two components, one in the range of L and one in the null space of L .
There are then two cases.
If F does not span the entire null space of L , the adversary can force an arbitrarily high approximation error by choosing h to be in the null space of L but not F .
On the other hand, if F spans the entire null space of L , the player can always perfectly approximate h , and so the adversary is forced to maximize cost by using u .
In particular, they will pick
[Equation]
where A is the matrix (cid:0) L † (cid:1) 1 / 2 D 1 / 2 A ( I − R )
D 1 / 2 A (cid:0)
L † (cid:1) 1 / 2 .
The optimal choice for u is an eigenvector of A with maximal eigenvalue, and the cost is then that maximal eigenvalue.
But observe that M is similar to the following:
[Equation]
Similar matrices have the same eigenvalues, so the maximum cost attainable by the adversary is the maximal eigenvalue of
A ′ .
Finally, we consider step 1.
Which R should our player choose to minimize this maximum cost?
They should first ensure the cost is finite, by choosing R to span the null space of D − 1 / 2 A
LD − 1 / 2 A .
(Note that if d is less than the dimension of this null space, there is no choice that ensures a finite cost; in this case every representation has unbounded worst-case approximation error.)
Afterward, they should ensure that A ′ has the smallest maximum eigenvalue.
The sorted vector of eigenvalues of A ′ is bounded below by the vector obtained by matching the largest eigenvalues of (cid:16) D − 1 / 2 A
LD − 1 / 2 A (cid:17) † with the smallest of ( I − R )
(e.g. the largest of R ) (Bhatia, 2013, exercise III.6.14) 4 .
Let d ∗ be the dimension of the null space of D − 1 / 2 A
LD − 1 / 2 A .
Then I − R has d zero eigenvalues, of which d ∗ are used to span this null space, and the remaining d − d ∗ (if any) are used to reduce the eigenvalues of A ′ .
Thus, the largest eigenvalue of A ′ is always at least as big as the ( d − d ∗ + 1) -th largest eigenvalue of (cid:16) D − 1 / 2 A
LD − 1 / 2 A (cid:17) † .
We can attain this bound by setting R to exactly capture the ( d − d ∗ ) -dimensional subspace of (cid:16) D − 1 / 2 A
LD − 1 / 2 A (cid:17) † spanned by its top eigenspaces, along with the d ∗ -dimensional null space of D − 1 / 2 A
LD − 1 / 2 A .
But the combination of the null space of D − 1 / 2 A
LD − 1 / 2 A and the top eigenspace of (cid:16) D − 1 / 2 A
LD − 1 / 2 A (cid:17) † is just the space spanned by the d eigenvectors of D − 1 / 2 A
LD − 1 / 2 A with the smallest eigenvalues.
Furthermore,
[Equation]
so we are looking for the eigenvectors of M with the largest eigenvalues, where M is the matrix described in Appendix B.
Thus, the player should choose F such that D 1 / 2 A F spans the top d -dimensional eigenspace of M , e.g. they should choose functions of the form D − 1 / 2 A v i where the v i are the eigenvectors of M with largest eigenvalue.
But these are exactly the left eigenvectors f i of the positive pair Markov chain, which is how r d ∗ is defined.
We conclude that F r d ∗ is the optimal choice for the player, and thus Equation 7 holds.
Indeed, we can conclude something further: if λ d +1 is the ( d + 1) th eigenvalue of the positive pair Markov chain (the variance along the ( d + 1) th principal component of the contrastive kernel), then as long as d ≥ d ∗ , (1 − λ d +1 ) − 1 is the ( d − d ∗ + 1) th eigenvalue of (cid:16) D − 1 / 2 A
LD − 1 / 2 A (cid:17) † , which is exactly the worst-case approximation error for F r d ∗ against any function with E
[( g ( v 1 ) − g ( v 2 )) 2 ] = 2 g T Lg = 2 .
Scaling by ε , if E [( g ( v 1 ) − g ( v 2 )) 2 ]
= ε then the worst case error is 12 ε/ (1 − λ d +1 ) .
In other words,
[Equation]
We now return our attention to Equation 6.
This equation can also be formulated as an adversarial game:
Player chooses a rank- d subspace F ⊂ A → R of functions.
Adversary chooses a function ˆ g ∈ F with unit norm E [ˆ g ( v ) 2 ]
= ˆ g T D
A ˆ g = 1 to maximize E
[(ˆ g ( v 1 ) − ˆ g ( v 2 )) 2 ]
= 2ˆ g T L ˆ g .
We can again identify the choice of F with the choice of the orthogonal projection matrix R on D 1 / 2 A F .
We know ˆ g ∈ F , so we can write ˆ g =
D − 1 / 2 A
RD 1 / 2 A ˆ g .
Also note that for any h (not even necessarily in F ), D − 1 / 2 A
RD 1 / 2 A h ∈ F .
Now suppose we choose an h so that E [ h ( a ) 2 ]
=
h T D
A h
= 1 , and define ˆ g = D − 1 / 2 A
RD 1 / 2 A h .
Then
[Equation]
because R has eigenvalues at most 1.
So, the following are equivalent:
• choosing ˆ g ∈ F with E [ˆ g ( v ) 2 ] ≤• choosing h ∈ R A with E [ h ( v ) 2 ] ≤ 1 and letting ˆ g = D − 1 / 2 A
[Equation]
We also note that there is no advantage to the adversary from picking a function such that E [ˆ g ( v ) 2 ] < 1 .
So we can reframe step 2 as choosing h so that E [ h ( v ) 2 ] = 1 , to maximize
[Equation]
We can further reparameterize by letting h
= D − 1 / 2 A u , so that E [ h ( v ) 2 ]
= 1 is equivalent to ∥ u ∥
= 1 .
We then have cost
[Equation]
The choice that maximizes the cost is then an eigenvector of
[Equation]
with maximal eigenvalue [Equation] and the cost is 2 times the maximum eigenvalue of B .
But note that B has the same eigenvalues as
[Equation]
since R 2 = R .
And B ′ is similar to
[Equation]
so B and B ′′ have the same eigenvalues.
We now consider step 1.
What should the player choose for R ?
By a similar eigenvalue-of-product argument as used for
Equation 7, regardless of the choice of R the largest eigenvalue of B ′′ must always be at least as big as the d th smallest eigenvalue of D − 1 / 2 A
LD − 1 / 2 A , because R has eigenvalue 1 with multiplicity d .
We can attain this minimum cost by choosing
R to project into the eigenspace spanned by the d eigenvectors of D − 1 / 2 A
[Equation]
with the smallest eigenvalues.
But observe that the smallest eigenvalues and corresponding eigenvectors of D − 1 / 2 A
[Equation]
− M are exactly the largest eigenvalues and corresponding eigenvectors of M = D − 1 / 2 A
[Equation]
= I − L , which as we argued above, is exactly the set of eigenfunctions
f i used to construct r d ∗ .
In this case, the optimal cost itself is determined by the largest eigenvalue of D − 1 / 2 A
[Equation]
(times two), so we obtain
[Equation]
C.2. Generalization bound for linear prediction with the eigenfunction representation
Proposition 4.2.
Let ( A i , Y i ) ni
=1 be i.i.d.
samples, choose R ≥ 0 , and consider the constrained empirical risk minimizer ˆ β R ∈ argmin
∥ β ∥ 2 ≤ R n − 1 (cid:80) ni
=1 |⟨ β, r ∗ d ( A i ) ⟩ − Y i | .
Then the expected excess risk of ˆ β R is bounded by:
[Equation]
Proof.
We start by decomposing the excess risk as:
[Equation]
The first term is the estimation error, which we can readily bound by first noting that:
[Equation]
then noticing that our loss is 1 -Lipschitz, and finishing with the standard Rademacher complexity argument for constrained linear classes (Kakade et al., 2008) to get:
[Equation]
The second term is an approximation error term due to the use of a constrained linear class instead of the full linear class.
Define ˜ β ∗ := β ∗ max {∥ β ∗ ∥ 2 /R, 1 } .
Then we can bound this second term by:
[Equation]
where the second inequality follows from the 1 -Lipschitznes of the loss, the third by Cauchy-Schwartz inequality, and the last by Jensen’s inequality.
The third and last term is an approximation error term due to the use of a function class which is contained in the span of the first d eigenfunctions ( f i ) di =1 .
We can bound it as follows:
[Equation]
where the first inequality follows from the 1 -Lipschitznes of the loss, the second from Jensen’s inequality, and the last by first noticing that the function h ( a ) := (cid:10) r d ∗ ( a ) , β ∗ (cid:11) satisfies h = argmin g ∈F rd ∗ E (cid:2) ( g ( A ) − g ∗ ( A )) 2 (cid:3) (since it is the projection of g ∗ onto F r d ∗ under the norm
∥ x ∥ 2 := E (cid:2) x 2 ( A ) (cid:3) ), then appealing to the proof of Proposition 4.1.
Combining the bounds of equations (8), (9), and (10), we obtain the stated generalization bound.
D. Experimental details D.1. Task: MNIST with Multinomial Augmentations
Our goal in designing our task was to construct distributions p ( X ) and p ( A | X ) such that the exact contrastive kernel could be computed, and so that the Markov chain would mix between different unperturbed dataset examples x ∈ X without changing the label too frequently.
To this end, we constructed our task as follows:
• Define X to be a particular subset of the MNIST dataset, and choose p ( X ) to be the uniform distribution over X . We consider two choices for X : randomly selecting 512 images from each of the ten digit classes (used for comparisions between models), and randomly selecting 1024 images from the digits 0, 1, and 2 (used for visualizations of the eigenfunctions).
• For each image, generate 64 pertubed copies, by randomly blurring, translating, and rotating digits by a small amount.
Add a small constant to each pixel so that no pixel has value zero, then normalize each such copy so that its pixel values sum to 1.
• To generate an augmentation of an image x ∈ X according to p ( A | X = x ) , first choose one of the 64 copies of x , then sample a set of k pixel locations with replacement from the distribution represented by that copy, where k determines the augmentation strength.
This means we are more likely to sample pixel locations which were within the original digit, although due to the perturbations described above the pixels may be scattered around the digit.
Our set A is thus the set of all 28 × 28 images for which all pixels have a nonnegative integer value, and the total of all pixels is k .
(Due to the low pixel density, to improve visibility in our figures we render each pixel as a box 5x its original size, with shading indicating overlap of these boxes.
However, when giving input to the model, we directly input this sparse pixel reprsentation, without the 5x multiplier.)
Given a particular augmented example a , we can compute p ( a | x ) for any x ∈ X by summing over each of the 64 copies of x and using the closed-form PDF of a multinomial distribution.
We can then compute p ( x | a ) by normalizing over all possible x , and use this to exactly compute the contrastive kernel feature map ϕ Ctr .
We selected the perturbation parameters such that there was nontrivial uncertainty in x given each a ; in other words, we made sure the positive pair Markov chain mixed well between examples.
D.2. Ten-class MNIST models
For our main experimental results on all ten MNIST digit classes, we used three-block variants of a ResNet-18 model followed by a two 128-dimensional fully-connected layers and a final output layer.
• Linear kernel: We used kernel parameterization (cid:98) K θ ( a 1 , a 2 )
=
h θ ( a 1 ) T h θ ( a 2 ) and the spectral loss, with h θ as the output of the final layer. We set the dimension of the final layer to 512. • Temperature-normalized exponential kernel: We used kernel parameterization exp ( h θ ( a 1 ) T h θ ( a 2 ) /τ ) , where h θ is computed by normalizing the output of the final layer to lie on the unit hypersphere.
We set the dimension of the final layer to 32.
For the loss function, we used a weighted combination of 0.9 times the NT-XEnt loss and 0.1 times the NT-Logistic loss.
(Including the NT-Logistic loss ensures that there is a unique function that minimizes the total loss, without requiring us to estimate proportionality constants.)
We optimized the temperature τ jointly with the model parameters.
• Temperature-normalized exponential kernel with scale: We used kernel parameterization exp
( h θ ( a 1 ) T h θ ( a 2 ) /τ ) · s θ ( a 1 ) s θ ( a 2 ) .
We set the dimension of the final layer to 33, and defined h θ by taking the first 32 entries and normalizing them to lie on the unit hypersphere.
We then defined s θ to be exp(5 × tanh( x )) where x is the 33rd entry of the final layer.
We again optimized the temperature τ jointly with the model parameters, using a weighted combination of 0.9 times the NT-XEnt loss and 0.1 times the NT-Logistic loss.
The scale parameter allows the model to adjust the magnitude of the kernel, which can be used to scale the eigenvalues of the principal components or to correct for differences in likelihood between points. We trained our models using the Adam optimizer over 50,000 training iterations and a batch size of 4096 positive pairs per iteration, implemented using the JAX and FLAX libraries (Bradbury et al., 2018; Heek et al., 2020). We used batch normalization for the first 35,000 iterations, then interpolated between the current batch statistics and the average from previous batches for 2,000 more iterations, and finally trained for 13,000 iterations using frozen batch norm statistics alone (e.g. in “inference” mode); we found that this increased the quality of the extracted principal components.
Figures 1 and 3, we trained a ResNet-18 model on only the MNIST digits 0, 1, and 2, using a scaled two-dimensional rational quadratic kernel: [Equation] Here f θ : A → R 2 embeds inputs into the plane, s θ : A → R + is a learned scale function, and α is a learned shape parameter.
We set the output dimension of the ResNet-18 model to 3, and took the first two elements as f θ ; s θ was defined as exp(5 × tanh( x )) where x is the third element.
We also parameterize α = exp( γ ) and learn the scalar parameter γ .
The model has a base hidden dimension of 128.
The model was trained for 50,000 training iterations. We used the cross entropy InfoNCE loss (as described in Appendix A.1) and the Adam optimizer, with a batch size of 4096 positive pairs per iteration.
D.4. Extraction and analysis of principal components
To estimate the eigenfunctions of the true kernel, we performed PCA using the explicit form ϕ Ctr of the contrastive kernel feature map, where the population covariance was approximated by averaging over 256 augmentations for each of the images in X .
We then constructed the principal component projection functions using that covariance.
For our approximate kernels (cid:98) K θ , we first constructed an approximation of the feature map for (cid:98) K θ using the Nyström method (Williams and Seeger, 2000): we sampled a subset S of augmentations by randomly selecting 25% of X and sampling one augmentation for each image, computed the Gram matrix (cid:98) K θ ( S, S ) for that subset, then set
[Equation]
) where (cid:98) K θ ( S, a ) is the vector of similarities of a to each reference augmentation in S .
The result is a feature map such that ˆ ϕ ( a 1 )
⊤ ˆ ϕ ( a 2 ) ≈ (cid:98) K
θ ( a 1 , a 2 ) .
We then again performed PCA using this feature map, using 256 samples per example in X to compute the covariance.
We normalized all principal component projection functions to have unit uncentered variance, e.g. E [ f i ( a ) 2 ]
= 1 and E [ ˆ f i ( a ) 2 ] = 1 . We then computed alignments by taking the squared covariance E [ f i ( a ) ˆ f j ( a )] 2 . Note that by Theorem 3.
this is equivalent to the square of the coefficient c i for the function ˆ f j expanded in terms of the {+basis of eigenfunctions f i ; consequently we have (cid:80) i E [ f i ( a ) ˆ f j ( a )] 
 =+}
E
[ ˆ f j ( a ) 2 ] = 1 . (Note that the sign of a principal component is not uniquely determined, so squaring the covariance ensures that f and − f are considered to be aligned with each other.) We estimated the positive-pair discrepancy for each principal component function by taking the sample average of (cid:0) f i ( a 1 ) − f i ( a 2 ) (cid:1) 2 over 16 randomly sampled augmentation pairs for each image in X .
D.5. Downstream supervised learning
To generate our supervised training and validation sets, we sampled one augmentation of 16 random images from each digit class, labeled with the original label, a total of 160 labeled augmentations in each set.
For the test set, we tookdistinct images from each digit class and generated one augmentation from each image, without overlap with the training or validation sets, for a total of 1700 labeled augmentations.
For the classification task, we fit a logistic regression classifier on the training set using SciKit Learn. For PCA representations, we swept over 40 logarithmically-spaced L2 regularization strengths from 10 − 4 to 10 1 , and also swept over representation dimension d , taking the first d principal components for d between 1 and 256. For ResNet embedding representations, we swept over 150 logarithmically-spaced L2 regularization strengths from 10 − 4 to 10 10 ; we found that higher regularization strengths were necessary to attain a good solution.
We selected the hyperparameters based on which setting gave the highest top-1 accuracy on the validation set. For the regression task, we used a centered version of the one-hot indicator vector, e.g. the target vector for an example from digit 2 was [ − 0 . 1 , − 0 .
9 , − 0 . 1 , − 0 . 1 , − 0 . 1 , − 0 .
1 , − 0 . 1 , − 0 . 1 , − 0 . The purpose of this centering was to ensure that the expected value of the label vector was the zero vector.
We then fit a predictor using ridge regression (L2-regularized least-squares regression) in Numpy. As for the classification task, for PCA representations we swept over 40 logarithmically-spaced L2 regularization strengths from 10 − 4 to 10 1 and over each representation dimension d between 1 and 256, and for ResNet embedding representations we swept over 150 L regularization strengths from 10 − 4 to 10 10 .
We selected the hyperparameters based on which setting gave the lowest squared error on the validation set. D.6. Alignment and discrepancy relationships across regularization strengths Figure 4 in the main paper shows the alignment and discrepancy relationships for the ten-digit MNIST task for k =sampled pixels per digit. In this section we show the corresponding plots for k = 10 (stronger augmentations) and k =(weaker augmentations.)
E. Additional visualizations of three-digit MNIST task
In this section, we provide a few additional visualizations of the three-digit variant of the MNIST task, the contrastive kernel for that task, and our learned rational quadratic model.