PaperlayerView exact revision on arXiv

Source-linked AI summary

LpWM: A Case for Sparse Representations in World Models

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

arXiv:2608.22764v1cs.LG

TL;DR

It is unclear whether dense features are the most favorable geometry for modeling controlled dynamics. LpWM uses RDMReg to learn sparse codes, and on PushT it outperforms dense LeWM by up to 57% in success rate with intermediate-capacity predictors.

  • Problem

    It remains unclear whether dense features are the most favorable geometry for modeling controlled dynamics.

    Source: §1 · p. 2

  • Method

    LpWM uses RDMReg to match rectified encoder features to Rectified Generalized Gaussian targets, producing non-negative sparse codes.

    Source: §1 · p. 2

  • Results

    On PushT, sparse LpWM outperforms dense LeWM by up to 57% in success rate with intermediate-capacity predictors.

    Source: §1 · p. 3

  • Takeaways & Limitations

    Sparse representations can reduce predictor complexity for control while yielding interpretable dynamics whose support encodes discrete regimes and magnitudes capture within-regime state.

    Source: §1 · p. 2

    Vanilla LpWM sparsity alone may not make support transitions semantically meaningful, and without temporal structure support can behave largely as a motion detector.

    Source: §4.2 · p. 8

Abstract

from arXiv · show

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.

1 Introduction

The introduction asks whether sparse rather than dense latent geometry can make action-conditioned dynamics easier to model. It presents LpWM as a JEPA using sparse codes and reports simpler-to-learn dynamics and interpretable mode structure.

Empirical motivation

  • Up to 57% in success rate: sparse LpWM outperforms dense LeWM with intermediate-capacity predictors, while performing on par with LeWM using high-capacity DiT predictors.The introduction also reports that predictors failing over dense codes succeed over sparse ones.Source: Figure 1 · p. 2, §1 · pp. 2–3

Sparse representations for latent linearization

  • Sufficiently high-dimensional one-hot codes approximate Lipschitz controlled dynamics with action-conditioned linear latent dynamics, with rollout error vanishing as dimension grows.This idealized result motivates distributed sparse codes as a learnable relaxation of one-hot representations.Source: §1 · p. 2

LpWorldModel

  • LpWM uses RDMReg [Kuang et al., 2026b] to match rectified features to Rectified Generalized Gaussian targets and learn non-negative sparse codes.The targets arise from Generalized Gaussian distributions that are maximum-entropy under expected ℓp-norm constraints for p ≤1.Source: §1 · p. 2, Figure 1 · p. 2

Mode-factored latent dynamics

  • Sparse representations are mode-factored: their support encodes discrete dynamical regimes, while feature magnitudes capture continuous within-regime state.The introduction presents this organization as an interpretable structure of learned dynamics.Source: §1 · pp. 2–3

2 Method

The method pairs action-conditioned JEPA prediction with RDMReg and rectification, then uses learned latent rollouts for goal-directed MPC.

2.1 LpWM: JEPA with Rectified Distribution Matching Regularization (RDMReg)

  • LpWM jointly trains an action-conditioned JEPA with RDMReg anti-collapse regularization, matching features to a rectified generalized Gaussian via 2-Wasserstein distance; defaults yield rectified Laplace [Kuang et al., 2026b].With µ=0, σ=1, p=2 and no ReLU, the target recovers the dense isotropic Gaussian used in LeWM.Source: §2.1 · p. 3

2.2 Architecture and Rectification

  • The ViT encoder and AdaLN-zero-conditioned Transformer predictor use three-layer MLP projectors; RepReLU yields exact-zero forward outputs while routing gradients through GeLU.RDMReg also trains without collapse using standard ReLU; RepReLU is an optimization safeguard against dead neurons, not essential for collapse prevention.Source: §2.2 · p. 3

2.3 Planning with the Learned World Model

  • At test time, MPC uses CEM to find an action sequence minimizing the latent rollout goal cost C = ∥ẑT − zg∥2 [Rubinstein, 1999, Zhou et al., 2025].The current and goal images are encoded to z0 and zg, respectively.Source: §2.3 · p. 3

