SplitJEPA: Learning Invariant and Variant Latent Worlds without Reconstruction

1 Introduction

• Robotic representations should retain shared task content while separating variation in configuration, motion, viewpoint, illumination, and appearance.The motivation is to use stable information without discarding variation when it is useful.Source: §1 · p. 1
• Reconstruction objectives add modeling and computational costs, can allocate capacity to noise and unnecessary details, and do not by themselves determine the latent partition.Invertible nonlinear reparameterizations can preserve the observation distribution while mixing latent factors.Source: §1 · p. 1
• JEPAs learn representations through prediction without reconstruction, but existing empirical results do not guarantee recovery of the ground-truth invariant–variant decomposition.LeJEPA identifies a complete latent process under specified conditions, yet residual linear mixing can still entangle the two blocks.Source: §1 · p. 2
• SplitJEPA combines latent-space prediction with invariant-pair consistency to identify invariant and variant structure without reconstructing observations.The paper reports theoretical identification and experiments on controlled nonlinear systems and robotic manipulation.Source: §1 · p. 2

2 Problem Background and Formulation

• The latent state is divided into shared factors z_H and potentially changing factors z_L, with observations produced by an unknown smooth embedding.The encoder maps observations into a representation, while the latent parameterization remains undetermined.Source: §2 · p. 3
• Predictive pairs have context-controlled correlation and equal marginal distributions, while invariant correspondences preserve z_H and allow z_L to vary.The extent and directions of variation across invariant pairs determine whether they can identify the partition.Source: §2 · p. 3
• The target is identification up to independent within-block transformations, rather than recovery of individual coordinates within each block.This block-orthogonal target allows transformations within the invariant and variant blocks.Source: §2 · p. 3
• Observation-only learning cannot identify the invariant–variant decomposition because a smooth nonlinear transformation can mix variant information into the first latent block while preserving observations.The resulting parameterizations are observationally equivalent but are not related by independent within-block transformations when the mixing derivative is nonzero.Source: §2 · pp. 3–4

3.1 Predictive Identification and the Remaining Ambiguity

• Predictive identification recovers the complete latent state up to a global orthogonal transformation, but that rotation can still mix invariant and variant directions.The remaining cross-block components are the ambiguity the partition-identification step must eliminate.Source: §3.1 · p. 4

3.2 Identifying the Partition through Variation

• Under predictive dynamics and full-rank variant changes, SplitJEPA identifies invariant and variant subspaces up to independent within-block orthogonal transformations without reconstruction.The full-rank condition means no nonzero linear direction in the variant block remains unchanged across all available pairs.Source: §3.2 · p. 5
• Complementary changes across matched pairs can jointly span the variant subspace, allowing invariant-block agreement to remove dependence on variant factors.A single pair need not vary every coordinate; the aggregate changes must cover the variant directions.Source: §3.2 · p. 6

3.3 Robustness under Variant Perturbations

• A predictor using only the invariant representation is locally insensitive to variant perturbations, provided the perturbed state remains within the modeled support.The guarantee applies only to changes confined to the variant factor; changes to shared task content fall outside it.Source: §3.3 · pp. 6–7

3.4 Estimation and Normalization

• The estimation objective combines predictive alignment, invariant-block agreement, and variance preservation, while finite-weight variance regularization does not impose exact whitening.The variance regularizer discourages collapse, and invariant consistency encourages agreement across matched pairs.Source: §3.4 · p. 7

4.1 Synthetic Evaluation

• With the predictive encoder fixed, SplitJEPA achieves near-perfect synthetic block recovery with little cross-block leakage, while LeJEPA retains block mixing and shuffled pairs disrupt separation.Global R² is approximately 0.9977 ± 0.0016 for all methods, so the comparison concerns organization rather than globally available information.Source: §4.1 · p. 8, Table 1 · p. 8
• Ablations show that prediction supports global recovery, invariant-pair recovery supports block separation, and variance regularization with ZCA normalization supports covariance control.Scaling experiments report that more predictive training improves global recovery, more recovery samples improve block separation with the encoder fixed, and larger latent spaces make complementary-block estimation more sensitive.Source: §4.1 · p. 8, Figure 5 · p. 8

4.2 Robotic Manipulation Evaluation

• Both proposed representations improve validation action prediction over baselines, while closed-loop gains are clearest under visual shifts and especially pronounced on PushCube.PickCube remains more challenging, and whether retaining both blocks helps depends on the shift.Source: §4.2 · p. 9, Table 2 · p. 9, Figure 6 · p. 9
• Robotic ablations find that removing invariance most degrades PushCube camera-shift success, while removing diversity has task-dependent effects across PushCube and PickCube.The population guarantee requires informative variation among exact invariant correspondences; practical cross-trajectory pairing trades variation coverage against approximate alignment quality.Source: §4.2 · p. 9

