Title: How JEPA Avoids Noisy Features: The Implicit Bias of Deep Linear Self Distillation Networks
Venue: NeurIPS 2024
Accepted: yes
Keywords: SSL

JEPA
Abstract: Two competing paradigms exist for self-supervised learning of data representations. 
    Joint Embedding Predictive Architectures (JEPAs) 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 with a lightweight predictor network. This is contrasted with the Masked Auto Encoder (MAE) paradigm, where an encoder and decoder are trained to reconstruct missing parts of the input in ambient 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 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, or 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.

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

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.

1 Introduction

Representation learning has arguably been one of the most promising prospects and longest-standing goals of deep learning. In simple terms, representation learning refers to the process of automatically discovering and extracting meaningful features or representations from raw data. The learned representations can then be used to solve various downstream tasks, such as classification, and regression, or as general-purpose representations in robotics and embodied agents. In recent years, the practice of representation learning has witnessed remarkable strides, particularly in the domain of Self-Supervised Learning (SSL) [1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12]. The SSL paradigm has garnered significant attention due to its ability to exploit vast amounts of unlabeled data, freeing practitioners from the burden of expensive annotation efforts. While many SSL approaches have been proposed in the literature, two paradigms have emerged as particularly successful, driving the bulk of the recent breakthroughs in the field:

Masked AutoEncoders (MAE) The original MAE [8] and its successive variants (e.g. [13, 14]) introduced a training objective that seeks to reconstruct missing pieces of data from its partially masked input. This approach, while simple, has proven efficient and scalable [15]. Formally, the MAE objective uses an encoder function f W ( X ) : R d → R d and a decoder function g V ( f ) : R d → R d to predict a target y given input x :

where we define a joint distribution over inputs and targets p = p ( x, y ) with x, y ∈ R d . Here we assume tied dimensions for simplicity.

In general, we can think of x and y as different inputs sharing the same underlying semantic information. One common practice is to let the input x correspond to a partially masked input, and the target y corresponds to the masked-out portions of the input. Note that a critical design choice of the MAE objective is that the loss penalizes the reconstruction error directly in the target's input space. This implies that, as perhaps some parts of the target y cannot be predicted given x , the quality of the encoding f W crucially depends on how x, y are sampled.

Joint Embedding Predictive Architectures (JEPA) The JEPA family of models does not attempt to predict y directly and instead opts to predict the latent representation of y given by some encoding function. When the encoder function for x and y is shared, a simple JEPA objective is given by:

Unlike the objective in (1), the objective given by (2) is susceptible to collapse: it can be minimized trivially by employing an encoder and decoder that map all inputs to the same vector. Contrastive methods [16, 6, 17] avoid this issue by including an additional loss term that pushes apart representations for negative pairs (two different samples). The drawback to this approach is that it can require comparing a sample to many negative examples to work effectively. This has led to a rise in the popularity of non-contrastive SSL. With non-contrastive methods, it becomes imperative to implicitly or explicitly bias the representation mapping f W ( x ) to avoid its collapse to a trivial solution (e.g., f W ( x ) = 0 , ∀ x ). In this paper, we leverage the widely adopted [11] StopGrad operator, preventing the flow of gradients through the f W ( y ) branch. Prominent examples of this foundational architecture include BYOL [9], data2vec [10], SimSiam [11] and the more recent I-JEPA [18] and V-JEPA [19] models. Consequently, the JEPA models under consideration fall under non-contrastive self-distillation methods. We will refer to such methods as JEPA for the remainder of this paper.

While both paradigms have found practical use, characterizing the implicit bias of each remains a major research question. Recently, it was shown in [20] that in the vision domain, perceptual features mostly reside in the data subspace which explains a relatively small portion of the observed variance. This suggests that reconstruction losses are sub-optimal for perception tasks since they tend to prioritize learning subspaces that explain the bulk of the input variance. On the other hand, the JEPA approach appears to be more suitable for perception tasks: empirically, JEPA often achieves better downstream classification performance with fewer optimization steps [3, 19].

In this work, we seek to uncover the mechanisms behind these observations by analyzing tractable deep linear neural networks. To that end, we introduce a fully solvable toy model admitting a complete characterization of its training dynamics for both objectives. Our analysis is inspired by previous work on the greedy learning dynamics of gradient descent, and SSL in particular [21, 22, 23]. However, going beyond previous work, here we present a more nuanced analysis showing how the greedy nature of learning is affected by the choice of objective, model, and the properties of the data. Our analysis reveals a distinct implicit bias of JEPAs: the propensity of the JEPA objective towards learning influential features, defined as features with a large regression coefficient, a property for which the MAE objective is mostly agnostic. Moreover, we show that this distinct bias takes hold only when the encoder is deep and diminishes substantially when the encoder is shallow. This allows JEPAs to steer away from the top subspace of the data which explains its observed variance, and focus on the more semantically rich subspace. Our work therefore uncovers a fundamental implicit bias, shedding light on both the advantages and drawbacks of predicting in latent space.

To summarize, our specific contributions are:

We analytically solve the learning dynamics of JEPA and MAE, in the tractable theoretical setting of deep linear networks.

In this setting, we prove a significant difference between JEPA and MAE: informally, JEPA prioritizes 'influential features' (features which are most informative in prediction), whereas MAE only prioritizes highly covarying features (even if they are noisy and thus less informative).

We show that our theoretical setting is rich enough to capture two more realistic generative processes, based on static and dynamic linear models.

Our results help clarify the precise advantages of joint-embedding architectures, and complement the recent results of [20] on limitations of reconstruction-based methods.

Preliminaries

In this paper, we consider linear over-parameterized encoders and a shallow linear decoder. This tractable setting allows the non-linear training dynamics [24, 25, 26, 27] of the deep linear encoder network to be solved exactly under a certain set of assumptions. Let { W a } a =1 ...L ∈ R d × d and V ∈ R d × d denote the set of weights for a linear encoder f ( x ) of depth L , and that of a linear decoder g ( x ) respectively, and consider the linear parameterization

where ∏ L a =1 W a indicates the conventional right-to-left product of matrices. As for the input data distribution, we will take x, y to be samples from a centered distribution with covariances:

For the purpose of deriving learning dynamics, we will later assume that Σ xx , Σ xy , Σ yy ∈ R d × d are diagonal. Let SG ( · ) denote the StopGrad operator applied on · . Thus, the JEPA and MAE objectives become

We will denote the diagonal elements of the diagonalized covariance and correlations matrix by σ 2 i = Σ xx ii and λ i = Σ yx ii respectively, and let ρ i = λ i σ 2 i denote the regression coefficients.

To motivate our theoretical analysis, we next show empirically that depending on the particular values of { λ i } and { ρ i } , optimizing the objectives in eq. (5) via gradient descent may learn different encoders entirely.

(a)

(b)

(C)

Distribution 1

MAE, Distribution 1

JEPA，Distribution 1

2.0

300

0.7

300

p

0.6

Iteration

0.5

150

0.4

150

0.3

2

0.2

0.1

0.5

50

100

50

100

0

50

100

(d)

(e)

(f)

Distribution2

JEPA，Distribution 2

2.0

入

300

0.7

00E

0.6

p

Iteration

0.5

-0.5

0.4

150

0.3

150

0.3

0.2

02

0.1

0.1

0.5

Feature Index (i)

50

100

0

0

Feature Index (i)

50

100

Feature Index (i)

50

100

2 The Implicit Bias of JEPA

Before presenting the theoretical results, we describe a simple experiment using the setup in section 1 which will flesh out a surprising key difference between the JEPA and MAE objectives in eq. (5).

First, note that each coordinate x i of the input data represents an independent feature characterized by the covariance and regression coefficient λ i , ρ i , which we refer to as feature parameters . Let e i ∈ R d denote the i -th standard basis vector. We say that an encoder ¯ W has learned the i -th feature if its projection on e i , namely ∥ ¯ We i ∥ is large. Intuitively, one should expect that if ρ i is small, ∥ ¯ We i ∥ should be small post training as well. Note that this should not be overly surprising since, at least in the case of the MAE objective, the global minimizer is achieved with the regression coefficients V ¯ W = Σ yx (Σ xx ) -1 . In this investigation, however, we seek an understanding of the dynamics of training, by studying how the MAE and JEPA objectives prioritize different features in training, based on their feature parameters.

To this end, we track how the projections ∥ ¯ We i ∥ evolve during training for each component i by training models on Gaussian data with feature parameters λ i , ρ i according to the following two distributions (see Figure 1(a) and (d) for an illustration):

Distribution 1 : Set both { λ i } , { ρ i } to be monotonically increasing (i.e ∀ i>j , λ i > λ j , ρ i > ρ j ).

Distribution 2 : Set { λ i } to be monotonically increasing, and { ρ i } to be monotonically decreasing.

For each distribution, we initialize all the weights with a Gaussian initialization and train the models using both the MAE and JEPA objectives using gradient descent on batches of randomly sampled data. We illustrate the dynamics of training by measuring the magnitude of the encoder components ∥ ¯ We i ∥ during training in Figure 1. On the first distribution, we observe clear greedy learning dynamics for both objectives, where components with a higher λ i are learned earlier on in training. This behavior in greedy settings is in line with the findings in [21] which showed the same behavior in similar settings. However, on the second distribution, we observe a distinction: while the MAE objective retains its order of learning as in distribution 1, the JEPA objective exhibits a complete inversion of the learning order where features with a higher ρ i , rather than a high λ i are learned first. This implies that in practical scenarios with a finite training budget, JEPA and MAE objectives can potentially learn entirely different features with little to no overlap, as illustrated in Figure 1. To better understand these phenomena, we turn to a dynamical analysis of a simple model that fully recovers our observations.

3 Dynamical Analysis

In the following theoretical analysis, we train the corresponding models employing objectives eq. (5) by optimizing all weights simultaneously starting at some initial weights configuration { W a (0) } a =1 ...L and V (0) , using the gradient flow equations

where L is either JEPA or MAE objective. Generically, the two systems of ODEs are intractable without any simplifying assumptions. A special class of initializations discussed below permits analytical treatment of the problem, allowing us to derive a depth separation result indicative of differences in the inductive bias of the two objectives. In the following, we describe our set of assumptions on the encoder and decoder weights at initialization, and data distribution formally:

Assumption 3.1.

Σ xx , Σ yx are diagonal and positive definite.

All weight matrices W a (0) , a = 1 ...L, V (0) are initialized as W a = ϵ 1 L U a , V (0) = ϵ 1 L U v where { U a } a , U v are orthogonal matrices, and 0 < ϵ ≪ 1 is an initialization scale.

