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

The paper investigates whether sparse rather than dense latent geometry makes action-conditioned dynamics easier to model and what structure sparse codes learn. It introduces LpWM, combining JEPA prediction with RDMReg, and reports lower predictor complexity for planning plus mode-factored representations. On PushT, LpWM improves success over dense LeWM by up to 57% at intermediate predictor capacities.

  • Problem

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

  • Method

    LpWM uses a JEPA world model with RDMReg to learn non-negative sparse codes, motivated by high-dimensional one-hot linearization of nonlinear controlled dynamics.

  • Results

    Up to 57% higher success rate is reported for sparse LpWM than dense LeWM on PushT with intermediate-capacity predictors, and sparse codes are mode-factored.

  • Takeaways & Limitations

    Sparse representations can reduce the predictor complexity required for control while yielding more interpretable dynamics.

  • Takeaways & Limitations

    Without temporal regularization, support may track fast-varying motion rather than semantically meaningful dynamical events; temporal regularization for more complex dynamics is left for future work.

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 paper asks whether sparse latent geometry makes action-conditioned dynamics easier to model than dense representations. It introduces LpWM and finds lower predictor complexity and interpretable mode-factored dynamics.

  • Dense anti-collapse objectives yield representations in which nearly all latent coordinates are nonzero, leaving their suitability for controlled dynamics unclear.
  • High-dimensional one-hot representations can approximate nonlinear controlled dynamics with action-conditioned linear transitions, motivating distributed sparse codes.For Lipschitz controlled dynamics on compact state spaces, rollout error vanishes as representation dimension grows.
  • LpWM is a JEPA model using RDMReg to learn non-negative sparse latent representations for dynamics modeling and planning.
  • Up to 57% higher success rate is reported for sparse LpWM than dense LeWM with intermediate-capacity predictors on PushT.The comparison spans MLP◦LTV(k), MLP◦LTI(k), and LTI(k) predictor variants.
  • Sparse representations are mode-factored: support encodes discrete dynamical regimes, while feature magnitudes encode continuous within-regime state.

2 Method

LpWM combines JEPA prediction with RDMReg and rectification to produce sparse, non-negative latents. Its training and control pipeline compares predicted and encoded target latents, then uses latent MPC for planning.

  • The action-conditioned JEPA trains an encoder and predictor by comparing predicted next latents with encoded next observations, alongside an anti-collapse regularizer.RDMReg is the chosen regularizer, weighted by λRDMReg.
  • RDMReg matches rectified feature distributions to Rectified Generalized Gaussian targets, with distribution discrepancy measured using the 2-Wasserstein distance.The default target uses µ = 0, σ = 1/2, and p = 1, reducing to a Rectified Laplace distribution.
  • The encoder uses a ViT and MLP projector, while the predictor uses a Transformer with AdaLN-zero action conditioning and an MLP output projector.
  • RepReLU produces exact zeros in the forward pass while routing backward gradients through GeLU to mitigate dying-ReLU behavior.The paper also reports that standard ReLU with RDMReg trains without collapse and yields exactly sparse representations.
  • At test time, MPC searches for an action sequence minimizing the distance between the predicted terminal latent and the encoded goal latent using CEM.The cost is C = ∥ˆzT − zg∥2.

3 Does Sparsity Simplify Latent Dynamics?

The paper motivates sparse latent codes by showing that high-dimensional one-hot representations can linearize nonlinear controlled dynamics, then tests whether learned distributed sparsity reduces predictor complexity. On PushT, sparse LpWM improves planning over dense representations at intermediate capacities, while gains diminish for sufficiently expressive predictors.

  • 3.1 Linearization of Nonlinear Dynamics with Sparse Representations: High-dimensional one-hot sparse representations can approximate nonlinear controlled dynamics with exactly linear latent transitions, with rollout error vanishing as dimension grows.This theoretical result motivates distributed sparse codes as a practical relaxation of exact one-hot representations.
  • 3.1 Linearization of Nonlinear Dynamics with Sparse Representations: Distributed sparse codes are proposed because exact one-hot encodings are impractical and require exponentially large feature dimensions for exact latent linearization.The one-step approximation error decreases as O(N^-1/d), exposing a curse-of-dimensionality trade-off.
  • 3.3 Results: LpWM is evaluated against dense LeWM and VICReg representations across predictor families, latent dimensions, and Wall and PushT environments.The predictor ladder spans linear state-space models, MLP-augmented linear models, and AdaLN Transformer predictors.
  • 3.3 Results: Nearly 100% closed-loop CEM success with LTI(1) is achieved in Wall for both dense and sparse embeddings, leaving little room for sparsity to improve performance.This indicates that Wall dynamics are already sufficiently captured by a simple linear predictor.
  • 3.3 Results: 24%–57% higher planning success is obtained by sparse LpWM over dense LeWM with MLP◦LTI(k) predictors on PushT at intermediate capacity.The corresponding gains are 36%–45% with MLP◦LTV(k) and 11%–23% with LTI(k).
  • 3.3 Results: LpWM outperforms VICReg across Deep-AdaLN(k), MLP◦LTV(k), and MLP◦LTI(k) predictors, including the high-capacity Deep-AdaLN(k) setting.The result extends the observed sparse-versus-dense advantage beyond LeWM’s specific dense isotropic Gaussian targets.
  • 3.3 Results: Sparsity gains depend on predictor capacity relative to dynamics complexity: they are substantial at intermediate PushT capacities but diminish with sufficiently expressive predictors.The paper leaves whether the advantage extends to increasingly expressive predictors as dynamics become more complex for future work.