3 Does Sparsity Simplify Latent Dynamics?

The section tests whether sparse latents reduce the predictor capacity needed for planning. It shows that one-hot sparse codes can make Lipschitz dynamics exactly linear, while learned sparse codes outperform dense baselines on PushT at intermediate predictor capacities.

3.1 Linearization of Nonlinear Dynamics with Sparse Representations

  • For compact controlled systems with uniformly state-Lipschitz dynamics, finite one-hot codes yield exactly linear action-conditioned latent transitions and arbitrarily accurate decoded one-step predictions; rollout error vanishes as dimension grows.One-step approximation error decreases as O(N^-1/d), suffering from the curse of dimensionality; this motivates distributed sparse codes as a practical relaxation.Source: §3.1 · pp. 4–5

3.2 Experiment Designs

  • Experiments compare sparse LpWM and dense LeWM across predictor capacities and latent dimensions 384–4096 on Wall and PushT, with a VICReg-based dense LeWM variant as an additional baseline.They use maximum update parameterization to stabilize training across widths.Source: §3.2 · p. 5, Table 1 · p. 5

3.3 Results

  • Figure 3 shows LpWM outperforming VICReg across Deep-AdaLN(k), MLP◦LTV(k), and MLP◦LTI(k) predictors, including the high-capacity Deep-AdaLN(k) predictor.This extends the sparsity benefit beyond comparison with LeWM’s isotropic Gaussian targets; the authors hypothesize that VICReg’s second-order moment constraints leave the target distribution underspecified.Source: Figure 3 · p. 6, §3.3 · p. 7
  • Wall’s LTI(1) predictor reaches nearly 100% closed-loop CEM success for both methods, whereas PushT’s lowest-capacity LTI(1) fails for both and sparse gains diminish at highest capacities.The authors conclude that sparsity’s benefit depends on predictor capacity relative to dynamics complexity, and that gains may extend to more expressive predictors as dynamics become more complex.Source: §3.3 · pp. 5–7, Figure 3 · p. 6
  • On PushT, sparse LpWM improves planning success over dense LeWM by 24%–57% with MLP◦LTI(k), 36%–45% with MLP◦LTV(k), and 11%–23% with LTI(k) at intermediate capacity.The performance gaps were verified across extensive hyperparameter sweeps.Source: §3.3 · p. 6

4 Does Sparsity Lead to Interpretable Latent Dynamics?

The section tests whether sparse codes become mode-factored, with support identifying discrete regimes and feature magnitudes capturing within-regime state. This structure emerges in Piecewise; on OGBench-Cube, tracking contact-related regime changes requires temporal regularization.

4.1 Discovering Piecewise-Affine Dynamics

  • 84.7% versus 65.3% success rate: LpWM outperforms LeWM on Piecewise random goals, while both methods saturate on goals sampled from evaluation episodes.The comparison is reported for planning success rate.Source: §4.1 · p. 7
  • In Piecewise, support overlap stays high within zones and drops at boundaries—even without visual cues—showing that action-conditioned prediction recovers dynamics modes.The pattern indicates that active features organize around the underlying regime rather than appearance alone.Source: §4.1 · p. 8
  • Piecewise support decodes the ground-truth zone at near-perfect accuracy, while embedding magnitudes decode continuous agent position better.The support performs as accurately as the full embedding for zone decoding, separating discrete mode from within-zone state.Source: §4.1 · p. 8
  • Sparse LpWM outperforms dense LeWM at every Piecewise planning horizon, with the gap widening as horizon H increases.Source: §4.1 · p. 8

4.2 From Navigation to Contact-Rich Dynamics

  • Without temporal regularization, support instability tracks effector motion (r ≈0.87) more than contact (r ≈0.05), but TJ raises cube-motion correlation from rcube = 0.21 to 0.80.This suggests sparsity alone may not make support transitions semantically meaningful; with a temporal prior, sparse support can track physical regime changes such as object interaction and contact.Source: §4.2 · p. 8

5 Conclusion · A Related Work

