SplitJEPA: Learning Invariant and Variant Latent Worlds without Reconstruction

Ruijin Hua; Zichuan Liu; Zhuokai Zhao; Yujia Zheng

cs.LGOct 8, 2026

SplitJEPA: Learning Invariant and Variant Latent Worlds without Reconstruction

Authors: Ruijin Hua, Zichuan Liu, Zhuokai Zhao, Yujia Zheng

Organizations: University of Illinois Urbana-Champaign · Carnegie Mellon University · University of Chicago

Abstract

Understanding a dynamical world calls for more than a latent state that summarizes its observations: the state should also be organized into the factors that stay shared across related observations and the factors that vary between them. For example, a robot pushing a cube to a goal should take the same action when the camera shifts or the lights dim, since nothing in the scene has moved. Existing approaches to this decomposition commonly obtain it through reconstruction, so the latent variables must first explain the entire observational world before their organization can be trusted. Joint embedding predictive architectures (JEPAs) model the latent state directly and never reconstruct, yet no existing result recovers the invariant and variant parts of the state they learn. How to learn the invariant-variant structure of the latent world without paying for its reconstruction therefore remains open. To close this gap, we introduce SplitJEPA, a JEPA that jointly recovers the latent state and its invariant and variant organization directly in representation space, without any reconstruction. We prove that, under stationary Gaussian predictive dynamics and a full-rank variation condition, SplitJEPA identifies the invariant and variant subspaces up to independent block-wise isometries, without introducing an observation decoder. Since the guarantee needs no decoder, the result extends reconstruction-free latent recovery to invariant-variant block identification. Experiments on synthetic nonlinear systems and robotic manipulation tasks support the theoretical results and show their practical value for both robustness and efficiency.

Figures & tables

Figure 1: Block identifiability without observation reconstruction. Reconstruction-based identifiable models recover structured latent factors through an encoder-decoder pathway together with additional structural assumptions. SplitJEPA combines prediction in latent space with block consistency to recover invariant and variant structure directly in representation space.

Figure 2: Nonlinear mixing of the latent blocks.

Figure 3: Overview of reconstruction-free block identifiability. (a) Matched PushCube trajectories share invariant task content while varying in execution and environmental factors. (b) The ground-truth latent state consists of invariant and variant blocks that jointly generate the observations. (c) Predictive alignment and variance preservation identify the complete latent state, while block consistency determines its partition without an observation decoder. (d) The resulting representation recovers the invariant and variant subspaces. (e) Reconstruction alone permits observationally equivalent latent parameterizations, and prediction alone leaves a global orthogonal ambiguity; additional consistency structure resolves this ambiguity and yields reconstruction-free block identification.

Figure 4: Complementary changes span the variant subspace.

Method	Invariant $R^{2}$	Variant $R^{2}$	$\lVert A_{12}\rVert_{F}$	$\lVert A_{21}\rVert_{F}$	Block Purity
LeJEPA	$0.9613\pm 0.0334$	$0.9612\pm 0.0328$	$0.3292\pm 0.2145$	$0.3299\pm 0.2137$	$0.9635\pm 0.0318$
Shuffled Pairs	$0.7134\pm 0.0714$	$0.7102\pm 0.0736$	$1.0600\pm 0.1415$	$1.0606\pm 0.1394$	$0.7142\pm 0.0734$
SplitJEPA (Ours)	$\mathbf{0.9979\pm 0.0014}$	$\mathbf{0.9973\pm 0.0019}$	$\mathbf{0.0031\pm 0.0018}$	$\mathbf{0.0236\pm 0.0021}$	$\mathbf{0.99993\pm 0.00001}$

Table 1: Block recovery on the synthetic OU system. All methods reuse the same frozen reconstruction-free encoder and differ only in the information available to the spectral recovery stage. Results are the mean

\pm

standard deviation over five seeds. Global

R^{2}

is approximately

0.9977\pm 0.0016

for all methods.

Figure 5: Ablation of the predictive, invariance, and variance-preserving components. Bars show mean

\pm

standard deviation over five seeds. (a) Removing predictive alignment reduces global linear recovery. (b)–(c) Removing invariant-pair recovery leaves global identification unchanged but degrades block recovery and increases cross-block leakage. (d) Removing variance regularization and ZCA normalization increases the representation covariance error.

