Denoising with a Joint-Embedding Predictive Architecture

Authors: Dengsheng Chen, Jie Hu, Xiaoming Wei, Enhua Wu
Primary area: generative models
Keywords: AIGC, JEPA, Diffusion Model, Flow Matching, Image Synthetic

Abstract:
Joint-embedding predictive architectures (JEPAs) have shown substantial promise in self-supervised representation learning, yet their application in generative modeling remains underexplored. Conversely, diffusion models have demonstrated significant efficacy in modeling arbitrary probability distributions. In this paper, we introduce Denoising with a Joint-Embedding Predictive Architecture (D-JEPA), pioneering the integration of JEPA within generative modeling. By recognizing JEPA as a form of masked image modeling, we reinterpret it as a generalized next-token prediction strategy, facilitating data generation in an auto-regressive manner. Furthermore, we incorporate diffusion loss to model the per-token probability distribution, enabling data generation in a continuous space. We also adapt flow matching loss as an alternative to diffusion loss, thereby enhancing the flexibility of D-JEPA. Empirically, with increased GFLOPs, D-JEPA consistently achieves lower FID scores with fewer training epochs, indicating its good scalability. Our base, large, and huge models outperform all previous generative models across all scales on ImageNet conditional generation benchmarks. Beyond image generation, D-JEPA is well-suited for other continuous data modeling, including video and audio.

Decision: accept (poster)
Reviewer scores: 6, 5, 8, 6

Reviews and discussion:
<thread>
creation date:
1729971883931

thread_type:
official_review

<thread_object>
creator: reviewer_5mfa

<content>
summary:
The paper, D-JEPA, explores a novel approach by merging joint-embedding predictive architectures with generative modeling. It applies masked image modeling as a generalized next-token prediction strategy for autoregressive data generation. D-JEPA employs a combination of diffusion losses and MAE loss to model token distributions effectively. It aims to enhancing generative tasks across various domains like images, videos, and audio. Extensive experiments demonstrate the model's scalability and effectiveness, claiming state-of-the-art performance on ImageNet class-conditional benchmarks.

strengths:
The ideas tha integrating joint-embedding predictive architectures(e.g. MAE and MIM) with diffusion-based generative modeling makes sense. 

The experiments in the paper are robust, covering a wide range of applications, which showcases the versatility of the D-JEPA model. This extensive experimental validation strengthens the paper's claims about the model's effectiveness across different types of generative tasks.

weaknesses:
However, the paper presents a significant methodological oversight—it appears as a mere combination of existing MAE (Masked Autoencoder) and MAR (Masked Autoregressive Model) methods without a critical analysis or ablation study showing how the MAE branch impacts the AR branch. Specifically, the absence of experiments where the MAE branch is removed casts doubt on the incremental benefit of this integration.

The representation learning capability of D-JEPA is also questionable. According to Table 7, the best ImageNet accuracy is obtained when the model is trained on pixel-level data, while the generative experiments focus on a D-JEPA trained on VAE latent space. This discrepancy points to a lack of a truly unified model approach, as the performance seems to depend heavily on the training specifics rather than the model architecture itself.

questions:
1. Could the authors provide an ablation study that isolates the MAE branch to elucidate its specific impact on the AR branch's performance?

2. How does the model's performance on representation learning tasks vary between training in pixel space and latent space, and what does this indicate about the model's ability to serve as a unified architecture?

3. It’s interesting to see from Tables 1 and 2 that D-JEPA has more parameters than MAR. I’m curious—shouldn’t these models have similar architectures? Can the authors explain why D-JEPA is larger? What additional components does D-JEPA include that MAR doesn't?
</content>
</thread_object>

</thread>
<thread>
creation date:
1730212632292

thread_type:
official_review

<thread_object>
creator: reviewer_3tna

<content>
summary:
The author describes a JEPA based architecture which combines a prediction loss and a diffusion loss.
* The prediction loss learns to predict the representation of the mask tokens
* The diffusion loss learns to predict the raw pixels (or the latent representation of the patch in case of VAE)

