🤖 machine learning

LpWM: A Case for Sparse Representations in World Models

This paper introduces LpWM, a Joint-embedding predictive architecture that utilizes sparse, non-negative latent representations regularized by Rectified Distribution Matching to significantly reduce predictor complexity and improve planning success compared to traditional dense representations, while simultaneously revealing interpretable, mode-factored dynamical structures.

Original authors: Yilun Kuang, Yash Dagade, Quentin Le Lidec, Lucas Maes, Randall Balestriero, Yann LeCun

Published 2026-08-25

📖 5 min read🧠 Deep dive

CC BY 4.0

Original authors: Yilun Kuang, Yash Dagade, Quentin Le Lidec, Lucas Maes, Randall Balestriero, Yann LeCun

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

To understand how a machine learns to move through the world, imagine teaching a child to navigate a room. The child does not just memorize a static map; they build a mental model of how the room changes when they push a chair or open a door. In the field of artificial intelligence, this mental model is called a "world model." For a computer to plan a path or perform a task, it needs to predict what will happen next based on its current state and the actions it takes. A common approach involves translating the complex visual world into a simpler, internal language of numbers, known as a "latent space." In this hidden space, the computer predicts the future by moving these numbers forward. However, a persistent problem has been that these internal representations often become too dense, with every single number in the set holding some value, making the prediction task incredibly difficult and prone to errors.

Researchers have long wondered if there is a better way to structure this internal language. What if, instead of filling every slot with information, the computer used a sparse language, where most of the numbers are zero and only a few carry the meaning? This idea suggests that by forcing the system to be selective, the rules governing how the world changes might become simpler to learn. This is the central question explored in a new study by a team of researchers from institutions including NYU, Duke University, and Brown University. They set out to test whether this sparse approach could make it easier for an artificial intelligence to understand and control its environment, specifically looking at how the geometry of these internal numbers affects the difficulty of the task.

The team introduced a new system called LpWorldModel, designed to learn these sparse representations. Unlike previous methods that encouraged every part of the internal code to be active, this system uses a specific mathematical constraint to ensure that most of the numbers in the code remain exactly zero. They trained this system on two different environments: a simple navigation task where a dot moves through a doorway, and a more complex manipulation task where a robotic arm must push a T-shaped block to a target location. In the simpler environment, the system worked well regardless of whether the internal code was dense or sparse, because the rules of movement were straightforward enough for even a basic predictor to handle. However, the results changed dramatically in the more difficult pushing task.

When the researchers tested the system with predictors of varying complexity, they found that the sparse approach offered a distinct advantage. In the middle range of complexity, where the predictive model was not powerful enough to handle the dense, cluttered representations, the sparse system succeeded where the dense one failed. Specifically, on the pushing task, the sparse system improved the success rate of planning by up to 57 percent compared to the dense system when using predictors with intermediate capacity. This suggests that by organizing the information into a sparse format, the system reduced the complexity required to model the dynamics of the world. The dense system struggled because it had to learn a complicated web of interactions, whereas the sparse system could rely on a simpler, more direct set of rules.

Beyond just improving performance, the researchers discovered that the sparse representations organized the information in a way that was surprisingly interpretable. The system naturally separated the "mode" of the world from the specific details of the state. In a test environment where the physics changed depending on which zone the agent was in, the system used the presence or absence of specific numbers (the "support") to identify which zone it was in, effectively acting as a switch for the type of physics being applied. The actual values of the non-zero numbers then captured the continuous details, such as the exact position of the agent within that zone. This separation allowed the system to understand that the rules of the world had changed, rather than just seeing a jumble of changing pixels.

