An interactive essay

JEPA & World Models

Modern AI is converging on a single bet: that the right way to learn about the world is to predict — but in an abstract embedding space, not in pixels. This essay walks through that idea, lets you poke at every part of it, and watches a tiny model collapse and recover in real time.

Binh Van • May 2026 • Built from notes in /topics & /code
FIGURE 1. Hover the image. The bars on the right are a small hand-built feature vector (mean RGB, brightness, saturation, hue, edge energy) of the pixels under the patch — an illustrative animation that changes as the patch moves over the sun, the sea, the sand. The actual content-aware prediction widget — with real encoder, predictor, and target comparison — is in §3 below.
Section 01

What is a world model?

When you reach for a glass of water without spilling it, your brain is silently running a tiny physics simulation: how heavy the glass is, where your hand needs to go, how the water will slosh. You don't try the motion in reality first — you imagine it, then act. Predictive coding (Rao & Ballard 1999, Friston) hypothesizes the cortex maintains generative models of its sensorimotor input and minimizes prediction error. A world model is the same thing built inside a computer: the part of an agent that says “if I do X, what happens?”.

Formally, a world model is a learned function $$ f_\theta:\; (s_t,\, a_t) \;\to\; s_{t+1} $$ that maps state and action to the next state. The interesting word is “state” — almost nobody works in raw pixel space anymore.

Below is a 2D world with a moving ball. Drag the ball, drag the prediction horizon, toggle the noise. The prediction cone shows where the ball could end up. This is the intuition for why pure pixel-prediction fails over time: long-horizon futures fan out, and forcing the model to pick one pixel-perfect outcome wastes capacity.

World model demo — 2D ball interactive
14
0.06
35°

Three families of world models are being pursued in 2026: predict embeddings (the JEPA family), predict pixels (Sora, Genie, GAIA), and learn a latent dynamics + reward model (the Dreamer line). This essay is mostly about the first — but section 8 puts them side by side.

Section 02 — the centerpiece

The core JEPA idea

JEPA — Joint-Embedding Predictive Architecture — trains by predicting one embedding from another, not by predicting raw pixels. Hide half a photo of a forest. There are billions of valid pixel arrangements that all look like “a forest”. A pixel-predictor has to commit to one and is punished for every wrong leaf. A JEPA punts: “I'll predict an embedding — a compressed summary that says forest, dappled light, late afternoon. That's the part that matters.”