Panel of Figure 6 — Figure 6: Offline action prediction performance. Validation action MSE during offline behavior cloning training on PickCube and PushCube.

	PushCube Manipulation					PickCube Manipulation
Method	ID		OOD			ID		OOD
	CONTACT	SR@10cm	Camera	Lighting	Visual	GRASP	SR@10cm	Camera	Lighting	Visual
Raw	$100.0\!\pm\!0.0$	$98.0\!\pm\!1.0$	$12.8\!\pm\!6.8$	$98.3\!\pm\!1.5$	$12.7\!\pm\!7.0$	$23.0\!\pm\!1.6$	$14.0\!\pm\!1.4$	$10.3\!\pm\!0.5$	$11.7\!\pm\!0.9$	$10.7\!\pm\!0.5$
R3M	$100.0\!\pm\!0.0$	$97.7\!\pm\!1.2$	$48.0\!\pm\!7.5$	$90.7\!\pm\!2.1$	$25.8\!\pm\!1.8$	$73.0\!\pm\!1.6$	$32.0\!\pm\!1.4$	$14.3\!\pm\!0.5$	$10.0\!\pm\!0.0$	$10.0\!\pm\!0.0$
LeJEPA	$90.3\!\pm\!14.2$	$67.3\!\pm\!20.1$	$25.0\!\pm\!2.6$	$31.0\!\pm\!31.2$	$28.7\!\pm\!17.5$	$33.3\!\pm\!3.3$	$16.3\!\pm\!0.9$	$10.0\!\pm\!0.0$	$10.0\!\pm\!0.0$	$10.0\!\pm\!0.0$
Ours- $r_{H}$	$100.0\!\pm\!0.0$	$99.3\!\pm\!0.6$	$71.5\!\pm\!10.4$	$\mathbf{99.7\!\pm\!0.6}$	$64.5\!\pm\!12.7$	$\mathbf{97.0\!\pm\!1.4}$	$45.0\!\pm\!2.4$	$\mathbf{30.3\!\pm\!6.1}$	$42.7\!\pm\!4.5$	$28.3\!\pm\!3.9$

Table 2: Robotic manipulation performance. Results on PickCube and PushCube under ID and OOD conditions. CONTACT and GRASP are reported for ID; all OOD columns report SR@10cm.

Appendix figures & tables13 assets

Supplementary material from the paper’s appendix.

Appendix

Figure 7: Global linear identification under three scaling regimes. Higher Global

R^{2}

indicates more accurate linear identification of the complete latent representation.

Figure 8: Recovery of the invariant latent block under three scaling regimes. Invariant

R^{2}

measures the linear recoverability of the ground-truth invariant factors from the learned invariant representation.

Figure 9: Recovery of the variant latent block under three scaling regimes. Variant

R^{2}

measures the linear recoverability of the ground-truth variant factors from the learned variant representation.

Figure 10: Normalized cross-block leakage

A_{12}

under three scaling regimes. The normalized

A_{12}

score quantifies leakage from the variant latent block into the recovered invariant block and is defined as

\lVert A_{12}\rVert_{F}/\sqrt{d_{H}}

Figure 11: Normalized cross-block leakage

A_{21}

under three scaling regimes. The normalized

A_{21}

score quantifies leakage from the invariant latent block into the recovered variant block and is defined as

\lVert A_{21}\rVert_{F}/\sqrt{d_{L}}

Figure 12: Simulation efficiency comparison.

Method	Invariant $R^{2}$	Variant $R^{2}$	$\lVert A_{12}\rVert_{F}$	$\lVert A_{21}\rVert_{F}$	Block Purity
LeJEPA	$0.6778\pm 0.0360$	$0.6769\pm 0.0387$	$1.0940\pm 0.0696$	$1.0910\pm 0.0710$	$0.6936\pm 0.0393$
Shuffled Pairs	$0.6598\pm 0.0555$	$0.6606\pm 0.0567$	$1.1243\pm 0.0997$	$1.1220\pm 0.1013$	$0.6750\pm 0.0581$
SplitJEPA (Ours)	$\mathbf{0.9770\pm 0.0003}$	$\mathbf{0.9769\pm 0.0002}$	$\mathbf{0.0056\pm 0.0007}$	$\mathbf{0.0187\pm 0.0037}$	$\mathbf{0.99995\pm 0.00002}$