In a more complex scenario involving a robotic arm interacting with a cube, the researchers found that the sparse system could track the physical state of the object, such as whether the arm was touching the cube. However, without additional guidance, the system tended to focus on the fastest-moving parts of the scene, like the arm itself, rather than the slower, more significant event of contact. By adding a simple constraint that encouraged the system to be stable over time, they were able to shift the focus so that the sparse code tracked the contact between the arm and the cube. This demonstrated that the sparse representation could be tuned to highlight the most physically meaningful events, turning the internal code into a readable map of the world's dynamics.

The study concludes that sparse representations offer a powerful alternative to the dense methods currently used in artificial intelligence. By forcing the system to be selective, the researchers showed that it is possible to make the task of predicting the future significantly easier, allowing simpler models to achieve better results. The findings suggest that the geometry of the internal language matters deeply; a sparse, structured approach can reveal the underlying rules of a system more clearly than a dense, unstructured one. While the system still requires careful tuning to ensure it tracks the right physical events, the results indicate that sparsity is a promising path toward more efficient and interpretable world models, potentially allowing machines to learn complex behaviors with less computational power.

Technical Summary: LpWM: A Case for Sparse Representations in World Models

Problem Statement

Joint-embedding predictive architectures (JEPAs) have emerged as a viable approach for model-based control from pixels, where an encoder maps observations to a latent space, and a predictor rolls out latent dynamics under actions for planning via Model-Predictive Control (MPC). To prevent representation collapse (degeneration to constant or low-rank subspaces), prior methods employ anti-collapse losses that maximize information, typically matching features to maximum-entropy distributions like isotropic Gaussians. This process yields dense representations, where all latent coordinates are non-zero.

However, it remains unclear whether dense geometry is optimal for modeling controlled dynamics. The central question addressed in this work is whether sparse representations offer a more favorable geometry for modeling action-conditioned latent dynamics, potentially simplifying the predictor required for successful planning and revealing interpretable dynamical structures.

Methodology

Theoretical Motivation: Linearization via Sparsity

The authors propose that sufficiently high-dimensional sparse representations can approximate nonlinear controlled dynamics with linear latent transitions.

Proposition 1: For a compact state space with Lipschitz continuous dynamics, there exists a finite-dimensional one-hot sparse encoding such that the latent dynamics are exactly linear ( $z_{t+1} = P_\epsilon(a_t)z_t$ ). The decoding error is bounded by $(L+1)\epsilon$ , where $L$ is the Lipschitz constant.
Corollary 1: As the encoding dimension $N$ grows, the rollout error vanishes ( $O(N^{-1/d})$ ).
Implication: While exact one-hot codes are impractical to optimize, this motivates distributed sparse representations (where a non-trivial fraction of coordinates are exactly zero) as a learnable relaxation that may simplify the dynamics model.

LpWorldModel (LpWM) Architecture

The paper introduces LpWM, a JEPA model designed to learn non-negative, sparse latent codes.

Encoder ( $f_\theta$ ): A Vision Transformer (ViT) mapping observations to latent embeddings.
Predictor ( $g_\phi$ ): A Transformer with AdaLN-zero action conditioning, predicting the next latent state given the current state and action.
Rectified Distribution Matching Regularization (RDMReg): To prevent collapse and induce sparsity, the model minimizes a loss consisting of a prediction error term and a regularization term $R(z)$ $R (z)$ .
- $R(z)$ matches the marginal distribution of the latent features to a Rectified Generalized Gaussian (RGG) distribution.
- By setting the shape parameter $p=1$ and location $\mu=0$ , the target distribution becomes a Rectified Laplace distribution. This maximum-entropy distribution under $\ell_1$ constraints encourages sparsity.
- The loss uses the 2-Wasserstein distance (via random projections) to match the encoder's output distribution to the target.
Activation Function: Both encoder and predictor projectors use a reparameterized ReLU (RepReLU). In the forward pass, it behaves as standard ReLU (yielding exact zeros); in the backward pass, gradients flow through a GeLU component to mitigate the "dying ReLU" problem.

Planning