For JEPA, we set U v to be diagonal, and for MAE we set U v such that U v ( ∏ a U a ) is diagonal.

A few notes on assumption 3.1 before stating our results. The assumption on orthogonal weights implies that the weights are 'balanced' at initialization, as is typically assumed in prior works on deep linear networks [28, 22]. The extra technical constraint on the uniformity of the initialization eigenvalues is necessary to deal with the non-uniformity of the data eigenvalues and approximately

holds in typical initializations with large Gaussian matrices. Likewise, we initialize V differently for JEPA and MAE to achieve alignment of eigenvectors between the model and data at initialization. Although somewhat unrealistic, our special initialization schemes will prove informative to the general case. Applying assumption 3.1 to JEPA and MAE objectives in eq. (5) respectively, and using the gradient flow equations in eq. (6), one arrives at the following theorem:

Theorem 3.2 (ODE Equivalence) . Suppose { W a } a =1 ··· L and V are initialized according to assumption 3.1. Let ¯ w i = ∥ ¯ We i ∥ , where e i s are the standard basis. Furthermore, assume the JEPA objective in eq. (5) is optimized using gradient flow according to eq. (6) . Then, we have

Similarly, the MAE objective eq. (5) is optimized using gradient flow according to eq. (6) yielding:

Remarkably, the only difference between the JEPA and MAE dynamical equations is in an exponent. For a derivation of these dynamical equations see appendix B.1. Although the assumption 3.1 enables us to characterize the learning dynamics fully, we will show via numerical simulations (see section 4), that the observations made in this paper will not change qualitatively under the more general and realistic initialization. Since ODEs for each component i decouple, for the remainder of this paper, we drop the subscript i and consider a single ODE for ¯ w parameterized by the encoder and data feature parameters { λ, ρ } . A direct corollary to theorem 3.2 is that the training dynamics of MAE for L >> 1 as described in eq. (8) approach those of JEPA in eq. (7) for L = 1 . Additionally, a depth-dependent difference is apparent in the fixed-point solutions between the two objectives. This is formalized in the following corollary:

Corollary 3.3. Let ¯ w MAE ( t, L ) , ¯ w JEPA ( t, L ) denote the solutions to eqs. (7) and (8) for depth L encoders and at time t , given initial condition ¯ w (0) = ϵ , we have

In addition, it holds that the dynamics of a 1-layer JEPA model match an infinite-depth MAE

When L > 1 , JEPA will suppress the encoding of noisier directions with a lower regression coefficient and this suppression increases with the network depth. At large depths, the encoder becomes low-rank to the point a single dominating eigenvalue corresponding to the feature with the largest regression coefficient remains. Note that corollary 3.3 also indicates that the JEPA solution for L > 1 is not reachable by the MAE objective, and vice versa. While interesting in itself, the difference between the two objectives runs deeper than their fixed-point solutions and is found by analyzing the training dynamics.

In the small initialization regime, the evolution of the weights during training for both objectives exhibits incremental learning dynamics, in which the encoder learns features progressively. We define the critical time , denoted as t ∗ , to be the time it takes for the encoder projection ¯ w to reach a positive finite-but-close-to-1 fraction p of its final fixed point value. It is a quantity we wish to track since it captures an important data-dependent difference between JEPA and MAE training dynamics, concerning the order in which feature learning proceeds in the encoder:

The dependence of this definition of the critical time on a choice for p introduces a degree of arbitrariness in the actual value of t ∗ . However, as long as p is not too close to 1 or 0 , the leading order in the Laurent expansion for t ∗ = t ∗ ( ϵ, λ, ρ ) in ϵ , as we show in section 3.1, is not affected by the specific values of p . The temporal ordering of t ⋆ ( λ, ρ ) s across all features, allows us to observe which feature takes priority in learning according to each objective.

3.1 Critical Time Analysis

To derive the critical time we first solve the dynamical equations in eqs. (7) and (8) in the following. The JEPA dynamics given by eq. (7) can be solved implicitly in a closed form as given by theorem B.7 with a full proof in the supplementary material section.

(a)

(b)

(c)

JEPA,p= 1

JEPA,入 = 1

JEPA

1.0

入0.65

1.0

p=0.65

入0.7

p=0.7

入0.75

p=0.75

入=0.8

p=0.8

0.5

入0.85

入=0.9

0.5

p=0.85

p=0.9

入0.95

p=0.95

0.4

0.0

0’0

0.0

3000

6000

3000

6000

3000

0009

(d)

(e)

(f)

MAE,p=

1

MAE,入 = 1

MAE

1.0

入=0.65

p=0.65

0’T

入=0.7

p=0.7

入=0.75

p=0.75

入=0.8

p=0.8

入=0.85

p=0.85

0.5

=0.9

0.5

p=0.9

0.5

入=0.95

p=0.95

0.0

0.0

0.0

3000

6000

3000

6000

3000

6000

Iterations

Iterations

Iterations

The closed form solution implicitly describes the full trajectory of ¯ w ( t ) for any time t . From this solution, it is clear that the encoder projection in the standard basis ¯ w , starting from its initial value ϵ at t = 0 , will reach its corresponding asymptotic value at t = ∞ . The critical time t ∗ , defined in (11), is the time scale for which ¯ w exits the dynamical region near t = 0 and reaches a finite fraction of its asymptotic value. Theorem 3.4 gives the critical time for JEPA with arbitrary depth L encoders:

Theorem 3.4 (JEPA critical time) . The critical time t ∗ in the small initialization regime ¯ w ( t = 0) = ϵ ≪ 1 for JEPA is given by

as long as p is not too close, as defined in eq. (81) , to zero or one.

In the small initialization regime, the JEPA solution given in theorem B.7 then describes a step-wise learning process of features where each feature is learned in the time scale given by eq. (12). We also derive an analogous theorem B.10 applied to the MAE objective. For technical reasons, we assume L > 1 in the following and provide the equivalent theorems for L = 1 in Appendix B.4. From the MAE dynamics summarized in , eq. (12) we arrive at the following result for MAE critical time:

Theorem 3.5 (MAE critical time) . The MAE critical time t ∗ in the small initialization regime of ¯ w ( t = 0) = ϵ ≪ 1 and L > 1 is given by

as long as p is not too close, as defined in eq. (109) , to zero or one.

3.2 Comparing Learning Dynamics: JEPA vs. MAE

theorem 3.4 and theorem 3.5 reveal a crucial distinction between the two objectives. To understand how each objective prioritizes features, we note the functional form of the leading orders in ϵ in the critical time t ⋆ . For MAE, we observe that the leading order term depends principally on the inverse of λ . Crucially, we note that the next to leading order term (NLO) is small in comparison for any encoder depth L . This implies step-wise learning dynamics where the step ordering is predominantly set by the feature covariance λ . Conversely, features with identical λ will be learned in the same

timescale. In the case of JEPA, we observe that the leading order term is again principally a function of λ , however, the NLO term, which depends inversely on the regression coefficient ρ , is increasingly large with the encoder depth L . That is, features with identical λ but with different ρ will be learned in different timescales, where the feature with the highest regression coefficient is prioritized. Moreover, the priority towards large ρ increases with depth. This implies a meaningful and distinct implicit bias of the JEPA objective towards learning features with a high regression coefficient, increasingly so with depth. This is summarized in the following corollary:

Corollary 3.6. Let t ∗ jepa ( ϵ, λ, ρ ) , t ∗ mae ( ϵ, λ, ρ ) denote the critical time for the JEPA and MAE objectives given feature parameters ϵ, λ, ρ . WLOG assume ρ ′ > ρ , let ∆ ρ = 1 ρ -1 ρ ′ . Then, it holds that

We note the strong dependency on depth L and ∆ ρ in the JEPA case through the term ∆ ρ ϵ 1 L in eq. (14), indicating greedy learning dynamics with respect to ρ , a property absent in the case of MAE. We next turn to empirical simulations to test the validity of our theory, as well as its generalizability in less stringent settings.

(a)

(b)

(c)

(d)

(e)

0.4

2.0

Y = 99

1.6

A1

- 0.3

1.5

Y2 = .95

1.2

0.3

1.4

Az

0.2

- 0.1

0.5

1.0

0.8

1.0

0.2

1.0

1.2

P2

- 0.0

0.1

0.2

0.5

1.0

0.0

0.4

0.6

0.2

0.1

0.0

0.8

0.6

0.3

1.5

0.0

0.1

0.4

0.4

2.0

0.2

0.2

3.0

3.5

4.0

4.5

5.0

5.5

logno(T)

4 Numerical Experiments

We also conduct experiments directly verifying our theoretical findings. To that end, we numerically simulate the ODEs in eqs. (7) and (8) for different distributions and depths starting from a small initialization. In Figure 2, we fix the encoder depth L = 5 while varying the data parameters. As predicted by the theory, we observe a clear greedy learning process for both objectives. However, only the JEPA objective exhibits greedy dynamics with respect to ρ , where components with larger ρ are learned first. Moreover, we observe that the JEPA objective can reverse the feature learning order relative to MAE when ρ and λ have the opposite trends, qualitatively reproducing the observation in section 2. In Figure 4, we run further simulations by varying the depth parameter L . As predicted, we observe that deeper encoders introduce a wider temporal spread between learning features with the same λ but different ρ , in contrast to MAE which learns all features at the same time independent of the depth. Moving beyond gradient flow and the assumptions on initialization, in Figure 5 we train randomly initialized deep linear neural networks using stochastic gradient descent on randomly sampled batches of Gaussian data with varying values for ρ, λ , using both objectives. As is evident, our results carry over to this setup as well, where greedy learning of features can be seen with an objective dependent ordering according to ρ, λ .

4.1 Linear Generative Models

We consider generative models which satisfy our simultaneously-diagonalizable assumption (in the large sample limit). This is a setting in which our theory holds exactly and datasets with this property allow both JEPA and MAE models to learn the same modes, although not necessarily in the same

order. We analyze the simplest linear generative models which still lead to non-trivial learning differences between JEPAs and MAEs for two settings: random masking and a temporal model . The full details of these different settings can be found in Appendix C. In the first case, our views are random masks of the data, this corresponds to a static setting and we derive closed-form expressions for λ and ρ (see C.1). For the remainder of this section, we will focus in on the temporal model due to its simplicity and intuitive interpretation.

The temporal model considers the scenario where our views x, y correspond to our data z 1 , ..., z T at two consecutive times: x = z t and y = z t +1 . The goal is to predict the next frame from the previous one (this setting is consistent with comparing VideoMAE and V-JEPA video models). We now define a simple linear model corresponding to combining independently time-varying images { v a } to the model simultaneously with random noise. We can write this mathematically as:

