Joint‑Embedding vs Reconstruction: Provable Benefits of Latent Space Prediction for Self‑Supervised Learning

We would like to thank the reviewer for the insightful feedback and comments.

Can the analysis be extended to non-Gaussian augmentations such as random crops, color jitter or CutMix? In particular, how might you handle the patch-masking used in MAE?

We appreciate the reviewer’s question regarding extending our Gaussian covariance-based perturbation analysis, giving us the opportunity to clarify an important point.

Let us consider the first-order Taylor expansion around the mean augmentation. Defining $\bar{\tau}(\mathbf{x}) := \mathbb{E}_{\tau \sim \mathcal{T}}[\tau(\mathbf{x})]$, for a given function $f$, we approximate:

$$f(\tau(\mathbf{x})) \approx f(\bar{\tau}(\mathbf{x})) + Df(\bar{\tau}(\mathbf{x}))[\delta\tau(\mathbf{x})], \quad\text{with}\quad \delta\tau(\mathbf{x}) = \tau(\mathbf{x}) - \bar{\tau}(\mathbf{x}).$$

where $D$ is the Jacobian operator. From this linear approximation, the induced covariance of $f(\tau(\mathbf{x}))$ is given by:

$$\mathrm{Cov}(f(\tau(\mathbf{x}))) \approx Df(\bar{\tau}(\mathbf{x})) \mathbf{\Sigma} Df(\bar{\tau}(\mathbf{x}))^\top, \quad\text{where}\quad \mathbf{\Sigma} = \mathbb{E}[\delta\tau(\mathbf{x})\delta\tau(\mathbf{x})^\top].$$

as defined in equation (Cov) in the paper. When $f$ is linear, as considered in this work for analytical tractability (consistent with the SSL literature focused on tractability [a][b]), this approximation becomes exact as no higher-order terms remain. In that case, we recover the same mean and covariance for $f(\tau(\mathbf{x}))$ as if we used the additive Gaussian noise model with covariance $\mathbf{\Sigma}$ to define $\tau$.

Thus, modeling the augmentation distribution as Gaussian provides a natural second-order approximation, matching exactly the first two statistical moments. Notably, several practically relevant augmentations admit closed-form covariance expressions, as explicitly derived (see page 11 in [c]). These include Pepper noise injection, Random Cutout, and unbiased random masking. This shows that masking, in particular, can be captured through a tractable second-moment approximation, directly addressing the reviewer's question about its compatibility with our framework. Concretely, to incorporate such augmentations, one can simply substitute the corresponding closed-form covariance operator into our formalism (as Gaussian covariance), effectively mimicking the effect of the augmentation through its second-order moment.

While finite-moment approximations of data augmentations and noise are established in supervised learning contexts [d], they remain under-explored in the SSL literature. Current analytical results in SSL [e] rarely exploit the structural properties of the data augmentation distribution explicitly, despite its central role. We believe leveraging finite-moment approximations is a promising direction for deriving meaningful theoretical insights for SSL methods.

We will insist on these aspects in the revised version to clarify our approach.

[a] Cabannes, Vivien, et al. "The ssl interplay: Augmentations, inductive bias, and generalization." ICML 2023

[b] Etai Littwin, et al. "How jepa avoids noisy features: The implicit bias of deep linear self distillation networks." NeurIPS 2024

[c] Lin, Chi-Heng, et al. "The good, the bad and the ugly sides of data augmentation: An implicit spectral regularization perspective." JMLR 25.91 (2024): 1-85.

[d] Balestriero, R., I. Misra, and Y. LeCun. "A Data-Augmentation Is Worth A Thousand Samples: Exact Quantification From Analytical Augmented Sample Moments. NeurIPS 2022

[e] Balestriero, Randall, and Yann LeCun. "Contrastive and non-contrastive self-supervised learning recover global and local spectral embedding methods." NeurIPS 2022

How could one estimate alignment or noise magnitude from raw, unlabelled data in practice?

We thank the reviewer for this important question related to the point raised above.

Our work provides precise analytical quantities in settings where second-order moment approximations hold and where noise statistics can be estimated from data (bounds of proposition 4.3 and 4.4). While the practical estimation of these quantities in real-world systems is indeed valuable, it falls beyond the scope of this study. Our focus is on establishing a controlled, theoretically grounded comparison between joint embedding and reconstruction-based SSL methods.

