The Annotated JEPA

Elon Litman 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 starts from paired, semantically related views of the world. In the most general form, imagine triples \((x, y, z)\): \(x\) is what you observe, \(y\) is what you want to predict, and \(z\) is an optional latent capturing unknown factors that make the prediction multimodal. The pairs \((x, y)\) are drawn from some joint distribution over observations. In image pretraining, \(x\) might be visible patches and \(y\) the masked patches of the same image; in video, \(x\) is a clip prefix and \(y\) is the continuation; in cross-modal settings, \(x\) could be audio and \(y\) the corresponding video. The only structural requirement is that knowing \(x\) should constrain what \(y\) can be. LeCun's formulation adds an optional latent \(z\) to handle multimodality in the predictive relationship. When multiple values of \(y\) are consistent with the same \(x\), the predictor conditions on \(z\) to select among them. This matters for temporal prediction, where the future is genuinely uncertain, less so for masked image modeling, where the context typically determines the target up to noise. The template has three parts. Encode both observations into a shared representation space: $$ s_x = f_\theta(x), \quad s_y = f_{\bar\theta}(y) $$ The encoders \(f_\theta\) and \(f_{\bar\theta}\) can differ in architecture or parameter sharing. When \(x\) and \(y\) live in different modalities, they must differ; when they are the same modality, weight sharing is a design choice that trades inductive bias against flexibility.

Predict the target representation from the context representation: $$\widehat{s}_y = g_\phi(s_x, z)$$ Minimize prediction error in representation space: $$ \mathcal{L}(\theta, \bar\theta, \phi) = D(\widehat{s}_y, s_y) $$ The rest of this post instantiates these abstractions: defining \(x\) and \(y\) concretely for images, choosing architectures for the encoders and predictor, specifying the distance function \(D\), and implementing the mechanism that prevents the trivial solution where all representations collapse to a constant. Prediction vs. Reconstruction Why predict in representation space rather than pixel space? Consider two image patches that are semantically identical, same object, same meaning, but differ in hundreds of pixels due to lighting, texture, or JPEG artifacts. A reconstruction objective must explain all those differences. The model wastes capacity modeling high-entropy noise. With JEPA, the target encoder can learn to output representations that discard nuisance details. When you regress to that representation, you push the model to capture what the encoder preserves, not what pixels happen to contain. I‑JEPA makes this explicit: the architecture resembles generative models, but the loss lives in embedding space, not input spaceAnalogous to perceptual losses (LPIPS) vs. pixel losses (MSE). Representation-space prediction lets the model focus on semantically meaningful features.. Collapse The easiest way to minimize a matching loss? Output the same vector for everything. Then, loss goes to zero. But the representations become useless. This is the central failure mode of all joint embedding methods, and I‑JEPA addresses it directly: they use an asymmetric design between encoders. The literature offers many anti-collapse strategiesSimCLR uses contrastive negatives. VICReg uses variance-invariance-covariance regularization. BYOL and SimSiam use stop-gradient. DINO/MoCo use EMA teachers. Each trades off differently.: contrastive negatives (expensive), explicit covariance constraints (tricky to tune), stop-gradient (surprisingly effective), or teacher-student encoders with exponential moving average updates (the I‑JEPA choice). I‑JEPA trains the context encoder and predictor by gradient descent, but updates the target encoder by exponential moving average (EMA) of the context encoder weights.

Here is why this prevents collapse. Consider what happens if both encoders are trained by gradient descent. The loss is \(D(\hat{s}_y, s_y)\). If the target encoder can change freely, it will learn to output representations that are easy to predict, which means constant. The context encoder follows. Both converge to outputting the same vector for all inputs: loss zero, representations useless. EMA breaks this co-adaptation. The target encoder updates as $$\theta_{\text{target}} \leftarrow m \theta_{\text{target}} + (1-m) \theta_{\text{context}}$$ with \(m\) close to 1. This means the target encoder lags behind the context encoder. The target representations are stable on the timescale of a gradient step, so the context encoder cannot exploit fast-moving targets to find degenerate solutions. The context encoder must actually learn to predict what the (slowly evolving) target encoder outputs, which forces it to extract meaningful features. Later we discuss LeJEPA, which tries to make collapse prevention principled rather than heuristic. For now, we focus on the classic I‑JEPA design. Specializing JEPA to images A JEPA needs paired views \(x\) and \(y\) that share semantics. For I‑JEPA, both come from the same image: \(x\) is a large context block of patches, and \(y\) consists of several smaller target blocks elsewhere in the image.