Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

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

Hugues Van Assel, Mark Ibrahim, Tommaso Biancalani, Aviv Regev, Randall Balestriero

TL;DR

The work analyzes reconstruction-based and joint-embedding SSL through closed-form solutions for linear models, linking augmentation design to representation quality. It shows that supervised learning can overcome augmentation-noise misalignment with enough data, while SSL requires sufficiently aligned augmentations to suppress irrelevant features, with distinct thresholds for reconstruction vs. joint-embedding. A key finding is that reconstruction is preferable in low-noise scenarios, whereas joint-embedding offers strictly weaker alignment requirements under high-noise conditions, explaining empirical success of JE methods on challenging datasets. Practically, the results guide practitioners to choose SSL paradigms based on noise characteristics and to design augmentations that align with the underlying nuisance structure, potentially improving robustness in real-world applications.

Abstract

Reconstruction and joint embedding have emerged as two leading paradigms in Self Supervised Learning (SSL). Reconstruction methods focus on recovering the original sample from a different view in input space. On the other hand, joint embedding methods align the representations of different views in latent space. Both approaches offer compelling advantages, yet practitioners lack clear guidelines for choosing between them. In this work, we unveil the core mechanisms that distinguish each paradigm. By leveraging closed form solutions for both approaches, we precisely characterize how the view generation process, e.g. data augmentation, impacts the learned representations. We then demonstrate that, unlike supervised learning, both SSL paradigms require a minimal alignment between augmentations and irrelevant features to achieve asymptotic optimality with increasing sample size. Our findings indicate that in scenarios where these irrelevant features have a large magnitude, joint embedding methods are preferable because they impose a strictly weaker alignment condition compared to reconstruction based methods. These results not only clarify the trade offs between the two paradigms but also substantiate the empirical success of joint embedding approaches on real world challenging datasets.

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

TL;DR

The work analyzes reconstruction-based and joint-embedding SSL through closed-form solutions for linear models, linking augmentation design to representation quality. It shows that supervised learning can overcome augmentation-noise misalignment with enough data, while SSL requires sufficiently aligned augmentations to suppress irrelevant features, with distinct thresholds for reconstruction vs. joint-embedding. A key finding is that reconstruction is preferable in low-noise scenarios, whereas joint-embedding offers strictly weaker alignment requirements under high-noise conditions, explaining empirical success of JE methods on challenging datasets. Practically, the results guide practitioners to choose SSL paradigms based on noise characteristics and to design augmentations that align with the underlying nuisance structure, potentially improving robustness in real-world applications.

Abstract

Reconstruction and joint embedding have emerged as two leading paradigms in Self Supervised Learning (SSL). Reconstruction methods focus on recovering the original sample from a different view in input space. On the other hand, joint embedding methods align the representations of different views in latent space. Both approaches offer compelling advantages, yet practitioners lack clear guidelines for choosing between them. In this work, we unveil the core mechanisms that distinguish each paradigm. By leveraging closed form solutions for both approaches, we precisely characterize how the view generation process, e.g. data augmentation, impacts the learned representations. We then demonstrate that, unlike supervised learning, both SSL paradigms require a minimal alignment between augmentations and irrelevant features to achieve asymptotic optimality with increasing sample size. Our findings indicate that in scenarios where these irrelevant features have a large magnitude, joint embedding methods are preferable because they impose a strictly weaker alignment condition compared to reconstruction based methods. These results not only clarify the trade offs between the two paradigms but also substantiate the empirical success of joint embedding approaches on real world challenging datasets.
Paper Structure (49 sections, 13 theorems, 86 equations, 10 figures, 1 table)

This paper contains 49 sections, 13 theorems, 86 equations, 10 figures, 1 table.

Table of Contents

  1. Introduction
  2. The Reconstruction-based approach
  3. The Joint-embedding approach
  4. Contributions
  5. Background
  6. Augmentation-Aware Closed-Form Solutions
  7. Reconstruction-Based Self-Supervised Learning
  8. Joint-Embedding-Based Self-Supervised Learning
  9. Augmentation Alignment Requirements in Self-Supervised Learning
  10. Data, Noise and Augmentation
  11. Supervised Learning Consistency Regardless of Augmentations
  12. Self-Supervised Learning Requires Aligned Augmentation and Noise
  13. Comparison of Joint-Embedding and Reconstruction-Based Methods via Augmentation Alignment Requirement
  14. Experiments
  15. Experiments With Linear Models
  16. ...and 34 more sections

