The Annotated JEPA

surprisetalk2 pts0 comments

The Annotated JEPA | Elements of a Vector Space

"/>

"/>

"/>

"/>

--><br>This post is a step-by-step, annotated, from-scratch walkthrough of Joint Embedding Predictive Architectures, or JEPAs. The goal is to do for JEPA what The Annotated Transformer did for the Transformer: build the full object, explain every moving part, and end with a working training loop.

JEPA is Yann LeCun's proposed answer to a fundamental question in self-supervised learning: how do you train a model to understand the world without labels, without collapsing to trivial solutions, and without wasting capacity on irrelevant details?

The answer, elegant in principle and subtle in practice, is prediction in representation or latent space.

To keep the discussion concrete, the main running example is I‑JEPA, the image instantiationWhy images and video rather than text? LeCun argues that language is already a highly compressed, discrete representation of knowledge; predicting the next token requires modeling human communication patterns, not physical reality. Visual prediction, by contrast, demands understanding of persistence, occlusion, and dynamics. JEPA is designed for domains where pixel-level reconstruction wastes capacity on irrelevant details, a problem that does not arise in the same way for discrete tokens. We return to this near the end., introduced as a self-supervised method that learns semantic image representations by predicting representations of masked regions from visible contextSee Self-Supervised Learning from Images with a Joint-Embedding Predictive Architecture (2023). I-JEPA is a non-generative approach that avoids hand-crafted data augmentations entirely.. We will build I‑JEPA from scratch, then discuss its extension to video with V‑JEPA and V‑JEPA 2See V-JEPA: Latent Video Prediction for Visual Representation Learning (2024), and V-JEPA 2: Self-Supervised Video Models Enable Understanding, Prediction and Planning (2025)., and then examine LeJEPA, the latest attempt to replace engineering heuristics with a distributional regularizerSee LeJEPA: Provable and Scalable Self-Supervised Learning Without the Heuristics (2025)..

What follows is meant to be pedagogical. The implementation omits FlashAttention, gradient checkpointing, mixed-precision, and the batching strategies that make large-scale training feasible. These are engineering choices that would dominate a production codebase but are trivially separable from the mathematics.

The Problem

Self-supervised representation learning asks: how do you learn useful features without labels? You need an objective that captures meaningful structure and, without labels, finding one that actually works is the central difficulty of the field.

JEPA's answer: train by prediction, but predict in representation space. Why should this work at all, though?

Suppose you see part of an image, the context \(x\), and want to learn representations. Somewhere else in the image is a target region \(y\) that you cannot see. An encoder maps \(y\) to a representation \(s_y\). A predictor takes your encoding of the context and outputs \(\hat{s}_y\), its guess at what \(s_y\) should be. Training minimizes the distance \(D(\hat{s}_y, s_y)\).

Now ask: when can the predictor succeed? Only when the context encoding \(s_x\) contains enough information to determine what \(s_y\) must be. If you saw the hood of a car, predicting the representation of the wheels requires that your encoding of the hood captures this is a car. If you saw a face, predicting the representation of the hair requires that your encoding captures identity, pose, and lighting. The predictor cannot hallucinate structure that the context encoding lacks.

This is the forcing function. The context encoder must learn to extract features from \(x\) that are predictive of \(y\)'s representation. These are exactly the semantic, structural features: object identity, spatial relationships, physical constraints. Pixel-level noise in \(x\) does not help predict \(s_y\), so the encoder learns to ignore it. What remains is what generalizes.

The target encoder has a complementary pressure. Its output \(s_y\) must be predictable from context. If \(s_y\) encoded random high-frequency texture, no amount of context would help predict it. So the target encoder learns to output representations that capture the shared structure between \(x\) and \(y\), the structure that makes prediction possible, rather than idiosyncratic details of \(y\) alone.

LeCun's position paper frames this as an energy-based formulationSee A Path Towards Autonomous Machine Intelligence (2022), OpenReview. LeCun frames JEPA as the foundation for world models that plan in latent space.: encode \(x\) and \(y\) into representations, predict one from the other, define energy as prediction error in that abstract space. The architecture factors into components we can implement and analyze: two encoders, one predictor, one distance function.

The JEPA Template

A JEPA...

jepa representation from prediction context self

Related Articles