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

Authors: etai  littwin
omid  saremi
madhu  advani
vimal  thilak
preetum  nakkiran
chen  huang
joshua  susskind apple
Venue: NeurIPS_cc_2024_Conference
Published: 2024-05-15 20:41:54

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.

Main paper: 
 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:EQUATIONwhere 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:EQUATION)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:1. We analytically solve the learning dynamics of JEPA and MAE, in the tractable theoretical setting of deep linear networks. 2. 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). 3. 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.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 parameterizationEQUATIONwhere 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:EQUATION)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 becomeEQUATION)We will denote the diagonal elements of the diagonalized covariance and correlations matrix byσ 2 i = Σ xxii 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. In (e),(f) the MAE objective maintains the same learning order as in (b) on distribution 2, however, the JEPA objective reverses the learning order, due to sensitivity to ρ i . 
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 ∥ W e i ∥ is large. Intuitively, one should expect that if ρ i is small, ∥ W e 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 coefficientsV 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 ∥ W e 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):1. Distribution 1: Set both {λ i }, {ρ i } to be monotonically increasing (i.e ∀ i>j , λ i > λ j , ρ i > ρ j ).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 ∥ W e 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. 
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 equationsEQUATIONwhere 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.1. Σ xx , Σ yx are diagonal and positive definite.2. All weight matrices W a (0), a = 1...L, V (0) are initialized asW a = ϵ 1 L U a , V (0) = ϵ 1 L U vwhere {U a } a , U v are orthogonal matrices, and 0 < ϵ ≪ 1 is an initialization scale.3. For JEPA, we set U v to be diagonal, and for MAE we setU 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 wi = ∥ W e 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 haveEQUATIONSimilarly, the MAE objective eq. ( 5) is optimized using gradient flow according to eq. (6) yielding:EQUATIONRemarkably, 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 wMAE (t, L), wJEPA (t, L) denote the solutions to eqs. (7) and (8) for depth L encoders and at time t, given initial condition w(0) = ϵ, we haveEQUATIONIn addition, it holds that the dynamics of a 1-layer JEPA model match an infinite-depth MAE lim L ′ →∞ wMAE (t, L ′ ) = wJEPA (t, 1).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:EQUATIONThe 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.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. As can be seen, both objectives exhibit greedy learning dynamics with respect to λ, however, only JEPA exhibits greedy dynamics with respect to ρ. (c), (f) darker curves correspond to higher λ but lower ρ. In this case, the order of learning is inverted between the objectives due to the different trends in ρ, λ.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 byEQUATIONas 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 byEQUATION)as long as p is not too close, as defined in eq. (109), to zero or one.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 t * jepa (ϵ, λ, ρ) t * jepa (ϵ, λ, ρ ′ ) = 1 + 2L -1 2L -2 ∆ ρ ϵ 1 L + Θ(ϵ 2 L ), t * mae (ϵ, λ, ρ) t * mae (ϵ, λ, ρ ′ ) = 1 + Θ ϵ L-1 L (14)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. 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. 
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 ρ, λ.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:EQUATIONNote 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:EQUATION)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.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. , 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. 
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.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 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). 
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. 
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.NeurIPS Paper Checklist 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. 
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. 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 wi = ∥ W e 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 haveEQUATIONSimilarly, the MAE objective eq. ( 5) is optimized using gradient flow according to eq. ( 6) yielding:EQUATIONProof. Let us begin with the derivation of the JEPA dynamical equation. The gradient flow equations are given byEQUATIONEvaluating the gradient using the JEPA loss in eq. ( 5), one obtainsEQUATIONEQUATIONSubstituting for the expected values usingEQUATIONEQUATIONWith the gradient flow equations eq. ( 17) in mind, we get the following system of ODEs for {w a , V } to solveEQUATIONEQUATIONSimilarly, using the MAE loss given by eq. ( 5), the gradient flow equations for MAE lead to the following system of ODEsEQUATIONEQUATIONFor 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 followsEQUATION- L+1 b=a+1 W b (t) ⊤ W (t)Σ xx a-1 b=1 W b (t) ⊤ .For V (t) = Ẇ L+1 (t) we obtainEQUATIONFrom the product rule for derivative, we haveEQUATIONSubstituting the expression for Ẇ a (t) in (28) givesEQUATION- L+1 a=1 L+1 b=a+1 W b (t) L+1 b=a+1 W b (t) ⊤ W (t)Σ xx a-1 b=1 W b (t) ⊤ a-1 b=1 W b (t) .Similarly, the dynamical equations for MAE can be re-expressed for W = W (t)EQUATION- L+1 b=a+1 W b (t) ⊤ W (t)Σ xx a-1 b=1 W b (t) ⊤ .For V (t) = Ẇ L+1 (t) we getEQUATIONFor MAE, plugging the expression for Ẇ a (t) into (28) results inEQUATION- L+1 a=1 L+1 b=a+1 W b (t) L+1 b=a+1 W b (t) ⊤ W (t)Σ xx a-1 b=1 W b (t) ⊤ a-1 b=1 W b (t) .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: 26) and ( 27) or (30) and (31).EQUATIONProof. 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)EQUATION- L+1 b=a+1 W b (t) ⊤ W (t)Σ xx a b=1 W b (t) ⊤ .Rewriting (30) for a + 1 ≤ L layer index and multiplying by W a+1⊤ from the left results inEQUATION- L+1 b=a+1 W b (t) ⊤ W (t)Σ xx a b=1 W b (t) ⊤ = Ẇ a (t)W a⊤ (t).Transposing (34) , one obtainsEQUATIONUsing ( 27) and (31) , it can also be shown thatEQUATIONEQUATIONwhich can be written asEQUATIONat 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 impliesEQUATIONat 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 followsEQUATION- L+1 a=1 [W (t)W (t) ⊤ ] 1-a L+1 W (t)Σ xx [W (t) ⊤ W (t)] a-1 L+1 ,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 knowEQUATIONTheorem B.4. For MAE, suppose ω(t) = ω L (t) with a diagonal ω(t). The following ansatzEQUATIONW L+1 (t) = ω(t)U ⊤ ,where U is a constant orthogonal matrix, solves the dynamical equations ( 40) and (31), ifEQUATIONwith ω(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 ansatzEQUATIONPlugging this expression into the equation (40) givesEQUATIONwhere the fact that Σ yx and Σ xx are diagonal was used. Simplifying the above expression, one obtainsEQUATION)One can rewrite this ODE in terms of ω = ω(t). Multiply both sides of (60) Lω L-1 (t) givesEQUATIONNow let us turn to (31) . Using the ansatz eq. ( 56), we can write (31) asEQUATIONwhich 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 wi (t) = ωii (t). Observe thatEQUATIONUsing the above, one arrives at the following ODE for MAEEQUATIONwhere the constant L was absorbed into the definition of the gradient flow time. At this point, we can drop index i to avoid clutterEQUATIONSimilarly, for JEPA one arrives atEQUATIONwhere 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 wMAE (t, L), wJEPA (t, L) denote the solutions to eqs. (7) and (8) for depth L encoders and at time t, given initial condition w(0) = ϵ, we haveEQUATIONIn addition, it holds that the dynamics of a 1-layer JEPA model match an infinite-depth MAE lim L ′ →∞ wMAE (t, L ′ ) = wJEPA (t, 1).Proof. The fixed-point solutions for a finite L can be easily derived by setting ẇ = 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 wMAE (t, L = ∞) is the solution toEQUATIONwith 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.4In 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: Lemma B.6. DefineEQUATIONthenEQUATION)where C is a constant of integration.Proof. We use induction to prove (69). For L = 1, the integral is elementary given byEQUATIONhence, 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 thatEQUATION)Integrating both sides with respect to ψ gives the following recursion relationEQUATIONAssuming the statement of the lemma holds for some L ̸ = 1, one will haveEQUATIONEQUATION= -EQUATIONHence, 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 identityEQUATION)This concludes the proof.We are now in a position to state and prove the main theorem:Theorem B.7 (JEPA dynamics). Suppose ψ = 1 ρ w 1 L . The JEPA dynamics given by eq. ( 7) admits the following implicit solution in ψEQUATION)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 byEQUATIONas 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 toEQUATIONTo fix C, we need to impose the initial condition at t = 0. Using lemma B.6 and notingEQUATIONTo 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 toEQUATIONAs long as ψ * is not too close to 0 or 1 in an ϵ-dependent way, meaning ψ * ∼ Θ(ϵ β ) where β < 1 L , or 1 -ψ * ∼ Θ exp(-ϵ -β ) where β < 2L-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 sayEQUATIONThis implies the following Laurent expansion for the critical time holdsEQUATIONUsing ρ = λ σ 2 , the above equation can be rewritten asEQUATIONThis 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:This fixes C to beEQUATIONTo compute the critical time, we start with (97) and ( 105)EQUATIONTo 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 isEQUATIONIf 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 -givesEQUATION) 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 * = pwhere 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 ϵ pEQUATION)This means in the limit ϵ → 0 the leading term in the Laurent expansion in ϵ of the critical time is given byEQUATIONThe sub-leading Θ(1) contribution to the critical that depends on the choice of uEQUATIONwhere γ = Lρ EQUATIONwhere Φ := Φ(z, a, q) is the Lerch transcendent function defined byEQUATIONand C is an integration constant.From eq. ( 112) we can extract the critical time:Proof. To prove this theorem, note that (97) givesEQUATIONwhich reduces to the statement of the theorem after rearranging terms and constants, where C is given by (105). This proves the theorem.For a single-layer encoder we have (64)EQUATIONwhich has a closed-form solutionEQUATIONSolving for w(t) and imposing the initial condition leads toEQUATIONEQUATIONEQUATIONQualitatively, the dynamics of w(t) for ϵ ≪ 1 can be described as follows:1. w(t) remains indistinguishable 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 writtenEQUATIONThe leading contribution to the critical time in the limit ϵ → 0 isEQUATIONIn 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. 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.The framework we will use is to assume a data sample z ∈ R d is generated as:EQUATION)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:x i = z i m i y i = z i (1 -m i ) m i =1 with probability f 0 with probability 1 -fUnder 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 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:EQUATION)λ i = f (1 -f )(E s 2 i -E s 2 ).(125) From this, we can see thatEQUATIONWe 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:EQUATIONIt follows that:EQUATIONDiag(•) here is an operator removing all non-diagonal elements in the matrix •. Similarly we can derive thatE [m i (1 -m j )] = f (1 -f )(1 -δ ij ) and E [(1 -m i )(1 -m j )] = f (1 -f )δ ij + (1 -f ) 2to show that:EQUATIONE yy T = (1 -f ) 2 E zz T + (f -f 2 )Diag(E zz T ).(130) 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:EQUATION)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:EQUATIONIn 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:lim γd→∞ E z 2 i = E s 2 + σ 2 ϵ (133)Thus, we have simultaneous diagonalizability.C.2 Temporal modelEQUATION)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:E (u a (t)) 2 = 1, E [u a (t)u a (t + 1)] = γ a , E [u a (t)] = 0.E u a (t)u b (t) = 0, E u a (t)u b (t + 1) = 0.(136) This will hold for instance if we generate u a (1) = η a 0 , u a (t + 1) = γ a u a (t) + 1 -γ 2 a η a t , η a t ∼ N (0, 1).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:EQUATION)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. Simultaneous diagonalizability is demonstrated by measured by mean of squared error in off-diagonal elements when attempting to diagonalize Σxy using the eigenbasis of Σxx . The plot show 2 standard deviation with 10 runs.We then have thatEQUATIONHere δ 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:λ a = γ a ∥v a ∥ 2 , a = 1, ..., M, λ a = 0 otherwise. (140) ρ a = γ a ∥v a ∥ 2 σ 2 a + ∥v a ∥ 2 , a = 1, ..., M, ρ a = 0 otherwise.(141) 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. 
 
 Appendix:  