Moreover, our experiments on real image datasets confirm the theoretical trend: reconstruction-based methods are significantly more sensitive to noise perturbations than joint embedding approaches, even beyond controlled or synthetic settings. This phenomenon has been observed in several prior studies, and our work provides a theoretical explanation for it. We believe this connection is valuable for the SSL community.

Do you have intuition or experiments for settings where signal and nuisance features overlap in the same subspace?

Yes, we do cover such cases in our experiments. In Section 5, we evaluate models on corrupted images from ImageNet-C, including transformations such as Pixelate, Gaussian noise, and Zoom Blur. These corruptions affect the same subspace as core signal features.

Despite this overlap, we observe the same consistent trend: joint embedding methods remain more robust than reconstruction-based methods as noise severity increases. This suggests that the robustness of joint embedding extends even to more challenging settings where signal and nuisance features are entangled.

Limitations

We thank the reviewer for pointing out that the limitations of our work should be more clearly highlighted. We will address this in the revised version of the paper by more explicitly discussing the key assumptions and scope.

We appreciate the thoughtful discussion and remain at your disposal to elaborate on any part of our submission.

We would like to thank the reviewer for the constructive and insightful feedback.

1. Are linear models used in section 3 a good proxy for the complex non-linear models that are used in practice? How do we ensure that the findings transfer to these more complex non-linear models?

Thanks a lot for this important question. We appreciate the opportunity to clarify why we chose a linear‑parameter framework for our analysis.

Modeling our system as linear in its parameters yields closed‑form solutions that can be derived and analyzed exactly. This approach is standard in theoretical studies of self‑supervised learning (SSL): for instance, prior work on the interplay of augmentations and inductive bias [a], on implicit bias in self‑distillation networks [b], and on SSL dynamics without contrastive pairs [c] all adopt linear‑parameter models to perform precise derivations.

Importantly, “linear” in our work refers only to dependence on the parameters, not to the representation power of the model. We still allow arbitrary nonlinear feature maps, akin to kernel regression. Moreover, it is common to linearize around initialization via a first‑order Taylor expansion to study deep neural networks. In the infinite‑width limit, the Neural Tangent Kernel (NTK) framework [d] shows that gradient descent training behaves exactly like kernel regression with a fixed kernel, implying that the dynamics are linear in the parameters. Thus, our linear‑parameter analysis can be understood as capturing the same training regime as very wide deep networks.

In supervised learning, recent work has derived closed‑form regularization effects of data augmentation [e]. Our contribution is the analogous analysis for SSL, where augmentations play an even more central role.

Finally, as emphasized in [f], many fundamental aspects of (implicit or explicit) regularization can be understood in the linear‑model formalism. We believe that the fact that we can explain the key distinction between reconstruction‑based and joint‑embedding objectives with such a simple standard formalism highlights the generality and clarity of our findings.

[a] Cabannes, Vivien, et al. "The ssl interplay: Augmentations, inductive bias, and generalization." ICML 2023

[b] Etai Littwin, et al. "How jepa avoids noisy features: The implicit bias of linear self distillation networks." NeurIPS 2024

[c] Tian, Yuandong, Xinlei Chen, and Surya Ganguli. "Understanding self-supervised learning dynamics without contrastive pairs." ICML 2021

[d] Jacot, Arthur, Franck Gabriel, and Clément Hongler. "Neural tangent kernel: Convergence and generalization in neural networks." NeurIPS 2018.

[e] Balestriero, Randall, Ishan Misra, and Yann LeCun. "A data-augmentation is worth a thousand samples: Analytical moments and sampling-free training." NeurIPS 2022.

[f] Bach, Francis. Learning theory from first principles. MIT press, 2024.

2. I understand how increasing alpha in Eq. 7 results in a more aligned augmentation with the noise added. However, how do the authors control for “unaligned” noise? It seems to me that the extremes of alpha either result in a dominant aligned augmentation with high alpha or just the theta term with alpha as 0. Shouldn’t we also control for the magnitude of theta?

Thanks for the question. The covariance $\mathbf{\Theta}$ can be any positive semidefinite matrix, potentially arbitrarily far from the true noise covariance $\mathbf{\Gamma}$. Our theoretical results hold in full generality for any choice of $\mathbf{\Theta}$. Concretely, if we choose $\mathbf{\Theta}$ to assign large eigenvalues to irrelevant noise components and small eigenvalues to informative data components, then we reproduce the case of a large $\alpha$ where we match the data augmentation to the noise. However this is not the scenario we are interested in. This is why in the experiments we consider no correlation between the default augmentation structure ($\mathbf{\Theta}$) and the actual noise ($\mathbf{\Gamma}$).