5 Conclusion

• SplitJEPA recovers invariant and variant latent structure without reconstruction, with predictive representation learning and invariant-block consistency jointly yielding block identifiability under sufficient variation.The recovered structure isolates computations based on the invariant block from variant perturbations.Source: §5 · p. 9
• The experiments cover controlled nonlinear systems and robotic manipulation, but their scale limits the work; larger, complex real-world scenarios remain a next step.Source: §5 · p. 9

A Related Work

• Prior predictive recovery identifies the complete latent state only up to a global orthogonal transformation; paired correspondences resolve residual mixing between invariant and variant subspaces.The resulting subspaces are identified up to independent block-wise isometries, without an observation decoder.Source: A Related Work · p. 12

• A smooth nonconstant function of the variant factor can be added to the invariant factor, making the alternative invariant block depend on the variant factor.The derivative of the added function is nonzero when that dependence is present.Source: §B.1 · p. 12
• The construction defines a global smooth diffeomorphism and an alternative observation-generating map for the transformed latent variables.Source: §B.1 · p. 12
• Pushing the latent distribution through this transformation yields the same observed marginal distribution under the original and alternative latent models.Source: §B.1 · p. 13
• Thus, the observation law alone cannot distinguish the true invariant factor from an alternative containing nonlinear information from the variant factor.An encoder's extracted invariant representation therefore cannot be identified with the true factor from that law alone.Source: §B.1 · p. 13

• Under the predictive model and stated assumptions, a global population minimizer is constrained to the centered and whitened representation class.The proof uses stationary Gaussian predictive dynamics, with conditional correlation strictly between zero and one.Source: §B.2 · p. 13
• Hermite spectral analysis shows that all variance at a global optimum lies in degree one, so the representation is linear in the latent state.The covariance normalization then makes the square linear map orthogonal.Source: §B.2 · p. 14
• Shared invariant correspondences remove cross-block dependence from the initially mixed linear representation.The proof strategy combines predictive dynamics with shared trajectory structure, then uses orthogonality to remove residual leakage.Source: §B.2 · pp. 13–14
• The invariant and variant latent subspaces are identified up to independent within-block orthogonal transformations, without observation reconstruction.Source: §B.2 · pp. 13, 15

B.3 Proof of Theorem 3

• A differentiable predictor that depends only on the invariant representation has zero local sensitivity to perturbations in the variant latent factor.Consequently, such perturbations cannot propagate through the learned invariant representation.Source: §B.3 · p. 15

C.1 Synthetic Setting Details

• The invariant block is recovered from the bottom d_H eigenspace of the matched-pair difference moment, with its orthogonal complement used as the variant block.Under the exact block model, invariant correspondences satisfy Δz_H = 0, making the population difference moment's null space the invariant subspace.Source: §C.1 · p. 16
• When paired variant states are equal, the difference covariance has no informative direction, so recovery defaults to the identity head and the invariant subspace is not uniquely determined.In this implementation, that boundary case is operationally identical to LeJEPA.Source: §C.1 · p. 16

C.2 Loss Ablations

• Without matched-pair recovery, Global R2 remains 0.9977 ± 0.0016, while invariant and variant R2 fall and leakage rises to 0.3320 ± 0.2152 and 0.3300 ± 0.2149.Removing predictive alignment decreases Global R2; removing variance preservation leaves mean block scores near the full model but increases covariance error and A21 variability.Source: §C.2 · pp. 16–17

C.3 Scaling Diagnostics

• At d = 64, normalized ∥A21∥F reaches 0.0325, while both block-level R2 values remain above 0.994 and normalized ∥A12∥F stays below 1.5 × 10^-3.Global R2 exceeds 0.9999 for d ≥ 16; increasing dimension makes estimation of the complementary block more sensitive.Source: §C.3 · p. 17, Figure 7 · p. 17
• With encoder training budgets from 2.56 × 10^5 to 5.12 × 10^6 OU pairs, Global R2 increases, with a clear transition around 1.28 × 10^6 and limited improvement beyond 2.56 × 10^6.Block-level leakage decreases in parallel.Source: §C.3 · p. 17, Figure 7 · p. 17
• With 250 versus 30,000 head-fitting samples, normalized ∥A21∥F falls from 0.156 to 0.012, while normalized ∥A12∥F remains below 0.005 throughout.Global and block-level R2 curves remain comparatively stable; more samples primarily improve whitening and complementary-block estimation.Source: §C.3 · p. 17, Figure 7 · p. 17

