Source-linked AI summary
An Analysis of Self-supervised Pre-training with Dependent Samples
Maximilian Fleissner, Debarghya Ghoshdastidar, Samory Kpotufe
TL;DR
Self-supervised learning often pools multiple dependent augmentations to estimate an invariant subspace, while existing theory commonly avoids those dependencies. This paper compares pooling with a partitioning baseline and finds that pooling is never statistically worse and can improve rates as the number of augmentations grows, subject to scope limitations.
Problem
Existing theory largely avoids dependent views, leaving unclear why pooling many augmentations works well for estimating invariant representations.
Method
The paper models the representation as a subspace V of a feature space and contrasts naive pooling of all L augmentations with partitioning that uses one augmentation per sample.
Results
Pooling achieves a subspace estimation rate no worse than the one-augmentation-per-sample baseline, and cross-correlations can produce strictly faster rates for common augmentations.
Takeaways & Limitations
Using more augmentations can improve estimation of V in supported settings, helping explain the practical success of pooling augmented samples in self-supervised pre-training.
Takeaways & Limitations
The analysis does not apply to next-token prediction in language models and assumes, in its main presentation, a shared pre-training and marginal input distribution with whitened features.
Abstract
from arXiv · showhide
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 $\cal V$ for downstream tasks. In practice, such augmentations $\{ φ_l(x_i) \}$ are pooled together to learn $\cal 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 $\cal 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.
1 Introduction
The paper studies why pooling dependent augmentations can work well for self-supervised representation learning, despite theory often avoiding such dependencies. It models the learned representation as a subspace and shows pooling can match or outperform dependency-avoiding partitioning.
- Self-supervised learning seeks low-dimensional representations of unlabeled data for downstream prediction with fewer labeled examples.
- Existing theory often treats the number of views L as small or constant, avoiding interactions among dependent augmentations of the same datapoint.
- The paper represents the desired embedding as a low-dimensional subspace V of a feature space G induced by a feature map g.
- The analysis compares naive pooling of L dependent views per datapoint with a baseline that partitions datapoints so each receives one view.
- Pooling estimates the population V at a rate never worse than the partitioning baseline, whose bound has the usual √m rate.
- For common augmentations including masking, salt-and-pepper, cropping, and solarization, increasing L can yield strictly faster estimation rates than the baseline.The improvement depends on cross-correlation terms and does not imply taking L arbitrarily large, because predictive features may be lost when L greatly exceeds feature dimension d.
2 Setup and Notation
The paper formalizes V as a subspace orthogonal to feature changes induced by augmentations and estimates it from bounded augmentation differences. Its setup uses fixed-dimensional representations and shallow neural-network feature maps, with invariance tied to downstream irrelevance.
- The pre-training sample consists of i.i.d. datapoints and transformations such as pixel masking or noise added to speech components.
- Augmentations are intended to preserve predictive features while exposing aspects of the input that are irrelevant to downstream tasks.
- The target V is a k-dimensional subspace of the feature space G that makes augmentation-induced feature differences nearly orthogonal to V.
- The covariance of augmentation differences identifies V through its bottom eigenspaces, where v^TΣv measures expected projected augmentation discrepancy.
- The analysis assumes k is chosen in advance and assumes bounded augmentation-difference vectors, with bounded inputs and Lipschitz augmentations providing one sufficient example.
- When augmentation differences lie in downstream-irrelevant directions, subspaces containing Null(Σ) are invariant or nearly invariant for linear and suitable nonlinear prediction heads.
- The model uses shallow networks f(x)=w^Tσ(Wx), taking g(x)=x or g(x)=σ(Wx) as the feature map.
3 Overview of Results
The paper compares naive pooling of dependent augmentations with a data-splitting baseline for estimating a low-dimensional subspace V. Pooling is never worse in general and can achieve faster rates for several practical augmentations as the number of views increases.
- Estimator comparison: Naive pooling reuses every augmentation for every sample, whereas the baseline partitions samples so each receives one augmentation.The two estimators are compared through their induced bottom eigenspaces and sin Θ subspace distance.
- Estimator comparison: The pooling estimation rate is never worse than the baseline rate under the stated eigengap and concentration assumptions.The comparison follows from contrasting variance terms for the two estimators.
- Common augmentations: The analysis covers masking, salt-and-pepper noise, cropping, and solarization, while noting that very large L may produce a learned subspace that no longer preserves predictive features.The practical scope therefore includes both rate improvements and a constraint on increasing the number of augmentations indefinitely.
- Common augmentations: For salt-and-pepper noise with a shallow neural-network feature map, pooling achieves faster rates than baseline when the dimension is sufficiently large.The bound compares pooled and baseline variance up to logarithmic factors and under distributional and initialization conditions.
- Common augmentations: Random masking yields strictly smaller pooling variance than baseline variance for any L > 1, so dependence helps in this setting.The result is stated for g(x) = x under the masking model.
- Common augmentations: Solarization improves estimation as L increases: pooling and baseline variance terms coincide, while the subspace distance decreases with both m and L.The population covariance is independent of L in the stated uniform-input setting.
4 Detailed Analysis
The detailed analysis derives augmentation-specific estimation rates and shows when naive pooling benefits from increasing the number of augmentations. Across masking, salt-and-pepper noise, cropping, and solarization, the results identify conditions under which dependence is benign or improves estimation.
- Deterministic and block masking: For block masking with L = Θ(d) and bounded overlap parameters, the baseline is worse than naive pooling by a factor of order Θ(L).The comparison assumes r, C, and s are O(1), while the baseline scale satisfies c_u = d.
- Random masking: For block masking, the variance statistic is small enough to cancel the eigengap when each coordinate appears only O(1) times and pairwise block overlaps are bounded.The baseline lacks this cancellation under the same setting.
- Random masking: Random masking with L = Θ(d) and ∥Γ∥op = Θ(1) has variance decaying as O(1/d), cancelling the small eigengap.This avoids the curse of dimensionality induced by that eigengap despite dependent views.
- Salt and pepper: Salt-and-pepper augmentations admit faster naive-pooling rates when the representation is the first hidden layer of a shallow neural network.The analysis randomizes the network's first-layer weights and assumes a twice-differentiable activation with bounded first and second derivatives.
- Cropping: For cropping, the pooled estimator has strictly smaller variance than the baseline whenever L > 1.The result states ∥∆∥op < ∥∆♭∥op for any L > 1.
- Solarization: Solarization has an eigengap independent of L, yet its subspace estimation rates improve with L.The stated setting is PX = Unif([0, 1]d) with g(x) = x.
5 Conclusion
The paper concludes that naive pooling all augmentations across samples is never worse than using one augmentation per sample. Cross-correlations can be benign enough to produce strictly faster rates, and estimation can improve with both sample size and augmentation count.
- Conclusion: Naive pooling achieves a subspace estimation rate no worse than the data-splitting baseline using one augmentation per sample.The conclusion concerns estimating the low-dimensional subspace V from unlabeled pre-training data.
- Conclusion: Cross-correlations between augmented samples can yield strictly faster rates for commonly used augmentations.The conclusion attributes the improvement to sufficiently benign cross-correlations.
- Conclusion: Estimation of V can improve with both the number of original samples m and the number of augmentations L.The paper identifies this as a further consequence of its analysis.
- Generalized setting: The downstream analysis relates subspace estimation to excess risk under C-induced geometry and approximately augmentation-invariant linear predictors.It uses C-orthogonal projections and a Lipschitz or linear downstream-function analysis.
- Generalized setting: The generalized framework defines V using the bottom generalized eigenvectors of C−1Σ when the downstream covariance is a positive definite matrix C.This reduces to the usual bottom eigenspace of Σ when C = Id.
D Proofs for Proposition 1 and Theorem 1 in the General Case
The general-case proofs whiten the covariance geometry, establish concentration for the pooled and split covariance estimators, and then transfer these bounds to subspace error. The resulting comparison proves that pooling is never worse and can improve by a factor of L when correlations are sufficiently mild.
- Proof strategy: The proof bounds the whitened covariance error with matrix Bernstein and converts it into projection error using Wedin’s bound in the C-inner-product space.The relevant operators are C−1bΣ and C−1Σ, which are self-adjoint under the C-induced inner product.
- Pooled concentration: The pooled estimator groups dependent augmentations within each original sample, producing independent zero-mean sums across the m original samples.This preserves sample-level independence while retaining all L augmentations per sample.
- Split concentration: The split estimator instead uses one augmentation per sample, yielding independent summands whose concentration is controlled separately.Its variance bound is expressed through the corresponding split-estimator statistic ∆♭.
- Pooling versus splitting: Naive pooling is never worse than data splitting because R ≤ cu and ∆ ⪯ L · ∆♭.These inequalities are the key comparison between the pooled and split concentration parameters.
- Pooling versus splitting: When ∆ ≲ ∆♭, pooling improves asymptotically over splitting by a factor of L.This is the stronger regime identified after the general comparison.
E Proof of Theorem 2
The proof analyzes block masking through coordinate-wise masking frequency and pairwise block overlap. These quantities determine the pooled and split variance terms and explain the resulting advantage of pooling.
- Alignment term: The alignment bound satisfies R ≤ r/d for the block-masking construction.The bound follows from the maximum coordinate masking frequency r.
- Correlation term: The correlation and variance calculations remain diagonal because the masking matrices are diagonal and their products preserve diagonality.The operator norm is therefore controlled by the largest diagonal entry.
F Proof of Lemma 2
The proof derives a fourth-moment identity for uniformly sampled directions on the sphere and applies it to the masking covariance structure with symmetric Γ.
- Application to Γ: Substituting M = Γ and using Γ's symmetry yields identity (14): (Trace(Γ)Γ + 2Γ^2)/(d(d + 2)).The final step follows directly from the general fourth-moment formula.
- Spherical fourth moments: For a uniform spherical vector U, E[UU^⊤MUU^⊤] = (Trace(M)I_d + M + M^⊤)/(d(d + 2)).The proof obtains this identity through a Gaussian representation of U and fourth-moment calculations.
G Proof of Theorem 6
The proof extends Theorem 6 to general covariance C, represents the relevant matrix through sampled spherical directions, and controls it with high-probability concentration bounds.
- General covariance: The general setting uses C = E[XX^⊤] and learns the bottom generalized eigenspace of C^-1Σ, with view differences whitened accordingly.This extends the analysis beyond the identity-covariance case.
- Matrix representation: The matrix Δ admits a product representation involving W, whose columns are the sampled directions, and an entrywise-defined matrix H.This representation is the starting point for operator-norm control.
- Concentration: Lemma 11 provides a high-probability operator-norm bound for a matrix with independent uniform-sphere columns.The spherical samples are analyzed through Gaussian normalization and concentration.
- Bounding H: The covariance matrix H is positive semidefinite because it is a Hadamard product of positive semidefinite matrices, enabling an operator-norm bound.The proof combines this structural property with Cauchy–Schwarz bounds on covariance terms.
H Proof of Theorem 7
The proof analyzes nonlinear representations by separating main and remainder terms, then shows that diagonal augmentation effects dominate cross terms with high probability.
- Setup: The nonlinear setting studies a single-layer network g(X) = σ(W^⊤X) with bounded first and second derivatives and identity downstream second moment C = I_d.The analysis targets the bottom-k generalized eigenspaces of the resulting covariance structure.
- Main difficulty: Nonlinearity makes bounding Δ and Δ^♭ substantially harder than in the linear model, where augmentation differences are data-independent.The proof therefore introduces a main-term/remainder decomposition.
- Main terms: For distinct coordinates l and k, the diagonal main term is expected to dominate the cross term because the l = k contribution squares a weight coordinate.Proposition 5 formalizes this comparison under assumptions on the coefficients and random Gaussian weights.
- High-probability controls: The proof controls random-weight norms and weighted sums using Gaussian concentration, subexponential Bernstein bounds, and union bounds over coordinates.These events provide the high-probability controls needed for the main-term comparison.
- Remainders: When d1 > (log d)^2, the remainder terms are asymptotically smaller than the main terms and hence negligible for the final operator-norm comparison.The remainder bound depends on the activation, augmentation strength, and an almost-sure bound on ||X||∞.
I Proof of Theorem 8
The proof computes the covariance contributions of event-based augmentations and shows that cross-covariance terms vanish, making the pooled quantity smaller than its split counterpart for L > 1.
- Covariance calculation: For event-based augmentations, the view difference is expressed through the indicator of an event E_l(X) multiplied by the shared representation φ(X).Independence between the event indicator and φ(X) is used in the covariance calculation.
- Diagonal and cross terms: Cross terms depend on the joint occurrence of E_l(X) and E_k(X), while diagonal terms use the single-event probability.The proof separates these cases when evaluating the covariance matrix.
- Final comparison: For any L > 1, the resulting operator-norm comparison is ||Δ||op < ||Δ^♭||op.The inequality follows from the vanishing cross terms in the event-based setting.
J Proof of Theorem 9
The proof develops concentration and perturbation ingredients for bounding random positive semidefinite matrix averages. It handles both positive-definite and rank-deficient cases using matrix concentration and Loewner-order arguments.
- Concentration tools: The proof invokes matrix Bernstein concentration for independent Hermitian random matrices.The stated lemma assumes zero mean, bounded operator norm, and controls the matrix sum.
- Auxiliary results: Additional ingredients include subgaussian quadratic-form concentration and eigenvector perturbation results for compact operators.The perturbation result applies to projections onto leading singular-vector subspaces and also covers self-adjoint operators.
- Positive semidefinite averages: A matrix Chernoff argument bounds the average of independent positive semidefinite matrices around its expectation.The proof normalizes by A^-1/2, obtains identity expectation, and applies a union bound.
- Positive semidefinite averages: With probability at least 1 − δ, the normalized average satisfies 0.5I_d ⪯ B ⪯ 2I_d under the stated conditions.The Loewner-order inequality is transferred back by multiplying with A^1/2 on both sides.
- Rank-deficient case: The proof extends the concentration argument to singular A using the Moore–Penrose pseudoinverse and a reduced-rank representation.The resulting construction uses the nonzero eigenspace and expectation I_r.
L Subspace Perturbation Results
This section states a general perturbation result for leading singular-vector subspaces and notes its specialization to self-adjoint operators. It also introduces a synthetic random-masking experiment comparing naive pooling with a baseline.
- Subspace perturbation: The perturbation lemma compares projections onto the top-s left singular-vector subspaces of two compact operators.The bound is stated in operator norm on a separable Hilbert space.
- Subspace perturbation: For self-adjoint operators, the same perturbation inequality applies to eigenvalues and eigenvectors.This provides the corresponding symmetric-operator interpretation of the lemma.
- Synthetic experiment: The experiment reports results for naive pooling and the baseline across L values of 10, 15, and 20.The displayed estimator table lists both estimators alongside these augmentation counts.
- Synthetic experiment: Figure 5 applies random masking to a synthetic Gaussian dataset and compares subspace distances for naive pooling and a baseline.The figure is labeled as showing faster rates for naive pooling on synthetic data.