In our simulated experiments, we independently and uniformly sample the eigenvalue spectra of both $\mathbf{\Theta}$ and $\mathbf{\Gamma}$. This ensures there is no correlation between augmentation and noise covariances. Setting $\alpha = 0$ recovers the case where there is no correlation between actual noisy features and augmentation (because the augmentation follows the structure given by $\mathbf{\Theta}$), while larger $\alpha$ interpolates towards aligning the augmentation with the noise (discarding the effect of $\mathbf{\Theta}$ as $\alpha$ increases).

This protocol also mirrors our real‑data image experiments, where we inject corruption noise (from Imagenet-C) not present in the standard augmentation pipeline. Here, $\mathbf{\Theta}$ represents the default augmentation pipeline, unaware of the injected noise, thus creating a natural mismatch.

3. Propositions 4.2-4: How do we conclude that the limit holds (almost surely) as either alpha or n goes to positive infinity? It is unclear to me from the main text without any pointers to appendix.

We apologize for omitting explicit pointers to the proofs in the main text. In the revised version, we will add direct links before Propositions 4.2–4 to guide readers to Appendix A.3–A.5.

The proofs in Appendix A.3–A.5 proceed by first invoking the closed‑form solutions from Theorems 3.1 and 3.2. We then decompose each solution into two parts: the components that cover important data features and the other components that cover the irrelevant noise features. We then analyze each part in the limits of infinite sample size ($n\to\infty$) and perfect alignment ($\alpha\to\infty$). The infinite‑sample limit uses the law of large numbers to establish almost sure convergence of empirical covariances to their population counterparts, while the perfect‑alignment limit isolates the effect of the noise structure in the global augmentation covariance. We then compare the solutions obtained for clean and corrupted data settings.

We will also refine our notation and expand intermediate steps in the appendix to further improve readability. Please let us know if you would like any additional details.

4. Figure 2, bottom right (Strong noise, reconstruction). Do the authors have insight why the model with strong alignment lags behind at a low sample size despite high alignment between the noise and the augmentation?

This is indeed a consistent observation and an interesting point also raised by reviewer CuYw. We provide the same answer below for consistency.

We believe this discrepancy stems from finite‑sample effects. By the multivariate Central Limit Theorem, the variance of the empirical noise‑covariance spectrum decays like $1/\sqrt{n}$. In low‑sample regimes, this high variance causes the spectrum to deviate from the true noise covariance $\mathbf{\Gamma}$, thus worsening probing error. Because the excess variance scales with $\alpha^2$, its impact is greatest at large $\alpha$. As $n$ increases, these fluctuations vanish, and strong alignment consistently outperforms weaker settings, as predicted by Proposition 4.3. Importantly, we do not observe this behavior with joint embeddings: because both views undergo the same augmentation pipeline, they maintain identical variance levels for each $\alpha$, even in small‑sample scenarios.

We will incorporate this discussion in the revised version, thanks a lot for pointing this out.

5. For Table 1, I would highly suggest color-coding or adding a symbol to help readers separate joint embedding (BYOL, DINO) methods from reconstruction (MAE) methods at a glance.

Thank you for this excellent suggestion. We will do this in the revised version to make the results clearer for readers.

We’re grateful for the reviewers’ comments and are ready to provide additional clarification or expand on any point you wish.

We would like to thank the reviewer for the constructive review and very positive feedback.

1. One of the main assumptions of the theoretical analysis is that the labels are considered to be continuous in order to frame a least squares task. While this obviously allows for an in-depth theoretical exploration, how would it affect the performance of SSL models downstream on categorical matching tasks?

Indeed, probing is often done through discrete labels hence this question is very relevant for practical scenarios. We agree that treating labels as continuous values via MSE is primarily a theoretical convenience.

In practice, MSE can be applied to classification by representing categorical labels via one-hot encoding, where each class corresponds to a unique basis vector. This approach enables the use of regression-based objectives even when the downstream task is inherently categorical. Importantly, recent empirical work [a] has demonstrated that training classifiers with squared loss on one-hot targets can match, and in some cases even exceed, the performance of models trained with cross-entropy loss. These findings suggest that squared loss is not only theoretically convenient but also practically viable for classification.