Key Result

Theorem 3.1

Let $\overline{{\mathbf{x}}}_i \coloneqq \mathbb{E}_{\tau \sim \mathcal{T}}[\tau({\mathbf{x}}_i)]$ for each $i\in{[\![n]\!]}$, and define $\overline{{\mathbf{X}}} \coloneqq (\overline{{\mathbf{x}}}_1, \dots, \overline{{\mathbf{x}}}_n)^\top$. Assume that $\tfrac{1}{n} \overline{{\mathbf{X}}}^\top \ov where ${\mathbf{R}} \in \mathbb{R}^{d \times d}$ and ${\mathbf{P}} \in \mathbb{R}^{d \times d}$ are

Figures (10)

  • Figure 1: Two self-supervised learning paradigms studied in this work. Left: Reconstruction problem of \ref{['eq:reconstruction_problem']}: a random augmentation $\tau \sim \mathcal{T}$ is applied to ${\mathbf{x}}$ to form $\tau({\mathbf{x}})$. An encoder $f_{\mathbf{E}}$ together with a decoder $f_{\mathbf{D}}$ is trained to recover ${\mathbf{x}}$ from $\tau({\mathbf{x}})$. Right: Joint embedding problem of \ref{['eq:ssl']}: two independent augmentations $\tau_1,\tau_2 \sim \mathcal{T}$ of the same ${\mathbf{x}}$ are mapped by $f_{\mathbf{W}}$ to nearby representations, while embeddings of different inputs are pushed apart.
  • Figure 2: Injecting corruption-aligned noise into data augmentation improves SSL representation quality on corrupted CIFAR-10. Thus aligning augmentations with the irrelevant components in the data is crucial in SSL. t-SNE Van_Assel_TorchDR visualizations of (left to right): (1) Supervised features (penultimate layer), clean data. (2) Supervised features (penultimate layer), fog-corrupted (severity 5). (3) VICReg representations, clean data. (4) VICReg representations, fog-corrupted (severity 5). (5) VICReg representations, fog-corrupted (severity 5) with fog noise (severity 1) injection during augmentation. Unlike supervised features, VICReg representations degrade significantly under corruption (compare 3 and 4). Injecting noise in the data augmentation (5) enhances class separability.
  • Figure 3: Performance of linear supervised and SSL models (\ref{['sec:reconstruction_ssl', 'sec:je_ssl', 'prop:closed_form_reconstruction', 'prop:closed_form_ssl']}) on MNIST corrupted with synthetic Gaussian noise (\ref{['sec:setup']}) with various augmentation alignment $\alpha$\ref{['sec:theory']}. Each subplot's y-axis is the absolute difference of supervised linear probing loss (on clean vs. corrupted data) and its x-axis is the sample size $n$. This figure highlights that joint-embedding is preferable to reconstruction in the presence of strong irrelevant noise features. Conversely, reconstruction requires less tailored augmentation when dealing with weak irrelevant noise features. Weak noise corresponds to $\lambda^{\bm{\Gamma}}_{\mathrm{max}} = 10^3$ and strong noise to $\lambda^{\bm{\Gamma}}_{\mathrm{max}} = 10^6$ (details in \ref{['sec:exp_linear_models']}).
  • Figure 4: Performance of linear supervised and SSL models (\ref{['sec:reconstruction_ssl', 'sec:je_ssl', 'prop:closed_form_reconstruction', 'prop:closed_form_ssl']}) with synthetic noise (\ref{['sec:setup', 'sec:exp_linear_models']}) on MNIST. Each subplot's y-axis is the absolute difference of supervised linear probing loss (on clean vs. corrupted data) and its x-axis is the sample size $n$.
  • Figure 5: Performance of linear supervised and SSL models (\ref{['sec:reconstruction_ssl', 'sec:je_ssl', 'prop:closed_form_reconstruction', 'prop:closed_form_ssl']}) with synthetic noise (\ref{['sec:setup', 'sec:exp_linear_models']}) on Fashion-MNIST. Each subplot's y-axis is the absolute difference of supervised linear probing loss (on clean vs. corrupted data) and its x-axis is the sample size $n$.
  • ...and 5 more figures

Theorems & Definitions (21)

  • Theorem 3.1
  • Theorem 3.2
  • Remark 4.1
  • Proposition 4.1
  • Proposition 4.1
  • Proposition 4.1
  • Corollary 4.1
  • Theorem A.1
  • proof
  • Theorem A.1
  • ...and 11 more