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

1 Introduction

• The paper identifies a gap: JEPA latents capture slow features but lack an interpretable progression axis or a shared progression manifold across episodes of the same task.It asks whether structured progression can be added without interference with LEWM’s SIGReg anti-collapse mechanism, and whether planning improves.Source: §1 · p. 2
• SD-JEPA splits the latent into a low-dimensional progression subspace shaped by a cosine-margin triplet loss and a high-dimensional content subspace regularised by SIGReg.The split assigns the two subspaces disjoint roles.Source: §1 · p. 2
• The authors prove that SIGReg on content and triplet loss on progression have disjoint gradient supports, so the anti-collapse forces act on separate latent coordinates.They also connect the losses through sample/dimension-contrastive duality [2023] and report a subspace-ablation falsifier.Source: §1 · p. 2
• SD-JEPA matches or improves on LEWM on most control benchmarks at matched compute, using a planning cost that separates content matching from optional angular progression matching.The paper reports that the 1-D angular coordinate tracks task progress, regresses during genuine task regression, and separates surprise from post-perturbation relocalisation.Source: §1 · p. 2

2.1 Notation and setting

• Offline trajectories provide pixel observations and action labels; an encoder maps observations to latents, a predictor models latent dynamics, and CEM plans toward a goal embedding.Planning minimizes a latent-space cost to the goal embedding.Source: §2.1 · p. 3

2.2 LeWM: next-embedding prediction with SIGReg

• LEWM jointly trains its encoder and predictor with next-embedding prediction and SIGReg, which constrains the marginal latent distribution toward an isotropic Gaussian without affecting the conditional dynamics distribution.SIGReg aggregates univariate Epps–Pulley normality statistics over random unit directions using the Cramér–Wold theorem.Source: §2.2 · p. 3

2.3 Duality of contrastive and covariance-based regularizers

• Under double-L2 normalization, the sample- and dimension-contrastive criteria are related by Lnc = Lc + N − M; SIGReg and cosine-triplet belong to the broader variance-repulsion family.This connection motivates combining the two regularisers in the method.Source: §2.3 · p. 3

3.1 Subspace decomposition of the latent

• The latent is decomposed into orthogonal progression and content subspaces, with the default canonical split assigning the first k coordinates to progression and the remaining D−k to content.The default progression dimension is k = 2.Source: §3.1 · p. 3

3.2 Four-term training objective

• The four-term objective combines full-latent next-step prediction, SIGReg on content, cosine-margin triplet ordering on progression, and optional temporal straightening.Triplet positives are sampled near an anchor in time, while negatives fall outside the temporal window or come from another trajectory; the LEWM base uses SIGReg on the full latent.Source: §3.2 · pp. 3–4

3.3 Latent-space indicators as predictor conditioning

• Angular and radial readouts summarize the state on the progression manifold and are supplied to the predictor alongside action embeddings as additional conditioning.The angle is encoded as sine and cosine to avoid wrap-around discontinuities before embedding and summation with the action embedding.Source: §3.3 · p. 4

3.4 Planning with an angular goal-matching cost

• The planning cost separates content, radial, and angular progression matching, making the progression match scale-invariant and independent of ambient content geometry.The receding-horizon CEM solver otherwise follows LEWM, sampling 300 action sequences per step and retaining the top 30 as elites; it runs 30 iterations on Push-T and 10 elsewhere.Source: §3.4 · p. 4

4 Theoretical analysis

• SIGReg applied to content and a loss applied to progression have gradients confined to their respective orthogonal subspaces.The result follows because each loss depends on the latent only through its own projection.Source: §4 · p. 4
• Removing either regulariser cannot be compensated for by increasing the other, because its gradients do not reach the affected coordinates.The proposition concerns latent-level gradients under the orthogonal split.Source: §4 · p. 4
• The ablation reports that removing the split and applying triplet loss to the full latent returns the gain to the LEWM baseline, while encoder-level conflict is not formally precluded.The encoder’s parameter gradients still sum both loss contributions through the chain rule.Source: §4 · p. 5

5.2 Main results

• At matched 10-epoch compute, planning success rises on three of four tasks: Two-Room 87→90, Reacher 86→88, Push-T 96→97.3, but Cube falls 74→72.The reported values are 3-seed means; the best progression dimension varies by environment.Source: Table 1 · p. 5, §5.2 · p. 5

