LpWM: A Case for Sparse Representations in World Models

Yilun Kuang; Yash Dagade; Quentin Le Lidec; Lucas Maes; Randall Balestriero; Yann LeCun

cs.LGAug 24, 2026

LpWM: A Case for Sparse Representations in World Models

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

Organizations: NYU, AMI Labs · Duke University · Mila, AMI Labs · Brown University, AMI Labs

Abstract

Joint-embedding predictive architectures (JEPAs) learn latent dynamics for planning and avoid representation collapse by matching features to maximum-entropy distributions such as isotropic Gaussians, yielding dense representations. However, it is unclear whether dense representations are the most favorable geometry for modeling dynamics. In this work, we ask whether a different geometry, sparse representations, can make action-conditioned latent dynamics easier to model, and what dynamical structure emerges from such representations. We first show that nonlinear Lipschitz dynamics can be approximated arbitrarily well by action-conditioned linear dynamics in a sufficiently high-dimensional one-hot latent space, with rollout error vanishing as the dimension grows. This motivates distributed sparse representations as a practical relaxation of one-hot sparsity. We introduce LpWorldModel (LpWM), a JEPA model regularized with Rectified Distribution Matching Regularization (RDMReg) to match encoder features to a Rectified Generalized Gaussian distribution, yielding non-negative sparse codes. Empirically, sparsity lowers the predictor complexity required for successful planning: on PushT, sparse LpWM outperforms dense LeWM by up to 57% in planning success at intermediate predictor capacities. This advantage also extends beyond Gaussian distribution matching, with LpWM outperforming dense VICReg representations across multiple predictor families. We further find that the learned sparse representations are mode-factored, with support encoding discrete dynamical regimes and feature magnitudes capturing continuous within-regime state. Together, these results suggest that sparse representations can reduce the predictor complexity required for control while revealing interpretable structure.

Figures & tables

Panel of Figure 1 — Figure 1: LpWorldModel: a JEPA model with sparse representations. (a) An encoder $f$ maps observation $\mathbf{o}_{t}$ to non-negative, sparse latents $\mathbf{z}_{t}$ via a terminal (Rep)ReLU function. The predictor $g$ takes inputs $\mathbf{z}_{t}$ and actions $\mathbf{a}_{t}$ to predict $\hat{\mathbf{z}}_{t+1}$ , compared against the encoded target $\mathbf{z}_{t+1}=f(\mathbf{o}_{t+1})$ . RDMReg ( 24 ) matches per-timestep latent marginals to a Rectified Generalized Gaussian, preventing collapse and inducing sparsity. (b). Arrows indicate the success-rate difference between LpWM and LeWM: $\uparrow$ denotes higher performance for LpWM, while $\downarrow$ denotes lower performance. LpWM outperforms LeWM with more linear predictor variants (MLP $\circ$ LTV $(k)$ , MLP $\circ$ LTI $(k)$ , and LTI $(k)$ . See full specifications in Table 1 ) by up to $57\%$ in success rate, while performing on par with LeWM using the high-capacity DiT predictors. This suggests that sparse representations induce a latent space with simpler dynamics to learn compared to their dense counterparts.

Panel of Figure 2 — Figure 2: The sparse support recovers the piecewise structure of the dynamics. (a) Support Jaccard heatmap over a $20\times 20$ grid: the Jaccard index of each cell’s support with a fixed reference cell (white $+$ ) is high within the reference zone and drops across zone boundaries (dashed), for a model trained with rendered zones (top) and one with no visual cues (bottom). (b) The binary support decodes the ground-truth zone identity as accurately as the full embedding, while agent position is decoded better from the magnitudes than from the support. (c) Planning success vs. horizon (closed-loop, 3-seed mean $\pm$ std). Goal locations of the dot agent are sampled randomly. Both sparse variants beat dense LeWM at every horizon.

