Orthogonal JEPA: Factorized Predictive States for Latent World Models

Abstract: 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.

• The paper introduces Orthogonal JEPA, a framework that decomposes target states into distinct, linearly independent coordinate subspaces, enhancing predictive performance by ensuring...],

Latent world models and the monolithic capacity problem

Latent world models formulate state spaces optimized for anticipating unobserved consequences of context observations. Joint-embedding predictive architectures (JEPAs) construct these states by operating directly within representation space, discarding fine-grained, high-variance raw observation details that do not contribute to downstream predictive utility. Standard JEPA formulations channel the entire predictable target state through a monolithic representation vector and a single prediction pathway. In complex natural and physical systems characterized by multi-scale dynamics, heterogeneous entity interactions, and asymmetric variance profiles, this monolithic structure creates a fundamental capacity-allocation failure. Dominant or easily predictable signals monopolize representation capacity, leading to redundant latent coordinates and weak or conflicting gradient updates for subtle predictive components.

Orthogonal JEPA (2608.20065) resolves this bottleneck by framing latent predictive modeling as an orthogonal capacity-allocation problem. Rather than regressing a single unstructured target vector, the framework decomposes target states into distinct, linearly independent coordinate subspaces via learned basis matrices. Each subspace is predicted via a dedicated branch before being synthesized back into a complete latent state. By enforcing both intra-factor and inter-factor orthogonality alongside variance regularizers, the model ensures that the latent capacity is distributed across diverse predictive factors without collapsing into redundant target directions.

Orthogonal predictive factorization framework

The architecture interfaces with diverse data modalities through a domain adapter that maps raw observations into content tokens and optional structural descriptors, such as spatial coordinates, timestamps, or physical entity identifiers. Given a sampled context index set and a target index set, an online encoder parameterized by weights $\theta$ maps the context to a representation $z_c$, while an exponential moving average (EMA) target encoder parameterized by $\bar{\theta}$ produces target embeddings $z_t \in \mathbb{R}^d$. Crucially, gradients are blocked from propagating directly into the target encoder output $\widetilde{z}_t = \operatorname{sg}(z_t)$.

Orthogonal predictive factorization divides the latent space into $K$ distinct factors, each of rank $r$ such that $K r = d$. A set of learned basis matrices $B_k \in \mathbb{R}^{d \times r}$ projects the target embedding into factorized coordinates:

$$z_t^{(k)} = B_k^\top \widetilde{z}_t, \quad k = 1, \dots, K$$

A dedicated prediction branch $z_c$0 maps the shared context representation $z_c$1 and target descriptor $z_c$2 to the corresponding target subspace estimate $z_c$3. The individual predictions are concatenated into $z_c$4, and the full latent state $z_c$5 is synthesized using the Moore–Penrose pseudoinverse of the concatenated basis $z_c$6:

$z_c$7

When exact orthogonality is achieved, $z_c$8, reducing synthesis to an additive reconstruction $z_c$9 where total energy is preserved ($\bar{\theta}$0).

$B_k \in \mathbb{R}^{d \times r}$2

The overall training objective combines direct predictive regression with algebraic and empirical regularizers:

$\bar{\theta}$1

Predictive regression uses a mean squared error loss over all factor branches:

$\bar{\theta}$2

The structural regularizer $\bar{\theta}$3 penalizes deviations from intra-factor orthonormality and inter-factor orthogonality:

$\bar{\theta}$4

To prevent informational collapse, the factor-activity loss $\bar{\theta}$5 penalizes projected target coordinates whose standard deviation falls below a threshold $\bar{\theta}$6, and the encoder variance loss $\bar{\theta}$7 enforces coordinate-wise variance above $\bar{\theta}$8 directly on the online context representation.

Experimental evaluation across predictive modalities

Controlled visual binding and compositional readout

The visual evaluation assesses whether factorized representations mitigate binding failures—specifically disentangling spatial location ("where") from visual operation ("what") under block masking on MuJoCo visual scenes (2608.20065). Readouts are evaluated using an injective Hungarian alignment metric on held-out cells (INJ), representation collapse rate (Coll.), and grid recovery (Rec.) across DINOv3 (Siméoni et al., 13 Aug 2025) and SigLIP 2 (Tschannen et al., 20 Feb 2025) backbones.