LpWM uses RDMReg to learn exactly sparse, non-negative latents, where a shallow predictor plans on PushT at intermediate capacity when it fails on dense codes. The paper also reports mode-factored sparse codes and places LpWM within JEPA world-model research.

5 Conclusion

  • At intermediate predictor capacity on PushT, a shallow predictor plans over sparse LpWM codes but fails over dense codes.The conclusion presents this as evidence that sparsity simplifies action-conditioned latent dynamics.Source: §5 · p. 8
  • LpWM combines action-conditioned JEPA with RDMReg to produce non-negative, exactly sparse latents.The method is motivated by a one-hot linearization result.Source: §5 · p. 8
  • Sparse-code support encodes discrete dynamics regimes, decodable at 94–99% on Piecewise; with a temporal-Jaccard prior, it tracks contact on OGBench-Cube.The conclusion characterizes these representations as mode-factored, with interpretable dynamics.Source: §5 · p. 8
  • LpWM’s sparse representations can reduce predictor complexity for control while yielding more interpretable dynamics.Source: §5 · p. 8

A Related Work

D Deep-AdaLN(k) Shallow-AdaLN(k) MLP◦LTV(k) MLP◦LTI(k) LTI(k) LTI(1) … C.3 Visualization of Environments

The supplied sections position sparse continuous codes against prior latent-regularization and discrete-state approaches, describe predictor designs, and specify the environments used to study them. They also distinguish fixed linear predictors from a structured state-dependent variant.

D Deep-AdaLN(k) Shallow-AdaLN(k) MLP◦LTV(k) MLP◦LTI(k) LTI(k) LTI(1)

  • JEPA regularization addresses collapse and planning geometry through stop-gradient and EMA methods, information-maximization and distribution matching, among other approaches.The passage also names DINO-WM and V-JEPA 2 AC as using stop-gradient and EMA techniques.Source: D Deep-AdaLN(k) Shallow-AdaLN(k)… · p. 12
  • LpWM studies distributed, non-negative sparse codes whose support can reflect dynamical regimes while magnitudes retain continuous state, unlike fixed categorical groupings [Hafner et al., 2022] [Jain et al., 2022].The paper asks whether this geometry reduces action-conditioned transition-model complexity.Source: D Deep-AdaLN(k) Shallow-AdaLN(k)… · p. 12

B.1 Overall Designs

  • The compared predictors range from six- and one-block AdaLN Transformers to linear time-invariant maps and an MLP readout over a linear core.LpWM uses a RepReLU link with LTI(k) and LTI(1), so those are not exact linear dynamics models.Source: §B.1 · p. 13

B.2 MLP◦LTV(k)

  • MLP◦LTV(k) uses a state-emitted gate to modulate fixed low-rank corrections, costing O(Dr) rather than requiring an O(D3) matrix-valued map.The gate is shared across history lags and the action term, preserving an affine action input; r=16 in the experiments.Source: §B.2 · p. 13

C.1 TwoRoom, PushT, and OGBench Cube

  • TwoRoom, PushT, and OGBench Cube use saved datasets; planning selects evaluation episodes, sets goals 25 steps away, and uses five-step action blocks.Figure 5 depicts the Wall, PushT, Piecewise, and OGBench-Cube environments.Source: §C.1 · p. 13, Figure 5 · p. 14

C.2 Piecewise

  • Piecewise divides an open room into 4 or 9 zones with known, locally linear but globally piecewise-affine motion, enabling direct tests of whether sparse support recovers modes.An exact expert policy inverts the motion equation.Source: §C.2 · p. 14

C.3 Visualization of Environments

  • The environment visualizations include a Piecewise version without background-color rendering, as shown in Figure 2a.Figure 5 provides example environment snapshots.Source: §C.3 · p. 14

D Experimental Details … F Maximal-Update Parameterization (µP) for JEPA

The appendix specifies two complementary LpWM experimental pipelines, support-similarity and temporal-smoothness measures, and width scaling used to stabilize feature learning. It also notes evaluation limits and that µP learning-rate transfer is not universal.