Appendix

Table 3: Reconstruction-free block recovery under joint nonlinear MLP mixing (mean

\pm

standard deviation over five seeds). All methods reuse the same frozen MLP encoder and differ only in the information available during spectral recovery. Global

R^{2}

is approximately

0.977

for all methods.

Coupling Strength	Global $R^{2}$	$\lVert Q^{\top}Q-I\rVert_{F}/\sqrt{d}$	$\lVert A_{12}\rVert_{F}$	$\lVert A_{21}\rVert_{F}$
$0.5$	$\mathbf{0.9979}$	$\mathbf{0.0199}$	$\mathbf{0.0026}$	$\mathbf{0.0155}$
$1.0$	$0.9565$	$0.0586$	$0.0572$	$0.0210$
$2.0$	$0.8765$	$0.1507$	$0.1393$	$0.0289$

Appendix

Table 4: Sensitivity to nonlinear coupling strength. Seed-

0

diagnostic under the default synthetic configuration. Increasing the coupling strength progressively degrades the global identification stage.

Panel of Figure 13 — Figure 13: Sample efficiency of learned representations. Performance of different representations under varying amounts of training data on PushCube and PickCube, evaluated under ID and OOD settings.

	PushCube			PickCube
Intervention	$r_{H}+r_{L}$	$r_{H}$	$r_{H}$ advantage	$r_{H}+r_{L}$	$r_{H}$	$r_{H}$ advantage
Zero $r_{L}$	$93.7\pm 3.4\%$	$99.9\pm 0.2\%$	$+6.1$ pp	$61.4\pm 5.1\%$	$99.7\pm 0.5\%$	$+38.4$ pp
Noise $0.5\sigma$	$99.0\pm 1.0\%$	$100.0\pm 0.0\%$	$+1.0$ pp	$87.9\pm 0.8\%$	$100.0\pm 0.0\%$	$+12.1$ pp
Noise $1.0\sigma$	$90.2\pm 0.5\%$	$99.9\pm 0.2\%$	$+9.7$ pp	$62.2\pm 4.6\%$	$100.0\pm 0.0\%$	$+37.8$ pp
Noise $2.0\sigma$	$44.1\pm 4.0\%$	$99.9\pm 0.2\%$	$+55.8$ pp	$29.6\pm 0.8\%$	$100.0\pm 0.0\%$	$+70.4$ pp

Appendix

Table 5: Performance retention under representation interventions

Panel of Figure 14 — Figure 14: Empirical validation of invariant-variant block specialization. Representation responses to matched viewpoint and trajectory variations on PickCube and PushCube. The learned invariant block $r_{H}$ changes substantially less than the variant block $r_{L}$ under the same variations.

	PushCube					PickCube
Method	ID		OOD			ID		OOD
	SR@10cm	CONTACT	Camera	Lighting	Visual	SR@10cm	GRASP	Camera	Lighting	Visual
Full ( $r_{H}$ )	$99.3\pm 0.6$	$100.0\pm 0.0$	$71.5\pm 10.4$	$99.7\pm 0.6$	$64.5\pm 12.7$	$45.0\pm 2.4$	$97.0\pm 1.4$	$\mathbf{30.3\pm 6.1}$	$42.7\pm 4.5$	$28.3\pm 3.9$
Full ( $r_{H}+r_{L}$ )	$99.3\pm 0.6$	$100.0\pm 0.0$	$75.3\pm 8.1$	$98.3\pm 1.5$	$\mathbf{70.2\pm 12.9}$	$\mathbf{46.3\pm 4.9}$	$96.7\pm 0.5$	$27.7\pm 0.9$	$43.3\pm 2.9$	$\mathbf{30.3\pm 1.7}$
w/o Predictive ( $r_{H}$ )	$98.7\pm 0.6$	$100.0\pm 0.0$	$65.7\pm 6.4$	$99.7\pm 0.6$	$54.0\pm 14.3$	$39.7\pm 4.5$	$92.7\pm 1.2$	$20.7\pm 0.9$	$35.7\pm 3.9$	$19.0\pm 0.8$
w/o Predictive ( $r_{H}+r_{L}$ )	$98.7\pm 0.6$	$100.0\pm 0.0$	$65.3\pm 8.8$	$98.0\pm 1.0$	$59.8\pm 12.4$	$39.3\pm 2.5$	$94.0\pm 1.4$	$18.3\pm 1.7$	$39.3\pm 3.4$	$18.7\pm 1.2$