The architecture itself is based on a three ViTs:
* The context encoder takes the unmasked tokens and outputs Y
* The feature predictor takes Y, pad it, and learns to reconstruct the full representation Z
* The target encoder takes all tokens (no masking) and produces the target representation Z'
* The prediction loss is a smooth l1 loss between Z and Z'
* The diffusion loss takes Z as input and tries to predict the pixels / latents
* The target encoder is itself an EMA of the context encoder

The paper explores this architecture and shows that:
1) The diffusion loss prevents the collapse of traditional JEPA architectures
2) This architecture is scalable, as the results improves with increased compute budget
3) At inference, instead of predicting all unmasked tokens, we can predict a subset of them, as do MaskGIT, to provide better results
4) They reach SOTA performance on class conditional generation on ImageNet 256x256, equating the results of MAR with CFG

strengths:
* Interesting combination of JEPA representation learning with generative AI, showing that representation learning can help generative AI
* Strong SOTA results on ImageNet 256x256, better or equal than MAR
* When applied for representation learning, the context encoder achieves good results on ImageNet classification
* Experiments showing generalization to text conditioned image generation (and not just class conditioned)
* Experiments showing that it works with audio as well, class conditioned video generation, and CFG

weaknesses:
* It would be interesting to have generation results at higher resolution than 256x256 px to see if inference speed suffers when bi-directional attention meets lots of tokens

questions:
n/a
</content>
</thread_object>

</thread>
<thread>
creation date:
1730544432830

thread_type:
official_review

<thread_object>
creator: reviewer_dqp1

<content>
summary:
The proposed D-JEPA is the combination of Joint-embedding Predictive Architecture (I-JEPA) with MAR's dfiffusion parts, ie, leveraging a feautre from JEPA as the condition of diffusion parts.  Good generative performance is achieved by the proposed model.
And authors provide the exps over multi-modalities.

strengths:
1. Suffieicnt experiments, including generative reuslts on Imagenet-256, sufficient abaltion studies and the exps about representation learning.

2. Authors present the effectiveness of D-JEPA  over multi-modalities including videos, images and audio.

3. Authors provide a theretical support about the proposed models.

4. Beyond the generative results, authors provides the empirical study about the linear performance of the proposed D-JEPA with the pixel/latent-level inputs.

weaknesses:
1. The organization and the structure of the current version should be improved. The writing of chapter 3 is very confusing for me. And some important results and discussions should be re-located in the main paper not the appendix.

2. The novelity is limited:
 It seems that this work just replace the MAGE parts of MAR. The simple combination of I-JEPA and MAR's dfiffusion parts. It doesn't solve the core issue between the gap of representation learning and generative modeling.
Such an archecture cannot bring huge improvement on representation learning with latent-level input (Results in AppendixE over  down-stream tasks are good, just for pixel-level inputs), and the corresponding performance are very poor.
Could you please provide the similar settings for MAR? Use both the pixel-level and latent-level inputs to check the representation performance? I think your design with JEPA cannot solve the inherent limitation (latent input -> bad representation results) of such an TOKEN learning + diffusion framework.

3. The core of I-JEPA is the small-block-wise masks to perform the feature-level augmentation to achieve the better linear performance without the help of data-level augmentation. But in your work, the random mask strategy with high ratio is applied. Could you please provide some empirical evidence and intuitions about such a design？


4. Line 200： “It is important to note that training in latent space is not a necessity for D-JEPA; it can be trained in raw space and still achieve excellent results.” WHERE is the corresponding generative results？ I haven't found it.

5. It is unfair to compare Huge scale models in Table1 and 2. The parameters of the proposed D-JEPA is not aligned with MAR-H. 

6. How about the D-JEPA-H linear and finetuning performance in Table 7/8.

7. Please provide the exps on Imagenet-512x512 to compare with  SOTAs in a more comoutation situation.

