# Understanding Augmentation-based
Self-Supervised Representation Learning via RKHS Approximation and Regression

Runtian Zhai, Bingbin Liu, Andrej Risteski, Zico Kolter, Pradeep Ravikumar
Carnegie Mellon University
<EMAIL_ADDRESS>
###### Abstract

Data augmentation is critical to the empirical success of modern self-
supervised representation learning, such as contrastive learning and masked
language modeling. However, a theoretical understanding of the exact role of
augmentation remains limited. Recent work has built the connection between
self-supervised learning and the approximation of the top eigenspace of a
graph Laplacian operator, suggesting that learning a linear probe atop such
representation can be connected to RKHS regression. Building on this insight,
this work delves into a statistical analysis of augmentation-based
pretraining. Starting from the isometry property, a geometric characterization
of the target function given by the augmentation, we disentangle the effects
of the model and the augmentation, and prove two generalization bounds that
are free of model complexity. Our first bound works for an arbitrary encoder,
where the prediction error is decomposed as the sum of an estimation error
incurred by fitting a linear probe with RKHS regression, and an approximation
error entailed by RKHS approximation. Our second bound specifically addresses
the case where the encoder is near-optimal, that is it approximates the
top-$d$ eigenspace of the RKHS induced by the augmentation. A key ingredient
in our analysis is the augmentation complexity, which we use to quantitatively
compare different augmentations and analyze their impact on downstream
performance.

##  1 Introduction

It is widely acknowledged that better data augmentation techniques have been a
major driving force in many recent breakthroughs in self-supervised
representation learning. For example, contrastive learning (Chen et al., 2020)
with aggressive random cropping and strong color distortion greatly improves
the performance of vision tasks, and masked prediction with random masking and
replacement is among the state-of-the-art representation learning methods for
both natural language processing (Devlin et al., 2019) and vision (He et al.,
2022). Due to their extraordinary empirical successes, developing a rigorous
theoretical understanding of augmentation-based representation learning is an
important open problem, and is crucial for inventing new, better and
principled augmentation techniques.

A key advance was recently made by HaoChen et al. (2021), who analyzed a
variant of contrastive representation learning, termed spectral contrastive
learning, by connecting it to learning eigenfunctions of a Laplacian operator
over a population augmentation graph. This has been followed up by multiple
works which extended these results to more general contrastive learning
approaches (Saunshi et al., 2022; Johnson et al., 2023; Cabannes et al.,
2023). However, the generalization guarantees in these works have to
explicitly grapple with the function class being used to learn the
eigenfunctions (typically deep neural networks in practice). Such a dependency
was even deemed necessary by Saunshi et al. (2022) who argued that “function-
class-agnostic analysis leads to vacuous guarantees.” Consequently, the
generalization bounds such as HaoChen et al. (2021, Theorem 4.2) need to
depend on the Rademacher or other complexities of this encoder function class,
yet finding proper complexity measures for flexible function classes such as
neural networks remains an open problem. Moreover, with the effects of the
augmentation and the encoder function class intertwined, these analyses cannot
discern the exact role of augmentation in self-supervised representation
learning.

In this work, we factor out the effect of the encoder function class
completely. Our starting point is the observation that target functions that
are “soft invariant” to the augmentation lie in a specific RKHS
${{\mathcal{H}}_{\Gamma}}$ which we call the (augmentation) induced RKHS.
Meanwhile, the $d$-dimensional feature encoder we pretrain also spans an RKHS,
and the closer this RKHS is to the induced RKHS, the better the encoder. Thus,
pretraining can be viewed as approximating the induced RKHS with unlabeled
data. Then, learning a linear probe atop corresponds to regression on the
encoder’s RKHS with labeled data. This perspective elucidates what roles the
two stages of representation learning are playing:

  * •

Upstream: The self-supervised pretraining stage can be viewed as doing RKHS
approximation. It learns, with unlabeled data, an encoder $\hat{\Psi}$ whose
learned RKHS ${\mathcal{H}}_{{\hat{\Psi}}}$ (Eqn. 9) approximates
${{\mathcal{H}}_{\Gamma}}$.

  * •

Downstream: The supervised stage is doing RKHS regression over
${\mathcal{H}}_{{\hat{\Psi}}}$ with labeled data.

Figure 1 provides an illustration of this RKHS approximation/regression
perspective. The learning performance is measured by the prediction error of
the downstream predictor, which decomposes into two parts based on the above
framework: the approximation error upstream entailed by RKHS approximation,
and the estimation error downstream entailed by RKHS regression.

Thus, using classical function analytic tools, this leads us to our first set
of results: For an arbitrary pretrained encoder, we provide a generalization
bound on the $L^{2}$ distance between the predictor and the target function
(Theorem 1). We emphasize that the encoder can be arbitrary, i.e. any
architecture and size. As a result, Theorem 1 starkly contrasts prior work in
that it disentangles the effects of the model and the augmentation: our
generalization bound (a) is nonparametric, (b) does not depend on any model
complexity or model inductive bias, and hence (c) allows the user to choose
any model class for the encoder. Our bound depends on two key quantities: (a)
the augmentation complexity $\kappa$, which is smaller for stronger and
“better behaved” augmentations; (b) the trace gap, which serves as a proxy of
how well the encoder approximates the induced RKHS.

The bound for an arbitrary encoder is great, but one might also be interested
in proving guarantees for the optimal encoder that minimizes the worst-case
approximation error. We will show that the optimal $d$-dimensional encoder
consists of the top-$d$ eigenfunctions of the induced RKHS (4). However, these
eigenfunctions are not accessible given only finite samples, so we instead
study the near-optimal $d$-dimensional encoder, which is a Monte-Carlo
approximation of the optimal one. Our second set of results provide a
generalization bound for the near-optimal encoder, by bounding its trace gap
(which captures its quality) with high probability in terms of some spectral
aspects of the induced RKHS. The bound shows that: As the number of unlabeled
samples goes to infinity, the near-optimal encoder will become optimal.

${L^{2}(P_{\mathcal{X}})}$${{\mathcal{H}}_{\Gamma}}$${\hat{{\mathcal{H}}}_{\Gamma}}$${\hat{{\mathcal{H}}}_{d}}$$f^{*}$${f_{\hat{\Psi}}}$$\hat{f}$${L^{2}(P_{\mathcal{X}})}$:
Large space of $L^{2}$ bounded functions, containing the target function
$f^{*}$Induced RKHS ${{\mathcal{H}}_{\Gamma}}$: Functions with soft invariance
to augmentation, $f^{*}\in{{\mathcal{H}}_{\Gamma}}$Empirical RKHS
${\hat{{\mathcal{H}}}_{\Gamma}}$ : Obtained with $N$ samples to approximate
inaccessible ${{\mathcal{H}}_{\Gamma}}$Empirical top-$d$ RKHS
${\hat{{\mathcal{H}}}_{d}}$ : Induced by the pretrained $d$-dimensional
encoder Approximation error $\|f^{*}-{f_{\hat{\Psi}}}\|$ : Distance from
$f^{*}$ to ${\hat{{\mathcal{H}}}_{d}}$ (3, Theorem 2) Estimation error
$\|{f_{\hat{\Psi}}}-\hat{f}\|$: Entailed by downstream regression (2)  Figure
1: Overall RKHS approximation/regression framework illustration and
commentary.

A practical implication of our framework is that the _augmentation complexity_
$\kappa$, defined as the $L^{\infty}$ norm of the kernel of
${{\mathcal{H}}_{\Gamma}}$, can be used as a tool to quantitatively compare
different augmentations and analyze their impact on downstream performance. In
Section 5, we demonstrate this point with mask-type augmentations on synthetic
and real datasets, and show that (i) $\kappa$ depends on both the augmentation
strength and the augmentation strategy; (ii) a smaller $\kappa$ (e.g. stronger
augmentation) leads to a smaller generalization gap, but an overly strong
augmentation causes poor training performance. Thus, there is a “sweet spot”
in the middle with the best test performance. Overall, we believe that this
work places modern representation learning on a more scientific footing and
can inspire future empirical advancements. Finally, it is worth emphasizing
that this work is about studying self-supervised learning from a kernel
perspective, rather than studying kernel methods.

##  2 Problem Setup, Geometry of the Augmentation Induced RKHS

Let ${\mathcal{X}}\subset\mathbb{R}^{d_{\mathcal{X}}}$ be the data space, and
${P_{\mathcal{X}}}$ be the data distribution. Let ${L^{2}(P_{\mathcal{X}})}$
be the $L^{2}$ function space, such that any $f\in{L^{2}(P_{\mathcal{X}})}$
satisfies $\mathbb{E}_{{P_{\mathcal{X}}}}[f(X)^{2}]<\infty$. Let $\langle
f_{1},f_{2}\rangle_{P_{\mathcal{X}}}=\int
f_{1}(x)f_{2}(x)d{P_{\mathcal{X}}}(x)$ and $\|f\|_{P_{\mathcal{X}}}=\langle
f,f\rangle_{P_{\mathcal{X}}}^{1/2}$ be the inner product and the norm of the
Hilbert space ${L^{2}(P_{\mathcal{X}})}$. Let
$f^{*}\in{L^{2}(P_{\mathcal{X}})}$ be the target function we want to learn.
This work studies the least-squares regression problem (see Györfi et al.
(2002) for an introduction), which is formally stated as follows:

Problem. Given unlabeled samples $x_{1},\cdots,x_{N}$ and labeled samples
$\tilde{x}_{1},\cdots,\tilde{x}_{n}$ i.i.d. sampled from ${P_{\mathcal{X}}}$,
and labels $\tilde{y}_{k}=f^{*}(\tilde{x}_{k})+\nu_{k}$ for $k\in[n]$ and
random noise $\nu_{k}$, find a predictor $\hat{f}\in{L^{2}(P_{\mathcal{X}})}$
with a low prediction error
$\textnormal{err}(\hat{f},f^{*}):=\|\hat{f}-f^{*}\|_{P_{\mathcal{X}}}^{2}=\mathbb{E}_{{P_{\mathcal{X}}}}[(\hat{f}(X)-f^{*}(X))^{2}]$.

###  2.1 The Augmentation Operator $\Gamma$: A Source of Prior Knowledge

The core problem in self-supervised representation learning is how to leverage
the unlabeled data $x_{1},\cdots,x_{N}$. Without labels, the pretraining task
must be built upon some prior knowledge we have about the target function
$f^{*}$. Data augmentation based pretraining tasks are built upon the
following prior knowledge: Two augmentations of the same sample $x$ should be
similar to each other.

Formally, let ${\mathcal{A}}$ denote an augmented space that contains the
augmentations of $x\in{\mathcal{X}}$. Data augmentation induces a joint
distribution ${P_{{\mathcal{A}}{\mathcal{X}}}}$, with marginal distributions
${P_{\mathcal{A}}}$ and ${P_{\mathcal{X}}}$. We will use capital letters $A$
and $X$ to denote random variables from ${P_{\mathcal{A}}}$ and
${P_{\mathcal{X}}}$. We define an augmentation operator
$\Gamma={\Gamma_{x\rightarrow
a}}:{L^{2}(P_{\mathcal{X}})}\rightarrow{L^{2}(P_{\mathcal{A}})}$ as
$({\Gamma_{x\rightarrow a}}f)(a)=\mathbb{E}[f(X)|a]$ for all
$f\in{L^{2}(P_{\mathcal{X}})}$, and define its adjoint operator
$\Gamma^{*}={\Gamma_{a\rightarrow
x}}:{L^{2}(P_{\mathcal{A}})}\rightarrow{L^{2}(P_{\mathcal{X}})}$ as
$({\Gamma_{a\rightarrow x}}g)(x)=\mathbb{E}[g(A)|x]$ for all
$g\in{L^{2}(P_{\mathcal{A}})}$. We can see that ${\Gamma_{a\rightarrow x}}$ is
indeed the adjoint of ${\Gamma_{x\rightarrow a}}$, since for all
$f\in{L^{2}(P_{\mathcal{X}})}$ and $g\in{L^{2}(P_{\mathcal{A}})}$, there is
$\langle{\Gamma_{x\rightarrow a}}f,g\rangle_{P_{\mathcal{A}}}=\iint
f(x)g(a)p(a,x)dadx=\langle f,{\Gamma_{a\rightarrow
x}}g\rangle_{P_{\mathcal{X}}}$. Intuitively, the augmentation operator
connects any single sample $x$ to the distribution of its augmentations
$p(a|x)$.

Example: In BERT, suppose we use 15% random masking. Then, ${\mathcal{X}}$ is
the space of original sentences, and ${\mathcal{A}}$ is the space of 15%
masked sentences; ${P_{\mathcal{X}}}$ is the distribution over original
sentences, $A\sim p(\cdot|x)$ is the 15% randomly masked version of an
original sentence $x$, and
${P_{{\mathcal{A}}{\mathcal{X}}}}(a,x)={P_{\mathcal{X}}}(x)p(a|x)$. Thus,
$({\Gamma_{a\rightarrow x}}g)(x)$ is the mean of $g$ over all 15% randomly
masked versions of $x$.

Now we define the “soft invariance” w.r.t. the augmentation. Denote the
function class ${\mathcal{F}}(\Gamma;\epsilon)$ to be the set of function
$f^{*}$ such that $\exists g^{*}\in{L^{2}(P_{\mathcal{A}})}$ such that
$f^{*}(x)=({\Gamma_{a\rightarrow x}}g^{*})(x)=E[g^{*}(A)|x]$, and

$\frac{1}{2}\mathbb{E}_{X\sim{P_{\mathcal{X}}}}\mathbb{E}_{A,A^{\prime}\sim
p(\cdot|X)}\left[(g^{*}(A)-g^{*}(A^{\prime}))^{2}\right]\leq\epsilon\|g^{*}\|_{P_{\mathcal{A}}}^{2},$
(1)

for some small positive constant $\epsilon$. Eqn. (1) is a regularization
constraint induced by the random walk normalized Laplacian over the
augmentation graph defined in HaoChen et al. (2021), where the edge weight
between $a$ and $a^{\prime}$ is given by ${P_{A}^{+}}(a,a^{\prime})$ in Eqn.
(2). Compared to Assumption 1.1 in Johnson et al. (2023), Eqn. 1 has an
additional $\|g^{*}\|_{P_{\mathcal{A}}}^{2}$ term on the right-hand side so it
is homogeneous. For any $f\in{\mathcal{F}}(\Gamma;\epsilon)$, if
$\|f\|_{P_{\mathcal{X}}}\leq B$ for some constant $B$, then we denote
$f\in{\mathcal{F}}_{B}(\Gamma;\epsilon)$.

######  Assumption 1 (See discussions in Appendix A).

We assume that $f^{*}\in{\mathcal{F}}_{B}(\Gamma;\epsilon)$ for given
$B,\epsilon$.

Let the range of $\Gamma^{*}$ be
$R(\Gamma^{*})={\left\\{f={\Gamma_{a\rightarrow
x}}g\;\middle|\;g\in{L^{2}(P_{\mathcal{A}})}\right\\}}$. Then $f^{*}\in
R(\Gamma^{*})$ by Assumption 1.

###  2.2 Induced RKHS ${{\mathcal{H}}_{\Gamma}}$ and the Isometry Property

Now, we show that $R(\Gamma^{*})$ is an RKHS. Define the following kernel over
${\mathcal{A}}\times{\mathcal{A}}$:

${K_{A}}(a_{1},a_{2}):=\frac{d{P_{A}^{+}}}{d({P_{\mathcal{A}}}\otimes{P_{\mathcal{A}}})}=\frac{{P_{A}^{+}}(a_{1},a_{2})}{{P_{\mathcal{A}}}(a_{1}){P_{\mathcal{A}}}(a_{2})},\
{P_{A}^{+}}(a_{1},a_{2}):=\int p(a_{1}|x)p(a_{2}|x)d{P_{\mathcal{X}}}(x).$ (2)

Here, ${P_{A}^{+}}$ is called the associative distribution, and ${K_{A}}$ is
the Radon-Nikodym derivative of ${P_{A}^{+}}$ w.r.t.
${P_{\mathcal{A}}}\otimes{P_{\mathcal{A}}}$, also called the positive-pair
kernel in Johnson et al. (2023). Next, we define a dual kernel as:

${K_{X}}(x_{1},x_{2}):=\frac{d{P_{X}^{+}}}{d({P_{\mathcal{X}}}\otimes{P_{\mathcal{X}}})}=\frac{{P_{X}^{+}}(x_{1},x_{2})}{{P_{\mathcal{X}}}(x_{1}){P_{\mathcal{X}}}(x_{2})}=\int\frac{p(a|x_{1})p(a|x_{2})}{{P_{\mathcal{A}}}(a)}da,$
(3)

which is constructed by swapping $a$ and $x$ in ${K_{X}}$ and then applying
the Bayes rule. It holds that:

$\vspace{-.05 in}\begin{cases}({\Gamma_{a\rightarrow x}}{\Gamma_{x\rightarrow
a}}f)(x)=(\Gamma^{*}\Gamma
f)(x)=\int{K_{X}}(x,x^{\prime})f(x^{\prime}){P_{\mathcal{X}}}(x^{\prime})dx^{\prime};\\\
({\Gamma_{x\rightarrow a}}{\Gamma_{a\rightarrow
x}}g)(a)=(\Gamma\Gamma^{*}g)(a)=\int{K_{A}}(a,a^{\prime})g(a^{\prime}){P_{\mathcal{A}}}(a^{\prime})da^{\prime}.\end{cases}$

See App. D for derivation. Let $\lambda_{i}\in\mathbb{R}$ and
$\psi_{i}\in{L^{2}(P_{\mathcal{X}})}$ be the sorted eigenvalue and
eigenfunction of $\Gamma^{*}\Gamma$, i.e.
$\Gamma^{*}\Gamma\psi_{i}=\lambda_{i}\psi_{i}$, with
$\lambda_{1}\geq\lambda_{2}\geq\cdots\geq 0$. Suppose
$\int{K_{X}}(x,x^{\prime})^{2}d{P_{\mathcal{X}}}(x)d{P_{\mathcal{X}}}(x^{\prime})<\infty$,
so that $\Gamma^{*}\Gamma$ is a compact Hilbert-Schmidt integral operator.
Then, by Hilbert-Schmidt theorem, we can choose $\psi_{1},\psi_{2},\cdots$
that form an orthonormal basis of ${L^{2}(P_{\mathcal{X}})}$, such that
$\langle\psi_{i},\psi_{j}\rangle_{P_{\mathcal{X}}}=\delta_{i,j}$, and any
$f\in{L^{2}(P_{\mathcal{X}})}$ can be written as $f=\sum_{i}u_{i}\psi_{i}$ for
some $\\{u_{i}\\}_{i}$.

Note that $\lambda_{i}\in[0,1]$, since
$\|g\|_{P_{\mathcal{A}}}^{2}\geq\|\Gamma^{*}g\|_{P_{\mathcal{X}}}^{2}$ for any
$g\in{L^{2}(P_{\mathcal{A}})}$ by Jensen’s inequality. It is easy to check
that $\lambda_{1}=1$ and $\psi_{1}\equiv 1$ is always a pair of eigenvalue and
eigenfunction of $\Gamma^{*}\Gamma$. Moreover, when $d$ is sufficiently large,
$\lambda_{d}$ should be close to $0$. Similarly, we can define a set of
orthonormal eigenfunctions of $\Gamma\Gamma^{*}$ denoted by $\\{\phi_{i}\\}$,
and we can relate these two sets of eigenfunctions by:

######  Proposition 1 (Duality).

Operators $\Gamma\Gamma^{*}$ and $\Gamma^{*}\Gamma$ share the same non-zero
eigenvalues, and there exist eigenfunctions $\\{\phi_{i}\\}$ of
$\Gamma\Gamma^{*}$ that form an orthonormal basis of
${L^{2}(P_{\mathcal{A}})}$, such that for any $\lambda_{i}>0$,

$\psi_{i}=\lambda_{i}^{-1/2}\Gamma^{*}\phi_{i}=\lambda_{i}^{-1/2}{\Gamma_{a\rightarrow
x}}\phi_{i}\quad\text{and}\quad\phi_{i}=\lambda_{i}^{-1/2}\Gamma\psi_{i}=\lambda_{i}^{-1/2}{\Gamma_{x\rightarrow
a}}\psi_{i}.$

Moreover, we have the following spectral decomposition of the Radon-Nikodym
derivative:

$\frac{dP_{{\mathcal{A}}{\mathcal{X}}}}{d({P_{\mathcal{A}}}\otimes{P_{\mathcal{X}}})}=\frac{P_{{\mathcal{A}}{\mathcal{X}}}(a,x)}{{P_{\mathcal{A}}}(a){P_{\mathcal{X}}}(x)}=\sum_{i}\lambda_{i}^{1/2}\phi_{i}(a)\psi_{i}(x).$
(4)

