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

Q. Would the main theoretical results hold for deep predictor/decoder?

A. Our analysis can be directly extended to deep linear decoders/predictors without changing the results in any qualitative way. Nonlinear predictors however are non-tractable, and therefore beyond the scope of this work. We opted to use a shallow rather than a deep linear predictor in our analyzed model due to 1) a deeper predictor would add to an already notation heavy presentation, without adding or changing the results in any meaningful way. 2) JEPA models typically use a lightweight predictor relative to the encoder, and even linear predictors (see [1]). And 3) Shallow linear predictors have been analyzed theoretically in [1] and [2] and have been shown to work even in practice (though not necessarily producing SOTA results).

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

A. The difference comes from masking or mapping that is applied in the self-supervised model. In the case of random masking one place the difference occurs is on the diagonal elements: because the same element will always be zero in either x_i or y_i, the diagonal elements of C_{xy} will be zero, and this is not true of C_{xx}. For a full derivation and the full impact please see Appendix C.1. One intuitive impact of this difference is that increasing iid pixel noise will not impact diagonal elements of C_{xy} but will impact diagonal elements of C_{xx} leading to a suppression of regression coefficients.

We thank the reviewer for their time and effort in reviewing our paper. If we have sufficiently addressed the reviewers concerns, we kindly ask them to consider raising their score.

[1] Tian et al: Understanding Self-Supervised Learning Dynamics Without Contrastive Pairs“ [2] Richemond et al: “The Edge of Orthogonality: A Simple View of What Makes BYOL Tick”

Q. “More explanations would be helpful to understand the toy setting.”

A. Thank you for pointing this out, we plan to add more elaboration in the rebuttal to make it more accessible.

Q. “what's the motivation of choosing these two distributions?”

A. The motivation for choosing these distributions is to illustrate how changing the structure of a distribution affects the learning speed of different features, as predicted by our theory. Note that, for a linear model, the only relevant aspects of the training data population statistics are its first and second moments. Hence, when it comes to characterizing data features, \lambda_i and \rho_i are the only relevant quantities to consider (assuming a centered distribution with independent components). In other words, the time it takes to learn the i'th feature has to depend on the feature parameters \lambda_i,\rho_i. The distributions in section 3 are meant to illustrate how varying these feature parameters affects the order of learning for each objective. We will add a note to help clarify this in section 2.

Q. “On Section 4, line 165, the subscript is dropped.”

A. Thank you for noticing this.

Q. “The connection between "influential/covarying features" and the parameter is not clear to me.”

A. By “influential features” we simply mean features with a high regression coefficient \rho. In the linear regression literature these features are sometimes referred to as influential/predictive/significant features. By highly covarying features we simply mean features with a high covariance parameter \lambda. Is this explanation clear to the reviewer? We will make this clearer in the intro section, and will include references to this usage in prior work.

Q. “...I believe it'd be helpful to provide some examples on non-linear cases.”

A. We generally agree with this statement, hence we have conducted the following experiment: one prediction that directly follows from our theory is that the JEPA objective allows the model to focus on the bottom subspace of the observed data variance where, at least for natural images, most of the perceptual features tend to reside [1] (see lines 59 - 62 and 74 - 76 in the intro of the paper). We therefore design an experiment that directly tests this hypothesis (see figures in the added pdf). Additionally, we would like to highlight the fact that beyond analytical tractability, linear models are especially appealing as a testbed in our case for the reason that both the MAE and JEPA objectives eventually learn the same features (perhaps with different amplitudes, see Corollary 4.3 in the paper) if we train the models for sufficiently long. This allows us make fair comparisons on the timescale of learning these features in each objective.

Q. “Figure 3 is too small and not annotated”

A. Thanks for pointing this out, we will fix it for the final version.

We thank the reviewer for their time and effort in reviewing our paper. If we have sufficiently addressed the reviewers concerns, we kindly ask them to consider raising their score.

[1]: Randall et al: “Learning By reconstruction Produces Uninformative Features for Perception“