Predictor	Layers	History ( $k$ )	Functional form
Deep-AdaLN $(k)$	$6$	$3$	$\hat{\mathbf{z}}_{t+1}=\sigma\!\left(\mathrm{AdaLN}^{(6)}(\mathbf{z}_{t-k+1:t},\,\mathbf{a}_{t})\right)$
Shallow-AdaLN $(k)$	$1$	$3$	$\hat{\mathbf{z}}_{t+1}=\sigma\!\left(\mathrm{AdaLN}^{(1)}(\mathbf{z}_{t-k+1:t},\,\mathbf{a}_{t})\right)$
MLP $\circ$ LTV $(k)$	$1$	$3$	$\hat{\mathbf{z}}_{t+1}=\sigma\!\left(\mathbf{W}\,\mathrm{ReLU}(\textstyle\sum_{i=0}^{k-1}\mathbf{A}_{i}(\mathbf{z}_{t})\,\mathbf{z}_{t-i}+\mathbf{B}(\mathbf{z}_{t})\,\mathbf{a}_{t})\right)$
MLP $\circ$ LTI $(k)$	$1$	$3$	$\hat{\mathbf{z}}_{t+1}=\sigma\!\left(\mathbf{W}\,\mathrm{ReLU}(\textstyle\sum_{i=0}^{k-1}\mathbf{A}_{i}\,\mathbf{z}_{t-i}+\mathbf{B}\,\mathbf{a}_{t})\right)$
LTI $(k)$	$1$	$3$	$\hat{\mathbf{z}}_{t+1}=\sigma\!\left(\textstyle\sum_{i=0}^{k-1}\mathbf{A}_{i}\,\mathbf{z}_{t-i}+\mathbf{B}\,\mathbf{a}_{t}\right)$
LTI $(1)$	$1$	$1$	$\hat{\mathbf{z}}_{t+1}=\sigma\!\left(\mathbf{A}\,\mathbf{z}_{t}+\mathbf{B}\,\mathbf{a}_{t}\right)$

Table 1: Predictors with Different Levels of Complexity.

\hat{\mathbf{z}}_{t+1}

is the predicted latent,

\mathbf{z}_{t-i}

past latents, and

\mathbf{a}_{t}

the action. AdaLN predictors are Transformer-based models with Adaptive Layer Normalization for action conditioning, i.e. DiT architecture ( 30 ) . The Deep-AdaLN

(k)

predictor is the one used in LeWM ( 28 ) . LTI denotes a linear time-invariant predictor with fixed operators

\mathbf{A}_{i},\mathbf{B}

, while LTV denotes a linear time-varying predictor with state-dependent operators

\mathbf{A}_{i}(\mathbf{z}_{t}),\mathbf{B}(\mathbf{z}_{t})

. We use MLP

\circ

to denote one extra nonlinear transform. The history

k

denotes context length, with

1

the single-frame case. The output link is

\sigma=\mathrm{identity}

for dense LeWM and

\sigma=\mathrm{RepReLU}

for sparse LpWM. See Appendix B for details.

Predictor	$D{=}384$	$D{=}768$	$D{=}1536$	$D{=}2048$	$D{=}4096$
Deep-AdaLN $(k)$	$25.8$ M	$62.2$ M	$166.9$ M	$260.2$ M	$822.4$ M
Shallow-AdaLN $(k)$	$5.6$ M	$13.0$ M	$33.1$ M	$50.4$ M	$151.1$ M
MLP $\circ$ LTV $(k)$	$0.81$ M	$3.10$ M	$12.1$ M	$21.4$ M	$84.7$ M
MLP $\circ$ LTI $(k)$	$0.74$ M	$2.95$ M	$11.8$ M	$21.0$ M	$83.9$ M
LTI $(k)$	$0.59$ M	$2.36$ M	$9.44$ M	$16.8$ M	$67.1$ M
LTI $(1)$	$0.30$ M	$1.18$ M	$4.72$ M	$8.39$ M	$33.6$ M

Table 2: Predictor parameter counts across latent dimensions

D

. Counts are for the predictor alone for each

D

used in our experiments. LTI, MLP

\circ

LTI, and MLP

\circ

LTV are attention-free and far smaller than the AdaLN Transformers.

Panel of Figure 3 — Figure 3: Linear predictor is sufficient for the wall environment regardless of dense or sparse embeddings, and LpWM outperforms VICReg-based JEPA across predictor settings in PushT. (a) We plot the planning success rate of the linear LTI $(1)$ predictor against the nonlinear Deep-AdaLN $(k)$ , where each dot denotes one setting of method and feature dimension. For simplicity, we use $\mu=*$ to denote LpWM with the target distribution parameter $\mu$ . We show the open-loop evaluation on top and closed-loop results at bottom. Both are ran using CEM planner. We observe that the closed-loop success rate for LTI $(1)$ predictor using SIGReg is already approaching $100\%$ , indicating that the wall environment can already be described well by a linear dynamics model in the latent space. (b) We have already shown LeWM vs. LpWM in Figure 1(b) . Here we compare LpWM with a variant of LeWM where we replaced SIGReg with VICReg in the PushT environment across predictor settings. As oppposed to SIGReg, we observe consistent improvement from LpWM over VICReg even for the Deep-AdaLN $(k)$ predictor, and the gap between sparse LpWM and dense methods (i.e. VICReg here) continues to hold for MLP $\circ$ LTV $(k)$ and MLP $\circ$ LTI $(k)$ predictors. See full hyperparameter sweeps in Appendix H.2 .