Define a Hilbert space
${{\mathcal{H}}_{\Gamma}}:={\left\\{f=\sum_{i}u_{i}\psi_{i}\in{L^{2}(P_{\mathcal{X}})}\;\middle|\;\sum_{i}\lambda_{i}^{-1}u_{i}^{2}<\infty\right\\}}$,
whose inner product is $\langle
f_{1},f_{2}\rangle_{{{\mathcal{H}}_{\Gamma}}}=\sum_{i}\lambda_{i}^{-1}u_{i}v_{i}$
for $f_{1}=\sum_{i}u_{i}\psi_{i}$ and $f_{2}=\sum_{i}v_{i}\psi_{i}$. It
satisfies (proof in Appendix D):

  1. (i)

${K_{X}}$ is the reproducing kernel of ${{\mathcal{H}}_{\Gamma}}$, such that
for all $f\in{{\mathcal{H}}_{\Gamma}}$, $f(x)=\langle
f,{K_{X}}(x,\cdot)\rangle_{{{\mathcal{H}}_{\Gamma}}}$.

  2. (ii)

${{\mathcal{H}}_{\Gamma}}=R(\Gamma^{*})$. We call this the augmentation
induced RKHS.

  3. (iii)

${{\mathcal{H}}_{\Gamma}}$ is isometric to
$\textnormal{span}({\left\\{\phi_{i}\right\\}}_{\lambda_{i}>0})$, a subspace
of ${L^{2}(P_{\mathcal{A}})}$, and
$\|f\|_{{{\mathcal{H}}_{\Gamma}}}=\inf_{g:f=\Gamma^{*}g}\|g\|_{P_{\mathcal{A}}}$.

  4. (iv)

For any $f^{*}\in{\mathcal{F}}_{B}(\Gamma;\epsilon)\subset R(\Gamma^{*})$, let
$f^{*}=\sum_{i}u_{i}\psi_{i}$. Define
$g_{0}:=\sum_{i}\lambda_{i}^{-1/2}u_{i}\phi_{i}$. Then, $g_{0}$ must satisfy
Eqn. 1, so we can choose $g^{*}=g_{0}$, in which case Eqn. 1 is equivalent to:

$\vspace{-.05 in}\langle
g^{*},(I-\Gamma\Gamma^{*})g^{*}\rangle_{P_{\mathcal{A}}}\leq\epsilon\|g^{*}\|_{P_{\mathcal{A}}}^{2}\;\Leftrightarrow\;\sum_{i}\frac{1-\lambda_{i}}{\lambda_{i}}u_{i}^{2}\leq\epsilon\sum_{i}\frac{1}{\lambda_{i}}u_{i}^{2}.$
(5)

Moreover, Eqn. (5) is equivalent to a simple and intuitive isometry property:

$(1-\epsilon)\|f^{*}\|_{{{\mathcal{H}}_{\Gamma}}}^{2}\leq\|f^{*}\|_{P_{\mathcal{X}}}^{2}\leq\|f^{*}\|_{{{\mathcal{H}}_{\Gamma}}}^{2}.$
(6)

Understanding the isometry property: This property essentially says that the
operator $\Gamma^{*}\Gamma$ preserves most variance of the target function
$f^{*}$ in the infinite-dimensional functional space. Thus naturally, the
optimal $d$-dimensional encoder that minimizes the approximation error should
keep as much variance as possible under $\Gamma^{*}\Gamma$. This isometry
property can be considered as a counterpart of the restricted isometry
property (RIP) in infinite-dimensional functional spaces: In Section 4, we
will show that the optimal encoder consists of the top-$d$ eigenfunctions of
$\Gamma^{*}\Gamma$. Analogously, for a finite-dimensional vector space, PCA
finds the $d$ principal components that keep the most variance under a linear
transformation $T$, and they consist of the top-$d$ eigenvectors of $T$.

##  3 Learning Guarantees for an Arbitrary Encoder

We now present generalization bounds for an arbitrary encoder. But before
that, we need to clarify which encoder we are talking about. In practice,
people don’t directly pretrain ${\hat{\Psi}}$ during upstream. Instead, the
common practice is to first pretrain an encoder
${\hat{\Phi}}=[\hat{\phi}_{1},\cdots,\hat{\phi}_{d}]$ on the augmented space
with $\hat{\phi}_{i}\in{L^{2}(P_{\mathcal{A}})}$, and then transform it into
${\hat{\Psi}}$, on top of which a linear probe is learned downstream. For
example, in contrastive learning, the encoder is trained on views of images
instead of original images; in masked language modeling, the encoder is
trained on masked sentences instead of full sentences. In practice, people
usually apply ${\hat{\Phi}}$ to downstream tasks directly without
transformation, but for theoretical analysis we need to explicitly write out
the transformation since ${\hat{\Psi}}$ and ${\hat{\Phi}}$ work on different
spaces. We consider the commonly used average encoder (Saunshi et al., 2022,
Eqn. (4)):

${\hat{\Psi}}(x)=\mathbb{E}[{\hat{\Phi}}(A)|x]=\int{\hat{\Phi}}(a)p(a|x)da,\vspace{-.05
in}$ (7)

which is equivalent to ${\hat{\Psi}}=\Gamma^{*}{\hat{\Phi}}$, and thus
$\hat{\psi}_{i}\in R(\Gamma^{*})={{\mathcal{H}}_{\Gamma}}$ for all $i\in[d]$.
In this section, ${\hat{\Phi}}$ can be an arbitrary function, even with
infinite dimensions, so that the model can have any architecture and size;
whereas ${\hat{\Psi}}$ is always the average encoder of ${\hat{\Phi}}$.
Moreover, we consider the following predictor:

######  Definition 1.

The final predictor is the nonparametric least-squares estimate defined as
$\vspace{-.05
in}\hat{f}:=\operatorname*{arg\,min}_{f:f=w^{\top}\hat{\Psi}\in{\mathcal{H}}_{{\hat{\Psi}}},\|f\|_{{{\mathcal{H}}_{\Gamma}}}\leq\frac{B}{\sqrt{1-\epsilon}}}{\left\\{\frac{1}{n}\sum_{k=1}^{n}(\tilde{y}_{k}-f(\tilde{x}_{k}))^{2}\right\\}}.$
(8)

Here, note that by Eqn. (6) and $\|f^{*}\|_{{P_{\mathcal{X}}}}\leq B$, there
is $\|f^{*}\|_{{{\mathcal{H}}_{\Gamma}}}\leq\frac{B}{\sqrt{1-\epsilon}}$.

We next introduce two key ingredients crucial to our analyses. The first
ingredient is what we term the _augmentation complexity_ , which uniformly
bounds the kernel and is typically required in RKHS generalization analyses;
for instance, see Schölkopf & Smola (2002, Section 12.1.3)).

######  Definition 2.

Define the augmentation complexity as $\kappa:=\|K_{X}\|_{\infty}^{1/2}$, i.e.
for ${P_{\mathcal{X}}}$-almost all $x$,
${K_{X}}(x,x)=\sum_{i}\lambda_{i}\psi_{i}(x)^{2}=\int\frac{p(a|x)^{2}}{{P_{\mathcal{A}}}(a)}da=D_{\chi^{2}}({P_{\mathcal{A}}}(\cdot|x)\parallel{P_{\mathcal{A}}})+1\leq\kappa^{2}.\vspace{-.05
in}$

Here, $D_{\chi^{2}}(P\parallel Q):=\int(\frac{dP}{dQ}-1)^{2}dQ$ is the
$\chi^{2}$-divergence. Let $S_{\lambda}(d):=\sum_{i=1}^{d}\lambda_{i}$ and
$S_{\lambda}:=S_{\lambda}(\infty)=\sum_{i=1}^{\infty}\lambda_{i}$. Wang et al.
(2022b) showed that
$S_{\lambda}=D_{\chi^{2}}(P_{{\mathcal{A}}|{\mathcal{X}}}\parallel
P_{\mathcal{A}})+1$. By convexity of $D_{\chi^{2}}$, we have
$S_{\lambda}\leq\kappa^{2}$, i.e. $\Gamma^{*}\Gamma$ is a trace-class operator
(or $\kappa^{2}\geq\int K_{X}(x,x)d{P_{\mathcal{X}}}(x)=S_{\lambda}$). Thus,
$\kappa\geq 1$ as $\lambda_{1}=1$. We will provide examples of this
augmentation complexity in Section 5.

Our second ingredient is the trace gap that captures the quality of the
encoder. It is based on the notion of the ratio trace. Given an encoder
${\hat{\Phi}}$, without loss of generality, suppose it is full-rank.

######  Definition 3.

Define covariance matrices ${\bm{F}},{\bm{G}}$ as
${\bm{F}}(i,j)=\langle\hat{\psi}_{i},\hat{\psi}_{j}\rangle_{P_{\mathcal{X}}}=\langle\Gamma^{*}\hat{\phi}_{i},\Gamma^{*}\hat{\phi}_{j}\rangle_{P_{\mathcal{X}}}$
and
${\bm{G}}(i,j)=\langle\hat{\phi}_{i},\hat{\phi}_{j}\rangle_{P_{\mathcal{A}}}$.
Then, the ratio trace is defined as
$\operatorname{Tr}({\bm{G}}^{-1}{\bm{F}})$, if ${\bm{G}}^{-1}$ is well-
defined.

Ratio trace is a classical quantity in linear discriminant analysis (LDA)
(Wang et al., 2007) and, as we will show, controls the approximation error.
The largest ratio trace of any $d$-dimensional $\hat{\Phi}$ is
$\lambda_{1}+\cdots+\lambda_{d}$, and can be achieved by the top-$d$
eigenspace of ${{\mathcal{H}}_{\Gamma}}$. Then, define the learned kernel as

${\hat{K}_{{\hat{\Psi}}}}(x,x^{\prime})=\langle\Gamma^{*}({\bm{G}}^{-1/2}{\hat{\Phi}})(x),\Gamma^{*}({\bm{G}}^{-1/2}{\hat{\Phi}})(x^{\prime})\rangle,$
(9)

which is the reproducing kernel of
${\mathcal{H}}_{{\hat{\Psi}}}=\textnormal{span}(\hat{\psi}_{1},\hat{\psi}_{2},\cdots)$,
a subspace of ${{\mathcal{H}}_{\Gamma}}$. Here ${\bm{G}}^{-1/2}$ is used for
normalization. The ratio trace can be viewed as the trace of
${\mathcal{H}}_{{\hat{\Psi}}}$. Then, define the trace gap as:

$\tau^{2}:=\inf_{d^{\prime}\leq
d}\inf_{h_{1},\cdots,h_{d^{\prime}}}S_{\lambda}(d^{\prime}+1)-\operatorname{Tr}({\bm{G}}_{h}^{-1}{\bm{F}}_{h}),\vspace{-.05
in}$ (10)

where $\tau\geq 0$, $h_{i}=w_{i}^{\top}\hat{\Phi}$,
${\bm{G}}_{h}=\left(\langle
h_{i},h_{j}\rangle_{P_{\mathcal{A}}}\right)_{i,j\in[d^{\prime}]}$, and
${\bm{F}}_{h}=\left(\langle\Gamma^{*}h_{i},\Gamma^{*}h_{j}\rangle_{P_{\mathcal{X}}}\right)_{i,j\in[d^{\prime}]}$.
Note that for any $d^{\prime}\leq d$ there is
$\operatorname{Tr}({\bm{G}}_{h}^{-1}{\bm{F}}_{h})\leq
S_{\lambda}(d^{\prime})$, so $\tau^{2}$ is always lower bounded by
$\lambda_{d+1}$. And by choosing $h_{i}=\hat{\phi}_{i}$ for $i\in[d]$, we can
see that $\tau^{2}\leq
S_{\lambda}(d+1)-\operatorname{Tr}({\bm{G}}^{-1}{\bm{F}})$.

We now state our first result. For simplicity, we assume that the random noise
$\nu_{1},\cdots,\nu_{n}$ are i.i.d. ${\mathcal{N}}(0,\sigma^{2})$ variates for
some $\sigma>0$, though this can be relaxed (Wainwright, 2019, Chapter 13).

######  Theorem 1.

Let $\nu_{1},\cdots,\nu_{n}$ be i.i.d. ${\mathcal{N}}(0,\sigma^{2})$ variates.
Let $\hat{f}$ be defined by 1. If $\hat{\Phi}$ has $d$ dimensions and $\tau<1$
($d$ can be $\infty$), then there are universal constants $c_{0},c_{1},c_{2}$
such that with probability at least
$1-c_{1}\exp\left(-\frac{c_{2}\sqrt{2nS_{\lambda}(d+1)}}{\kappa}\right)-\exp\left(-\sqrt{\frac{2n\kappa^{2}B^{2}}{1-\epsilon}}\right)$,
there is
$\|\hat{f}-f^{*}\|_{P_{\mathcal{X}}}^{2}\leq\frac{9\tau^{2}(\tau+\epsilon)B^{2}}{(1-\tau^{2})(1-\epsilon)}+\frac{c_{0}\kappa(B^{2}+\sigma
B)}{1-\epsilon}\sqrt{\frac{S_{\lambda}(d+1)}{n}}\quad\text{for all
}f^{*}\in{\mathcal{F}}_{B}(\Gamma;\epsilon).$

######  Remark.

The first term in the bound controls the approximation error entailed by the
limited capacity of the $d$-dimensional encoder $\hat{\Phi}$. This term may
not vanish as $N,n\rightarrow\infty$, since $\tau^{2}$ is lower bounded by
$\lambda_{d+1}$ which can be positive. For instance, if $d$ is too small for
$\hat{\Phi}$ to represent $f^{*}$, then the approximation error cannot be zero
regardless of the number of samples available. Nevertheless, we can show that
with a proper pretraining algorithm, the trace gap $\tau^{2}$ can be very
close to $\lambda_{d+1}$ as $N\rightarrow\infty$ (Theorem 2). A larger ratio
trace leads to a smaller approximation error, and indeed many existing
objectives are strongly connected to maximizing
$\operatorname{Tr}({\bm{G}}^{-1}{\bm{F}})$ (see Appendix C). The second term
bounds the downstream estimation error, which vanishes as
$n\rightarrow\infty$. Finally, this result requires full access to $p(a|x)$
for any $x$, which is available in theory since the augmentation scheme is of
our choice. Practical considerations for estimating $p(a|x)$ is outside the
scope of this work.

➢ Proof Sketch of Theorem 1: This result follows from two bounds:

######  Lemma 2 (Estimation error bound).

Suppose $\nu_{1},\cdots,\nu_{n}$ are i.i.d. ${\mathcal{N}}(0,\sigma^{2})$
variates. If $\hat{\Phi}$ has $d$ dimensions ($d$ can be $\infty$), then we
have the following uniform bound over all
$f^{*}=\Gamma^{*}g^{*}\in{\mathcal{F}}_{B}(\Gamma;\epsilon)$:

$\displaystyle\underset{\tilde{x}_{i},\nu_{i}}{{\mathbb{P}}}\left[\forall
f^{*}\in{\mathcal{F}}_{B}(\Gamma;\epsilon),\|\hat{f}-f^{*}\|_{P_{\mathcal{X}}}^{2}\leq
9\|{f_{\hat{\Psi}}}-f^{*}\|_{P_{\mathcal{X}}}^{2}+\frac{c_{0}\kappa(B^{2}+\sigma
B)}{1-\epsilon}\sqrt{\frac{S_{\lambda}(d+1)}{n}}\right]$ $\displaystyle\geq$
$\displaystyle
1-c_{1}\exp\left(-\frac{c_{2}\sqrt{2nS_{\lambda}(d+1)}}{\kappa}\right)-\exp\left(-\sqrt{\frac{2n\kappa^{2}B^{2}}{1-\epsilon}}\right),\vspace{-.05
in}$

where ${f_{\hat{\Psi}}}=\Gamma^{*}(\Pi_{{\hat{\Phi}}}g^{*})$ is the projection
of $f^{*}$ onto ${\mathcal{H}}_{{\hat{\Psi}}}$ w.r.t.
$\langle\cdot,\cdot\rangle_{{\mathcal{H}}_{\Gamma}}$, and $c_{0},c_{1},c_{2}$
are universal constants. Note that $\hat{f}$ depends on $f^{*}$. Moreover,
$S_{\lambda}(d+1)\leq\min{\left\\{d+1,\kappa^{2}\right\\}}$.

######  Remark.

The constant $9$ in the bound is loose and can be tightened arbitrarily close
to 1 (at the cost of increasing other terms in the bound) by straightforward
modifications to the proof, which we omit here. Such a greater-than-one
constant is inevitable for any uniform deviation bound.

######  Lemma 3 (Approximation error, upper bound).

If $\tau<1$, then for any $f^{*}\in{\mathcal{F}}_{B}(\Gamma;\epsilon)$, there
is

$\|{f_{\hat{\Psi}}}-f^{*}\|_{P_{\mathcal{X}}}^{2}\leq\frac{\tau^{2}}{1-\tau^{2}}\frac{\tau+\epsilon}{1-\epsilon}B^{2}.$
(11)

##  4 Analyses for the Near-Optimal Encoder

In this section, we consider the special case where $\hat{\Psi}$ is near-
optimal, i.e. the Monte-Carlo approximation of the optimal $d$-dimensional
encoder that minimizes the worst-case approximation error. First, we show that
the optimal encoder spans the top-$d$ eigenspace, which is the linear span of
the top-$d$ eigenfunctions $\psi_{1},\cdots,\psi_{d}$. Define the worst-case
approximation error over ${\mathcal{F}}_{B}(\Gamma;\epsilon)$ as:

$\textnormal{err}({\hat{\Psi}};{\mathcal{F}}_{B}(\Gamma;\epsilon)):=\sup_{f\in{\mathcal{F}}_{B}(\Gamma;\epsilon)}\min_{w\in\mathbb{R}^{d}}\textnormal{err}(w^{\top}{\hat{\Psi}},f)=\sup_{f\in{\mathcal{F}}_{B}(\Gamma;\epsilon)}\min_{w\in\mathbb{R}^{d}}\|w^{\top}{\hat{\Psi}}-f\|_{P_{\mathcal{X}}}^{2}.$

######  Proposition 4 (Approximation error, lower bound).

For any ${\hat{\Psi}}=[\hat{\psi}_{1},\cdots,\hat{\psi}_{d}]$ where
$\hat{\psi}_{i}\in{L^{2}(P_{\mathcal{X}})}$,

$\textnormal{err}({\hat{\Psi}};{\mathcal{F}}_{B}(\Gamma;\epsilon))\geq\frac{\lambda_{d+1}}{1-\lambda_{d+1}}\frac{\epsilon}{1-\epsilon}B^{2}\quad\text{given
that}\quad\frac{\lambda_{d+1}}{1-\lambda_{d+1}}\frac{\epsilon}{1-\epsilon}\leq\frac{1}{2}.$

To attain equality, it is sufficient for ${\hat{\Psi}}$ to span the top-$d$
eigenspace, and also necessary if $\lambda_{d+1}<\lambda_{d}$.

One can use kernel PCA to recover top-$d$ eigenspace (Belkin & Niyogi, 2003;
Johnson et al., 2023), which is however not scalable. Alternatively, one can
use variational objectives that attain minimum when $\hat{\Phi}$ spans the
top-$d$ eigenspace (see Appendix C for examples). Also, when $\lambda_{d}>0$,
the top-$d$ eigenspace lies within $R(\Gamma^{*})$, which is the function
class of $\Psi$ when using the average encoder.

The optimal $d$-dimensional encoder is inaccessible since
$\psi_{1},\cdots,\psi_{d}$ cannot be extracted with only finite samples. What
one can do instead is extracting the empirical top-$d$ eigenspace
${\hat{{\mathcal{H}}}_{d}}$. Given samples $x_{1},\cdots,x_{N}$, define two
Hilbert spaces ${L^{2}({\hat{P}_{\mathcal{X}}})}$ and
${L^{2}({\hat{P}_{\mathcal{A}}})}$, where $\langle
f_{1},f_{2}\rangle_{\hat{P}_{\mathcal{X}}}=\frac{1}{N}\sum_{k=1}^{N}f_{1}(x_{k})f_{2}(x_{k})$,
and $\langle g_{1},g_{2}\rangle_{\hat{P}_{\mathcal{A}}}=\int
g_{1}(a)g_{2}(a)d{\hat{P}_{\mathcal{A}}}(a)$ for
${\hat{P}_{\mathcal{A}}}(a)=\frac{1}{N}\sum_{k=1}^{N}p(a|x_{k})$. With these
finite samples, $\Gamma^{*}$ will remain the same, but $\Gamma$ will become
the empirical version $\bar{\Gamma}$:

$(\bar{\Gamma}f)(a)=\frac{1}{N}\sum_{k=1}^{N}\frac{f(x_{k})p(a|x_{k})}{{\hat{P}_{\mathcal{A}}}(a)}.$

Let the eigenvalues and eigenfunctions of $\Gamma^{*}\bar{\Gamma}$ be
$\\{(\bar{\lambda}_{i},\bar{\psi}_{i})\\}$. Let $\bar{\phi}_{i}$ be the
eigenfunctions of $\bar{\Gamma}\Gamma^{*}$. $\bar{\phi}_{i}$ and
$\bar{\psi}_{i}$ have the duality property similar to Proposition 1. Let
${\hat{{\mathcal{H}}}_{\Gamma}}$ be the RKHS associated with
$\Gamma^{*}\bar{\Gamma}$, which is the span of
${\left\\{\bar{\psi}_{i}\right\\}}_{\bar{\lambda}_{i}>0}$ and hence is a
subspace of ${{\mathcal{H}}_{\Gamma}}=R(\Gamma^{*})$. Let
${\hat{{\mathcal{H}}}_{d}}$ be the top-$d$ eigenspace of
${\hat{{\mathcal{H}}}_{\Gamma}}$. Let $\lambda_{\max}({\bm{G}})$,
$\lambda_{\min}({\bm{G}})$ be the largest and smallest eigenvalue of
${\bm{G}}$. For the case of interest where
${\mathcal{H}}_{{\hat{\Psi}}}={\hat{{\mathcal{H}}}_{d}}$, we have our second
main learning guarantee that bounds the trace gap:

######  Theorem 2.

Suppose $\hat{\phi}_{i}=\bar{\phi}_{i}$ for $i\in[d]$. Let
$\gamma_{{\bm{G}}}:=\lambda_{\max}({\bm{G}})/\lambda_{\min}({\bm{G}})$ be the
condition number of ${\bm{G}}$. Then, for any $\delta>0$, it holds with
probability at least $1-\delta$ that $\tau^{2}\leq
S_{\lambda}(d+1)-\operatorname{Tr}({\bm{G}}^{-1}{\bm{F}})\leq\lambda_{d+1}+\left(2+\sqrt{2\log\frac{2}{\delta}}\right)\frac{(\lambda_{d}^{-1}+\bar{\lambda}_{d}^{-1}\gamma_{{\bm{G}}}^{1/2}+2)\kappa^{2}}{\sqrt{N}}d.$

######  Remark.

Combining this result with the first main Theorem 1 leads to the bound for
${\mathcal{H}}_{{\hat{\Psi}}}={\hat{{\mathcal{H}}}_{d}}$. Since
$\tau^{2}\geq\lambda_{d+1}$, this result says that the gap between $\tau^{2}$
and its optimal value is $O(N^{-1/2})$, ignoring potential dependence of
$\gamma_{{\bm{G}}}$ and $\bar{\lambda}_{d}^{-1}$ on $N$. Comparing the upper
bound in 3 \+ Theorem 2 to the lower bound in 4, we note that the upper bound
is near tight: The only difference is that in Eqn. 11, we have
$\frac{\tau+\epsilon}{1-\epsilon}$ instead of $\frac{\epsilon}{1-\epsilon}$.
Unlike the bounds in HaoChen et al. (2021); Saunshi et al. (2022), this bound
only depends on $\lambda_{d}^{-1}$, $\bar{\lambda}_{d}^{-1}$ and
$\gamma_{{\bm{G}}}$, but not $(\lambda_{d^{\prime}}-\lambda_{d+1})^{-1}$ for
some $d^{\prime}\leq d$, i.e. it does not require the separability of the
top-$d$ eigenfunctions, because it does not need the top-$d$ eigenspace to be
close to the empirical top-$d$ eigenspace. Instead, it only requires the
traces of the two top-$d$ eigenspaces to be close. This is a big improvement
since the eigenvalues can have high multiplicity as pointed out by Saunshi et
al. (2022).

➢ Proof Sketch of Theorem 2: Denote
$S_{\bar{\lambda}}(d):=\sum_{i=1}^{d}\bar{\lambda}_{i}$. We define two
empirical covariance matrices ${\hat{{\bm{F}}}}$ and ${\hat{{\bm{G}}}}$ as:
${\hat{{\bm{F}}}}(i,j)=\langle\hat{\psi}_{i},\hat{\psi}_{j}\rangle_{\hat{P}_{\mathcal{X}}}$,
and
${\hat{{\bm{G}}}}(i,j)=\langle\hat{\phi}_{i},\hat{\phi}_{j}\rangle_{\hat{P}_{\mathcal{A}}}$.
Since any invertible linear transformation to $\hat{\Phi}$ does not change
$\operatorname{Tr}({\bm{G}}^{-1}{\bm{F}})$ or the empirical ratio trace
$\operatorname{Tr}({\hat{{\bm{G}}}}^{-1}{\hat{{\bm{F}}}})$, we can see that
for ${\mathcal{H}}_{{\hat{\Psi}}}={\hat{{\mathcal{H}}}_{d}}$ there is
$\operatorname{Tr}({\hat{{\bm{G}}}}^{-1}{\hat{{\bm{F}}}})=S_{\bar{\lambda}}(d)$
(simply consider $\hat{\phi}_{i}=\bar{\phi}_{i}$). Thus, it suffices to bound
$|\operatorname{Tr}({\hat{{\bm{G}}}}^{-1}{\hat{{\bm{F}}}})-\operatorname{Tr}({\bm{G}}^{-1}{\bm{F}})|$,
and $S_{\lambda}(d)-S_{\bar{\lambda}}(d)$.

We start with bounding the gap between empirical and real ratio traces for any
encoder $\hat{\Phi}$:

######  Lemma 5.

Suppose there exists a constant $C>0$ such that
$\mathbb{E}_{P_{\mathcal{A}}}[g^{4}]\leq C^{2}\|g\|_{P_{\mathcal{A}}}^{2}$,
for all $g=w^{\top}\hat{\Phi}$ where $\|g\|_{P_{\mathcal{A}}}\leq 1$. Then,
for any $\delta>0$, it holds with probability at least $1-\delta$ that

$|\operatorname{Tr}({\hat{{\bm{G}}}}^{-1}{\hat{{\bm{F}}}})-\operatorname{Tr}({\bm{G}}^{-1}{\bm{F}})|\leq\left(2+\sqrt{2\log\frac{2}{\delta}}\right)\frac{C\kappa+\kappa^{2}}{\sqrt{N}}d.$
(12)

Lemma 5 considers a general $\hat{\Phi}$, and requires a fourth-moment control
assumption (Wainwright, 2019, Eqn. (14.22a)). For the specific case
$\hat{\phi}_{i}=\bar{\phi}_{i}$ studied in this section, we can prove that
this assumption holds (that is why it does not appear in Theorem 2), i.e. the
top-$d$ empirical eigenfunctions can be proved to be delocalized (Erdős et
al., 2009). In particular, we prove the following:

######  Lemma 6.

Suppose $\hat{\phi}_{i}=\bar{\phi}_{i}$ for $i\in[d]$. Let
$\gamma_{{\bm{G}}}:=\lambda_{\max}({\bm{G}})/\lambda_{\min}({\bm{G}})$, which
is the condition number of ${\bm{G}}$. Then, for any $\delta>0$, both

$\sum_{j=1}^{d}\bar{\lambda}_{j}\geq\sum_{i=1}^{d}\lambda_{i}-\left(2+\sqrt{2\log\frac{2}{\delta}}\right)\frac{(\lambda_{d}^{-1}+1)\kappa^{2}}{\sqrt{N}}d$

and Eqn. (12) with $C=\kappa\bar{\lambda}_{d}^{-1}\gamma_{{\bm{G}}}^{1/2}$
hold simultaneously for
${\mathcal{H}}_{{\hat{\Psi}}}={\hat{{\mathcal{H}}}_{d}}$ with probability at
least $1-\delta$.

##  5 Estimating and Exploiting the Augmentation Complexity

An important quantity from our analysis is the augmentation complexity
$\kappa$ — both approximation and estimation error bounds get smaller when
$\kappa$ is reduced. A natural way to reduce $\kappa$ is via a stronger
augmentation, which has indeed been helpful in practice (Chen et al., 2020;
Wettig et al., 2023).

In this section, we take a closer look at $\kappa$ of mask-type augmentations.
We study $\kappa$ of different masking schemes on the hypercube data model
introduced in Saunshi et al. (2022), and show that $\kappa$ depends on both
the mask ratio (augmentation strength) and the masking strategy. Then, we
empirically estimate $\kappa$ on the real-world NLP dataset wikipedia-simple,
and demonstrate its correlation with the generalization gap on real downstream
tasks. Our results indicate that the augmentation complexity $\kappa$ offers a
means to compare different augmentations quantitatively and implies a tradeoff
in the choice of augmentations.

###  5.1 Theoretical Hypercube Data Model

3\. Block Mask + Flip2\. Block Mask1\. Random MaskOriginal
Sample11-1-11-11100-101-11010000-111100-1111-1MaskedFlipped Figure 2:
Augmentation illustration.

Like Cabannes et al. (2023), we study three random masking methods (Figure 2):
(i) Independent random masking, (ii) Cutout-like block masking, and (iii)
BERT-like masking. Let $\alpha$ denote the mask rate. Consider the data space
${\mathcal{X}}=\\{-1,1\\}^{d_{\mathcal{X}}}$ with ${P_{\mathcal{X}}}$ being
the uniform distribution over ${\mathcal{X}}$. Denote the corresponding
$\kappa$ of these three methods to be $\kappa_{r},\kappa_{c}$, and
$\kappa_{b}$. We will show that these $\kappa$’s will be very different even
with the same $\alpha$.

Let’s first consider the independently random masking (derivations are
deferred to Appendix G):

######  Example 1.

Consider a _random masking_ augmentation, i.e. for any $x\in{\mathcal{X}}$,
each coordinate $x^{(i)}$ is randomly and independently masked to be $0$ (i.e.
$0$ denotes the [MASK] token) with probability $\alpha\in(0,1)$ . Then, its
augmentation complexity is given by
$\kappa_{r}^{2}=(2-\alpha)^{d_{\mathcal{X}}}$.

Next, consider the Cutout-like block masking (DeVries & Taylor, 2017), which
has been successfully applied in practice such as ViT (Dosovitskiy et al.,
2021) and MAE He et al. (2022). Compared to the previous example, Cutout-like
block masking leads to a smaller $\kappa$ with the same $\alpha$:

######  Example 2.

Consider _random block masking_ , which masks
$x^{(i)},x^{(i+1)},\cdots,x^{(i+r-1)}$ for $r=\lceil\alpha
d_{{\mathcal{X}}}\rceil$ and a uniformly random $i\in[d_{{\mathcal{X}}}-r]$,
for any $x\in{\mathcal{X}}$. Then,
$\kappa_{c}^{2}\leq[2^{(1-\alpha)}]^{d_{{\mathcal{X}}}}$.

Our third and final example resembles the 80-10-10 strategy in BERT (Devlin et
al., 2019): First randomly mask some tokens, and then randomly replace some
unmasked tokens with other tokens.

######  Example 3.

Consider _random block masking with flipping_ , where for any
$x\in{\mathcal{X}}$, first mask $x^{(i)},\cdots,x^{(i+r-1)}$ to be $0$ for
$r=\lceil\alpha d_{{\mathcal{X}}}\rceil$ and a uniformly random
$i\in[d_{{\mathcal{X}}}-r]$, then randomly flip the sign of each remaining
coordinate w.p. $\frac{\alpha}{2}$ independently. Then,
$\kappa_{b}^{2}\leq\left[(\alpha^{2}-2\alpha+2)^{(1-\alpha/2)}\right]^{d_{{\mathcal{X}}}}$.

While $\kappa_{r},\kappa_{c},\kappa_{b}$ all have an exponential dependency on
$d_{{\mathcal{X}}}$, their bases are different, as shown in Figure 3(a). This
means these methods have different $\kappa$ despite having the same mask rate,
which highlights the importance of the masking strategy.

##### On the exponentiality of $\kappa$ in $d_{{\mathcal{X}}}$.

We suspect the exponential dependency on $d_{\mathcal{X}}$ is a manifestation
of the typical curse of dimensionality in high-dimensional statistics in the
absence of strong inductive bias. Bengio et al. (2013, Section 3.2) had a
discussion on this point. What our analysis makes salient is that the
inductive bias can be expressed via controlled augmentation complexity. One
approach to obtain a polynomial augmentation complexity is by using very
strong augmentations, such as using ImageNet pretraining where ${\mathcal{A}}$
is a finite set of size $1000$, i.e. the ImageNet label set. However, such
simple strong augmentations usually lead to substantially worse performance in
practice. We posit developing more interesting augmentations with bounded
augmentation complexity as an important open problem facing the field of
representation learning.

We postulate moreover that despite the worst-case bounds above come with
exponential dependency on data dimension, the empirical success of existing
augmentation-based self-supervised learning suggest that they implicitly adapt
to the inherent low-dimensional manifold structure in real-world data. We
conjecture that as long as the augmentation captures such a structure, it can
evade the curse.

###  5.2 Augmentations on Real Datasets

The augmentation complexity $\kappa$ can be estimated on real datasets and
used to quantitatively analyze and compare augmentations, which we demonstrate
on the NLP dataset wikipedia-simple. We study masked language modeling, where
$x$ is a full sentence and $a$ is a masked sentence. Recall that $\kappa^{2}$
upper bounds $\int\frac{p(x|a)}{p(x)}p(a|x)da$. Empirically, we replace the
integration with sample average; namely, for $x=[x^{(1)},\cdots,x^{(l)}]$
where $x^{(i)}$ is the $i^{th}$ token, we have

$\log p\left(x\middle|a\right)=\log p\left(x^{(1)}\middle|a\right)+\log
p\left(x^{(2)}\middle|a,x^{(1)}\right)+\cdots+\log
p\left(x^{(l)}\middle|a,x^{(1)},\cdots,x^{(l-1)}\right).$

We leverage a bi-directional MLM such as a BERT, and compute
$p(x^{(i)}|a,x^{(<i)})$ auto-regressively: For any $i\in[l]$, compute
$p(x^{(i)}|a,x^{(<i)})$ with the BERT, and then replace $a^{(i)}$ with
$x^{(i)}$ for $i+1$. We compute this for a random subset of samples. We report
the $99^{\text{th}}$ percentile of ${K_{X}}(x,x)$ instead of
$\sup_{x}{K_{X}}(x,x)$ due to the presence of outliers. Experimental details
are deferred to Appendix H.

##### Estimating $\kappa$.

We study Random masking, Block masking, Random masking + Flipping, and Block
masking + Flipping. For Masking + Flipping, we randomly mask $\alpha/2$ of the
tokens, and replace another $\alpha/2$ of the tokens with random tokens; this
replace rate is much higher than standard practice, since we want to highlight
the effect of flipping. Figure 3(b) plots the estimated $\log\kappa^{2}$ of
the four augmentations w.r.t. the mask rate $\alpha$, where $\log\kappa^{2}$
is calculated as the average of five runs with different random seeds. The
complexity drops as $\alpha$ increases as predicted by theory. One observation
is that the “Random + Flip” curve intersects with “Block” and “Block + Flip”,
suggesting that block masking has a stronger effect when $\alpha$ is small,
whereas flipping has a stronger effect when $\alpha$ is large.

$0$$0.5$$1$$1$$1.5$$2$Mask Ratio
$\alpha$$y=2-\alpha$$y=2^{1-\alpha}$$y=(\alpha^{2}-2\alpha+2)^{1-\frac{\alpha}{2}}$
(a) Bases of three $\kappa^{2}$.

$0$$0.2$$0.4$$0.6$$0.8$$100$$200$$300$Mask Ratio
$\alpha$$\log\kappa^{2}$RandomRandom + FlipBlockBlock + Flip (b) Estimated
$\log\kappa^{2}$.

$0$$0.2$$0.4$$0.6$$0.8$$70$$80$$90$$(0.15,86.8)$Mask Ratio $\alpha$Test
Accuracy$0$$2$$4$QNLI$0$$0.2$$0.4$$0.6$$0.8$$80$$85$$90$$95$$(0.40,92.4)$Mask
Ratio $\alpha$$6$$9$$12$Train-Test GapSST-2 (c) Results on QNLI and SST-2.

Figure 3: Plots for Section 5. In (b), $\log\kappa^{2}$ is estimated on
wikipedia-simple.

##### Downstream performance.

We also study how the mask ratio $\alpha$ affects the downstream performance,
using QNLI (Wang et al., 2018) and SST-2 (Socher et al., 2013). For
pretraining, we train roberta-large models with different mask ratios using
the fast pretraining recipe in Wettig et al. (2023). We focus on Random
masking and do not apply the 80-10-10 strategy. For downstream, we fine-tune
the encoder together with the linear head following common practice. To align
better with our theory, we explicitly use the average encoder Eqn. 7,
estimated by sampling 16 $a$’s for each $x$.

We evaluate the train/test accuracies of the models, and plot the test
accuracy (blue solid) and the train-test accuracy gap (green dashed) in Figure
3(c). The highest test accuracy is achieved at $\alpha=0.15$ on QNLI and at
$\alpha=0.40$ on SST-2 (marked in red). The test accuracy is low when $\alpha$
is too small due to the large generalization gap, and also low when $\alpha$
is too large due to low training accuracy. Regarding the train-test gap, QNLI
shows a monotonic decrease in the gap as the mask ratio grows, but the gap on
SST-2 is U-shaped, with the lowest point at $\alpha=0.40$. This is likely
because with $\alpha>0.40$ is too strong an augmentation for SST-2 that
Assumption 1 is broken, in which case our theoretical results will not hold.
Thus, these results align with our theory that while augmentations should be
sufficiently robust, they must not be so strong as to violate Assumption 1.
This suggests the presence of a “sweet spot”, which is also supported by
evidence in prior work (Tian et al., 2020).

##  6 Conclusion

This work establishes an RKHS approximation/regression framework that
completely disentangles the effects of the model and the augmentation. Our
framework (i) clarifies the roles of the two stages of representation
learning, as well as the role of augmentation with the isometry property Eqn.
6; (ii) leads to nonparametric statistical guarantees free of model complexity
or inductive bias, so that they are valid for an arbitrary encoder; and (iii)
formulates the augmentation complexity $\kappa$ as a tool to quantitatively
analyze the effect of augmentations. We also refer our readers to two recent
works that used similar techniques but derived quite different results: Wang
et al. (2022b) studied an RKHS similar to ${{\mathcal{H}}_{\Gamma}}$, but in a
completely different context (they studied conditional moment models);
Cabannes et al. (2023) studied a similar problem, but they used a pre-defined
kernel to define the function class and derived guarantees, while the RKHS
${{\mathcal{H}}_{\Gamma}}$ in this work is completely induced by the data
augmentation. We defer discussions for more related work to Appendix B.

##### Limitations and open problems.

The major limitation of this work is that $\kappa$ can be exponential in
$d_{{\mathcal{X}}}$, which as we pointed out could be a natural consequence of
the curse of dimensionality, and that the low-dimensional structure of real-
world data could be a potential reason of circumventing the exponential
dependency. Another limitation is that we studied the average encoder, while
usually people directly apply the original encoder to downstream. The latter
case requires the model to perform some type of implicit transformation as the
upstream and downstream have different input spaces (${\mathcal{A}}$ and
${\mathcal{X}}$), and this transformation cannot be studied within the
architecture-free framework of this work. Additionally, an open problem is how
to deal with the common upstream-to-downstream data distribution shift in
practice, for which downstream fine-tuning could be critical. This has been
discussed in Cabannes et al. (2023), but a more general analysis is desired.

#### Code

The code of Section 5 can be found at
https://colab.research.google.com/drive/1loSZLLI-
qfoKE7BCIi1SWJKgruU6i4ku?usp=sharing.

#### Acknowledgments

We would like to thank Yutong He, Di He and Elan Rosenfeld for their very
useful discussions, and Jeremy Cohen, Amrith Setlur, Xiaoyu Huang, Samuel
Sokota, Zhili Feng, Swaminathan Gurumurthy and Zhengyang Geng for their
feedback on the early draft of this work. We are also very grateful to our
anonymous ICLR reviewers, with whose help this work has been greatly improved.
We acknowledge the support of NSF via IIS-2211907, IIS-1909816, CCF-2238523,
ONR via N00014-23-1-2368, and Amazon.