Note that we are allowing the amplitude of the images to vary with time according to a scalar function u ( t ) ; we make this choice to study the simplest non-trivial temporal variability. The full details of the setup and derivation for simultaneous diagonalizability can be found in Appendix C.2, where we derive the correlation coefficient for each feature:

Here σ a represents the noise amplitude applied to mode a . See Figure 3 for a demonstration of our theory matched to simulations and an example of higher correlation modes with lower correlation coefficient in Figure 3(e). Additionally, a simulation demonstrating an approach to simultaneously diagonalizable empirical covariance and correlation ˆ Σ xx and ˆ Σ xy can be found in Figure 6. Combining (16) with our theory (14) demonstrates how JEPAs can be slower to learn noisy features in a way that MAEs are not. JEPA will also learn noisy features by converging to lower amplitude weights (9), making such features potentially less likely to be used downstream from the embedding.

5 Related Work

Self-supervised learning. One successful SSL paradigm is MAE [8] and its variants [13, 14] which learn representations via input space reconstruction. However, recent works (e.g., [20]) show that such generative paradigms often learn uninformative features for perception, since MAE-like methods tend to learn the data subspace that explains the bulk of observed variance and can include unhelpful information for perception. By contrast, the JEPA paradigm [9, 10, 11, 18, 19] learns to predict the representations of similar image views in the latent space. Such objectives have been found to prioritize semantic features over detailed pixel information, leading to superior performance for perception tasks. To prevent the feature collapse of JEPA, stop-gradients, and other architectural tricks are widely used. More recent works propose techniques like spectrum modulation [29] and weight regularization [30]. Another popular collapse prevention method is based on contrastive loss against negative image pairs [2, 5, 6, 11, 7]. In this work, we focus on MAE and JEPA methods, seeking to understand their implicit bias that leads to different training dynamics. To measure representation quality in JEPAs, [31] devised a metric that effectively counts the number of directions with a large regression coefficient in latent space. Our work can be seen as further motivation for this approach.

Theoretical analysis of SSL. There have been several recent works attempting to understand the success of non-contrastive SSL (JEPA) from various perspectives. [32] studies how the self-distillation in JEPA avoids representation collapse, finding the key role of eigenspace alignment between the predictor and input correlation matrix. The authors of [33] further provide a theoretical bridge between contrastive and non-contrastive objectives towards global and local spectral embedding methods respectively. It is shown that all these methods can be deployed successfully if the pairwise relations during SSL are aligned with the downstream task. While for MAE-based methods, [20] shows they tend to learn uninformative features by pixel reconstruction, unlike the JEPA objective that prioritizes semantic features. In this paper, we seek to understand the mechanism behind this empirical observation by characterizing the implicit bias of both methods using deep linear models.

Learning dynamics in linear networks. Deep linear neural networks have been used in [24, 25, 26, 27, 34, 28, 35, 36] to study the nonlinear dynamics of learning in various settings applying different assumptions, including the saddle point behavior as well as the step-wise behavior where distinct features are learned at different time scales. This is also in line with the observed greedy learning dynamics for modern architectures like Transformers [22, 23]. In the context of SSL, the authors of [21] observed a similar step-wise nature which corresponds to the learning of eigenmodes of the contrastive kernel Σ yx +Σ xy . However, we go beyond this observation by characterizing the distinctive implicit bias of JEPA relative to MAE, i.e., they learn the same features in a step-wise fashion but not necessarily in the same order. Surprisingly, this difference in behavior is only apparent when the encoder network is deep. Our findings about JEPA are related to empirical observations in [37] which show JEPA-based methods focus on 'slow features' that vary slowly over time. This is confirmed in our studies where JEPA-based methods will prioritize 'slow features' for which the marginal distribution has the smallest variance.

6 Limitations

Our theory has several clear limitations. The theoretical results presented in this paper are restricted to conditions in assumption 3.1, which include a deep linear MLP with some restrictions on the initialization scheme, as well as a simultaneously diagonalizable covariances Σ yx , Σ xx . It is worth discussing the implications of these assumptions, and the generalizability of our results to more typical settings. On the model side, empirical simulations with deep linear models initialized using the default initialization [38] indicate our results qualitatively generalize to popular initialization schemes. However, a potentially stronger assumption we make is simultaneously diagonalizability of Σ yx and Σ xx , which allows JEPA and MAE to learn the same features. We provide generative frameworks that adhere to this property, however we expect data distributions of interest to significantly deviate from this assumption. Moreover, the generalizability of our claims to fully practical scenarios yet remains unexplored, and precise characterization of the implicit biases pertaining to more general data distributions and architectures is a direction for future work.

7 Discussion

We have introduced toy models of two of the most popular self-supervised algorithms, namely JEPA and MAE, where we can completely characterize the learning process. We have uncovered a novel learning principle that separates the two objectives: while MAE focuses on highly co-varying features early in training, the JEPA objective focuses on highly influential features (indicated by high values of the regression coefficient). This implicit bias of the JEPA objective allows it to focus on features that are predictive, yet contain minimal noise, as measured by their variance across the training set. This new understanding of the learning dynamics of JEPA sheds light on recent empirical observations and opens the door to new avenues of research strengthening these results. Certainly, one must first generalize these results to non-simultaneously-diagonalizable data distributions, as well as the more imposing challenge of non-linear models. Still, a less than rigorous leap to practical settings may still provide valuable insights. For example, an implicit bias towards predictive yet low-variance features may explain the tendency of JEPA to learn more semantic features, which are inherently less noisy. Conversely, it may shed light on JEPAs vulnerability to spurious or slow features. Additionally, the feature prioritization distinction between the objectives may bear on the efficiency of JEPA in learning semantic features quickly, as these features tend to reside in the low variance subspace of the data, as shown recently [20]. Finally, it is also worth investigating whether the mechanism uncovered here generalizes to other joint embedding architectures not relying on self-distillation.

8 Acknowledgments

We thank Dan Busbridge, Arwen Bradley, Stefano Cosentino, and Shuangfei Zhai for stimulating discussions and useful feedback on the research and writing of this paper.

References

Randall Balestriero, Mark Ibrahim, Vlad Sobal, Ari Morcos, Shashank Shekhar, Tom Goldstein, Florian Bordes, Adrien Bardes, Gregoire Mialon, Yuandong Tian, Avi Schwarzschild, Andrew Gordon Wilson, Jonas Geiping, Quentin Garrido, Pierre Fernandez, Amir Bar, Hamed Pirsiavash, Yann LeCun, and Micah Goldblum. A cookbook of self-supervised learning, 2023.

Ting Chen, Simon Kornblith, Mohammad Norouzi, and Geoffrey E. Hinton. A simple framework for contrastive learning of visual representations. ArXiv , abs/2002.05709, 2020.

Mahmoud Assran, Quentin Duval, Ishan Misra, Piotr Bojanowski, Pascal Vincent, Michael Rabbat, Yann LeCun, and Nicolas Ballas. Self-supervised learning from images with a jointembedding predictive architecture. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition , pages 15619-15629, 2023.

Mathilde Caron, Hugo Touvron, Ishan Misra, Herv'e J'egou, Julien Mairal, Piotr Bojanowski, and Armand Joulin. Emerging properties in self-supervised vision transformers. 2021 IEEE/CVF International Conference on Computer Vision (ICCV) , pages 9630-9640, 2021.

Adrien Bardes, Jean Ponce, and Yann LeCun. Vicreg: Variance-invariance-covariance regularization for self-supervised learning. ArXiv , abs/2105.04906, 2021.

Kaiming He, Haoqi Fan, Yuxin Wu, Saining Xie, and Ross Girshick. Momentum contrast for unsupervised visual representation learning. In Proceedings of the IEEE/CVF conference on computer vision and pattern recognition , pages 9729-9738, 2020.

Chen Huang, Hanlin Goh, Jiatao Gu, and Josh Susskind. MAST: Masked augmentation subspace training for generalizable self-supervised priors. In ICLR , 2023.

Kaiming He, Xinlei Chen, Saining Xie, Yanghao Li, Piotr Doll´ ar, and Ross Girshick. Masked autoencoders are scalable vision learners. In Proceedings of the IEEE/CVF conference on computer vision and pattern recognition , pages 16000-16009, 2022.

Jean-Bastien Grill, Florian Strub, Florent Altch´ e, Corentin Tallec, Pierre Richemond, Elena Buchatskaya, Carl Doersch, Bernardo Avila Pires, Zhaohan Guo, Mohammad Gheshlaghi Azar, et al. Bootstrap your own latent-a new approach to self-supervised learning. Advances in neural information processing systems , 33:21271-21284, 2020.

Alexei Baevski, Wei-Ning Hsu, Qiantong Xu, Arun Babu, Jiatao Gu, and Michael Auli. Data2vec: A general framework for self-supervised learning in speech, vision and language. In International Conference on Machine Learning , pages 1298-1312. PMLR, 2022.

Xinlei Chen and Kaiming He. Exploring simple siamese representation learning. In Proceedings of the IEEE/CVF conference on computer vision and pattern recognition , pages 15750-15758, 2021.

R. Devon Hjelm, Alex Fedorov, Samuel Lavoie-Marchildon, Karan Grewal, Adam Trischler, and Yoshua Bengio. Learning deep representations by mutual information estimation and maximization. ArXiv , abs/1808.06670, 2018.

Zhan Tong, Yibing Song, Jue Wang, and Limin Wang. Videomae: Masked autoencoders are data-efficient learners for self-supervised video pre-training. Advances in neural information processing systems , 35:10078-10093, 2022.

Yanghao Li, Haoqi Fan, Ronghang Hu, Christoph Feichtenhofer, and Kaiming He. Scaling language-image pre-training via masking. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition , pages 23390-23400, 2023.

Limin Wang, Bingkun Huang, Zhiyu Zhao, Zhan Tong, Yinan He, Yi Wang, Yali Wang, and Yu Qiao. Videomae v2: Scaling video masked autoencoders with dual masking. 2023 IEEE/CVF Conference on Computer Vision and Pattern Recognition (CVPR) , pages 14549-14560, 2023.

Ting Chen, Simon Kornblith, Mohammad Norouzi, and Geoffrey Hinton. A simple framework for contrastive learning of visual representations. In International conference on machine learning , pages 1597-1607. PMLR, 2020.

Aaron van den Oord, Yazhe Li, and Oriol Vinyals. Representation learning with contrastive predictive coding. arXiv preprint arXiv:1807.03748 , 2018.

