DSReg Makes JEPA World Models Interpretable: Each Latent Dimension Maps to a Real-World Variable

Problem: JEPA Learns Predictive but Entangled Representations

Joint Embedding Predictive Architecture (JEPA) predicts future representations instead of reconstructing pixels, avoiding expensive decoders and becoming a popular world-modeling approach. However, without a reconstruction anchor, the learned representations are only identifiable up to an unknown linear transformation. In Linear JEPA (LeJEPA), the representation z relates to the true world state s by z = A s where A is an unknown invertible matrix. Consequently, each latent dimension mixes multiple physical factors (e.g., a single dimension may combine cup position and knife position), preventing the model from truly understanding the generative mechanism. This entanglement is harmless for pure prediction but dangerous for control: a robot trying to move only the cup might inadvertently move the knife if both are encoded in the same latent dimension.

Solution: DSReg — Dependency-Sparsity Regularization

DSReg is a lightweight post-processing step applied to an already-trained LeJEPA checkpoint. After whitening the LeJEPA representations, the remaining ambiguity is an orthogonal rotation R (since whitening fixes scaling and shear). Crucially, any rotation preserves the JEPA training loss. DSReg finds the unique rotation that disentangles the latents by exploiting the dependency structure between latents and observations.

Three-Step Algorithm

Estimate structure: Select anchor points in the data. For each anchor, fit a ridge regression on its nearest neighbors to estimate the local Jacobian J = ∂x/∂z mapping latents to observations.

Find rotation: Search over orthogonal matrices R to minimize the sparsity of the rotated Jacobians J R. The intuition: each true latent affects only a subset of observations; a mixed latent affects the union of its components' observation sets. The rotation that makes each latent depend on the fewest observations is the disentangling rotation.

Output variables: The rotated dimensions z' = z R now correspond one-to-one with real-world latent variables (up to permutation and sign).

The entire procedure requires no decoder, no reconstruction loss, and no labels. Local Jacobians are computed on-the-fly and need not be stored, enabling scaling to 8192 latent dimensions on a single 48 GB GPU.

Theoretical Guarantees

DSReg requires only Structural Diversity : for any two distinct latent variables, the sets of observations they influence cannot be identical. This condition is extremely weak — observation sets may overlap arbitrarily, and one may even fully contain the other. Under this assumption, the authors prove that DSReg recovers all latent variables up to permutation and sign flip, yielding the first provably identifiable JEPA. They further prove robustness to observation noise and formally verified the theorems in the Lean proof assistant.

Experimental Results

Synthetic data: LeJEPA latents remain mixed; adding DSReg achieves near-perfect recovery across all dimensions, reaching the theoretical upper bound.

Pixel-trained encoders: DSReg matches the supervised upper bound obtained by fitting latents with ground-truth labels. Classical rotation methods (PCA, Varimax, FastICA) show almost no improvement.

3DShapes: DSReg lifts performance from mixed to near the upper bound.

dSprites: DSReg outperforms β-VAE and β-TCVAE, which are explicitly designed for disentanglement.

Downstream Task Improvements

Because each dimension now corresponds to a semantically meaningful world variable, DSReg significantly improves six downstream tasks while preserving the original predictive performance:

Visual editing (manipulating individual factors)

Sparse control (acting on one factor without side effects)

Rollout prediction (long-horizon forecasting)

Anomaly detection (identifying unexpected changes)

Two additional control-oriented tasks

Scalability

DSReg scales linearly with latent dimension and runs on a single GPU. Experiments show no degradation in disentanglement quality when increasing from 64 to 8192 latents.

Implications for Robotics and Generalization

The authors argue that for robots, "知行合一" (unity of knowledge and action) is critical: an LLM's misunderstanding may cause a wrong utterance, but a robot's misunderstanding can cause catastrophic physical consequences. DSReg enables world models that not only predict accurately but also represent the true causal structure of the environment, providing a foundation for reliable generalization to novel situations. The method is available at github.com/kunwuz/dsreg with paper at arXiv:2610.09457 and project page dsreg.github.io.