## References

  * Allen-Zhu et al. (2019a) Zeyuan Allen-Zhu, Yuanzhi Li, and Yingyu Liang.  Learning and generalization in overparameterized neural networks, going beyond two layers.  _Advances in neural information processing systems_ , 32, 2019a.
  * Allen-Zhu et al. (2019b) Zeyuan Allen-Zhu, Yuanzhi Li, and Zhao Song.  A convergence theory for deep learning via over-parameterization.  In _International Conference on Machine Learning_ , pp. 242–252. PMLR, 2019b.
  * Arora et al. (2018) Sanjeev Arora, Rong Ge, Behnam Neyshabur, and Yi Zhang.  Stronger generalization bounds for deep nets via a compression approach.  In _International Conference on Machine Learning_ , pp. 254–263. PMLR, 2018.
  * Arora et al. (2019a) Sanjeev Arora, Nadav Cohen, Wei Hu, and Yuping Luo.  Implicit regularization in deep matrix factorization.  _Advances in Neural Information Processing Systems_ , 32, 2019a.
  * Arora et al. (2019b) Sanjeev Arora, Simon Du, Wei Hu, Zhiyuan Li, and Ruosong Wang.  Fine-grained analysis of optimization and generalization for overparameterized two-layer neural networks.  In _International Conference on Machine Learning_ , pp. 322–332. PMLR, 2019b.
  * Arora et al. (2019c) Sanjeev Arora, Simon S Du, Wei Hu, Zhiyuan Li, Russ R Salakhutdinov, and Ruosong Wang.  On exact computation with an infinitely wide neural net.  _Advances in neural information processing systems_ , 32, 2019c.
  * Bardes et al. (2022) Adrien Bardes, Jean Ponce, and Yann LeCun.  VICReg: Variance-invariance-covariance regularization for self-supervised learning.  In _International Conference on Learning Representations_ , 2022.  URL https://openreview.net/forum?id=xm6YD62D1Ub.
  * Bartlett et al. (2017) Peter L Bartlett, Dylan J Foster, and Matus J Telgarsky.  Spectrally-normalized margin bounds for neural networks.  _Advances in neural information processing systems_ , 30, 2017.
  * Bartlett et al. (2023) Peter L. Bartlett, Philip M. Long, and Olivier Bousquet.  The dynamics of sharpness-aware minimization: Bouncing across ravines and drifting towards wide minima.  _Journal of Machine Learning Research_ , 24(316):1–36, 2023.  URL http://jmlr.org/papers/v24/23-043.html.
  * Belkin & Niyogi (2003) Mikhail Belkin and Partha Niyogi.  Laplacian eigenmaps for dimensionality reduction and data representation.  _Neural computation_ , 15(6):1373–1396, 2003.
  * Belkin & Niyogi (2004) Mikhail Belkin and Partha Niyogi.  Semi-supervised learning on riemannian manifolds.  _Machine learning_ , 56:209–239, 2004.
  * Bengio et al. (2004) Yoshua Bengio, Olivier Delalleau, Nicolas Le Roux, Jean-François Paiement, Pascal Vincent, and Marie Ouimet.  Learning eigenfunctions links spectral embedding and kernel pca.  _Neural computation_ , 16(10):2197–2219, 2004.
  * Bengio et al. (2013) Yoshua Bengio, Aaron Courville, and Pascal Vincent.  Representation learning: A review and new perspectives.  _IEEE transactions on pattern analysis and machine intelligence_ , 35(8):1798–1828, 2013.
  * Bousquet & Elisseeff (2002) Olivier Bousquet and André Elisseeff.  Stability and generalization.  _The Journal of Machine Learning Research_ , 2:499–526, 2002.
  * Cabannes et al. (2023) Vivien Cabannes, Bobak Kiani, Randall Balestriero, Yann Lecun, and Alberto Bietti.  The SSL interplay: Augmentations, inductive bias, and generalization.  In Andreas Krause, Emma Brunskill, Kyunghyun Cho, Barbara Engelhardt, Sivan Sabato, and Jonathan Scarlett (eds.), _Proceedings of the 40th International Conference on Machine Learning_ , volume 202 of _Proceedings of Machine Learning Research_ , pp. 3252–3298. PMLR, 23–29 Jul 2023.
  * Chan et al. (2022) Kwan Ho Ryan Chan, Yaodong Yu, Chong You, Haozhi Qi, John Wright, and Yi Ma.  Redunet: A white-box deep network from the principle of maximizing rate reduction.  _The Journal of Machine Learning Research_ , 23(1):4907–5009, 2022.
  * Chen et al. (2020) Ting Chen, Simon Kornblith, Mohammad Norouzi, and Geoffrey Hinton.  A simple framework for contrastive learning of visual representations.  In _International conference on machine learning_ , pp. 1597–1607. PMLR, 2020.
  * Chizat et al. (2019) Lenaic Chizat, Edouard Oyallon, and Francis Bach.  On lazy training in differentiable programming.  _Advances in neural information processing systems_ , 32, 2019.
  * Chung (1997) Fan RK Chung.  _Spectral graph theory_ , volume 92.  American Mathematical Soc., 1997.
  * Dai et al. (2022) Xili Dai, Shengbang Tong, Mingyang Li, Ziyang Wu, Michael Psenka, Kwan Ho Ryan Chan, Pengyuan Zhai, Yaodong Yu, Xiaojun Yuan, Heung-Yeung Shum, et al.  Ctrl: Closed-loop transcription to an ldr via minimaxing rate reduction.  _Entropy_ , 24(4):456, 2022.
  * Dao et al. (2019) Tri Dao, Albert Gu, Alexander Ratner, Virginia Smith, Chris De Sa, and Christopher Ré.  A kernel theory of modern data augmentation.  In _International conference on machine learning_ , pp. 1528–1537. PMLR, 2019.
  * Devlin et al. (2019) Jacob Devlin, Ming-Wei Chang, Kenton Lee, and Kristina Toutanova.  BERT: Pre-training of deep bidirectional transformers for language understanding.  In _Proceedings of the 2019 Conference of the North American Chapter of the Association for Computational Linguistics: Human Language Technologies, Volume 1 (Long and Short Papers)_ , pp. 4171–4186, Minneapolis, Minnesota, June 2019. Association for Computational Linguistics.  doi: 10.18653/v1/N19-1423.  URL https://aclanthology.org/N19-1423.
  * DeVries & Taylor (2017) Terrance DeVries and Graham W Taylor.  Improved regularization of convolutional neural networks with cutout.  _arXiv preprint arXiv:1708.04552_ , 2017.
  * Dosovitskiy et al. (2021) Alexey Dosovitskiy, Lucas Beyer, Alexander Kolesnikov, Dirk Weissenborn, Xiaohua Zhai, Thomas Unterthiner, Mostafa Dehghani, Matthias Minderer, Georg Heigold, Sylvain Gelly, Jakob Uszkoreit, and Neil Houlsby.  An image is worth 16x16 words: Transformers for image recognition at scale.  In _International Conference on Learning Representations_ , 2021.  URL https://openreview.net/forum?id=YicbFdNTTy.
  * Du & Lee (2018) Simon Du and Jason Lee.  On the power of over-parametrization in neural networks with quadratic activation.  In _International conference on machine learning_ , pp. 1329–1338. PMLR, 2018.
  * Du et al. (2019a) Simon Du, Jason Lee, Haochuan Li, Liwei Wang, and Xiyu Zhai.  Gradient descent finds global minima of deep neural networks.  In _International conference on machine learning_ , pp. 1675–1685. PMLR, 2019a.
  * Du et al. (2019b) Simon S. Du, Xiyu Zhai, Barnabas Poczos, and Aarti Singh.  Gradient descent provably optimizes over-parameterized neural networks.  In _International Conference on Learning Representations_ , 2019b.  URL https://openreview.net/forum?id=S1eK3i09YQ.
  * Dziugaite & Roy (2017) Gintare Karolina Dziugaite and Daniel M Roy.  Computing nonvacuous generalization bounds for deep (stochastic) neural networks with many more parameters than training data.  In Gal Elidan, Kristian Kersting, and Alexander T. Ihler (eds.), _Proceedings of the Thirty-third Conference on Uncertainty in Artificial Intelligence_ , 2017.
  * Erdős et al. (2009) László Erdős, Benjamin Schlein, and Horng-Tzer Yau.  Local semicircle law and complete delocalization for wigner random matrices.  _Communications in Mathematical Physics_ , 287(2):641–655, 2009.
  * Fan (1949) Ky Fan.  On a theorem of weyl concerning eigenvalues of linear transformations i.  _Proceedings of the National Academy of Sciences_ , 35(11):652–655, 1949.
  * Foret et al. (2021) Pierre Foret, Ariel Kleiner, Hossein Mobahi, and Behnam Neyshabur.  Sharpness-aware minimization for efficiently improving generalization.  In _International Conference on Learning Representations_ , 2021.  URL https://openreview.net/forum?id=6Tm1mposlrM.
  * Goldblum et al. (2023) Micah Goldblum, Marc Finzi, Keefer Rowan, and Andrew Gordon Wilson.  The No Free Lunch Theorem, Kolmogorov Complexity, and the Role of Inductive Biases in Machine Learning.  _arXiv e-prints_ , art. arXiv:2304.05366, April 2023.  doi: 10.48550/arXiv.2304.05366.
  * Györfi et al. (2002) László Györfi, Michael Kohler, Adam Krzyzak, Harro Walk, et al.  _A distribution-free theory of nonparametric regression_ , volume 1.  Springer, 2002.
  * Hahn & Meeker (2011) Gerald J Hahn and William Q Meeker.  _Statistical intervals: a guide for practitioners_ , volume 92.  John Wiley & Sons, 2011.
  * HaoChen & Ma (2023) Jeff Z. HaoChen and Tengyu Ma.  A theoretical study of inductive biases in contrastive learning.  In _The Eleventh International Conference on Learning Representations_ , 2023.  URL https://openreview.net/forum?id=AuEgNlEAmed.
  * HaoChen et al. (2021) Jeff Z HaoChen, Colin Wei, Adrien Gaidon, and Tengyu Ma.  Provable guarantees for self-supervised deep learning with spectral contrastive loss.  _Advances in Neural Information Processing Systems_ , 34:5000–5011, 2021.
  * HaoChen et al. (2022) Jeff Z. HaoChen, Colin Wei, Ananya Kumar, and Tengyu Ma.  Beyond separability: Analyzing the linear transferability of contrastive representations to related subpopulations.  In Alice H. Oh, Alekh Agarwal, Danielle Belgrave, and Kyunghyun Cho (eds.), _Advances in Neural Information Processing Systems_ , 2022.  URL https://openreview.net/forum?id=vmjckXzRXmh.
  * Hardt et al. (2016) Moritz Hardt, Ben Recht, and Yoram Singer.  Train faster, generalize better: Stability of stochastic gradient descent.  In _International conference on machine learning_ , pp. 1225–1234. PMLR, 2016.
  * He et al. (2022) Kaiming He, Xinlei Chen, Saining Xie, Yanghao Li, Piotr Dollár, and Ross Girshick.  Masked autoencoders are scalable vision learners.  In _Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition_ , pp. 16000–16009, 2022.
  * Hochreiter & Schmidhuber (1997) Sepp Hochreiter and Jürgen Schmidhuber.  Flat minima.  _Neural computation_ , 9(1):1–42, 1997.
  * Huang et al. (2023) Weiran Huang, Mingyang Yi, Xuyang Zhao, and Zihao Jiang.  Towards the generalization of contrastive self-supervised learning.  In _The Eleventh International Conference on Learning Representations_ , 2023.  URL https://openreview.net/forum?id=XDJwuEYHhme.
  * Huang et al. (2020) Xiao Shi Huang, Felipe Perez, Jimmy Ba, and Maksims Volkovs.  Improving transformer optimization through better initialization.  In Hal Daumé III and Aarti Singh (eds.), _Proceedings of the 37th International Conference on Machine Learning_ , volume 119 of _Proceedings of Machine Learning Research_ , pp. 4475–4483. PMLR, 13–18 Jul 2020.
  * Jacot et al. (2018) Arthur Jacot, Franck Gabriel, and Clément Hongler.  Neural tangent kernel: Convergence and generalization in neural networks.  _Advances in neural information processing systems_ , 31, 2018.
  * Jiang et al. (2020) Yiding Jiang, Behnam Neyshabur, Hossein Mobahi, Dilip Krishnan, and Samy Bengio.  Fantastic generalization measures and where to find them.  In _International Conference on Learning Representations_ , 2020.  URL https://openreview.net/forum?id=SJgIPJBFvH.
  * Jing et al. (2022) Li Jing, Pascal Vincent, Yann LeCun, and Yuandong Tian.  Understanding dimensional collapse in contrastive self-supervised learning.  In _International Conference on Learning Representations_ , 2022.  URL https://openreview.net/forum?id=YevsQ05DEN7.
  * Johnson et al. (2023) Daniel D. Johnson, Ayoub El Hanchi, and Chris J. Maddison.  Contrastive learning can find an optimal basis for approximately view-invariant functions.  In _The Eleventh International Conference on Learning Representations_ , 2023.  URL https://openreview.net/forum?id=AjC0KBjiMu.
  * Kawaguchi (2016) Kenji Kawaguchi.  Deep learning without poor local minima.  _Advances in neural information processing systems_ , 29, 2016.
  * Kearns & Ron (1997) Michael Kearns and Dana Ron.  Algorithmic stability and sanity-check bounds for leave-one-out cross-validation.  In _Proceedings of the tenth annual conference on Computational learning theory_ , pp. 152–162, 1997.
  * Keskar et al. (2017) Nitish Shirish Keskar, Dheevatsa Mudigere, Jorge Nocedal, Mikhail Smelyanskiy, and Ping Tak Peter Tang.  On large-batch training for deep learning: Generalization gap and sharp minima.  In _International Conference on Learning Representations_ , 2017.  URL https://openreview.net/forum?id=H1oyRlYgg.
  * Kolmogorov (1963) Andrei N Kolmogorov.  On tables of random numbers.  _Sankhyā: The Indian Journal of Statistics, Series A_ , pp. 369–376, 1963.
  * Langford & Seeger (2001) John Langford and Matthias Seeger.  _Bounds for averaging classifiers_.  School of Computer Science, Carnegie Mellon University, 2001.
  * Lee et al. (2019) Jaehoon Lee, Lechao Xiao, Samuel Schoenholz, Yasaman Bahri, Roman Novak, Jascha Sohl-Dickstein, and Jeffrey Pennington.  Wide neural networks of any depth evolve as linear models under gradient descent.  _Advances in neural information processing systems_ , 32, 2019.
  * Lee et al. (2021) Jason D Lee, Qi Lei, Nikunj Saunshi, and Jiacheng Zhuo.  Predicting what you already know helps: Provable self-supervised learning.  _Advances in Neural Information Processing Systems_ , 34:309–323, 2021.
  * Li & Yuan (2017) Yuanzhi Li and Yang Yuan.  Convergence analysis of two-layer neural networks with relu activation.  _Advances in neural information processing systems_ , 30, 2017.
  * Li et al. (2018) Yuanzhi Li, Tengyu Ma, and Hongyang Zhang.  Algorithmic regularization in over-parameterized matrix sensing and neural networks with quadratic activations.  In _Conference On Learning Theory_ , pp. 2–47. PMLR, 2018.
  * Liu et al. (2022) Bingbin Liu, Daniel Hsu, Pradeep Ravikumar, and Andrej Risteski.  Masked prediction tasks: a parameter identifiability view.  _Advances in Neural Information Processing Systems_ , 2022.
  * Liu et al. (2021) Hong Liu, Jeff Z. HaoChen, Adrien Gaidon, and Tengyu Ma.  Self-supervised learning is more robust to dataset imbalance.  In _NeurIPS 2021 Workshop on Distribution Shifts: Connecting Methods and Applications_ , 2021.  URL https://openreview.net/forum?id=vUz4JPRLpGx.
  * Lotfi et al. (2022) Sanae Lotfi, Marc Anton Finzi, Sanyam Kapoor, Andres Potapczynski, Micah Goldblum, and Andrew Gordon Wilson.  PAC-bayes compression bounds so tight that they can explain generalization.  In Alice H. Oh, Alekh Agarwal, Danielle Belgrave, and Kyunghyun Cho (eds.), _Advances in Neural Information Processing Systems_ , 2022.  URL https://openreview.net/forum?id=o8nYuR8ekFm.
  * McAllester (1999) David A McAllester.  Pac-bayesian model averaging.  In _Proceedings of the twelfth annual conference on Computational learning theory_ , pp. 164–170, 1999.
  * Mohri et al. (2018) Mehryar Mohri, Afshin Rostamizadeh, and Ameet Talwalkar.  _Foundations of machine learning_.  MIT press, 2018.
  * Mroueh et al. (2015) Youssef Mroueh, Stephen Voinea, and Tomaso A Poggio.  Learning with group invariant features: A kernel perspective.  _Advances in neural information processing systems_ , 28, 2015.
  * Mukherjee et al. (2006) Sayan Mukherjee, Partha Niyogi, Tomaso Poggio, and Ryan Rifkin.  Learning theory: stability is sufficient for generalization and necessary and sufficient for consistency of empirical risk minimization.  _Advances in Computational Mathematics_ , 25:161–193, 2006.
  * Nagarajan & Kolter (2019) Vaishnavh Nagarajan and Zico Kolter.  Deterministic PAC-bayesian generalization bounds for deep networks via generalizing noise-resilience.  In _International Conference on Learning Representations_ , 2019.  URL https://openreview.net/forum?id=Hygn2o0qKX.
  * Negrea et al. (2020) Jeffrey Negrea, Gintare Karolina Dziugaite, and Daniel Roy.  In defense of uniform convergence: Generalization via derandomization with an application to interpolating predictors.  In _International Conference on Machine Learning_ , pp. 7263–7272. PMLR, 2020.
  * Neyshabur et al. (2017) Behnam Neyshabur, Srinadh Bhojanapalli, David McAllester, and Nati Srebro.  Exploring generalization in deep learning.  _Advances in neural information processing systems_ , 30, 2017.
  * Neyshabur et al. (2018) Behnam Neyshabur, Srinadh Bhojanapalli, and Nathan Srebro.  A PAC-bayesian approach to spectrally-normalized margin bounds for neural networks.  In _International Conference on Learning Representations_ , 2018.  URL https://openreview.net/forum?id=Skz_WfbCZ.
  * Pokle et al. (2022) Ashwini Pokle, Jinjin Tian, Yuchen Li, and Andrej Risteski.  Contrasting the landscape of contrastive and non-contrastive learning.  In Gustau Camps-Valls, Francisco J. R. Ruiz, and Isabel Valera (eds.), _International Conference on Artificial Intelligence and Statistics, AISTATS 2022, 28-30 March 2022, Virtual Event_ , volume 151 of _Proceedings of Machine Learning Research_ , pp. 8592–8618. PMLR, 2022.  URL https://proceedings.mlr.press/v151/pokle22a.html.
  * Radford et al. (2021) Alec Radford, Jong Wook Kim, Chris Hallacy, Aditya Ramesh, Gabriel Goh, Sandhini Agarwal, Girish Sastry, Amanda Askell, Pamela Mishkin, Jack Clark, et al.  Learning transferable visual models from natural language supervision.  In _International Conference on Machine Learning_ , pp. 8748–8763. PMLR, 2021.
  * Raj et al. (2017) Anant Raj, Abhishek Kumar, Youssef Mroueh, Tom Fletcher, and Bernhard Schölkopf.  Local group invariant representations via orbit embeddings.  In _Artificial Intelligence and Statistics_ , pp. 1225–1235. PMLR, 2017.
  * Rivasplata et al. (2019) Omar Rivasplata, Vikram M Tankasali, and Csaba Szepesvári.  Pac-bayes with backprop.  _arXiv preprint arXiv:1908.07380_ , 2019.
  * Saunshi et al. (2019) Nikunj Saunshi, Orestis Plevrakis, Sanjeev Arora, Mikhail Khodak, and Hrishikesh Khandeparkar.  A theoretical analysis of contrastive unsupervised representation learning.  In _International Conference on Machine Learning_ , pp. 5628–5637. PMLR, 2019.
  * Saunshi et al. (2022) Nikunj Saunshi, Jordan Ash, Surbhi Goel, Dipendra Misra, Cyril Zhang, Sanjeev Arora, Sham Kakade, and Akshay Krishnamurthy.  Understanding contrastive learning requires incorporating inductive biases.  In _International Conference on Machine Learning_ , pp. 19250–19286. PMLR, 2022.
  * Schölkopf & Smola (2002) Bernhard Schölkopf and Alexander J Smola.  _Learning with kernels: support vector machines, regularization, optimization, and beyond_.  MIT press, 2002.
  * Shalev-Shwartz et al. (2010) Shai Shalev-Shwartz, Ohad Shamir, Nathan Srebro, and Karthik Sridharan.  Learnability, stability and uniform convergence.  _The Journal of Machine Learning Research_ , 11:2635–2670, 2010.
  * Shen et al. (2022) Kendrick Shen, Robbie M Jones, Ananya Kumar, Sang Michael Xie, Jeff Z HaoChen, Tengyu Ma, and Percy Liang.  Connect, not collapse: Explaining contrastive learning for unsupervised domain adaptation.  In _International Conference on Machine Learning_ , pp. 19847–19878. PMLR, 2022.
  * Socher et al. (2013) Richard Socher, Alex Perelygin, Jean Wu, Jason Chuang, Christopher D. Manning, Andrew Ng, and Christopher Potts.  Recursive deep models for semantic compositionality over a sentiment treebank.  In _Proceedings of the 2013 Conference on Empirical Methods in Natural Language Processing_ , pp. 1631–1642, Seattle, Washington, USA, October 2013. Association for Computational Linguistics.  URL https://www.aclweb.org/anthology/D13-1170.
  * Tian et al. (2020) Yonglong Tian, Chen Sun, Ben Poole, Dilip Krishnan, Cordelia Schmid, and Phillip Isola.  What makes for good views for contrastive learning?  _Advances in Neural Information Processing Systems_ , 33:6827–6839, 2020.
  * Tian (2022) Yuandong Tian.  Deep contrastive learning is provably (almost) principal component analysis.  _Advances in Neural Information Processing Systems_ , 2022.
  * Tian et al. (2021) Yuandong Tian, Xinlei Chen, and Surya Ganguli.  Understanding self-supervised learning dynamics without contrastive pairs.  In Marina Meila and Tong Zhang (eds.), _Proceedings of the 38th International Conference on Machine Learning, ICML 2021, 18-24 July 2021, Virtual Event_ , volume 139 of _Proceedings of Machine Learning Research_ , pp. 10268–10278. PMLR, 2021.  URL http://proceedings.mlr.press/v139/tian21a.html.
  * Tosh et al. (2021a) Christopher Tosh, Akshay Krishnamurthy, and Daniel Hsu.  Contrastive estimation reveals topic posterior information to linear models.  _J. Mach. Learn. Res._ , 22:281–1, 2021a.
  * Tosh et al. (2021b) Christopher Tosh, Akshay Krishnamurthy, and Daniel Hsu.  Contrastive learning, multi-view redundancy, and linear models.  In _Algorithmic Learning Theory_ , pp. 1179–1206. PMLR, 2021b.
  * Wainwright (2019) Martin J Wainwright.  _High-dimensional statistics: A non-asymptotic viewpoint_ , volume 48.  Cambridge university press, 2019.
  * Wang et al. (2018) Alex Wang, Amanpreet Singh, Julian Michael, Felix Hill, Omer Levy, and Samuel R. Bowman.  Glue: A multi-task benchmark and analysis platform for natural language understanding.  _BLACKBOXNLP@EMNLP_ , 2018.  doi: 10.18653/v1/W18-5446.
  * Wang et al. (2007) Huan Wang, Shuicheng Yan, Dong Xu, Xiaoou Tang, and Thomas Huang.  Trace ratio vs. ratio trace for dimensionality reduction.  In _2007 IEEE Conference on Computer Vision and Pattern Recognition_ , pp. 1–8. IEEE, 2007.
  * Wang et al. (2020) Xiang Wang, Chenwei Wu, Jason D Lee, Tengyu Ma, and Rong Ge.  Beyond lazy training for over-parameterized tensor decomposition.  _Advances in Neural Information Processing Systems_ , 33:21934–21944, 2020.
  * Wang et al. (2022a) Yifei Wang, Qi Zhang, Yisen Wang, Jiansheng Yang, and Zhouchen Lin.  Chaos is a ladder: A new theoretical understanding of contrastive learning via augmentation overlap.  In _International Conference on Learning Representations_ , 2022a.  URL https://openreview.net/forum?id=ECvgmYVyeUz.
  * Wang et al. (2023) Yifei Wang, Qi Zhang, Tianqi Du, Jiansheng Yang, Zhouchen Lin, and Yisen Wang.  A message passing perspective on learning dynamics of contrastive learning.  In _The Eleventh International Conference on Learning Representations_ , 2023.  URL https://openreview.net/forum?id=VBTJqqWjxMv.
  * Wang et al. (2022b) Ziyu Wang, Yucen Luo, Yueru Li, Jun Zhu, and Bernhard Schölkopf.  Spectral representation learning for conditional moment models.  _arXiv preprint arXiv:2210.16525_ , 2022b.
  * Wei et al. (2021) Colin Wei, Sang Michael Xie, and Tengyu Ma.  Why do pretrained language models help in downstream tasks? an analysis of head and prompt tuning.  _Advances in Neural Information Processing Systems_ , 34:16158–16170, 2021.
  * Wen et al. (2023) Kaiyue Wen, Tengyu Ma, and Zhiyuan Li.  How sharpness-aware minimization minimizes sharpness?  In _The Eleventh International Conference on Learning Representations_ , 2023.  URL https://openreview.net/forum?id=5spDgWmpY6x.
  * Wen & Li (2021) Zixin Wen and Yuanzhi Li.  Toward understanding the feature learning process of self-supervised contrastive learning.  In Marina Meila and Tong Zhang (eds.), _Proceedings of the 38th International Conference on Machine Learning, ICML 2021, 18-24 July 2021, Virtual Event_ , volume 139 of _Proceedings of Machine Learning Research_ , pp. 11112–11122. PMLR, 2021.  URL http://proceedings.mlr.press/v139/wen21c.html.
  * Wen & Li (2022) Zixin Wen and Yuanzhi Li.  The mechanism of prediction head in non-contrastive self-supervised learning.  _Advances in Neural Information Processing Systems_ , 2022.
  * Wettig et al. (2023) Alexander Wettig, Tianyu Gao, Zexuan Zhong, and Danqi Chen.  Should you mask 15% in masked language modeling?  In _Proceedings of the 17th Conference of the European Chapter of the Association for Computational Linguistics: Main Volume_ , 2023.
  * Wolpert & Macready (1997) David H Wolpert and William G Macready.  No free lunch theorems for optimization.  _IEEE transactions on evolutionary computation_ , 1(1):67–82, 1997.
  * Xiong et al. (2020) Ruibin Xiong, Yunchang Yang, Di He, Kai Zheng, Shuxin Zheng, Chen Xing, Huishuai Zhang, Yanyan Lan, Liwei Wang, and Tieyan Liu.  On layer normalization in the transformer architecture.  In _International Conference on Machine Learning_ , pp. 10524–10533. PMLR, 2020.
  * Yu et al. (2020) Yaodong Yu, Kwan Ho Ryan Chan, Chong You, Chaobing Song, and Yi Ma.  Learning diverse and discriminative representations via the principle of maximal coding rate reduction.  _Advances in Neural Information Processing Systems_ , 33:9422–9434, 2020.
  * Zbontar et al. (2021) Jure Zbontar, Li Jing, Ishan Misra, Yann LeCun, and Stéphane Deny.  Barlow twins: Self-supervised learning via redundancy reduction.  In _International Conference on Machine Learning_ , pp. 12310–12320. PMLR, 2021.
  * Zhai et al. (2024) Runtian Zhai, Rattana Pukdee, Roger Jin, Maria-Florina Balcan, and Pradeep Ravikumar.  Spectrally transformed kernel regression.  In _International Conference on Learning Representations_ , 2024.  URL https://openreview.net/forum?id=OeQE9zsztS.
  * Zhang et al. (2017) Chiyuan Zhang, Samy Bengio, Moritz Hardt, Benjamin Recht, and Oriol Vinyals.  Understanding deep learning requires rethinking generalization.  In _International Conference on Learning Representations_ , 2017.  URL https://openreview.net/forum?id=Sy8gdB9xx.
  * Zhang et al. (2021) Chiyuan Zhang, Samy Bengio, Moritz Hardt, Benjamin Recht, and Oriol Vinyals.  Understanding deep learning (still) requires rethinking generalization.  _Communications of the ACM_ , 64(3):107–115, 2021.
  * Zhou et al. (2019) Wenda Zhou, Victor Veitch, Morgane Austern, Ryan P. Adams, and Peter Orbanz.  Non-vacuous generalization bounds at the imagenet scale: a PAC-bayesian compression approach.  In _International Conference on Learning Representations_ , 2019.  URL https://openreview.net/forum?id=BJgqqsAct7.

