DSReg: Provably Recovering Individual World Latents without Reconstruction

Yujia Zheng; David Klindt; Randall Balestriero; Bernhard Schölkopf

cs.LGOct 7, 2026

DSReg: Provably Recovering Individual World Latents without Reconstruction

Authors: Yujia Zheng, David Klindt, Randall Balestriero, Bernhard Schölkopf

Organizations: University of Illinois Urbana-Champaign · Cold Spring Harbor Laboratory · Brown University · Max Planck Institute for Intelligent Systems · ELLIS Institute Tübingen

Abstract

Methods that recover individual latent variables of the world, from nonlinear ICA to dictionary learning and causal representation learning, anchor the latents to observations through reconstruction, auxiliary supervision, or distributional asymmetries such as non-Gaussianity. Methods without these anchors, including joint-embedding predictive architectures (JEPAs), identify the latent state only up to a linear transformation, so individual latents remain mixed. We close this gap: individual world latents can be provably recovered with no reconstruction, no decoder, and no labels. The key condition is Structural Diversity: different latents leave distinct dependency footprints on observations, just as no two snowflakes are alike. Building on the linear identifiability that LeJEPA provides, we prove that under Structural Diversity, DSReg (Dependency-Sparsity Regularization) recovers individual world latents up to signed permutation, without reconstruction or a decoder. It applies post hoc to any linearly identified representation, reusing trained checkpoints at no loss over joint training, and establishes the first fully identifiable JEPA that recovers every world latent. Moreover, as a condition on dependency footprints, Structural Diversity is strictly weaker than all structural conditions of prior identifiable latent variable models. Across synthetic regimes, world model probes, learned visual encoders, and external renderers, DSReg preserves dense prediction while improving individual-latent recovery and downstream use with scales.

Figures & tables

Figure 1: DSReg turns a mixed latent state into individual latents. World latents $z$ generate observations $x=g(z)$ . LeJEPA identifies the latent state up to a linear transformation, $\hat{z}=Az$ , so each learned variable can remain a mixture of ground-truth latents. DSReg recovers individual latents up to signed permutation.

Figure 2: Dependency footprints. Each column of $\partial x/\partial z$ marks which observed variables one latent affects.

Panel of Figure 3 — Figure 3: DSReg recovers individual latents wherever Structural Diversity holds, and only there. (a) DSReg (solid) is near ceiling at every $N$ , twenty runs. LeJEPA (dashed) stays mixed. (b) On analytic orbits $h=Qz$ , recovery succeeds wherever the condition holds and falls back to the shared pair subspace where it fails.

Figure 4: DSReg improves all six sparse-use probes. DSReg in teal, LeJEPA in gray. Tasks are (1) editing (top) and (2) control, prediction, and monitoring (bottom).

Panel of Figure 5 — Figure 5: DSReg matches the supervised Procrustes oracle on encoders trained from pixels. (a) Latent-only rotations leave the variables mixed. DSReg closes the gap to the label-fitted oracle (fifteen seeds). (b) With $\dim z=8$ fixed, the match holds within $0.001$ as the estimate widens to $\dim\tilde{z}=32$ (five seeds).

Appendix figures & tables20 assets

Supplementary material from the paper’s appendix.

Appendix

		MCC			step time (ms)			peak memory (GB)
$d$	$n$	mat.	fact.	samp.	mat.	fact.	samp.	mat.	fact.	samp.
128	$20$ k	$0.853$	$0.853$	$0.850$	$7.5$	$4.9$	$42.6$	$0.14$	$0.11$	$0.11$
256	$20$ k	$0.843$	$0.844$	$0.840$	$7.4$	$6.8$	$43.9$	$0.47$	$0.31$	$0.30$
512	$24.6$ k	$0.833$	$0.831$	$0.829$	$13.1$	$34.0$	$43.8$	$1.76$	$1.07$	$1.02$
1024	$49.2$ k	$0.827$	$0.828$	$0.823$	$101.6$	$231.6$	$53.5$	$6.97$	$4.20$	$3.26$
2048	$98.3$ k	$0.821$	$0.820$	$0.809$	$699.8$	$1933.2$	$92.5$	$27.76$	$13.73$	$11.42$

