Joint Embedding Predictive Architecture (JEPA)

A Joint Embedding Predictive Architecture (JEPA) is a self-supervised learning design in which a network learns by predicting the abstract representation of a hidden part of its input from a visible part. Given an image with some regions masked, a JEPA predicts the embeddings of the masked regions, without reconstructing their pixels, so its training signal concentrates on predictable, semantic content such as object shape and layout and skips unpredictable detail such as leaf texture or sensor noise. Yann LeCun proposed the architecture in his 2022 position paper A Path Towards Autonomous Machine Intelligence as a building block for world models, and Meta's I-JEPA and V-JEPA models put it into practice for images and video.

How JEPA Works#

A JEPA has three networks. A context encoder, usually a vision transformer (ViT), embeds the visible patches. A target encoder embeds the full input, and its embeddings at the hidden locations become the targets. A predictor, a narrower transformer, takes the context embeddings plus mask tokens that mark where each target lies and outputs its estimate of the target embeddings. The loss function is the average L2 distance between predicted and actual target embeddings, and only the context encoder and the predictor receive gradients. In I-JEPA, each training image yields four target blocks that each cover 15–20% of the image and one context block that covers 85–100%, with any overlap with the targets removed, so the model must infer large, semantic regions from sparse context.

Avoiding Representation Collapse#

Predicting embeddings opens a shortcut: if both encoders output the same constant vector for every input, the prediction is perfect and the loss is zero, a failure known as representation collapse. I-JEPA and V-JEPA prevent it with asymmetry. No gradient flows into the target encoder, and its weights instead track an exponential moving average (EMA) of the context encoder's weights, with a momentum that starts at 0.996 and rises linearly to 1.0 during I-JEPA training. The same EMA update smooths YOLO26 training in Ultralytics, where the trainer keeps a moving average of the model weights with a decay that ramps toward 0.9999 and saves that averaged copy as the checkpoint. LeJEPA, published in 2025 by Randall Balestriero and LeCun, replaces the EMA teacher and stop-gradient with a regularizer called SIGReg that pushes embeddings toward an isotropic Gaussian distribution.

The PyTorch loop below trains a small JEPA on synthetic vectors, in which a context view and a target view share a hidden 8-dimensional cause.

import copy

import torch
from torch import nn

torch.manual_seed(0)
latent = torch.randn(512, 8)  # hidden cause shared by both views
x, y = latent @ torch.randn(8, 64), latent @ torch.randn(8, 64)  # context view x, target view y

encoder = nn.Sequential(nn.Linear(64, 128), nn.ReLU(), nn.Linear(128, 32))
target_encoder = copy.deepcopy(encoder).requires_grad_(False)  # updated only by EMA
predictor = nn.Linear(32, 32)
optimizer = torch.optim.Adam([*encoder.parameters(), *predictor.parameters()], lr=1e-3)

for step in range(1, 501):
    # Predict the target embedding from the context embedding; raw inputs are never reconstructed
    loss = nn.functional.mse_loss(predictor(encoder(x)), target_encoder(y))
    optimizer.zero_grad()
    loss.backward()
    optimizer.step()
    with torch.no_grad():  # EMA update of the target encoder with momentum 0.996
        for target_param, param in zip(target_encoder.parameters(), encoder.parameters()):
            target_param.lerp_(param, 0.004)
    if step in {1, 500}:
        print(f"step {step}: latent loss {loss.item():.4f}, target embedding std {target_encoder(y).std(0).mean():.3f}")

It prints step 1: latent loss 0.5372, target embedding std 0.547 and step 500: latent loss 0.0060, target embedding std 0.595. The prediction error in embedding space falls about 90-fold while the spread of the target embeddings holds steady, so the loss dropped without the embeddings collapsing to a constant.

JEPA vs. Related Approaches#

• Masked autoencoders (MAE): Mask patches in the same way and reconstruct the missing pixels, which spends model capacity on low-level detail. JEPA predicts the missing content in representation space.
• Contrastive learning: Pulls together the embeddings of two augmented views of an image and pushes apart those of other images, so it needs negative examples and hand-crafted augmentations such as crops and color jitter. I-JEPA uses neither.
• Self-distillation methods such as DINO: Also train a student against an EMA teacher, matching embeddings across augmented crops of the same image. JEPA instead predicts one region's embedding from another region, which teaches spatial and temporal structure.
• Generative world models: Predict future frames in pixel space, as video generators do. A JEPA world model predicts future embeddings and can plan directly in that latent space.

JEPA Models#

• I-JEPA (2023): The image model trained a ViT-Huge/14 on ImageNet with 16 A100 GPUs in under 72 hours, and its features transfer to linear classification, object counting, and depth prediction without hand-crafted data augmentations.
• V-JEPA (2024): Trained on 2 million videos with feature prediction as the only objective, its largest model, a ViT-H/16 kept frozen, reaches 81.9% on Kinetics-400, 72.2% on Something-Something v2, and 77.9% on ImageNet-1k.
• V-JEPA 2 (2025): Pretrained on over 1 million hours of internet video, it reaches 77.3% top-1 accuracy on Something-Something v2 and 39.7 recall-at-5 on Epic-Kitchens-100 action anticipation. Its action-conditioned variant, V-JEPA 2-AC, was post-trained on less than 62 hours of unlabeled robot video from the DROID dataset and then picked and placed objects zero-shot on Franka robot arms in two labs by planning toward image goals.