C.4 Efficiency at Matched Recovery Quality

• The implementations use the same latent dimension, five seeds, and 256 atomic observations per optimization step, but each method uses its native grouping structure.SplitJEPA uses matched pairs with shared content and varying style; V3 uses groups with shared style and varying content. Runtime includes each method's training and recovery stages, with model construction and held-out evaluation excluded.Source: §C.4 · p. 18
• SplitJEPA reaches the target in all five runs versus three for V3, with 2.70× lower capped runtime and a 6.33× reduction in observation cost.Capped runtime is 3.16 ± 1.14 seconds for SplitJEPA and 8.54 ± 3.05 seconds for V3; observation cost is 0.307 ± 0.114 million versus 1.946 ± 0.668 million.Source: Figure 12 · p. 18, §C.4 · pp. 18–19
• SplitJEPA uses 64 trainable parameters, compared with 272 for V3.Both are evaluated at checkpoints across a 10,000-step budget; runs that do not reach the target are assigned full-budget runtime and observation count.Source: §C.4 · pp. 18–19

C.5 Additional Robustness

• Under joint nonlinear MLP mixing, SplitJEPA achieves the best invariant and variant R2 values, with ∥A12∥F = 0.0056 ± 0.0007 and ∥A21∥F = 0.0187 ± 0.0037.The compared methods reuse the same frozen MLP encoder; LeJEPA retains larger cross-block leakage, while Shuffled Pairs further degrades block alignment and purity.Source: §C.5 · p. 19, Table 3 · p. 20
• At coupling strength 2.0, Global R2 falls to 0.8765, with orthogonality error and cross-block terms increasing further.The experiment links departures from the globally identifiable linear solution class under finite optimization with increased orthogonality error and cross-block leakage.Source: §C.5 · pp. 19–20, Table 4 · p. 20

D.1 Robotic Manipulation Setting Details

• PickCube and PushCube datasets contain 800 successful trajectories per task, grouped into 100 sets of eight with shared object identity, initial pose, and goal pose.Cross-view pairs match the same physical state across cameras; cross-trajectory pairs approximately align frames by normalized execution progress and relax exact state alignment.Source: §D.1 · p. 20
• The robotic representation uses an action-conditioned predictor, applies the invariance loss only to r_H, and uses variance and covariance regularization to prevent collapse.The learned partition has equal-sized blocks: 64 dimensions for PickCube and 128 for PushCube.Source: §D.1 · pp. 20–21
• Policies are trained by behavior cloning on ID demonstrations and evaluated across ID and three OOD visual conditions without OOD fine-tuning or adaptation.The evaluation uses 100 closed-loop episodes per condition; visual changes include camera azimuth rotation, reduced lighting intensity, or both.Source: §D.1 · p. 21

Controlled Intervention

• With environment state and other policy inputs fixed, the 𝑟𝐻-only policy retains approximately 100% of clean performance under 𝑟𝐿 corruption, while 𝑟𝐻+𝑟𝐿 becomes progressively more sensitive.The intervention is confined to 𝑟𝐿; exposing that block adds a pathway through which perturbations can affect downstream decisions.Source: Table 5 · p. 22, §D.2 · pp. 22–23

Block Validation

• On PickCube, viewpoint responses are 0.18 in 𝑟𝐻 versus 1.51 in 𝑟𝐿 (8.4×); matched-trajectory responses differ 2.1×, versus 14.0× and 3.1× on PushCube, respectively.Figure 14 validates block specialization using normalized representation distances for matched viewpoint and trajectory variations.Source: Figure 14 · p. 22, §D.2 · p. 22

Sample Efficiency

• Across evaluated budgets, both Ours-𝑟𝐻 and Ours-𝑟𝐻+𝑟𝐿 outperform baselines; on PushCube, both exceed 90% ID SR@10cm with only 10% of demonstrations.On PickCube, baseline OOD performance reaches a low plateau as training data increases, while the proposed representations consistently outperform baselines; neither input dominates uniformly.Source: §D.2 · pp. 21–22

Temporal Transferability

• Across both tasks and all temporal gaps, 𝑟𝐻 changes less than 𝑟𝐿 between state-matched asynchronous observations, while both blocks change more as timestep gaps increase.Pairs come from different validation trajectories and normalized timesteps; results are averaged over three seeds without additional training, and the smaller shift makes 𝑟𝐻 a more transferable prediction interface.Source: Figure 15 · p. 23, §D.2 · p. 23