5.3 The subspace split is load-bearing

• On single-seed Push-T, removing the split returns A2_full to A0 at 96, while applying the triplet to the full latent with the split yields 92.The ablation identifies the split, rather than the triplet alone, as the load-bearing ingredient.Source: §5.3 · p. 5, Table 2 · p. 6
• Optimal kprog varies by task: Push-T rises 94→96→97.3 for k={2, 4, 8}, while Two-Room is non-monotone and Reacher favors k=4.The reported pattern is that the appropriate progression dimension matches task complexity, rather than higher k always winning.Source: §5.4 · p. 6, Table 3 · p. 6

5.5 The latent acts as a scene-aware compass

• On Push-T, zcont separates by episode while zprog arcs partly mix; per-episode probes also predict target distance across four environments using 4.2% of the latent.The probe R2 values are 0.91, 0.91, 0.95, and 0.72 on Cube, Push-T, Reacher, and Two-Room, respectively.Source: §5.5 · pp. 6, 8, Table 5 · p. 8
• θ is environment-dependent: Cube target-distance correlation exceeds the clock (|ρ|=0.59 vs 0.47), while Push-T clock correlation falls from 0.91 at k=2 to 0.56 at k=8.The authors describe θ as partly clock-like and partly task-coupled, with the balance changing by environment and kprog.Source: §5.5 · pp. 6–7
• In Push-T teleport-and-continue, |∆θt| spikes about 75× baseline drift, then θ relocalises to a sector for the new solving mode; MSE-VoE reports only the spike.Splice, one-frame teleport, and reverse perturbations also reproduce this pattern.Source: §5.5 · p. 7, Figure 3 · p. 8
• On Cube, |∆θ| beats z-MSE at every contact-event tolerance, winning on 97.5% of episodes at the tightest tolerance.In the episode-500 visualization, z-MSE peaks midmanipulation while |∆θt| peaks at a contact transition.Source: Table 4 · p. 7, §5.5 · p. 7, Figure 2 · p. 7

5.6 Cross-environment summary and limitations

• Per-environment kprog tuning means results report best-k values rather than a universal k; the fixed orthogonal split and cosine geometry also bound the design choices explored.The authors identify learned orthogonal decomposition and alternative distance metrics or readouts as natural extensions.Source: §5.6 · p. 8

Latent world models

• SD-JEPA inherits LeWM’s encoder-predictor stack and SIGReg, adding a fixed orthogonal split and triplet term with no new learnable weights (18.04M parameters).The passage names LeWM among feature-space JEPAs, citing Maes et al. [2026].Source: §6 · p. 9

Progression geometry and temporal straightening

• The cosine-margin triplet on zprog belongs to temporal-ordering work [Dwibedi et al., 2019, Schneider et al., 2023], while its straightening regularizer matches Wang et al. [2026].The passage also notes latent straightening’s OOD-detection applications [Internò et al., 2025].Source: §6 · p. 9
• At matched protocol on Push-T, SD-JEPA at kprog=8 reaches 97.3 ± 1.2 versus 91.3 ± 2.5 for Wang et al. [2026]; adding curvature regularization costs 14 points.The authors relate that cost to their principle that anti-collapse forces should not compete on the same coordinates.Source: §6 · p. 9
• The no-double-counting result is grounded in the SSL duality framework of Garrido et al. [2023].The passage places this connection alongside the paper’s argument about complementary anti-collapse forces.Source: §6 · p. 9

7 Discussion

• θ advances with task progress, regresses when the agent backtracks, and under perturbation jumps to a semantically appropriate new sector.The authors describe it as a state-based navigation signal rather than merely a chronological counter; they make no mechanistic claim about the parallel to head-direction cells.Source: §7 · p. 9
• Planning gains are non-uniform: Reacher improves +3.3 pp and Cube +2.0, while Two-Room and Push-T gain nothing; θ-coupling is necessary but insufficient.The discussion attributes the Cube result to contact-rich rollout error and Push-T’s result to its near-ceiling score.Source: §7 · p. 9
• Unlike scalar prediction error, the angular readout reports a signed phase sector that is comparable across environments, and the tests find it complementary to prediction error.The authors describe broader geometric readouts of structured-latent world models as an open agenda.Source: §7 · p. 9

8 Conclusion