Q. “Can you offer qualitative prediction from the theory to real-world systems implementing JEPA or MAE? What would be such prediction if the same deep architecture was trained on both JEPA and MAE loss, what would you expect to see differently in terms of the learned features?”

A. (Reproduced from our General Response to All Reviewers): The main qualitative insight from the toy model is that noisy (or high variance) features are learned more slowly and with lower amplitude when using the JEPA loss. A direct prediction that follows from our theory is that JEPA will tend to focus on the lower subspace of the data variance (PCA space) where most of the perceptual features reside in natural images (see lines 59 - 62 and 74 - 76 in the intro of the paper). We have conducted additional experiments on realistic models/data providing evidence for this claim (see attached pdf). As additional evidence for this in the literature we would like to point the reviewer to [1] which shows how reconstruction losses tend to focus on the upper part of the PCA space, and [2] which shows that JEPA tends to learn “slow features” (low variance). Our work can be seen as a first principled analysis of these claims in toy settings. Notably, the toy linear setting not only allows for tractable training dynamics but also has the added benefit that both setups learn the same features. This allows us to focus entirely on comparing the schedule according to which the features are learned for the two setups.

Q. What is the spectrum of self-supervised algorithms through the lens of deep linear network dynamics?

A. We are not sure we understood the question; could you please clarify what you meant by “spectrum of SSL algorithms”? (Do you mean the linear-algebraic notion of “spectrum”, or are you asking about how the variety of different SSL algorithms manifest in the deep linear setting, beyond MAE and JEPA?)

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

A. In Figure 3, we demonstrate that the data distribution considered in our paper (along with its assumed constraints) can in principle be realized — that is, our assumptions are not vacuous. Moreover, it provides a simple data-generative model to help guide intuition about how JEPA and MAE will learn differently in deep linear networks. We had included two candidates generative processes: one for the static data and the other for time-varying data. Fig 3. focuses on the time-varying (“video-like”) data generation variant. The generative process is given by eq. (17). To summarize, $v^{a}$s are static images, while $u^a=u^a(t)$ are stochastic processes characterized by the autocorrelation coefficients $\gamma^a$s. The Figure 3.a depicts a sample set of $v$s. Fig 3.b shows two sample realizations of the stochastic functions $u^a$s for two different autocorrelation values. Fig 3.{c, d} show the resulting diagonal and off-diagnal pieces of the covariance matrix generated by the process eq.(17). Fig 3e demonstrates how rapidly one converges to the expected values of $\rho$ and $\lambda$ eq (18) derived in the Appendix (as a function of the log of the “mixing time”). The last plot on the right shows diagonal-ity condition is satisfied. These plots together illustrate that the process in eq.(17) generates a data distribution that satisfies our assumptions 4.1.1 and show us how the key parameters of correlation coefficient ($\rho$) and correlation vary with noise and autocorrelation coefficient $\gamma$ in each mode. We do include theory, but not a similar plot for the random masking task (Appendix C.1) . We put the discussion of this in the appendix because we wanted to keep the formulation simple in the main paper and avoid confusing readers, but it is possible to construct a similar formulation for the random masking case with some additional assumptions.

Q. Can you demonstrate the effect of the number of layers is different under JEPA and MAE, as the theory predicts?

A. Outside of figure 4 which verifies the effect in the linear setting, we have not done comprehensive experiments studying the effect of depth for non-linear networks.

We thank the reviewer for their time and effort in reviewing our paper. If we have sufficiently addressed the reviewers concerns, we kindly ask them to consider raising their score.

[1]: Balestriero et al: “Learning By reconstruction Produces Uninformative Features for Perception“

[2] Sobal et al: “Joint Embedding Predictive Architectures Focus On Slow Features”

Q. “How well do the theoretical results generalize to non-linear models and more complex data distributions?”

