MolJEPA-World v1

An E(3)-equivariant joint-embedding predictive architecture (JEPA) world model for protein conformational dynamics. It encodes a coarse-grained CA frame into a 32-dimensional latent and advances that latent with a learned Koopman operator, so future conformations are predicted in representation space rather than in coordinates.

• Weights: model.safetensors (2.32 M parameters, 9.3 MB)
• Corrections: two PRST operator corrections (prst/)
• Code: self-contained PyTorch inference code (no training dependencies)
• Training data: 236 mdCATH domains, 944,000 CA frames
• Licence: Apache-2.0

This is v1 of a research model. It is a latent-dynamics model, not a force field and not a structure predictor.

Newer release: AdagioBiology/world-v2 keeps this encoder but replaces the Koopman operator with multi-horizon displacement heads trained on segment-aware pairs. On the corrected segment-aware protocol it is 11–21% better than the v1 operator and beats v1 + PRST at every horizon (100% win rate over persistence on 23 held-out domains). Start there if you are evaluating the model today.

---

What it does

CA frame (N,3)  +  ESM-2 per-residue embeddings  +  temperature (K)
        │
   EGNN encoder (3 layers, E(3)-equivariant, kNN graph k=16)
        │
   pooled latent z ∈ R^32
        │
   Koopman operator K (12 complex modes + 3 state-gated rank-2 experts)
        │
   z(t + n·275 ps)  ──►  heads: distogram → 3D, spatial descriptors,
                                  forces, electrostatics

The operator is trained with a semigroup-consistency loss, a VICReg anti-collapse term, a spectral/timescale term, and an EMA target encoder (JEPA asymmetry). The released weights are the student encoder and operator; the EMA teacher is a training artifact and is not shipped.

Released files

Which correction should I use?

• default: round2 (full chain) — best at 40/80/120 frames (11/22/33 ns)
• round1 — best at 20 frames (5.5 ns); essentially tied at 40
• no PRST — the raw learned operator

The correction is a 32×32 matrix applied after every 20-frame operator step:

z = K^20(z) + (z_prev - z_mean) @ C.T

load_pretrained(..., prst="round2") attaches it and advance() applies it with exactly the semantics used in the held-out evaluation.

Quickstart

pip install torch numpy safetensors huggingface_hub
python example_inference.py

import torch
from modeling_moljepa import load_pretrained, prepare_coords, advance

model = load_pretrained()                    # this folder
# model = load_pretrained("AdagioBiology/world-v1", prst="round1")

x = prepare_coords(coords_angstrom)          # (B,N,3), centred and / 8.5
z = model.encode(atom_types, x, atom_mask=mask,
                 esm_emb=esm, sim_temp=temperature)

z_future = advance(model, z, frames=120)     # 120 × 275 ps = 33 ns
coords = model.generate_3d(atom_types, x, z=z_future) * model.coord_scale

Inputs

One model step is 275 ps (the training loader keeps 4,000 of mdCATH's 11,030 frames). The published horizons are 20/40/80/120 steps = 5.5/11/22/33 ns.

Visuals

Left: PCA of 800 latents across five mdCATH temperatures. Right: the 12 learned Koopman relaxation times; the dashed line is the 33 ns evaluation horizon.

Left: segment-aware held-out error (24 domains, 3,495 windows) for the operator and the full PRST chain. Right: the legacy uniform-window protocol, in which 69.7% of 120-frame windows cross a replica/temperature boundary — it makes v1 look 5–10% better at short horizons.

CPU inference latency (measured while a training job shared the machine, so these are upper bounds).

v32 training run: loss and latent quality versus wall-clock time, plus throughput. 30,000 steps in 506 minutes on one RTX 5060 Ti.

The 3D structure gallery (visuals/structure_rollout.png) and the animated rollout (visuals/rollout_3d.gif) are included for inspection but are not headline evidence: they show a hard 450 K transition, where the decoder is well outside its comfort zone.

Evaluation

Protocol: frozen segment-aware manifest, 24 mdCATH domains never used in training, 10 segments per domain × 5 starts × 3 seeds = 3,495 windows, with start + horizon constrained to stay inside one replica/temperature segment. Metric: per-domain latent MSE of the rollout divided by the MSE of persistence (z_t copied forward); the table reports the median over domains with a 95% bootstrap CI. This is the corrected protocol: the earlier numbers used uniform windows, 69.7% of which crossed a segment boundary and which included 8 training domains in the aggregate.

Win rate versus persistence: 71 / 83 / 75 / 88% for the base operator, and 100 / 100 / 96 / 92% for the full PRST chain. The eval/ folder contains the 24 held-out domains; eval_manifest_v2.json and the per-domain CSV are available from the project repository.

Learnability checks on the training domain: latent cosine at 5.5 ns ≈ 0.998; the learned slowest timescales are ≈ 196, 68, 13, 11, 7, 5 ns.

Read the metric carefully. vs persistence is an aggregate over a heterogeneous corpus (5 temperatures × 5 replicas per domain). Within a single replica at one temperature, consecutive frames are highly correlated and persistence is an unusually strong baseline; on the shipped 320 K demo the raw operator has cosine 0.998 but ~100× the persistence MSE, and PRST brings that to ~1.7×. The metric measures forecast skill on the full distribution, not physical accuracy at a fixed state point.

Training details

The PRST corrections were fit with the frozen encoder on 59,994 real and 17,820 model-generated, teacher-verified pairs from training domains, and evaluated only on held-out domains.

Limitations

• CA-only and coarse-grained. No side chains, hydrogens, ligands, ions, solvent or protonation states. It is not an all-atom force field.
• Not a structure predictor. The 3D path is distogram → SMACOF; it has reflection ambiguity and typically lands within a few percent of Rg and CA-CA bond length (demo: Rg 13.98 vs 13.83 Å, CA-CA 3.90 vs 3.83 Å). Use it for visualisation, not for atomic modelling.
• Domain coverage. Trained on single-domain mdCATH proteins. Multi-domain proteins, membrane proteins, disordered regions, bound ligands and temperatures outside 320–450 K are out of distribution.
• Conditioning matters. The model expects ESM-2 embeddings and the simulation temperature. Omitting them feeds the encoder out-of-distribution inputs and degrades every metric.
• Latent forecasts are not physical trajectories. The operator is a learned linear model with a small non-normal correction; long rollouts can drift and do not conserve energy or obey detailed balance.
• Not for clinical, therapeutic or drug-discovery decisions.

Versioning

The repository name is world-v1, but the underlying training run is v32 of the internal MolJEPA series. v1 denotes the first public release of this line; future releases will use v2, v3, ….

Citation

@misc{moljepa_world_v1,
  title  = {MolJEPA-World: an E(3)-equivariant joint-embedding predictive
            world model for protein dynamics},
  author = {AdagioBiology},
  year   = {2026},
  url    = {https://huggingface.co/AdagioBiology/world-v1}
}

Licence

Apache-2.0. mdCATH data is subject to its own licence; only derived weights and a 41-frame demo excerpt are distributed here.