Appendix

Table 1 : Scaling across latent dimension. The full estimated procedure, three seeds each. At

d\geq 1024

the sampled columns use a compute-matched budget whose wall-clock stays at or below the materialized one; at smaller

d

the standard budget already suffices, and sampling holds no advantage there since its per-step cost is nearly

d

-independent while the materialized cost is what grows. All three modes optimize the same objective, and the materialized and factored columns agree on MCC to within

0.002

everywhere, which is a live check that the factorization is exact. The factored mode trades time for memory, roughly halving peak memory at

d=2048

at about

2.8\times

the step time, since it recomputes the neighbor differences on each pass. Axes without a panel of their own, each measured on its own sweep with its own baseline rather than on the rows above: raising the sample count to

10^{6}

leaves recovery and step time flat to within seed noise (

0.826\to 0.833

d=1024

49

ms throughout;

0.822

d=2048

), and growing the observed dimension to

2^{17}

holds the sampled step at

47

ms while MCC improves from

0.806

0.884

across the full doubling sweep.

Figure 6 : Scaling the full estimated procedure. (a) Recovery declines gently with

d

; open markers are the single-GPU frontier, where anchor count shrinks to fit memory. (b) Recovery against the neighborhood ratio

k/d

at three scales; dashed lines mark the analytic score of an unrotated mixing,

\sqrt{2\ln d/d}

. (c) Sixteen anchors already match

256

d=2048

. (d) The sampled step cost is nearly flat in

d

while the materialized cost grows two orders of magnitude, with the curves crossing between

d=512

and

d=1024

Figure 7 : Footprint regimes across dimension. Recovery follows Structural Diversity at every dimension: diverse, nested, and minimal-difference footprints stay at ceiling while identical footprints sit at the shared-subspace level, far above the decaying baseline of the unrotated representation; the identical pair’s span is still recovered, with CCA at least

0.998

at every dimension of the full regime sweep.

Figure 8 : Recovery is insensitive to the optimizer’s choices. Recovery is insensitive to the sparsity surrogate and to the anchor budget; only the minimal-difference regime rewards more anchors, exactly where the evidence is scarcest. The two rightmost bars replace the relaxation by a hard support count at

\tau=0.05

, minimized by coordinate search from the same six starts, and by that count used to refine the

\ell_{1}

solution afterwards. Both series are DSReg selections, with color denoting the underlying footprint regime.

Schedule	Latent MCC
LeJEPA (no rotation)	$0.640\pm 0.041$
Joint	$0.906\pm 0.014$
Joint, then decoupled fit	$0.922\pm 0.005$
Decoupled (default)	$0.921\pm 0.005$

Appendix

Table 2 : Training-schedule ablation on Gaussian 3DShapes. Latent MCC (mean

\pm

std), twenty runs per arm on a shared encoder trajectory; the rotation head never feeds back into the encoder.

Figure 9 : Corrupted Jacobians leave recovery at ceiling. Diverse footprints stay above

0.99

through noise

0.20

, and identical footprints stay at their predicted shared-subspace level. Nonlinear observation rows, so Assumption 2 holds;

N=16

, five seeds. LeJEPA applies no rotation, so its curve is flat by construction.

Figure 10 : Sparse access improves while dense content is unchanged. At identical dense recovery (

R^{2}=0.988

for both methods), the rotation transforms individual-latent recovery, few-shot readout from eight labels, and sparse editing, each moving decisively in DSReg’s favor on the MLP benchmark at

N=8