Panel of Figure 4 — Figure 4: Temporal-Jaccard (TJ) makes the cube support track contact, not effector motion. (a) As TJ strength increases, support-instability correlation with effector motion drops ( $0.87{\to}0.40$ ) while cube-motion ( $0.26{\to}0.80$ ) and gripper-contact ( $0.05{\to}0.61$ ) rise—at no change in planning success. (b) One example episode: without TJ (top) support instability tracks effector displacement ( $r_{\mathrm{eff}}{=}0.86$ ); with TJ (bottom) it tracks cube displacement ( $r_{\mathrm{cube}}{=}0.80$ ), spiking inside the shaded contact windows. (c) Support raster: without TJ the support rapidly transitions across frame; with TJ it is smoothly varying in time, reorganizing almost during the shaded contact window.

Appendix figures & tables13 assets

Supplementary material from the paper’s appendix.

Appendix

$D$	Deep-AdaLN $(k)$	Shallow-AdaLN $(k)$	MLP $\circ$ LTV $(k)$	MLP $\circ$ LTI $(k)$	LTI $(k)$	LTI $(1)$
$384$	$0.34$	$0.33$	$0.37$	$0.28$	$0.41$	$0.63$
$768$	$0.44$	$0.37$	$0.54$	$0.53$	$0.41$	$0.63$
$1536$	$0.49$	$0.39$	$0.49$	$0.54$	$0.41$	$0.60$
$2048$	$0.42$	$0.46$	$0.51$	$0.46$	$0.43$	$0.61$
$4096$	$0.48$	$0.49$	$0.47$	$0.43$	$0.44$	$0.56$

Appendix

Table 3: Sparsity of the LpWM encoder representations across predictor types and latent dimensions. Each entry is the fraction of non-zero coordinates

\mathbb{E}[\lVert\mathbf{z}\rVert_{0}]/D

in the sparse code for the target distribution with

\{\mu=0,\sigma=\sqrt{1/2},p=1\}

, measured on the validation set at the tuned-best cell for the LpWM run paired with each predictor at latent dimension

D

. Dense LeWM has active fraction

1.0

by construction. The learned code stays

\sim\!30

–

65\%

active regardless of predictor or dimension, confirming the codes underlying the planning comparison are genuinely sparse.

		Goals from Evaluation Set	Random Goals
Env	Method	$H=R=5$	$H=5,R=1$	$H=R=5$
Piecewise ( $2\times 2$ )	LpWM (Ours)	$\mathbf{99.33{\pm}1.15}$	$\mathbf{84.67{\pm}4.16}$	$59.33{\pm}2.31$
	LpWM+TJ (Ours)	$\mathbf{99.33{\pm}1.15}$	$82.67{\pm}5.03$	$\mathbf{68.67{\pm}4.62}$
	LeWM	$\mathbf{99.33{\pm}1.15}$	$65.33{\pm}4.16$	$36.00{\pm}3.46$
Piecewise ( $3\times 3$ )	LpWM (Ours)	$\mathbf{100.00{\pm}0.00}$	$58.67{\pm}1.15$	$56.00{\pm}4.00$
	LpWM+TJ (Ours)	$\mathbf{100.00{\pm}0.00}$	$\mathbf{60.00{\pm}2.00}$	$\mathbf{58.00{\pm}5.29}$

Appendix

Table 4: Success Rate (%) over 50 Test Samples in the Piecewise Environments with Varying Numbers of Zones . Piecewise (

n\times n

) denotes the environment with varying numbers of zones. Goal images are either sampled from the evaluation dataset or randomly sampled from all possible locations of the dot agent.

H

denotes the horizon and

R

represents the receding-horizon. Results are shown as mean

\pm

std over 3 seeds. Bold denotes the best performance.

Panel of Figure 5 — Figure 5: Environments. (a) Wall: 2D navigation through a doorway (red agent). (b) PushT: push a T-shaped block to a target pose (green). (c) Piecewise: an open room partitioned into force-field zones (colored) with a known ground-truth segmentation. (d) OGBench-Cube: a robot arm manipulating a cube.

	Section 3		Section 4
	Wall	PushT	Piecewise	OGBench-Cube
