YC Paper Club Paper 3 / 5 · World Models
  1. 1Speculative²
  2. 2Diffusion-MPC
  3. 3LeWorldModel
  4. 4DL Not Mysterious
  5. 5Infinite Compute

LeWorldModel — Stable JEPA from Pixels

May 30, 2026 · paperjuice · 27 min read · 5300 words advanced

LeWorldModel — A World Model That Refuses to Collapse.

paperjuice ml world-models jepa self-supervised

A world model is a model's internal simulator: a compact representation of how the world works, so an agent can imagine "if I do X, what happens?" without acting in the real world. Yann LeCun has argued for years that the right way to build one is in a latent space — predict the future as an abstract embedding, not as raw pixels — so the model spends its capacity on what matters and ignores unpredictable detail (the exact texture of every leaf).

That family of architectures is JEPA — Joint Embedding Predictive Architecture. The promise is huge. The practice is fragile. JEPAs are notorious for representation collapse: the model discovers it can make the prediction loss zero by mapping everything to the same constant embedding. Predict a constant, match a constant, loss is zero, model is useless.

To dodge collapse, existing JEPAs pile on machinery: multi-term losses, exponential moving average (EMA) target encoders, pre-trained encoders, auxiliary supervision. Lots of moving parts, lots of knobs. "LeWorldModel: Stable End-to-End Joint-Embedding Predictive Architecture from Pixels" by Lucas Maes, Quentin Le Lidec, Damien Scieur, Yann LeCun, and Randall Balestriero (arXiv:2603.19312) strips almost all of it away.

LeWM (LeWorldModel) trains end-to-end from raw pixels with only two loss terms and one tunable hyperparameter (down from six), ~15M parameters, on a single GPU in hours — and still doesn't collapse.

What "collapse" actually is, and why JEPAs invite it

The JEPA objective is: encode the present into an embedding, encode the (masked/future) target into an embedding, and train a predictor to map present-embedding → target-embedding. You measure success by how close the prediction is to the target embedding.

The trap: the target embedding is produced by the network you're training. So the network has two ways to lower the loss. The honest way — learn useful features that make the future predictable. The cheat — make every embedding identical, so prediction is trivial. The cheat is a global optimum with zero loss and zero information. Gradient descent loves it.

optional read — the collapse optimum

Let $f_\theta$ be the encoder and $g_\phi$ the predictor. The naive objective is

$$\mathcal{L} = \big\| g_\phi(f_\theta(x)) - \text{sg}[f_\theta(y)] \big\|^2$$

where $y$ is the target view. If $f_\theta(\cdot) \equiv c$ for any constant $c$, then both sides equal $c$ (with the predictor learning identity) and $\mathcal{L} = 0$. This is the collapse solution. Prior methods break it with stop-gradients plus EMA targets (BYOL-style), variance/covariance penalties (VICReg-style), or contrastive negatives — each adding terms and hyperparameters. LeWM's claim is that a single Gaussian-distribution regularizer on the embeddings is enough to make the constant solution unfavorable, leaving just two terms total: next-embedding prediction + that regularizer.

LeWM's recipe: two terms, one knob

The whole design is a subtraction. Instead of stacking anti-collapse tricks, LeWM uses:

  1. Next-embedding prediction loss — the actual learning signal: predict the future state's embedding from the current one.
  2. Gaussian-distribution regularization — a single term that pushes the embedding distribution toward a Gaussian (spread out, full-rank), which makes the collapse-to-a-constant solution unattractive without needing EMA encoders, negatives, or stop-gradient gymnastics.

That's it. Two terms. And critically, it's end-to-end from raw pixels — no frozen pre-trained vision encoder doing the heavy lifting underneath. The collapse-prevention is built into the objective rather than bolted on as architectural scaffolding.

LeWM — END-TO-END JEPA pixels x_t encoderf_θ predictorg_φ predictedembedding pixels x_{t+1} encoder f_θ targetembedding L1: predict next embedding (match) L2: Gaussian reg(stops collapse) 2 loss terms · 1 hyperparameter · no EMA · no pretrained encoder · no negatives ~15M params · single GPU · trains in hours

Fig 1 — LeWM predicts the next-step embedding and regularizes the embedding distribution to be Gaussian. Those two terms alone prevent collapse, no EMA target or pretrained encoder required.

What it buys you

  • One hyperparameter instead of six. The practical headline. JEPA tuning has been a dark art; LeWM makes it almost configuration-free, which matters far more for adoption than a fractional benchmark gain.
  • Tiny and cheap. ~15M parameters, trainable on a single GPU in hours. World models you can actually iterate on, not ones that need a cluster.
  • 48× faster planning than foundation-model-based approaches. The latent space is compact, so imagining futures and planning over them is cheap — the whole point of doing it in latent space rather than pixels.
  • Competitive control across 2D and 3D tasks despite the simplicity and small size.

