Loading…
All-atom protein dynamics simulator that bridges consecutive MD snapshots directly using SO(3)-equivariant two-sided stochastic interpolants.
No providers recorded yet. Browse all providers
A diffusion model asked for the next frame of a molecular dynamics trajectory begins from Gaussian noise and rebuilds the whole protein from scratch — even though the frame it was handed already differs from the answer by only a fraction of an Ångström per atom. The prior sits far from the data manifold, and most of the transport is spent re-deriving structure the model was given for free. EquiJump removes that detour. It uses two-sided stochastic interpolants, which bridge two arbitrary distributions rather than a Gaussian and a data distribution, so the source distribution is the current conformation itself and the target is the conditional distribution of the next snapshot given it. The transport path stays inside the space of physically plausible protein states for its entire length.
The model was built by MIT — the Center for Bits and Atoms and the Atomic Architects group — with NVIDIA and SISSA, and released as a preprint in October 2024, later presented at the ICLR 2025 GEM workshop. It sits alongside other learned MD surrogates in the catalog, including MDGen, which generates whole trajectories with a flow-matching transformer, and JAMUN, which samples ensembles by walk-jump in a noise-smoothed coordinate space. EquiJump differs from both in being an explicit step-to-step time-evolution operator trained on, and evaluated against, long-timescale trajectory statistics.
The network is SO(3)-equivariant throughout, built from a Self-Interaction module that mixes irreducible representations within a residue and a Spatial Convolution module in the Tensor Field Network style that passes messages between neighbours, with residual connections and an equivariant layer norm. Five such networks make up the model: one conditioner that embeds the sequence and source structure, and four prediction heads for the feature and coordinate components of drift and noise — the four-track design that lets the expensive conditioner run once per jump while only the heads iterate through the transport integration.
Training used the multi-millisecond all-atom MD dataset of Majewski et al. (2023): twelve fast-folding proteins of 10–80 residues, NVT runs at 350 K sampled every 100 ps. Because high-free-energy states are rare in unbiased trajectories, the authors reweight the training set by clustering in a TICA subspace and sampling clusters uniformly, an idea borrowed from umbrella sampling. Four capacities were trained for 500k steps at batch size 128: 6.5M, 25.4M, 100.8M and 391.1M parameters. Evaluation compares Jensen-Shannon divergences of MSM-reweighted observables — TICA components, Cα RMSD, GDT, radius of gyration, fraction of native contacts — against reference trajectories. The 391.1M model reaches 0.03 JS on the first two TICA components averaged over the twelve proteins, against 0.30 and 0.23 for the coarse-grained force field of Majewski et al., the only other transferable model covering the same set. On a single A100, a 100 ps jump for the largest system takes 0.34–1.05 s depending on capacity, which the authors estimate at a 3.4–34× speedup over the explicit-solvent MD that generated the training data.
Learned time-evolution operators like this one are aimed at the sampling problem: recovering the populations of folded, unfolded and intermediate states, and the transitions between them, without paying for microseconds of integration. Because EquiJump emits all heavy atoms rather than a Cα trace, its output can feed analyses that need side-chain geometry — dihedral distributions, contact maps, pocket occupancy — directly, with no backmapping step.
EquiJump's contribution is the demonstration that two-sided interpolants beat Gaussian-prior transport for this task: in a controlled comparison on Protein G with a shared backbone, it outperformed DDPM, flow matching, and a one-sided interpolant on every observable measured. The honest limit is scope. The transferable model is trained and evaluated on the same twelve fast-folding proteins at a single temperature, and the authors explicitly defer transferability and generalization to future work, so nothing here establishes zero-shot simulation of an arbitrary protein. No code or weights accompany the preprint, which leaves the results dependent on reimplementation to reproduce.
Much of this page is generated or calculated automatically. Flag anything that looks off and we will re-run it.