4 Does Sparsity Lead to Interpretable Latent Dynamics?

Sparse LpWM representations exhibit interpretable mode-factored structure in navigation and contact-rich dynamics, although vanilla sparsity does not reliably align support transitions with meaningful temporal events.

  • Planning performance: Sparse LpWM reaches 84.7% success versus 65.3% for dense LeWM on Piecewise with random goals, while both methods saturate on evaluation-episode goals.The planning comparison accompanies the observed factorization of discrete zones and continuous state.
  • Mode-factored representations: In Piecewise, support similarity remains high within zones and drops sharply across zone boundaries, including when zones lack visual cues.This indicates that action-conditioned prediction recovers dynamics regimes rather than relying only on appearance.
  • Mode-factored representations: LpWM’s binary support identifies discrete dynamics regimes, while feature magnitudes encode continuous within-regime state.In Piecewise, support decodes the ground-truth zone near-perfectly, whereas magnitudes better decode continuous agent position.
  • Planning performance: Sparse LpWM outperforms dense LeWM at every planning horizon on Piecewise, with the performance gap widening as horizon H increases.The result concerns strong long-horizon planning in the reported experiment.
  • Contact-rich dynamics: Without temporal regularization on OGBench-Cube, support instability correlates strongly with effector motion (r ≈0.87) but weakly with contact (r ≈0.05).RDMReg constrains per-frame marginals, so support tends to follow the dominant fast-varying signal.
  • Contact-rich dynamics: Increasing the temporal Jaccard loss weight shifts support instability toward slower physical changes, raising its cube-motion correlation from r_cube = 0.21 to 0.80.The resulting support can better track object interaction and contact, but the design of temporal regularization for more complex dynamics remains future work.

5 Conclusion

The paper introduces LpWM to study sparse latent geometry for action-conditioned world models. It reports simpler latent dynamics and mode-factored representations, while Table 3 confirms that the learned codes are genuinely sparse.

  • Conclusion: LpWM is an action-conditioned JEPA regularized with RDMReg to produce non-negative, exactly sparse latent representations.The model is intended for latent dynamics modeling and planning.
  • Conclusion: Sparse representations simplify action-conditioned latent dynamics: at intermediate predictor capacity, shallow predictors plan over sparse codes where they fail over dense codes on PushT.The paper motivates this empirically with a one-hot linearization result.
  • Conclusion: Sparse codes are mode-factored, with support encoding discrete dynamics regimes and, with a temporal-Jaccard prior, tracking contact on OGBench-Cube.On Piecewise, the discrete regime is 94–99% decodable.
  • Sparsity verification: The LpWM code remains ∼30–65% active across predictor types and latent dimensions, whereas dense LeWM has active fraction 1.0 by construction.The table measures E[∥z∥0]/D on validation data at each tuned-best LpWM configuration.

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

The paper situates sparse world-model representations among methods for improving latent dynamics and evaluates predictors spanning nonlinear Transformers, linear state-space models, and structured time-varying operators.

  • Prior world models use anti-collapse regularization and distribution matching to produce dense latent representations, but dense geometry may not favor controlled-dynamics modeling.
  • The predictor comparison ranges from 6-block and 1-block AdaLN Transformers to LTI models and MLP-augmented linear cores.MLP◦LTV(k) adds state-dependent low-rank corrections to a fixed linear operator.
  • Structured time-varying prediction costs O(Dr) rather than O(D^3) by using a small state-dependent gate over fixed low-rank modes.The action operator is gated identically while remaining affine in the action.
  • The experiments cover Wall, PushT, Piecewise, and OGBench-Cube environments, including analytically known piecewise-affine dynamics with ground-truth zone segmentation.

D Experimental Details