D Experimental Details

  • The Section 3 complexity and Section 4 interpretability studies use separate pipelines for different environments and configurations, while sharing the same action-conditioned JEPA and distribution-matching formulation.Dense LeWM uses identity links with an isotropic Gaussian target (p = 2); sparse LpWM uses (Rep)ReLU links with a Rectified Laplace target (p = 1).Source: D Experimental Details · p. 14
  • Section 3 uses k = 1 for Wall and k = 3 for PushT, while Section 4 fixes D = 192 and uses a depth-6 predictor with history k = 3.Section 3 sweeps the RDMReg weight and µP base learning rate for each predictor and latent dimension, reporting the best-performing grid cell; Section 4 trains for 10 epochs and does not use µP.Source: D Experimental Details · pp. 14–15
  • Both pipelines plan with CEM over latent-rollout MSE, using 300 candidates, retaining 30, and running 30 iterations with horizon H = 5.With frameskip 5, the goal lies 25 raw environment steps ahead; Section 3 reports open-loop success and closed-loop MPC success, evaluating Wall and PushT over 50 trajectories.Source: D Experimental Details · p. 15
  • Section 4 tests Piecewise with randomly sampled goal images because evaluation-set goals saturate performance, while OGBench-Cube’s 25-step goal horizon may limit benchmark discrimination.The authors leave longer-horizon and larger-frameskip evaluation to future work.Source: D Experimental Details · p. 15

E.1 Hard and Soft Jaccard index

  • Because nonnegative encoder features are not binary, the appendix replaces nondifferentiable set Jaccard with a soft relaxation that uses min and max and recovers exact Jaccard for binary inputs.The soft index measures support similarity; the hard set index is the intersection size divided by the union size.Source: §E.1 · pp. 15–16

E.2 Temporal Jaccard Loss

  • The optional temporal Jaccard loss penalizes support changes between consecutive frames, encouraging smoothly varying support without interfering with RDMReg’s global sparsity constraints.It minimizes mean support instability across consecutive frame pairs in each batch and sequence.Source: §E.2 · p. 16

F Maximal-Update Parameterization (µP) for JEPA

  • To compare latent widths from D = 384 to D = 4096 without width-specific optimization confounds, the study uses µP scaling for stable feature learning [Yang et al., 2022].For width-dependent dense maps, the reference width is D0 = 384; µP applies fan-in-based learning-rate and initialization scaling [Yang et al., 2024].Source: F Maximal-Update Parameterization (µP)… · p. 16
  • The authors use µP to maintain stable feature learning, not to reduce hyperparameter tuning, and report that learning-rate transfer across widths works most of the time but has exceptions.They note µP had not, to their best knowledge, previously been demonstrated for hyperparameter transfer across JEPA widths and leave further investigation to future work.Source: F Maximal-Update Parameterization (µP)… · p. 16

G Proofs of Linearization · G.1 Proof of Proposition 1

The proof constructs an action-conditioned linear latent model from a finite quantization of the compact state space, then bounds its decoded one-step and finite-horizon rollout errors using the Lipschitz assumption. The errors vanish as quantization becomes finer.

G.1 Proof of Proposition 1

  • Compactness gives a finite set of ε-covering centers for X, enabling a deterministic quantizer that assigns each state to a center.The proof uses this finite cover to define the discrete latent state space.Source: §G.1 · p. 17
  • Mapping each quantized state to its center’s standard basis vector yields a one-hot encoder, with the decoder mapping each basis vector back to its center.Every encoded state has exactly one nonzero coordinate.Source: §G.1 · p. 17
  • For each action, mapping each center to the quantized successor defines a transition matrix with one nonzero entry per column, so the latent transition is linear in z.The action-specific map sends each basis vector to the basis vector for its successor center.Source: §G.1 · p. 17
  • The decoded one-step error is uniformly bounded using Lipschitz continuity and the ε-quantization error, and the bound tends to zero as ε → 0.The proof uses that the Lipschitz constant L is fixed, so (L + 1)ε → 0.Source: §G.1 · p. 17
  • For the finite-horizon rollout, the initial decoded error is at most ε and each step satisfies e_{t+1} ≤ L e_t + ε.Here e_t denotes the distance between the ground-truth state and the decoded state at time t.Source: §G.1 · p. 18
  • Thus, for any fixed finite horizon, recursively propagated rollout error also vanishes as ε → 0.This follows from the finite-horizon bound under the same assumptions as Proposition 1.Source: §G.1 · p. 18

