The idea
Generative world models are trained to reconstruct future observations, which forces them to spend capacity on everything in the frame, including the texture, lighting and noise that no downstream task cares about. A Joint-Embedding Predictive Architecture avoids this by making its predictions in latent space: encode the current observation, predict the embedding of the next one, and never reconstruct a pixel.
The task was two-room navigation: an agent moving through a partitioned environment, 2.5 million frames of trajectories. Success was measured by probing the learned representation for the agent's (x, y) position, a state variable never provided during training. If the representation encodes it, the model has learned something about the world rather than about the renderer.
Architecture
Two design decisions did most of the work.
Separate encoders for agent and wall channels. Walls are static; the agent moves. Feeding both through one CNN forces a single representation to serve two very different jobs, and the static structure, being far easier to predict, dominates the loss. Splitting them into two CNN encoders lets the model keep a stable map of the environment while tracking a moving object against it.
An autoregressive state predictor. The recurrent predictor rolls the latent state forward conditioned on the action taken, so the model is trained to maintain a consistent state estimate over a trajectory rather than a good guess about the next single step.
The collapse problem
Latent-space prediction has a degenerate solution that is always available and always tempting: map every input to the same constant vector. Prediction error goes to zero, the loss is perfectly minimised, and the representation carries no information whatsoever. Left unchecked, this is where joint-embedding models go.
VICReg blocks the shortcut with two explicit terms: a variance term that keeps each embedding dimension above a floor of standard deviation, and a covariance term that decorrelates dimensions so they cannot collude into carrying one signal redundantly. Together they make the constant solution unreachable without any need for negative pairs, momentum encoders or stop-gradients.
In practice this is the part of the build that determines whether the model works at all. The variance and covariance weights are not decoration on the loss; they set whether training converges to a useful representation or to an expensive way of computing a constant. Monitoring embedding variance during training was more informative than watching the prediction loss, which looks healthy right up until it is describing nothing.
Training and results
2.5 million frames does not fit in memory. Data was loaded through memory-mapped NumPy arrays so the operating system paged frames on demand rather than holding the dataset resident, with training run in Singularity containers on HPC. Optimisation used Adam with cosine learning-rate scheduling.
| Item | Value |
|---|---|
| Training frames | 2,500,000 |
| Agent (x, y) prediction error | 1.89 MSE |
| Reconstruction loss terms | 0 |
The probe recovered agent position at 1.89 MSE from a representation trained without ever being shown a coordinate. The model had induced the state variable from the structure of the transitions alone.
What I took from it
- Architectural priors beat capacity. Separating the agent and wall channels was a small change that reflected a real fact about the domain, and it mattered more than making the encoders larger.
- Watch the representation, not the loss. Prediction loss is not a reliable health signal for joint-embedding models, because collapse minimises it. Embedding variance is the diagnostic that actually tells you what is happening.
- Engineering constraints shape the experiment. Memory-mapped loading was not an optimisation detail. Without it, the 2.5M-frame scale that made the result meaningful would have been out of reach.
Project at NYU Courant, March to May 2026, advised by Prof. Yann LeCun. Code on GitHub.