Chapter 81 — JEPA and world models: predicting meaning, not tokens

   context x  (the part you see:            target y  (the part hidden from
   e.g. video frames 1-8, or most            the model: frames 9-12, or the
   of an image)                              masked image blocks)
        |                                          |
   +----v----------+                         +-----v---------+
   | context       |                         | target        |  slow EMA copy
   | encoder       |                         | encoder       |  of the context
   +----+----------+                         +-----+---------+  encoder; NO
        | s_x                                      | s_y         gradients
   +----v----------+                               |
   | predictor     |  <-- optional: action a,      |
   |               |      or latent z              |
   +----+----------+                               |
        | predicted s_y                            |
        +-----------> distance( predicted , s_y ) <+
                         = the loss, measured in REPRESENTATION space

The catch: collapse

There's an obvious cheat. If the encoder maps every input to the same vector, the predictor can always predict it perfectly and the loss is zero. The model has learned nothing. This is called representation collapse, and every JEPA design is largely about preventing it:

• EMA target encoder + stop-gradient (I-JEPA, V-JEPA): the target encoder isn't trained by the loss at all; it's a slowly updated moving average of the context encoder. The cheat path, where both sides collapse together, is cut.
• Regularise the embeddings to stay spread out (the VICReg family): add a loss term that penalises dimensions whose variance drops toward zero. LeJEPA (Balestriero & LeCun, 2025) made this principled with a regulariser that pushes embeddings towards an isotropic Gaussian distribution (SIGReg), removing many of the hand-tuned tricks.

See it happen: a tiny JEPA (~80 lines, runs on a laptop CPU in ~10 s)

Toy world: noisy waves of random frequency. Each 16-step wave is split into 4 patches. The model sees the embeddings of patches 0–2 and must predict the embedding of patch 3. We train three versions: naive (collapses), EMA target (I-JEPA style), and a spread-regulariser (VICReg/LeJEPA style).

"""tiny_jepa.py -- a JEPA small enough to read in one sitting.

Data: 16-step noisy waves. Split each into 4 patches of 4 steps.
Task: from the embeddings of patches 0-2 (context), PREDICT THE EMBEDDING of
patch 3 (target). Never reconstruct the raw numbers.

Run three ways and watch the 'spread' column (std of embeddings):
  naive  -- target encoder = context encoder, gradients flow both ways -> COLLAPSE
  ema    -- target encoder is a slow EMA copy, no gradient through it (I-JEPA/V-JEPA)
  sigreg -- naive + a regulariser that keeps embeddings spread out (VICReg/LeJEPA idea)
Then a linear probe asks: did the embeddings learn the wave's FREQUENCY,
a thing we never trained on?
"""
import copy
import sys

import torch
import torch.nn as nn
import torch.nn.functional as F

torch.manual_seed(0)
D = 16                                   # embedding size


def make_batch(n):
    freq = torch.rand(n, 1) * 2.5 + 0.5                      # hidden cause
    phase = torch.rand(n, 1) * 6.28
    t = torch.arange(16).float().unsqueeze(0)
    x = torch.sin(freq * t * 0.4 + phase) + 0.1 * torch.randn(n, 16)
    return x.view(n, 4, 4), freq.squeeze(1)                 # (n, patches, steps)


encoder = nn.Sequential(nn.Linear(4, 64), nn.GELU(), nn.Linear(64, D))
predictor = nn.Sequential(nn.Linear(3 * D, 64), nn.GELU(), nn.Linear(64, D))


def train(mode, steps=3000):
    enc, pred = copy.deepcopy(encoder), copy.deepcopy(predictor)
    target_enc = copy.deepcopy(enc).requires_grad_(False)
    opt = torch.optim.AdamW(list(enc.parameters()) + list(pred.parameters()), lr=1e-3)
    for step in range(steps + 1):
        x, _ = make_batch(256)
        ctx = enc(x[:, :3]).flatten(1)                      # (n, 3*D)
        guess = pred(ctx)                                   # predicted embedding
        if mode == "ema":
            with torch.no_grad():
                target = target_enc(x[:, 3])                # no gradient here
        else:
            target = enc(x[:, 3])                           # gradient flows: cheat path
        loss = F.mse_loss(guess, target)
        if mode == "sigreg":                                # keep every dim spread out
            z = enc(x[:, 3])
            loss = loss + F.relu(1 - z.std(0)).mean()
        opt.zero_grad(); loss.backward(); opt.step()
        if mode == "ema":                                   # target slowly follows
            with torch.no_grad():
                for pt, pc in zip(target_enc.parameters(), enc.parameters()):
                    pt.mul_(0.99).add_(0.01 * pc)
        if step % 1000 == 0:
            spread = enc(x[:, 3]).std(0).mean().item()
            print(f"  {mode:6s} step {step:4d}  loss {loss.item():.4f}  spread {spread:.4f}")
    return enc