Figure 11 : Success strictly contains the guarantee. Success fraction of the thresholded criterion over the $(\tau,\varepsilon)$ grid, five seeds. Solid curve: the guarantee boundary $\tau+\varepsilon=\sigma\rho^{*}(\delta^{*})$ of Theorem 3 , with seed-range shading; circled cells satisfy the guarantee for every seed, and all of them succeed. The $\ell_{1}$ rotation used throughout the paper reproduces this map on all but four of the $108$ cells, again succeeding on every guaranteed cell.

Figure 12 : The excluded linear stratum is thin. MCC against the nonlinearity share

\gamma

with the measured no-cancellation margin

\sigma

overlaid (dotted, right axis). Where Assumption 2 fails (

\gamma=0

\sigma=0

) the criterion returns a mixed rotation; recovery crosses the

0.9

level (dashed) at

\gamma^{*}\approx 0.1

and reaches the ceiling by

\gamma\approx 0.35

at every dimension, in both regimes, with exact and estimated Jacobians in close agreement.

Benchmark	$N$	Dense $R^{2}$	MCC	Few-shot $R^{2}$	Sparse error $\downarrow$
Visual factors	8	$0.918\pm 0.002$	$0.625\to 0.958$	$-0.081\to 0.882$	$2.13\to 1.35$
Object scenes	8	$0.924\pm 0.002$	$0.591\to 0.961$	$-0.143\to 0.899$	$2.29\to 1.41$

Appendix

Table 3 : Split-audited learned visual encoders. The analytic signature survives learning: dense recovery is unchanged while individual-latent recovery, few-shot readout, and sparse use improve, with every score computed on a held-out split (the rotation is fit on a separate split of estimated latents and fixed pooled-pixel features). Arrows are ten-run means for LeJEPA

\to

DSReg; sparse error is lower better.

Figure 13 : The analytic signature survives pixel training. The gain concentrates where individual latents matter: component-wise recovery and single-latent readout improve sharply, and sparse use follows. A negative few-shot value underperforms the held-out mean; the sparse-use score is normalized so that LeJEPA sits at one. On the pixel probes, sparse control error drops from

0.63

0.52

, five-step sparse rollout

R^{2}

rises from

0.74

0.81

, and transition-surprise AUROC from

0.995

0.999

(LeJEPA / DSReg, ten runs).

Figure 14 : Training imbalance does not break the rotation. Long-tail and correlated training distributions leave the recovery gain intact, with the dense span fully preserved throughout.

Method	Latent MCC	DCI	MIG	SAP	Info. $R^{2}$	Dense $R^{2}$
LeJEPA	$0.640\pm 0.041$	$0.298\pm 0.044$	$0.069\pm 0.026$	$0.212\pm 0.074$	$0.854\pm 0.003$	$0.861\pm 0.002$
$\beta$ -VAE	$0.272\pm 0.025$	$0.318\pm 0.086$	$0.080\pm 0.038$	$0.102\pm 0.030$	$0.738\pm 0.076$	$0.191\pm 0.010$
$\beta$ -TCVAE	$0.277\pm 0.043$	$0.345\pm 0.117$	$0.102\pm 0.067$	$0.109\pm 0.052$	$0.751\pm 0.087$	$0.199\pm 0.027$
DSReg	$0.921\pm 0.005$	$0.901\pm 0.017$	$0.419\pm 0.013$	$0.842\pm 0.016$	$0.868\pm 0.003$	$0.861\pm 0.002$

Appendix

Table 4 : DSReg improves latent MCC on every run of Gaussian 3DShapes while matching LeJEPA’s dense recovery. Scores are latent MCC, DCI, MIG, SAP, and informativeness (test

R^{2}

of a boosted readout), twenty runs, on an external renderer. The LeJEPA row is the same frozen representation DSReg rotates, which is why the two share dense

R^{2}

; the

\beta

-VAE/

\beta

-TCVAE rows are external references with the matched trunk, trained on i.i.d. images at an untuned

\beta{=}4

; their low dense

R^{2}

reflects a linear readout only, with nonlinear informativeness staying high, so the gap measures axis alignment rather than lost information.