Decision: This paper investigates the implicit biases in two self-supervised learning paradigms: JEPAs and MAE. Through theoretical analysis using deep linear networks, the authors demonstrate that JEPA prioritizes predictive features with minimal noise, while MAE emphasizes high-variance features. These insights are supported by numerical experiments, offering valuable understanding of the distinct feature learning behaviors of JEPA and MAE.

The reviewers (Mqg6, 5fn3, U4uU, MHjv, Nsp4) largely agree that the paper’s theoretical contributions are solid and represent a meaningful addition to the field. Although Mqg6 and MHjv expressed concerns about potentially unrealistic assumptions, the empirical verification under relaxed conditions, as noted by Mqg6, supports the paper’s claims. Additionally, concerns about the limited scope of the experimental results (raised by Nsp4 and 5fn3) were addressed satisfactorily during the rebuttal period with additional experiments.

I recommend accepting this paper for NeurIPS.

Reviews: review: paper_summary: paper_summary: The 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).
strength_weakness: strength_weakness: * There are very strong assumptions that will almost never be realistic
* 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.
questions: questions: * 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.
* 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.
* Is the MLP used in line 478 nonlinear?
limitations: limitations: Yes
paper_summary: paper_summary: This 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.
strength_weakness: strength_weakness: Presentation 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.

