How JEPA Avoids Noisy Features: The Implicit Bias of Deep Linear Self Distillation Networks

TL;DR

This paper analyzes the training dynamics of deep linear self-distillation networks, revealing an implicit bias towards high-influence features in JEPA that explains its effectiveness in learning abstract representations.

Contribution

It uncovers the implicit bias of JEPA towards high-influence features through a simplified linear model analysis, explaining its empirical success.

Findings

JEPA biases learning towards high-influence features

Linear analysis reveals the mechanism behind JEPA's feature selection

JEPA's implicit bias favors abstract, high-impact features

Abstract

Two competing paradigms exist for self-supervised learning of data representations. Joint Embedding Predictive Architecture (JEPA) is a class of architectures in which semantically similar inputs are encoded into representations that are predictive of each other. A recent successful approach that falls under the JEPA framework is self-distillation, where an online encoder is trained to predict the output of the target encoder, sometimes using a lightweight predictor network. This is contrasted with the Masked AutoEncoder (MAE) paradigm, where an encoder and decoder are trained to reconstruct missing parts of the input in the data space rather, than its latent representation. A common motivation for using the JEPA approach over MAE is that the JEPA objective prioritizes abstract features over fine-grained pixel information (which can be unpredictable and uninformative). In this work, we seek to understand the mechanism behind this empirical observation by analyzing the training dynamics of deep linear models. We uncover a surprising mechanism: in a simplified linear setting where both approaches learn similar representations, JEPAs are biased to learn high-influence features, i.e., features characterized by having high regression coefficients. Our results point to a distinct implicit bias of predicting in latent space that may shed light on its success in practice.