Figure 15 : Per-factor recovery on Gaussian 3DShapes. DSReg improves every factor and alone recovers scale, shape, and orientation, where the VAE baselines collapse almost entirely.

Figure 16 : The rotation turns entangled responses into one selective response per factor. One run of Gaussian 3DShapes. Each row sweeps one factor with all others fixed (rendered frames on the left); the right panels show every learned latent’s standardized response in LeJEPA and in DSReg coordinates, on a shared scale per row. The flat gray curves carry the no-mixing claim across all six factors.

Figure 17 : Matched correlation heatmaps. For the run of Figure 16 , columns are permuted by each method’s own best assignment. The same encoder yields a smeared matrix in the unrotated coordinates and a near-diagonal one after rotation, in agreement with the per-factor responses of the figure above.

Figure 18 : DSReg beats both VAE baselines under a partially attained premise. Quarter-orientation dSprites, twenty runs, dense

R^{2}\approx 0.53

as the rotation-invariant ceiling; each VAE at its best

\beta

. One-sided Mann–Whitney tests are against

\beta

-VAE and

\beta

-TCVAE at their own best

\beta

; higher MCC is better.

Method	Latent MCC	$R^{2}(h\to z)$	Dependency sparsity $\downarrow$
LeJEPA	$0.65\pm 0.03$	$1$	$0.32\pm 0.02$
Random orbit	$0.64\pm 0.03$	$1$	–
PCA	$0.62\pm 0.04$	$1$	–
Varimax	$0.64\pm 0.03$	$1$	–
FastICA	$0.64\pm 0.05$	$1$	–
DSReg	$0.90\pm 0.05$	$1$	$0.18\pm 0.02$

Appendix

Table 5 : Latent-only rotation baselines. The observation-side signal rather than the span is what does the work: all methods start from the same linearly identifiable representation, but PCA, Varimax, and FastICA see only estimated-latent samples and cannot break the rotational symmetry, while the dependency criterion

\partial x/\partial(Rh)

can, because it also sees how the observed variables respond to the applied rotation.

Task	Retrieval cost LeJEPA / DSReg	Sparse err. LeJEPA / DSReg
TwoRoom	$1.31/1.31$	$0.81/0.58$
PushT	$1.45/1.45$	$1.45/1.10$
FetchSlide	$1.43/1.43$	$1.24/0.85$

Appendix

Table 6 : Dense versus sparse-use sanity check. Everything the rotation gains, it gains in sparse access: dense readout and ordinary retrieval are unchanged, with dense

R^{2}

1/1

(LeJEPA / DSReg) for every task, while sparse-control error drops whenever only a few estimated latents may be used.

Criterion	Gaussian 3DShapes	Quarter-orientation dSprites
Column-support sparse ICA	$0.76\pm 0.06$	$0.57\pm 0.05$
DSReg	$0.92\pm 0.01$	$0.60\pm 0.08$

Appendix

Table 7 : Sparse ICA on the same dependency signal recovers less. Latent MCC, mean

\pm

std over ten fresh runs per renderer benchmark; both criteria consume the same estimated anchor Jacobians, orthogonal parameterization, and restarts. Both rows share every setting except the criterion.

Explore similar work

May 25, 2026stat.ML

When Does LeJEPA Learn a World Model?

A representation that scrambles the true degrees of freedom of the world cannot support reliable planning or compositional generalization. We prove that LeJEPA (alignment plus Gaussian regularization) linearly recovers the world's latent variables from nonlinear observations, a property known as linear identifiability, in a broad class of worlds where latents evolve under stationary, additive-noise transitions. Our main result is that among all such worlds, the Gaussian is the unique latent distribution for which this guarantee holds. The forward direction rests on a spectral decomposition in which each degree of nonlinearity is strictly penalized by alignment, making the linear map the optimum; the converse rules out every non-Gaussian alternative. We further prove an approximate identifiability result where the guarantee degrades gracefully, and show that linear, orthogonal identifiability enables optimal latent-space planning. We validate the theory with experiments ranging from 2D examples to 1024-dimensional latents, including distributional ablations and pixel-based robotic control. Our theory turns an empirically successful recipe into a mathematical guarantee, providing the foundation for building World Models that provably recover the structure of the world.

