MDGen

A molecular dynamics trajectory is a movie, but most deep-learning surrogates train on it as if it were a pile of unrelated stills. Boltzmann generators learn only the equilibrium distribution; transition-density models learn a single step and are rolled out autoregressively. Either way the dynamical content of the saved simulation — which torsional barriers are slow, which state pairs carry flux, what happens between two saved frames — is discarded at training time. MDGen instead models the full time-series end to end, drawing an entire trajectory of 3D molecular structures in one sample, much as a video model produces a clip, not a frame.

That framing is what makes one model do several jobs. Once a generative model over complete trajectories exists, the choice of which frames to condition on selects the problem being solved: the first frame gives forward simulation, the first and last frames give transition path sampling, every M-th frame gives trajectory upsampling, and conditioning on part of the molecule while inpainting the rest designs residues consistent with an observed motion. Only the conditioning mask changes, and three of the four are inverse problems that more MD does not solve.

MDGen was developed by Bowen Jing, Hannes Stärk, Tommi Jaakkola and Bonnie Berger at MIT CSAIL and presented at NeurIPS 2024. It shares its ATLAS data and evaluation pipeline with AlphaFlow from the same group, but where AlphaFlow samples static conformations from a structure predictor, MDGen generates the time axis itself.

#Key Features

• Trajectory tokenization: An all-atom trajectory of T frames and L residues becomes a T×L array of SE(3)-invariant tokens, each holding roto-translational offsets from K chosen key frames plus seven torsion angles — sidestepping the memory-hungry pair representations of structure prediction.
• Task selection by conditioning: Forward simulation, interpolation, upsampling and inpainting are one framework under four conditioning masks, not four engineered methods.
• Zero-shot transition paths: Conditioned on two metastable endpoints, the model samples 1 ns paths for unseen peptides, with better Jensen-Shannon divergence against the reference Markov state model and a higher fraction of valid paths than any replicate MD run shorter than the 100 ns reference.
• Trajectory upsampling: Trajectories saved every 10 ps are reconstructed at a 100 fs timestep — a 100× upsampling factor — recovering sub-picosecond torsional oscillations the saved frames cannot represent.
• Dynamics-conditioned design: Given the trajectory of a tetrapeptide's two flanking residues, inpainting the inner residue identities recovers the ground-truth peptide 52.1% of the time on high-flux transitions, against 17.4% for an inverse-folding baseline given only the terminal frames.

#Technical Details

The velocity network is a Scalable Interpolant Transformer trained under the stochastic interpolants formulation of flow matching. Its trunk alternates DiT-style attention blocks across the residue index and across the time axis, with rotary position embeddings; invariant point attention layers embed the key-frame roto-translations into conditioning representations broadcast along time. Generating amino acid identities for inpainting adds Dirichlet flow matching over the 20-dimensional probability simplex. Replacing time attention with the Hyena operator scales single samples to 100k-frame trajectories spanning four orders of magnitude in timescale — a proof of concept obtained by overfitting one pentapeptide, which the authors note is not yet comparable to the main forward-simulation experiments.

Training data is all-atom MD in OpenMM with amber14: roughly 3,000 training tetrapeptides plus 100 validation and 100 test peptides, each simulated for 100 ns in explicit and implicit solvent, plus ATLAS — three 100 ns replicates for each of 1,390 protein monomers. On held-out tetrapeptides, forward-simulation rollouts reach a torsion-angle JSD of 0.109 against ground truth, better than 10 ns of replicate MD (0.125) and approaching the 100 ns oracle (0.076), at about 60 GPU-seconds per 100 ns-equivalent trajectory versus 3 GPU-hours for the simulation; predicted sidechain relaxation times correlate with MD at r = 0.97. Protein-scale ensembles from the separate ATLAS checkpoint reach a per-target RMSF correlation of 0.71 at 0.2 s per structure, between MSA subsampling with AlphaFold 2 and AlphaFlow in accuracy and far faster than either.

#Applications

The draw is cheap access to dynamical observables MD makes expensive: free energy surfaces and Markov state fluxes for short peptides, ensembles of transition paths between two known conformations without waiting for a rare event, and fast motion recovered from archived trajectories saved too coarsely to analyze. Inpainting points at a design use — choosing residues to scaffold a desired conformational change — with no clean formulation in conventional simulation.

#Impact

MDGen reframes MD surrogate modeling as generative modeling of trajectories rather than of structures or single steps, and has become a common baseline for later ensemble and dynamics generators. The paper states its limits plainly: the full capability set is validated on tetrapeptides, protein-scale results are first steps at intermediate accuracy, and the reliance on key frames means the model cannot generate unconditionally. The release is explicitly research-grade — reproduction scripts and separate task-specific checkpoints rather than one general-purpose model.