AdagioBiology/world-v1

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

File	What it is
model.safetensors	v32 final student weights (step 29,999)
config.json	architecture + preprocessing + training metadata
modeling_moljepa.py	standalone model class + load_pretrained / advance
moljepa_core/	EGNN, Koopman head, distogram head, decoders
prst/prst_C_round1.pt	PRST correction, round 1 (best at 5.5–11 ns)
prst/prst_C_round2_full.pt	full 3-round PRST chain (best at 22–33 ns)
examples/demo_1a0hA01.npz	41 frames of one mdCATH domain (runnable demo)
eval/heldout_domains.txt	the 24 domains never used in training
visuals/	latent geometry, corrected error curves, runtime and training figures

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

Input	Shape	Notes
pos	(B, N, 3)	CA coordinates in Å. Centre each frame and divide by 8.5 (prepare_coords). The model is E(3)-invariant, so global rotation/translation are arbitrary.
atom_types	(B, N)	mdCATH atom-type ids; every CA is type 6.
atom_mask	(B, N)	1 for real residues, 0 for padding.
esm_emb	(B, N, 480)	ESM-2 facebook/esm2_t12_35M_UR50D per-residue embeddings. Strongly recommended: without them the encoder has no per-residue identity.
sim_temp	(B,)	Simulation temperature in K (mdCATH range 320–450 K). Strongly recommended.

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.

System	20 f (5.5 ns)	40 f (11 ns)	80 f (22 ns)	120 f (33 ns)
v32 baseline (no PRST)	0.825 [0.681, 0.959]	0.797 [0.704, 0.876]	0.702 [0.649, 0.796]	0.698 [0.623, 0.834]
v32 + PRST round 1	0.676 [0.640, 0.715]	0.680 [0.626, 0.721]	0.648 [0.610, 0.739]	0.670 [0.592, 0.785]
v32 + PRST full chain	0.687 [0.657, 0.725]	0.670 [0.624, 0.703]	0.642 [0.601, 0.725]	0.643 [0.579, 0.749]

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

Data	mdCATH CA traces, 236 domains × 4,000 frames = 944,000 frames
Coverage	50–327 residues, 5 temperatures (320/348/379/413/450 K), 5 replicas
Model	3-layer EGNN, embed 64, hidden 128, latent 32, 12 Koopman modes, 3 experts (rank 2)
Parameters	2.32 M inference (3.23 M including the EMA teacher)
Optimisation	AdamW, lr 3e-4, wd 0.05, cosine schedule, 30,000 steps, batch 16, lag k=20
Losses	VICReg 30/2, prediction 2.0, disp 2.0, self-force 2.0 (×2), distogram 1.0, spectrum 1.0, semigroup consistency 0.1, expert balance 0.2
Hardware	one 16 GB consumer GPU, 506 min
Initialisation	warm-start from the previous cycle (v31)

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.