Where LeWM sits among world models

"World model" is a broad church. Placing LeWM helps:

  • Pixel-reconstruction world models (Dreamer, world-model RL) — predict future frames, or compress to a latent that's trained to reconstruct pixels. They work, but spend capacity modelling every visual detail, including stuff that doesn't matter for control (textures, lighting). Reconstruction is a heavy, often wasteful objective.
  • MuZero-style — learn a latent dynamics model trained purely to predict reward and value, no reconstruction. Elegant, but tied to a task's reward signal.
  • JEPA (LeWM) — predict future embeddings, self-supervised, no reconstruction and no reward needed. The dream: a task-agnostic world model that learns "how things move" in the abstract. The nightmare: collapse, which is why JEPAs historically needed scaffolding.

LeWM's pitch is to get JEPA's task-agnostic latent prediction without the scaffolding — making the most appealing branch of the tree also the simplest to train.

How everyone else dodges collapse

To appreciate the two-term simplicity, here's the menagerie of anti-collapse tricks LeWM is replacing:

MethodAnti-collapse mechanismCost
BYOLEMA target encoder + predictor + stop-gradientextra network, momentum knob
SimSiamstop-gradient + predictor headfragile, easy to mis-tune
VICRegvariance + invariance + covariance penaltiesthree loss terms, three weights
Contrastive (SimCLR)negative samples push embeddings apartlarge batches / memory bank
LeWMGaussian distribution regularizerone term, one knob

Every row above is a different patch for the same hole: stop the encoder from collapsing to a constant. LeWM argues a single distributional regularizer — pull embeddings toward a spread, full-rank Gaussian — closes the hole directly, so you don't need EMA targets, negatives, or a basket of weighted penalties.

Two flavours of collapse

Worth knowing there isn't one failure but two, and the Gaussian term addresses both:

  • Complete collapse — every input maps to the same point. The pathological zero-loss solution. A variance/spread term forbids it directly.
  • Dimensional collapse — embeddings spread out but only along a few axes; the rest of the latent space is unused (low-rank). Subtler, and it quietly caps how much the model can represent. A full-rank Gaussian target pushes the embedding distribution to use all its dimensions, not just a handful.

Planning in a compact latent

The 48× planning speedup deserves unpacking, because it's the practical reason to want a latent world model at all. Planning means rolling the model forward over candidate action sequences and scoring them. If "rolling forward" means generating high-resolution future frames (as in pixel world models or video-diffusion planners), each imagined step is enormously expensive. If it means stepping a small embedding forward, each step is a tiny matrix op. Same planning algorithm, radically cheaper per rollout — hence dozens of times faster search, which is what makes the model usable inside a control loop rather than just a research artifact.

The "it actually learned physics" evidence

Two results suggest LeWM's latent space captures real structure rather than a degenerate shortcut:

  1. Probing the latent space recovers meaningful physical structure. Train a small probe on the embeddings and you can read off physical quantities — evidence the model encoded the actual state of the world, not arbitrary features.
  2. Surprise detection works. Show the model a physically impossible event and its prediction error spikes — it "notices" when reality violates its learned physics. A collapsed model couldn't do this; it predicts a constant and is never surprised.
Surprise as a diagnostic is elegant: a model that has genuinely learned how the world moves will be wrong (and know it) when the world breaks its own rules. A collapsed model is never surprised because it never predicted anything.

Why this matters beyond JEPA

The broader lesson is a recurring one in ML: a lot of the scaffolding we treat as essential is really there to patch a badly-posed objective. LeCun's group has been pushing JEPA as the path to world models that reason in abstract space; the blocker has always been "but they're a pain to train and they collapse." LeWM is an existence proof that a clean two-term objective can be stable end-to-end from pixels — which makes the whole research direction far more accessible to anyone without a GPU farm.

The caveat: these are still relatively small-scale control tasks. "Stable, simple, and works on 2D/3D control benchmarks" is a real result; "scales to a foundation-level world model" is the next mountain, and this paper doesn't claim to have climbed it. But removing six knobs and the EMA machinery is exactly the kind of simplification that makes climbing it plausible.

arXiv:2603.19312LeWorldModel: Stable End-to-End Joint-Embedding Predictive Architecture from Pixels, Maes, Le Lidec, Scieur, LeCun, Balestriero. Presented at the YC Paper Club.

← prev: Diffusion-MPC next: DL Not Mysterious →
© cvam — written in plaintext, served warm