Backbone	Pretraining Strategy	Grid Supervision	INJ $\bar{\theta}$9	Coll. $z_t \in \mathbb{R}^d$0	Rec. $z_t \in \mathbb{R}^d$1
DINOv3	Frozen Checkpoint	Known Grid	0.452	0.667	—
DINOv3	Frozen Checkpoint	Learned Grid	0.569	0.433	0.643
DINOv3	Standard JEPA	Learned Grid	0.572	0.426	0.645
DINOv3	Orthogonal JEPA	Learned Grid	0.581	0.417	0.659
SigLIP 2	Frozen Checkpoint	Known Grid	0.476	0.656	—
SigLIP 2	Frozen Checkpoint	Learned Grid	0.484	0.511	0.676
SigLIP 2	Standard JEPA	Learned Grid	0.483	0.514	0.679
SigLIP 2	Orthogonal JEPA	Learned Grid	0.490	0.503	0.688

Orthogonal JEPA demonstrates consistent gains over monolithic JEPA pretraining, raising INJ accuracy to 0.581 on DINOv3 and 0.490 on SigLIP 2 while suppressing representation collapse rates to 0.417 and 0.503, respectively. The implication is that dedicating distinct predictive branches to orthogonal target subspaces prevents dominant visual features from washing out compositional spatial signals, yielding more linearly separable representations for downstream readouts.

Single-cell transcriptomics and perturbation response

In single-cell biology, masked gene expression profiles are predicted using an scGPT backbone (2608.20065) pretrained on roughly 800,000 human kidney cells. The model is evaluated on PBMC-10K cell clustering (AvgBIO) under finetuned and zero-shot settings, alongside post-perturbation response forecasting on the Norman and Adamson CRISPR datasets evaluated by Pearson correlation.

Model	JEPA Formulation	Orthogonal Factors	PBMC Finetuned (AvgBIO) $z_t \in \mathbb{R}^d$2	PBMC Zero-shot (AvgBIO) $z_t \in \mathbb{R}^d$3	Norman (Pearson) $z_t \in \mathbb{R}^d$4	Adamson (Pearson) $z_t \in \mathbb{R}^d$5
scGPT	None	No	0.7531	0.5288	0.631	0.905
Cell-JEPA	Monolithic	No	0.7830	0.7194	0.787	0.937
Orthogonal JEPA	Factorized	Yes	0.8001	0.7452	0.798	0.942

Factorizing the predictive cell target yields substantial zero-shot transfer gains, achieving an AvgBIO score of 0.7452 compared to 0.7194 for monolithic Cell-JEPA and 0.5288 for generative baseline scGPT. Similarly, perturbation forecasting accuracy reaches 0.798 on the Norman benchmark. This demonstrates that orthogonal target factorization prevents highly expressed housekeeping genes from dominating the loss gradients, preserving capacity for low-variance regulatory gene networks.

Longitudinal clinical event forecasting

Forecasting patient trajectories involves projecting time-ordered clinical records, demographics, and lab measurements into the future. A frozen latent state is decoded across a vocabulary of more than 1,000 discrete future clinical codes, evaluated using mean Precision-Recall Area Under the Curve (PRAUC).

Model Family	Specific Model Architecture	Mean PRAUC $z_t \in \mathbb{R}^d$6
Classical ML	Random Forest	0.602
Classical ML	XGBoost	0.667
LLM	Qwen2.5-0.5B	0.648
LLM	Qwen3-0.8B	0.651
Autoregressive Trajectory	Prophet	0.680
Autoregressive Trajectory	Delphi	0.689
Multimodal Fusion	Cross-Attention Fusion	0.702
Latent World Model	Standard JEPA (Monolithic)	0.711
Latent World Model	Orthogonal JEPA (Factorized)	0.718

Orthogonal JEPA achieves a mean PRAUC of 0.718, outperforming both general-purpose language modeling approaches and standard monolithic JEPA (0.711). By factorizing the future clinical target vector, the model successfully isolates independent disease progression pathways instead of collapsing distinct pathological timelines into an aggregated average clinical state.

Continuous control dynamics and planning

In model-based continuous control, latent world models are trained offline on 500 random trajectories from MuJoCo benchmarks and leveraged for planning via the Cross-Entropy Method (CEM) over an 8-step planning horizon (2608.20065).

Model Architecture	Walker2d-v5 $z_t \in \mathbb{R}^d$7	HalfCheetah-v5 $z_t \in \mathbb{R}^d$8	InvertedPendulum-v5 $z_t \in \mathbb{R}^d$9
Standard JEPA	$\widetilde{z}_t = \operatorname{sg}(z_t)$0	$\widetilde{z}_t = \operatorname{sg}(z_t)$1	$\widetilde{z}_t = \operatorname{sg}(z_t)$2
Orthogonal JEPA	$\widetilde{z}_t = \operatorname{sg}(z_t)$3	$\widetilde{z}_t = \operatorname{sg}(z_t)$4	$\widetilde{z}_t = \operatorname{sg}(z_t)$5