##  Appendix A Discussion

##### Why do we need Assumption 1 and can it be relaxed?

We first note that an assumption on the target function is necessary to
establish any non-vacuous learning theory. This is implied by the No Free
Lunch Theorem (Wolpert & Macready, 1997), which states that if the target
function is sampled uniformly at random from all possible target functions,
then any learner achieves the same performance as random guess.

Assumption 1 implies an isometry property. As pointed out by Goldblum et al.
(2023), real-world datasets are highly structured, and the structures are
largely shared across different domains. Assumption 1 and implied isometry
property suggest that a well-designed augmentation is a key to capturing such
a low-dimensional cross-domain structure. Indeed, better augmentation has been
the main driving force in the recent progress in contrastive learning and
masked language modeling. Our focus on augmentation is in stark contrast to
prior work on the generalization of overparameterized models (see Appendix B),
which can find the global minima fitting all samples, yet cannot capture the
invariant structure within the data and the target concept.

If we were to relax Assumption 1 (e.g. the target function $f^{*}$ can be
outside but very close to ${{\mathcal{H}}_{\Gamma}}$), additional assumptions
on $f^{*}$ would be required for our proof with uniform deviation bounds to
hold. Namely, $f^{*}$ can no longer be any $L^{2}$ function that is close to
${{\mathcal{H}}_{\Gamma}}$. For instance, we at least need to assume that
$|f^{*}(x)|$ is a.e. uniformly bounded by some constant $C$ which the final
bound will depend on, or assume that $f^{*}$ belongs to some predefined RKHS
similar to Cabannes et al. (2023). We feel that these additional assumptions
would complicate our analysis and obscure the role of augmentation, hence we
choose to assume $f^{*}\in{{\mathcal{H}}_{\Gamma}}$ instead.

##### The main theorem allows the encoder to use any model. Does it mean that
the inductive bias of the model is not useful as opposed to prior work?

No, the inductive bias is still useful. In fact, we believe that the model
family and data augmentation are two key ingredients of modern representation
learning. The former has been studied by prior work (Saunshi et al., 2022;
HaoChen & Ma, 2023), while this work studies the latter and introduces a class
of augmentation complexity based generalization bounds that is orthogonal to
the class of classical model complexity based bounds, in the sense that the
bounds are nonparametric and do not require any inductive biases from the
model family. This new class of guarantees deepens our understanding because
it clarifies the roles of these two ingredients. However, a downside of our
analysis is that it cannot leverage the geometric structure of the data. For
example, the data might be so good that it is linearly separable in the
Euclidean space, but our nonparametric framework cannot make use of this
information.

The inductive bias of the model has at least two roles: (i) It further reduces
the complexity on top of the data augmentation, leading to even tighter
bounds; (ii) While we used the average encoder in our analysis, in practice
people usually directly apply the pretrained encoder to the data in the
downstream task. In this case, the model architecture has to implicitly
perform some kind of transformation, since the pretraining and downstream
samples come from different spaces, in which case the inductive biases of the
model will be helpful.

##### Comparing to the claim in Saunshi et al. (2022).

Saunshi et al. (2022) argued that there exist setups where any learning
guarantees independent of the model inductive bias are necessarily vacuous.
They presented an example of a disjoint augmentation, such that any two
different original samples $x_{1}$ and $x_{2}$ cannot be augmented to the same
$a$. In this case, our $K_{X}$ defined in Eqn. 3 is zero everywhere, so
apparently our analysis will not work. The observation made by Saunshi et al.
(2022) is that even for this augmentation, if there is a good model (feature
map), then good generalization can still be achieved, and therefore they
claimed that model inductive bias is irreplaceable.

The analysis in this work is compatible with their observation. Our analysis
shows that if the augmentation is good enough (such that Assumption 1 is
satisfied and the augmentation complexity is small enough, which is in fact a
strong assumption), then we can achieve good generalization only with the
augmentation and without any help of the model. In other words, the disjoint
augmentation example showcases a poor choice of augmentation, which is
orthogonal to the emphasis of this paper.

##  Appendix B Related Work

###  B.1 Learning Guarantees for Big Models

Big models with millions and even billions of parameters have been proved to
work well on a wide variety of tasks. However, their empirical success cannot
be explained by classical generalization bounds that use the VC-dimension or
the Rademacher complexity. This discrepancy was made clear by Zhang et al.
(2017), which shows that while enjoying good generalization on real tasks,
modern neural networks have sufficient capacity to fit random labels. This
challenges the classical generalization theory, which appeals to the Occam’s
razor argument that big models cannot generalize well in general. There is a
large body of research on improving classical generalization bounds for neural
networks and big models, and we highlight a few seminal works below.

##### Algorithmic stability.

This line of research attributes the good generalization of complex models to
the stability of the (random) training algorithm. There is a folklore that one
key reason why neural networks generalize so well is that they are trained by
stochastic gradient descent (SGD). Let $f(x;A,S)$ be the predictor trained by
a random training algorithm $A$ on a finite training set $S$. $A$ is called
uniformly stable (Kearns & Ron, 1997; Bousquet & Elisseeff, 2002), if
$E_{A}|f(x;A,S)-f(x;A,S^{\prime})|$ is uniformly bounded when $S$ and
$S^{\prime}$ only differ by one training sample. $f$ is guaranteed to enjoy
good generalization if $A$ is uniformly stable. Then, Hardt et al. (2016)
proved that SGD is uniformly stable under certain conditions. There are also
weaker notions of algorithmic stability (Mukherjee et al., 2006; Shalev-
Shwartz et al., 2010). It has also been shown in Li et al. (2018); Arora et
al. (2019a) that gradient methods for overparameterized models provide an
implicit regularization effect that helps generalization. However, Zhang et
al. (2021) challenged this line of work by empirically showing that SGD on
real neural networks is not very stable, even with regularization.

##### PAC-Bayes and sharpness.

Compared to algorithmic stability, the PAC-Bayes line of analysis attributes
generalization to model stability, that is the empirical risk is insensitive
to perturbations in model parameters. If the model is stable, then the
empirical risk is small within a large region around the global minima,
meaning that the risk landscape has low sharpness (Keskar et al., 2017;
Neyshabur et al., 2017) around the global minima. The PAC-Bayes generalization
bound McAllester (1999); Langford & Seeger (2001) considers a prior
distribution $P$ over the model’s function space (usually Gaussian or
uniform), and a posterior stationary distribution $Q$ found by a training
algorithm. If $D_{KL}(Q\parallel P)$ is small, then good generalization is
guaranteed. Thus, the more models there are in the support of $Q$, the better
the generalization. So we can see that a lower sharpness around the global
minima leads to better generalization. Sharpness-based bounds are shown to be
empirically tighter than other bounds in Jiang et al. (2020), and inspire a
class of training method called sharpness-aware minimization (SAM) (Foret et
al., 2021; Bartlett et al., 2023; Wen et al., 2023). Other generalization
guarantees derived from the PAC-Bayes analysis with similar ideas include low
spectral norm (Bartlett et al., 2017; Neyshabur et al., 2018), flat minima
(Hochreiter & Schmidhuber, 1997), and derandomization (Negrea et al., 2020).

One prominent class of analysis within the PAC-Bayes paradigm uses compression
to obtain tighter, and empirically non-vacuous generalization bounds. For
instance, Dziugaite & Roy (2017) proposed to directly optimize the PAC-Bayes
bound for multi-layer perceptrons, and it was improved by Rivasplata et al.
(2019) with a relaxation on the bound. Then, Arora et al. (2018) formulated
the compression framework, and showed that models with noise stability are
more compressible and thus lead to tighter bounds (also discovered in
Nagarajan & Kolter (2019)). Zhou et al. (2019) then used the compression
framework to obtain non-vacuous generalization bounds for ImageNet. More
recently, Lotfi et al. (2022) compressed models by considering model
equivariance, and Goldblum et al. (2023) verified that large language models
favor simpler sequences. Both papers proposed to use generalization bounds
based on the Kolmogorov complexity (Kolmogorov, 1963).

##### Implicit bias of overparameterized models.

This line of analysis attributes the good generalization to the implicit bias
of a specific model architecture, i.e. with certain training algorithms such
as SGD, these models can find the global minima that can generalize. First, on
the optimization side, there is a series of work that shows that gradient
descent finds the global optima of overparameterized neural networks, starting
from the results in Kawaguchi (2016) for linear neural networks, and then Li &
Yuan (2017); Du et al. (2019b; a); Arora et al. (2019b) for two-layer fully-
connected ReLU networks, and Allen-Zhu et al. (2019b) for other architectures
such as CNNs and ResNets. The analysis on overparameterized fully-connected
networks with two or more layers is then extended to generalization in Du &
Lee (2018); Allen-Zhu et al. (2019a). There is also analysis based on neural
tangent kernels (NTK) (Jacot et al., 2018), including Lee et al. (2019); Arora
et al. (2019c). Most of these work requires the neural network to be
overparameterized, i.e. the number of neurons is polynomially large comparing
to the input size (Allen-Zhu et al., 2019b). One caveat of this line of
analysis, as argued in Chizat et al. (2019), is that these overparameterized
models might fall into the “lazy training” regime, meaning that the training
dynamics of these models resemble those of a linear model. They argued that
the phenomenal success of deep neural networks is unlikely to be due to their
similarity to linear models, and Wang et al. (2020) provided a counter-
argument.

###  B.2 Theoretical Study on Augmentation Based Representation Learning

Many recent studies have been devoted to theoretically understanding
representation learning. One line of research studies the effectiveness of
contrastive learning by showing its features are optimal for linear predictor
of certain downstream tasks (Saunshi et al., 2019; Tosh et al., 2021a; b),
robust to class imbalance (Liu et al., 2021), and suitable for unsupervised
domain adaptation (Shen et al., 2022; HaoChen et al., 2022). Masked prediction
tasks have been shown to be useful for reducing the downstream sample
complexity (Lee et al., 2021) and for parameter identifiability (Liu et al.,
2022). Relating to language applications, Saunshi et al. (2019) explained why
next-word prediction can benefit sentiment classification, and Wei et al.
(2021) studied the effect of prompt tuning through the lens of implicit
Bayesian inference. Regarding the optimization in representation learning,
there have been prior works on the training dynamics and loss landscapes of
contrastive learning (Wen & Li, 2021; Jing et al., 2022; Tian, 2022), non-
contrastive learning (Tian et al., 2021; Pokle et al., 2022; Wen & Li, 2022),
and masked prediction (Xiong et al., 2020; Huang et al., 2020).

Representation learning has adopted the idea of compression. Classical
representation learning aims to extract the low-dimensional manifold on which
the data resides (Belkin & Niyogi, 2003; 2004; Bengio et al., 2004). More
recently, Yu et al. (2020) proposes maximal coding rate reduction (MCR2),
which learns an explicitly meaningful low-dimensional representation by
maximizing the rate reduction, defined as the difference in coding rate, i.e.
the number of bits needed to encode the embedding before and after a certain
partition (e.g. according to different classes). Yu et al. (2020) showed that
MCR2 is closely related to contrastive learning. Follow-up work includes Chan
et al. (2022); Dai et al. (2022).

More closely related to this work is the line of work that formulates
contrastive learning as a Laplacian operator over the augmentation graph. The
idea of studying data augmentation from a kernel perspective was first
explored in Mroueh et al. (2015); Raj et al. (2017); Dao et al. (2019). The
augmentation graph was defined in HaoChen et al. (2021) in whose Theorem 4.2
proved a generalization bound for the spectral contrastive loss that includes
a Rademacher complexity term w.r.t. $a_{1},\cdots,a_{N}$, which is incurred by
sampling only one $a$ from each $x$ (see their Definition D.3). Then, Saunshi
et al. (2022) pointed out that this model-class-free bound could be vacuous
with a hypercube construction. One thing to note is that while Saunshi et al.
(2022) claimed that the bound in HaoChen et al. (2021) is “function class
independent”, actually it is not.***Theorem 4.2 in the NeurIPS version and the
arXiv [v7] version of HaoChen et al. (2021) have different forms. We have
confirmed with the authors that this theorem does depend on the complexity of
the function class. As a response to Saunshi et al. (2022), HaoChen & Ma
(2023) included the effect of the encoder’s inductive bias into their
generalization bound. Then, Wang et al. (2023) connected contrastive learning
to message passing on the augmentation graph.

