Understanding Augmentation-based Self-Supervised Representation Learning via RKHS Approximation and Regression
Authors: Runtian Zhai
Bingbin Liu
Andrej Risteski
J Zico Kolter
Pradeep Kumar Ravikumar
Venue: ICLR 2024 spotlight (2024)
Keywords: Learning Theory
Representation Learning
Self-supervised Learning
Data Augmentation
RKHS Approximation
RKHS Regression
Primary area: unsupervised, self-supervised, semi-supervised, and supervised representation learning
TL;DR: We establish an RKHS approximation/regression framework for analyzing self-supervised pretraining based on data augmentation, and derive nonparametric learning guarantees that disentangles the effects of the model and the augmentation.

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 the 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, and it is the sum of an estimation error bound incurred by fitting a linear probe, and an approximation error bound by RKHS approximation.
Our second bound specifically addresses the case
where the encoder extracts the top-d eigenspace of a finite-sample-based approximation of the underlying RKHS.
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.

Published as a conference paper at ICLR 2024
UNDERSTANDING AUGMENTATION--ASEE DELF-SUPERVISED REPRESENTATION LEARNING VIA RKHS APPROXIMATION AND REGRESSION
Runtian Zhai, Bingbin Liu, Andrej Risteski, Zico Kolter, Pradeep Ravikumar Carnegie Mellon University (rzhai, bingbinl, aristesk, zkolter, pradeepr (ccs cmu, edu
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 augmentation remains limited. Recent work has built the connection between sulf-sspe learning and the approximation the top eigenspace of a graph Laplacian operator, suggesting that learning linear prohe atop such representation regression. Building on this insight, this work delves into a statistical analysis augmentation-baln pretraining. Starting from the isometry property, 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 approximation error entailed by RKHS approximation. Our second bound specifically addresses the where the encoder is near-optimal, that is it approximates the top-d eigenspace the RKHS induced by the augmentation. A key ingredient our analysis is the augmentation complexity, which we use quantitatively compare different augmentations and analyze their impact on downstream performance.
be connected to RKHS
INTRODUCTION
It is widely acknowledged that better data augmentation techniques have been a major driving force in many recent breakthroughs self-supervised representation learning. For example, contrastive learning (Chen al., 2020) with aggressive random cropping and strong color distortion greatly im- proves the performance of vision tasks, and masked prediction with random masking and replacement is among the state-of-te state-of-the-art representation learning methods for both natural language process- ing (Devlin et al., 2019) and vision (He et al., 2022). Due to their extraordinary empirical successes, developing rigorous theoretical understanding of augmentation-base 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 variant of contrastive representation learning, termed spectral contrastive learning, by connecting it to learning eigenfunc- tions 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 explicitly grapple with the function class being used learn the eigenfunctions (typically deep neural networks in practice). Such a dependency was even deemed necessary by Saunshi et al. (2022) who argued that "funetion-class-agnoiiii analysis leads to vacuous guarantees." Consequently, the generalization bounds such 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 neural networks remains 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.

---

Published as a conference paper at ICLR 2024
L (Px): Large space of L bounded functions, containing the target function f" Induced RKHS Hr: Functions with soft invariance to augmentation, f E Hr Empirical RKHS Hr Obtained with N samples to approximate inaccessible Pr Empirical top-d RKHS Ha Induced by the pretrained d-dimensional encoder Approximation error IIf - fgll Distance from f" to Ha (Lemma 3, Theorem 2)
(Px)
Hr
Ha
14
5
Hr
Estimation error I f|:: Entailed by downstream regression (Lemma 2)
Figure 1: Overall RKHS approximattinn/egressioon framework illustration and commentary.
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 specific RKHS Hr which we call the (augmentation) induced RKHS. Meanwhile, the d-dimensional feature encoder 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 approxi- mation. It learns, with unlabeled data, an encoder v whose learned RKHS Hig (Eqn. (9)) approximates Hr-
Downstream: The supervised stage is doing RKHS regression over Hi with labeled data.
Figure provides illustration of this RKHS approximatiio/regreesiio 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 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 result, Theorem 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 k, 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 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- d-dimensional encoder consists of the top-d eigenfunctions of the induced RKHS (Proposition 4). However, these eigenfunctions are not accessible given only finite samples, so we instead study the near-optimal d-dimensional encoder, which is 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.
A practical implication of our framework is that the augmentation complexity k. defined as the L norm of the kernel of Hr, can be used as 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) λ depends on both the augmentation strength and the augmentation strategy: (ii) a smaller K (e.g. stronger augmentation) leads to a smaller generalization gap, but an overly strong augmentation causes poor training perfor- mance. 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 X C 4xx be the data space, and Px be the data distribution. Let (Pxx) be the L2 function space, such that any f E (P2) satisfies Ery lxx)) < co. Let (J1, f2) = [x(d))((()) and
2

---

Published as a conference paper at ICLR 2024
(J, be the inner product and the norm of the Hilbert space (Px). Let I (Px) 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 Tis" IN and labeled samples T1...- Jn. i.i.d. sampled from Px, and labels Vv = =#xx + us for k [n]] and random noise Vk. find predictor f (Pxx) with a low prediction error err( f*) : IIf 11 EExx
- (
2.1 THE AUGMENTATION OPERATOR D: A SOURCE OF PRIOR KNOWLEDGE
The core problem in self- supe-vueee representation learning is how leverage the unlabeled data 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.
.N.. T1.'..,
Formally, let A denote an augmented space that contains the augmentations of x X. Data augmentation induces a joint distribution PAx, with marginal distributions PA and Px. We will use capital letters A and X to denote random variables from PA and Px- We define an augmentation operator I' = T-=-a= (Pxx - (PA)) as (I' (a)
[[](())aa] for all f E (Pxx). and define its adjoint operator 1 = I o- (PA)) + (Px) as (T (=)=)))(a (x) = ([g(A)) [x] for all that ra-v is indeed the adjoint of r- since for all f E (Pxx and 9 E (PA), there is (==+affg)) = fJ f (x)gg())p(a x r)dadr= = (f. (aa+99)) Px- Intuitively, the augmentation operator connects any single sample x to the distribution of its augmentations p(alæ).
9 (PA)) We
Example: In BERT, suppose we use 15% random masking. Then, X' is the space of original sentences, and A is the space of 15% masked sentences; Px is the distribution over original sentences, A ~ p(-[z) is the 15% randomly masked version of an original sentence x. and PAx x) = (x()))(a(z)) Thus, (To-=g)) is the mcan of g over all 15% randomly masked versions of x.
Now we define the "soft invariance" w.r.. the augmentation. Denote the function class F(T; e)) to be the set of function f such that 3g* E (PPx)) such that I (xx) (T (Ta-rx9)) (x)) Elg" (A)(ar). and
1 Ex E (A (A) - g (A)) S < "lgg" 1p:
for some small positive constant c. 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' is given by P (a, a') in Eqn. (2). Compared to Assumption 1.1 in Johnson et al. (2023), Equ. (1) has an additional I19*1|?, term on the right-hand side so it is homogeneous. For any f E e), if || S B for some constant B, then we denote f E FB();e).
Assumption 1 (See discussions in Appendix A). We assume that f E FB( for given B,
Let the range of [ be R((T*) = [1 = T Pa E (P(())). Then f E R(T*) by Assumption 1.
2.2 INDUCED RKHS Hr AND THE ISOMETRY PROPERTY
Now, we show that R(T*) is an RKHS. Define the following kernel over A x A:
(t)))) (tta,, = plaale)lloos)))PP(() (2) K (a), = d(PA 2 PA) PA(a))
dPt
Here, pt is called the associative distribution, and K A is the Radon-Nikodym derivative P w.r.t. PA @ PA. also called the positive-pair kernel in Johnson et al. (2023). Next, define dual kernel as:
(f()1 (x(x) d(Px ② Px) Px(e2))x(22)
dPt
(1)==1)((((22)) (a(a)
da,
which is constructed by swapping and in Kx and then applying the Bayes rule. It holds that:
(Ta-xIzzaaf (T*))(()) = =
(T TT-gToa-2900 (T)
3

---

Published as a conference paper at ICLR 2024
See App. D for derivation. Let A, E R and 4/2 E (Pxx) be the sorted eigenvalue and eigenfunction of - Prr L.e. T = X,0%, with A 2 X ..... 0. Suppose f Xx(x, (')ddd(()))Px((( < oc. so that r is compact can choose 2/1, 2/2, ... that form an orthonormal basis of T.2 such that (Vs, 15)) = 81g and any f f (Pxx) can be written as f f= -= u.th, for some {uu
integral operator. Then, by Hilbert- theorem,
Note that X E [0, 1], since IIall 2 | for any g (PAA) by Jensen's inequality. It is easy to check that A = 1 and wh = is always pair of eigenvalue and eigenfunction of T* Moreover, when d is sufficiently large, da should be close to O. Similarly, we can define a set of orthonormal eigenfunctions of ITT* denoted by { and we can relate these two sets of cigenfunctions by:
Proposition 1 (Duality). Operators IT and I share the same non-zero eigenvalues. and there exist eigenfinctions { Tf that form an orthonormal basis of I.2 (PA). such that for any x >
Vi = Moreover, we have the following spectral decomposition of the Radon-Nikodyy derivative:
X; 1/1 Ta-aa and 6; -b=xk x,, 1 x, m.-a
dPAx
(ax((,,) Px(a) = Σ 2/22 (e(a)vs(x).
d(PA 80 Px)
Define Hilbert space Hr = Es=E, uath E (Px()x E, u < oo). whose inner product is [2) E. u,u, for S1 E; uebes and f2 E. U,dr.. It satisfies (proof in Appendix D):
(i) Kx is the reproducing kernel of Hr, such that for all f E Hr, f(x) (x,
(ii) Hr = =. We call this the augmentation induced RKHS.
(iii) Hr is isometric to span( (pan((6i)x, (6))), a subspace of (PA). and l|fll = inf.:/-r y 11l||p
(iv) For any f E EB) C. R(T"), let f = Ei ugbr. Define go := Ei>; 1/ Then, 9o
must satisfy Eqn. (1), so we can choose g" = 9o. in which case Eqg. (1) is equivalent to:
, (T") S <llg" IP:A
Σ
1 A A
Σ A
(5)
Moreover, Eqn. (5) is equivalent to a simple and intuitive isometry property:
(1 11 & 11 I 11
(6)
This property essentially says that the operator I'" pre- serves most variance of the target function f* in the naturally, the optimal d-dimensional encoder that minimizes the approximation error should keep as much variance as possible under 1"" This isometry property can be considered of the restricted isometry property (RIP) in infinitediimensioaal functional spaces: In Section 4, we will show that the optimal encoder consists of the top-d eigenfunctions of I''.. Analogously, for umite-dimensione vector space, PCA finds the principal components that keep the most variance under linear transformation T, and they consist of the top-d eigenvectors of T.
Understanding the isometry property:
functional space. Thus
counterpart
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 & during upstream. Instead, the common practice is to first pretrain an encoder & (d), ... &d] on the augmented space with à: < (PA).. and then transform into V 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 & to downstream tasks directly without transformation, but for theoretical analysis we need to explicitly write out the transformation since v and & work on different spaces. We consider the commonly used average encoder (Saunshi et al., 2022, Eqn. (4)):
(x) = E g (A))z) = / (a))p(a)z)da
(7)
4

---

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 & is A1 +... Aa, and can achieved by the top-d eigenspace Hp. Then, define the learned kernel as Published as a conference paper at ICLR 2024
which is equivalent to V r and thus di, R(TT) = Mr for all i [d]. In this section, @ can be an arbitrary function, even with infinite dimensions, so that the model can have any architecture and size; whereas always the average encoder of क Moreover, we consider the following predictor:
Definition 1. The final predictor is the nonparametric least-squares estimate defined as
f==
arg min f:f==" E
UK - f(2x))²
(8)
Here, note that by Eqn. (6) and || S B, there 1| S
B
VI-E*
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 K = Kx 11222 L.e. for Px-almost all x,
(x(x,2)= = Σ \ p(a]x)2 Pa(a) da da (x((((() d PA) +1
Di=1 d and Sx (x0) = Di= 7r Wang al. (2022h) showed that S) D>= PA) + By convexity of D22, we have SA < x², i.e. I a trace-class operator (or x² 2 [ (x, x))Px(()) = S2). Thus, d1 = I. We will provide examples of this augmentation complexity in Section 5.
Here, 2x= Q) 1 1)d is the di-dieergence. Let Sx(d) = Ei
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 क. without loss of generality, suppose full-rank.
Definition 3. Define covariance matrices F, G as F(i, = (1i,, r (T &i, I Tārrr and G(i.)) = (di, 6s) Then, the ratio trace is defined as Tr(G- if G- is well-defined.
112))())) (x, = (G (T"((--/)))())), (G-1))(())))))))))())))))))))))
which is the reproducing kernel of H.g span(Ei, ......)) ). subspace of Hr. Here G- is used for normalization. The ratio trace can be viewed as the trace of Hg. Then, define the trace gap as:
72 == inf inf Sx(d' + +) - -TrG Fr),
(10)
d' d'Sd hr
where 7 2 0, h1 = σ, G₁ = ((h,, hg) Pa) and F, = ((T /4, h22) Px 25jaa- Note that for any d' < there is Tr(G F,)) I (dd',, so 72 always lower bounded by Aat1- And by choosing h; = d. for i [d]. we can see that ? < Sx(d + 1) Tr
We now state our first result. For simplicity, we assume that the random noise v1, Vn are i.i.d. (O variates for some > O. though this can be relaxed (Wainwright, 2019, Chapter 13).
Theorem 1. Let PI, be i..dd. NO0, variates. Let f be defined by Definition 1. If & has d dimensions and (d can be xc), then there are universal constants Cos C1 C2 such that with probability at least exp
22 (d4xx
exp (-V
2nn B3
there is
97² +c CBn(B² + +BB) I 1- < (1
Sx(d +1)
for all J E'E Fn(iee).
1-6
Remark. The first term in the bound controls the approximation error entailed by the limited capacity of the d-dimensional encoder + This term may not vanish as N, - co, since 72 is lower bounded
5

---

Published as a conference paper at ICLR 2024
by Ae+1 which can be positive. For instance, if d is too small for & to represent f*, then the approximation error cannot be zero regardless of the number of samples available. Nevertheless, we can show that with proper pretraining algorithm. the trace gap T² can be very close to Ao+1 as - oo (Theorem 2). larger ratio trace leads to a smaller approximation error, and indeed many existing objectives are strongly connected to maximizing Tr(G-1 (see Appendix C). The second term bounds the downstream estimation error, which vanishes as n - oc. Finally, this result requires full access to p(a (x) for any æ, which available in theory since the augmentation scheme is of our choice. Practical considerations for estimating p(a[z) is outside the scope of this work.
> Proof Sketch of Theorem 1: Lemma 2 (Estimation error bound). Suppose Vi.-.. Dn are i.i.d. NN(0, o²) variates. If de has d dimensions (d can be oo), then we have the following uniform bound over all I D" E (T;
This result follows from two bounds:
CoR(B + oB)
5x(d +
P BLML
vf e FB(T:e))
f' I <911f - 1'
1 1-6
n
C2 2nSx(d+1
exp (-V 2nx2 B²
2 1 c1-ceexp exp
K
where f4 = (I1g9") is the projection of f onto H.g. w.r.f (-,)) Hr. and co. C1. C2 are universal constants. Note that f depends on f". Moreover, Sx + 1) S min {d+ }.
Remark. The constant 9 in the bound is loose and can be tightened arbitrarily close to (at the cost of increasing other terms in the bound) by straightforward modifications to the proof, which we omit here. Such greater-than-one constant is inevitable for any uniform deviation bound.
Lemma (Approximation error, upper bound). IfT < 1, thenfor any f E FB(:;e), there is
72 T T+ B²
lIf 1'
S 1 1-
(11)
4 ANALYSES FOR THE NEAR-OPTIMAL ENCODER
In this section, we consider the special case where & is near-optimal, i.e. the Monte-Carle approx- imation 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 vh,.... v.. Define the worst-case approximation error over FI; as:
Fn(;FF((TTe)) = sup min err( ww (f sup min w| V -ffliix fliixx
fEFa(T;:) WeRd
JEFe(Tie) MERd
Proposition (Approximation error, lower bound). For any & = Wh,,... .Vaa where $s E (Px)).).
det1 E FF((e)) 2 1- 1-- EB²
A641 6 < 2
B2 given that 1- 1
To attain equality, it is sufficient for V to span the top-d eigenspace, and also necessary if dat1 < da- 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 & spans the top-d eigenspace (see Appendix C for examples). Also, when Xa > O, the top-d eigenspace lies within R(["), ), which is the function class of V when using the average encoder.
The optimal d-dimensional encoder is inaccessible since v'......th cannot be extracted with only finite samples. What one can do instead is extracting the empirical top-d eigenspace Ha- Given samples ø1, ; N.,, define two Hilbert spaces L2((xx and L LPA)),, where (f1,2) = + M fi (z))) and (g1, 92) = gi (a))))))) for PA(a) = 1/ EN (x)) With these finite samples, I* will remain the same, but I' will become the empirical version I:
f(zx)p(((xx) Pa(a)
(T)) N Σ
Let the eigenvalues and eigenfunctions of i be ((),, k.)). Let ø; be the eigenfunctions of iT*
6s and yú have the duality property similar Proposition Let Hr be the RKHS associated with

---

Published as a conference paper at ICLR 2024
r*I, which is the span of { 3, and hence is subspace of Mr R(T*). Let Ha be the top-d eigenspace of Hy. Let Amax G), Amin (G) be the largest and smallest eigenvalue of G. For the case of interest where Hi Ha. have second main learning guarantee that bounds the trace gap:
Theorem 2. Suppose d. de for (d]. Let == Amax ( (G)/minn (G) be the condition number of G. Then, for any > 0, it holds with prohability at least - 5 that
7² S Tr(G Tr(G-1F) A041 + 2 V 22
(Ad
+
nd G
T-1 1/2
+2)x²
d.
VN
Remark. Combining this result with the first main Theorem leads the bound for Hig Hds Since 7² 2 Ad41. this result says that the gap between 7² and its optimal value is D(2V- ignoring potential dependence of G and X on N. Comparing the upper bound in Lemma 3 + Theorem to the lower bound in Proposition 4, note that the upper bound is near tight: The only difference is that in Eqn. (11), we have I+f instead I-c" Unlike the bounds in HaoChen et al. (2021); Saunshi I-e
al. (2022), this bound only depends X "x=' and "G. but not (Aar (4+1)=1 for some d' s 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 pointed out by Saunshi al. (2022).
> Proof Sketch of Theorem 2: Denote Sx(d) Σi= 74 We define two empirical covariance matrices F and G as: F(i, (v., uby) Px- and G(i,j) (ô., .5) P.- Since any invertible linear transformation to & does not change -r FF) the empirical ratio trace Tr(ô-1F) F)), we see that for H.s. = H.. there Tr F) Sx (simply consider å कi). Thus, suffices to bound| ( - Tr(G-1 F) and (x) - (d).
We start with bounding the gap between empirical and real ratio traces for any encoder d:
Lemma 5. Suppose there exists a constant C > 0 such that Ep, [g"]] s 1 p.. for all w where | < 1. Then, for any 8 > 0, it holds with probability at least 1 - 8 that
CK + x² VN
-
2
2 log
(12)
Lemma' 5 considers a general ④, and requires fourth momentt control assumption (Wainwright, 2019, Eqn. (14.22a)). For the specific case Φ; ē; studied in this section, we can prove that this assumption holds (that 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 òe for [d]. mamber of G. Then. for any 8 > 0, both
Let Na == Amax Amax(G) (iin (G), which is the condition
d
(2-
(A 1 + 1)x²
Σ
Σ A 1=1
2 2
VN
j=1
and Eqn. (12) with C = xZ
1/2 Naa
hold simultaneously for Hig Ha with probability at least 1 - 8.
5 ESTIMATING AND EXPLOITING THE AUGMENTATION COMPLEXITY
An important quantity from our analysis is the augmentation complexity x- both approximation and estimation error bounds get smaller when x is reduced. A natural way to reduce k 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 к of mask-type augmentations. We study k of different masking schemes on the hypereube data model introduced in Saunshi et al. (2022), and show that depends on both the mask ratio (augmentation strength) and the masking strategy. Then, we empirically estimate the real-world NL.P dataset wi k ped i a -i im and demonstrate
7

---

Published as a conference paper at ICLR 2024
its correlation with the generalization gap on real downstream tasks. Our results indicate that the augmentation complexity a offers a means to compare different augmentations quantitatively and implies tradeoff in the choice of augmentations.
5.1 THEORETICAL HYPERCUBE DATA MODEL
Original Sample
1
-1
1
-1
1
Like Cabannes et al. (2023), we study three random
1. Random Mask
0
0
-1
0
1
0
masking methods (Figure 2): (i) Independent ran- dom masking, (ii) Cutout-like block masking, and (iii) BERT-like masking. Let a denote the mask rate. Consider the data space X' = (-1,1)a 1)x with Px being the uniform distribution over X. Denote the corresponding K of these three methods to be Kep Ke, and Kg. We will show that these k's will be very different even with the same a.
2. Block Mask
0
0
3. Black Mask + Flip
0
0
Masked
Flipped
Figure 2: Augmentation illustration.
Let's first consider the independently random masking (derivations are deferred to Appendix G): Example 1. Consider random masking augmentation, i.e. for any x E X. each coordinate x( randomly and independently masked be (i.e. 0 denotes the [MASK] token) with probability E (0, 1). Then, its augmentation complexity is given by K. = (2 a))xx.
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 a smaller K with the same
Example 2. Consider random block masking augmentation, iee. masking 2.((), (6+1), ....... (1+r-1) forr = [orx] and a uniformly random E [dx r), for any E X. Then, 22 S [2(1-a)jdx.
Our third and final example resembles the 80-10-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 I E X, first mask xz(1) (it+r-1) to be 0 for r = [adx] 1 and uniformly random i E [dx - rl. then randomly flip a independently. Then, 1e? < (a 2a + ()-12
the sign of each remaining coordinate w.p. 2
While Ki, Ri, Ki all have an exponential dependency on dx, their bases are different, as shown in Figure 3a. This means these methods have different is despite having the same mask rate, which highlights the importance of the masking strategy.
On the exponentiality x in dx. We suspect the exponential dependency on dx 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 A is 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- suf--uueeeieed learning suggest that they implicitly adapt to the inherent -imeesniioonal 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 A can be estimated on real datasets and used to quantitatively analyze and compare augmentations, which we demonstrate the NLP dataset wi wikipedia-simpl study masked language modeling, where x is a full sentence and a is a masked sentence. Recall that p(z) p(a)ar)da Empirically, we replace the integration with sample average; namely, for = [z(1), ,x()] where z(1) is the jth token, we have
We
n² upper bounds f plala)
logp (x/aa = logp (c(:/a) + logg (2122) (2t20/a,
a, a(")) + ...
....4 + logg (I(0 a, (1)) " (
8

---

Published as a conference paper at ICLR 2024
300
Ranker
Tost Accancy
90
(0:16,
86.8)
95
(0.40 92.4)
12
Test Gap
Ranker Hip
1.5
log
201
Blac
80
Blck Fly
80
100
QNT
0.5
0.2 0.4 0.8
TO
0.4 0.6
0.8
0.2 0.4
Mask Katio
Mask Ratio
Rano
Musi Katio o
(a) Bases of three x²
(b) Estimated log x?.
(c) Results on QNLI and SST-2.
Figure 3: Plots for Section 5. In (b), log R² is estimated on wi wikipedia-simpl
We leverage bi-directional MLM such as BERT, and compute p(z:00 xl(0) auto-regressively:
For any E [1] compute p(x(*/a, x(<)) with the BERT, and then replace (t)) with xx(9) for i- 1. We compute this for a random subset of samples. We report the 99th percentile of Kx (α. x) instead of sup, Kx (a, æ) due to the presence of outliers. Experimental details are deferred Appendix H.
Estimating k. We study Random masking, Block masking, Random masking + Flipping, and Block masking Flipping. For Masking + Flipping, we randomly mask a/2 of the tokens, and replace another a/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 3b plots the estimated log x² of the four augmentations w.r.l. the mask rate a, where log x² is calculated as the average of five runs with different random seeds. The complexity drops as a increases as predicted by theory. One observation is that the "Random+ + Flip" curve intersects with "Block" and "Block + Flip", suggesting that block masking has stronger effect when a is small, whereas flipping has a stronger effect when a large.
Downstream performance. We also study how the mask ratio a affects the downstream perfor- mance, using QNLI (Wang et al., 2018) and SST-2 (Socher et al., 2013). For pretraining, we train roberta- aatartallaeee rgr 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 s'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 3c. The highest test accuracy is achieved at a = 0.15 on QNLI and at a 0.40 SST-2 (marked in red). The test accuracy is low when is too small due to the large generalization gap, and also low when a 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 a = 0.40. This is likely because with >> 0.40 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 This suggests the presence of "sweet spot", which is also supported by evidence in prior work (Tian et al., 2020).
6 CONCLUSION
This work establishes an RKHS approximattio/regresiion 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, that they are valid for an arbitrary encoder, and (iii) formulates the augmentation complexity i as a tool to quantitatively analyze the effect of augmentations. We also refer our readers two recent works that used similar techniques but derived quite different results: Wang et al. (2022b) studied an RKHS similar to Hr. but in a completely different context (they studied conditional moment models); Cabannes et al. (2023) studied a similar problem, but they used pre-defined kernel to define the function class and derived guarantees, while the RKHS Hr 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 x can be exponential in dx, which as we pointed out could be 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 the upstream and downstream have different input spaces (AA and X'), and this transformation cannot be studied within the architecture-Fi of this work. Additionally, an open problem is how to deal with the common upstream-to-dae towwonssrteammmm data distribution shift in practice, for which downstream fine- men-tnnnn could be critical. This has been discussed in Cabannes al. (2023), but more general analysis is desired.
framework

---

Published as a conference paper at ICLR 2024
CODE
The code of Section be found https: reserrch www.::c...oooll...ccccccooooooo 11oS2LLI- hro.aaa.-ppaataa.....oo.rrreaaar
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- 1909866, CCF-2238523, ONR via N00014-23-1 1-23-,1--6,, and Amazon.
REFERENCES
Zeyuan Allen-Zhu, Yuanzhi Li, and Yingyu Liang. Learning and generalization in overparameterized neural networks, going beyond two layers. Advances neural information processing systems, 32, 2019a.
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.
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.
Sanjeev Arora, Nadav Cohen, Wei Hu, and Yuping Luo. Implicit regularization in deep matrix factorization. Advances in Neural Information Processing Systems, 32, 2019a.
Sanjeev Arora, Simon Du, Wei Hu, Zhiyuan Li, and Ruosong Wang. Fine-gran Fine-grained analysis of optimization and generalization for overparameterized two-layer neural networks. International Conference Machine Learning, pp. 322-332. PMLR, 2019b.
Sanjeev Arora, Simon S Du, Wei Hu, Zhiyuan Li, Russ Salakhutdinov, and Ruosong Wang. On exact computation with an infinitely wide neural net. Advances in neural information processing systems, 32, 2019c.
Adrien Bardes, Jean Ponce, and Yann LeCun. VICReg: Variace-iivarianc--cvvaian regularization for self-supervised learning. In International Conference Learning Representationss, 2022. URL https: http///pppeeeiiieeeeem/ee net / forum? dididdmx1x0 www.D62D10ss...
Peter L Bartlett, Dylan J Foster, and Matus J Telgarsky.
margin bounds for
neural networks. Advances neural information processing systems, 30, 2017.
Peter L. Bartlett, Philip M. Long, and Olivier Bousquet. The dynamics of sharpness-awan awaree minimiza- tion: Bouncing across ravines and drifting towards wide minima. Journal of Machine Learning Research, 24(316):1-36, 2023. URL http: / http://llrrrrrr...
20444433333.. httml.
Mikhail Belkin and Partha Niyogi. Laplacian eigenmaps for dimensionality reduction and data representation. Neural computation, 15(6):1373-136, 2003.
Mikhail Belkin and Partha Niyogi. Semi-supervised learning riemannian manifolds. Machine learning, 56:209-239, 2004.
Yoshua Bengio, Olivier Delalleau, Nicolas Le Roux, Jean-François Paiement, Pascal Vincent, and Marie Ouimet. Learning eigenfunctions links spectral embedding and kernel pea. Neural computa- tion, 16(10):2197-2215 2004.
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-1888 1828, 2013.
10

---

Published as a conference paper at ICLR 2024
Olivier Bousquet and André Elisseeff. Stability and generalization. The Journal of Machine Learning Research, 2:499-522 526, 2002.
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-3296. PMLR, 23-29 Jul 2023.
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.
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.
Lenaic Chizat, Edouard Oyallon, and Francis Bach. On lazy training differentiable programming. Advances in neural information processing systems, 32, 2019.
Fan RK Chung. Spectral graph theory, volume 92. American Mathematical Soc., 1997.
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 ldr via minimaxing rate reduction. Entropy, 24(4):456, 2022.
Tri Dao, Albert Gu, Alexander Ratner, Virginia Smith, Chris De Sa, and Christopher Ré. A kernel theory of modern data augmentation. International conference on machine learning, pp. 1528- 1537. PMLR, 2019.
Jacob Devlin, Ming- Beig-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 (Long and Short Papers), pp. 4171-4186, Minneapolis, Minnesota, June 2019, Association for Computational Linguistics. doi: 10. 11853881//11911990 URL https: //acl aathooooyyyyy.....oooo org/oog///... 911-142.
Terrance Drieies and Graham W Taylor. Improved regularization of convolutional neural networks with cutout. arXiv preprint www.v::708.0000 04555, 2017.
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 scale. In International Conference on Learning Representations, 2021. URL https: //tp////eeeeeeeeieeeeeeee net/ Eorum? -- www.cbbdNTTTT
Simon Du and Jason Lee. On the power of over- peeemtrramttinntioonn in neural networks with quadratic activation. International conference on machine learning, pp. 1329-1338. PMLR, 2018
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.
Simon S. Du, Xiyu Zhai, Barnabas Poczos, and Aarti Singh. Gradient descent provably optimizes over- peremerrtttee neural networks, In International Conference on Learning Representations, 2019b. URL https: / /ttp:///pp/ee/eeeeeeeeee
net met/ffoummi forum? witiiiiisssee1999...
Gintare Karolina Dziugaite and Daniel 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.
11

---

Published as a conference paper at ICLR 2024
László Erdós, Benjamin Schlein, and Horng-Tzer Tzer Yau. Local semicircle law and complete delocal- ization for wigner random matrices. Communications in Mathematical Physics, 287 2009.
Ky Fan. On a theorem of weyl concerning eigenvalues linear transformations i. Proceedings of the National Academy of Sciences, 35(1 I 1:662-65555 1949.
Pierre Foret, Ariel Kleiner, Hossein Mobahi, and Behnam Neyshabur. Sharpness- aware minimization for efficiently improving generalization. In International Conference on Learning Representations, 2021. URL www.sttp///t.......
net/
1 rM.
Micah Goldblum, Marc Finzi, Keefer Rowan, and Andrew Gordon Wilson. The No Free Lunch Theorem, Kolmogorov Complexity, and the Role Inductive Biases in Machine Learning. arXiv e-prints, arXiv:2304.05366, April 2023. doi: 10.48550//r2iw.2344..555
László Györfi, Michael Kohler, Adam Krzyzak, Harro Walk, et al. A distribution-free theory of nonparametric regression, volume 1. Springer, 2002.
Gerald J Hahn and William Q Meeker. Statistical intervals: a guide for practitioners, volume 92.
John Wiley & Sons, 2011.
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: // //ppppneeveew.cco neteerrevieweee//eeeeeeeee forum?id AuuniiiiddAgEEeemeeeee
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.
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: wpwpreeeeeeeeeeee//e/peee. net worrrmmiiieo derrmmmidiiijipp www.dccwwww.aa...cc.
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.
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.
Sepp Hochreiter and Jürgen Schmidhuber. Flat minima. Neural computation, 9(1::1442, 1997.
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:tt///co wpwpp/peeeeew..eeee/eeee net www.:oiddddddmuuttt..oo.e
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.
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.
Yiding Jiang, Behnam Neyshabur, Hossein Mobahi, Dilip Krishnan, and Samy Bengio. Fantas- tic generalization measures and where to find them. In International Conference on Learning Representations, 2020. URL https: opwnrevieee wtt.vvwwwww.wee.e.eeeee ltte://ooummiid-ssg1_1 JBFV..
2id=
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: opwnrevieeeee...e//e net fo fo f
12

---

Published as a conference paper at ICLR 2024
Daniel D. Johnson, Ayoub El Hanchi. and Chris J. Maddison. Contrastive learning find an optimal basis for approximately cew-invariant functions. In The Eleventh International Confer- ence Learning Representations, 2023. URL https:/ http::/:///peneeeee...oo nwt. www./eeen//iii......ooo AjCOKB iMu.
Kenji Kawaguchi. Deep learning without poor local minima. Advances in neural information processing systems, 29, 2016.
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 Compulational learning theory, pp. 152-162, 1997.
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 Learning Representations, 2017. URL
net/ net/forumss forum? d= Ygg.
Andrei N Kolmogorov. On tables of random numbers. Sankhyā: The Indian Journal of Statistics, Series A, pp. 369-376, 1963.
John Langford and Matthias Seeger. Bounds for averaging classifiers. School of Computer Science, Carnegie Mellon University, 2001.
Jachoon Lee, Lechao Xiao, Samuel Schoenholz, Yasaman Bahri, Roman Novak, Jascha Sohl- Dickstein, and Jeffrey Pennington. Wide neural networks of any depth evolve linear models under gradient descent. Advances in neural information processing systems, 32, 2019.
Jason D Lee, Qi Lei, Nikunj Saunshi, and Jiacheng Zhuo Predicting what you already know helps: Provable superviee learning. Advances Neural Information Processing Systems, 34: 309-323, 2021.
Yuanzhi Li and Yang Yuan. Convergence analysis of two-layer neural networks with relu activation. Advances in neural information processing systems, 30, 2017.
Yuanzhi Li, Tengyu Ma, and Hongyang Zhang. Algorithmic regularization in wer-paramettrrized matrix sensing and neural networks with quadratic activations. In Conference On Learning Theory, pp. 2-47. PMLR, 2018.
Bingbin Liu, Daniel Hsu, Pradeep Ravikumar, and Andrej Risteski. Masked prediction tasks: parameter identifiability view. Advances in Neural Information Processing Systems, 2022.
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: / wwwpreview.wwep/////ppp. wwt.:/ww.eeeeeeeee.eeeeeeeeeeeee
d-vUz JJP JRRLGX.
Sanae Lotfi, Mare Anton Finzi, Sanyam Kapoor, Andres Potapezynski, Micah Goldblum, and An- drew 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: openrevieeeeeeeeeeeeeeeeeee net/ forum? Ydd=8nYuRSekFm
David A McAllester. Pac-bayesian model averaging. In Proceedings of the twelfth annual conference on Computational learning theory, pp. 164-170, 1999.
Mehryar Mohri, Afshin Rostamizadeh, and Ameet Talwalkar. Foundations of machine learning. MIT press, 2018
Youssef Mrouch. Stephen Voinea, and Tomaso A Poggio. Learning with group invariant features: kemel perspective. Advances in neural information processing systems, 28, 2015.
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.
13

---

Published as a conference paper at ICLR 2024
Vaishnavh Nagarajan and Zico Kolter. Deterministic PAC-bayesian generalization bounds for deep networks via generalizing resslseers In International Conference on Learning Representa- tions, 2019. URL https: wwtp//eeeeeeeeeeeeeeeeeeeeeeeeee httprriwwwweeeeeeeeeee www//eeruueeeeee 2d=Hygno00gKX.
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.
Behnam Neyshabur, Srinadh Bhojanapalli, David McAllester, and Nati Srebro. Exploring generaliza- tion in deep learning. Advances in neural information processing systems, 30, 2017.
Behnam Neyshabur, Srinadh Bhojanapalli, and Nathan Srebro. A PAC-bayesian approach to spectrally-normalizeed margin bounds for neural networks. In International Conference Learning Representations, 2018. URL htw.s://////. / opwnreeeeveeeeeeeeee nww.//t/ee...m forum? ? lisssk_wfbCC
Ashwini Pokle, Jinjin Tian, Yuchen Li, and Andrej Risteski. Contrasting the landscape of contrastive and non-contrastive learning. In Gustau Camps- 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: / /tte:://ppppeeee pttppeeiecengn..ce/e mlr. pwwssppsseeeeeppe....ee / vell/vlllll/// http///poeeaaa...eo
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.
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.
Omar Rivasplata, Vikram M Tankasali, and Csaba Szepesvári. Pac-bayes with backprop. arXiv preprint Ar0iv: 073800 2019.
Nikunj Saunshi, Orestis Plevrakis, Sanjeev Arora, Mikhail Khodak, and Hrishikesh Khandeparkar. A theoretical analysis of contrastive unsupervised representation learning. In International Confer- ence on Machine Learning, pp. 56285622666 5637. PMLR, 2019.
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.
Bernhard Schölkopf and Alexander J Smola. Learning with kernels: support vector machines, regularization, optimization, and beyond. MIT press, 2002.
Shai Shalev-Shwartz, Ohad Shamir, Nathan Srebro, and Karthik Sridharan. Learnability, stability and uniform convergence. The Journal of Machine Learning Research, 1 11:2635-2670, 2010.
Kendrick Shen, Robbie M Jones, Ananya Kumar, Sang Michael Xie, Jeff 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.
Richard Socher, Alex Perelygin, Jean Wu, Jason Chuang, Christopher D. Manning, Andrew Ng, and Christopher Potts. Recursive deep models for semantic compositionality over 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: / wwww./wwom.oooo..ooo.... acl web. org/
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.
Yuandong Tian. Deep contrastive learning is provably (almost) principal component analysis. Advances in Neural Information Processing Systems, 2022.
14

---

Published as a conference paper at ICLR 2024
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 volume 139 of Proceedings of Machine Learning Research, pp. 10268-10278. PMLR, 2021. URL http::/proceet ngs mlr. wwl..press/// / 2t/119//iaaie
Machine Learning, ICML 2021, 18-24 July 2021, Virtual Event,
html.
Christopher Tosh, Akshay Krishnamurthy, and Daniel Hsu. Contrastive estimation reveals topic posterior information to linear models. J. Mach. Learn. Res., 22:281-1, 20221,,,,021a.
Christopher Tosh, Akshay Krishnamurthy, and Daniel Hsu. Contrastive learning, multi- redun- dancy, and linear models. In Algorithmic Learning Theory, pp. 1179-1206. PMLR, 2021b.
Martin J Wainwright. High- digesdimmesional statistics: A non-asymptotic viewpoint, volume 48, Cam- bridge university press, 2019.
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. BLACK- BOXNLP@ P.NXLP EMNLP, 2018. doi: 10. 1111663/11/11854440
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.
Xiang Wang, Chenwei Wu, Jason D Lee, Tengyu Ma, and Rong Ge. Beyond lazy training for over-parameteriize tensor decomposition. Advances in Neural Information Processing Systems, 33:21934-21944, 2020.
Yifei Wang, Qi Zhang, Yisen Wang, Jiansheng Yang, and Zhouchen Lin. Chaos is a ladder:
new theoretical understanding of contrastive learning via augmentation overlap. In International Conference on Learning Representations, 2022a. URL https:: / /ttp///ppp/////e///////ee/e net forum?
id-ECvgmYVyeUz.
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: https:://opereeeeeeee //opnneeeiiewe net/ forum?
id=VBTJqqW jxMv.
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.
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 2021.
Processing Systems,
Kaiyue Wen, Tengyu Ma. and Zhiyuan Li. How sharpness minimization minimizes sharpness? In The Eleventh International Conference on Learning Representations, 2023. URL https:
/oppneeview.... n forum? d=5
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, ICMI. 2021, 18-24 July 2021, Virtual Event, volume 139 of Proceedings of Machine Learning Research, pp. 11112- 111222 PMLR, 2021. URL http: / /pproceedings mll .wr.sspes////.....
ptteess/vi9/// vv19///1//////// wen2lc. h1/13///wn21..mm html.
Zixin Wen and Yuanzhi Li. The mechanism of prediction head in non-contrastive self-supervised learning. Advances in Neural Information Processing Systems, 2022.
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.
David H Wolpert and William G Macready. No free lunch theorems for optimization. 1EEE transactions on evolutionary computation, 1(1):67-82, 1997.
15

---

Published as a conference paper at ICLR 2024
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.
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.
Jure Zbontar, 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.
Runtian Zhai, Rattana Pukdee, Roger Jin, Maria-Florina Balcan, and Pradeep Ravikumar. Spectrally transformed kemne regression. In International Conference Learning Representations, 2024.
URL httph:tt///// hwwp///weeeeeeeeeeeeeeeeeeeeeeep net forum? id-00ee92ozttatt
Chiyuan Zhang, Samy Bengio, Moritz Hardt, Benjamin Recht, and Oriol Vinyals. Understand- ing deep learning requires rethinking generalization. In International Conference on Learning Representations, 2017. URL https: openreview wet.:riiwweeeeeeeee.eeeeee.eee forum Sy=
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.
Wenda Zhou, Victor Veitch, Morgane Austern, Ryan P. Adams, and Peter Orbanz. Non-vacuous gen- cralization bounds at the imagenet scale: a PAC-bayesian compression approach. In International Conference on Learning Representations, 2019, URL https: /wwp///eeee/...e///////e net forum? id=BJggqsAct7.
16

---

Published as a conference paper at ICLR 2024
A DISCUSSION
Why do we need Assumption 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 isometry property. As pointed out by Goldblum al. (2023), real-world datasets are highly structured, and the structures are largely shared across different domains. As- sumption and implied isometry property suggest that a well-designed augmentation is a key to capturing such a fow-dimensions 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 (e.g. the target function f" can outside but very close to Hr), additional assumptions on f* would be required for our proof with uniform deviation bounds to hold. Namely, f can no longer be any L² function that is close to Hr. For instance, we at least need to assume that If (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* E Hr 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 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 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 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 I1 and 22 cannot be augmented to the same a. In this case, our Kx 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 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 is satisfied and the augmentation complexity is small enough, which is in fact strong assumption), then with the augmentation and without any help of the model. In other words, the disjoint augmentation example showcases poor choice of augmentation, which is orthogonal to the emphasis of this paper.
achieve good generalization only
17

---

Published as a conference paper at ICLR 2024
B RELATED WORK
B.I 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 modern neural networks have sufficient capacity to fit random labels. This challenges the classical generalization theory, which appeals to the Occam's '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 few seminal works below.
real tasks,
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 trained by stochastic gradient descent (SGD). Let (a;; A, 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 F A, ff fæ; A, S') is uniformly bounded when S and S' 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). has also been shown in Li 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 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 (x || P) is small, then good generalization is guaranteed. Thus, the more models there are in the support of Q. the better the generalization. So can see that a lower sharpness around the global minima leads to better generalization. Sharpness-bud based bounds are shown to be empirically tighter than other bounds in Jiang al. (2020), and inspire class of training method called sharpness-awan 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 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 al. (2019b;a); Arora et al. (2019b) for two-layer fully-connected ReL.U
18

---

Published as a conference paper at ICLR 2024
networks, and Allen-Zhu et al. (2019b) for other architectures such 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- 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. Le. 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 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 optimal for linear predictor of certain downstream tasks (Saunshi et al., 2019; Tosh et al., 202 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 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 (MCR²), which learns an explicitly meaningful low-dimensional representation by maximizing the rate reduction, defined the difference in coding rate, i.e. the number of bits needed encode the embedding before and after a certain partition (e.g. according to different classes). Yu et al. (2020) showed that MCR² 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 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.l. "1,'"" aN., which is incurred by sampling only one from each x (see their Definition D.3). Then, Saunshi et al. (2022) pointed out that this model-class-fie 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." 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 Hr defined in this work. In partieular, the kernels kg and kg in their work similar to the kernels
KA and Kx 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 strong innate connection, which alludes to the role of augmentation in self-supervised pretraining, namely, it defines ordered set of features as we articulated after Eqn. (6).
'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.
19

---

Published as a conference paper at ICLR 2024
Another related work is Cabannes al. (2023), who studied the same problem this work. The difference is that their work uses a pre-defined kernel, while the kernels KA and Kx 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 distribution
would also like to credit Cabannes et al. (2023) for studying the upstream-to--oowstrees 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.
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 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:
EscL( = 2E (A), (AA))) + IE (k(A), (A-)))
(13)
where (A, A-) is a positive pair that is jointly sampled from Px (a, a+), and A is negative sample that is sampled from PA independently of A. 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 Xa+1
Ad-
Here, "minimized by the top-d eigenspace" means that the objective is minimized when & spans the top-d eigenspace, and "umiquely minimized" means that this condition is necessary. The proof is based on the following Eckart- Young-Mirsky Theorem:
Theorem 4 (Eckart-Youuu -Mrrgyy),.,...... Given a matrix A E Rmxx, consider the following optimization problem:
min BeRmxn rank BYck
| - BI
(14)
where 1 1 k S. min(m, n). Let the SVD of A and B be = UEVT and B = -U-,VV where Σ diag(o1,"" .0 min(m ), and o 2 02 2 2 m.. Then is a minimizer of the above problem if is the first k columns of U, V1 is the first k columns of V, and Σn is the upper left ls x le block of Σ. Moreover, if fo ok, then this the unique minimizer.
To make this paper celf-contained, here we recap the proof of Theorem in HaoChen al. (2021):
Proof. Define matrix C as: $ C$* and matrix B as B C. C. Then, ( $ B Let the SVD of C be C = U.V Denote D2 = diag( (11, ..). Then, the spectral contrastive loss is equivalent
(sc(()) = -2 Σ Adbsi +
Σ b |B 3J
D D
(15)
20

---

Published as a conference paper at ICLR 2024
So it suffices to minimize | - Dx where rank(B) d. By Theorem 4, this is minimized if (ii,j) ō,, and Σ diag( (v)), which means that 4b Udiag(vAi V for some orthogonal matrix U E _dxdd Moreover, by Thm. this is the unique minimizer if Ad+1 Ad- O
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:
( ( - ( -
BM + (A)) (( ))
(16)
where & and E are encoders for different modalities (e.g. image and text). (A, X+) is a positive pair, while X is 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 Ad41 < Aa-
Proof. Let ② C5* and E SV. Let El = r -SD Φ. Then, we have
// (b(0). ((x) ( p(=la)p(aadadz
(((
/ ((a)) / =(()p(())))d p(a)da
/ (6(1), p(a)daa pla)da Tr ( (CTS
We also have
((( // Σ ((0)) pla)p(z)dadz
( @()))\(c)p((ada) ( ((L)))((o))o))))
Σ
Tr (CC Ss ||c $
Therefore, the spectral CLIP loss is equivalent to
(scui( = C $1 T (CT SD = CTS - D 1
D 1/2
(17)
F
F
Let the SVD of C be UEV. Let M ΣUT S. Then, CT VM. is well known in
linear regression that Eqn. (17) is minimized when M = VV -1 V DTD=² 1/2 = =V D 1/2
so
min CSCLIP E) = min [I [I- [ VV | $
By Theorem 4, this is minimized when the top d rows of V orthogonal matrix, and this condition also necessary A4+1 < ^d.. This proves the result. Moreover, we can see that E is the minimizer if it extracts the top-d eigenspace of up*
Regularized Barlow Twins. Barlow twins (Zbontar et al., 2021) non-contrastive learning method that trains without negative sampling. We consider the following regularized Barlow Twins loss:
CRBT( B)
(()))
+a Σ E ((A))
+
[|
2
(18)
When the regularization term 3 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:
2
min
E
$
s.t.
LRET (5: a. 0) = 0.
(19)
21

---

Published as a conference paper at ICLR 2024
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 >a41 < Xa-
Proof. Define matrix C' (Ci.) as: + CØ. Then, for all k, [d]. there is
(( = Σ A,
Thus, CRAT(O; 0) 0 is equivalent CDAC I. Meanwhile, E [|((4)) ||C|| So
2 2 the optimization problem is equivalent to minimizing || subject to CDACT =I.
I, where 9& [gk,1, Ik.a)- For all j E [d]. define M, = - gegk- Then, MMn) S j. and all eigenvalues of Mj belong to [o, 1].
Let G CvDx () then GG Σx
So the sum of eigenvalues of Mj is most j. which implies that
Σ9 TT(Oo 9)) = Σ Tr (g,9,) = Tr((M S.
1=1
k=
Denote Xo
O. Thus, by Abel transformation, have
|C
Σ 1 Σ d d 9, 9: 1
-d-
9h 9
E (d-j)
-
Xy
A3+1
X
A
sufficient condition of achieving the above equality is Ek= 9 gb for all j E [d]. and is also necessary if Ad+1 < Aa. This condition is equivalent to 91. Id forming an orthonormal basis of Rdd i.e. & 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 Definition 3. We start with the regularized
Barlow Twins loss, which we have shown is equivalent to minimizing ||C [ subject CD,C' = I. Now, observe that G = CCT, and F = CD,CT. Thus, regularized Barlow Twins is essentially minimizing Tr GG), which is similar but not equivalent to maximizing TF)
Now consider variant of VICReg (Bardes et al., 2022), also given by Cabannes et al. (2023,
Eqn. (2)):
LViCRey (3:) = E $ (AA) - r|
[|x(()) +(A+) 2
+ 3E
(20)
When the regularization term B is close to O, we have IE (AA{((A) =
= = And we minimize
E - ((A*+
which is equivalent to 22 Tr (AA) - -2 + c for
some constant c. Thus, this objective is equivalent to maximizing Tr(G F), the ratio trace, when 0. And when B Eqn. (20) is equivalent to the spectral contrastive loss.
D PROOFS FOR SECTION
I'T and * are integral operators.
( ((a=zIa=aa)(s)) (TT))(>)
Kx(x, (f(2))))
(21)
(T ((==a(a--=9)(1) = (a) = KA a jdd'
22

---

Published as a conference paper at ICLR 2024
Proof. We only show the first equation, and the second can be proved in the same way.
(TT]))))-) - (forrooddlaaat) = -J ( (ldd) p(a(x)dd
/ʃ Ma'hoaaaaspa la)dadr' / (x(x,x))(()) mod'
= = 1I Mat) ( p(o)))))))(())))) (dx')d p(a)
TT shares the same non-zero eigenvalues as TT and there exist eigenfunctions {0r} of TT that form an orthonormal basis of (PA), such that for any X; >> 0.
Proposition (Duality).
1 x, de and o; =, 1-
(22)
and we also have the following spectral decomposition of the Radon-Nikodyee derivative:
pla,x) PA 8 Px) p(a)p(a)
dPAx
Σ dc()ss)))))))
(23)
Proof. Suppose X; (1) is pair of eigenvalue and eigenfunction of [TT and A; 0. Then, have T = A, Vi which means that Thls is an eigenfunction of DT with eigenvalue A, The 1/2 x/1/2 is used for normalization. To see this, let o, =,," 1/2TT Then, have
(0:, ds) = >,,1/2x/1/2 (T y/s P
11/1/22,1//2 [ Twa wi Px
= 1// 13) da.-
We prove the reverse direction similarly. And for any fixed x. there is
x(a,2)
p(a,x) p(a)p(z) pla)p())" drla)p(a)da / plal))) (evi(a).
(())p())) : ) PA
(24)
which implies Eqn. (23).
Basic properties of Hr-
(i) Kx is the reproducing kernel Hr, such that for all f E Hr, f(x) = (f, Kx (2. MM
(ii) Hr = -(T*).
(iii) Hr is isometric to
subspace of (PA)) and [lf = infg:/=r'' 11ll|
(iv) For any fr E FB(TT;) C R(T"), let f E, U,ti; Define go Σ, 1/2 Then, go must satisfy Eqn. (1), so we can choose g" = 9n, in which case Eqn. (1) is equivalent to:
(g", - ((--) Pr S <lll" + Σ
A E'~Ai A
(25)
A
Proof.
(i) First, note that Hr = {Etr,> >o aee isomorphic to (((a))))))))) >0) and is thus a Hilbert space. Then, Kx(x,x') = (z))))))) For any f E Hr, let f = =, u4r,, then
Σ, a < 20) where e; = 1/ Wh, so it
(f(x"), Kx D' = Σ A (((Assl()))) = Σ un((() = f(x).
23

---

Published as a conference paper at ICLR 2024
(ii) For any E, 4 Hr, there is E, X, uu < so by definition. So for any X; 0, there must be u; = 0. Let >; A; 1/2 u,bi. Then, |1911 - E; AT < oo, meaning that E (P)). And there is f = - so f R(T"), which implies that Hr S R( Mcanwhile, for any f g E )( let E, 1; then Σ, v2 C0. Then, f Σ, X/22 0 by duality, 2: 1; (1, /2 0;)² 00, meaning that f E Hr, R(T*) S Hr.
(iii) For any f= u Hr. let g 2/,,, u.thr. By the proof of (ii) we know that f H g is bijective, and If ur = 1 Moreover, g E span( (pan((0;),n (0:)),0))
(iv) Let 9" Ei uØS, Then, since we have f g'gg E. uns. for any A 0, there is v2 u and for any A; 0, there is ue = O. Let g" = 90 + gt, where 90 = E, A, 1/2 u,or and 91 1 9o and I'r11 = O. By duality, I1 go belongs to the linear span of (d. 91 - 9u- As will show later, Eqn. (1) is equivalent Eqn. (25), which is the random walk normalized Laplacian over the augmentation graph (Chung, 1997, Section 1.2). This is equivalent (g" (I-IT" Pa I <1l which is further equivalent to (go. (I (T- ))o) Pa + ll S <(llool) +
( (note that IT 91 0). This implies that (go, (I - --" PA S ell.ool?... i.e. go satisfies Eqn. (1). Since we also have ʃ go.. 9u must satisfy Assumption 1, we can choose g" gu-
Next, to show the equivalence to Eqn. (25), We just need to show that (1-IT")))) = Ex E (A, [(0 (A) - g" ( And indeed, we have:
g" -f f al
PA
11901122 -1 o(a)g(a') ((d))))((d(())p()))) p(a')p(a)dada' p(a)p(a')
1
2 El (A' - -x (A)g (A')
ExxPx TA ( (AA)))]
desired.
E PROOFS FOR SECTION 3
E LOCAL GAUSSIAN COMPLEXITY AND LOCALIZED RADEMACHER COMPLEXITY
We first provide the definition of the complexities we will in our analysis. For a function f. 2 be its mean on the downstream samples.
let Il. == 3, -.= f
Definition 4. (Waimwright, 2019, Eqns. (13.16) & (14.3)) For any B,C O, define
Fo = { f2 fc
}
{
2B
S Ví=
E Hig
llfl S
(26)
vI
Then, the local Gaussian complexity around f at scale 8 > given by
sup (o) = Wis ... dehn [EFo
Σ wi wrs(
(27)
72
where w1, w,, are i.i.d. NN0, 1) variates. And define
F. f== a Hegllfallu I
B
(28)
VI-e
Then, the localized population Rademacher complexity of radius ō > 0 is given by
Ry (8:
E Bn
sup SCF-.Ili
aif(au)
(29)
G1. Cr.".,
71
where "1,"" on i...d. Rademacher variables taking values in (-1, +1} equiprobably.
24

---

Published as a conference paper at ICLR 2024
Our master plan is to apply Theorems 13.13 and 14.1 of Wainwright (2019) to fy. = (I
where I is the projection operator onto & in (Px),, and fa is the projection f onto Hi w.r.t
(t Pr.. Therefore, we need to bound F6; and (6; We start with the following uniform bound:
Proposition 7. If I and ll S T. then (f())) S KT for all x.
For any g = Σ, uid E (PA) Proof. By Eqn. (4), we have p(a|2)) E: (i)
such that |ą S (I* ()) = [ I'aappplll) = E (x)) Then, by Cauchy-Schwar inequality, we have for all 2, f(x)² (I)))(x)² s (O)))))))) uf) S R²T?.
This proposition immediately implies that f" and fa are uniformly bounded:
n² Corollary 8. For any 1* E eT) Eqn. (6) ensures that 119 S I- so [f(x)) S *B for all x. Moreover, ||| PA s llp, implies that f S VI-E" x.
Vi-c
B and (f)())) S cB
for for all
We will also use the following simple result in linear algebra:
Lemma 9. Let D = diag(An, (A1, )2, where A 2 X2 2 ... 2 0 and X - 0. Let Q be a matrix with d rows that are unit vectors. Then, Tr(QDxQ Q d + Xa-
Proof. Let qi be the i-th column of Q. Then for all j E [d]. there is Ei' q: qi J. And for j >
Ei ae qi S d. Thus, using Abel transformation, we have
8 Σ j=1
Tr(.0"" Xay qi
Σ qi qi
(A3 A5+1) S
A
which proves the assertion.
This result has many implications. For instance, for any rank-d subspace of Hr, its trace (the sum of its eigenvalues) is at most (x)d).
Now, let bound (6; Fo) with the following result:
Lemma 10. (Application of Wainwright (2019, Lemma 13.22)) Let H. be an RKHS with reproducing
kernel K. Given samples : _Tn, les K be the normalized kernel matrix with entries K (i,,) K (( in)/n. Let 41 2 ... 2 Un 2 O be the eigenvalues of K. Then for all > O. we have
1
2
E
sup
min 18 ,,,73,
(30)
[I ST.I
where Wi, W., are i.i.d. (0, 1) variates. We apply this result to K Kx. By Definition 2, all elements on the diagonal of K are at most /n, so Σ, u4; = Tr(K) S x2, Thus, we huve
8x2 B
(n(5;Fo) S
for any Hé.
(31)
n(1 c)
Regarding Rr (5): F. is also a subset of RKHS H., which is the linear span of d and f*, and is a subspace of Mr whose rank is at most (d + 1). By Lemma 9, the sum of eigenvalues of H., is at most Sx(d + 1). Since [If 11 < B_
Vi-e all f E F. satisfy [|f < 28 So we have the vI-c"
following bound for (o;
Lemma 11. (Application of Wainwright (2019, Corollary 14.5)) Let M2, be the eigenvalues of the RKHS H.. Since (Fn.) S rank(H))) + 1, we have
2
8B² Σ min { 82, 1-e n(1-c) (11-) Sx(d+1) (frank(H()) S dd (32)
B2
(6: S
and for an arbitrary Hip can simply replace Sx(d + with Sx-
25

---

Published as a conference paper at ICLR 2024
E.2 PROOFS
Lemma 2. Suppose V1,-"" are i.i.d. [O, o²) variates. If has d dimensions (d can be oo).
then have the following uniform bound all I -= Fn)ee)):
Con( +oB)
Sx(d +
P AA EF
S 59115 1-
MI
1-c
C22 (2o22nSS(d+1 - exp (-V2
2nk2 B2
2 1 - exp
-exp
1-c
where f& -I* is the projection of f onto Hi w.r.t (, and co. C1, C2 are universal constants. Moreover, Sx(d + 1) S min {d + k,
Proof. By Proposition 7, all functions in F. b-uniformly bounded, with b 2xB vI-r" And
obviously F. is star shaped, meaning that for all f E F, and all | E [0, 1]. 8f Fx. Let 8B2 m(1-c) Sx(d + 1) 2 6R, F.) Then, by Wainwright (2019, Theorem 14.1), we have
ta
t² P |11 - 11 1 I1 + 2 S exp (-ca
nt B2
for all f F
(33)
for universal constant C1, c2. We know that ff f F, and fé f which means that
[1 11 2 211 17 + +C) v (116 1 2 2 511 M + 1 =1
[1I (Ii-
(34)
n122
S c1 exp
b2
Let 83 20
Sx² B² nll-r) By Lemma 10, have 53 2 209 (on: And Fo is also star-shaped.
Thus, by setting 7 1/2 in Wainwright (2019, Theorem 13.13), we have"
no
P [ 1 fl
11. 320
exp
(35)
20t
Combining the two inequalities above with the union bound, we obtain the result.
Now we prove Lemma 3. Without loss of generality, suppose h,, ,, are linearly independent.
Let H. := span (... .hr Let g" go + 891, where go -I g" 91 Io. and || 1. So by Lemma we have:
Proposition 12. [ (G
1/2 h/2 +
(G 1/2 h/2 ) 91" 91 S A + ..... AM+1-
Proof. Let G 1/221,, h1, GT 1/2 hargi Q where Q a matrix with (d' + 1) orthonormal
rows. Then, GT 100-1/211)))
h1).... ... (G 1/2 1" =QD," 1/2 1/21/22" Thus, we have
|| (G 1/2 + ||" (G 1/2 //22//1 112
|[ I2, THQDAQT).
Then, applying Lemma 9 completes the proof.
Remark. This proposition is the functional version of Fan (1949, Theorem 1).
Notice that [T 1/2 With this, prove Lemma 3:
(G,11/ MI +I"." 1/2 1/2 M 1/" Tr(G F.).
Plecase refer the proof of Wainwright (2019, Theorem 13.13) for removing the universal constants in this thcorem.
26

---

Published as a conference paper at ICLR 2024
Lemma 3.
For any f" FB) there is
7² T T+
1 1' S
B².
1 - 1-6
Let a² = 1190ll Pa. and B = 1190 9" 0' By Corollary 8, + - B² I-c- Eqn. (6) implies
Proof.
that
- () +3² < ||'' fl S + = 8²72 20,3T,
since
|
- |1 p/ = = and | 91 12 S -² by Proposition 12. Thus,
B²
(1-13)3² S c(a2 + B2) + 20,37 S (e - () + (8) s (e + T-C
Thus, we have life f' = (g0 - gg 00) = [ gi'gi113 Px - 8372 which leads to the inequality we need to prove. Finally, by setting h4 = Φi we can see that 2 S - F). And for all S d. Tr(G Fb) S
(d'), so 72 Ad+1-
F PROOFS FOR SECTION 4
Proposition 4. For any v = [8],,... del where Vs E (2xx),, it holds that
err( FBFFFF(T))) Ad+1
Ad+1
1
B² given that
(36)
A4t1
1- - 1-11-- 2
To atlain equality, is sufficient for v span the top-d eigenspace, and also necessary if Aa+1 < Xa-
Proof. Necessity: Since V is at most rank-d, there must be function in spanfe/.. da41) that is orthogonal to L. Thus, we can find two functions fi. E span( dd411 such that: Il = px = 1. fi is orthogonal to V f2 = u v (which means that f2 1 fi). and 41 E sp Recall that A = 1, and tin 1. Let 1 - 01 + =2/22 then of + ag 1. Without loss of generality, suppose 01, 02 E [0. 1]. Let fo = a2fi - a1 Then, f = 1, fo 1 11. Note that we also have (v/1, fo)
0 by duality. Let $1. 32 [0. 1] be any value such
that f = 31411 + 32 satisfies || = Bf + $8 and / S T- This is satisfied as long
as 83 S - because [| S 87
Aet1
+ As+1
1-4411 B3
Moreover, we have
Aa+s
Bf E (T(
It is easy to show that F(a1) = 0n33 + 02 01.3 + VI 3 (01 E [0, 1)) first increases then decreases, so F(a))² min (F(0)2,
min the worst given that it is most 1. in which the prediction error of Bf is [[[ f2 = B² = - 1-
which can be T 1-ci-Xa 1-A0+7 A041
A1
in
Arra B². Thus, for any V we can find a function Bf e
FB(T:e) e) such that min err(w" T
Bf) 2 I-c 1-A41
+ Aara A B².
When Aa A441. to attain equality, we need 0 = 0, and || = BP
Xati which means that fo = 0-41- Thus, only fi = fo that1 is orthogonal to V, so v must span the top-d eigenspace.
0
Sufficiency: Suppose v spans the top-d eigenspace. For any f E FT((;)) such that f = E, u.t.. have Σ,u; ug B², and Σ; 1-c di
A uz < O. Let Σi2d+1 uf and b Ei= uf. Then,
min, err(W' (w) T v f), and a + b < B². So we have
d r -bb +
1
d0t1 a da+1
1-c-X
since
decreases with A
A
A
2 -e(2² a) + 1- 1
Ao+1
d4+1
Adt1
= -c + c)
Xdt1
which combined with the necessity part implies that e FB(Tee)
_ Ad11 B². II-II-041
27

---

Published as a conference paper at ICLR 2024
Lemma 5. Suppose there exists constant C 0 such that [g S Ca for all g = w T á where 1 $.. Then, for any 5 > 0, it holds with probability at least that
2
CK +K+
F)
2 + 2log
d.
VN
Proof. Since multiplying an invertible d x d matrix to & does not change either Tr F) or Tr(G-1F), for simplicity let us multiply G-1/2 to &b, so that (b).) = &i,j for all i, E [d] (ie. G I). Define F1 = {f E Hr I1]fller I 1}. Its Rademacher complexity is given by
N
1 N
Rx(Fi)
E
E
sup SEA
(37)
Z1 AN
By Mohri et al. (2018, Theorem 6.12), we have Ry (Fi) S K.N 1/2 Moreover, by Proposition 7, all
f Fi satisfy [f(x)) I for all z. Thus, by Wainwright (2019, Theorem 4.10), for any > 0, with probability at least 1 - 5/2, it holds for all f E Fi that
2 (xx) - [X)( < ( + + log S
K VN
2 + 2log
(38)
Define matrix M- = 2-112 G 1/2 = (ma.) (i se |M S 1, so
Σ Li=1
m 1 1 for all j E [d].
Consider Tr((I- - M). For any j E [d]. we have
((I GI- 35 Σ
m
Φ0 Σ
m
i=1
Pa
PA
Note that
E m
8.
11,,
s L. so
|ios (Ei= mi.
2i)
S
Bloti
(I
ma z
PA
C², which means that --2px we have
(0g
mig ) E Fi. So if Eqn. (38) holds, then for all j E [d].
N
Σ masea )) (1x)
massu )) (x))
((I G)M))))
E r
Σ i=1
=
N k=
i=1
2
Cr
<
2 + V
2 log
x)
VN
which implies that
Tr G f = Tr (T--)) - = Tr "I((I-G)))) < 2 V 22 log
Cnd
VN
Next, define =2 = Wfi22 f1, E Hr, ||ll 5 1 Hr S1.. By Proposition 13 (proved after this lemma), we have Ry (F2) < 112 And all f E Fg satisfy [f(z)) s x² for all by Proposition 7. So with probability at least 1 8/2, we have for all f f F2.
N
2 log 8
2 22og
x²
Σ (((x)-E[((x)) k=1
< 29 2
2 v
(39)
5
vNN
Note that 11 < 1. So under Eqn. (39), we have for all i E [d].
N
x²
(15,, 15)0) -6b, b5) = N
Σ
(ax))) -
$ 2 +++/0 22 8
VN
28

---

Published as a conference paper at ICLR 2024
2 which implies that Tr (f-G-IF)= (F ( s (2 + / log :)
x2d
VNN
Finally, applying the union bound completes the proof.
Proposition 13. Let F2 = [fifz f1. E Hr, llfaller S 1, f2lller S 1}. Then, RN( - vN
Proof. For any Mxx) (x)) E F2, let f1 = g' and f2 = -* 92, where 1 S and
l PA $ 1. Let 91 = =, u and 92 = Ei v. Let u [11, u22 and v [11, U2 :.
Then, ||u||2 $ and || $ 1. And we have fi = Σ, X u,,b, and f2 = Σ, 2//2,, e,
For any E X, let (x)) 2 (z)) 22 ()) Then, Si (1x)) u TuT (2x) and f2(x)
ot (xx). Denote V Ve (=x)). Then, V Ve x² for all l E [M]. So for any 5 {z1,"" ,NN ).
the empirical Rademacher complexity satisfies
Ug (F2)
sup |1
T T
V 4
N
k M V
N
1
E
Σ k VA
N
1/27
Σ Foneree)
E
Tr
Th U V
N
N
Tr
oKoN Y.T.
(Jensen)
N
R,1=1
Tr
E W LnNT
N
N
Tr
V&
V
N
1
R²
NR
N
VN
Then, since Ry (F2) Es s (F2)), we obtain the result.
Lemma 6. Suppose å. = o. for E [d]. Let Ya == >maxx (G)/Amin Amin which is the condition number of G. Then, for any 8 0, both
d
2 (A + 1)k2 2 220g 5
Σ
A 2
VN
/=1
1=1
and Eqn. (12) with C = x\j+nd 1 8.
1/2 76
hold simultaneously for Hig = Hi with probability at least
Proof. Denote : = 1d1. Øul and :; = 1Ø1. Tel. Let क; = P where P is matrix with d rows. Observe that for any = Σ; ug such that ||_|| < 1, we have =
Fr ( - Let u = (un..... ua), then there is u क5 u P so |PT ull2 $ 1.
Thus, have E. A f Pa =|pPT D Jull= P D - u
29

---

Published as a conference paper at ICLR 2024
equal to the square root of the largest eigenvalue of PT D 1 |P D (PPT((PP is
So we just need to show that IPPTDA (PPT)) S X 1/2
and by using two simple linear algebra exercises: (i) Amax (AB) S Amax (A) A (B) for positive definite matrices A and B, and (ii) AB and BA share the same non-zero eigenvalues (Sylvester's Theorem), and the fact that G = PP we can show that the largest eigenvalue of this matrix is at most Y GGG
Therefore, we have ||PT D implies that Ex 7, u.
Xa (PP ) Plle I 1d' which combined with || (Tulle S 1 1/2
Pa S 2/ 2/ By Proposition 7. - (E, A, u.):) (x) S Rt l (
1/2
for all I, we have [[T* (E; (( (a)l
(z)p(z|o)dr) & RtginG 1/2
for all a. This means that with C' = x probability at least 1 - ō, both Eqn. (38) and Eqn. (39) hold and they lead to Eqn. (12).
1/2
GG , satisfies the condition of Lemma 5. Therefore, with
Now let T = Q where Q is a matrix with d rows. Consider two matrices QQ QT, E Rdxd where D = diag(A1 ,2, ), for which we have
QQQ (i. (QQT) = (0b,, b3)
and (QD (i,1) = (T* Ss px)
We have (10:; PA Pa)sE1) = I and ((T "T"",, (x)))
D2a diag( (1,,". Aa). More- over, for any g- = 47 such that Px S 1, there is g = IT (Ei A u,,;,, and obviously p I A Thus, we can show that for all a, [g(aa) S k\ -1 which means that & satisfies the fourth-mone of Lemma 5, for all u E IRdd such that |/u||2 S 1, we can show that
control assumption in Lemma 5 with C² = ndd So similar to the proof
2 x² x²>=" 8
|u TQQ Tual (u D $3) PA
{u p3 u
F3) PA
S
2 1
log
VNN
2221" VN
which implies that lIQQ ||2 S 1+ (2+ / 2
It is easy to show that all non-zero
eigenvalues of Q Q are also eigenvalues of QQ so lIQ
Qll2 S +
(2 + v log 3)
" ~²'
vN
Moreover, similar the proof of Lemma 5, we can show that for all i E [d].
x²2 P7=
1QQ - 1) (1.5)) = 2 v
22og
(40)
VN
x²
[QD,00 - Dxa) (1.5) S 2
2
2og
(41)
8
VN
Let qi be the i-th column of Q. Then for all E [d]. q. q: S 11+ 2
220g 3
x² vN
And we
also have Ei g.f qi = Tr(Q Q) = TTQQQ ) S d + (2 + v 2 2
x² 3²x-'d 'd VN
Thus, we have
x²
Σ A/ -
2 + 2 log
d = Tr = QT Q)
VN
1=1
Σ 1=1
Neq. T q: = Σ
Σ
q T
(Λj - 1
j=1
1=1
x²>==
d
2 2 2 v
x²2
<
Σ
Ao
1+
2 / 200g
d +
log
VN
VN
1=1
1=1
which proves the assertion.
G PROOFS FOR SECTION 5
Example 1. Consider X' = (-1, where each x E X is a vector of -1 and with length dx- Let px be the uniform distribution over X. Consider random masking augmentation, where for any
30

---

Published as a conference paper at ICLR 2024
X. each coordinate x' is randomly masked to be 0 with probability E (0, independently, where a is the mask ratio (0 is the [MASK] token). Then, x²² = (2- aa)xx.
p(o) da, and by symmetry, the fight-hand-side the same for all Given an a, suppose has r coordinates masked and (dæx -r) coordinates unmasked. Then, there are 2" possible r that can be augmented to a. For each of these 2, p(a[z) = a" -a) So
Proof. We know that x² N l p(o]z)²
p(a) p(a)))p(z)dd = 2" xo-dxaal a))x= Thus, we have
ʃ p(alr)) p(a)
dx (
dxx'
a²r(1 ax)2x=22
x²
da 2
2 dx a² a dx -r
r
dx dx
x²(2 (2-2a)mx -2a)xx xx
= -aa +a 2 2a) = (2 - aa)ax
which completes the proof.
Example 2. Consider random block masking augmentation with ratio o, which for any a E X
masks x², xi+1, ,...,
,.... "...,,"trr"-".. x=tr-1 for a uniformly random i, and r = =[dx] 1. Then, x² s [2(1-a)ddx
1 Proof. For any a, we have p(a) dx -x-r+12de-1 and p(alz) dx dx-++1 if x can be augmented to So there always is pla m)
1
2dx < _2 ((1-a) Thus, have n²2 (11-o)dx.
p(a)
Consider random block masking + flipping with ratio a, where for any E X, firsi ,...,,ttr> are masked to be 0 for uniformly random and r = [gd and then each remaining coordinate is randomly slipped sign (1 - -1 and -1 - 1) with probability a
Example 3.
independently. Then, x2 & [[a2 - 2a + 20(1-a22)) dx
Proof. For any a, we have p(a) dx -I+I 2 unmasked (dx - r) coordinates, a and have k disagreeing coordinates. For a given k, there are (dx r 1)( 12x (x)")) possible a, and we have p(a[x)
Suppose a is augmented from x, and among the
dx dx-r+I (2) -dxx-r~ Thus, we have (dx (dx- ( 2 2 2adx
/ p(a)x)2 da =
dx
dx - r)
Σ (dx - r
p(a)
k
k=
dx-r+I -"x-
dx
dx -r)
-1) 2dx 2xx--
4 (a)) x
(1 (1-a+
a² dx dx-r-k
h:
k=
022
a= dx
-dx-r
+ aa
4
< (a² 2a + 2)
dx-r S (a² -2a + +2) 22-22-112/2 (1-a/2)dx
which completes the proof.
H EXPERIMENT DETAILS
This section contains the details of the experiments we conducted in Section 5.2.
H.1 ESTIMATING THE AUGMENTATION COMPLEXITY OF MASKED LANGUAGE MODELING
We estimate the of four augmentations, namely Random masking, Block masking, Random masking + Flipping, and Block masking + Flipping, on the widely used NLP dataset wik rikkpedia-se a
e. Since it is prohibitively expensive to iterate over the entire dataset, instead aim to estimate к with a subset of the dataset. In our experiments, instead of estimating sup,, K x (z, x), we estimate its 99th percentile. This is because of two reasons:
31

---

Published as a conference paper at ICLR 2024
(i) sup, (x, 2)) is statistically impossible estimate from subset of data without any extra assumptions on the distribution of Kx (x z). This means that without iterating over the entire dataset, cannot get finite confidence interval. The percentile, be estimated with a finite confidence interval via sampling regardless of the distribution of Kx 2.), which is well-known result in statistics (Hahn & Meeker, 201 1, Section 5.2).
(z,
the other hand, can
(ii) Almost all real datasets contain outliers, i.e. samples that are very different from the other samples. These samples will have very large [x æ). because their augmentations rarely overlap with the augmentations of other samples. Therefore, sup,, Kx (z, x) itself is not really meaningful because is too sensitive to outliers, and cannot show how the augmentation works most part of the population.
As demonstration of (ii), the following an example outlier the wikipedia-sseee -simmle dataset," and it has a very large Kx(æ,z);
Geroldsgrün is municipality in Hof, Bavaria, Germany. Geography
Boroughs
Dürrenwaid
Würrenwaiddrrhmmee
langenbachtal
Langenau
Langenbach
Lotharheil
Mühlleiten
Geroldsgrün
Geroldsreuth
Großenreuth
Hermesgrün
Hertwegsgrün
Hirschberglein
silberstein
Steinbach
Untersteinbach
References
We can see that this sample is very different from a typical sample in the dataset in that (a) It contains lots German words while this is an English dataset; (b) It mostly consists of individual words, while most other samples are comprised of sentences and paragraphs. Such outliers are rare in the dataset, but they result in a large Kx x (z. x).. Therefore, we instead estimate the 3-th percentile of Kx (z, æ)) for some 3 close to 100. This is a more robust estimator that can be computed with subset of data.
Specifically, sample The maximum length is set at 64. For each 2,, we estimate f pla; a)
1000 samples #1,--. Fm from the dataset uniformly at random.
p(x,) ((a]2x))da via sampling r = 255 i.i.d augmentations from p(a|z;) and computing the mean of p(aglas), So suffices to estimate p(zs(a,) and p(zz). To do this, we leverage bert pert-base-uncee
p(*,)
which is
bi-directional language model. Then, we can estimate p(z(a) for any x and a as follows: For a sample
x (1) ... ,0(0] (x(0) is token) and its augmentation a, there is
log (araa= logg (2(11) (a) + log (212)a a, a()" +... log (100 a, (
So we can estimate log (z|a) in an autoregressive fashion: First estimate log (z (1) (a), then replace (1) with (1) and estimate log (2(2)) a, x (1)), then replace (2) with (2) and estimate
Dataset available huggiingfaccccooddunntttttvviiie The index of the example outlier 199562.
32

---

Published as a conference paper at ICLR 2024
0.2
0.2
Density
Density
0.1
50
100
150
200
250
300
350
50
100
150
200
250
300
log
pirial p(a) f(a(z)da)
log
p(ale) M() ((a)z)dd
(J
p(zla) p(alz)dda) for random masking. Left: a = 0.15; Right: a = 0.4. p(a)
Figure 4: Plots of log
Table 99th percentile of estimated log Xx x) on wikipediaa-impll
Random Mask
Random Mask Flip
Block Mask
Block Mask Flip
a 0.05
353.94 + 4.63
350.99 + 3.20
346.77 + 4.47
350.46 1.43
Q 0) 0.10
348.45 + 2.94
349.14 + 1.15
339.86 + 5.62
343.73 + 5.76
a 0.15
346.40 + 1.94
342.77 4.17
327.13 + 2.33
336.81 5.27
0.20
339.03 +:2 2.88
335.01 2.65
313.90 2.55
322.87 3.55
Q 0.30
325.08 + 2.35
312.21 + 5.96
280.65 + 1.87
302.60 2.49
0.40
307.00 = 4.55
280.76 0.71
254.40 = 1.48
279.24 = 4.59
a 0.50
278.59 +: 2.46
241.59 + 4.50
217.35 + 2.81
240.73 + 1.58
a 0.60
250.01 + 3.43
198.15 4.01
187.60 2.45
205.45 + 2.63
a 0.70
214.14 + 0.88
151.58 3.32
152.70 + 1.60
164.48 + 2.59
a 0.80
160.24 + 2.44
84.51 = 2.35
113.86 + 2.39
108.44 = 1.43
log p (z(3) a, (()), .I))),, and so on. The bi-directionality of BERT is important for such estimation. To compute p(r), we set to be fully masked sentence, and then compute p(r(a) p(x).
plot the histogram of p(z) p(a)xx) (() estimated using the above method of random masking with mask ratio 0.15 and 0.4. The red dashed lines indicate the 99th percentile, and cut the tail of the data while preserving the bulk of the mass. Thus, we report the 99th percentile for this dataset.
First, we decide the percentile of Kx(x, x) report. In Figure
log (s plala)
that they
We study four masking augmentations: Random masking (mask r = [odæx] tokens uniformly at random), block masking (mask block of tokens uniformly at random). random masking and flipping (mask r/2 tokens uniformly at random and then replace r/2 remaining tokens with random tokens), and block masking and flipping (mask block of r/2 tokens uniformly at random and then replace r/2 remaining tokens with random tokens). In standard BERT, the number of randomly replaced tokens is much smaller than the number of masked tokens, but here we make them the same in order to magnify their difference.
In Table 1, we report the estimated 99th percentile of log x (z, w). Each experiment is run five times with different random seeds, and we report the average and standard deviation. We can see that:
decreases as a grows, showing that a stronger augmentation has a lower complexity.
(ii) Random masking always has the highest complexity. Regarding the effect of block masking and flipping, block masking has greater effect when a is small and it makes the complexity lower, whereas flipping has greater effect when is larger.
To conclude, long there is way to estimate from the dataset. Note that in our estimation, BERT is only used for estimating p(z(a). x itself is defined free of any model.
p(ala) p(z) p(a) we can estimate k via sampling
p(ala)
33

---

Published as a conference paper at ICLR 2024
H.2 ANALYZING THE EFFECT OF AUGMENTATION ON DOWNSTREAM PERFORMANCE
Our theoretical analysis implies that as long as Assumption 1 is satisfied, smaller k, i.e. a stronger augmentation, leads better generalization. However, this does not mean that stronger augmentation will always lead to a better test performance of tighter generalization gap, due to the following reasons:
(i) If the augmentation is too strong, then the model might have low training performance. In the extreme case, if the sentences are 100% masked, then the language model can learn nothing.
(ii) Stronger augmentations usually work better with larger models. For instance, Wettig et al. (2023) demonstrated that large models with strong augmentations can achieve higher performance than standard models and augmentations. However, for small models, stronger augmentations can lead to lower performance. One possible reason is that small models are not expressive enough, but we conjecture that this not the true cause since a Transformer with a moderate size is already sufficiently complex. We conjecture that the real reason lies in optimization, in that it would be harder for the optimizer to find a point close to the global minima with a small model and strong augmentation. The big model, the other hand, has more parameters and could be easier to optimize in practice.
(iii) If the augmentation is too strong, then Assumption 1 will be violated, and our results will not hold that the generalization gap can be large.
We study the effect of the augmentation on two real NLP downstream tasks, namely QNLI and SST-2. We study random masking with different mask ratios. Unlike the experiments in Wettig et al. (2023) that applied the pretrained encoder directly to downstream tasks, in our experiments we explicitly use the average encoder Eqn. (7), which estimated by sampling 16 augmentations for each sample and then averaging their embeddings. We do, however, fine-tune the encoder during downstream as people do in common practice, which has been proven useful, & even though our theory does not analyze fine-tuning. We acknowledge that there is a discrepancy between the theory and real practice, which we aim to address in the future.
Details of our experiments:
(i) We train roberta-large models using the fast training recipe provided by Wettig et al. (2023). For a 2 0.15, we directly use the Huggingface checkpoints they provide. For 0.05, 0.10, we pretrain new encoders using their source code without any modification. Note that these models use pre-layer normalization (see the original github repository regarding this detail). We use 8 NVIDIA A6000 GPUs for pretraining.
(ii) For downstream fine-tuning and evaluation, we use the official repository provided by Hug- gingface. The only modification we make is that we explicitly use the average encoder. The classifiers are trained for 3 epochs on QNLI and 6 epochs on SST-2. All hyperparameters are kept the same as the official repository. We use NVIDIA A6000 GPUs for downstream training and evaluation.
The results are plotted in Figure 3c in Section 5.2. As discussed there, the "sweet spot" where the model achieves the highest test performance is in the middle, where the augmentation is strong enough to have good generalization, yet not too strong to break Assumption 1 or to make the training performance too low.
'We find in our experiments that without fine-tuning, the downstream classifiers cannot achieve meaningful training accuracy, merely higher than 70% for a binary classification task. With the training accuracy so low, there is no empirical generalization gap, and the test accuracy is even higher than the training accuracy in many cases, these results are not meaningful.
Inttps://g wtw..tcoo/ppoo..eeeooeeo com/pr nccetonn11//inkkyrrr hhtps: / thub. wtw.b.ccoooooooo..oooooommo.ooth wtt.coo///////e htap///ggingfeeeeerrraffmmmrrrre n/examp 1ea/
pytorch/text--caass fication
34


Reviews:
[{"review_id": "tGwtlm1Awx", "reviewer": "Reviewer_tXrm", "summary": "The work provides a statistical analysis of augmentation-based pretraining using the so-called augmentation operator. This operator is naturally introduced and several properties of this operator are stated. Upper and lower error bounds, which apply to arbitrary encoders, are derived. The augmentation complexity is emphasized as a central object contributing to error analysis in both learning and augmentation.", "questions": "General remarks\n\n-The assumptions and notations should be pointed out more clearly. For instance\n\n       -it seems it is assumed, but I did not see it stated explicitly, that $P_\\mathcal{A}$ has a density, and that $p(a|x)$ and $p(x|a)$ exist.\n\n       -What is A in E(g(A)|x). Similarly, X is not defined as a random variable (it is clear from the context but should be defined).\n      \n       -The notational difference between $p$ and $p(\\cdot,\\cdot)$ is small.\n\n-Is $\\hat{\\psi}$ always assumed to be of the form (10) or is (10) just an example?\n\n-Can you explain the choice of the distribution $P_{AX}$ a bit more? Particularly, it could be emphasized how the conditional distribution $P(\\cdot |x)$ has to be given. You can refer to the given examples.\n\n-Can your approach be used to build or suggest an augmentation-based pretraining? Maybe under assumptions on prior information on the distribution $P_X$?\n\n-Proposition 1, (iv): What does it mean that g* satisfies Assumption 1? Assumption 1 is an Assumption on $f^* \\in \\mathcal{L}^2(P_X)$, or does the statement in (iv) refer to (1)? The proof of statement (v) is done in the proof of statement (iv); I think those parts can be moved to the proof of statement (v).\n\n-In (1), how are A and A' drawn from $p(\\cdot |x)$? A motivation for (1) would be helpful. The space $\\mathcal{F}_B(\\Gamma;\\epsilon)$ is a central object, it would be good to motivate it more.\n\n-Can Assumption 1 be verified; is it a restrictive assumption?\n\n-Can $\\epsilon$ or violation of Assumption 1 be inferred during the learning process?\n\n-Definition 3: Is $G^{-1}$ well-defined, even if $d = \\infty$?\n\n-Remark after Lemma 2: What is $\\|\\cdot\\|_n$?\n\n-How is $d$ chosen? Equation (20) would motivate to choose $d$ such that there is a large gap between $\\lambda_d$ and $\\lambda_{d+1}$, but your result improves such a guess for $d$ compared to HaoChen21 and Saunshi22.\n\n-Is $\\bar{\\phi}$ an eigenfunction of $\\Gamma^*\\bar{\\Gamma}$? If so, how restrictive is the condition $\\hat{\\phi} = \\bar{\\phi}$; is this just the choice of the encoder?\n\n-It is good to see a separation of model and augmentation in the error analysis. Do you think a combination, for instance, through augmentation based on the model, can be fruitful, too?\n\n-Does Lemma 9 need a condition on boundedness of the matrix $QDQ^T$ or absolute convergence of the sum $\\sum\\limits_{i = 1}^\\infty \\lambda_i q_i^T q_i$?\n\n\nMinor comments:\n\n-Did you compute the ``sweet spot\" mentioned in the introduction? Can it be computed?\n\n-In Section 2.1 some reference for conditional expectation and RKHS could be added.\n\n-Proof of Proposition 1: The proof should also contain showing that $\\mathcal{H}_\\Gamma$ is complete.\n\n-Are $K_X$ and $K_A$ square integrable? It seems so from (4) and necessary from the proof of the statement that $\\Gamma^*\\Gamma$ and $\\Gamma\\Gamma^*$ are integral operators. Are conditions on P_AX$ needed that (4) holds?\n\n-Below Definition 2: I don't think that convexity of $D_{\\chi^2}$ is needed because the bound $S_\\lambda\\leq \\kappa^2$ follows from (11), via $\\kappa^2 \\geq K_X(x,x) = S_\\lambda$.\n\n-First sentence of Section 3: Assumption 1 could be mentioned to recall the bound for $\\|f^*\\|$ in $H_{\\Gamma}$.\n\n-Lemma 5: Below Lemma 5, it is mentioned that the required 4th order bound holds. Why is the condition for such a bound included in the statement of the lemma?\n\n-Around Definition 2 you can refer to Examples 1,2,3 for examples of the augmentation complexity.\n\n-Are there exponential lower bounds on $\\kappa$ in Examples 2,3?\n\nTypo: Page 4: \"the the\", the word ``the\" was used two times", "limitations": "", "rating": 8, "confidence": 3, "soundness": 3, "presentation": 2, "contribution": 4, "strengths": "The paper gives a structured foundation for augmentation based learning. The augmentation operator is well motivated and its central role is well highlighted. The obtained error bounds show, in a good way, a separation in the contribution of the model and the augmentation. The analysis seems rigorous and, personally, I find the appearance of the spectral properties of the augmentation intriguing.", "weaknesses": "-I find the text very dense and compressed; and thus not easy to read. I acknowledge that the task of fitting the results and concepts in the 9 page limit is difficult. Therefore, I would suggest that the authors consider whether the text is better published in a journal with fewer page restrictions.\n-Some concepts could be motivated more. That is particularly helpful for readers less familiar with the topic\n-I find, that the assumptions and notations should be pointed out more clearly (see questions)", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "dD8WwpAQ52", "reviewer": "Reviewer_umfH", "summary": "The paper builds upon recent work on self-supervised representation learning to study the effect of augmentation on the generalization properties. The key insights are that 1) augmentation can be seen as introducing a natural kernel $K_X$ between the inputs, kernel which has to be estimated based on the unlabeled data, 2) the labeled data serves then to find a linear predictor, whose generalization performance is bounded by the authors, with a bound that crucially depends on some quantities related to the augmentation procedure. The main point for the proof is an assumption, called 'isometry' (w.r.t. the augmentation) by the authors, on the optimal solution. Numerical experiments on estimating these quantities are then performed.", "questions": "- Which formula corresponding to (1) can be found in HaoChen et al.?\n- Can you quote some articles that used a bound constraint on the final predictor as you do? In general imposing constraints simplifies proofs, is this reason of your choice?\n- Can you give examples of augmentations that do not lead to $\\int K_X(x,x) dp(x) <\\infty$? $p(a|x)=\\delta_x$ for instance? Similarly, give some examples where $\\Gamma \\Gamma^*$ is clearly Hilbert-Schmidt (Gaussian blur I guess)? I ask this because it is not a given that $K_X(x,x)\\in\\mathbb{R}$ in my opinion. There are kernel integral operators that are not related to RKHSs.\n- Is it true that (8) implies that $u_i \\approx 0$ for $\\lambda_i < 1-\\epsilon$? It seems to me to be the case. Eq.(8) can then be interpreted as requiring that $f^*$ is mostly encoded only on the first eigenvalues of $\\Gamma \\Gamma^*$, in other words that it is 'simple/low dimensional', and that has little to do with RIP I would say. Please discuss further this aspect, and if possible suggest a reformulation of Sec 2.1.\n\nI would upgrade my score if these questions and those of the other reviewers are answered.\n\nThe authors could underline that Cabanes et al. also restrict themselves to the quadratic loss despite the contrastive inspiration coming from classification, and use the arguments in Cabanes et al. to justify it through calibration.\n\nMinor:\n- p2 please add some context words, 'upstream/unlabeled', 'downstream/labeled' to simplify the reading for non-experts.\n- p3 last line $\\phi$ not $\\phi(a)$ which are values\n- p4 'the the'\n- Quoting Forrest Gump, ``my mama always said'' never to start a sentence with a mathematical expression.", "limitations": "", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 2, "contribution": 3, "strengths": "I appreciated the goal of the paper and that it improves on the presentation of some of its competitors. The discussions and remarks were enlightening and I chuckled when reaching the Appendix (A, B) to see the anticipation of some of my questions or critiques. The article clearly involved some work and the literature review appears quite adequate. The presence of the experiments nicely complements the theoretical part, justifying in particular the interest of introducing $\\kappa$ in the analysis.", "weaknesses": "The key proof element is Assumption 1. I wished Section 2.1 was written differently so as to improve its clarity, emphasizing the assumptions that really matter based on my understanding, that is 1) $\\Gamma \\Gamma^*$ is Hilbert-Schmidt 2) eq. (8) + $f^* \\in R(\\Gamma^*)$ rather than the more obscure (1) or Ass.1. The writing of Proposition 1 can be definitely improved.\n\nI think the comparison with the source condition of Cabannes et al. could be more developped since the authors' assumption ($f^* \\in R(\\Gamma^*)$) looks quite similar. I looked for the formula corresponding to (1) in HaoChen et al. but could not find it. Also the final predictor of Definition 1 used a bound constraint on the norm, I am not sure whether it is common in some literature since I am used rather to regularized versions. These formulations require better justification I would say. I also have doubts with the comparison with the RIP.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "7TGEzRJfwi", "reviewer": "Reviewer_kPdq", "summary": "The authors consider augmentation-based representation learning, where a feature map is learned from an unlabelled dataset augmented using some mechanism. \nThe paper provides a learning theory for this setup using a regression problem, where the regressor is linear in the learned features. \n\nThe authors consider an augmentation mechanism defined by a conditional distribution $p(a|x)$, which modifies a data point $x$ to an augmented sample $a$ stochastically. \nGiven this mechanism, the regression function is defined by the conditional expectation $\\mathbb{E}[g(A)|x]$ of some function $g$. \nFor this, the authors show that the problem may be described by an RKHS, and using this framework, they investigate the role of the feature map and augmentation.", "questions": "* (9) requires a suitable regularisation that depends on $B$ and $\\varepsilon$. In some sense, the analysis addresses a well-specified case. What happens if this is violated? \n* (12): Why is $G^{-1/2}$ required?\n* \"As shown in Wang et al. (2022b), there is...\": Is this an existence statement?", "limitations": "", "rating": 5, "confidence": 2, "soundness": 2, "presentation": 2, "contribution": 2, "strengths": "The theoretical framework introduced in the paper is interesting and original. \nThe authors introduce an RKHS structure to the representation learning problem. Specifically, the target function $\\mathbb{E}_A[g(A)|x]$ may be considered as an element of an RKHS whose kernel is given by the augmentation-induced joint distribution . \nThe framework thus enables us to quantitively analyse the performance of a downstream regressor using a kernel machinery, as well as establishes a criterion of good representation (approximation to the basis of the RKHS).", "weaknesses": "### Clarity \nA major issue is the lack of clarify. The paper is currently dense, not accessible to a general audience (as someone who is not well-versed in the area, I find the paper really challenging to read), and thus requires a major revision. Because of this, I have to admit that I have not fully understood the contribution of the paper. Specific comments are summarised below. \n \n### Significance \nThere are some elements that make the evaluation of significance difficult: \n* It is not clear how reasonable the assumption on the target function is. It seems that the assumption is conveniently introduced so that we can analyse the problem using a particular RKHS. Is there any motivating example for using $\\mathbb{E}_A[g^*(A)|x]$ as a target? It should be more straightforward to work with (functions on ) the augmentation space $\\mathcal{A}$ rather than $\\mathcal{X}$, as in Johnson et al., 2023.  \n* The encoder has to be of the form (10), as otherwise $\\hat{\\Psi}$ are not included in the RKHS. Therefore, the claim that the encoder $\\hat{\\Psi}$ can be arbitrary is not accurate to me. \n* The soft-invariance assumption is not straightforward to interpret. How does it measure \"invariance\"? Since $\\mathbb{E}[g(A)|x]$ and $g(a)$ are functions on different sets, the invariance is not well-defined.\n* The dimensionality dependency on the upper bound seems to be very pessimistic, as the authors admitted. I \n\nComments: \n* The problem is not well presented. It would have been more helpful if the authors had provided a high-level problem description and define symbols accordingly (e.g., as in the introduction of Johnson et al., 2023). For example, how representation learning is performed is not really mentioned using symbols in the paper (e.g., $\\hat{\\Phi}$ is not mentioned up to page 4). \n* Some symbols are loosely defined or mentioned before formally defined; e.g., $\\hat{\\mathcal{H}}_d$ is mentioned before the formal definition of \\hat{H}_{\\Gamma}. \n* Presentation of technical results could be improved. For instance, Proof Sketch of Theorem 1 only decomposes Theorem 1 into two parts, repeating what is mentioned in the remark, and thus could be omitted. \n* For Proposition 1, the authors might want to check Chapter 11 of \"An Introduction to the Theory of Reproducing Kernel Hilbert Spaces\" by Paulsen and Raghupathi.  \n* Using the same symbol for different density functions is confusing ($p(a)$ and $p(x)$). \n* Section 5. The term \"strong\" augmentation is obscure. One may consider an augmentation distribution that is independent of $x$, in which case the augmentation complexity is 1. Is the dependency between $A$ and $X$ considered a measure of \"strong\"?", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "JvbPfnS7iP", "reviewer": "Reviewer_dUJ6", "summary": "This paper gives a theoretical analysis of how the distribution of augmentations affects the accuracy of contrastive learning methods, building on previous work that has studied this problem through the lens of RKHS regression, and focusing on nonparametric encoder-complexity-independent bounds. The authors first discuss the augmentation operator and show that it induces a pair of dual operators with corresponding eigenfunction decompositions. They then introduce the concepts of *augmentation complexity* (based on the infinity norm of a kernel defined by the augmentations) and the *ratio trace* (which roughly characterizes \"how much\" of the eigenspaces of the corresponding RKHS are captured by a given encoder).\n\nThe authors then prove a number of results based on these quantities. First, they show that the error of a downstream predictor can be decomposed into an approximation-error term (how closely the function can be recovered from the learned encoder's representation) and a downstream-estimation-error term (how well the linear readout actually recovers this function). Second, they consider approximating the top-d eigenspace of the RKHS using only a finite pretraining dataset, and show that this can be used to bound the approximation error term.\n\nFinally, the authors consider how their \"augmentation complexity\" measurement translates to real augmentations, focusing on masking transformations. They give theoretical values for it for the hypercube examples discussed in Cabannes et al. (2023), and also show how to estimate it for masked language modeling on the `wikipedia-simple` dataset.", "questions": "**On inductive biases.** In the appendix, you state \"the nonparametric analysis established in this work clearly refutes the claim in Saunshi et al. (2022) that learning guarantees independent of the model inductive bias are necessarily vacuous\". I don't think this is true. As I understand it, Saunshi et al. show that when augmentations are approximately disjoint, the eigenspectrum of $\\Gamma\\Gamma^*$ decays extremely slowly, with many of the eigenvalues of $\\Gamma\\Gamma^*$ being close to 1 (or, under their notation, many eigenvalues of the normalized Laplacian $L_\\circ = I - A_\\circ$ being close to 0). This means that the approximation error of the top-d eigenspace is necessarily very large. I think this same analysis also applies to your results; Saunshi et al.'s approximately-disjoint augmentations would imply $\\lambda_{d+1} \\approx 1$ and so your $\\tau \\approx 1$, leading to a blowup in the first term of your equation 14.\n\nIf I'm correct, I think this deserves more discussion, and I don't think you can say that you've refuted the claims of Saunshi et al. I also think the paper could be improved by giving more context about the relationship of your results to those of Saunshi et al. For instance, is there any relationship between your augmentation complexity and Saunshi et al.'s approximately-disjoint augmentations? One might expect that approximately-disjoint augmentations would imply a high augmentation complexity; if so, that could be useful to discuss.\n\nAlso, what do you mean by \"disentangle the effects of the model and the augmentation\"? It's not clear to me how the presented results disentangle these effects, although I do think they give a clear explanation of the role of the augmentations in particular.\n\n\n**Interpreting the soft invariance criterion.** Can you give any intuition about what satisfying Assumption 1 actually means in practice? It seems unlikely to me that functions of practical interest could be *exactly* represented as expectations of some function $g$ on augmentations, and I also don't see any way to verify whether or not Assumption 1 holds.\n\n**Part (iv) of Proposition 1.** Proposition 1 part (iv) states that $g^*$ satisfies Assumption 1. This doesn't make sense to me; Assumption 1 is a statement about $f^*$, not $g^*$. Do you mean to say that $g^*$ satisfies Equation (1)?\n\nThe proof of part (iv) also seems potentially incorrect. In particular I'm not sure how $\\langle g^*, (I - \\Gamma \\Gamma^*) g^* \\rangle_{P_A} \\le \\epsilon ||g^*||^2_{P_A}$ implies $\\langle g_0, (I - \\Gamma \\Gamma^*) g_0 \\rangle_{P_A} + ||g_1||^2_{P_A}  \\le \\epsilon (||g_0||^2_{P_A} + ||g_1||^2_{P_A})$; shouldn't $||g_1||^2_{P_A}$ be $\\langle g_1, (I - \\Gamma \\Gamma^*) g_1 \\rangle_{P_A}$ instead? It also seems like your proof here implies that $R(\\Gamma^*) \\subseteq \\mathcal{F}_B(\\Gamma; \\epsilon)$ for all $\\epsilon$. But that doesn't seem like it could be true; perhaps I'm misunderstanding something.\n\n**Asymptotic behavior of Theorem 2.** In Theorem 2, don't the quantities $\\overline{\\lambda}_d^{-1}$ and $\\gamma_G$ potentially depend on $N$? It seems like that could affect the asymptotic behavior of the bound in equation (20). Can you say anything about the convergence or boundedness of these quantities?\n\n**Relationship of this analysis to that of Johnson et al. (2023).** I noticed that parts of your analysis resemble some of the statements of Johnson et al. (2023) at a high level; in particular, your soft invariance assumption is similar to their \"approximate view invariance\" assumption, and the decomposition in Theorem 1 is somewhat similar to their bound in Proposition 4.2. The details are different (their analysis focuses on approximating functions $\\mathcal{A} \\to \\mathbb{R}$ rather than $\\mathcal{X} \\to \\mathbb{R}$ and doesn't include the approximation error of the encoder) but it might still be worth including a brief discussion of the similarities.\n\n**Other suggestions / comments.**\n- Section 2 states \"Similar to RIP, this property clarifies the the role of augmentation: The augmentation defines a set of features (eigenfunctions) ordered by their a priori relevance to the target function.\" This wasn't clear to me; how does the isometry property relate the relevance of the eigenfunctions to the target function?\n- The definition of $||f||\\_n^2$ in section 3 doesn't seem to be used anywhere in the main paper.\n- It took me a while to locate the definition of $\\hat{f}$ when reading Theorem 1; perhaps you could add a reference back to equation (9) where it is defined.\n- Section 4 discusses the empirical eigenspace $\\hat{\\mathcal{H}}\\_d$ before actually defining what it is, which was a bit disorienting; it might make sense to define it at the start of section 4 instead of at the end of page 6.", "limitations": "", "rating": 8, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "**[S1]** The paper is overall clearly written and easy to follow, and the authors do a good job introducing the problem and explaining their analysis. I found the \"proof sketches\" for the main results to be very useful as well.\n\n**[S2]** The formalism of the dual operators and the isometry property are quite elegant. The authors build on previous work but present the previously-studied concepts in an insightful way, and clearly lay out the underlying structure produced by the augmentation distribution.\n\n**[S3]** The main results in sections 3 and 4 go further than previous work in characterizing how much information the augmentations alone can provide about the function of interest. In particular, Theorem 1 shows how the error of a function decomposes into approximation and estimation error, and furthermore identifies how the approximation error component depends on the \"trace gap\"; Theorem 2 then shows that finding the empirical top eigenspace can lead to small approximation error even if $\\lambda_{d+1}$ and $\\lambda_d$ are close.\n\n**[S4]** The analysis of the augmentation complexity of real augmentations is quite interesting, both for the theoretical values for particular masking strategies and for the empirical estimation of the augmentation complexity of real augmentations used for masked language modeling.", "weaknesses": "**[W1]** This work focuses on analyzing the effect of the augmentations alone, without considering the impact of inductive biases. I understand that this was an intentional choice, and I do think that this makes the analysis of augmentations more \"pure\". However, such an approach means that the theoretical results here may not have much explanatory power for real augmentations.\n\n~~Relatedly, I think the paper makes some unjustified claims regarding the power of their analysis for tasks where inductive biases are important.~~\n\n- ~~(a) The authors state that they \"disentangle the effects of the model and the augmentation\", but it's not clear to me what is being disentangled. This seems to be primarily an analysis of the augmentation distribution, with the model's role being just to achieve a good ratio trace with respect to those annotations. I'm not sure how easy it would be to extend this analysis to models with more structured inductive biases.~~\n- ~~(b) In the appendix, the authors state that their analysis \"clearly refutes the claim in Saunshi et al. (2022) that learning guarantees independent of the model inductive bias are necessarily vacuous\", but I disagree with this claim and don't think it is justified (see questions below).~~\n\n*Edit: The authors have addressed the unjustified claims. Also, I think the theoretical contributions of studying augmentations alone are valuable even if they don't fully explain the success of large-scale self-supervised learning, so this isn't a major concern.*\n\n**[W2]** The analysis starts from a strong assumption that the target function $f^*$ lies in a particular RKHS and satisfies a particular soft invariance condition. However, it's difficult to tell what this assumption actually means, or whether it would be true in practice. In particular, the RKHS membership seems quite strict, and adding even an arbitrarily small perturbation to $f^*$ might cause the assumption to no longer hold. (There is a bit of discussion in the appendix, but it's not clear to me what effect this would actually have on the main results.)\n\n*Edit: The authors have noted that assuming functions lie in a particular RKHS is standard for learning theory work and necessary to get statistical guarantees.*\n\n~~**[W3]** I think part of Proposition 1 is incorrect, or at least incorrectly stated. I also have a small concern about the asymptotic statements concerning Theorem 2. (These should hopefully be minor issues that are easy to address.)~~\n\n*Edit: Addressed in the revision.*\n\n**[W4]** The empirical results about real augmentations suggest that the true augmentation complexity of real-world datasets is *massive* (on the order of $\\kappa \\approx 10^{100}$). As such, it seems likely that the bounds in Theorem 1 and Theorem 2 would be effectively vacuous for real-world data.\n\n*Edit: As the authors point out, this is analogous to the standard curse of dimensionality for bounds based on model complexity, so it's not specific to their analysis.*", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (spotlight)