On 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?

On Section 4, line 165, the subscript $i$ is dropped. This should be said in the beginning of the section to avoid misunderstanding.

The 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.

The 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.

Figure 3 is too small and not anotated.
questions: questions: Please refer to the "Weakness".
limitations: limitations: Limitations are discussed in the paper.
paper_summary: paper_summary: Analyze 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.
strength_weakness: strength_weakness: *	Quality: 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).
 *	Significance: 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.
questions: questions: 1.	Can 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?
 2.	What is the spectrum of self-supervised algorithms through the lens of deep linear network dynamics?
 3.	What 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?
 4.	Can you demonstrate the effect of the number of layers is different under JEPA and MAE, as the theory predicts?
limitations: limitations: The 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.
paper_summary: paper_summary: The 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.
strength_weakness: strength_weakness: -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.
- 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.
questions: questions: - 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.
- 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?
-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?
-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?
limitations: limitations: The 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.
paper_summary: paper_summary: This 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).

The 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.

Overall, 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.
strength_weakness: strength_weakness: -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.

-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.

-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.
questions: questions: - Would the main theoretical results hold for deep predictor/decoder? 

- 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?
limitations: limitations: Authors adequately addressed the limitations.

Comments: title: response by Reviewer 5fn3
Response to authors
Acknowledgement
Thank your for the rebuttal.
comment: I have carefully reviewed the feedback from other reviewers, considered the author’s rebuttal, and followed the ensuing discussion. I appreciate the authors' thorough responses, particularly their additional experimental results (on W1) and answering my questions.