In Section 6 we highlighted two related recent papers. Wang et al. (2022b)
studied a similar RKHS as the ${{\mathcal{H}}_{\Gamma}}$ defined in this work.
In particular, the kernels $k_{x}$ and $k_{z}$ in their work are similar to
the kernels $K_{A}$ and $K_{X}$ in this work. However, their work studies the
conditional moment model, which is a completely different problem than ours.
Their bounds are also distinct from ours and are with stronger assumptions.
Nevertheless, the fact that they derived similar kernels to ours from a
completely different angle suggests a strong innate connection, which alludes
to the role of augmentation in self-supervised pretraining, namely, it defines
a ordered set of features as we articulated after Eqn. 6.

Another related work is Cabannes et al. (2023), who studied the same problem
as this work. The difference is that their work uses a pre-defined kernel,
while the kernels $K_{A}$ and $K_{X}$ in this work are completely induced by
the augmentation, i.e. there would be not be a kernel without the
augmentation. Their work studied regularized kernel regression, while this
work studies unregularized kernel regression. Moreover, the bounds derived in
their work and ours are significantly different. We would also like to credit
Cabannes et al. (2023) for studying the upstream-to-downstream distribution
shift which provides inspiring preliminary analysis. Last but not least, the
follow-up work by Zhai et al. (2024) applies the framework developed in this
paper to a more general semi-supervised kernel learning setting.

Finally, similar to Saunshi et al. (2022), we point out one common issue of
existing learning guarantees for contrastive learning: Most of them require
the samples from the same downstream class to have some type of “overlap” in
the augmentation graph, but we find that sometimes these assumptions could be
too strong to be reasonable. For instance, Wang et al. (2022a) assumed intra-
class connectivity, i.e. the subgraph of the augmentation graph given by each
downstream class is connected; Huang et al. (2023) assumed that each
downstream class has a “main part” of samples, in which any two samples can be
augmented to be close in Euclidean distance. We find these two assumptions
hard to be reasonable for any practical scenario.

##  Appendix C Variational Objectives that Extract the Top-d Eigenspace

In Proposition 4, we demonstrated the optimality of the top-$d$ eigenspace.
The question is: How to extract the top-$d$ eigenspace? We cannot do so via
explicit eigendecomposition of the kernel in most real tasks, where the kernel
is huge and sparse. Instead, we will show that several pretraining objectives
can extract the top-$d$ eigenspace assuming optimization converges to the
global minima, because these objective are uniquely optimized by the top-$d$
eigenspace.

##### Contrastive learning.

Let us warm up with contrastive learning, whose connection to kernels has been
studied in prior work, among which HaoChen et al. (2021) introduced the
spectral contrastive loss:

${\mathcal{L}}_{\textnormal{SCL}}(\hat{\Phi})=-2\mathbb{E}\left[\langle\hat{\Phi}(A),\hat{\Phi}(A^{+})\rangle\right]+\mathbb{E}\left[\langle\hat{\Phi}(A),\hat{\Phi}(A^{-})\rangle^{2}\right],$
(13)

where $(A,A^{+})$ is a positive pair that is jointly sampled from
$P_{X}(a,a^{+})$, and $A^{-}$ is a negative sample that is sampled from
${P_{\mathcal{A}}}$ independently of $A$. It has been shown that:

######  Theorem 3.

(HaoChen et al., 2021) Eqn. (13) is minimized by the top-$d$ eigenspace, and
is uniquely minimized by the top-$d$ eigenspace if
$\lambda_{d+1}<\lambda_{d}$.

Here, “minimized by the top-$d$ eigenspace” means that the objective is
minimized when $\hat{\Phi}$ spans the top-$d$ eigenspace, and “uniquely
minimized” means that this condition is necessary. The proof is based on the
following Eckart-Young-Mirsky Theorem:

######  Theorem 4 (Eckart-Young-Mirsky).

Given a matrix ${\bm{A}}\in\mathbb{R}^{m\times n}$, consider the following
optimization problem:

$\min_{{\bm{B}}\in\mathbb{R}^{m\times n},\textnormal{rank}({\bm{B}})\leq
k}\|{\bm{A}}-{\bm{B}}\|_{F}^{2}$ (14)

where $1\leq k\leq\min\\{m,n\\}$. Let the SVD of ${\bm{A}}$ and ${\bm{B}}$ be
${\bm{A}}={\bm{U}}{\bm{\Sigma}}{\bm{V}}^{\top}$ and
${\bm{B}}={\bm{U}}_{1}{\bm{\Sigma}}_{1}{\bm{V}}_{1}^{\top}$, where
${\bm{\Sigma}}={\textnormal{diag}}(\sigma_{1},\cdots,\sigma_{\min\\{m,n\\}})$,
and $\sigma_{1}\geq\sigma_{2}\geq\cdots\geq\sigma_{\min\\{m,n\\}}$. Then
${\bm{B}}$ is a minimizer of the above problem if ${\bm{U}}_{1}$ is the first
$k$ columns of ${\bm{U}}$, ${\bm{V}}_{1}$ is the first $k$ columns of
${\bm{V}}$, and ${\bm{\Sigma}}_{1}$ is the upper left $k\times k$ block of
${\bm{\Sigma}}$. Moreover, if $\sigma_{k+1}<\sigma_{k}$, then this is the
unique minimizer.

To make this paper self-contained, here we recap the proof of Theorem 3 in
HaoChen et al. (2021):

###### Proof.

Define matrix ${\bm{C}}$ as: $\hat{\Phi}={\bm{C}}{\bm{\Phi}^{*}}$, and matrix
${\bm{B}}$ as ${\bm{B}}={\bm{C}}^{\top}{\bm{C}}$. Then,
$\langle\hat{\Phi},\hat{\Phi}\rangle={\bm{\Phi}^{*}}^{\top}{\bm{B}}{\bm{\Phi}^{*}}$.
Let the SVD of ${\bm{C}}$ be ${\bm{C}}={\bm{U}}{\bm{\Sigma}}{\bm{V}}^{\top}$.
Denote
${\bm{D}}_{\lambda}={\textnormal{diag}}(\lambda_{1},\lambda_{2},\cdots)$.
Then, the spectral contrastive loss is equivalent to

${\mathcal{L}}_{\textnormal{SCL}}(\hat{\Phi})=-2\sum_{i}\lambda_{i}b_{i,i}+\sum_{i,j}b_{i,j}^{2}=\|{\bm{B}}-{\bm{D}}_{\lambda}\|_{F}^{2}-\|{\bm{D}}_{\lambda}\|_{F}^{2}.$
(15)

So it suffices to minimize $\|{\bm{B}}-{\bm{D}}_{\lambda}\|_{F}^{2}$ where
$\textnormal{rank}({\bm{B}})\leq d$. By Theorem 4, this is minimized if
${\bm{V}}(i,j)=\delta_{i,j}$ and
${\bm{\Sigma}}={\textnormal{diag}}(\sqrt{\lambda_{i}})$, which means that
$\hat{\Phi}={\bm{U}}{\textnormal{diag}}(\sqrt{\lambda_{1}},\cdots,\sqrt{\lambda_{d}}){\bm{\Phi}^{*}}$
for some orthogonal matrix ${\bm{U}}\in\mathbb{R}^{d\times d}$. Moreover, by
Thm. 4, this is the unique minimizer if $\lambda_{d+1}<\lambda_{d}$. ∎

##### CLIP.

CLIP (Radford et al., 2021) maximizes the similarity between two samples from
different modalities. There is an objective similar to the spectral
contrastive loss, which we term the spectral CLIP loss:

${\mathcal{L}}_{\textnormal{SCLIP}}(\hat{\Phi},\hat{\Xi})=-2\mathbb{E}[\langle\hat{\Phi}(A),\hat{\Xi}(X^{+})\rangle]+\mathbb{E}[\langle\hat{\Phi}(A),\hat{\Xi}(X^{-})\rangle^{2}],$
(16)

where $\hat{\Phi}$ and $\hat{\Xi}$ are encoders for different modalities (e.g.
image and text). $(A,X^{+})$ is a positive pair, while $X^{-}$ is a negative
sample independent of $A$. This loss also appeared in Wang et al. (2022b) (see
their Section 4.1). We will show that this loss also extracts the top-$d$
eigenspace.

######  Theorem 5.

Eqn. (16) is minimized by the top-$d$ eigenspace, and is uniquely minimized by
the top-$d$ eigenspace if $\lambda_{d+1}<\lambda_{d}$.

###### Proof.

Let $\hat{\Phi}={\bm{C}}{\bm{\Phi}^{*}}$ and
$\hat{\Xi}={\bm{S}}{\bm{\Psi}^{*}}$. Let
$\Xi^{\dagger}=\Gamma\hat{\Xi}={\bm{S}}{\bm{D}}_{\lambda}^{1/2}{\bm{\Phi}^{*}}$.
Then, we have

$\displaystyle\mathbb{E}[\langle\hat{\Phi}(A),\hat{\Xi}(X^{+})\rangle]$
$\displaystyle=\iint\left\langle\hat{\Phi}(a),\hat{\Xi}(x)\right\rangle
p(x|a)p(a)dadx$
$\displaystyle=\int\left\langle\hat{\Phi}(a),\int\hat{\Xi}(x)p(x|a)dx\right\rangle
p(a)da$
$\displaystyle=\int\left\langle\hat{\Phi}(a),\Xi^{\dagger}(a)\right\rangle
p(a)da=\operatorname{Tr}\left({\bm{C}}^{\top}{\bm{S}}{\bm{D}}_{\lambda}^{1/2}\right).$

We also have

$\displaystyle\mathbb{E}[\langle\hat{\Phi}(A),\hat{\Xi}(X^{-})\rangle^{2}]$
$\displaystyle=\iint\left(\sum_{k=1}^{d}\hat{\Phi}(a)_{k}\hat{\Xi}(x)_{k}\right)^{2}p(a)p(z)dadz$
$\displaystyle=\sum_{k,l=1}^{d}\left(\int\hat{\Phi}(a)_{k}\hat{\Phi}(a)_{l}p(a)da\right)\left(\int\hat{\Xi}(x)_{k}\hat{\Xi}(x)_{l}p(x)dx\right)$
$\displaystyle=\operatorname{Tr}\left({\bm{C}}{\bm{C}}^{\top}{\bm{S}}{\bm{S}}^{\top}\right)=\left\|{\bm{C}}^{\top}{\bm{S}}\right\|_{F}^{2}.$

Therefore, the spectral CLIP loss is equivalent to

${\mathcal{L}}_{\textnormal{SCLIP}}(\hat{\Phi},\hat{\Xi})=\left\|{\bm{C}}^{\top}{\bm{S}}\right\|_{F}^{2}-2\operatorname{Tr}\left({\bm{C}}^{\top}{\bm{S}}{\bm{D}}_{\lambda}^{1/2}\right)=\left\|{\bm{C}}^{\top}{\bm{S}}-{\bm{D}}_{\lambda}^{1/2}\right\|_{F}^{2}-\left\|{\bm{D}}_{\lambda}^{1/2}\right\|_{F}^{2}.$
(17)

Let the SVD of ${\bm{C}}$ be ${\bm{C}}={\bm{U}}{\bm{\Sigma}}{\bm{V}}^{\top}$.
Let ${\bm{M}}={\bm{\Sigma}}{\bm{U}}^{\top}{\bm{S}}$. Then,
${\bm{C}}^{\top}{\bm{S}}={\bm{V}}{\bm{M}}$. It is well known in linear
regression that Eqn. (17) is minimized when
${\bm{M}}=\left({\bm{V}}^{\top}{\bm{V}}\right)^{-1}{\bm{V}}^{\top}{\bm{D}}_{\lambda}^{1/2}={\bm{V}}^{\top}{\bm{D}}_{\lambda}^{1/2}$,
so

$\min_{\hat{\Phi},\hat{\Xi}}{\mathcal{L}}_{\textnormal{SCLIP}}(\hat{\Phi},\hat{\Xi})=\min_{{\bm{V}}}\left\|({\bm{I}}-{\bm{V}}{\bm{V}}^{\top}){\bm{D}}_{\lambda}^{1/2}\right\|_{F}^{2}-\left\|{\bm{D}}_{\lambda}^{1/2}\right\|_{F}^{2}.$

By Theorem 4, this is minimized when the top $d$ rows of ${\bm{V}}$ is an
orthogonal matrix, and this condition is also necessary if
$\lambda_{d+1}<\lambda_{d}$. This proves the result. Moreover, we can see that
$\hat{\Xi}$ is the minimizer if it extracts the top-$d$ eigenspace of
${\bm{\Psi}^{*}}$. ∎

##### Regularized Barlow Twins.

Barlow twins (Zbontar et al., 2021) is a non-contrastive learning method that
trains without negative sampling. We consider the following regularized Barlow
Twins loss:

${\mathcal{L}}_{\textnormal{RBT}}(\hat{\Phi};\alpha,\beta)=\sum_{k=1}^{d}\left(\mathbb{E}\left[\hat{\Phi}(A)_{k}\hat{\Phi}(A^{+})_{k}\right]-1\right)^{2}+\alpha\sum_{k\neq
l}\left(\mathbb{E}\left[\hat{\Phi}(A)_{k}\hat{\Phi}(A^{+})_{l}\right]\right)^{2}+\beta\mathbb{E}\left[\left\|\hat{\Phi}(A)\right\|_{2}^{2}\right].$
(18)

When the regularization term $\beta$ is close to 0 so that the last term is
much smaller than the first two, the above can be considered as a constrained
optimization task:

$\min_{\theta}\quad\mathbb{E}\left[\left\|\hat{\Phi}(A)\right\|_{2}^{2}\right]\qquad\text{s.t.}\qquad{\mathcal{L}}_{\textnormal{RBT}}(\hat{\Phi};\alpha,0)=0.$
(19)

######  Theorem 6.

The constrained optimization problem Eqn. (19) is minimized by the top-$d$
eigenspace, and is uniquely minimized by the top-$d$ eigenspace if
$\lambda_{d+1}<\lambda_{d}$.

###### Proof.

Define matrix ${\bm{C}}=(c_{i,j})$ as: $\hat{\Phi}={\bm{C}}{\bm{\Phi}^{*}}$.
Then, for all $k,l\in[d]$, there is

$\mathbb{E}[\hat{\Phi}(A)_{k}\hat{\Phi}(A^{+})_{l}]=\sum_{i}\lambda_{i}c_{k,i}c_{l,i}.$

Thus, ${\mathcal{L}}_{\textnormal{RBT}}(\hat{\Phi};\alpha,0)=0$ is equivalent
to ${\bm{C}}{\bm{D}}_{\lambda}{\bm{C}}^{\top}={\bm{I}}$. Meanwhile,
$\mathbb{E}\left[\left\|\hat{\Phi}(A)\right\|_{2}^{2}\right]=\|{\bm{C}}\|_{F}^{2}$.
So the optimization problem is equivalent to minimizing $\|{\bm{C}}\|_{F}^{2}$
subject to ${\bm{C}}{\bm{D}}_{\lambda}{\bm{C}}^{\top}={\bm{I}}$.

Let ${\bm{G}}={\bm{C}}\sqrt{{\bm{D}}_{\lambda}}=(g_{k,i})$, then
${\bm{G}}{\bm{G}}^{\top}=\sum_{k}{\bm{g}}_{k}{\bm{g}}_{k}^{\top}={\bm{I}}$,
where ${\bm{g}}_{k}=[g_{k,1},\cdots,g_{k,d}]$. For all $j\in[d]$, define
${\bm{M}}_{j}=\sum_{k=1}^{j}{\bm{g}}_{k}{\bm{g}}_{k}^{\top}$. Then,
$\textnormal{rank}({\bm{M}}_{j})\leq j$, and all eigenvalues of ${\bm{M}}_{j}$
belong to $[0,1]$. So the sum of eigenvalues of ${\bm{M}}_{j}$ is at most $j$,
which implies that

$\sum_{k=1}^{j}{\bm{g}}_{j}^{\top}{\bm{g}}_{j}=\sum_{k=1}^{j}\operatorname{Tr}\left({\bm{g}}_{j}^{\top}{\bm{g}}_{j}\right)=\sum_{k=1}^{j}\operatorname{Tr}\left({\bm{g}}_{j}{\bm{g}}_{j}^{\top}\right)=\operatorname{Tr}({\bm{M}}_{j})\leq
j.$

Denote $\frac{1}{\lambda_{0}}=0$. Thus, by Abel transformation, we have

$\|{\bm{C}}\|_{F}^{2}=\sum_{i}\frac{1}{\lambda_{i}}{\bm{g}}_{i}^{\top}{\bm{g}}_{i}=\sum_{j\geq
0}\left(d-\sum_{k=1}^{j}{\bm{g}}_{k}^{\top}{\bm{g}}_{k}\right)\left(\frac{1}{\lambda_{j+1}}-\frac{1}{\lambda_{j}}\right)\geq\sum_{j=0}^{d-1}(d-j)\left(\frac{1}{\lambda_{j+1}}-\frac{1}{\lambda_{j}}\right)=\sum_{i=1}^{d}\frac{1}{\lambda_{i}}.$

A sufficient condition of achieving the above equality is
$\sum_{k=1}^{j}{\bm{g}}_{k}^{\top}{\bm{g}}_{k}=j$ for all $j\in[d]$, and it is
also necessary if $\lambda_{d+1}<\lambda_{d}$. This condition is equivalent to
${\bm{g}}_{1},\cdots,{\bm{g}}_{d}$ forming an orthonormal basis of
$\mathbb{R}^{d}$, i.e. $\hat{\Phi}$ extracts the top-$d$ eigenspace. ∎

##### Maximizing the ratio trace.

Now we demonstrate the connection of contrastive learning and Barlow Twins to
maximizing the ratio trace defined in Definition 3. We start with the
regularized Barlow Twins loss, which we have shown is equivalent to minimizing
$\|{\bm{C}}\|_{F}^{2}$ subject to
${\bm{C}}{\bm{D}}_{\lambda}{\bm{C}}^{\top}={\bm{I}}$. Now, observe that
${\bm{G}}={\bm{C}}{\bm{C}}^{\top}$, and
${\bm{F}}={\bm{C}}{\bm{D}}_{\lambda}{\bm{C}}^{\top}$. Thus, regularized Barlow
Twins is essentially minimizing $\operatorname{Tr}({\bm{F}}^{-1}{\bm{G}})$,
which is similar but not equivalent to maximizing
$\operatorname{Tr}({\bm{G}}^{-1}{\bm{F}})$.

Now we consider a variant of VICReg (Bardes et al., 2022), also given by
Cabannes et al. (2023, Eqn. (2)):

${\mathcal{L}}_{\textnormal{VICReg}}(\hat{\Phi};\beta)=\left\|\mathbb{E}\left[\hat{\Phi}(A)\hat{\Phi}(A)^{\top}\right]-{\bm{I}}\right\|_{F}^{2}+\beta\mathbb{E}\left[\left\|\hat{\Phi}(A)-\hat{\Phi}(A^{+})\right\|_{2}^{2}\right].$
(20)

When the regularization term $\beta$ is close to 0, we have
$\mathbb{E}\left[\hat{\Phi}(A)\hat{\Phi}(A)^{\top}\right]={\bm{G}}={\bm{I}}$.
And we minimize
$\mathbb{E}\left[\left\|\hat{\Phi}(A)-\hat{\Phi}(A^{+})\right\|_{2}^{2}\right]$,
which is equivalent to
$-2\mathbb{E}\left[\operatorname{Tr}\left(\hat{\Phi}(A)\hat{\Phi}(A^{+})\right)\right]+c=-2\operatorname{Tr}({\bm{F}})+c$
for some constant $c$. Thus, this objective is equivalent to maximizing
$\operatorname{Tr}({\bm{G}}^{-1}{\bm{F}})$, the ratio trace, when
$\beta\approx 0$. And when $\beta=1$, Eqn. (20) is equivalent to the spectral
contrastive loss.

##  Appendix D Proofs for Section 2

#####  $\Gamma^{*}\Gamma$ and $\Gamma\Gamma^{*}$ are integral operators.