The loss is a regression in embedding space: $$ \mathcal{L}_{\text{JEPA}} \;=\; \big\|\, P_\psi(s_x, z) \;-\; \mathrm{sg}(s_y)\, \big\|_2^2 $$ where $s_x$ is the context embedding, $s_y$ is the target embedding produced by an EMA target encoder (no gradient flows back through it — that's $\mathrm{sg}$), and $z$ is an optional latent that absorbs irreducible uncertainty about the masked region.

Below is the centerpiece interactive of this essay. Pick an image, drag to draw a target block, and watch a real (toy) JEPA-style forward pass: tokenize → encode context → predict target embedding. Toggle the comparison with MAE-style pixel prediction to see what each is committing to.

Context Target Predictor Encoder

JEPA forward pass — draw a target block live math
10×10
0.00
FIGURE 2. A toy JEPA: the “encoder” is a tiny pre-set linear projection of patch pixel statistics (mean RGB, variance, edge energy); the “predictor” is a fixed MLP-like map. It is real arithmetic, but it isn't a trained network — it's chosen so the forward pass is visualizable. The point is the shape of the computation: two encoders, one predictor, an embedding-space loss.

JEPA is not contrastive (no negatives) and is not generative (no pixel decoder). That raises a question: what stops it from learning the trivial constant solution?

Section 03

The self-supervised landscape

JEPA sits in one specific corner of a busy field. Self-supervised learning splits into three big families, with characteristic tricks:

  • Contrastive — SimCLR, MoCo, CPC. Pull two views together; push others apart; use InfoNCE.
  • Non-contrastive joint-embedding — BYOL, DINO, SimSiam, VICReg, Barlow Twins. No negatives; collapse prevented by EMA / stop-grad / explicit regularizers.
  • Masked-prediction — MAE, BEiT, data2vec, I-JEPA. Mask, predict the missing chunk (pixels, discrete tokens, or embeddings).

JEPA borrows the EMA target trick from the non-contrastive family and the masking strategy from the masked-prediction family. Drag the nodes below to rearrange the map; click any method to read its one-liner.

SSL family map — drag the nodesdraggable
Tip: drag a node into a different cloud and the tooltip will explain how that family handles collapse.
Section 04 — runs real training in your browser

Watching collapse, live

The easiest way to win the JEPA loss is to output zero for every patch. Embeddings of context and target are now identical (both zero), loss is zero, and the model has learned nothing. This is representation collapse, the central problem of non-contrastive SSL. Every method has a clever fix.

The widget below trains a tiny 2-layer MLP in your browser on synthetic 2D data (four Gaussian blobs, one per quadrant) with a JEPA-style two-view loss. The two views are small random noise perturbations of the same point. You can switch the anti-collapse mechanism on the fly. Run with no tricks — the scatter implodes to a point and the singular-value spectrum drops to one nonzero value. Enable EMA + stop-grad — it doesn't.

Live collapse demo — tiny MLP, 2D synthetic data real training
0.030
0.980
0Step —Loss —Effective rank
FIGURE 4. Three panels, all live: the embedding scatter (blue = context view, orange = target view), the singular-value spectrum of the embedding matrix, and the loss curve. Effective rank is $\exp(\text{entropy of normalized singular values})$ — full rank is 16, 1 is total collapse.

Stop-grad here is the full SimSiam recipe: stop-grad on the target branch plus a small predictor MLP (16→32→16, ReLU) on the online branch. Without the predictor, stop-grad alone collapses (try the debug toggle). EMA target (BYOL-style, also with a predictor) is the trick I-JEPA inherits. VICReg uses explicit variance + covariance regularizers. SIGReg — the LeJEPA recipe — matches the standard-normal characteristic function at a few projection directions × evaluation points (a toy approximation of the Epps–Pulley test). None is magic: switch back to “none” mid-training and watch the scatter implode.

Debug: — turn this off to see the textbook SimSiam result: stop-grad without a predictor collapses faster than EMA.

interlude · build it yourself

Build it from scratch

You just saw collapse on a slider. Here's the same thing from the other end: four short Python files, ~600 lines total, that rebuild JEPA from nothing → naive → BYOL/I-JEPA → LeJEPA. Three of them run in your browser right now. Click Run on a cell and watch the scatter plot — the rank number is only one piece of the story; the actually-visible signal is whether the four colored clusters stay separated or smear together. Full Python in code/workshop/.

[1]

The simplest thing — and it collapses

Two views, one shared encoder. The loss is MSE directly between the two embeddings. No predictor. No EMA. No stop-grad. The simplest thing that could possibly work — and after ~1k steps the model has squashed its embedding into a low-rank subspace, regardless of what the loss says.

In [1]:
encoder = Encoder()             # 2 → 16 → 16
opt = AdamW(encoder.parameters(), lr=1e-3)

for step in range(2000):
    v1, v2, _ = make_batch()
    s1 = encoder(v1)            # shared encoder
    s2 = encoder(v2)            # gradient flows BOTH branches
    loss = mse(s1, s2)          # MSE directly in embedding space
    loss.backward(); opt.step()
Trains a tiny MLP on 4 clusters in your browser. ~2 s.
Out [1]: training log mode=naive 

                

                  
Embedding scatter (PCA→2D)

                  

                  

                    

                      
Effective rank

                      
— / 16

                      

                    

                  

                

              

            


            

          

        


        
        

          

            
[2]

            
Add an EMA target — and a predictor

          

          

            

              The BYOL/I-JEPA fix: split into online + target
              encoders, EMA-update the target, stop-gradient the target branch,
              and add a small predictor MLP on the online branch (SimSiam recipe).
              Three extra lines, totally different result.
            


            

              
In [2]:


online = Encoder()
target = copy.deepcopy(online)        # EMA target
for p in target.parameters(): p.requires_grad = False
predictor = Predictor()               # sits on the ONLINE branch

@torch.no_grad()
def ema_update(target, online, tau=0.996):
    for tp, op in zip(target.parameters(), online.parameters()):
        tp.data.mul_(tau).add_(op.data, alpha=1 - tau)

for step in range(2000):
    s1 = online(v1)                   # gradient flows
    with torch.no_grad():
        s2 = target(v2)               # stop-grad
    loss = mse(predictor(s1), s2)
    loss.backward(); opt.step()
    ema_update(target, online, tau=0.996)

            


            

              
              Same data, EMA + stop-grad enabled. Watch the 4 clusters come out.
            


            

              

                Out [2]: training log
                mode=ema
              

              

                

                

                  
Embedding scatter (PCA→2D)

                  

                  

                    

                      
Effective rank

                      
— / 16

                      

                    

                  

                

              

            


            

          

        


        
        

          

            
[3]

            
Now do it on images

          

          

            

              Same recipe, real data. Chop a 28×28 MNIST image into 49 patches
              of 4×4. Pick one patch at random to be the target; the
              rest is the context. Encoder takes a single patch and
              outputs an embedding; predictor takes (context summary, target
              position) → predicted target embedding. That's it.
            


            

              
In [3]:


patches = patchify(image)             # (B, 49, 1, 4, 4)
target_idx = torch.randint(0, 49, (B,))
target_patch = patches[arange(B), target_idx]      # (B, 1, 4, 4)
context = patches.clone()
context[arange(B), target_idx] = 0                 # zero out target slot

s_ctx = online(context)                            # (B, 49, 64)
with torch.no_grad():
    s_tgt = target(target_patch.unsqueeze(1)).squeeze(1)  # (B, 64)
s_hat = predictor(s_ctx, position_of(target_idx))  # (B, 64)
loss = mse(s_hat, s_tgt)                           # same as Attempt [2]

            


            

              Too slow to train in a browser tab — but the recipe is identical.
              Full implementation in code/workshop/03_masked_patches.py
              (runs in ~7 min on a laptop CPU).
            


            

              

              

                
~92–95%

                
linear-probe test accuracy on MNIST after 1 epoch (varies by seed)

                
Encoder: 4-layer CNN, EMB=64

                
Patch: 4×4 (49 per image)

                
EMA τ = 0.996

                
Pretrain labels used: 0

                
The encoder learned digit-discriminating features purely from masked-patch prediction.

              

            


            

              Same recipe, real images. Linear probe in the ~92–95% range with zero labels during pretraining. This is I-JEPA — just smaller. The real paper uses ViT-H/14 + 4 large rectangular target blocks + ImageNet, but the loss and the EMA update are identical.
            

          

        


        
        

          

            
[4]

            
Replace the EMA with SIGReg (LeJEPA)

          

          

            

              Back to Attempt [1]'s simplicity — one shared encoder, no
              predictor, no stop-grad, no EMA — but with one regularizer added.
              SIGReg pushes the joint embedding distribution toward an
              isotropic Gaussian via characteristic-function matching on K
              random 1-D projections. Same downstream behavior as Attempt [2],
              fewer moving parts, and now there's a proof.
              (This is a toy SIGReg — 4 fixed t-values instead
              of the Gaussian-weighted Epps–Pulley integral the LeJEPA
              paper uses; same spirit.)
            


            

              
In [4]:


def sigreg(embeds, K=8, ts=(0.5, 1.0, 1.5, 2.0)):
    e = embeds - embeds.mean(0)
    e = e / e.std(0).clamp_min(1e-6)               # standardize per-dim
    u = torch.randn(D, K); u = u / u.norm(dim=0)   # K random unit dirs
    p = e @ u                                      # project → (B, K)
    loss = 0
    for t in ts:
        cos_p = torch.cos(t*p).mean(0)
        sin_p = torch.sin(t*p).mean(0)
        target_re = math.exp(-t*t/2)               # N(0,1) char. func.
        loss += ((cos_p - target_re)**2 + sin_p**2).mean()
    return loss / len(ts)

# Training loop: shared encoder. NO predictor. NO stop-grad. NO EMA.
for step in range(2000):
    s1, s2 = encoder(v1), encoder(v2)
    loss = mse(s1, s2) + 1.0 * sigreg(torch.cat([s1, s2], 0))
    loss.backward(); opt.step()

            


            

              
              Same data again. No EMA. Watch the histogram below match the N(0,1) curve.
            


            

              

                Out [4]: training log
                mode=sigreg
              

              

                

                

                  
Embedding scatter (PCA→2D)

                  

                  
Standardized embed[:,*] vs N(0,1)

                  

                  

                    

                      
Effective rank

                      
— / 16

                      

                    

                  

                

              

            


            

          

        


        
        

          
What we built

          

            
              
#MechanismCluster separationLinear probe

            
            
              
[1]naive JEPA, no tricksarc / squasheddegrades on hard tasks

              
[2]+ EMA target + stop-grad + predictor4 clean clusters~99% on 4 clusters

              
[3][2] applied to MNIST 4×4 patches—~92–95% on 10 MNIST digits

              
[4]replace EMA with SIGReg4 clean clusters~99% on 4 clusters

            
          

          

            The whole arc — naive → BYOL → I-JEPA → LeJEPA — is ~600 lines of Python.
            Bigger encoders, bigger data, smarter masking, action-conditioning — those are quantitative,
            not conceptual. Files: code/workshop/0{1..4}_*.py.
            Narrative companion: topics/21-build-jepa-from-scratch.md.
          

          

            A note on numbers: the live widgets above use a tiny SGD trainer with limited
            steps to stay snappy in your browser, so effective-rank numbers stay around 3–5 / 16
            (the 4-cluster toy can't exercise more than that). The full .py files use
            AdamW for longer and reach higher rank. The visible signal in both is the same:
            naive smears the clusters; EMA and SIGReg keep them separated.
          

        


      

    


    
    

      
Section 05

      
I-JEPA in motion


      

        I-JEPA (Assran et al., CVPR 2023) is the first JEPA that worked at scale on images.
        It uses three ViTs: a context encoder, an EMA target encoder, and a small predictor.
        The masking strategy is unusual: instead of tiny patches, target blocks are 4 random large
        rectangles (~15-20% of the image each), forcing the predictor to commit to semantic
        rather than local prediction.
      


      
Step through the forward pass below. Each step highlights what's happening and what the
      relevant tensor shapes are.

    


    

      

        

          I-JEPA forward pass — six steps
          step-through
        

        

        

          

            
            

          

          
          
          
        

      

    


    
    

      


        
A field guide · plate 01–04

        
Anatomy&How the parts fit together

        

          The same architecture, drawn the way an old service manual would draw it.
          Four cutaway plates of the JEPA machine — pull the parts apart, watch the
          spring tension on the EMA teacher, follow a single image patch through the
          token matrix, and tilt the energy landscape until collapse is no longer a valley.
          Every diagram below is custom SVG; every slider drives the picture in real time.
        


        
        
Plate i. Exploded view


        

          

            
Fig. 01 Anatomy of a JEPA — context · target · predictor

            
Pull to assemble

          

          

            

          

          

            Context branch
            Target branch (EMA)
            Predictor
            EMA flow τ ≈ 0.996
            L2 in embedding space
          

          

            Plate 01.
            Three sub-assemblies. The context encoder (left, navy) and the target encoder (right, ochre)
            are weight-tied ViTs — except the target receives no gradient
            
            and updates only through an exponential moving average of the context encoder's weights.
            The predictor (small, rust) is a thin MLP that maps the context tokens plus the target's
            position → a guess at the target embedding. Drag the assembly slider; the
            three modules slide together along their alignment guides until the L2 loss node is fed.
          

          

            

              
              

                
                exploded
              

            

            

              
              

                
              

            

            

              
              

                
              

            

          

        


        
        
Plate ii. Patchify


        

          

            
Fig. 02 Standard ViT patch tokenization (generic pipeline)

            
Hover any patch

          

          

            

          

          

            Plate 02.
            Before anything else, a Vision Transformer chops the image into a grid of 16×16 pixel patches.
            Each patch (16·16·3 = 768 numbers) is flattened, multiplied by a learned linear projection,
            added to a positional embedding indexed by its grid position, and dropped into row
            i of an N×D token matrix.
            Hover any cell of the image to highlight its corresponding row — that row is the patch's
            life-long identity inside the network.
            


            Note on [CLS]: in classification-style ViTs (and at probe time)
            a learned [CLS] token is prepended at row 0 to act as a sequence-summary slot.
            I-JEPA / V-JEPA context encoders operate on patch tokens only — no learned
            [CLS] in the input sequence. We show the CLS row greyed-out to mark it as
            optional / classification-ViT only.
          

        


        
        
Plate iii. The damping mechanism


        

          

            
Fig. 03 Inside the EMA target encoder — spring & pulley

            
Drag the online weight

          

          

            

          

          

            Plate 03.
            Think of the target encoder as a heavy weight on a track, attached to the
            online encoder by a spring. As gradient descent yanks the online weight around,
            the spring drags the target along — but slowly. The constant τ is the
            smoothing factor / inertia: at τ = 0.996 the target only moves 0.4% of the way
            per step (so high τ = slow target = low effective spring stiffness;
            the “step toward online” weight is 1−τ).
            That lag is what stops the predictor from finding the trivial constant solution
            — you cannot trivially copy a slow-moving teacher when its update rule depends on you.
            Drag the online block; the spring stretches, then the target follows.
          

          

            

              
              

                
                0.996
              

            

            

              
              

                
                +0.60
              

            

            

              
              

                
                
              

            

          

        


        
        
Plate iv. The collapse landscape


        

          

            
Fig. 04 The collapse landscape — valleys lifted by regularizers

            
Three knobs reshape the floor

          

          

            

          

          

            Plate 04.
            A schematic / hand-tuned landscape illustrating how a regularizer can lift the
            central collapse valley — not the actual loss surface of a real JEPA.
            The deep central valley represents trivial collapse (every embedding equal to zero,
            loss equal to zero, knowledge equal to zero); the outer valleys are
            useful full-rank solutions.
            Each regularizer is drawn as a force that lifts the central valley.
            


            SIGReg vs EMA / stop-grad: these are mechanistically different and would deform
            the landscape in different ways. This widget only shows SIGReg's effect; EMA + stop-grad
            does not literally add a term to a loss surface in the same closed-form way.
          

          

            

              
              

                
                0.00
              

            

            

              
              

                
                0.00
              

            

            

              
              

                
                0.00
              

            

            

              
              

                
              

            

          

        


      

    


    
    

      
Section 06

      
Planning in latent space (V-JEPA 2 style)


      

        V-JEPA 2 (June 2025) does two things: pretrain on 1M+ hours of video with a JEPA objective,
        then train a small action-conditioned predictor on 62 hours of robot data.
        At inference time, give the robot a goal image, encode it, and plan a short action
        sequence by optimizing $a_{t:t+H}$ to minimize $\|\hat{s}_{t+H} - s_{\text{goal}}\|$
        through the predictor — model-predictive control
        with cross-entropy method in latent space.
      


      

        The widget below is that exact loop on a 2D point-mass world. Click anywhere to set the goal.
        The planner samples thousands of candidate action sequences in the model's latent space,
        keeps the elites, and executes the first action. The grey rectangle is a wall.
        It is the same code as code/toy_world_model.py, ported to JavaScript.
      

      

        Honest disclosure. This toy substitutes the true 2D environment dynamics for the
        learned latent model — the educational point is the CEM optimization loop itself.
        V-JEPA 2 plans by rolling out a learned embedding-space predictor instead, and the
        goal is a goal image encoded into the same embedding space (not a $(x,y)$ coordinate).
      

    


    

      

        

          Latent-space CEM planner — 2D point-mass
          live math
        

        

        

          

            
            
10

          

          

            
            
128

          

          

            
            
4

          

          

            
            

              
              
            

          

          
          
          
          0Step
          —Dist to goal
        

      

    


    

      

        This planner has no reward function, no policy network. The whole loop is:
        encode the current state, encode the goal image, sample latent action rollouts,
        keep the ones that end closest to the goal embedding, refit the action distribution to those elites.
        That's it. LeCun's vision in toy form.
      

    


    
    

      
Section 07

      
Three paradigms, side by side


      
Click a card to highlight its row in the comparison table and play its mini-animation.

    


    

      

      

        
PropertyPredict embeddings (JEPA)Predict pixels (Sora / Genie)Latent dynamics + reward (Dreamer)

        
          
What's predictedtarget embeddingnext-frame pixels or tokensnext latent + reward

          
Loss$\|\hat s - s\|^2$ in embedding spacediffusion or AR loss on pixels/tokensELBO: recon + KL + reward + continue

          
Decoder?noyes (huge)yes (small)

          
Reward built in?no (goal-image cost)no (separate)yes

          
Pretrain costlow–med (1B params, ~1k GPU-days)very high (10–50B params, >10k GPU-days)per-env (small)

          
Best atperception, planning, transfersynthesis, sim datasample-efficient RL

        
      

    


    
    

      
Section 08

      
The world-model races

      
A non-exhaustive timeline of the last eight years. Hover the cards for details.

    


    

      
Timeline 2018–2026hoverable

        

      

    


    

      
Who is building what?

      
Each cell links a lab to its bet across paradigms. Hover for the names of the systems.

    


    


    
    

      
Section 09

      
Open problems

      

        Even after V-JEPA 2, Genie 3, GAIA-3, π0, DreamerV4 — the field is stuck on
        a small number of hard, well-known problems. Spin the wheel.
      

    


    

      

        
Spin the wheelclick anywhere

        

          

            

            

              
Click the wheel or hit spin →

              
Each segment is one of the field's open frontiers.

            

          

        

        

          
        

      

    


    
    

      
Section 10

      
Teaching ladder

      

        The same idea — predict in embedding space — explained at five ability levels.
        Drag the slider.
      

    


    

      

        
One idea, five levelsslider

        

        

          

            
            

          

          High-schooler → undergrad → grad → PhD → researcher
        

      

    


    
    

      
References & further reading

      
Starred items are the must-reads.

      
    
        
  1. ★ LeCun (2022). A Path Towards Autonomous Machine Intelligence.
            OpenReview · The JEPA blueprint.
    2. 
        
  2. ★ Assran et al. (CVPR 2023). Self-Supervised Learning from Images with a Joint-Embedding
            Predictive Architecture (I-JEPA). arXiv:2301.08243.
    3. 
        
  3. ★ Bardes et al. (2024). Revisiting Feature Prediction for Learning Visual Representations
            from Video (V-JEPA). arXiv:2404.08471.
    4. 
        
  4. ★ Assran et al. (June 2025). V-JEPA 2: Self-Supervised Video Models Enable Understanding,
            Prediction and Planning. arXiv:2506.09985.
    5. 
        
  5. ★ Balestriero & LeCun (Nov 2025). LeJEPA: Provable and Scalable Self-Supervised
            Learning Without the Heuristics. arXiv:2511.08544.
    6. 
        
  6. ★ Hafner et al. (Nature 2025). Mastering Diverse Domains through World Models (DreamerV3).
            arXiv:2301.04104.
    7. 
        
  7. Ha & Schmidhuber (2018). World Models.
            arXiv:1803.10122.
    8. 
        
  8. Bruce et al. (ICML 2024 best paper). Genie: Generative Interactive Environments.
            arXiv:2402.15391.
    9. 
        
  9. He et al. (2022). Masked Autoencoders Are Scalable Vision Learners (MAE).
            arXiv:2111.06377.
    10. 
        
  10. Grill et al. (2020). Bootstrap Your Own Latent (BYOL).
            arXiv:2006.07733.
    11. 
        
  11. Caron et al. (2021). Emerging Properties in Self-Supervised Vision Transformers (DINO).
            arXiv:2104.14294.
    12. 
        
  12. Bardes, Ponce, LeCun (2022). VICReg.
            arXiv:2105.04906.
    13. 
        
  13. Hafner et al. (2019). PlaNet: Learning Latent Dynamics for Planning from Pixels.
            arXiv:1811.04551.
    14. 
        
  14. Tian, Chen, Ganguli (2021). Understanding Self-Supervised Learning Dynamics without
            Contrastive Pairs. arXiv:2102.06810.
    15. 
        
  15. OpenAI (Feb 2024). Sora — video generation models as world simulators.
            openai.com.
    16. 
        
  16. DeepMind (Aug 2025). Genie 3: a new frontier for world models.
            deepmind.google.
    17. 
        
  17. Wayve (Dec 2025). GAIA-3.
            wayve.ai.
    18. 
        
  18. Physical Intelligence (Oct 2024). π0: our first generalist policy.
            pi.website.
    19. 
        
  19. NVIDIA (Jan 2025). Cosmos World Foundation Model Platform for Physical AI.
            research.nvidia.com.
    20. 
        
  20. Survey: Understanding World or Predicting Future? A Comprehensive Survey of World Models,
            arXiv:2411.14499 (Dec 2025).
    21.