JEPA-Anything: Learning Predictive Models across Different Worlds

AI summary

Overview

Research area: Self-supervised representation learning and latent world modeling, applied across vision, biology, clinical medicine, control, molecular dynamics, physical fields, and weather (submitted under the arXiv category cs.CL).

Technical level: Advanced. The paper combines joint-embedding predictive architectures, subspace decomposition, and domain-specific scientific evaluation; readers need working familiarity with JEPA-style self-supervised learning and latent dynamics.

Scope: A single paper introducing orthogonal predictive factorization (OPF) as a domain-agnostic predictive core, instantiated and evaluated across seven domains spanning perceptual, dynamical, and scientific tasks.

What This Paper Is About

Joint-embedding predictive architectures (JEPA) learn by predicting the latent representation of a target from a context, but they typically compress everything there is to predict into one monolithic target embedding. When a system contains local and global structure, multiple entities, or changes at different scales, that single embedding forces weakly predictable structure to compete with dominant or high-variance structure for the same predictive capacity. The authors ask whether one shared latent world-model interface can instead organize the predictive state into multiple complementary components, so that the same learning principle works across radically different systems — from images and cells to molecules, physical fields, and weather.

Key Contributions

1. A common context–target state interface for latent world modeling. Domain adapters handle observation geometry (patches, timestamps, graph positions, entity identities) and produce tokens plus structural descriptors; everything downstream follows one algorithm.

2. Orthogonal predictive factorization (OPF). The target latent state is decomposed into K learned, approximately orthogonal subspaces, each with its own predictor, and the factor predictions are recombined via pseudoinverse synthesis into a complete latent state.

3. A composite regularized objective. Within-factor and cross-factor orthogonality terms, a per-factor activity floor, and an online encoder variance term are added to each domain's original base loss, yielding stable state synthesis plus factor-level diagnostics.

4. Cross-domain instantiation and validation. The same core is applied to vision, single-cell biology, clinical trajectories, intervention-conditioned control, molecular dynamics, PDE/physical fields, weather, and orbital dynamics, including biological wet-lab validation of a model-nominated intervention.

Main Findings

• Consistent gains on matched dynamics tasks: Against matched JEPA baselines, JEPA-Anything improves reported metrics on all 10 dynamics tasks in the benchmark.

• Large intervention-prediction improvement: On CITRIS Interventional Pong, single-intervention one-step MSE drops from 0.009541 to 0.006218 (a 34.8% reduction), combined-intervention MSE from 0.009441 to 0.008223 (12.9%), and six-step free-rollout MSE from 0.009478 to 0.008665 (8.6%).

• Best molecular rollout accuracy: Across four molecular systems (water, quartz, paracetamol, benzene), the method achieves the lowest one-step and 100-step errors among compared methods, with 100-step free rollouts showing improved stability.

• Better visual binding: With the readout held fixed, factorized pretraining raises injective held-out-cell accuracy and grid recovery while lowering collapse for both DINOv3 ViT-S and SigLIP2 Base encoders.

• Stronger single-cell representations: On PBMC-10K, AvgBIO reaches 0.8301 finetuned and 0.7752 zero-shot, versus 0.7830/0.7194 for Cell-JEPA and 0.7531/0.5288 for scGPT. Perturbation-response Pearson correlation reaches 0.814 (Norman) and 0.942 (Adamson), the highest among compared models.

• Broad clinical forecasting: Predicting orthogonal factors and synthesizing one future patient state yields higher mean PRAUC than a matched monolithic health-state JEPA across a vocabulary of more than 1,000 clinical events.

• Biological hypothesis confirmed in the lab: A factor-nominated intervention received experimental support in cell co-cultures, patient-derived organoids, tumor fragments, and mice.

• Recovery of a physical law: Latent orbital modes recover the Keplerian scaling exponent with a fitted slope of −1.4991, close to the theoretical −1.5.

• Theoretical guarantee for synthesis: Under exact orthogonality, the factor spaces form an orthogonal direct sum, preserve the target norm, and reconstruct the state exactly; the synthesis map has condition number 1, so prediction errors in factors translate one-to-one into latent-state errors. Unconstrained multi-head variants offer no such guarantee.

Methodology in Plain English

The framework separates what varies by domain from what stays the same.

Each domain supplies an adapter that turns raw data into tokens and structural descriptors, plus a view sampler that picks a context and a target — for example visible patches predicting masked patches, patient history predicting a future patient state, or a current molecular configuration predicting the next one. From there, a single procedure runs: an online encoder summarizes the context, an exponential-moving-average target encoder (with no gradients) provides the target latent, and the target is projected onto K learned subspaces of equal width.