$\begin{cases}({\Gamma_{a\rightarrow x}}{\Gamma_{x\rightarrow
a}}f)(x)=(\Gamma^{*}\Gamma
f)(x)=\int{K_{X}}(x,x^{\prime})f(x^{\prime})p(x^{\prime})dx^{\prime};\\\
({\Gamma_{x\rightarrow a}}{\Gamma_{a\rightarrow
x}}g)(a)=(\Gamma\Gamma^{*}g)(a)=\int{K_{A}}(a,a^{\prime})g(a^{\prime})p(a^{\prime})da^{\prime}.\end{cases}$
(21)

###### Proof.

We only show the first equation, and the second one can be proved in the same
way.

$\displaystyle(\Gamma^{*}\Gamma f)(x)$ $\displaystyle=\Gamma^{*}\left(\int
f(x^{\prime})p(x^{\prime}|a)dx^{\prime}\right)=\int\left(\int
f(x^{\prime})p(x^{\prime}|a)dx^{\prime}\right)p(a|x)da$ $\displaystyle=\iint
f(x^{\prime})p(a|x)p(x^{\prime}|a)dadx^{\prime}=\iint
f(x^{\prime})\frac{p(a|x)p(a|x^{\prime})}{p(a)}p(x^{\prime})dadx^{\prime}$
$\displaystyle=\int{K_{X}}(x,x^{\prime})f(x^{\prime})p(x^{\prime})dx^{\prime}.$

∎

##### Proposition 1 (Duality).

$\Gamma\Gamma^{*}$ shares the same non-zero eigenvalues as $\Gamma^{*}\Gamma$,
and there exist eigenfunctions $\\{\phi_{i}\\}$ of $\Gamma\Gamma^{*}$ that
form an orthonormal basis of ${L^{2}(P_{\mathcal{A}})}$, such that for any
$\lambda_{i}>0$,

$\psi_{i}=\lambda_{i}^{-1/2}\Gamma^{*}\phi_{i}\quad\text{and}\quad\phi_{i}=\lambda_{i}^{-1/2}\Gamma\psi_{i},$
(22)

and we also have the following spectral decomposition of the Radon-Nikodym
derivative:

$\frac{dP_{{\mathcal{A}}{\mathcal{X}}}}{d({P_{\mathcal{A}}}\otimes{P_{\mathcal{X}}})}=\frac{p(a,x)}{p(a)p(x)}=\sum_{i}\lambda_{i}^{1/2}\phi_{i}(a)\psi_{i}(x).$
(23)

###### Proof.

Suppose $\lambda_{i},\psi_{i}(x)$ is a pair of eigenvalue and eigenfunction of
$\Gamma^{*}\Gamma$, and $\lambda_{i}>0$. Then, we have
$\Gamma\Gamma^{*}\Gamma\psi_{i}=\lambda_{i}\Gamma\psi_{i}$, which means that
$\Gamma\psi_{i}$ is an eigenfunction of $\Gamma\Gamma^{*}$ with eigenvalue
$\lambda_{i}$. The $\lambda_{i}^{-1/2}$ is used for normalization. To see
this, let $\phi_{i}=\lambda_{i}^{-1/2}\Gamma\psi_{i}$. Then, we have

$\displaystyle\langle\phi_{i},\phi_{j}\rangle_{P_{\mathcal{A}}}$
$\displaystyle=\lambda_{i}^{-1/2}\lambda_{j}^{-1/2}\langle\Gamma\psi_{i},\Gamma\psi_{j}\rangle_{P_{\mathcal{A}}}$
$\displaystyle=\lambda_{i}^{-1/2}\lambda_{j}^{-1/2}\langle\Gamma^{*}\Gamma\psi_{i},\psi_{j}\rangle_{P_{\mathcal{X}}}$
$\displaystyle=\lambda_{i}^{-1/2}\lambda_{j}^{-1/2}\langle\lambda_{i}\psi_{i},\psi_{j}\rangle_{P_{\mathcal{X}}}=\delta_{i,j}.$

We can prove the reverse direction similarly. And for any fixed $x$, there is

$\left\langle\frac{p(a,x)}{p(a)p(x)},\phi_{i}\right\rangle_{P_{\mathcal{A}}}=\int\frac{p(a,x)}{p(a)p(x)}\phi_{i}(a)p(a)da=\int
p(a|x)\phi_{i}(a)da=\sqrt{\lambda_{i}}\psi_{i}(x).$ (24)

which implies Eqn. (23). ∎

##### Basic properties of ${{\mathcal{H}}_{\Gamma}}$.

  1. (i)

${K_{X}}$ is the reproducing kernel of ${{\mathcal{H}}_{\Gamma}}$, such that
for all $f\in{{\mathcal{H}}_{\Gamma}}$, $f(x)=\langle
f,{K_{X}}(x,\cdot)\rangle_{{{\mathcal{H}}_{\Gamma}}}$.

  2. (ii)

${{\mathcal{H}}_{\Gamma}}=R(\Gamma^{*})$.

  3. (iii)

${{\mathcal{H}}_{\Gamma}}$ is isometric to
${\textnormal{span}}({\left\\{\phi_{i}\right\\}}_{\lambda_{i}>0})$, a subspace
of ${L^{2}(P_{\mathcal{A}})}$, and
$\|f\|_{{{\mathcal{H}}_{\Gamma}}}=\inf_{g:f=\Gamma^{*}g}\|g\|_{P_{\mathcal{A}}}$.

  4. (iv)

For any $f^{*}\in{\mathcal{F}}_{B}(\Gamma;\epsilon)\subset R(\Gamma^{*})$, let
$f^{*}=\sum_{i}u_{i}\psi_{i}$. Define
$g_{0}:=\sum_{i}\lambda_{i}^{-1/2}u_{i}\phi_{i}$. Then, $g_{0}$ must satisfy
Eqn. 1, so we can choose $g^{*}=g_{0}$, in which case Eqn. 1 is equivalent to:

$\langle
g^{*},(I-\Gamma\Gamma^{*})g^{*}\rangle_{P_{\mathcal{A}}}\leq\epsilon\|g^{*}\|_{P_{\mathcal{A}}}^{2}\;\Leftrightarrow\;\sum_{i}\frac{1-\lambda_{i}}{\lambda_{i}}u_{i}^{2}\leq\epsilon\sum_{i}\frac{1}{\lambda_{i}}u_{i}^{2}.$
(25)

###### Proof.

  1. (i)

First, note that
${{\mathcal{H}}_{\Gamma}}={\left\\{\sum_{i:\lambda_{i}>0}a_{i}{\bm{e}}_{i}\;\middle|\;\sum_{i}a_{i}^{2}<\infty\right\\}}$
where ${\bm{e}}_{i}=\lambda_{i}^{-1/2}\psi_{i}$, so it is isomorphic to
$\ell^{2}((a_{i})_{i:\lambda_{i}>0})$ and is thus a Hilbert space. Then,
${K_{X}}(x,x^{\prime})=\sum_{i}\lambda_{i}\psi_{i}(x)\psi_{i}(x^{\prime})$.
For any $f\in{{\mathcal{H}}_{\Gamma}}$, let $f=\sum_{i}u_{i}\psi_{i}$, then

$\langle
f(x^{\prime}),{K_{X}}(x,x^{\prime})\rangle_{{{\mathcal{H}}_{\Gamma}}}=\sum_{i}\frac{1}{\lambda_{i}}u_{i}(\lambda_{i}\psi_{i}(x))=\sum_{i}u_{i}\psi_{i}(x)=f(x).$

  2. (ii)

For any $f=\sum_{i}u_{i}\psi_{i}\in{{\mathcal{H}}_{\Gamma}}$, there is
$\sum_{i}\lambda_{i}^{-1}u_{i}^{2}<\infty$ by definition. So for any
$\lambda_{i}=0$, there must be $u_{i}=0$. Let
$g=\sum_{i}\lambda_{i}^{-1/2}u_{i}\psi_{i}$. Then,
$\|g\|_{P_{\mathcal{A}}}^{2}=\sum_{i}\lambda_{i}^{-1}u_{i}^{2}<\infty$,
meaning that $g\in{L^{2}(P_{\mathcal{A}})}$. And there is $f=\Gamma^{*}g$, so
$f\in R(\Gamma^{*})$, which implies that ${{\mathcal{H}}_{\Gamma}}\subseteq
R(\Gamma^{*})$. Meanwhile, for any $f=\Gamma^{*}g\in R(\Gamma^{*})$, let
$g=\sum_{i}v_{i}\phi_{i}$, then $\sum_{i}v_{i}^{2}<\infty$. Then,
$f=\sum_{i}\lambda_{i}^{1/2}v_{i}\psi_{i}$ by duality, so
$\sum_{i}\lambda_{i}^{-1}(\lambda_{i}^{1/2}v_{i})^{2}<\infty$, meaning that
$f\in{{\mathcal{H}}_{\Gamma}}$, so
$R(\Gamma^{*})\subseteq{{\mathcal{H}}_{\Gamma}}$.

  3. (iii)

For any $f=\sum_{i}u_{i}\psi_{i}\in{{\mathcal{H}}_{\Gamma}}$, let
$g=\sum_{i}\lambda_{i}^{-1/2}u_{i}\psi_{i}$. By the proof of (ii) we know that
$f\mapsto g$ is bijective, and
$\|f\|_{{\mathcal{H}}_{\Gamma}}=\|g\|_{P_{\mathcal{A}}}$. Moreover,
$g\in{\textnormal{span}}({\left\\{\phi_{i}\right\\}}_{\lambda_{i}>0})$.

  4. (iv)

Let $g^{*}=\sum_{i}v_{i}\phi_{i}$. Then, since we have
$f^{*}=\Gamma^{*}g^{*}=\sum_{i}\lambda_{i}^{1/2}v_{i}\psi_{i}$, for any
$\lambda_{i}>0$, there is $v_{i}=\lambda^{-1/2}u_{i}$; and for any
$\lambda_{i}=0$, there is $u_{i}=0$. Let $g^{*}=g_{0}+g_{1}$, where
$g_{0}=\sum_{i}\lambda_{i}^{-1/2}u_{i}\phi_{i}$, and $g_{1}\perp g_{0}$ and
$\Gamma^{*}g_{1}=0$. By duality, $\Gamma\Gamma^{*}g_{0}$ belongs to the linear
span of $\\{\phi_{i}\\}_{\lambda_{i}>0}$, so
$g_{1}\perp\Gamma\Gamma^{*}g_{0}$. As we will show later, Eqn. (1) is
equivalent to Eqn. 25, which is the random walk normalized Laplacian over the
augmentation graph (Chung, 1997, Section 1.2). This is equivalent to $\langle
g^{*},(I-\Gamma\Gamma^{*})g^{*}\rangle_{P_{\mathcal{A}}}\leq\epsilon\|g^{*}\|_{P_{\mathcal{A}}}^{2}$,
which is further equivalent to $\langle
g_{0},(I-\Gamma\Gamma^{*})g_{0}\rangle_{P_{\mathcal{A}}}+\|g_{1}\|_{P_{\mathcal{A}}}^{2}\leq\epsilon(\|g_{0}\|_{P_{\mathcal{A}}}^{2}+\|g_{1}\|_{P_{\mathcal{A}}}^{2})$
(note that $\Gamma\Gamma^{*}g_{1}=0$). This implies that $\langle
g_{0},(I-\Gamma\Gamma^{*})g_{0}\rangle_{P_{\mathcal{A}}}\leq\epsilon\|g_{0}\|_{P_{\mathcal{A}}}^{2}$,
i.e. $g_{0}$ satisfies Eqn. (1). Since we also have $f^{*}=\Gamma^{*}g_{0}$,
$g_{0}$ must satisfy Assumption 1, so we can choose $g^{*}=g_{0}$.

Next, to show the equivalence to Eqn. 25, We just need to show that $\langle
g^{*},(I-\Gamma\Gamma^{*})g^{*}\rangle_{P_{\mathcal{A}}}=\frac{1}{2}\mathbb{E}_{X\sim{P_{\mathcal{X}}}}\mathbb{E}_{A,A^{\prime}\sim
p(\cdot|X)}\left[(g^{*}(A)-g^{*}(A^{\prime}))^{2}\right]$. And indeed, we
have:

$\displaystyle\langle
g^{*},(I-\Gamma\Gamma^{*})g^{*}\rangle_{P_{\mathcal{A}}}$
$\displaystyle=\left\langle g^{*},g^{*}-\int
g^{*}(a^{\prime}){K_{A}}(\cdot,a^{\prime})p(a^{\prime})da^{\prime}\right\rangle_{P_{\mathcal{A}}}$
$\displaystyle=\|g^{*}\|_{P_{\mathcal{A}}}^{2}-\iint
g(a)g(a^{\prime})\frac{\int
p(a|x)p(a^{\prime}|x)p(x)dx}{p(a)p(a^{\prime})}p(a^{\prime})p(a)dada^{\prime}$
$\displaystyle=\frac{1}{2}\mathbb{E}[g^{*}(A)^{2}]+\frac{1}{2}\mathbb{E}[g^{*}(A^{\prime})^{2}]-\frac{1}{2}\mathbb{E}_{X\sim{P_{\mathcal{X}}}}\mathbb{E}_{A,A^{\prime}\sim
p(\cdot|X)}[2g^{*}(A)g^{*}(A^{\prime})]$
$\displaystyle=\frac{1}{2}\mathbb{E}_{X\sim{P_{\mathcal{X}}}}\mathbb{E}_{A,A^{\prime}\sim
p(\cdot|X)}\left[(g^{*}(A)-g^{*}(A^{\prime}))^{2}\right],$

as desired. ∎

##  Appendix E Proofs for Section 3

###  E.1 Local Gaussian Complexity and Localized Rademacher Complexity

We first provide the definition of the two complexities we will use in our
analysis. For a function $f$, let
$\|f\|_{n}^{2}:=\frac{1}{n}\sum_{i=1}^{n}f(\tilde{x}_{i})^{2}$ be its mean on
the downstream samples.

######  Definition 4.

(Wainwright, 2019, Eqns. (13.16) & (14.3)) For any $B,\epsilon>0$, define

${\mathcal{F}}_{0}:={\left\\{f_{1}-f_{2}\;\middle|\;f_{i}\in{\mathcal{H}}_{{\hat{\Psi}}},\|f_{i}\|_{{{\mathcal{H}}_{\Gamma}}}\leq\frac{B}{\sqrt{1-\epsilon}}\right\\}}={\left\\{f\in{\mathcal{H}}_{{\hat{\Psi}}}\;\middle|\;\|f\|_{{{\mathcal{H}}_{\Gamma}}}\leq\frac{2B}{\sqrt{1-\epsilon}}\right\\}}.$
(26)

Then, the local Gaussian complexity around ${f_{\hat{\Psi}}}$ at scale
$\delta>0$ is given by

${\mathcal{G}}_{n}(\delta;{\mathcal{F}}_{0}):=\underset{\omega_{1},\cdots,\omega_{n}}{\mathbb{E}}\left[\sup_{f\in{\mathcal{F}}_{0},\|f\|_{n}\leq\delta}\left|\frac{1}{n}\sum_{i=1}^{n}\omega_{i}f(\tilde{x}_{i})\right|\right],$
(27)

where $\omega_{1},\cdots,\omega_{n}$ are i.i.d. ${\mathcal{N}}(0,1)$ variates.
And define

${\mathcal{F}}_{*}:={\left\\{f=f_{1}+\alpha
f^{*}\;\middle|\;\alpha\in[-1,1],f_{1}\in{\mathcal{H}}_{{\hat{\Psi}}},\|f_{1}\|_{{{\mathcal{H}}_{\Gamma}}}\leq\frac{B}{\sqrt{1-\epsilon}}\right\\}}.$
(28)

Then, the localized population Rademacher complexity of radius $\delta>0$ is
given by

${\bar{\mathfrak{R}}}_{n}(\delta;{\mathcal{F}}_{*}):=\underset{\sigma_{1},\cdots,\sigma_{n},x_{1},\cdots,x_{n}}{\mathbb{E}}\left[\sup_{f\in{\mathcal{F}}_{*},\|f\|_{P_{\mathcal{X}}}\leq\delta}\left|\frac{1}{n}\sum_{i=1}^{n}\sigma_{i}f(x_{i})\right|\right],$
(29)

where $\sigma_{1},\cdots,\sigma_{n}$ are i.i.d. Rademacher variables taking
values in ${\left\\{-1,+1\right\\}}$ equiprobably.

Our master plan is to apply Theorems 13.13 and 14.1 of Wainwright (2019) to
${f_{\hat{\Psi}}}=\Gamma^{*}(\Pi_{{\hat{\Phi}}}g^{*})$, where
$\Pi_{{\hat{\Phi}}}$ is the projection operator onto ${\hat{\Phi}}$ in
${L^{2}(P_{\mathcal{X}})}$, and ${f_{\hat{\Psi}}}$ is the projection of
$f^{*}$ onto ${\mathcal{H}}_{{\hat{\Psi}}}$ w.r.t.
$\langle\cdot,\cdot\rangle_{{\mathcal{H}}_{\Gamma}}$. Therefore, we need to
bound ${\mathcal{G}}_{n}(\delta;{\mathcal{F}}_{0})$ and
${\bar{\mathfrak{R}}}_{n}(\delta;{\mathcal{F}}_{*})$. We start with the
following uniform bound:

######  Proposition 7.

If $f=\Gamma^{*}g$, and $\|g\|_{P_{\mathcal{A}}}\leq T$, then
$|f(x)|\leq\kappa T$ for all $x$.

###### Proof.

By Eqn. (4), we have
$p(a|x)=\sum_{i}\sqrt{\lambda_{i}}\phi_{i}(a)\psi_{i}(x)p(a)$. For any
$g=\sum_{i}u_{i}\phi_{i}\in{L^{2}(P_{\mathcal{A}})}$ such that
$\|g\|_{P_{\mathcal{A}}}\leq T$, $(\Gamma^{*}g)(x)=\int
g(a)p(a|x)da=\sum_{i}\sqrt{\lambda_{i}}u_{i}\psi_{i}(x)$. Then, by Cauchy-
Schwarz inequality, we have for all $x$,
$f(x)^{2}=(\Gamma^{*}g)(x)^{2}\leq(\sum_{i}\lambda_{i}\psi_{i}(x)^{2})(\sum_{i}u_{i}^{2})\leq\kappa^{2}T^{2}$.
∎

This proposition immediately implies that $f^{*}$ and ${f_{\hat{\Psi}}}$ are
uniformly bounded:

######  Corollary 8.

For any $f^{*}\in{\mathcal{F}}_{B}(\Gamma;\epsilon)$, Eqn. (6) ensures that
$\|g^{*}\|_{P_{\mathcal{A}}}^{2}\leq\frac{B^{2}}{1-\epsilon}$, so
$|f^{*}(x)|\leq\frac{\kappa B}{\sqrt{1-\epsilon}}$ for all $x$. Moreover,
$\|\Pi_{{\hat{\Phi}}}g^{*}\|_{P_{\mathcal{A}}}\leq\|g^{*}\|_{P_{\mathcal{A}}}$
implies that
$\|{f_{\hat{\Psi}}}\|_{{\mathcal{H}}_{\Gamma}}\leq\frac{B}{\sqrt{1-\epsilon}}$,
and $|{f_{\hat{\Psi}}}(x)|\leq\frac{\kappa B}{\sqrt{1-\epsilon}}$ for all $x$.

We will also use the following simple result in linear algebra:

######  Lemma 9.

Let ${\bm{D}}_{\lambda}={\textnormal{diag}}(\lambda_{1},\lambda_{2},\cdots)$
where $\lambda_{1}\geq\lambda_{2}\geq\cdots\geq 0$ and $\lambda_{i}\rightarrow
0$. Let ${\bm{Q}}$ be a matrix with $d$ rows that are unit vectors. Then,
$\operatorname{Tr}({\bm{Q}}{\bm{D}}_{\lambda}{\bm{Q}}^{\top})\leq\lambda_{1}+\cdots+\lambda_{d}$.

###### Proof.

Let ${\bm{q}}_{i}$ be the $i$-th column of ${\bm{Q}}$. Then for all $j\in[d]$,
there is $\sum_{i=1}^{j}{\bm{q}}_{i}^{\top}{\bm{q}}_{i}\leq j$. And for $j>d$,
$\sum_{i=1}^{j}{\bm{q}}_{i}^{\top}{\bm{q}}_{i}\leq d$. Thus, using Abel
transformation, we have

$\displaystyle\operatorname{Tr}({\bm{Q}}{\bm{D}}_{\lambda}{\bm{Q}}^{\top})$
$\displaystyle=\operatorname{Tr}({\bm{D}}_{\lambda}{\bm{Q}}^{\top}{\bm{Q}})=\sum_{i=1}^{\infty}\lambda_{i}{\bm{q}}_{i}^{\top}{\bm{q}}_{i}=\sum_{j=1}^{\infty}\left(\sum_{i=1}^{j}{\bm{q}}_{i}^{\top}{\bm{q}}_{i}\right)(\lambda_{j}-\lambda_{j+1})\leq\sum_{i=1}^{d}\lambda_{i},$

