Denoising with a Joint-Embedding Predictive Architecture
<|reference_start|>Denoising with a Joint-Embedding Predictive Architecture: 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 class-conditional ImageNet benchmarks. Beyond image generation, D-JEPA is well-suited for other continuous data modeling, including video and audio.<|reference_end|>
@article{chen2024denoising,
    title={Denoising with a Joint-Embedding Predictive Architecture},
    author={Dengsheng Chen, Jie Hu, Xiaoming Wei, Enhua Wu},
    journal={arXiv preprint arXiv:2410.03755},
    year={2024},
    archivePrefix={arXiv},
    eprint={2410.03755},
    primaryClass={cs.LG cs.CV}
}