JEPA-Anything

A unified predictive framework across domains

Note: This project is under active development. Beyond the experiments reported in the paper, we will also continue to add experiments and methods that are not yet covered in the paper but that we have found to work well in practice, and share them with the community.

Why predictive world models across domains?

A world model builds an internal state that can be used to anticipate another state of the same system. The observation and the target can take very different forms: visible image patches can predict a masked region, a control state and action can predict the next state, a cellular profile and perturbation can predict a response, and a patient history can predict a future health state. Molecular configurations, physical fields, and weather maps define the same kind of question over different structures and time scales.

JEPA-Anything treats these problems through a common context-to-target interface. Each domain supplies an adapter that turns its raw data into content tokens and structural descriptors such as position, time, entity identity, graph connectivity, action, intervention, or external forcing. A context encoder summarizes what is observed, an exponential-moving-average target encoder defines the state to predict, and prediction takes place in representation space. This focuses learning on structure that is useful for future queries and downstream tasks.

Why factorize the predictive state?

A standard Joint-Embedding Predictive Architecture uses one target embedding and one prediction pathway. Complex targets often contain several predictable components at once: local and global visual changes, multiple interacting entities, biological identity and response, or physical modes evolving at different rates. When all of them share one pathway, high-variance or easier components can dominate the learning signal, multiple latent directions can learn similar roles, and weaker modes can receive conflicting gradients.

JEPA-Anything introduces Orthogonal Predictive Factorization (OPF) to organize this capacity. OPF analyzes the target representation into learned subspaces, assigns a dedicated predictor to each factor, and encourages the subspaces to cover complementary directions. Factor-activity and encoder-variance objectives keep the predictive coordinates informative. The factor predictions are then recombined into a complete latent world state for decoding, planning, intervention prediction, or the next rollout step.

The shared element is the predictive interface and OPF core. Domain adapters retain the geometry and semantics of images, sequences, graphs, sets, fields, and multivariate records, so each task can use the encoder family and readout that fit its data.

How the learned state is used

1. Terminal readout: the learned encoder state supports a downstream task such as controlled visual binding, cell-state clustering, perturbation-response prediction, or disease-risk readout.
2. Operational world modeling: predicted factors are synthesized into the next state and reused for intervention-conditioned prediction, model-based planning, and multi-step rollout.
3. Scientific analysis: retained factor coordinates expose predictive modes that can be tested against biological interventions, dynamical structure, symmetries, and scaling laws.

The experiments follow the same organization. Group I studies terminal readout, Group II studies latent dynamics and rollout stability, and Group III studies factor-level scientific analysis.

Research domains

Domain	Context to target	What the model is used for	Repository entry
Vision and controlled binding	Visible patches or a source scene to masked regions and controlled scene changes	Learn visual states that preserve both where a change occurs and which operation occurred, including new support-operation combinations	Shapes3D models
Biology and single cells	Masked expression to complete cell state; control and perturbed cells to a latent response	Cell-state representation, clustering, and gene-expression response prediction	Norman single-cell model
Clinical and longitudinal health	Time-ordered patient history to a future latent health state	Disease forecasting, treatment-response modeling, and broad future-event risk prediction	Progressive checkpoint release
Control and robotics	Observation plus action or intervention to the next state and future trajectory	Compositional intervention prediction, learned simulation, CEM planning, and long-horizon control	Interventional dynamics, control tasks, and locomotion
Molecular systems	Atomic positions, species, velocities, and geometry to future molecular configurations	Geometry-aware one-step forecasting and free autoregressive trajectory generation	TiTo molecular forecasting
Physical fields	Current discretized field and known forcing to the next field	PDE forecasting, out-of-distribution prediction, and recurrent rollout over nonlinear dynamics	PDEBench collection, APEBench, and long-horizon Burgers
Weather and climate	Current atmospheric field and temporal context to a future field	Large-scale spatiotemporal forecasting with repeated state prediction	WeatherBench2
Scientific analysis and discovery	Learned factor coordinates to candidate interventions or interpretable dynamical modes	Form biological hypotheses and compare predictive modes with established physical relations	Scientific discovery

Checkpoint release

All checkpoints developed in JEPA-Anything will be made publicly available progressively. Releases are organized domain by domain and paired with the corresponding model code, input description, and reproducible inference command. The current repository provides checkpoints for vision, single-cell biology, interventional dynamics, control, PDEs, locomotion, weather, and molecular forecasting; the remaining research directions will be added through the same release stream.

Quick start

Most state-transition and field checkpoints use the shared dynamics entry point:

python test/infer_dynamics.py \
  --checkpoint <experiment>/checkpoints/<checkpoint>.pt \
  --branch model \
  --input <preprocessed_transitions.npz> \
  --output <output_directory>/predictions.npz

The input NPZ contains x and either action for one-step prediction or actions for a recurrent rollout. Add y or targets to obtain overall and per-step MSE in the JSON report. Paired checkpoints expose standard and jepa_anything branches through --branch.

Visual encoder checkpoints use:

python test/infer_vision.py \
  --checkpoint <experiment>/checkpoints/<checkpoint>.pt \
  --input <preprocessed_images.npz> \
  --output <output_directory>/tokens.npz

Norman single-cell prediction, corrected CITRIS evaluation, and TiTo molecular rollout have dedicated commands in their model pages.

Repository map

• JEPA-anything-final/: reusable JEPA-Anything core package.
• model/: shared checkpoint reconstruction and inference models.
• test/: common prediction entry points.
• Domain directories: background, input format, checkpoint links, and task-specific tests.

Contact and collaboration

JEPA-Anything is under active development, and its framework, experiments, and domain-specific releases will continue to evolve as new results and checkpoints are added. If you are interested in collaborating on a particular domain or have questions about the project, please contact me at cty21@tsinghua.org.cn.