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

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.

Strengths

- 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. - The numerical simulations and Section 3 effectively support the theory while also highlighting the qualitative differences between MAE and JEPA.

Weaknesses

-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

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

Rating

8

Confidence

2

Soundness

3

Presentation

3

Contribution

3

Limitations

Authors adequately addressed the limitations.

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.

Strengths

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

Weaknesses

-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

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

Rating

7

Confidence

5

Soundness

3

Presentation

3

Contribution

3

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.

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.

Strengths

* Originality: 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. * Quality: 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. * Clarity: the manuscript is clearly written and well organized, introducing all the relevant background in a concise manner. * Significance: 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.

Weaknesses

* 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

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?

Rating

7

Confidence

4

Soundness

4

Presentation

4

Contribution

3

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.

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.

Strengths

1. 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. 2. The analyis is comprehensive, including both the learning dynamics and the critical time. 3. The difference between the behaviors of JEPA and MAE observed in this paper is interesting. 4. The theoretical results are supported by the numerical experiments.

Weaknesses

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

Please refer to the "Weakness".

Rating

6

Confidence

2

Soundness

3

Presentation

3

Contribution

3

Limitations

Limitations are discussed in the 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).

Strengths

* 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. * The theoretical results are insightful, and the experimental results strongly support the theoretical results. * 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.

Weaknesses

* 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

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

Rating

7

Confidence

3

Soundness

4

Presentation

3

Contribution

4

Limitations

Yes

Thank your for the rebuttal.

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!

Acknowledgement

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.

Response to authors

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.

response by Reviewer 5fn3

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.