questions:
Listed in Weakness.
</content>
</thread_object>

</thread>
<thread>
creation date:
1730606840925

thread_type:
official_review

<thread_object>
creator: reviewer_gtwo

<content>
summary:
This paper extends the JEPA framework—a masked-image-based self-supervised representation learning model—to image generation by incorporating an additional MLP head that operates on image patches. This design unifies representation learning and image generation within a single framework, demonstrating improved image generation quality over recent state-of-the-art methods in specific model configurations.

strengths:
(1) This paper introduces a straightforward and effective approach to enhance generation quality by bridging representation learning and image generation.
(2) In the appendix, the authors offer an in-depth analysis of the model design, including additional insights on representation learning as well as applications in video and audio generation, showcasing the versatility of the proposed methods across multiple tasks.
(3) The presentation is clear, and the ideas are easy to understand.

weaknesses:
(1) The performance is comparable to similar approaches like MAR, which does not use the JEPA loss. In Table 1, the proposed method requires more training epochs (1400 vs. 800 for D-JEPA-B VS ) to achieve similar results to MAR-B. While D-JEPA-L/H outperforms MAR-L/H, it also involves more parameters. Similar trends are observed in Table 2.
(2) There is no comparison to baseline methods, such as the effect of removing the JEPA loss.
(3) How does the model perform in unconditional generation tasks or on more complex datasets, such as COCO?

questions:
My main concerns with this paper relate to its performance compared to the similar approach, MAR. The absence of ablation studies on model components and evaluation on other datasets makes it challenging to assess the impact of each design choice.
</content>
</thread_object>

</thread>
<thread>
creation date:
1731465123081

thread_type:
official_comment

<thread_object>
creator: authors

<content>
title:
Preface

comment:
# Preface

We would like to extend our heartfelt thanks to each reviewer for their valuable comments. Before we proceed with a detailed, point-by-point response, I hope that the reviewers and the Area Chair could spare a moment to browse through this blog content. It primarily addresses potential misunderstandings regarding our paper and provides comprehensive comparisons with recent methodologies (those published on ArXiv after the ICLR submission deadline).
</content>
</thread_object>

</thread>
<thread>
creation date:
1732332346015

thread_type:
official_comment

<thread_object>
creator: authors

<content>
title:
About the Revision and Supplementary Material

comment:
Thank you very much for continued attention to our work and for the valuable feedback provided. 

Based on the feedback received during the rebuttal period, we have submitted a revised version of the paper along with additional supplementary materials. Below are some of the major revisions, and we hope that the reviewers and the Area Chair will take a fresh look at the latest version:

1. We have added comparisons with VAR (Visual Autoregressive Modeling: Scalable Image Generation via Next-Scale Prediction) in Table 2.
2. Following the suggestions from Reviewer 5MfA, we have revised the inappropriate expressions in the Related Work section of Appendix A. We have also made careful revisions to the introduction and experiment sections of the paper, emphasizing the performance of generative modeling.
3. In response to Reviewer dQp1's comments, we have included more experimental data in Appendices C.3, D.2, and E.2, including generative results at 512 resolution, results in raw space, and ablation studies on the JEPA Loss.

**Compared with Fluid** Additionally, during the rebuttal phase, we noticed that several reviewers were particularly interested in comparisons between D-JEPA and MAR. Recently, we have observed new work based on MAR, namely Fluid (SCALING AUTOREGRESSIVE TEXT-TO-IMAGE GENERATIVE MODELS WITH CONTINUOUS TOKENS), which extends MAR to a model size of 10.5B for text-to-image tasks. To better address concerns about D-JEPA's performance in more complex tasks, we have provided comparison results in the latest supplementary materials between images generated by the D-JEPA (2.6B) model and the Fluid (10.5B) model in text-to-image tasks. It is evident that the image quality generated by D-JEPA is significantly higher than that of Fluid. 