The experiments use two complementary pipelines: one tests how predictor complexity and latent width affect planning, while the other probes sparse-support structure and interpretability.

  • Experimental design: The complexity study varies predictor capacity and latent dimension, whereas the interpretability study fixes configuration and analyzes sparse support in Piecewise and OGBench-Cube.
  • Training: Action-conditioned JEPAs match per-timestep latent marginals to transformed Generalized Gaussian targets using 2-Wasserstein distance.
  • Planning and evaluation: Planning uses CEM with 300 candidates per iteration, retains 30, runs 30 iterations, and uses horizon H = 5.
  • Planning and evaluation: Open-loop evaluation measures accumulated model-prediction quality, while closed-loop evaluation replans after each five-action block for up to 10 iterations.
  • Limitations: OGBench-Cube goals are only 25 raw environment steps ahead, which can saturate performance and limit benchmark discrimination.Longer-horizon and larger-frameskip evaluation is left for future work.
  • Representation analysis: Temporal Jaccard regularization penalizes rapid support changes between consecutive frames and encourages smoothly varying sparse supports.
  • Width scaling: µP is used to stabilize feature learning across latent widths, although learning-rate transfer works only most of the time and has occasional exceptions.

G.1 Proof of Proposition 1

The proof constructs a one-hot quantization of a compact state space and defines action-conditioned linear transitions that map each representative cell to the quantized next state.

  • Compactness provides a finite ε-cover of the state space, enabling finitely many representative centers and a deterministic quantization map.
  • The encoder maps each state to a standard basis vector, so every latent code is exactly one-hot.
  • For each action, the transition matrix maps each basis vector to the basis vector associated with the quantized successor of its representative.
  • Lipschitz continuity bounds one-step prediction error by propagating current-state approximation error through the dynamics and adding quantization error.
  • Because the bound is uniform over states and actions and tends to zero with ε, the constructed linear latent dynamics approximate the controlled dynamics arbitrarily well.
  • The finite-horizon rollout argument initializes error at most ε and recursively bounds later errors under arbitrary action sequences.

G.2 Proof of Corollary 1

The corollary specializes the construction to a uniform grid, relating the number of one-hot cells to covering error and showing that finite-horizon rollout error vanishes as grid resolution increases.

  • Partitioning each coordinate interval into n subintervals creates N = n^d axis-aligned hypercubes.
  • The cell-center construction bounds the distance from any state to its representative using the cube side length and Euclidean geometry.
  • The one-step decoded prediction error is bounded using the covering radius together with the Lipschitz dynamics constant.
  • Substituting the covering radius into the finite-horizon rollout bound shows convergence to zero for fixed horizon, state dimension, and Lipschitz constant.

H.1 LpWM and LeWM Sweeps in the PushT Environment

The PushT sweeps systematically vary regularization, learning rate, feature dimension, model, target distribution, and predictor settings. Compute constraints led to sequential sweeps, one training seed, three planning-evaluation seeds, and selective grid reduction after poor settings were identified.

  • Evaluation protocol: The sweep procedure is reported to support full reproducibility and transparency for Figure 1b..
  • Sweep design: The study uses sliced Wasserstein distance for isotropic Gaussian targets instead of the Epps-Pulley loss.Both losses are described as Cramer-Wold theorem instantiations that differ in one-sample versus two-sample distribution matching.
  • Evaluation protocol: Compute constraints limit training to one seed, while planning evaluations with CEM average over three seeds.Sweeps are run sequentially, and some predictor types receive more thorough grid searches than others.
  • Evaluation protocol: Learning rates 0.05 and 0.005 almost always produce 0% success, so later sweeps selectively remove those grid portions to save compute.The selective reductions apply to the sweeps shown in Figures 6, 7, 8, 10, and 11.

H.2 VICReg Sweep in the PushT Environment

The VICReg comparison replaces the LeWM setting’s SIGReg regularizer with VICReg and evaluates several predictor families through PushT hyperparameter sweeps. The reported sweep outputs measure open-loop planning success across RDMReg and learning-rate settings, averaged over three seeds.

  • VICReg setup: The VICReg model is formed by taking the LeWM setting and replacing SIGReg with VICReg.The resulting model has no separate name and is called VICReg in this section.
  • Predictor families: VICReg sweeps cover Deep-AdaLN(k), MLP◦LTV(k), and MLP◦LTI(k) predictors.Results are shown in Figures 12, 13, and 14, respectively.
  • Sweep outputs: The Deep-AdaLN(k) sweep reports open-loop planning success over RDMReg and learning rate as mean ± standard deviation across three seeds..
  • Sweep outputs: The Shallow-AdaLN(k) sweep reports open-loop planning success over RDMReg and learning rate as mean ± standard deviation across three seeds..
  • Sweep outputs: The MLP LTV(k) sweep reports open-loop planning success over RDMReg and learning rate as mean ± standard deviation across three seeds..
  • Sweep outputs: The MLP LTI(k), LTI(1), and LTI(k) sweeps likewise report open-loop planning success over RDMReg and learning rate across three seeds.These sweeps correspond to Figures 9, 10, and 11, respectively.
Loading 2608.22764v1…