The Annotated JEPA
The Annotated JEPA
By Elon Litman | January 27, 2026 | Filed under: Deep Learning
This guide serves as a comprehensive, from-scratch, annotated walkthrough of Joint Embedding Predictive Architectures (JEPAs).
The objective here is to mirror the impact of The Annotated Transformer by meticulously constructing the architecture, clarifying every internal mechanism, and concluding with a functional training loop.
The Core Philosophy
JEPA represents Yann LeCun's strategic response to a pivotal challenge in self-supervised learning (SSL): How can a model acquire a deep understanding of the world without human-provided labels, while avoiding the "collapse" into trivial solutions and ignoring useless noise?
The solution is as elegant as it is nuanced: perform prediction within the representation (latent) space rather than the input space.
Why Visuals over Text?
While many models focus on LLMs, JEPA emphasizes images and video. The reasoning is as follows:
- Language is already a discrete, highly compressed abstraction of knowledge. Predicting the next token is more about modeling human communication than physical reality.
- Visuals require a model to grasp concepts like persistence, occlusion, and physical dynamics.
Pixel-level reconstruction is viewed as a waste of computational capacity because it forces the model to obsess over irrelevant high-frequency details.
Roadmap of the Implementation
- Build the foundational I-JEPA architecture.
- Explore extensions to video via V-JEPA and V-JEPA 2.
- Analyze LeJEPA's distributional regularizers.
- Implement the final training loop.
Evolution of the Architecture
| Model | Focus | Key Contribution | Reference |
|---|---|---|---|
| I-JEPA | Images | Semantic representations via masked region prediction | (2023) |
| V-JEPA | Video | Latent video prediction for representation learning | (2024) |
| V-JEPA 2 | Video | Advanced understanding, prediction, and planning | (2025) |
| LeJEPA | General | Replacing heuristics with distributional regularizers | (2025) |
Note on Implementation: To keep the mathematics transparent, this walkthrough omits production-level engineering such as
FlashAttention, mixed-precision training, gradient checkpointing, and complex batching strategies.
The Fundamental Problem
The central struggle of self-supervised representation learning is finding an objective function that captures meaningful structure without labels. JEPA solves this by predicting representations.
The Intuition
Imagine you are presented with a partial image (the context ) and must predict a hidden region (the target ).
- A Predictor takes the encoding of the context and guesses the representation .
- The system only succeeds if the context encoding contains enough semantic information to constrain the possibilities of .
Examples of Semantic Necessity:
- If you see a car's hood, you can only predict the wheels if your encoder understands the concept: "This is a car."
- If you see a face, predicting the hair requires an understanding of identity, pose, and lighting.
Because the predictor cannot "hallucinate" details missing from the context, the encoder is forced to ignore pixel-level noise and focus on structural features (object identity, spatial constraints, etc.).
Theoretical Framework
LeCun frames this within an Energy-Based Model (EBM) context, viewing JEPA as the bedrock for "World Models" that can plan within a latent space.
The JEPA Template
A JEPA operates on paired, semantically linked views of the world. Generally, these are triples :
- : The observed context.
- : The target to be predicted.
- : An optional latent variable used to handle multimodality (when one could lead to multiple valid outcomes).
Application Domains
- Images: = visible patches = masked patches.
- Video: = clip prefix = future continuation.
- Cross-Modal: = audio = corresponding video.
The Mathematical Formulation
Both observations are mapped into a shared latent space:
Where and are the encoders. These may share parameters or utilize different architectures depending on the specific implementation.
Conceptual Code Structure
def jepa_loss(context, target):
# Encode both views
s_x = context_encoder(context)
s_y = target_encoder(target)
# Predict target representation from context
s_y_hat = predictor(s_x)
# Calculate distance (Energy) in latent space
loss = distance_function(s_y_hat, s_y)
return loss