Appendix

Table 6: Robot manipulation representation ablation results. Values report the mean

\pm

standard deviation over three training seeds. All variants retain the same partitioned architecture and downstream pipeline as the Full model and remove only the indicated component. ID is evaluated using SR@10cm and the task-specific interaction rate (CONTACT for PushCube and GRASP for PickCube); all OOD columns report SR@10cm.

Figure 15: Temporal transferability. Normalized representation changes between state-matched observations from different trajectories and timesteps averaged over three seeds.

Explore similar work

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 Architecture Orthogonal Representation Learning Predictive Representation Learning Latent World Models

Ask PDF

May 29, 2026cs.LG

Subspace-Decomposed JEPAs: Disentangling Progression and Content in Latent World Models

Joint-Embedding Predictive Architectures (JEPAs) learn compact latent world models by predicting future embeddings, but no single coordinate of the latent is designated to encode task progression. We carve the JEPA latent into two orthogonal subspaces with disjoint roles: a low-dimensional progression subspace shaped by a cosine-margin triplet loss, and a high-dimensional content subspace regularised by the existing SIGReg objective of LeWM. We prove that the two anti-collapse forces act on disjoint coordinates, so they compose additively rather than competing on the same dimensions. Our method, SD-JEPA improves over the LeWM baseline on the majority of its control benchmarks at matched compute, and outperforms the strongest non-LeWM JEPA baseline on Push-T; a subspace-ablation falsifier confirms the split is the load-bearing ingredient. Beyond planning, the resulting 1-D angular progression coordinate functions as a scene-aware compass on the latent. It advances with task progress, regresses when the agent backtracks, and under controlled perturbations both spikes and relocalises to a semantically appropriate new task-phase sector, separating the moment of surprise from its meaning in a way that prediction-error scalars cannot.

Lucas Thil, Jesse Read, Rim Kaddah +1

LIX, École Polytechnique IRT SystemX Palaiseau, France · IRT SystemX Palaiseau, France · Safran Tech Chateaufort, France

Disentangled Representation Learning Representation Learning Joint-Embedding Predictive Architecture Latent World Models

Ask PDF

Aug 6, 2026cs.CV

PhyLatent: Learning Dynamics-Relevant Representations for JEPA World Models

We propose PhyLatent, a dynamics-relevant training objective for Joint-Embedding Predictive Architecture (JEPA) world models. Our key observation is that preventing global latent collapse does not necessarily ensure that the learned representation preserves physically meaningful state and action relationships. We identify three failure modes: sensitivity to appearance changes that leave the physical state unchanged, insufficient separation of distinct physical states, and insufficient separation of different action-conditioned futures. We refer to these as Physical Invariance Collapse, Physical Distinguishability Collapse, and Counterfactual Dynamics Collapse, respectively. PhyLatent targets these failures through three coordinated training pathways, implemented with static visual invariance, physical state grounding, future representation alignment, counterfactual branch separation, and latent denoising. On OGBench-Cube, PhyLatent reduces the three collapse rates by 43.9%, 27.9%, and 47.1%, respectively, while improving model predictive control (MPC) success by 12.0 percentage points (17.2% relative). Across four visual-control tasks, average planning success increases by 6.62 percentage points (8.3% relative). These results show that global non-collapse alone is insufficient for learning a reliable JEPA world-model state space, and that explicitly preserving dynamics-relevant structure can improve closed-loop planning.

Xi Zeng, Haojie Ren, Ziying Song +2

School of Mechanical and Aerospace Engineering, Nanyang Technological University, Singapore · The University of Sheﬃeld, Sheﬃeld, United Kingdom · School of Artiﬁcial Intelligence (School of Software), Yanshan University, Qinhuangdao, China

Representation Learning World Model Learning World Model-Based Planning Joint-Embedding Predictive Architecture State Representation Learning

Ask PDF