Unlike standard JEPA, which uses one predictor against one monolithic target embedding, JEPA-Anything assigns one predictor per factor and then recombines the predicted factors with a pseudoinverse of the projection matrix to rebuild the full latent state. Training adds four loss terms to whatever loss the domain normally uses: a per-factor regression loss, an orthogonality loss that keeps directions within a factor orthonormal and across factors non-overlapping, a factor-activity floor that prevents any projected coordinate from going dead, and an encoder variance floor that discourages representation collapse.

This produces two usage modes. For ordinary downstream tasks, only the online encoder is kept, and each domain attaches its own readout (pooling, probing, decoding). For world-modeling tasks — future-state prediction, intervention forecasting, planning, autoregressive rollout — the projectors and predictors are retained, and the synthesized latent state is fed back into the next transition step. Factor coordinates can also be read out directly for scientific analysis, such as identifying candidate biological interventions or extracting latent spectral modes.

Evaluation is organized into three groups: terminal readout (vision, single cells, clinical forecasting), latent world dynamics (intervention prediction, ten dynamics tasks, PDE audit, long rollouts, locomotion planning, molecular trajectories), and scientific analysis (wet-lab biology and orbital scaling). In direct JEPA comparisons, the adapter, encoder, sampler, data splits, optimization budget, and readout are all held fixed, so the only difference is the factorized core and its regularizers.

Why This Matters

The paper argues that latent world modeling need not be reinvented per domain. If a single predictive core can serve images, cells, patients, robots, molecules, and weather, then representation learning, planning, and scientific hypothesis generation can share infrastructure rather than diverging into separate silos. The factorized target also turns the learned latent space into something inspectable: factors are assigned to distinct predictive directions, so they can be probed for physical laws or used to nominate interventions that are then tested experimentally.

Real-world applications:

• Drug discovery and cell therapy: perturbation-response prediction and factor-based nomination of interventions, validated here in organoids, tumor fragments, and mice.
• Clinical risk forecasting: simultaneously predicting risks for over a thousand downstream events from a synthesized future patient state.
• Scientific simulation: stable long-horizon rollouts of molecular dynamics, physical fields, and weather that could reduce reliance on expensive numerical solvers.
• Robotics and control: latent world models that support compositional intervention reasoning and model-based planning under held-out action combinations.

Industry relevance spans pharmaceutical and biotech research pipelines, clinical decision-support systems, weather and climate modeling, materials discovery, and robotics — all areas where learning a reusable predictive state is more valuable than fitting a narrow, task-specific predictor.

Future Directions

• Selecting factor count and width automatically. The number of factors K and their width r are currently configured per domain; a principled criterion for choosing them from data would remove a manual tuning step.
• Interpreting factor semantics. Factors emerge from predictability rather than predefined meaning. Understanding what each factor captures — and when it aligns with a human-meaningful variable — remains open.
• Closing the loop between model and experiment. The paper demonstrates one factor-nominated biological intervention with wet-lab support; scaling this into an iterative hypothesis–experiment cycle is a natural extension.
• Scaling and broader domains. Testing whether OPF continues to help at much larger model and data scales, and in audio, language, or multi-agent settings, would test the "Anything" claim more aggressively.

Target Audience

Machine learning researchers working on self-supervised learning, JEPA-style predictive architectures, and latent world models; domain scientists in computational biology, clinical informatics, physics, and weather modeling who need reusable predictive representations; and applied engineers building model-based planning or simulation systems who want a training principle that transfers across data modalities without redesigning the architecture per domain.

Authors’ abstract

World modeling enables intelligence to anticipate consequences, guide interventions, and learn from interaction. Yet predictive models remain domain-specific: can a common learning principle support world modeling across radically different systems? We introduce JEPA-Anything, a domain-agnostic framework based on orthogonal predictive factorization (OPF). Extending joint-embedding predictive architectures, OPF decomposes latent targets into complementary factors, learns them through dedicated pathways, and recombines them within a shared predictive design. We evaluate JEPA-Anything across seven domains: vision, biology, clinical trajectories, control, molecular dynamics, physical fields, and weather. Experiments span representation learning, intervention prediction, out-of-distribution generalization, and long-horizon dynamics, including 10 matched dynamics tasks, forecasting of over 1,000 clinical events, and 100-step molecular rollouts across four systems. Against matched JEPA baselines, JEPA-Anything improves reported metrics on all 10 dynamics tasks and reduces single-intervention prediction error on Interventional Pong by 34.8%. It achieves the lowest one-step and 100-step molecular errors among compared methods in all four systems. Beyond prediction, a factor-nominated biological intervention receives experimental support in cell co-cultures, patient-derived organoids, tumor fragments, and mice; latent orbital modes recover the Keplerian scaling exponent with a fitted slope of -1.4991. These results support a common factorized predictive principle across heterogeneous worlds, connecting world modeling with intervention and experimentally grounded scientific discovery. Code: https://github.com/Gen-Verse/JEPA-Anything