Thanks for the question. We will clarify our contributions by specifying that we focus on MSE loss in our study of supervised learning.

[a] Hui, Like, and Mikhail Belkin. "Evaluation of neural architectures trained with square loss vs cross-entropy in classification tasks." ICLR 2021

2. In Equation 7, where the augmentation distribution is formalized, is my understanding correct that the directions that are unimportant to the task are augmented more, i.e., there is more noise added, as alpha grows, in the directions with $\lambda_i^{\mathbf{Γ}} = 0$. Is this correct?

Yes, your intuition is exactly right, the directions that are unimportant to the task (which we refer to as irrelevant noise features) are indeed augmented more as $\alpha$ increases. Concretely, these correspond to components with $\lambda_i^{\mathbf{Γ}} > 0$ (rather than $\lambda_i^{\mathbf{Γ}} = 0$), since $\mathbf{Γ}$ denotes the covariance of the noisy, task irrelevant signal components. By increasing $\alpha$, our augmentation process injects proportionally more noise into these irrelevant directions, thereby encouraging the model to become invariant to them.

Thank you for pointing this out, we will clarify this intuition in the revised manuscript.

3. As a minor element and more of a suggestion, I believe the proofs would benefit more people looking to understand theoretical research in SSL if some of the identities, e.g., regarding matrix traces and norms, were expanded more in the first few derivations. This greatly helps the readability of the more technical material.

Thank you very much for this helpful suggestion. We will expand the derivations to guide readers more smoothly through the proof. We’ll also include a comprehensive glossary of notation to further enhance clarity.

We would be happy to address any remaining questions if needed.

We would like to thank the reviewer for the insightful review and very positive feedbacks.

1. In Appendix Figure 3, the probing error for the reconstruction-based method under high alignment [...]

This is indeed a consistent and interesting point to discuss.

We believe this discrepancy stems from finite‑sample effects. By the multivariate Central Limit Theorem, the variance of the empirical noise‑covariance spectrum decays like $1/\sqrt{n}$. In low‑sample regimes, this high variance causes the spectrum to deviate from the true noise covariance $\mathbf{\Gamma}$, thus worsening probing error. Because the excess variance scales with $\alpha^2$, its impact is greatest at large $\alpha$. As $n$ increases, these fluctuations vanish, and strong alignment consistently outperforms weaker settings, as predicted by Proposition 4.3.

Importantly, we do not observe this behavior with joint embeddings: because both views undergo the same augmentation pipeline, they maintain identical variance levels for each $\alpha$, even in small‑sample scenarios.

We will incorporate this clarification into the revised manuscript. Thank you for your feedback.

2. Could the authors clarify what latent dimensionality (i.e., value of k) was used in the synthetic experiments? [...]

We apologize for the lack of clarity. In the synthetic experiments, we set the latent dimensionality to $k=50$ informative components (the top 50 principal components of the input data) plus 50 pure‐noise dimensions. Although this setup is described in Appendix C, we will state it explicitly in the main text.

We selected $k=50$ heuristically, and we observed that the qualitative patterns hold for other choices of $k$ as well. In practice, the relative magnitudes of the data components matter more than the exact number of dimensions.

3. The experimental settings for the reconstruction-based method do not mention the use of masking, which is a common component [...]

In section 5.2, we indeed experimented with MAE that implements masking as data augmentation, as provided by the lightly SSL package : https://github.com/lightly-ai/lightly. We use the hyperparameters for ImageNet as provided in the original MAE paper [a]. Results are reported in Table 1 and show that MAE is much more sensible to data corruption than Joint-Embedding approaches, thus confirming the insights of Corollary 4.5.

Moreover, if the reviewer is interested, we refer to our answer to reviewer ufsf for explanations on how the covariance operator of masking (as defined in [b]) can be leveraged to derive theoretical results similar to those obtained in the Gaussian case (section 4 in the paper). Exploring specific covariances related to known data augmentation functions, rather than generic ones, is a promising direction for future work.

[a] He, Kaiming, et al. "Masked autoencoders are scalable vision learners." Proceedings of the IEEE/CVF conference on computer vision and pattern recognition. 2022.

[b] Lin, Chi-Heng, et al. "The good, the bad and the ugly sides of data augmentation: An implicit spectral regularization perspective." JMLR 25.91 (2024): 1-85.

We would be happy to address any remaining questions if needed.