Mahmoud Assran, Quentin Duval, Ishan Misra, Piotr Bojanowski, Pascal Vincent, Michael Rabbat, Yann LeCun, and Nicolas Ballas. Self-supervised learning from images with a jointembedding predictive architecture. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition , pages 15619-15629, 2023.

Adrien Bardes, Quentin Garrido, Jean Ponce, Xinlei Chen, Michael Rabbat, Yann LeCun, Mahmoud Assran, and Nicolas Ballas. Revisiting feature prediction for learning visual representations from video. arXiv preprint arXiv:2404.08471 , 2024.

Randall Balestriero and Yann LeCun. Learning by reconstruction produces uninformative features for perception. arXiv preprint arXiv:2402.11337 , 2024.

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.

Zhiyuan Li, Yuping Luo, and Kaifeng Lyu. Towards resolving the implicit bias of gradient descent for matrix factorization: Greedy low-rank learning. ArXiv , abs/2012.09839, 2020.

Enric Boix-Adser` a, Etai Littwin, Emmanuel Abbe, Samy Bengio, and Joshua M. Susskind. Transformers learn through gradual rank increase. ArXiv , abs/2306.07042, 2023.

Andrew M Saxe, James L McClelland, and Surya Ganguli. Exact solutions to the nonlinear dynamics of learning in deep linear neural networks. arXiv preprint arXiv:1312.6120 , 2013.

Madhu S Advani, Andrew M Saxe, and Haim Sompolinsky. High-dimensional dynamics of generalization error in neural networks. Neural Networks , 132:428-446, 2020.

Bethany Lusch, J Nathan Kutz, and Steven L Brunton. Deep learning for universal linear embeddings of nonlinear dynamics. Nature communications , 9(1):4950, 2018.

Lukas Braun, Cl´ ementine Domin´ e, James Fitzgerald, and Andrew Saxe. Exact learning dynamics of deep linear networks with prior knowledge. Advances in Neural Information Processing Systems , 35:6615-6629, 2022.

Sanjeev Arora, Nadav Cohen, Wei Hu, and Yuping Luo. Implicit regularization in deep matrix factorization. In Neural Information Processing Systems , 2019.

Xi Weng, Yunhao Ni, Tengwei Song, Jie Luo, Rao Muhammad Anwer, Salman Khan, Fahad Shahbaz Khan, and Lei Huang. Modulate your spectrum in self-supervised learning. International Conference on Learning Representations (ICLR) , 2024.

Ali Saheb Pasand, Reza Moravej, Mahdi Biparva, and Ali Ghodsi. WERank: Towards rank degradation prevention for self-supervised learning using weight regularization. ArXiv , abs/2402.09586, 2024.

Vimal Thilak, Chen Huang, Omid Saremi, Laurent Dinh, Hanlin Goh, Preetum Nakkiran, Josh Susskind, and Etai Littwin. Lidar: Sensing linear probing performance in joint embedding ssl architectures. ArXiv , abs/2312.04000, 2023.

Yuandong Tian, Xinlei Chen, and Surya Ganguli. Understanding self-supervised learning dynamics without contrastive pairs. In International Conference on Machine Learning , pages 10268-10278. PMLR, 2021.

Randall Balestriero and Yann LeCun. Contrastive and non-contrastive self-supervised learning recover global and local spectral embedding methods. arXiv preprint arXiv:2205.11508 , 2022.

Andrew M Saxe, James L McClelland, and Surya Ganguli. A mathematical theory of semantic development in deep neural networks. Proceedings of the National Academy of Sciences , 116(23):11537-11546, 2019.

Raphael Berthier. Incremental learning in diagonal linear networks. J. Mach. Learn. Res. , 24:171:1-171:26, 2022.

Scott Pesme and Nicolas Flammarion. Saddle-to-saddle dynamics in diagonal linear networks. ArXiv , abs/2304.00488, 2023.

Vlad Sobal, V JyothirS, Siddhartha Jalagam, Nicolas Carion, Kyunghyun Cho, and Yann LeCun. Joint embedding predictive architectures focus on slow features. ArXiv , abs/2211.10831, 2022.

Kaiming He, X. Zhang, Shaoqing Ren, and Jian Sun. Delving deep into rectifiers: Surpassing human-level performance on imagenet classification. 2015 IEEE International Conference on Computer Vision (ICCV) , pages 1026-1034, 2015.

Alexey Dosovitskiy, Lucas Beyer, Alexander Kolesnikov, Dirk Weissenborn, Xiaohua Zhai, Thomas Unterthiner, Mostafa Dehghani, Matthias Minderer, Georg Heigold, Sylvain Gelly, Jakob Uszkoreit, and Neil Houlsby. An image is worth 16x16 words: Transformers for image recognition at scale. In International Conference on Learning Representations , 2021.

Jia Deng, Wei Dong, Richard Socher, Li-Jia Li, Kai Li, and Li Fei-Fei. Imagenet: A largescale hierarchical image database. In 2009 IEEE conference on computer vision and pattern recognition , pages 248-255. Ieee, 2009.

Ari Morcos, Maithra Raghu, and Samy Bengio. Insights on representational similarity in neural networks with canonical correlation. Advances in neural information processing systems , 31, 2018.

NeurIPS Paper Checklist

1. Claims

Question: Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?

Answer: [Yes]

Justification: The central claims made in the paper are stated in the Abstract and the Introduction. Additionally, we enumerate our contributions in the Introduction.

Guidelines:

The answer NA means that the abstract and introduction do not include the claims made in the paper.

The abstract and/or introduction should clearly state the claims made, including the contributions made in the paper and important assumptions and limitations. A No or NA answer to this question will not be perceived well by the reviewers.

The claims made should match theoretical and experimental results, and reflect how much the results can be expected to generalize to other settings.

It is fine to include aspirational goals as motivation as long as it is clear that these goals are not attained by the paper.

2. Limitations

Question: Does the paper discuss the limitations of the work performed by the authors?

Answer: [Yes]

Justification: We discuss the limitations in Section 6 in the main text.

Guidelines:

The answer NA means that the paper has no limitation while the answer No means that the paper has limitations, but those are not discussed in the paper.

The authors are encouraged to create a separate 'Limitations' section in their paper.

The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (e.g., independence assumptions, noiseless settings, model well-specification, asymptotic approximations only holding locally). The authors should reflect on how these assumptions might be violated in practice and what the implications would be.

The authors should reflect on the scope of the claims made, e.g., if the approach was only tested on a few datasets or with a few runs. In general, empirical results often depend on implicit assumptions, which should be articulated.

The authors should reflect on the factors that influence the performance of the approach. For example, a facial recognition algorithm may perform poorly when image resolution is low or images are taken in low lighting. Or a speech-to-text system might not be used reliably to provide closed captions for online lectures because it fails to handle technical jargon.

The authors should discuss the computational efficiency of the proposed algorithms and how they scale with dataset size.

If applicable, the authors should discuss possible limitations of their approach to address problems of privacy and fairness.

While the authors might fear that complete honesty about limitations might be used by reviewers as grounds for rejection, a worse outcome might be that reviewers discover limitations that aren't acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play an important role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed to not penalize honesty concerning limitations.

3. Theory Assumptions and Proofs

Question: For each theoretical result, does the paper provide the full set of assumptions and a complete (and correct) proof?

Answer: [Yes]

Justification: The assumptions and main theoretical results are included in the main text. Detailed proofs are available in the Appendix.

Guidelines:

The answer NA means that the paper does not include theoretical results.

All the theorems, formulas, and proofs in the paper should be numbered and crossreferenced.

All assumptions should be clearly stated or referenced in the statement of any theorems.

The proofs can either appear in the main paper or the supplemental material, but if they appear in the supplemental material, the authors are encouraged to provide a short proof sketch to provide intuition.

Inversely, any informal proof provided in the core of the paper should be complemented by formal proofs provided in appendix or supplemental material.

Theorems and Lemmas that the proof relies upon should be properly referenced.

4. Experimental Result Reproducibility

Question: Does the paper fully disclose all the information needed to reproduce the main experimental results of the paper to the extent that it affects the main claims and/or conclusions of the paper (regardless of whether the code and data are provided or not)?

Answer: [Yes]

Justification: Experimental setup details are included in the appendix.

Guidelines:

The answer NA means that the paper does not include experiments.

If the paper includes experiments, a No answer to this question will not be perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data are provided or not.

If the contribution is a dataset and/or model, the authors should describe the steps taken to make their results reproducible or verifiable.

Depending on the contribution, reproducibility can be accomplished in various ways. For example, if the contribution is a novel architecture, describing the architecture fully might suffice, or if the contribution is a specific model and empirical evaluation, it may be necessary to either make it possible for others to replicate the model with the same dataset, or provide access to the model. In general. releasing code and data is often one good way to accomplish this, but reproducibility can also be provided via detailed instructions for how to replicate the results, access to a hosted model (e.g., in the case of a large language model), releasing of a model checkpoint, or other means that are appropriate to the research performed.

While NeurIPS does not require releasing code, the conference does require all submissions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example

(a) If the contribution is primarily a new algorithm, the paper should make it clear how to reproduce that algorithm.

(b) If the contribution is primarily a new model architecture, the paper should describe the architecture clearly and fully.

(c) If the contribution is a new model (e.g., a large language model), then there should either be a way to access this model for reproducing the results or a way to reproduce the model (e.g., with an open-source dataset or instructions for how to construct the dataset).

(d) We recognize that reproducibility may be tricky in some cases, in which case authors are welcome to describe the particular way they provide for reproducibility. In the case of closed-source models, it may be that access to the model is limited in some way (e.g., to registered users), but it should be possible for other researchers to have some path to reproducing or verifying the results.

5. Open access to data and code

Question: Does the paper provide open access to the data and code, with sufficient instructions to faithfully reproduce the main experimental results, as described in supplemental material?

Answer: [No]

Justification: The major contributions of the paper are theoretical in nature. The experiments included in the paper are used to support and illustrate analytical results. Consequently, the models used are simple linear models trained with synthetic data. The paper describes how to generate data used in the experiments as well as other setup information needed to run and reproduce experiments.

Guidelines:

The answer NA means that paper does not include experiments requiring code.

Please see the NeurIPS code and data submission guidelines ( https://nips.cc/ public/guides/CodeSubmissionPolicy ) for more details.

While we encourage the release of code and data, we understand that this might not be possible, so 'No' is an acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (e.g., for a new open-source benchmark).

The instructions should contain the exact command and environment needed to run to reproduce the results. See the NeurIPS code and data submission guidelines ( https: //nips.cc/public/guides/CodeSubmissionPolicy ) for more details.

The authors should provide instructions on data access and preparation, including how to access the raw data, preprocessed data, intermediate data, and generated data, etc.

The authors should provide scripts to reproduce all experimental results for the new proposed method and baselines. If only a subset of experiments are reproducible, they should state which ones are omitted from the script and why.

At submission time, to preserve anonymity, the authors should release anonymized versions (if applicable).

Providing as much information as possible in supplemental material (appended to the paper) is recommended, but including URLs to data and code is permitted.

6. Experimental Setting/Details

Question: Does the paper specify all the training and test details (e.g., data splits, hyperparameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?

Answer: [Yes]

Justification: The paper specifies all details used to construct datasets used for training.

Guidelines:

The answer NA means that the paper does not include experiments.

The experimental setting should be presented in the core of the paper to a level of detail that is necessary to appreciate the results and make sense of them.

The full details can be provided either with the code, in appendix, or as supplemental material.

7. Experiment Statistical Significance

Question: Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?

Answer: [No]

Justification: Our experiments are used to verify predictions made by the analytical model developed in the paper. The theory developed in this paper suggests strong effects that we illustrate via experiments. These effects are strong and do not require making fine-grained comparisons via statistical significance testing.

Guidelines:

The answer NA means that the paper does not include experiments.

The authors should answer 'Yes' if the results are accompanied by error bars, confidence intervals, or statistical significance tests, at least for the experiments that support the main claims of the paper.

The factors of variability that the error bars are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, or overall run with given experimental conditions).

The method for calculating the error bars should be explained (closed form formula, call to a library function, bootstrap, etc.)

The assumptions made should be given (e.g., Normally distributed errors).

It should be clear whether the error bar is the standard deviation or the standard error of the mean.

It is OK to report 1-sigma error bars, but one should state it. The authors should preferably report a 2-sigma error bar than state that they have a 96% CI, if the hypothesis of Normality of errors is not verified.

For asymmetric distributions, the authors should be careful not to show in tables or figures symmetric error bars that would yield results that are out of range (e.g. negative error rates).

If error bars are reported in tables or plots, The authors should explain in the text how they were calculated and reference the corresponding figures or tables in the text.

8. Experiments Compute Resources

Question: For each experiment, does the paper provide sufficient information on the computer resources (type of compute workers, memory, time of execution) needed to reproduce the experiments?

Answer: [Yes]

Justification: This information is included with experimental details.

Guidelines:

The answer NA means that the paper does not include experiments.

The paper should indicate the type of compute workers CPU or GPU, internal cluster, or cloud provider, including relevant memory and storage.

The paper should provide the amount of compute required for each of the individual experimental runs as well as estimate the total compute.

The paper should disclose whether the full research project required more compute than the experiments reported in the paper (e.g., preliminary or failed experiments that didn't make it into the paper).

9. Code Of Ethics

Question: Does the research conducted in the paper conform, in every respect, with the NeurIPS Code of Ethics https://neurips.cc/public/EthicsGuidelines ?

Answer: [Yes]

Justification: We confirm that the work presented in this paper is performed in a manner consistent with NeurIPS Ethics Guidelines.

Guidelines:

The answer NA means that the authors have not reviewed the NeurIPS Code of Ethics.

If the authors answer No, they should explain the special circumstances that require a deviation from the Code of Ethics.

The authors should make sure to preserve anonymity (e.g., if there is a special consideration due to laws or regulations in their jurisdiction).

10. Broader Impacts

Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?

Answer: [Yes]

Justification: We discuss the broader impact of our research in Appendix E.

Guidelines:

The answer NA means that there is no societal impact of the work performed.

If the authors answer NA or No, they should explain why their work has no societal impact or why the paper does not address societal impact.

Examples of negative societal impacts include potential malicious or unintended uses (e.g., disinformation, generating fake profiles, surveillance), fairness considerations (e.g., deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations.

The conference expects that many papers will be foundational research and not tied to particular applications, let alone deployments. However, if there is a direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that an improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point out that a generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster.

The authors should consider possible harms that could arise when the technology is being used as intended and functioning correctly, harms that could arise when the technology is being used as intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology.

If there are negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition to attacks, mechanisms for monitoring misuse, mechanisms to monitor how a system learns from feedback over time, improving the efficiency and accessibility of ML).

11. Safeguards

Question: Does the paper describe safeguards that have been put in place for responsible release of data or models that have a high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?

Answer: [NA]

Justification: Our work is theoretical in nature so this question is not applicable.

Guidelines:

The answer NA means that the paper poses no such risks.

Released models that have a high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere to usage guidelines or restrictions to access the model or implementing safety filters.

Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images.

We recognize that providing effective safeguards is challenging, and many papers do not require this, but we encourage authors to take this into account and make a best faith effort.

12. Licenses for existing assets

Question: Are the creators or original owners of assets (e.g., code, data, models), used in the paper, properly credited and are the license and terms of use explicitly mentioned and properly respected?

Answer: [Yes]

Justification: We include this information in experimental details.

Guidelines:

The answer NA means that the paper does not use existing assets.

The authors should cite the original paper that produced the code package or dataset.

The authors should state which version of the asset is used and, if possible, include a URL.

The name of the license (e.g., CC-BY 4.0) should be included for each asset.

For scraped data from a particular source (e.g., website), the copyright and terms of service of that source should be provided.

If assets are released, the license, copyright information, and terms of use in the package should be provided. For popular datasets, paperswithcode.com/datasets has curated licenses for some datasets. Their licensing guide can help determine the license of a dataset.

For existing datasets that are re-packaged, both the original license and the license of the derived asset (if it has changed) should be provided.

If this information is not available online, the authors are encouraged to reach out to the asset's creators.

13. New Assets

Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?

Answer: [NA]

Justification: Paper does not release new assets.

Guidelines:

The answer NA means that the paper does not release new assets.

Researchers should communicate the details of the dataset/code/model as part of their submissions via structured templates. This includes details about training, license, limitations, etc.

The paper should discuss whether and how consent was obtained from people whose asset is used.

At submission time, remember to anonymize your assets (if applicable). You can either create an anonymized URL or include an anonymized zip file.

14. Crowdsourcing and Research with Human Subjects

Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given to participants and screenshots, if applicable, as well as details about compensation (if any)?

Answer: [NA]

Justification: Paper does not involve crowdsourcing nor research with human subjects.

Guidelines:

The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.

Including this information in the supplemental material is fine, but if the main contribution of the paper involves human subjects, then as much detail as possible should be included in the main paper.

According to the NeurIPS Code of Ethics, workers involved in data collection, curation, or other labor should be paid at least the minimum wage in the country of the data collector.

15. Institutional Review Board (IRB) Approvals or Equivalent for Research with Human Subjects

Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or an equivalent approval/review based on the requirements of your country or institution) were obtained?

Answer: [NA]

Justification: paper does not involve crowdsourcing nor research with human subjects.

Guidelines:

The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.

Depending on the country in which research is conducted, IRB approval (or equivalent) may be required for any human subjects research. If you obtained IRB approval, you should clearly state this in the paper.

We recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors to adhere to the NeurIPS Code of Ethics and the guidelines for their institution.

For initial submissions, do not include any information that would break anonymity (if applicable), such as the institution conducting the review.

Contents

Appendix A Additional Figures

A.1 Experimental Details

For Figure 1 and Figure 5, we train MLPs using stochastic gradient descent on random Gaussian data. We use a constant d = 100 width network and fix the encoder depth to be L = 5 . We sample x, y from a jointly Gaussian distribution with diagonal covariances, with different values of ρ, λ per feature, as illustrated in the first column of each figure. At each iteration, we randomly sample a minibatch of 1000 samples and train for 5000 iterations.

JEPA, L=2

JEPA, L=5

JEPA, L=6

JEPA, L=7

02'0

SE O

0.35

0.00

0.00

0.00

0.00

Iterations

500000

Iterations

500000

0

Iterations

500000

Iterations

500000

MAE, L2

MAE, L=5

MAE, L=6

MAE, L=7

08 0

0.80

08 0

0.80

0.00

0.90

0.00

0.00

Iterations

500000

Iterations

500000

Iterations

500000

Iterations

500000

Distribution1

MAE, Distribution 1

JEPA,Distribution 1

Encoder Magnitud Well

400

400

Iteration

Iteration

200

200

K

p

.5

Feature Index (i)

50

100

Feature Index (i)

50

100

Feature Index (i)

50

100

Distribution 2

JEPA,Distribution 2

Encoder Magnitude Well

.0 -

898388

400

400

Iteration

Iteration

200

200

K

02

p

0.1

Feature Index (i)

50

Feature Index (i)

Feature Index (i)

100

Distribution 3

MAE,Distribution 3

Encoder Magnitud IWell

JEPA，Distribution 3

1.0

400

400

p

Iteration

Iteration

200

200

20

0.5

Feature Index (i)

50

100

Feature Index (i)

50

100

Feature Index (i)

50

100

Distribution 4

入

400

400

p

Iteration

Iteration

200

200

20

0.5

50

100

50

100

50

100

Feature Index (i)

Feature Index (i)

Feature Index (i)

Distribution 5

MAE, Distribution 5

Encoder Magnitud We;l

JEPA, Distribution 5

888338

入

400

400

P

Iteration

Iteration

200

Feature Index (i)

50

100

0

Feature Index (i)

100

Feature Index (i)

50

100

Appendix B Proofs

In the following sections, we derive dynamical equations for JEPA and MAE and characterize their training dynamics. We also write down closed-form implicit solutions for their dynamics, all presented in the form of proofs for the theorems/corollaries mentioned in the main paper.

B.1 Proofs of theorem 3.2 and corollary 3.3

Theorem 3.2 (ODE Equivalence) . Suppose { W a } a =1 ··· L and V are initialized according to assumption 3.1. Let ¯ w i = ∥ ¯ We i ∥ , where e i s are the standard basis. Furthermore, assume the JEPA objective in eq. (5) is optimized using gradient flow according to eq. (6) . Then, we have

Similarly, the MAE objective eq. (5) is optimized using gradient flow according to eq. (6) yielding:

Proof. Let us begin with the derivation of the JEPA dynamical equation. The gradient flow equations are given by

Evaluating the gradient using the JEPA loss in eq. (5), one obtains

Substituting for the expected values using E [ xx ⊤ ] = Σ xx and E [ yx ⊤ ] = Σ yx , one gets

With the gradient flow equations eq. (17) in mind, we get the following system of ODEs for { w a , V } to solve

Similarly, using the MAE loss given by eq. (5), the gradient flow equations for MAE lead to the following system of ODEs

For ease of presentation, in this appendix we will use V ( t ) = W L +1 ( t ) , and denote W ( t ) = W L +1 ( t ) ¯ W ( t ) . With this notation, the dynamical equations for JEPA can be rewritten for W = W ( t ) and ¯ W = ¯ W ( t ) as follows

For ˙ V ( t ) = ˙ W L +1 ( t ) we obtain

From the product rule for derivative, we have

Substituting the expression for ˙ W a ( t ) in (28) gives

Similarly, the dynamical equations for MAE can be re-expressed for W = W ( t )

For ˙ V ( t ) = ˙ W L +1 ( t ) we get

For MAE, plugging the expression for ˙ W a ( t ) into (28) results in

In the following, we study the relevant consequences of these dynamical equations and the assumption 3.1. Let us begin by noting that the above assumptions on initialization imply that the weights are 'balanced' at initialization, i.e., W a +1 ⊤ (0) W a +1 (0) = W a (0) W a ⊤ (0) , for a = 1 ...L . In the following lemma, we show during training, the weights will remain balanced:

Lemma B.1. W a +1 ⊤ ( t ) W a +1 ( t ) = W a ( t ) W a ⊤ ( t ) for a = 1 ...L , if W a ( t ) satisfies the ODE sets (26) and (27) or (30) and (31) .

Proof. Multiplying both sides of the MAE dynamic equation (30) from the right by W a ⊤ ( t ) for a = 1 · · · L -1 gives (after absorbing W a ⊤ in the relevant products)

Rewriting (30) for a +1 ≤ L layer index and multiplying by W a +1 ⊤ from the left results in

Transposing (34), one obtains

Therefore, for a = 1 · · · L

which can be written as

Using (27) and (31), it can also be shown that

at all times. Consequently, W a +1 ⊤ ( t ) W a +1 ( t ) -W a ( t ) W a ⊤ ( t ) = C where C is a matrix with constant entries. On the other hand, according to assumption 3.1, W a s a = 1 · · · L +1 are scaled orthogonal matrices at initialization, therefore C = 0 . This implies

at all times for a = 1 · · · L , which proves the claim for MAE. The proof for the JEPA case is similar. This concludes the proof of the lemma.

Using the above lemma one can write the following alternative form for (29) and (32)

Lemma B.2. The dynamical equations for JEPA and MAE share the same general form as follows

where Ξ jepa = ¯ W ( t )Σ yx and Ξ mae = Σ yx .

Proof. Let W a ( t ) = Q a ( t ) ω a ( t ) P a ( t ) denote its SVD decomposition, where Q a , P a are the left and right eigenbasis of the corresponding SVD decomposition. Using lemma B.1, we know

for a = 1 · · · L . Substituting the corresponding SVD decomposition for W a s implies

and the conclusion that all ω s are equal

which leads to

It is easy to see that this implies

On the other hand

Using the above relations, the equations (29) and (32) result in the statement of the lemma. This concludes the proof.

Now we are in a position to state and prove the following theorem:

Theorem B.3. For JEPA, suppose ¯ ω ( t ) = ω ( t ) L with a diagonal ω ( t ) . For any time t , the following ansatz

where U is a constant orthogonal matrix, solves dynamical equations for JEPA, if ¯ ω ( t ) obeys the following differential equation

with ¯ ω (0) = ϵ 1 at initialization.

Proof. First note that the ansatz satisfies conditions 2 and 3 of the assumption 3.1 at initialization. Plugging the ansatz into (40) give

which is the claimed ODE after rearrangement of the terms and noting U is non-singular. Also, note that the equation for V ( t ) is satisfied. In fact, from the ansatz one obtains

Substituting this in (27), we get

Rearranging the terms and expanding the left-hand side gives

which is satisfied if the claimed ODE in the theorem is satisfied. This concludes the proof.

Theorem B.4. For MAE, suppose ¯ ω ( t ) = ω L ( t ) with a diagonal ω ( t ) . The following ansatz

where U is a constant orthogonal matrix, solves the dynamical equations (40) and (31) , if

with ¯ ω (0) = ϵ 1 at initialization.

Proof. Let us start by noticing that the ansatz is consistent with conditions 2 and 3 of the assumption 3.1 at initialization. From the ansatz

Plugging this expression into the equation (40) gives

where the fact that Σ yx and Σ xx are diagonal was used. Simplifying the above expression, one obtains

One can rewrite this ODE in terms of ¯ ω = ¯ ω ( t ) . Multiply both sides of (60) Lω L -1 (t) gives

Now let us turn to (31). Using the ansatz eq. (56), we can write (31) as

which is the same as (60) given Σ yx , Σ xx and ω are diagonal. This concludes the proof.

Keeping in mind ¯ ω (0) is proportional to identity at initialization and Σ xx and Σ xx are assumed diagonal, ¯ ω = ¯ ω ( t ) remains diagonal at all times during its evolution according to (51) and eq. (57). This also means the assumption ω = ω ( t ) diagonal used during derivation of eq. (51) and eq. (57) is a self-consistent one. Let us denote ¯ w i ( t ) = ¯ ω ii ( t ) . Observe that

Using the above, one arrives at the following ODE for MAE

where the constant L was absorbed into the definition of the gradient flow time. At this point, we can drop index i to avoid clutter

Similarly, for JEPA one arrives at

where the definition of the regression coefficient ρ was used. This completes the proof.

We may now proceed to prove corollary 3.3.

Corollary B.5. Let ¯ w MAE ( t, L ) , ¯ w JEPA ( t, L ) denote the solutions to eqs. (7) and (8) for depth L encoders and at time t , given initial condition ¯ w (0) = ϵ , we have

In addition, it holds that the dynamics of a 1-layer JEPA model match an infinite-depth MAE

Proof. The fixed-point solutions for a finite L can be easily derived by setting ˙ ¯ w = 0 for both the JEPA and MAE equivalent ODEs and solving for ¯ w . To show that the limit of the MAE fixed-point at L = ∞ equals the JEPA fixed-point at L = 1 , we note that the right-hand side of eq. (8) is a smooth function of ζ = 1 L for L > 0 . Hence, we have that ¯ w MAE ( t, L = ∞ ) is the solution to

with the initial condition ¯ w (0) , which is identical to the JEPA dynamics in eq. (7).

We will now proceed to discuss the results on JEPA dynamics next:

B.2 JEPA dynamics: proofs of theorem B.7 and theorem 3.4

In this section, we solve the JEPA training dynamics given by the ODE in (66) for general depth L . Let us start by stating and proving the following lemma:

̸

This concludes the proof.

We are now in a position to state and prove the main theorem:

where C is a constant of integration.

Proof. We use induction to prove (69). For L = 1 , the integral is elementary given by

hence, the statement of the lemma holds. Next, we show if the statement of the lemma holds for an arbitrary L = 1 , it will also hold for L +1 . Note that

Integrating both sides with respect to ψ gives the following recursion relation

̸

Assuming the statement of the lemma holds for some L = 1 , one will have

Hence, the claimed relation (69) holds for L +1 as well. An alternative proof strategy would have been based on integrating both sides of the following identity

Theorem B.7 (JEPA dynamics) . Suppose ψ = 1 ρ ¯ w 1 L . The JEPA dynamics given by eq. (7) admits the following implicit solution in ψ

where ψ ( t = 0) = 1 ρ ϵ L , and C is a constant of integration.

Theorem 3.4 (JEPA critical time) . The critical time t ∗ in the small initialization regime ¯ w ( t = 0) = ϵ ≪ 1 for JEPA is given by

as long as p is not too close, as defined in eq. (81) , to zero or one.

Proof. A change for variable of the form ¯ w = ( ρψ ) L in the JEPA dynamical equation given by (66) leads to

To fix C , we need to impose the initial condition at t = 0 . Using lemma B.6 and noting ψ | t =0 = 1 ρ ϵ 1 L

To compute the critical time, take ψ ∗ = ψ ∗ ( t ∗ ) to be the ψ value at which we measure the critical time following the definition and arguments in (11). Putting (78) and (79) together leads to

As long as ψ ∗ is not too close to 0 or 1 in an ϵ -dependent way, meaning ψ ∗ ∼ Θ( ϵ β ) where β < 1 L , or 1 -ψ ∗ ∼ Θ [ exp( -ϵ -β ) ] where β < 2 L -1 L , the first three terms on the right-hand side in the above equation are sub-leading compared to the rest of the terms. Given that ψ ∗ = p 1 L , these conditions for ψ ∗ translate to constraints on the scaling of p , defined in eq. (11). That is to say

This implies the following Laurent expansion for the critical time holds

Using ρ = λ σ 2 , the above equation can be rewritten as

This concludes the proof.

Embedded in the steps of the theorem's proof is an implicit closed-form solution to the JEPA dynamical equation in (66), which for completeness we state as a theorem below:

Theorem B.7 (JEPA dynamics) . Suppose ψ = 1 ρ ¯ w 1 L . The JEPA dynamics given by eq. (7) admits the following implicit solution in ψ

where ψ ( t = 0) = 1 ρ ϵ L , and C is a constant of integration.

Proof. Substituting ¯ w = ( ρψ ) L in (66), one obtains

From lemma B.6, we have

where

which, using the definition of the regression coefficient, can be rewritten as

This completes the proof.

Now we move on to studying the MAE dynamics. It turns out, it makes sense to discuss MAE dynamics for L = 1 and L > 1 cases, separately:

B.3 MAEdynamics: proofs of theorem B.10 and theorem 3.5 for L > 1

Now we turn to analyzing the MAE dynamics for an arbitrary encoder depth L . Let us start with the following lemma:

Lemma B.8. Suppose Ω = [ δ, u ] , where δ is a constant and u , δ belong to (0 , 1) , then the following holds

where α > 2 and C is an integration constant.

Proof. The integrand admits the following Laurent series expansion convergent within the integration domain Ω

therefore, the series expansion can be integrated term by term

where C δ is some delta-dependent constant. In passing, we also recognize that the last expression above can be written in terms of the Lerch transcendent function Φ = Φ( z, s, a ) for the special value of s = 1

where Φ = Φ( z, s, a ) can be represented as a series over the complex plane

It is convergent for | z | < 1 or on the unit circle | z | = 1 when R ( s ) > 1 . This concludes the proof of the lemma.

Next, we proceed to state the main theorem of this section:

Theorem 3.5 (MAE critical time) . The MAE critical time t ∗ in the small initialization regime of ¯ w ( t = 0) = ϵ ≪ 1 and L > 1 is given by

as long as p is not too close, as defined in eq. (109) , to zero or one.

Proof. Let us make the following change of the dependent variable in (64)

Then it can be integrated to give the following

where C is an integration constant. The integral I can be performed using lemma (B.8)

where Φ is the Lerch transcendent function . An initial condition needs to be imposed to fix the integration constant. Recall that

under the small initialization assumption where ϵ → 0 . From (97)

Note that for L > 1

Together with (100), we deduce

for some exponent r . Given

This allows us to write

This fixes C to be

To compute the critical time, we start with (97) and (105)

To further proceed, we make the following observation, stated as a lemma:

Lemma B.9. The second term on the right-hand side of (106) is always sub-leading compared to the first unless u ∗ is chosen so that u ∗ ∼ Θ( ϵ β ) or 1 -u ∗ ∼ Θ [ exp( -ϵ -β ) ] , where β > L -1 L .

Proof. Singular points of I = I ( u ) are at u = 0 or u = 1 . The leading term in the Taylor expansion near u = 0 is

If u ∗ ∼ Θ( ϵ β ) with β < L -1 L , then I ( u ∗ ) would be sub-leading compared to the first term on the right-hand side of (106). In addition, I ( u ) contains a logarithmically singular point at u = 1 . Taylor-expanding around u = 1 -gives

If 1 -u ∗ ∼ Θ [ exp( -ϵ -β ) ] then I ( u ∗ ) ∼ ϵ -β which is sub-leading compared to the first term of the right-hand side of (106) for β < L -1 L . This completes the proof.

Therefore, for any choice of u ∗ , unless it is too close to either of the two singular points at u = 0 and u = 1 (as quantified in the above lemma), the I term in (106) is sub-leading. Given u ∗ = p L -1 L where p is the fraction in the definition of the critical time, defined in eq. (11), the conditions on u ∗ become constraints on p not scaling too close to zero or one with small ϵ

This means in the limit ϵ → 0 the leading term in the Laurent expansion in ϵ of the critical time is given by

The sub-leading Θ(1) contribution to the critical that depends on the choice of u ∗

where γ = Lρ 1 -L 1+ L I ( u ∗ ) , which is u ∗ -dependent. This concludes the proof.

Let us state one of the intermediate results (closed-form solution to MAE dynamics) in proving the above theorem:

Theorem B.10 (MAE dynamics) . Suppose L > 1 and let ψ = ρ 1 -L L +1 ¯ w L -1 L . The MAE dynamics given by eq. (8) admits the following implicit solution in ψ

where Φ := Φ( z, a, q ) is the Lerch transcendent function defined by

and C is an integration constant.

From eq. (112) we can extract the critical time:

Proof. To prove this theorem, note that (97) gives

which reduces to the statement of the theorem after rearranging terms and constants, where C is given by (105). This proves the theorem.

B.4 MAE: shallow encoder L = 1

For a single-layer encoder we have (64)

which has a closed-form solution

Solving for ¯ w ( t ) and imposing the initial condition leads to

Qualitatively, the dynamics of ¯ w ( t ) for ϵ ≪ 1 can be described as follows:

¯ w ( t ) remains indistinguishable from zero until a critical time t ∗ = t ⋆ ( ϵ ) is reached.

2. At t > t ⋆ ( ϵ ) , ¯ w ( t ) grows rapidly until it converges to its asymptotic value given by ¯ w ( ∞ ) = √ ρ .

We can derive the critical time t ∗ defined in (11). The solution (116) can be written

The leading contribution to the critical time in the limit ϵ → 0 is

In summary, for L = 1 and small initialization, learning of projections of the encoder in the standard basis occurs in a step-wise manner over the time scale t ∗ which is controlled by λ only .

B.5 Proof of corollary 3.6

Corollary B.11. Let t ∗ jepa ( ϵ, λ, ρ ) , t ∗ mae ( ϵ, λ, ρ ) denote the critical time for the JEPA and MAE objectives given feature parameters ϵ, λ, ρ . WLOG assume ρ ′ > ρ , let ∆ ρ = 1 ρ -1 ρ ′ . Then, it holds that

Proof. The statement of the corollary is easily seen to hold if one forms JEPA/MAE critical time ratios using the results in theorem 3.4 and theorem 3.5) for some fixed λ but two values of the regression coefficient, denoted as ρ and ρ ′ . Taylor-expanding the resulting expression around ϵ = 0 up to sufficient order proves the corollary.

Appendix C Linear Generative Models

C.1 Random masking with isotropic noise

The framework we will use is to assume a data sample z ∈ R d is generated as:

where η i is noise and s k are model coefficients for important factors (rows of B ). We will also find it useful to define E [ s 2 ] = 1 d ∑ d k =1 E [ s 2 k ] and will assume that the variance of these coefficients and noise is isotropic: E [ η 2 i ] = E [ η 2 ] .

We then sample x and y from masked versions of z as:

Under assumptions that B ki ∼ N (0 , 1 d ) and in the high dimensional limit of infinite d , infinite sample size and some assumptions on the distribution of model coefficients (for instance that they are γ -sparse with finite γ ), we can show:

Theorem C.1. Σ xx and Σ xy are simultaneously diagonalizable in the basis B and have diagonal elements:

From this, we can see that

We can see from this that the order in which components are learned does not depend on the masking ratio, but it does impact the final solution learned and the precise timing in which components are learned. This derivation can also be applied to show that all the parameters in a sparse basis will eventually be learned, but the weighting will depend on the sparsity and noise level of the data and which architecture and loss function we choose.

In order to derive these results, we compute the covariances and correlation terms of this model since our results in the main paper are computed in terms of these:

It follows that:

Diag ( · ) here is an operator removing all non-diagonal elements in the matrix · . Similarly we can derive that E [ m i (1 -m j )] = f (1 -f )(1 -δ ij ) and E [(1 -m i )(1 -m j )] = f (1 -f ) δ ij +(1 -f ) 2 to show that:

Note that all of these matrices have the same form: a weighted sum of the full data covariance and the diagonal of this covariance. It is helpful therefore to compute this covariance directly:

where Σ s ∈ R d × d is a diagonal with components Σ s jk = E [ s 2 k ] δ jk . Unfortunately, the diagonal of the data covariance matrix above is not generally going to be exactly simultaneously diagonalizable with the data covariance unless it is proportional to an identity matrix. We can however show that this will occur in a certain limit. To see this consider the diagonal elements of the full data covariance:

In the limit where ∑ k B 2 ki s 2 k becomes a self-averaged quantity meaning that it does not fluctuate as we draw new instances of B and s from the dataset, we will have that these diagonal elements are also the same for all i and it is therefore proportional to an identity matrix.

To keep our derivation simple we will assume that the coefficients are γ -sparse, meaning [ γ · d ] are drawn from a distribution P s and the rest are zero ( [ · ] rounds to the nearest integer). It follows that:

Thus, we have simultaneous diagonalizability.

C.2 Temporal model

This corresponds to combining a set of images { v a } but flickering them with a time-varying intensity u ( t ) . We will assume the following properties in the large T limit and then construct appropriate functions:

This will hold for instance if we generate

We must ensure T is large enough to allow for good mixing of values for simulations. Note that the mixing time will depend on γ a .

We will also assume that the noise is uncorrelated in time so that: E [ ξ t i ξ t +1 i ] = 0 . It follows that in the large T limit:

For the purposes of this experiment, we assume the images take up different parts of the input space so that we can easily tune the level of noise that shows up in each to explore the trade-offs between temporal correlation and noise. In this case, all vectors v a are orthogonal by definition so we immediately know that the above expression for C xy is diagonal. We then assume the noise is ξ i ∼ N (0 , σ 2 a ( i ) ) where a ( i ) denotes the image that pixel i belongs to.

We then have that

̸

Here δ ij refers to the Kroneker delta function which is zero when i = j and one otherwise. It follows that v a are all rescaled eigenvectors and these matrices are simultaneously diagonalizable with the corresponding parameters:

Appendix D Imagenet Experiments

Our results in the paper predict that JEPA will tend to focus on the lower subspace of data variance where most of the perceptual features reside as claimed in [20]. To test this prediction in a realistic setting, we use the empirical setup described in [20] to study the differences in the feature learning dynamics of I-JEPA [3] and MAE [8]. We first describe the empirical setup and then make our observations.

Setup : We train a vision transformer (ViT) [39] encoder using I-JEPA [3] and MAE [8] on ImageNet-1K [40] data resized to 64 × 64 pixels. We use MAE-style masking to ensure parity with the masking function and copy all other hyperparameters described in I-JEPA [3] and MAE [8]. We build additional datasets by removing certain principal components of the full data subspace as described in [20] bottom refers to the dataset that preservers bottom 25% of explained variance while top denotes the dataset with top 75% of explained variance. We use the original ( full ) images as well the filtered images described above and extract representations using encoders trained with IJEPA and MAE. We then compare the representation between full-bottom and full-top via canonical correlation analysis (CCA) [41] for several checkpoints that are gathered during training.

Observe that I-JEPA's features obtained from the full and bottom datasets show higher similarity compared to MAE throughout training while the trend is reversed for full and top images. This implies that JEPA learns features from the bottom portion of the PCA space faster than MAE. Balestrieo and Lecun [20] suggest that the bottom portion of the spectrum contain features useful for discrimination tasks which is what JEPA focuses on during optimization.

0

200

400

600

800

epochs

0.10

0.15

0.20

0.25

0.30

0.35

0.40

Mean CCA

full_bottom mean CCA

JEPA

MAE

0

200

400

600

800

epochs

0.10

0.11

0.12

0.13

0.14

0.15

0.16

0.17

Mean CCA

full_top mean CCA

JEPA

MAE

Appendix E Broader Impact

This work is a theoretical contribution to shed light on techniques in self-supervised learning to try to understand their implicit biases. The analyses reveal differences in the kinds of features that get

learned, which can potentially be helpful to practitioners in selecting which technique to use when. These benefits may be important in deriving better models for a variety of downstream tasks that people care about such as classification, planning, and decision-making.

Peer reviews:
{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper presents a theoretical analysis of the JEPA and MAE SSL objectives for deep linear networks. Under a somewhat restrictive diagonal covariance assumption, the authors demonstrate that the critical time for learning a feature is dependent only on the input variance for MAE, while JEPA prioritizes learning features with high regression coefficients, which are predictive yet have minimal noise as measured by their variance across the training set (for large depth encoders).\n\nThe authors provide various numerical simulations to validate their theory and show that their findings are robust with respect to the network initialization scheme, even when using a more restrictive initialization necessary for the theory.\n\nOverall, this paper offers valuable insights into the behavior of JEPA and MAE SSL objectives in deep linear networks, and provides a solid foundation for future research in this area.\n\nSTRENGTHS:\n- This work is the first to provide theoretical insights into the empirical observation that JEPA-based approaches tend to learn 'abstract' features more efficiently than MAE. Although the theory is limited to a restrictive case of diagonal linear networks, it calls for further research to generalize these findings.\n\n- The numerical simulations and Section 3 effectively support the theory while also highlighting the qualitative differences between MAE and JEPA.\n\nWEAKNESSES:\n-While the paper is generally clear, the presentation of the theoretical results could be improved. Specifically, it would be helpful to expand on Theorems 4.4 and 4.6 in the main paper, as their validity is not immediately clear without referring to the appendix.\n\n-In addition to the diagonal assumption, the authors focus on the case where the predictor/decoder is linear. However, it is important to use a deep predictor/decoder for both JEPA and MAE methods in practice.\n\n-It would be useful to provide more intuition on what covariance_{x,x} and covariance_{x,y} represent for usual pretraining tasks. This would help readers better understand the underlying concepts and appreciate the significance of the theoretical results.\n\nQUESTIONS:\n- Would the main theoretical results hold for deep predictor/decoder? \n\n- Given that MAE predicts a corrupted input in the same input space, what does it mean for covariance_{x,y} to differ from covariance_{x,x}, and how does the corruption affect this difference?\n\nLIMITATIONS:\nAuthors adequately addressed the limitations.", "TITLE": "Official Review", "RECOMMENDATION": 8, "REVIEWER_CONFIDENCE": 2, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2024-07-02"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe paper investigates the implicit bias of predictive self-supervised learning methods, specifically focusing on the Joint-Embedding Predictive Architecture (JEPA) and comparing it with the Masked Autoencoders (MAE). The study presents a theoretical analysis of the learning dynamics of these methods, revealing how different objectives lead to varied implicit biases in feature learning. It also includes numerical experiments with linear generative models to illustrate the theoretical findings.\n\nSTRENGTHS:\n- The paper provides a rigorous theoretical framework for understanding the implicit bias of JEPA and MAE. This is valuable for the community as it offers insights into why these methods might prioritize certain features over others during training.\n- By comparing JEPA and MAE, the paper helps delineate the strengths and weaknesses of each method. This comparative approach can guide practitioners in choosing the appropriate method for their specific tasks.\n\nWEAKNESSES:\n-The paper would benefit from experiments on more diverse datasets, including real-world data. This would demonstrate the practical implications of the theoretical findings and validate their robustness in more varied scenarios.\n- The numerical experiments are limited in scope. They primarily focus on linear generative models, which may not fully capture the complexities of real-world data. Expanding the experiments to include more diverse datasets and model architectures would strengthen the findings.\n\nQUESTIONS:\n- How well do the theoretical results generalize to non-linear models and more complex data distributions? The paper's findings are based on linear generative models, which may not fully represent the behavior of JEPA and MAE in practical settings.\n- what are the practical implications of the implicit bias observed in JEPA and MAE? How should practitioners account for these biases when applying these methods to real-world tasks?\n-The paper suggests that JEPA may be more efficient in learning semantic features. Can this efficiency be quantified in practical scenarios, and how does it impact downstream tasks such as classification and object detection?\n-aree there strategies to mitigate the negative effects of the implicit bias observed in JEPA and MAE? For instance, can architectural modifications or additional regularization techniques help balance the feature learning dynamics?\n\nLIMITATIONS:\nThe paper offers valuable theoretical insights into the implicit bias of JEPA and MAE in predictive self-supervised learning. However, it falls short in terms of related work, novelty,and experimental scope. Addressing these weaknesses and answering the important questions raised would significantly enhance the impact and applicability of the research.", "TITLE": "Official Review", "RECOMMENDATION": 7, "REVIEWER_CONFIDENCE": 5, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2024-07-11"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nAnalyze learning in two self-supervised paradigms, JEPA and MAE, through the lens of learning dynamics in deep linear networks. Report on a qualitative difference in the order in which features are learned, thus demonstrating their different implicit bias.\n\nSTRENGTHS:\n*\tOriginality:  the analysis of deep linear networks has been a prolific line of research in terms of deep learning theory, but the application to the self-supervised setting is novel and promising. Related work is adequately cited, as far as I can see.\n *\tQuality: the submission is very solid technically, with the assumptions of the theory clearly written and interesting results are derived, demonstrating a clear distinction between the two paradigms.\n *\tClarity: the manuscript is clearly written and well organized, introducing all the relevant background in a concise manner.\n *\tSignificance: theory provides here a unique view, where the implications of design choices made in development of different algorithms are made visible and thus can educate development of better algorithms. This premise is not fully fulfilled here, see weaknesses.\n\nWEAKNESSES:\n*\tQuality: this is a theory paper with very limited experimental support (beyond the numerical evaluation of the theory). The random masking and temporal model presented in section 5.1 are very briefly presented, and it is unclear how the simulation results relate to the theoretical predictions. I would expect a clear distillation of an experiment designed to highlight the difference between JEPA and MAE, the qualitative theoretical prediction and how the experimental results support it (or not).\n *\tSignificance: the implications of the results beyond the theory community are unclear. I would expect the authors to be able to offer (i) some characterization of JEPA and MAE, which is useful for practitioners, or (ii) some classification of possible behaviours of self-supervised algorithms, which can inspire the development of new algorithms with different properties than JEPA and MAE.\n\nQUESTIONS:\n1.\tCan you offer qualitative prediction from the theory to real-world systems implementing JEMA or MAE? What would be such prediction if the same deep architecture was trained on both JEMA and MAE loss, what would you expect to see differently in terms of the learned features?\n 2.\tWhat is the spectrum of self-supervised algorithms through the lens of deep linear network dynamics?\n 3.\tWhat is the result presented in Figure 3? How does it demonstrate a correspondence between theory and the experimental results? Can you offer a similar plot for the random masking task?\n 4.\tCan you demonstrate the effect of the number of layers is different under JEPA and MAE, as the theory predicts?\n\nLIMITATIONS:\nThe authors adequately addressed the limitations of their approach, namely the assumptions about a common diagonalization of the covariance matrices' diagonal dynamics in the low-variance initialization. It would be great if the authors provided experimental support for the applicability of their theory to real-world systems implementing JEPA and MAE or even just provided qualitative predictions derived from the theory.", "TITLE": "Official Review", "RECOMMENDATION": 7, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 4, "DATE": "2024-07-11"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper aims to understand the implicit bias of on two paradigms of self-supervised learning, Joint Embedding Predictive Architectures (JEPAs) and Masked Auto Encoder (MAE). The authors introduce a tractable setting of deep diagonal linear networks and charaterize the learning dynamics of two objectives on the toy problem. Through theoretical analysis, the authors show different behavoirs of two objectives: JEPA prioritizes \"influences features\" while MAE prioritizes highly-covarying features. These observations are supported numerical experiments on the toy model and the Linear Generative Models.\n\nSTRENGTHS:\n1. This paper is well written the presentation is nice. The symbols and math formulations are clear. The demonstration on Section 3 is helpful on understanding the paper. \n2.  The analyis is comprehensive, including both the learning dynamics and the critical time. \n3. The difference between the behaviors of JEPA and MAE observed in this paper is interesting.\n4. The theoretical results are supported by the numerical experiments.\n\nWEAKNESSES:\nPresentation on the setting (Section 2) could be improved. More explanations on $\\rho_i, \\lambda_i, \\sigma_i$ would be helpful to understand the toy setting.\n\nOn Section 3, more explanation on the settings are needed, e.g, what's the motivation of choosing these two distributions? What senarios do the distributions represent?\n\nOn Section 4, line 165, the subscript $i$ is dropped. This should be said in the beginning of the section to avoid misunderstanding.\n\nThe conclusion needs to be clearly elaborated. The authors say that JEPA prioritizes \"influential features\", whereas MAE prioritizes \"highly-covarying features\". The connection between \"influential/covarying features\" and the parameter $\\rho_i, \\lambda_i, \\sigma_i$ is not clear to me.\n\nThe study is limited to diagonal linear networks. The analyis on the toy setting provides interesting insights about the learning dynamics of JEPA and MEA. But it's not sure whether the conclusions are applicable to more complex scenarios. I believe it'd be helpful to provide some examples on non-linear cases.\n\nFigure 3 is too small and not anotated.\n\nQUESTIONS:\nPlease refer to the \"Weakness\".\n\nLIMITATIONS:\nLimitations are discussed in the paper.", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 2, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2024-07-11"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe authors study a characteristic of two common approaches towards visual modeling in SSL--particularly MAE and JEPA. Previous works have demonstrated or identified empirically that JEPA architecture are more prone towards lower variance features whereas MAE optimize towards higher variance features. This work focuses on a theoretical understanding of why this occurs through deep linear encoders. The authors also verify their framework empirically, with both linear and nonlinear encoders (I believe they verify with nonlinear encoders, please see questions for more details).\n\nSTRENGTHS:\n* The theoretical contribution is very important, as the question as to whether raw pixel or feature reconstruction works better is widely debated. Thus, the theory can be used to justify specific architectures in different circumstances.\n* The theoretical results are insightful, and the experimental results strongly support the theoretical results. \n* The authors experimented with relaxed assumptions, and had the same or similar empirical results, providing evidence that the theoretical results, while derived under unrealistic assumptions, may still hold true.\n\nWEAKNESSES:\n* There are very strong assumptions that will almost never be realistic\n* I am not convinced that the magnitude of encoders is a proper metric for the training dynamics. There were not many details or justification for this. Please see the questions section for more details.\n\nQUESTIONS:\n* There was no mention of the moving average commonly used in JEPA architectures to update the teacher encoder (the encoder with the stop grad operator). It would be nice to know how different approaches to train/get the teacher encoder affect the training dynamics.\n* I am a bit confused by the figures involving the encoder magnitudes, particularly about the values of encoder magnitudes and how they were gotten for each feature. For the first layer, or if the encoder had a single layer, this metric makes sense. However, as linear layers are fully connected, feature i in weight layer 1 does not necessarily correspond strongly with feature i in weight layer 3. Thus, if the average magnitude across each layer for feature i is being used that seems like an invalid metric.\n* Is the MLP used in line 478 nonlinear?\n\nLIMITATIONS:\nYes", "TITLE": "Official Review", "RECOMMENDATION": 7, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 4, "DATE": "2024-07-13"}