• At matched 10-epoch compute, SD-JEPA improves on three of four LEWM benchmarks: Two-Room +3, Reacher +2, Push-T +1.3, and Cube −2.The reported evaluation uses three seeds.Source: §8 · p. 9
• Across 40 cube episodes, |∆θt| beats z-MSE for semantic-event localisation (+0.18 pooled AUROC at ±1-step tolerance; 39/40 per-episode wins); a within-episode probe yields zprog R2 0.72–0.95 across four environments.The probe used 40 episodes per environment and only 4.2% of the latent; the angular readout adds information the scalar prediction-error metric does not supply.Source: §8 · p. 10

A Extended theoretical analysis

• The appendix includes further results on straightening consistency, progression-axis identifiability, and conditions under which the subspace split fails.It also expands the theoretical analysis and duality discussion.Source: A Extended theoretical analysis · p. 11

A.1 Non-conflict under the subspace split

• Orthogonal projections confine SIGReg’s gradient to the content subspace and a progression-only loss’s gradient to the progression subspace, preserving SIGReg’s marginal constraint.The argument is specific to the marginal-distributional term on zcont and a geometric term acting only on zprog; full-latent prediction loss can have content-subspace components.Source: §A.1 · p. 12

A.2 Relation to the sample/dimension-contrastive duality

• Triplet loss and SIGReg belong to a shared variance-repulsion family, yet their anti-collapse gradients are complementary rather than redundant on orthogonal subspaces.Removing either term cannot be compensated by increasing the other because its gradients do not act in the affected coordinate directions.Source: §A.2 · p. 12

A.3 Consistency with emergent straightening

• The straightening term softly raises mean consecutive-velocity cosine without forcing perfect straightness; its non-antagonism claim rests on prior reports and is identified as the weakest empirically.The proposition assumes LEWM’s reported increase toward a limit c⋆<1.Source: §A.3 · pp. 12–13

A.4 Identifiability of the progression axis

• θ is identifiable up to global rotation, and episodes with different durations or solving modes need not share a reference direction.The rotation freedom is harmless for the rotation-invariant canonical zcont-MSE planning cost.Source: §A.4 · p. 13

A.5 When does the subspace split fail?

• The split assumes coherent episode phase: exploratory rollouts lack useful ordering, branching goals may not fit one manifold, and oversized k can absorb content.Goal-conditioned variants are proposed for branching progressions.Source: §A.5 · p. 13

B Implementation notes

• SD-JEPA builds on LEWM’s encoder, predictor, action-conditioning, and SIGReg modules; its triplet samples positives and negatives within each batch with negligible added training-step overhead.Positives come from the anchor trajectory within a temporal window; negatives come from outside that window or another trajectory in the batch.Source: B Implementation notes · p. 13

C Hyperparameter summary

• Base training uses AdamW for 10 epochs, batch size 128, 224 × 224 images, and 192-dimensional embeddings; default SIGReg and triplet weights are 0.09 and 0.10.The learning rate is 5×10−5 with weight decay 10−3.Source: C Hyperparameter summary · pp. 13–14

D Additional auxiliary-term explorations

• In this setup, full-latent straightening costs about 14 Push-T points, polar conditioning adds a roughly 2-point drop, and the tested extras do not improve canonical A2.The auxiliary exploration also tested an angular planning cost and a first-anchor regularizer.Source: D Additional auxiliary-term explorations · p. 14

E Extended latent geometry across environments

• The extended geometry appendix reports per-environment t-SNE projections, state-space θ overlays, progress-proxy correlations, and selected frame strips from held-out episodes.The held-out episodes span the dataset’s full episode-id range.Source: E Extended latent geometry across… · p. 14

E.1 Push-T at kprog=2 (canonical A2)

• On Push-T A2 at kprog=2, θ is strongly clock-correlated except during regression, while zprog mixes episodes and zcont separates them.Episode 8000 is the regression case; the reported cross-episode geometry matches the disjoint-supports prediction.Source: §E.1 · p. 15, Figure 5 · p. 15, Figure 4 · p. 15

E.2 Push-T at kprog=8

• At kprog=8 on Push-T, θ spans 6.1 rad and is less clock-like than at kprog=2, with clock and task-proxy correlations both near 0.56.The progression radius is r ∈[0.62, 0.77] with per-episode standard deviation ∈[0.04, 0.07].Source: §E.2 · p. 16

