AdagioBiology/world-v2

MolJEPA-World v2

The same E(3)-equivariant JEPA encoder as world-v1, with the operator replaced. v1 rolled the latent with a learned Koopman operator; v2 rolls it with direct multi-horizon displacement heads (one MLP per horizon: 1, 20, 40, 80, 120 frames) that were trained on segment-aware pairs — every training and evaluation pair has its start and target frame inside the same replica and temperature, which is the physical continuity the data actually has.

On a frozen, segment-aware held-out manifest this is worth 11–21% lower rollout error than the v1 Koopman operator, and it also beats v1 + its PRST correction at every horizon, with a 100% win rate over persistence on all four horizons.

• Weights: model.safetensors (2.36 M parameters: the v1 encoder plus the new heads)
• Code: self-contained PyTorch inference code (no training dependencies)
• Training data: 184 mdCATH domains, segment-aware pairs
• Licence: Apache-2.0

---

Visuals

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

Left: segment-aware held-out error for the operator and the full PRST chain (24 domains, 3,495 windows, shaded 95% CI). 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.

What changed vs world-v1

	world-v1	world-v2
Rollout operator	Koopman (12 modes + 3 experts)	direct displacement heads per horizon
Training pairs	uniform starts (≈12–25% crossed replicas)	segment-aware (start and target in one replica)
Post-hoc correction	PRST (a 32×32 matrix fitted offline)	none needed
Held-out error at 33 ns	0.700× persistence	0.623× persistence

The v1 failure mode was diagnosed rather than guessed: the v1 operator is not contractive (effective spectral radius 1.14), and the legacy evaluation drew 69.7% of its 120-frame windows across replica/temperature boundaries, where the conditional mean beats persistence and nothing physical is being tested. v2 fixes both: the heads predict displacements on continuous trajectories, and the evaluation manifest is frozen and segment-aware.

Results

Frozen segment-aware manifest: 24 mdCATH domains never used in training (23 with a usable backbone cache), 10 segments × 5 starts × 3 seeds = 3,495 windows, start + horizon constrained to one replica. Metric: per-domain latent MSE of the rollout divided by the MSE of persistence, median over domains. All four systems share the encoder and the windows; only the operator differs.

Horizon	v2 heads	v1 Koopman	v1 + PRST (r1)	v1 + PRST (full)	v2 win rate
20 f (5.5 ns)	0.667	0.845	0.679	0.688	100%
40 f (11 ns)	0.671	0.811	0.682	0.675	100%
80 f (22 ns)	0.641	0.719	0.664	0.647	100%
120 f (33 ns)	0.623	0.700	0.673	0.645	100%

33 ns decoded rollout

Reference (top), v1 Koopman (middle) and v2 heads (bottom) at 0, 5.5, 11, 22 and 33 ns on a held-out domain, decoded with the distogram/SMACOF decoder and aligned by Kabsch.

Read this honestly: the latent forecast is measurably better, but the 3D decode is still the weak link — the SMACOF path was trained on teacher latents, so decoding a rolled latent is out of distribution and both systems land around 14 Å RMSD at 33 ns. A rollout-conditioned morphing decoder is the next build; it is not included here.

How good can a 3D forecast be? The replicate-MD oracle

mdCATH ships five independent MD replicas per protein per temperature. Two replicas of the same protein at the same elapsed time differ by a finite amount — the irreducible spread of the thermal ensemble, which no forecast can beat. Measuring it on held-out domains (CA RMSD after alignment):

horizon	320 K: replica oracle / persistence	450 K: replica oracle / persistence
5.5 ns	8.7 / 7.3 Å	12.6 / 14.9 Å
11 ns	8.7 / 7.0 Å	14.1 / 12.8 Å
22 ns	9.0 / 9.2 Å	14.3 / 13.4 Å
33 ns	9.9 / 8.9 Å	16.5 / 15.4 Å

Consequences, which matter for how the next model should be evaluated:

• At 320 K the protein is nearly rigid: persistence is already at (or better than) the replica spread, so a point forecast cannot show much.
• At 450 K the ensemble is broad and a forecast is effectively a draw from it — the 3D RMSD of a 33 ns prediction is dominated by thermal noise.
• Our 3D forecast sits at that noise floor, and the same RMSD is obtained when decoding the true future latent. A 3D point forecast at these horizons therefore cannot discriminate between latent systems.

This is why the headline metric here is a latent one, and why the next evaluation must move to ensemble statistics (state occupancies, TICA/PC distributions, coverage of the reference ensemble) rather than point RMSD.

Quickstart

pip install torch numpy safetensors huggingface_hub
python example_inference.py

from modeling_moljepa import load_pretrained, prepare_coords, advance

model = load_pretrained()                       # this folder
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_33ns = advance(model, z, frames=120)          # direct head for 120 frames
z_free = advance(model, z, frames=120, mode="rollout")   # free-running 1-step

advance(..., mode="direct") uses the head trained for that horizon (the numbers in the table above). mode="rollout" applies the 1-step head repeatedly, which is what the movie above uses.

Inputs

Input	Shape	Notes
pos	(B, N, 3)	CA coordinates in Å. Centre each frame and divide by 8.5 (prepare_coords).
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.
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).

What we tried that did not work

Recorded because the negative results are as informative as the positive one:

Idea	Outcome
Fine-tuning the v1 encoder on the displacement objective	Destroys the operator: the Koopman path degrades from 0.70× to 1.61× at 33 ns and the new heads land at 0.734× (vs 0.623× frozen). A single loss term does not preserve the representation's other properties — adapting the encoder needs the full loss suite.
64% more proteins (302 size-matched domains vs 184)	No gain: 0.656× at 33 ns vs 0.623×. The heads are not data-limited at this scale.
EMA teacher for the heads	No difference (0.654× vs 0.656×) — with a frozen encoder the teacher is the student.
A 5.5 M per-residue transformer (local-frame tokenizer, pair-distance bias, distance profile)	Collapses to a static latent (within-protein std 0.006 vs norm 13.5). Do not scale it without a per-residue rollout path.
Global-latent displacement decoders (MLP and EGNN-coupled)	Plateau at 9–12 Å per-residue error: a 32-d global latent cannot specify ~165 DOF.
Per-residue latent rollout	Reaches persistence level, no better — consistent with the oracle result below: at these horizons the predictable 3D component is smaller than the thermal spread.

Limitations

• The encoder is still the v1 encoder (0.78 M parameters of EGNN plus 1.5 M of heads). v2 changes the operator and the data protocol, not the representation; a scaled backbone-frame model is the next step.
• 3D decoding is the weak link (see above). Latent quality and decoded structure quality are not the same thing, and only the first is measured here. A 3D point forecast at 5–33 ns also sits at the MD-replica noise floor, so 3D RMSD cannot discriminate between latent systems — the next evaluation must use ensemble statistics (state occupancies, TICA/PC distributions, coverage), not point RMSD.
• CA-only, no side chains, ligands, ions or solvent.
• 184 training domains, single-chain mdCATH proteins, 320–450 K. Nothing outside that.
• Latent rollouts are not physical trajectories: they do not conserve energy or obey detailed balance.
• Not for clinical, therapeutic or drug-discovery decisions.

Versioning

world-v2 is the second public release of this line. world-v1 remains available with the v1 Koopman operator and PRST corrections.

Citation

@misc{moljepa_world_v2,
  title  = {MolJEPA-World v2: multi-horizon displacement heads for an
            E(3)-equivariant JEPA of protein dynamics},
  author = {AdagioBiology},
  year   = {2026},
  url    = {https://huggingface.co/AdagioBiology/world-v2}
}

Licence

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