**More results** Furthermore, to address concerns about D-JEPA's image generation capabilities beyond 256 resolution, we have included more high-resolution (1k) generation results and results across arbitrary resolutions (ranging from 64 to 1024) in the latest supplementary materials. As these results extend beyond the scope of the current work, we have not included them in the paper's appendix but are providing them as [supplementary materials](https://openreview.net/attachment?id=d4njmzM7jf&name=supplementary_material) for the reviewers' consideration.
</content>
</thread_object>

</thread>
<thread>
creation date:
1732580180289

thread_type:
official_comment

<thread_object>
creator: authors

<content>
title:
Rebuttal Summary

comment:
Dear Area Chair and Reviewers,

We would like to express our heartfelt gratitude for the time and effort you have dedicated to reviewing our manuscript. We are especially thankful to reviewers dQp1, 3tNA, and 5MfA for their active participation throughout the rebuttal phase. We sincerely hope that reviewers dQp1 and gTwo will re-evaluate the contents of our revised version, and if there are still any concerns, we encourage you to provide feedback promptly. We are committed to addressing any issues as quickly as possible. Your insightful comments and constructive feedback have been invaluable in guiding our revisions and improving the clarity and quality of our work.

**Key Revisions and Responses:**

1. **Ablation Studies and Additional Results:**
   - We have included more detailed results and analyses of the ablation studies in Appendix C.3 and Table 6.
   - Additional experiments and analyses on ImageNet 512x512 have been added to Appendix D.2 and Table 8.
   - Experiments on D-JEPA's generative modeling in raw pixel space are now included in Appendix E.2 and Table 11.

2. **Other Modifications to the Paper:**
   - We have made more precise revisions to the main content of the paper to eliminate potential misunderstandings.
   - Appropriate modifications have also been made to the Related Work section in Appendix A.

**Broader Impact and Potential:**

Beyond ImageNet experiments, D-JEPA demonstrates significant potential in speech, video, and multimodal understanding and generation tasks, as showcased in Appendix G. The model's versatility and effectiveness in these domains highlight its broader impact and the value it brings to the field.

Considering the reviewers' concerns about D-JEPA's performance on more complex tasks, we have provided additional experimental results and comprehensive comparisons in the [supplementary materials](https://openreview.net/attachment?id=d4njmzM7jf&name=supplementary_material) and [Preface](https://openreview.net/forum?id=d4njmzM7jf&noteId=B3X0OHeuS9) sections. These include comparisons with recent works such as Fluid-10.5B, a continuation of MAR, and REPA, which have recently been posted on ArXiv.

We hope that our detailed responses and the revisions made to the manuscript address your concerns and demonstrate the robustness and significance of our work.

Thank you once again for your invaluable feedback and support.

Warm regards,

The Authors
</content>
</thread_object>

</thread>


Metareview:
metareview:
The paper introduces a new architecture called Joint-Embedding Predictive Architecture (JEPA), which aims to enhance generative modeling tasks. The authors have demonstrated that JEPA achieves state-of-the-art performance on the ImageNet benchmarks. Moreover, its potential extends to various domains such as speech, video, and multimodal tasks, highlighting its versatility and scalability.

The paper's primary strength lies in its innovative approach of incorporating JEPA into generative modeling, marking a significant departure from traditional methods. The extensive experimentation and the results further support the model's capabilities in different domains.

The paper has faced criticism regarding the perceived incremental nature of integrating Masked Autoencoders (MAE) and Masked Autoregressive (MAR) methods. Additionally, there was some initial confusion about the model's approach to representation learning and generation. Concerns were also raised about the fairness of experimental comparisons due to discrepancies in training specifics and parameter sizes.

The authors have effectively addressed the concerns about the integration of MAE and MAR by providing additional results and explanations. The confusion regarding the model's unified approach has been cleared up with a revised focus on generative capabilities. Furthermore, the authors have discussed the parameter size differences in detail and provided supporting results, which helps to validate the model's performance.