Assuming that the insights from these discussions will be included in the final paper, I recommend the paper for acceptance as it provides interesting insights and has the potential to contribute to the ML community and I will raise my score from 5 to 7.
I have read the reviewer's responses and am satisfied with their rebuttal. They answered all the questions I had. 

I am keeping my score as a 7 as I believe the paper will have a high impact in the field of Self-Supervised Learning (SSL), helping to theoretically explain the learning phenomena observed in MAE vs. JEPA. While the paper is impactful in systematically explaining this phenomenon, it's not revolutionary in that it was not the first paper to observe and explain this trend. Regardless, having a theoretical backbone is important, and can be used to motivate future research directions within SSL.
It seems that I had some technical difficulty submitting a response to your rebuttal, and my detailed answer is now lost.

In that beautifully written response, I said I was satisfied with your answer, especially with the empirical demonstration of the qualitative prediction in non-linear self-supervised learning, and that I would raise my score (which I did).

Also, I tried to explain better what I meant in my puzzling comment on the "spectrum of self-supervised algorithms", but this explanation was not very important just a few hours before the discussion deadline.
Thank you for the response. My concerns are addressed and I believe this paper would be valuable for the SSL community. Thus I'd raise my score to 6.
The paper highlights a fundamental difference between two main SSL paradigm (MAE and JEPA).  By presenting empirical evidence and theoretical insights, it suggests that JEPA may be more effective for learning semantic features in images, but faces challenges with low-level features required for fine-grained tasks. The authors' rebuttal successfully addressed my initial concerns and provided additional supporting evidence from ImageNet, which strengthened the paper's contributions and led me to revise my score.

I believe this work will be highly relevant and of interest to the SSL community. Thank you for your contribution!