def probe(enc):
    """Linear regression from frozen embeddings to the hidden frequency."""
    def feats(n):
        x, f = make_batch(n)
        with torch.no_grad():
            return enc(x).flatten(1), f
    Xtr, ytr = feats(2000)
    Xte, yte = feats(500)
    Xtr1 = torch.cat([Xtr, torch.ones(len(Xtr), 1)], 1)
    w = torch.linalg.lstsq(Xtr1, ytr.unsqueeze(1)).solution
    pred = torch.cat([Xte, torch.ones(len(Xte), 1)], 1) @ w
    return 1 - ((pred.squeeze() - yte) ** 2).mean() / yte.var()   # R^2


if __name__ == "__main__":
    print(f"random (untrained) encoder probe R^2 = {probe(encoder):.2f}")
    for mode in (sys.argv[1:] or ["naive", "ema", "sigreg"]):
        enc = train(mode)
        print(f"  -> {mode} probe R^2 for frequency = {probe(enc):.2f}\n")

Real output:

random (untrained) encoder probe R^2 = 0.97
  naive  step    0  loss 0.0474  spread 0.1351
  naive  step 1000  loss 0.0000  spread 0.0024
  naive  step 2000  loss 0.0000  spread 0.0014
  naive  step 3000  loss 0.0000  spread 0.0010
  -> naive probe R^2 for frequency = 0.10

  ema    step    0  loss 0.0458  spread 0.1408
  ema    step 1000  loss 0.3189  spread 2.5054
  ema    step 2000  loss 0.4728  spread 3.1958
  ema    step 3000  loss 0.5264  spread 3.2800
  -> ema probe R^2 for frequency = 0.97

  sigreg step    0  loss 0.9050  spread 0.1467
  sigreg step 1000  loss 0.0594  spread 1.0202
  sigreg step 2000  loss 0.0519  spread 1.0973
  sigreg step 3000  loss 0.0385  spread 1.0270
  -> sigreg probe R^2 for frequency = 0.97

What to take from it:

• Naive: the loss goes to exactly 0 and the spread collapses to ~0.001. A perfect loss and a useless model. The embeddings no longer carry the wave's frequency (probe R² 0.10). A falling loss is not evidence of learning in self-supervised setups; always check the spread.
• EMA and the regulariser both prevent collapse, and the embeddings keep the frequency information (R² 0.97). Note that the EMA version's loss rises: its target keeps moving as the encoder learns, so the loss value isn't comparable across methods.
• An honest caveat: at this toy scale, a random encoder also scores R² 0.97, because random features of 4 numbers keep most of the information a linear probe needs. This demo shows the mechanics (prediction in embedding space, collapse, and its cures), not the payoff. The payoff appears on hard, high-dimensional data like video, where V-JEPA-style features beat pixel-reconstruction features on motion understanding by wide margins, and random features are useless.

Planning with a world model

Once you have an action-conditioned predictor, you can plan by imagining, in embedding space, before acting:

   goal   = encoder(image of the desired end state)
   state  = encoder(current camera image)
   repeat every step:
       sample 300 candidate action sequences (say 5 actions each)
       for each: roll the predictor forward 5 times -> imagined final embedding
       score = distance(imagined final embedding, goal)
       keep the best 30, resample around them, repeat a few times   (CEM)
       EXECUTE ONLY THE FIRST ACTION of the best sequence
       observe the real result; re-encode; plan again             (MPC)

This is model-predictive control, an old idea from robotics, now applied in a learned representation space. Because imagined rollouts are just embedding arithmetic, not rendered video, they're cheap enough to run many times per step.

The timeline: from position paper to a billion-dollar bet

Where world models fit (and where JEPA sits among them)

"World model" covers several competing designs in 2026:

• JEPA-style (non-generative): predict embeddings; aimed at understanding and planning. AMI Labs, plus academic groups.
• Generative world models: models that generate the future as video, sometimes interactively. Examples: Google DeepMind's Genie 3 (2025) generating explorable worlds in real time, NVIDIA's Cosmos world foundation models for robotics and autonomous-vehicle simulation, and large video generators described as "world simulators".
• LLMs as world models: the argument that enough text (and images) gives an implicit model of the world. It's partially true, and the reason LLMs can reason about everyday physics at all.

The honest state of play: LLMs remain far ahead for language, code, and abstract reasoning, and agents (Parts 12–20) are built on them. World models are where the action is for robotics, autonomous driving, video understanding, and physical planning: domains where token prediction struggles and where a cheap, imagined rollout is valuable. Many researchers expect hybrids: an LLM for language and high-level plans, a world model for perception and physical consequences. Whether JEPA specifically wins that role is one of the most interesting open bets in the field.

If you've read this far, you've now covered the entire guide — the core path (Parts 1–16) and all five deep-dives (17–21). From here, the Appendices are pure reference material: a glossary, formulas, setup troubleshooting, cheat sheets, and a reading list. Nobody reads them front to back — bookmark this page and come back whenever a term or formula needs a refresher.