Encoder	from-scratch ViT, $12$ layers, width $384$		from-scratch ViT-Tiny, width $192$
Predictor	See Table 1		Deep-AdaLN $(k)$
Dimension ( $D$ )	$\{384,768,1536,2048,4096\}$		$192$
History $k$	$1$	$3$	$3$	$3$
Frameskip	$5$	$5$	$5$	$5$

Appendix

Table 5: Experimental configurations for Section 3 and Section 4 . We list out the experimental pipelines in the following table.

Panel of Figure 6 — Figure 6: Deep-AdaLN $(k)$ Predictor. Hyperparameter Sweeps for the PushT Environment.

Panel of Figure 7 — Figure 7: Shallow-AdaLN $(k)$ Predictor. Hyperparameter Sweeps for the PushT Environment.

Panel of Figure 8 — Figure 8: MLP $\circ$ LTV $(k)$ Predictor. Hyperparameter Sweeps for the PushT Environment.

Panel of Figure 9 — Figure 9: MLP $\circ$ LTI $(k)$ Predictor. Hyperparameter Sweeps for the PushT Environment.

Panel of Figure 10 — Figure 10: LTI $(1)$ Predictor. Hyperparameter Sweeps for the PushT Environment.

Panel of Figure 11 — Figure 11: LTI $(k)$ Predictor. Hyperparameter Sweeps for the PushT Environment.

Panel of Figure 12 — Figure 12: Deep-AdaLN $(k)$ Predictor. Hyperparameter Sweeps for VICReg-based JEPA models in the PushT Environment.

Panel of Figure 13 — Figure 13: MLP $\circ$ LTV $(k)$ Predictor. Hyperparameter Sweeps for VICReg-based JEPA models in the PushT Environment.

Panel of Figure 14 — Figure 14: MLP $\circ$ LTI $(k)$ Predictor. Hyperparameter Sweeps for VICReg-based JEPA models in the PushT Environment.

World Model Learning

World model learning builds predictive models of how environments change under actions to support planning, simulation, and decision-making.

See all 419 papers

Anisotropic Representations Improve Planning in JEPA World Models
Sep 29, 2026Mingu Kang, Yoori Oh, Sookyung Kim +1Representation Learning World Model-Based Planning
Seoul National University · Ewha Womans University
LRC-JEPA: Disentangling Dynamics and Residual Context for Efficient World Models
Sep 28, 2026Luzhe Huang, Lei Chu, Jingyi Liang +1Disentangled Representation Learning World Models for Robotics
Independent researcher
Adaptive Latent Capacity for World Models
Sep 26, 2026Idan Achituve, Lior Dikstein, Idit Diamant +2World Model Learning Latent Dynamics Modeling
Arm Research Israel

Sep 29, 2026J/K move · Enter open · S save

Anisotropic Representations Improve Planning in JEPA World Models

Mingu Kang, Yoori Oh, Sookyung Kim, Joonseok Lee

Seoul National University · Ewha Womans University

Representation Learning World Model-Based Planning Action-Conditioned World Models Model-Based Planning Predictive Representation Learning Latent World Models+1

Open paper AskPDF arXiv

Latent world models learn action-conditioned dynamics in representation space and often score candidate actions by Euclidean distance to a goal representation. Joint training typically regularizes the representation to prevent collapse, but the resulting representation geometry also determines how terminal errors are weighted during planning. We show that accurate prediction and noncollapsed representations do not guarantee a task-aligned latent planning cost: isotropic Gaussian regularization can induce a geometry that ranks feasible outcomes differently from the task cost. To address this mismatch, we introduce AnisoWM with

Λ

Reg, which replaces the fixed isotropic Gaussian target with a learnable diagonal covariance under fixed-trace and anisotropy constraints. The prediction objective, predictor architecture, and Euclidean planner remain unchanged; the target is used only during training. Our analysis characterizes the prediction-driven allocation of target variance, its dependence on the training distribution, and the conditions under which the induced metric reduces planning regret. Across four visual control environments, AnisoWM improves planning success over LeWorldModel in all four. Its latent planning cost also shows better agreement with task outcomes. Project website: https://rkdrn79.github.io/AnisoWM-page/

Abstract

World Model Learning

Explore similar work

Anisotropic Representations Improve Planning in JEPA World Models

LRC-JEPA: Disentangling Dynamics and Residual Context for Efficient World Models

Adaptive Latent Capacity for World Models