TLDR: Joint Embedding Predictive Architectures (JEPAs), known for learning powerful data representations, have an anti-collapse mechanism that also implicitly estimates the underlying data density. This new finding, called JEPA-SCORE, allows JEPAs to be used for tasks like outlier detection and data curation, traditionally requiring generative models, without explicit generative training. The method is theoretically proven and empirically validated across various datasets and JEPA models.
In the rapidly evolving landscape of artificial intelligence, foundation models are at the forefront, capable of tackling a multitude of tasks with minimal prior training. A key player in this domain is Self-Supervised Learning, and within it, a powerful family of methods known as Joint Embedding Predictive Architectures, or JEPAs.
JEPAs are designed to learn rich data representations by combining two core objectives. First, they aim for ‘latent-space prediction,’ meaning that the representation of a slightly altered version of a data sample should be predictable from the original sample’s representation. Second, they incorporate an ‘anti-collapse’ term, which ensures that not all data samples end up with identical representations, preventing the model from learning trivial features.
While the anti-collapse term has traditionally been viewed as a straightforward solution to prevent representation collapse, new research from Randall Balestriero, Nicolas Ballas, Mike Rabbat, and Yann LeCun at Meta-FAIR and Brown University uncovers a much deeper secret. Their paper, titled “Gaussian Embeddings: How JEPAs Secretly Learn Your Data Density,” reveals that this anti-collapse mechanism does far more than just prevent collapse—it provably estimates the underlying data density.
This groundbreaking finding means that any successfully trained JEPA model can implicitly provide insights into the probability of a given data sample. This capability opens up exciting new avenues for applications such as data curation, where less probable or outlier data can be identified, and outlier detection, where unusual samples can be flagged. Essentially, JEPAs can now be used for density estimation, a task traditionally reserved for generative models, but without the need for explicit input space reconstruction or a parametric model for the data density.
The researchers introduce a method called JEPA-SCORE to extract this learned density. This score can be computed efficiently and in closed-form using the model’s Jacobian matrix for any given sample. The theory behind JEPA-SCORE is robust, applying regardless of the dataset or specific architecture used within the JEPA family.
Empirical validations across various datasets, including synthetic, controlled, and large-scale Imagenet, support these theoretical findings. The study examined different Self-Supervised Learning methods that fall under the JEPA umbrella, such as I-JEPA and DINOv2, and even multimodal models like MetaCLIP. For instance, in Imagenet, JEPA-SCORE could distinguish between high and low probability samples within a class, often identifying flying birds as high probability and seated birds as low probability for bird classes.
Furthermore, experiments showed that datasets not seen during pretraining, like Galaxy images when a model was trained on Imagenet, consistently yielded much lower JEPA-SCOREs. This demonstrates the model’s implicit understanding of its training data distribution and its ability to identify out-of-distribution samples.
Also Read:
- The Iterative Advantage: Implicit Models and Their Growing Expressive Capabilities
- Glocal Information Bottleneck: A New Approach for Time Series Imputation
The implications of JEPA-SCORE are significant. It establishes a novel connection between JEPAs and score-based methods, bridging two families of AI models previously considered unrelated. This discovery not only enhances our understanding of how JEPAs learn but also provides a powerful, non-parametric density estimation tool in high-dimensional spaces. The authors hope that JEPA-SCORE will pave the way for advancements in outlier detection and more robust model assessment for downstream tasks. You can read the full research paper here: Gaussian Embeddings: How JEPAs Secretly Learn Your Data Density.