David Klindt, Yann LeCun, Randall Balestriero

Cold Spring Harbor Laboratory · New York University · Brown University

Latent World Models Low-Latency Latent Planning World Models Identifiability World Regularization

Ask PDF

Aug 20, 2026cs.LG

Orthogonal JEPA: Factorized Predictive States for Latent World Models

World models construct latent states that support prediction, planning, and reasoning about an underlying system. Joint-embedding predictive architectures (JEPAs) offer a direct way to learn such states by predicting targets in representation space instead of reconstructing every detail of the observation. Standard JEPAs, however, organize all predictable content through one target embedding and one prediction pathway. In complex systems, this monolithic state can allocate redundant capacity to dominant signals while providing weak or conflicting gradients to less dominant predictive structure. We introduce \method, a latent world-modeling framework based on orthogonal predictive factorization. Learned basis matrices analyze each target state into multiple components, and a dedicated prediction branch estimates each component from a shared context representation. Predictive regression preserves the factor magnitudes required for state synthesis, an orthogonality objective discourages repeated directions, factor-activity regularization maintains variation in projected targets, and online variance regularization discourages coordinate-wise encoder collapse. Predicted components are synthesized into a complete latent state that can be used by a readout, decoder, planner, or autoregressive rollout. The same predictive-state mechanism applies when the target is temporally future, spatially hidden, or another partial observation of the same system. Experiments on controlled vision, single-cell transcriptomics, longitudinal health records, continuous control, and molecular dynamics evaluate representation quality, forecasting, planning, and long-horizon stability.

Taoyong Cui, Pheng Ann Heng, Wanli Ouyang

The Chinese University of Hong Kong (CUHK)

Joint-Embedding Predictive Architectures Next-State Prediction Latent World Models World Models Regularization Single-Cell Rna

Ask PDF

Jul 15, 2026cs.LG

The SIGReg Objective as Variational Free Energy: A Theoretical Active-Inference Account of JEPA World Models

Joint-Embedding Predictive Architectures (JEPAs) are the dominant design for latent world models, yet they are usually justified by empirical performance rather than a normative principle. We show that the choice of anti-collapse regulariser determines whether a JEPA's training objective, a prediction loss plus a weighted embedding regulariser, is a valid Active Inference (AIF) variational free energy. We organise four non-contrastive regularisers (VICReg, LogDet, PairDist, and SIGReg) into an entropy-estimator hierarchy indexed by a prior-miscalibration gap, and show that the gap's sign, whether the estimator bounds the latent entropy from above or below, decides whether the AIF surprise bound survives: VICReg and LogDet are unsafe upper bounds, PairDist a safe lower bound, and SIGReg eliminates the gap. We then prove a correspondence theorem: under the standard constant-noise encoder model and successful SIGReg enforcement (isotropic-Gaussian embeddings), the gap vanishes, the objective becomes an exact information bottleneck, the surprise bound is preserved, and the latent goal cost becomes an exact proxy for AIF pragmatic value, whereas VICReg leaves an irreducible second-order anisotropy term. We extend the correspondence to multi-step expected free energy, ensemble epistemic value, and a learned-policy regime, and we identify the one AIF term no current JEPA world model computes: the state-epistemic value, a future-state coverage signal. The predictions differ in kind, not degree, and are stated here as theoretical consequences left for empirical test in separate work; full proofs are in Appendix A, and the algebraic core of every result is machine-verified in Lean 4 (Appendix D).

Fabio Arnez, Alexandra Gomez-Villa

Université Paris-Saclay, CEA, List F-91120, Palaiseau, France · Computer Vision Center Barcelona, Spain

Joint-Embedding Predictive Architectures Variational Free Energy Principle Signsgd Collapse

Ask PDF