G.2 Proof of Corollary 1 · H Hyperparameter Sweeps

The proof partitions the unit cube into equal cells, bounds prediction error using their covering geometry, and applies one-step and finite-horizon bounds to obtain vanishing rollout error as the partition grows. The supplied passages contain no hyperparameter-sweep findings.

G.2 Proof of Corollary 1

  • Partitioning each coordinate interval into n equal subintervals creates n^d axis-aligned hypercubes of side length h.Source: §G.2 · p. 18
  • Within each hypercube, every coordinate of a point differs from the cell center by at most h/2.Source: §G.2 · p. 18
  • The proof then expresses the resulting distance bound using the Euclidean norm.Source: §G.2 · p. 18
  • For fixed d and L, Proposition 1 bounds the one-step decoded prediction error using the partition’s covering radius.Source: §G.2 · pp. 18–19
  • With H, d, and L fixed, the proof concludes that the finite-horizon rollout error vanishes as the partition grows.The bound follows after substituting the covering radius into the finite-horizon bound from Lemma 1.Source: §G.2 · p. 19

H.1 LpWM and LeWM Sweeps in the PushT Environment

The PushT sweeps report the hyperparameter search underlying Figure 1b, with later sweeps narrowed to save compute after high learning rates repeatedly performed poorly.

H.1 LpWM and LeWM Sweeps in the PushT Environment

  • For each fixed feature dimension d ∈{384, 768, 1536, 2048, 4096}, the grid varied λRDMReg ∈{0.01, 0.1, 1.0, 10.0}, ℓmuP ∈{0.05, 0.005, 0.0005, 0.00005}, method, and predictor setting.The compared methods were LeWM and LpWM with target distribution Qd^(1/2).Source: §H.1 · p. 19
  • For isotropic Gaussian targets, the sweeps used sliced Wasserstein distance rather than SIGReg’s Epps-Pulley loss [Balestriero and LeCun, 2025] for standardized comparison.Both losses instantiate the Cramer-Wold theorem and differ in one-sample versus two-sample distribution matching [Kuang et al., 2026b].Source: §H.1 · p. 19
  • Training used one seed, while planning evaluations averaged three seeds with Cross-Entropy Methods (CEM).Sweeps ran sequentially, and hyperparameters found to perform poorly were dropped, so predictor-type figures cover different subsets of the full grid.Source: §H.1 · p. 19
  • Learning rates ℓmuP ∈{0.05, 0.005} almost always produced a 0% success rate in the sweeps.Subsequent sweeps selectively removed portions of the grid to save compute.Source: §H.1 · p. 19

H.2 VICReg Sweep in the PushT Environment

This section reports PushT hyperparameter sweeps for a VICReg-regularized JEPA formed by replacing SIGReg in the LeWM setting. The sweeps cover Deep-AdaLN(k), MLP◦LTV(k), and MLP◦LTI(k) predictors.

H.2 VICReg Sweep in the PushT Environment

  • The VICReg model uses the LeWM setting with VICReg replacing SIGReg, and its PushT sweeps cover three predictor families.The model had no name yet, so the paper refers to it as VICReg.Source: §H.2 · p. 19
  • The PushT VICReg sweeps also include MLP◦LTV(k) and MLP◦LTI(k) predictors.Figures 13 and 14 present their respective hyperparameter sweeps.Source: §H.2 · p. 19, Figure 13 · p. 27, Figure 14 · p. 28
  • The PushT VICReg sweep includes the Deep-AdaLN(k) predictor.Figure 12 presents its hyperparameter sweeps.Source: §H.2 · p. 19, Figure 12 · p. 26

Related papers

Loading 2608.22764v1…