which proves the assertion. ∎

This result has many implications. For instance, for any rank-$d$ subspace of
${{\mathcal{H}}_{\Gamma}}$, its trace (the sum of its eigenvalues) is at most
$S_{\lambda}(d)$.

Now, let us bound ${\mathcal{G}}_{n}(\delta;{\mathcal{F}}_{0})$ with the
following result:

######  Lemma 10.

(Application of Wainwright (2019, Lemma 13.22)) Let ${\mathcal{H}}$ be an RKHS
with reproducing kernel $K$. Given samples
$\tilde{x}_{1},\cdots,\tilde{x}_{n}$, let ${\bm{K}}$ be the normalized kernel
matrix with entries ${\bm{K}}(i,j)=K(\tilde{x}_{i},\tilde{x}_{j})/n$. Let
$\mu_{1}\geq\cdots\geq\mu_{n}\geq 0$ be the eigenvalues of ${\bm{K}}$. Then
for all $\delta>0$, we have

$\mathbb{E}\left[\sup_{\|f\|_{\mathcal{H}}\leq
T,\|f\|_{n}\leq\delta}\left|\frac{1}{n}\sum_{i=1}^{n}\omega_{i}f(\tilde{x}_{i})\right|\right]\leq\sqrt{\frac{2}{n}}\sqrt{\sum_{j=1}^{n}\min\\{\delta^{2},\mu_{j}T^{2}\\}},$
(30)

where $\omega_{1},\cdots,\omega_{n}$ are i.i.d. ${\mathcal{N}}(0,1)$ variates.
We apply this result to $K={K_{X}}$. By Definition 2, all elements on the
diagonal of ${\bm{K}}$ are at most $\kappa^{2}/n$, so
$\sum_{j}\mu_{j}=\operatorname{Tr}({\bm{K}})\leq\kappa^{2}$. Thus, we have

${\mathcal{G}}_{n}(\delta;{\mathcal{F}}_{0})\leq\sqrt{\frac{8\kappa^{2}B^{2}}{n(1-\epsilon)}}\qquad\text{for
any }{\mathcal{H}}_{{\hat{\Psi}}}.$ (31)

Regarding ${\bar{\mathfrak{R}}}_{n}(\delta;{\mathcal{F}}_{*})$,
${\mathcal{F}}_{*}$ is also a subset of RKHS $\hat{{\mathcal{H}}}_{*}$, which
is the linear span of $\hat{\Psi}$ and $f^{*}$, and is a subspace of
${{\mathcal{H}}_{\Gamma}}$ whose rank is at most $(d+1)$. By Lemma 9, the sum
of eigenvalues of $\hat{{\mathcal{H}}}_{*}$ is at most $S_{\lambda}(d+1)$.
Since $\|f^{*}\|_{{{\mathcal{H}}_{\Gamma}}}\leq\frac{B}{\sqrt{1-\epsilon}}$,
all $f\in{\mathcal{F}}_{*}$ satisfy
$\|f\|_{{{\mathcal{H}}_{\Gamma}}}\leq\frac{2B}{\sqrt{1-\epsilon}}$. So we have
the following bound for ${\bar{\mathfrak{R}}}_{n}(\delta;{\mathcal{F}}_{*})$:

######  Lemma 11.

(Application of Wainwright (2019, Corollary 14.5)) Let
$\mu_{1},\mu_{2},\cdots$ be the eigenvalues of the RKHS
$\hat{{\mathcal{H}}}_{*}$. Since
${\textnormal{rank}}(\hat{{\mathcal{H}}}_{*})\leq{\textnormal{rank}}({\mathcal{H}}_{{\hat{\Psi}}})+1$,
we have

${\bar{\mathfrak{R}}}_{n}(\delta;{\mathcal{F}}_{*})\leq\sqrt{\frac{2}{n}}\sqrt{\sum_{j=1}^{\infty}\min{\left\\{\delta^{2},\frac{4\mu_{j}B^{2}}{1-\epsilon}\right\\}}}\leq\sqrt{\frac{8B^{2}}{n(1-\epsilon)}S_{\lambda}(d+1)}\quad\text{if
}{\textnormal{rank}}({\mathcal{H}}_{{\hat{\Psi}}})\leq d,$ (32)

and for an arbitrary ${\mathcal{H}}_{{\hat{\Psi}}}$, we can simply replace
$S_{\lambda}(d+1)$ with $S_{\lambda}$.

###  E.2 Proofs

##### Lemma 2.

Suppose $\nu_{1},\cdots,\nu_{n}$ are i.i.d. ${\mathcal{N}}(0,\sigma^{2})$
variates. If $\hat{\Phi}$ has $d$ dimensions ($d$ can be $\infty$), then we
have the following uniform bound over all
$f^{*}=\Gamma^{*}g^{*}\in{\mathcal{F}}_{B}(\Gamma;\epsilon)$:

$\displaystyle\underset{\tilde{x}_{i},\nu_{i}}{{\mathbb{P}}}\left[\forall
f^{*}\in{\mathcal{F}}_{B}(\Gamma;\epsilon),\|\hat{f}-f^{*}\|_{P_{\mathcal{X}}}^{2}\leq
9\|{f_{\hat{\Psi}}}-f^{*}\|_{P_{\mathcal{X}}}^{2}+\frac{c_{0}\kappa(B^{2}+\sigma
B)}{1-\epsilon}\sqrt{\frac{S_{\lambda}(d+1)}{n}}\right]$ $\displaystyle\geq$
$\displaystyle
1-c_{1}\exp\left(-\frac{c_{2}\sqrt{2nS_{\lambda}(d+1)}}{\kappa}\right)-\exp\left(-\sqrt{\frac{2n\kappa^{2}B^{2}}{1-\epsilon}}\right),$

where ${f_{\hat{\Psi}}}=\Gamma^{*}(\Pi_{{\hat{\Phi}}}g^{*})$ is the projection
of $f^{*}$ onto ${\mathcal{H}}_{{\hat{\Psi}}}$ w.r.t.
$\langle\cdot,\cdot\rangle_{{\mathcal{H}}_{\Gamma}}$, and $c_{0},c_{1},c_{2}$
are universal constants. Moreover,
$S_{\lambda}(d+1)\leq\min{\left\\{d+1,\kappa^{2}\right\\}}$.

###### Proof.

By Proposition 7, all functions in ${\mathcal{F}}_{*}$ are $b$-uniformly
bounded, with $b=\frac{2\kappa B}{\sqrt{1-\epsilon}}$. And obviously
${\mathcal{F}}_{*}$ is star-shaped, meaning that for all
$f\in{\mathcal{F}}_{*}$ and all $\beta\in[0,1]$, $\beta
f\in{\mathcal{F}}_{*}$. Let
$t^{2}=b\cdot\sqrt{\frac{8B^{2}}{n(1-\epsilon)}S_{\lambda}(d+1)}\geq
b{\bar{\mathfrak{R}}}_{n}(\delta;{\mathcal{F}}_{*})$. Then, by Wainwright
(2019, Theorem 14.1), we have

${\mathbb{P}}\left[\left|\|f\|_{n}^{2}-\|f\|_{P_{\mathcal{X}}}^{2}\right|\geq\frac{1}{2}\|f\|_{P_{\mathcal{X}}}^{2}+\frac{t^{2}}{2}\right]\leq
c_{1}\exp\left(-c_{2}\frac{nt^{2}}{b^{2}}\right)\qquad\text{for all
}f\in{\mathcal{F}}_{*}$ (33)

for universal constant $c_{1},c_{2}$. We know that
$\hat{f}-f^{*}\in{\mathcal{F}}_{*}$ and
${f_{\hat{\Psi}}}-f^{*}\in{\mathcal{F}}_{*}$, which means that

$\displaystyle{\mathbb{P}}\left[\left(\|\hat{f}-f^{*}\|_{P_{\mathcal{X}}}^{2}\geq
2\|\hat{f}-f^{*}\|_{n}^{2}+t^{2}\right)\vee\left(\|{f_{\hat{\Psi}}}-f^{*}\|_{n}^{2}\geq\frac{3}{2}\|{f_{\hat{\Psi}}}-f^{*}\|_{P_{\mathcal{X}}}^{2}+\frac{t^{2}}{2}\right)\right]$
(34) $\displaystyle\leq c_{1}\exp\left(-c_{2}\frac{nt^{2}}{b^{2}}\right).$

Let $\delta_{n}^{2}=2\sigma\sqrt{\frac{8\kappa^{2}B^{2}}{n(1-\epsilon)}}$. By
Lemma 10, we have $\delta_{n}^{2}\geq
2\sigma{\mathcal{G}}_{n}(\delta_{n};{\mathcal{F}}_{0})$. And
${\mathcal{F}}_{0}$ is also star-shaped. Thus, by setting $\gamma=1/2$ in
Wainwright (2019, Theorem 13.13), we have†††Please refer to the proof of
Wainwright (2019, Theorem 13.13) for removing the universal constants in this
theorem.

${\mathbb{P}}\left[\|\hat{f}-f^{*}\|_{n}^{2}\geq
3\|{f_{\hat{\Psi}}}-f^{*}\|_{n}^{2}+32\delta_{n}^{2}\right]\leq\exp\left(-\frac{n\delta_{n}^{2}}{2\sigma^{2}}\right).$
(35)

Combining the two inequalities above with the union bound, we obtain the
result. ∎

Now we prove Lemma 3. Without loss of generality, suppose
$h_{1},\cdots,h_{d^{\prime}}$ are linearly independent. Let
$\hat{{\mathcal{H}}}_{d^{\prime}}:={\textnormal{span}}{\left\\{h_{1},\cdots,h_{d^{\prime}}\right\\}}$.
Let $g^{*}=g_{0}+\beta g_{1}$, where
$g_{0}=\Pi_{\hat{{\mathcal{H}}}_{d^{\prime}}}g^{*}$, $g_{1}\perp g_{0}$, and
$\|g_{1}\|_{P_{\mathcal{A}}}=1$. So by Lemma 9, we have:

######  Proposition 12.

$\|\Gamma^{*}({\bm{G}}_{h}^{-1/2}h_{1})\|_{P_{\mathcal{X}}}^{2}+\cdots+\|\Gamma^{*}({\bm{G}}_{h}^{-1/2}h_{d^{\prime}})\|_{P_{\mathcal{X}}}^{2}+\|\Gamma^{*}g_{1}\|_{P_{\mathcal{X}}}^{2}\leq\lambda_{1}+\cdots+\lambda_{d^{\prime}+1}$.

###### Proof.

Let
$\left[{\bm{G}}_{h}^{-1/2}h_{1},\cdots,{\bm{G}}_{h}^{-1/2}h_{d^{\prime}},g_{1}\right]={\bm{Q}}{\bm{\Phi}^{*}}$,
where ${\bm{Q}}$ is a matrix with $(d^{\prime}+1)$ orthonormal rows. Then,
$\left[\Gamma^{*}({\bm{G}}_{h}^{-1/2}h_{1}),\cdots,\Gamma^{*}({\bm{G}}_{h}^{-1/2}h_{d^{\prime}}),\Gamma^{*}g_{1}\right]={\bm{Q}}{{\bm{D}}_{\lambda}}^{1/2}{\bm{\Phi}^{*}}$.
Thus, we have

$\|\Gamma^{*}({\bm{G}}_{h}^{-1/2}h_{1})\|_{P_{\mathcal{X}}}^{2}+\cdots+\|\Gamma^{*}({\bm{G}}_{h}^{-1/2}h_{d^{\prime}})\|_{P_{\mathcal{X}}}^{2}+\|\Gamma^{*}g_{1}\|_{P_{\mathcal{X}}}^{2}=\operatorname{Tr}({\bm{Q}}{{\bm{D}}_{\lambda}}{\bm{Q}}^{\top}).$

Then, applying Lemma 9 completes the proof. ∎

######  Remark.

This proposition is the functional version of Fan (1949, Theorem 1).

Notice that
$\|\Gamma^{*}({\bm{G}}_{h}^{-1/2}h_{1})\|_{P_{\mathcal{X}}}^{2}+\cdots+\|\Gamma^{*}({\bm{G}}_{h}^{-1/2}h_{d^{\prime}})\|_{P_{\mathcal{X}}}^{2}=\operatorname{Tr}({\bm{G}}_{h}^{-1/2}{\bm{F}}_{h}{\bm{G}}_{h}^{-1/2})=\operatorname{Tr}({\bm{G}}_{h}^{-1}{\bm{F}}_{h})$.
With this, we can prove Lemma 3:

##### Lemma 3.

For any $f^{*}\in{\mathcal{F}}_{B}(\Gamma;\epsilon)$, there is

$\|{f_{\hat{\Psi}}}-f^{*}\|_{P_{\mathcal{X}}}^{2}\leq\frac{\tau^{2}}{1-\tau^{2}}\frac{\tau+\epsilon}{1-\epsilon}B^{2}.$

###### Proof.

Let $\alpha^{2}=\|g_{0}\|_{P_{\mathcal{A}}}^{2}$, and
$\beta^{2}=\|g_{0}-g^{*}\|_{P_{\mathcal{A}}}^{2}$. By Corollary 8,
$\alpha^{2}+\beta^{2}\leq\frac{B^{2}}{1-\epsilon}$. Eqn. (6) implies that

$(1-\epsilon)(\alpha^{2}+\beta^{2})\leq\|\Gamma^{*}(g_{0}+\beta
g_{1})\|_{P_{\mathcal{X}}}^{2}\leq\alpha^{2}+\beta^{2}\tau^{2}+2\alpha\beta\tau,$

since
$\|\Gamma^{*}g_{0}\|_{P_{\mathcal{X}}}^{2}\leq\|g_{0}\|_{P_{\mathcal{A}}}^{2}=\alpha^{2}$,
and $\|\Gamma^{*}g_{1}\|_{P_{\mathcal{X}}}^{2}\leq\tau^{2}$ by Proposition 12.
Thus,

$(1-\tau^{2})\beta^{2}\leq\epsilon(\alpha^{2}+\beta^{2})+2\alpha\beta\tau\leq(\epsilon+\tau)(\alpha^{2}+\beta^{2})\leq(\epsilon+\tau)\frac{B^{2}}{1-\epsilon}.$

Thus, we have
$\|{f_{\hat{\Psi}}}-f^{*}\|_{P_{\mathcal{X}}}^{2}=\|\Gamma^{*}(g_{0}-g^{*})\|_{P_{\mathcal{X}}}^{2}=\beta^{2}\|\Gamma^{*}g_{1}\|_{P_{\mathcal{X}}}^{2}\leq\beta^{2}\tau^{2}$,
which leads to the inequality we need to prove. Finally, by setting
$h_{i}=\hat{\phi}_{i}$, we can see that $\tau^{2}\leq
S_{\lambda}(d+1)-\operatorname{Tr}({\bm{G}}^{-1}{\bm{F}})$. And for all
$d^{\prime}\leq d$, $\operatorname{Tr}({\bm{G}}_{h}^{-1}{\bm{F}}_{h})\leq
S_{\lambda}(d^{\prime})$, so $\tau^{2}\geq\lambda_{d+1}$. ∎

##  Appendix F Proofs for Section 4

##### Proposition 4.

For any ${\hat{\Psi}}=[\hat{\psi}_{1},\cdots,\hat{\psi}_{d}]$ where
$\hat{\psi}_{i}\in{L^{2}(P_{\mathcal{X}})}$, it holds that

$\textnormal{err}({\hat{\Psi}};{\mathcal{F}}_{B}(\Gamma;\epsilon))\geq\frac{\lambda_{d+1}}{1-\lambda_{d+1}}\frac{\epsilon}{1-\epsilon}B^{2}\quad\text{given
that}\quad\frac{\lambda_{d+1}}{1-\lambda_{d+1}}\frac{\epsilon}{1-\epsilon}\leq\frac{1}{2}.$
(36)

To attain equality, it is sufficient for ${\hat{\Psi}}$ to span the top-$d$
eigenspace, and also necessary if $\lambda_{d+1}<\lambda_{d}$.

###### Proof.

Necessity: Since $\hat{\Psi}$ is at most rank-$d$, there must be a function in
${\textnormal{span}}{\left\\{\psi_{1},\cdots,\psi_{d+1}\right\\}}$ that is
orthogonal to $\hat{\Psi}$. Thus, we can find two functions
$f_{1},f_{2}\in{\textnormal{span}}{\left\\{\psi_{1},\cdots,\psi_{d+1}\right\\}}$
such that: $\|f_{1}\|_{P_{\mathcal{X}}}=\|f_{2}\|_{P_{\mathcal{X}}}=1$,
$f_{1}$ is orthogonal to $\hat{\Psi}$, $f_{2}={\bm{u}}^{\top}\hat{\Psi}$
(which means that $f_{2}\perp f_{1}$), and
$\psi_{1}\in{\textnormal{span}}{\left\\{f_{1},f_{2}\right\\}}$. Recall that
$\lambda_{1}=1$, and $\psi_{1}\equiv 1$. Let
$\psi_{1}=\alpha_{1}f_{1}+\alpha_{2}f_{2}$, then
$\alpha_{1}^{2}+\alpha_{2}^{2}=1$. Without loss of generality, suppose
$\alpha_{1},\alpha_{2}\in[0,1]$. Let $f_{0}=\alpha_{2}f_{1}-\alpha_{1}f_{2}$.
Then, $\|f_{0}\|_{P_{\mathcal{X}}}=1$, $f_{0}\perp\psi_{1}$. Note that we also
have $\langle\psi_{1},f_{0}\rangle_{{\mathcal{H}}_{\Gamma}}=0$ by duality. Let
$\beta_{1},\beta_{2}\in[0,1]$ be any value such that
$f=\beta_{1}\psi_{1}+\beta_{2}f_{0}$ satisfies
$\|f\|_{P_{\mathcal{X}}}^{2}=\beta_{1}^{2}+\beta_{2}^{2}=1$, and
$\|f\|_{{\mathcal{H}}_{\Gamma}}^{2}\leq\frac{1}{1-\epsilon}$. This is
satisfied as long as
$\beta_{2}^{2}\leq\frac{\epsilon}{1-\epsilon}\frac{\lambda_{d+1}}{1-\lambda_{d+1}}$,
because
$\|f\|_{{\mathcal{H}}_{\Gamma}}^{2}\leq\beta_{1}^{2}+\frac{\beta_{2}^{2}}{\lambda_{d+1}}=1+\frac{1-\lambda_{d+1}}{\lambda_{d+1}}\beta_{2}^{2}\leq\frac{1}{1-\epsilon}$.
Moreover, we have $Bf\in{\mathcal{F}}_{B}(\Gamma;\epsilon)$.

It is easy to show that
$F(\alpha_{1})=\alpha_{1}\beta_{1}+\alpha_{2}\beta_{2}=\alpha_{1}\beta_{1}+\sqrt{1-\alpha_{1}^{2}}\beta_{2}\
(\alpha_{1}\in[0,1])$ first increases then decreases, so
$F(\alpha_{1})^{2}\geq\min{\left\\{F(0)^{2},F(1)^{2}\right\\}}=\min{\left\\{\beta_{1}^{2},\beta_{2}^{2}\right\\}}$,
which can be
$\frac{\epsilon}{1-\epsilon}\frac{\lambda_{d+1}}{1-\lambda_{d+1}}$ in the
worst case given that it is at most $\frac{1}{2}$, in which case the
prediction error of $Bf$ is
$\|B(\alpha_{1}\beta_{1}+\alpha_{2}\beta_{2})f_{1}\|_{P_{\mathcal{X}}}^{2}=F(\alpha_{1})^{2}B^{2}=\frac{\epsilon}{1-\epsilon}\frac{\lambda_{d+1}}{1-\lambda_{d+1}}B^{2}$.
Thus, for any $\hat{\Psi}$, we can find a function
$Bf\in{\mathcal{F}}_{B}(\Gamma;\epsilon)$ such that
$\min_{w}\textnormal{err}(w^{\top}\hat{\Psi},Bf)\geq\frac{\epsilon}{1-\epsilon}\frac{\lambda_{d+1}}{1-\lambda_{d+1}}B^{2}$.

When $\lambda_{d}>\lambda_{d+1}$, to attain equality, we need $\alpha_{1}=0$,
and
$\|f\|_{{\mathcal{H}}_{\Gamma}}^{2}=\beta_{1}^{2}+\frac{\beta_{2}^{2}}{\lambda_{d+1}}$,