At test time, the trained world model serves as the dynamics model for MPC. The system encodes current and goal observations into latents $z_0$ and $z_g$ , then uses the Cross-Entropy Method (CEM) to search for an action sequence that minimizes the distance between the predicted rollout and the goal latent.

Key Contributions

LpWorldModel (LpWM): A JEPA framework utilizing RDMReg to learn non-negative, exactly sparse latent representations for dynamics modeling and planning.
Sparse Representations for Latent Linearization: Formalization of the intuition that high-dimensional sparse codes can approximate nonlinear dynamics with linear transitions, motivating distributed sparse codes as a practical relaxation.
Reduced Predictor Complexity: Empirical evidence that sparsity lowers the predictor capacity required for successful planning. On the PushT environment, sparse LpWM significantly outperforms dense LeWM (using SIGReg) with intermediate-capacity predictors (up to 57% higher success rate), while performing comparably with high-capacity predictors.
Mode-Factored Latent Dynamics: Discovery that the learned sparse representations organize dynamics into interpretable factors:
- Support (Binary): Encodes discrete dynamical regimes (e.g., contact states or spatial zones).
- Magnitudes: Capture continuous within-regime state (e.g., precise position).

Experimental Results

Environments

Experiments were conducted on:

Wall: A 2D navigation task through a doorway (simple dynamics).
PushT: A manipulation task pushing a T-shaped block to a target pose (complex, contact-rich dynamics).
Piecewise: A synthetic 2D navigation environment with known piecewise-affine dynamics (zones with different force fields).
OGBench-Cube: A robot arm manipulating a cube (contact-rich).

Performance Findings

Predictor Capacity vs. Sparsity:
- In the Wall environment, even simple linear predictors (LTI) achieved near-100% success for both dense and sparse models, indicating the dynamics were already simple enough that sparsity offered no additional advantage.
- In the PushT environment, where dynamics are harder to model, sparse LpWM substantially outperformed dense LeWM when using intermediate-capacity predictors (e.g., MLP $\circ$ LTI, MLP $\circ$ LTV). The success rate gap ranged from 24% to 57%.
- With high-capacity predictors (Deep-AdaLN), the performance gap narrowed, suggesting that expressiveness can compensate for dense geometry, but sparsity provides a significant advantage when model capacity is constrained.
Generalization to Other Regularizers: LpWM outperformed a dense JEPA variant using VICReg (another anti-collapse method) across multiple predictor families, indicating the benefit of sparsity extends beyond the specific Gaussian targets used in LeWM.
Interpretability (Piecewise & OGBench-Cube):
- In Piecewise, the binary support of the sparse code decoded the ground-truth zone identity with near-perfect accuracy (94–99%), while magnitudes decoded the continuous agent position. The support Jaccard index remained high within zones and dropped sharply at boundaries, recovering the underlying dynamical regime structure.
- In OGBench-Cube, without additional constraints, the support tended to track fast effector motion. However, adding a Temporal Jaccard (TJ) loss encouraged support stability, causing the support to track slower physical regime changes (e.g., contact events) rather than just motion.

Significance and Claims

The paper claims that sparse representations offer a distinct advantage in the geometry of latent spaces for control:

Simplification of Dynamics: Sparsity can make action-conditioned latent dynamics easier to model, allowing simpler (lower-capacity) predictors to achieve successful planning where they would fail with dense representations.
Interpretability: Sparse codes naturally factorize dynamics into discrete regimes (encoded by the support) and continuous states (encoded by magnitudes), providing a more interpretable structure than dense embeddings.
Practicality: The proposed LpWM framework achieves these benefits without requiring discrete categorical latent spaces (which can be hard to optimize) but rather through continuous, non-negative sparse codes regularized by distribution matching.

The authors conclude that while dense representations work well with sufficiently expressive predictors, sparse representations are a preferable default for reducing predictor complexity and revealing interpretable dynamical structures, particularly in environments with complex, contact-rich, or piecewise dynamics.