A. We have conducted additional experiments on ImageNet (in attached pdf), which are consistent with aspects of our theoretical predictions. (Please see the pdf for details of the setup). Regarding our focus on linear models: The question of how to quantify the efficiency of learning semantic features in practical scenarios may require different metrics than the ones used in this paper. Beyond analytical tractability, concentrating on deep linear networks allows us to separate the question of which features are learned by MAE and JEPA models from the timing/order with which these features are learned, the later being the main focus of this work. Because there is a leap in going between the linear and non-linear setting, we focused on deriving exact results and theoretical guarantees in this paper. Given the long and prolific literature on deep linear models, we view this contribution as self contained, and leave for future work the construction of metrics (of which there may be many) that allow us to analyze the non-linear case where both the features and order are different between the two models.

Q. what are the practical implications of the implicit bias observed in JEPA and MAE?

A. The main qualitative insight from our results is that noisy (or high variance) features are learned more slowly and with lower amplitude when using the JEPA loss, since a large feature variance across the dataset (denoted as \sigma_i in the paper) reduces the regression coefficient for a fixed cross covariance \lambda_i. A direct prediction that follows from this is that JEPA will tend to focus on the lower subspace of the data variance (PCA space) where most of the perceptual features reside in natural images, as claimed in [1] (see lines 59 - 62 and 74 - 76 in the intro of the paper).

We have conducted additional experiments on ImageNet providing evidence for this claim (see attached pdf). As additional evidence for this in the literature we would like to point the reviewer to [1] which shows how unlike JEPA, reconstruction losses tend to focus on the upper part of the PCA space, and [2] which shows that JEPA tends to learn “slow features” (low variance). Our work can be seen as a first principled analysis of these claims in a toy settings. Additionally, our results provide an intuition to why JEPA objectives are perhaps inefficient for learning features suited for fine-grained pixel level tasks, as those features tend to be noisy (features that would correspond to a low regression coefficient in the linear setting). Finally, our results point to a fundamental limit of the efficiency of the MAE objective in learning semantic features since depth does not meaningfully change the feature learning dynamics (see theorem 4.7 and figure 4 in the paper), unlike the JEPA objective. Questions such as how practitioners should account for these insights and limitations in practice we consider out of scope and left to future work.

[1]: Randall et al: “Learning By reconstruction Produces Uninformative Features for Perception“

[2] Sobal et al: “Joint Embedding Predictive Architectures Focus On Slow Features”

We thank the reviewer for their time and effort in reviewing our paper. If we have sufficiently addressed the reviewers concerns, we kindly ask them to consider raising their score.

Q. “There was no mention of the moving average commonly used in JEPA...”

A. Indeed EMA (exponential moving average) is often used in practice to boost performance in a variety of SSL methods that employ self distillation, however we argue that the stop gradient operator is the crucial design choice preventing encoder collapse, rather than EMA, which mainly contributes to increased training stability. Indeed, these claims were argued and verified in the SimSiam [1] paper, showing that EMA can be removed altogether without sacrificing performance. Having said that, we agree EMA should be mentioned, and will add a short discussion on it in the paper.

Q. “I am a bit confused by the figures involving the encoder magnitudes...”

A. Let us clarify. We do not compute the average magnitude across layers. Rather, for each feature i (feature i is the i’th component in the input), we measure the magnitude of the projection of the full encoder on e_i (i’th standard basis), which simply corresponds to the norm of the i’th column of the encoder (see line 127 in the paper). Since we analyze deep linear models, the full encoder is a matrix given by the product of all the layers belonging to the encoder (see \bar{W} in eq 3 in the paper). This makes sense as a metric since the amplitude of the projection of e_i on the encoder is the amplitude of the response to a unit input in feature direction i, which corresponds exactly to how sensitive the encoder latent (output) is to feature i. A zero projection would indicate, for example, invariance to the feature, making it a useless feature for any downstream task that uses the encoder latent.

Q. “Is the MLP used in line 478 nonlinear?”

A. No, it is a deep liinear MLP

We thank the reviewer for their time and effort in reviewing our paper. If we have sufficiently addressed the reviewers concerns, we kindly ask them to consider raising their score.

[1] Chen et al: “Exploring Simple Siamese Representation Learning”