The numerical gap in control is pronounced: on Walker2d-v5, Orthogonal JEPA achieves a mean return of 45.1 compared to 4.9 for monolithic JEPA, with corresponding improvements on InvertedPendulum-v5 ($\widetilde{z}_t = \operatorname{sg}(z_t)$6 vs. $\widetilde{z}_t = \operatorname{sg}(z_t)$7). In continuous control, errors in high-frequency velocity coordinates frequently corrupt action conditioned trajectory rollouts when merged with low-frequency spatial coordinates. Orthogonal factorization isolates these dynamic scales into separate predictive branches, preventing gradient cancellation and producing stable latent rollouts during CEM trajectory optimization.

Equivariant molecular dynamics forecasting

Autoregressive molecular rollout requires predicting future atomic states while respecting $\widetilde{z}_t = \operatorname{sg}(z_t)$8 Euclidean symmetry. Building upon TrajCast (2608.20065), Orthogonal JEPA is applied across compatible irreducible representations and evaluated across liquid water, crystalline quartz, gas-phase paracetamol, and benzene systems. Metrics comprise one-step displacement Mean Absolute Error (MAE) and median final-position Root-Mean-Square Deviation (RMSD) after 100 unconstrained autoregressive steps.

System	Metric	Scratch Baseline	TrajCast-JEPA (Monolithic)	Orthogonal JEPA (Factorized)
Water ($\widetilde{z}_t = \operatorname{sg}(z_t)$9)	1-step MAE ($K$0) $K$1	0.00387	0.00452	0.00376
Water ($K$2)	100-step RMSD ($K$3) $K$4	3.331	2.536	2.459
Quartz ($K$5)	1-step MAE ($K$6) $K$7	0.01080	0.01043	0.01011
Quartz ($K$8)	100-step RMSD ($K$9) $r$0	2.089	1.912	1.877
Paracetamol	1-step MAE ($r$1) $r$2	0.00901	0.00777	0.00765
Paracetamol	100-step RMSD ($r$3) $r$4	3.155	1.868	1.846
Benzene	1-step MAE ($r$5) $r$6	$r$7	$r$8	$r$9
Benzene	100-step RMSD ($K r = d$0) $K r = d$1	0.0958	0.0701	0.0699

Standard monolithic TrajCast-JEPA degrades one-step MAE on liquid water relative to training from scratch ($K r = d$2 vs. $K r = d$3), despite improving long-horizon RMSD. Orthogonal JEPA eliminates this trade-off, securing the lowest error across both short-term displacement ($K r = d$4) and 100-step trajectory drift ($K r = d$5). This indicates that maintaining factorized equivariant channels prevents high-frequency vibrational modes from degrading macro-structural geometric drift during iterative rollout.

Assumptions, limitations, and open questions

The theoretical foundation and empirical behavior of Orthogonal JEPA rest on specific structural assumptions and leave several operational questions open:

• Geometric orthogonality versus statistical independence: The basis penalty strictly enforces linear geometric separation ($K r = d$6). This does not guarantee statistical independence, non-linear disentanglement, or causal modularity among the learned factors, leaving open the question of whether true causal isolation requires non-linear coordinate projections.
• Marginal variance versus covariance rank: The variance regularizers $K r = d$7 and $K r = d$8 penalize coordinate-wise standard deviations along marginal axes. They do not penalize off-diagonal sample covariance within an individual factor branch, leaving the latent space vulnerable to partial rank deficiency within sub-factors.
• Synthesis conditioning during iterative rollouts: Reconstructing the complete latent state depends on the pseudoinverse $K r = d$9. If $B_k \in \mathbb{R}^{d \times r}$0 drifts from strict orthogonality during training, ill-conditioning can amplify synthesis errors. This creates stability concerns when unconstrained autoregressive rollouts extend beyond hundreds of steps.
• Deterministic prediction and multimodal distributions: The prediction heads $B_k \in \mathbb{R}^{d \times r}$1 implement deterministic point estimates, meaning the framework does not explicitly model stochastic branching or multimodal distributions in future states.
• Domain coverage: While evaluated on continuous control, molecular trajectories, and tabular-token systems, the empirical validation omits high-dimensional pixel-based closed-loop visual control, continuous partial differential equation fields, and explicit causal discovery benchmarks.

Conclusion

Orthogonal JEPA addresses capacity misallocation in joint-embedding world models by decomposing monolithic target states into linearly independent, factorized prediction channels. Enforcing structural orthogonality and variance constraints across dedicated prediction heads prevents high-variance signals from drowning out subtle predictive dynamics. Across controlled visual binding, single-cell perturbation response, longitudinal health forecasting, continuous planning, and equivariant molecular dynamics, this factorized predictive state design systematically improves representation stability and rollout fidelity over monolithic target baselines.