E.3 Reacher at kprog=8

• On Reacher, θ correlates similarly with elapsed time and end-effector-to-target distance, but is essentially uncorrelated with joint speed.The mean |ρ| values are 0.53 for the clock, 0.49 for target distance, and approximately 0.08 for joint speed.Source: §E.3 · p. 16

E.4 Two-Room at kprog=8, the cyclic-phase example

• In Two-Room, episode 3000 traces a U-shaped θ trajectory despite Spearman ρ≈0, illustrating how a monotonic correlation measure can miss cyclic phase.Across episodes, coupling is bimodal: some are clock-like, some task-coupled, and one is cyclic.Source: §E.4 · p. 16, Figure 8 · p. 17

E.5 OGBench-Cube at kprog=8

• OGBench-Cube is the only tested environment where θ correlates more with a task-physical signal than with elapsed time, and its encoded phase varies across episodes.The reported mean |ρ| is 0.59 for block_target_dist versus 0.47 for step_idx; episodes show spatial-distance, gripper-state, and U-cycle patterns.Source: §E.5 · p. 17

• The head-direction-cell analogy is explicitly qualitative: both involve a circular phase coordinate and remapping after sensory restructuring, but no biological circuit rediscovery is claimed.The comparison concerns geometric similarity, not an asserted shared mechanism.Source: F Supplementary perturbation analyses · pp. 17–18
• The three perturbation modes use the same Push-T A2 encoder without test-time fine-tuning, isolating responses to splice, single-frame teleport, and reversed-window inputs.The main teleport-and-continue experiment is the reference for separating the surprise moment from post-perturbation relocalisation.Source: F Supplementary perturbation analyses · p. 18
• ≈0.17 rad is the splice-boundary surprise spike, followed by elevated noise and a new θ path consistent with the donor episode’s progression.The donor was not selected for maximum angular contrast, making the spike more modest than in teleport-and-continue.Source: F Supplementary perturbation analyses · p. 19, Figure 11 · p. 19
• ≈0.20 rad is the single-frame teleport spike, and θ returns to baseline within one step when the original episode resumes.The transient perturbation produces transient relocalisation.Source: F Supplementary perturbation analyses · p. 19, Figure 12 · p. 20
• ≈0.11 rad and ≈0.13 rad spikes mark the boundaries of a reversed 10-frame window, although the frames themselves are unchanged.The paired responses support interpreting the signal as sensitivity to violated dynamics rather than scene novelty.Source: F Supplementary perturbation analyses · p. 19, Figure 13 · p. 20

G Phase-event localisation vs. classical surprise: extended results

• On 40 Cube episodes, |∆θ| better localizes semantic events, z-MSE better detects raw magnitude anomalies, and combining them improves neither test.The phase-event comparison uses gripper-contact transitions as ground-truth semantic boundaries.Source: G Phase-event localisation vs.… · p. 20, Figure 14 · p. 21

G.1 Phase-alignment summary across all 40 episodes

• At ±1-step tolerance, both AUROC distributions sit below 0.5; widening the window brings |∆θ| above 0.5 before z-MSE.Only approximately 11% of steps are positive labels, and z-MSE peaks during stable contact rather than at transitions.Source: §G.1 · p. 21

G.2 Action-corruption AUROC: the magnitude-anomaly negative control

• On action-shuffle corruption, z-MSE wins because the magnitude anomaly favors the 192-dimensional error signal over a one-dimensional phase scalar; the combined score does not beat it.This test rewards information bandwidth rather than semantic alignment.Source: §G.2 · p. 21

G.3 Regime-change segmentation (change-point detection)

• At loose tolerance, angular velocity beats both prediction-error metrics in change-point F1, while BinSeg locates z-MSE spike-and-return boundaries more accurately at tight tolerance.The useful signal depends on whether the target is a smoother regime shift or a sharp magnitude spike.Source: §G.3 · pp. 21–22, Figure 15 · p. 22

G.4 Linear-probe information density: full per-feature table

• Per-episode zprog reaches mean R2 +0.905 with positive R2 in 100% of episodes, while pooled probes show it is not a globally calibrated absolute-distance estimate.The per-episode results come from leave-one-step-out cross-validation across 40 episodes; zcont and the full latent overfit in the pooled regime.Source: §G.4 · pp. 22–23, Figure 16 · p. 22