Generative Modeling of Molecular Dynamics Trajectories

Authors: bowen  jing
hannes  stark
tommi  jaakkola
bonnie  berger
Venue: NeurIPS_cc_2024_Conference
Published: 2024-05-15 19:25:33

Abstract: Molecular dynamics (MD) is a powerful technique for studying microscopic phenomena, but its computational cost has driven significant interest in the development of deep learning-based surrogate models. We introduce generative modeling of molecular trajectories as a paradigm for learning flexible multi-task surrogate models of MD from data. By conditioning on appropriately chosen frames of the trajectory, we show such generative models can be adapted to diverse tasks such as forward simulation, transition path sampling, and trajectory upsampling. By alternatively conditioning on part of the molecular system and inpainting the rest, we also demonstrate the first steps towards dynamics-conditioned molecular design. We validate the full set of these capabilities on tetrapeptide simulations and show preliminary results on scaling to protein monomers. Altogether, our work illustrates how generative modeling can unlock value from MD data towards diverse downstream tasks that are not straightforward to address with existing methods or even MD itself. Code is available at https://github.com/bjing2016/mdgen.

Main paper: 
 introduction: 
 Numerical integration of Newton's equations of motion at atomic scales, known as molecular dynamics (MD), is a widely-used technique for studying diverse molecular phenomena in chemistry, biology, and other molecular sciences (Alder & Wainwright, 1959; Rahman, 1964; Verlet, 1967; McCammon et al., 1977) . While general and versatile, MD is computationally demanding due to the large separation in timescales between integration steps and relevant molecular phenomena. Thus, a vast body of literature spanning several decades aims to accelerate or enhance the sampling efficiency of MD simulation algorithms (Ryckaert et al., 1977; Darden et al., 1993; Sugita & Okamoto, 1999; Laio & Parrinello, 2002; Anderson et al., 2008; Shaw et al., 2009) . More recently, learning surrogate models of MD has become an active area of research in deep generative modeling (Noé et al., 2019; Zheng et al., 2023; Klein et al., 2024; Schreiner et al., 2024; Jing et al., 2024) . However, existing training paradigms fail to fully leverage the rich dynamical information in MD training data, restricting their applicability to a limited set of downstream problems.In this work, we propose MDGEN, a novel paradigm for fast, general-purpose surrogate modeling of MD based on direct generative modeling of simulated trajectories. Different from previous works, which learn the autoregressive transition density or equilibrium distribution of MD, we formulate end-to-end generative modeling of full trajectories viewed as time-series of 3D molecular structures. Akin to how image generative models were extended to videos (Ho et al., 2022) , our framing of the problem augments single-structure generative models with an additional time dimension, opening the door to a larger set of forward and inverse problems to which our model can be applied. When provided (and conditioned on) the initial "frame" of a given system, such generative models serve as familiar surrogate forward simulators of the reference dynamics. However, by providing other kinds of conditioning, these "molecular video" generative models also enable highly flexible applications to a variety of inverse problems not possible with existing surrogate models. In sum, we formulate and showcase the following novel capabilities of MDGEN:• Forward simulation-given the initial frame of a trajectory, we sample a potential time evolution.• Interpolation-given the frames at the two endpoints of a trajectory, we sample a plausible path connecting the two. In chemistry, this is known as transition path sampling and is important for studying reactions.• Upsampling-given a trajectory with timestep ∆t between frames, we upsample the "framerate" to ∆t/M .• Inpainting-given part of a molecule and its trajectory, we generate the rest of the molecule (and its time evolution) to be consistent with the known part of the trajectory. This ability could be applied to design molecules to scaffold desired dynamics. These tasks are conceptually illustrated in Figure 1 . While the forward simulation task aligns with the typical modeling paradigm of approximating the data-generating process, the others represent novel capabilities on scientifically important inverse problems not straightforward to address even with MD itself. As such, our framework presents a new perspective on how to unlock value from MD simulation with machine learning towards diverse downstream objectives. We highlight further exciting possibilities opened up by our framework in Section 4.We evaluate MDGEN on the forward simulation, interpolation, upsampling, and inpainting tasks on tetrapeptides in a transferable setting (i.e., generalizing to test peptides). Our model accurately reproduces free energy surfaces and dynamical content such as torsional relaxation and Markov state fluxes, provides realistic transition paths between arbitrary pairs of metastable states, and recovers fast dynamical phenomena below the sampling threshold of coarsetimestep trajectories. In preliminary steps toward dynamicsscaffolded design, we show that molecular inpainting with MDGEN obtains much higher sequence recovery than inverse folding methods based on one or two static frames. Finally, we evaluate MDGEN on forward simulation of protein monomers and find that its ensembles' statistical properties surpass the fidelity of ensembles from multiple sequence alignment (MSA) subsampling with AlphaFold (Del Alamo et al., 2022) . 
tokenizing molecular trajectories: 
 For further background on MD and Stochastic Interpolants, please see Appendix A. Given a chemical specification of a molecular system with N atoms, our aim is to learn a generative model over time-series χ ≡ [X 1 , . . . X T ] of cor-responding molecular structures X i ∈ R 3N for some trajectory length T . In this work, we specialize to MD trajectories of short peptides (Sections 3.1-D.1) or single-chain proteins (D.1); thus, our chemical specifications are amino acid sequences A = {1 . . . 20} L . We adopt a SE(3)-based parameterization of molecular structures (Jumper et al., 2021; Yim et al., 2023) , such that the all-atom coordinates of each amino acid residue are described by a roto-translation (i.e., element of SE(3)) and seven torsion angles χ l t = ((R, t), (ψ, ϕ, ω, χ 1 . . . χ 4 )), χ ∈ SE(3) × T 7 L T where sub-/superscripts indicate time/residue indices.To learn a generative model over this space of rototranslations and torsion angles, one natural choice is to employ diffusion or flow-based models for protein structures explicitly designed for such representations (Yim et al., 2023; Lin & AlQuraishi, 2023) . However, these architectures rely on expensive edge-based and IPA-based updates, which can be memory-intensive to replicate across a large number of trajectory frames. Instead, we leverage the fact that we are concerned with conditional trajectory generation-meaning that there always exists at least one frame in the trajectory with un-noised roto-translations which we do not need to generate and can reference in the modeling process. Inspired by analogy to video compression, we refer to such frames as key frames and parameterize the rototranslations of remaining structures as offsets relative to the key frames. Specifically, given K key frames i 1 . . . i K we tokenize residue j in frame t as:EQUATIONwhere g j ∈ SE(3) represents the roto-translation and τ j the torsion angles of residue j. Namely, we convert the relative roto-translational offsets g -1 i g j to unit quaternions with positive real part Q+ ⊂ R 4 (representing the rota- tions) and translation vectors in R 3 , and convert torsion angles to points on the unit circle, obtaining a (7K + 14)dimensional token for each residue in every frame. To untokenize to an all-atom frame X i ∈ R 3N , we read off the torsion angles from the unit circle and apply the key frame roto-translations g i k to the generated offsets to obtain absolute roto-translations ĝj . The offsets and torsion angles are SE(3)-invariant; thus, we obtain a representation of molecular trajectories as an (T × L)-array of SE(3)-invariant tokens: χ ∈ R T ×L×(7K+14) . This representation allows us to directly apply a flow model architecture to generate the tokens (details in Appendix B.1.We present the precise specifications of conditioning settings in Table 2 . Depending on the task, we choose the key frames to be the first frame g 1 or the first and last frames g 1 , g T . Each task is further characterized by providing the groundtruth tokens of known frames or residues as additional inputs to the model. Meanwhile, mask tokens are provided for the unknown frames and residues that the model generates. For example, in the upsampling setting, we provide ground-truth tokens every M frames, while mask tokens are provided for all other frames. We note that in the inpainting setting, the model accesses the roto-translations g of designed residues at the trajectory endpoints via the key frames, constituting a slight departure from the full inpainting setting. 9 ). (Bottom Right) Statistics averaged over 100 test peptides and 1000 paths for each of them. Shown are JSD, fraction of drawn paths that are valid transition paths, and average path likelihood of our discretized transitions under the reference MSM compared to discrete transitions drawn from the reference MSM or alternative MSMs built from replica simulations of varying lengths. representation of MD (Prinz et al., 2011; Noé et al., 2013) . Appendix C provides further details on MSMs with more results in Appendix D. 
forward simulation: 
 Distributional similarity. We train a model to sample 10 ns trajectories conditioned on the first frame. By chaining together successive model rollouts at inference time, we obtain 100 ns trajectories for each peptide to compare with ground-truth simulations. We report the Jensen-Shannon divergence (JSD) between the ground-truth and emulated trajectories along various collective variables shown in Fig- ure 2 and Table 2 . The first set of these are the individual torsion angles (backbone and sidechains) in each tetrapeptide. The second set of variables are the top independent components obtained from time-lagged independent components analysis (TICA), representing the slowest dynamic modes of the peptide. By each of these collective variables, MDGEN demonstrates excellent distributional similarity to the ground truth, approaching the accuracy of replicate 100-ns simulations. To more stringently assess the ability to locate and populate modes in the joint distribution over state space, we build Markov State Models (MSMs) for each test peptide using the MD trajectory, extract the corresponding metastable states, and compare the ground-truth and emulated distributions over metastable states. Our model captures the relative ranking of states (Figure 2D ). Dynamical content. We compute the dynamical properties of each tetrapeptide in terms of the decorrelation time of each torsion angle from the MD simulation and from our sampled trajectory. Intuitively, this assesses if our model can discriminate between slow-and fast-relaxing torsional barriers. The correlation between true and predicted relaxation timescales is plotted in Figure 2F , showing excellent agreement for sidechain torsions. To assess coarser but higher-dimensional dynamical content, we compute the flux matrix between all pairs of distinct metastable states using ground-truth and sampled trajectories and find substantial Spearman correlation between their entries (mean ρ = 0.67 ± 0.01; Figure 8 ).Sampling speed. Averaged across test peptides, our model samples 100 ns-equivalent trajectories in ≈60 GPU-seconds, compared to ≈3 GPU-hours for MD. To quantify the speedup more rigorously, we compute the decorrelation wall-clock times along the slowest independent component from TICA, capturing how quickly the simulation traverses the highest barriers in state space. These times are plotted in Figure 2E , showing that our model achieves an speedup of 10x-1000x over the MD simulation for 78 out of 100 peptides (the other 22 peptides did not fully decorrelate).In the interpolation or transition path sampling setting, we train a model to sample 1 ns trajectories conditioned on the first and last frames. For evaluation, we identify the two most well-separated states (i.e., with the least flux between them) for each test peptide and sample an ensemble of 1000 transition paths between them. Figure 3 shows an example of such a sampled path, which passes through several intermediate states on the free energy surface.To evaluate the accuracy of these sampled transitions, we cannot directly compare with MD trajectories since, in most cases, there are zero or very few 1-ns transitions between the two selected states (by design, the transition is a rare event). Thus, we instead discretize the trajectory over MSM metastable states and evaluate the path likelihood under the transition path distribution from the reference MSM (details in Appendix C.3). We also report the fraction of non-zero probability paths and the JSD between the distribution of visited states from our path distribution versus the transition path distribution of the reference MSM. For baselines, we sample transition paths from MSMs constructed from replicate MD simulations of varying lengths and compute the metrics for their path ensembles under the reference MSM.As shown in Figure 3 , our paths have higher likelihoods than those sampled from any replicate MD MSM shorter than 100ns, which is the length of the reference MD simulation itself. Moreover, MDGEN's ensembles have the best JSDs to the distribution of visited states of the reference MD MSM and the highest fraction on valid non-zero probability paths. Hence, our model enables zero-shot sampling of trajectories corresponding to arbitrary rare transitions.In the upsampling setting, we train MDGEN to upsample trajectories saved with timestep 10 ps to a finer timestep of 100 fs, representing a 100x upsampling factor. To evaluate if the upsampled trajectories accurately capture the fastest dynamics, we compute the autocorrelation function ⟨cos(θ t -θ t+∆t )⟩ of each torsion angle in the test peptides as a function of lag time ∆t ranging from 100 fs to 100 ps.Representative examples of ground truth, subsampled, and reconstructed autocorrelation functions for two test peptides are shown in Figure 4 (further examples in Figure 10 ). We further compute the dynamical content as the negative derivative of the autocorrelation with respect to logtimescale, which captures the extent of dynamic relaxations occuring at that timescale (Shaw et al., 2009) . These visualizations highlight the significant dynamical information absent from the subsampled trajectory and which are accurately recovered by our model. In particular, our model distinctly recovers the oscillations of certain torsion angles as seen in the non-monotonicity of the autocorrelation function at sub-picosecond timescales; these features are completely missed at the original sampling frequency. 
discussion: 
 Opportunities. Similar to the foundational role of video generative models for understanding the macroscopic world (Yang et al., 2024) Molecular dynamics. At a high level, the aim of molecular dynamics is the integrate the equations of motionM i ẍi = -∇ xi U (x 1 . . . x N ) for each particle i in a molecular configuration (x 1 . . . x N ) ∈ R 3N, where M i is the mass and U is the potential energy function (or force field) U : R 3N → R. However, these equations of motion are often modified to include a thermostat in order to model contact with surroundings at a given temperature. For example, the widely-used Langevin thermostat transforms the equations of motion into a stochastic diffusion process of the equations dx i = p i /M i dt anddp i = -∇ xi U dt-γp i dt+ √ 2M i γkT dw, where p i are the momenta. By design, this process converges to the Boltzmann distribution of the system p(x 1 . . . x N ) ∝ e -U/kT . To incorporate interactions with solvent molecules-ubiquitous in biochemistry-one includes a box of surrounding solvent molecules as part of the molecular system (explicit solvent) or modifies the force field U to model their effects (implicit solvent). In either case, only the positions x i of non-solvent atoms are of interest, and their time evolution constitutes (for our purposes) the MD trajectory. Thus, MD can be mathematically described as a partially observed Markovian diffusion process.Deep learning for MD. An emerging body of work seeks to approximate the distributions over configurations X = (x 1 . . . x N ) arising from MD with deep generative models. Fu et al. (2023) , Timewarp (Klein et al., 2024) , and ITO (Schreiner et al., 2024) learn the transition density p(X t+∆t | X t ) and emulate MD trajectories via simulation rollouts of the learned model. On the other hand, Boltzmann generators (Noé et al., 2019; Köhler et al., 2021; Garcia Satorras et al., 2021; Midgley et al., 2022; 2024) directly approximate the stationary Boltzmann distribution, forgoing any explicit modeling of dynamics. In particular, Boltzmann-targeting diffusion models trained with frames from MD trajectories have demonstrated promising scalability and generalization to protein systems (Zheng et al., 2023; Jing et al., 2024) . However, these works have focused exclusively on forward simulation and have not explored joint modeling of entire trajectories (X t . . . X t+N ∆t ) or the inverse problems accessible under such a formulation.Stochastic interpolants. We build our MD trajectory generative model under the stochastic interpolants framework: Given a continuous distribution p 1 ≡ p data over R n , stochastic interpolants (Albergo & Vanden-Eijnden, 2022; Albergo et al., 2023; Lipman et al., 2022; Liu et al., 2022) , provide a method for learning continuous flow-based models dx = v θ (x, t) dt transporting a prior distribution p 0 (e.g., p 0 ≡ N (0, I)) to the data p 1 . To do so, one defines intermediate distributions x t ∼ p t , t ∈ (0, 1) via x t = α t x 1 + σ t x 0 where x 0 ∼ p 0 and x 1 ∼ p 1 and the interpolation path satisfies α 0 = σ 1 = 0 and α 1 = σ 0 = 1. A neural network v θ : R n × [0, 1] → R n is trained to approximate the time-evolving flow field v θ (x t , t) ≈ v(x t , t) ≡ E x0,x1|xt [ αt x 1 + σt x 0 ] which satisfies the transport equation ∂p t /∂t + ∇ • (p t v t ) = 0. Hence, at convergence, noisy samples x 0 ∼ p 0 can be evolved under v θ to obtain data samples x 1 ∼ p 1 . When parameterized with transformers (Vaswani et al., 2017) , stochastic interpolants are state-of-the-art in image generation (Esser et al., 2024) . In particular, we adopt the notation, architecture, and training framework of Scalable Interpolant Transformer (SiT) (Ma et al., 2024) , to which we refer to for further exposition.Our base modeling task is to generate a distribution over R T ×L×(7K+14) conditioned on roto-translations of one or more key frames g i1 . . . g i K , and (in most settings) amino acid identities A. To do so, we learn a flow-based model via the stochastic interpolant framework described in SiT (Ma et al., 2024) and parameterize a velocity network 7K+14) . To condition on the key frames and amino acids, we first provide the sequence embedding to several IPA layers (Jumper et al., 2021) that embed the key frame roto-translations; these conditioning representations (which are SE(3)-invariant) are broadcast across the time axis and added to the input embeddings. The main trunk of the network consists of alternating attention blocks across the residue index and across time, with the construction of each block closely resembling DiT (Peebles & Xie, 2023) . Sidechain torsions and roto-translation offsets, when available, are directly provided to the model as conditioning tokens.v θ (• | g i1 . . . g i K , A) : R T ×L×(7K+14) × [0, 1] → R T ×L×(In the molecular inpainting setting where we also generate the amino acid identities, we additionally require a generative framework over these discrete variables. While several formulations of discrete diffusion or flow-matching are available (Hoogeboom et al., 2021; Austin et al., 2021; Campbell et al., 2022; 2024) , we select Dirichlet flow matching (Stark et al., 2024) as it is most compatible with the continuous-space, continuous-time stochastic interpolant framework used for the positions. Specifically, we place the amino acid identities on the 20-dimensional probability simplex (one per amino acid), augment the token representations with these variables, and regress against a T × L × (7K + 14 + 20)-dimensional vector field. Further details in Appendix B.2. Key frames Generate Conditioned on Token dim.Forward simulationg 1 g 1•••T , τ 1•••T g 1 , τ 1 , A 21 Interpolation g 1 , g T g 1•••T , τ 1•••T g 1,T , τ 1,T , A 28 Upsampling g 1 g 1•••T , τ 1•••T g 1+{1,2,••• }M , τ 1+{1,2,••• }M , A 21 Inpainting g 1 , g T g 1•••T , A g known 1•••T 7 (+20)Notation. Here, as in the main text, we use the following notation:• T : number of trajectory frames• L: number of amino acids• K: number of key frames, with indicies i 1 . . . i KIn Algorithms 1-3 below, we modify the architectures of DiffusionTransformerAttentionLayer and DiffusionTransformerFinalLayer from DiT (Peebles & Xie, 2023) . Elements from these layers are also then incorporated into our custom InvariantPointAttentionLayer.Input: noisy tokens χ ∈ R T ×L×(7K+14) , conditioning tokens χ cond ∈ R T ×L×(7K+14) , key frame rototranslations g i1 . . . g i K ∈ SE(3) L K , flow matching time t, amino acid identities A ∈ {1, . . . 20} L , conditioning mask m ∈ {0, 1} T ×L×(7K+14) Output:flow velocity v ∈ R T ×L×(7K+14) t ← Embed(t) for k ← 1 to K do x k = Embed(A) + k ′ Linear(g -1 i k g i k ′ ) for l ← 1 to num ipa layers do x k = InvariantPointAttentionLayer(x, g k , t) end for end for x = k x k + Linear(χ) + Linear(χ cond ⊙ m) + Embed(m) for l ← 1 to num transf ormer layers do x = DiffusionTransformerAttentionLayer(x, t) end for return DiffusionTransformerFinalLayer(x, t) Algorithm 2 DiffusionTransformerAttentionLayer Input: x ∈ R T ×L×C , time conditioning t (α, β, γ) t,ℓ,f = Linear(t) x += g ℓ ⊙ AttentionWithRoPE(γ ℓ ⊙ LayerNorm(x) + β ℓ , dim = 1) x += g t ⊙ AttentionWithRoPE(γ t ⊙ LayerNorm(x) + β t , dim = 0) x += g m ⊙ MLP(γ m ⊙ LayerNorm(x) + β m ) return xA Markov State Model (MSM) is a representation of a system's dynamics discretized into r states s ∈ {1 . . . r} and a discrete timesteps separated by time lag τ such that the dynamics are approximately Markovian (Husic & Pande, 2018; Chodera & Noé, 2014; Pande et al., 2010 ). An MSM is parameterized with a vector π that assigns each state a stationary probability and a matrix T containing the probabilities for transitioning from state s t to s t+1 after one timestep, i.e., T i,j = p(s t+1 = j | s t = i).To build a Markov state model, we use PyEMMA (Scherer et al., 2015; Wehmeyer et al.) and its accompanying tutorials. Briefly, we first featurize molecular trajectories with all torsion angles as points on the unit circle, obtaining a 2m-dimensional invariant trajectory where m is the number of torsion angles. We run TICA on these trajectories with kinetic scaling and then run k-means clustering with k = 100 over the first few (5-10 chosen by PyEMMA) TICA coordinates. We then estimate an MSM over these 100 states and use PCCA+ spectral clustering (Röblitz & Weber, 2013) to further group these into 10 metastable states. Our final MSM is built from the discrete trajectory over these 10 metastable states. In all cases we use timelag τ = 100 ps.Unconditionally sampling an MSM. To unconditionally sample a trajectory of length N from an MSM, we first sample the start state from the stationary distribution, i.e., s 1 ∼ π. We then iteratively sample each subsequent state as s t+1 ∼ T st,: . Sampling an MSM conditioned on a start state. To sample a trajectory of length N conditioned on a starting state s 1 , we iteratively sample each subsequent state as s t+1 ∼ T st,: . Sampling an MSM conditioned on a start and end state. For our transition path sampling evaluations in Section 3.2, we employ replica transition paths sampled from an MSM by conditioning on a start state s 1 and end state s N . To do so, we iteratively sample each state between the conditioning states by utilizing the probabilityEQUATIONFirstly, the term p(s t+1 = j | s t = i) is available in out transition matrix as T i,j . Secondly, we obtain p(s N = k | s t = i) as an entry of the (N -t)th matrix exponential of the transition matrix. Specifically p(sN = k | s t = i) = T (N -t) i,kwhere the superscript denotes a matrix exponential. Lastly, we obtain the term p(s N = k | s t+1 = j, s t = i) by realizing that under the Markov assumption p(sN = k | s t+1 = j, s t = i) = p(s N = k | s t+1 = j). Further, p(s N = k | s t+1 = j) = T (N -t)-1 j,k .Replacing the terms in Equation 4 results inEQUATIONThus, we sample states s 2 . . . s N -1 iteratively asEQUATIONWe run all-atom molecular dynamics simulations in OpenMM (Eastman et al., 2017) using the amber14 force field parameters with gbn2 implicit solvent or tip3pfb water model. Initial structures are generated with PyMOL, prepared with pdbfixer, and protonated at neutral pH. For explicit solvent, we prepare a solvent box with 10 Å padding and neutralize the system with sodium or chloride ions. All simulations are integrated with Langevin thermostat at 350K with hydrogen bond constraints, timestep 2 fs, and friction coefficient 0.3 ps -1 (explicit) or 0.1 ps -1 (implicit). For explicit solvent, nonbonded interactions are cut off at 10 Å with long-range particle-mesh Ewald. We first minimize the energy with L-BFGS and then equilibrate the system in the NVT ensemble for 20 ps. We then run 100 ns of production simulation in the NVT ensemble (implicit) or NPT ensemble with Monte Carlo barostat at 1 bar (explicit). We write heavy atom positions every 100 fs.For explicit-solvent settings (forward simulation, interpolation, inpainting), we run simulations for 3109 training, 100 validation, and 100 test peptides. For implicit-solvent settings (upsampling), we run simulations for 2646 training, 100 validation, and 100 test peptides. All peptides are randomly chosen and split. Additionally, 5195 training and 100 validation implicit solvent simulations are run for the pentapeptide MDGEN.Trajectory Featurization We featurize trajectories by selecting the sine and cosine of all torsion angles as the collective variables. Specifically, we featurize ψ, ϕ backbone angles and all χ sidechain torsion angles for each peptide. We then reduce dimensionality with Time-lagged Independent Components Analysis (TICA) (Pérez-Hernández et al., 2013) in PyEMMA (Scherer et al., 2015) .Jensen-Shannon Divergence We compute the JSD as implemented in scipy, i.e.,D(p | m) + D(q | m) 2 (7)where m = (p + q)/2. For the 1-dimenional JSD over torsion angles, we discretize the range [-π, π] into 100 bins. For the 1-dimensional JSD over TIC-0, we discretize the range spanning the maximum and minimum values into 100 bins. For the 2-dimensional JSD over TIC-0,1 we discretize the space into 50 × 50 bins.Autocorrelation The autocorrelation of torsion angle θ at time lag ∆t is defined as ⟨cos(θ t -θ t+∆t )⟩, corresponding to the inner product of θ t , θ t+∆t on the unit circle. To compute the decorrelation time of a torsion angle, we subtract the baseline inner product ⟨cos θ⟩ 2 + ⟨sin θ⟩ 2 , this is analogous to removing the mean of a real-valued time series before computing the autocorrelation. The decorrelation time is then defined as the time required for the autocorrelation to fall below 1/e of its initial value (which is always unity), with the subtracted baseline computed from the reference trajectory.In a small number of cases (21 torsions), the MDGEN trajectory did not decorrelate within 1000 frames (10 ns), and we exclude the angle from Figure 2F .To compute the decorrelation time for TIC-0, we now define the autocorrelation asEQUATIONwhere µ, σ are computed from the reference trajectory. Hence, when computed for a sampled MDGEN trajectory, the autocorrelation may not start at unity and may not decay to zero. We report a decorrelation time if starts above and falls below 0.5 within 1000 frames (10 ns), which happens in 74 out of 100 cases as shown in Figure 2E .Interpolation In our interpolation or transition path sampling experiments, we sample 1000 trajectories of length 1ns for each of our 100 test tetrapeptides. We first select a start state s 1 and an end state s N that exhibits non-trivial transitions. To do so, we consider a reference MD simulation of 100 ns for the tetrapeptide and obtain an MSM as described in Appendix C.1. From the MSM's transition matrix T and stationary distribution π, we compute the flux matrix F = T ⊙ P i where P i is the square matrix with π in each column. The chosen start and end state is the row and column index of the smallest non-zero entry in F .With the start state s 1 and end state s N selected, we sample 1000 start frames x 1 and end frames x N from the states. The 1000 start frames are sampled from all frames in the reference MD simulation that belong to state s 1 . Analogously, the end frames are sampled from all frames belonging to state s N . Using the 1000 pairs of start and end frames, we condition MDGEN on them and generate trajectories of 100 frames (1 ns). For evaluation, we discretize these trajectories under the 10-state clustering determined by the MSM of the reference MD simulation as described in Appendix C.1. Note that with the MSM lag time of 100 ps, these discrete trajectories are of length 10.MD baselines. To sample transition paths of 1 ns between our selected start and end states, we employ MSMs built from replica MD simulations of varying lengths. For instance, for a replica MD simulation of 100ns, we first discretize its trajectory with the cluster assignments of the reference MD simulation (the same cluster assignments as we use to discretize the MDGEN ensemble and that we use for evaluation). Next, we estimate an MSM from the discretized trajectory. We then proceed to sample 1000 transition paths from the MSM as described in C.1 where the path is conditioned on an end and start state. In the event that the replica MSM has zero transition probability for transitioning out of the start state or zero probability for transitioning into the end state (this occurs if the replica MD simulation never visited the start or end state), we treat all 1000 paths of the replica MD as having zero probability for our evaluation metrics which are further detailed in the following.Computing TPS metrics. As described above, we obtain ensembles of 1000 discretized 1ns paths of 100 frames for both MDGEN and the replica MD simulations. For these, in Figure 3 , we show a JSD, the rate of valid paths, and the average path probability. These metrics are computed with respect to the MSM of the reference MD simulation of length 100 ns.• To compute the JSD, we draw 1000 discrete transition paths from the reference MD simulation and compute the probability of visiting each state from the frequency with which each state is visited. We do the same for the transition path ensemble of MDGEN (or the baseline) and compute the JSD between the categorical distributions as described above.• The average path probability for an ensemble is the average of its paths' probabilities for transitioning from the start to the end state under the reference MSM. This probability can be computed as described in Appendix C.1.• The valid path rate is the fraction of paths that have a non-zero probability.Inpainting In our inpainting experiments, we set out to design tetrapeptides that transition between two states. Considering the residue indices 1, 2, 3 and 4, we call the residues 1, 4 the flanking residues which we condition on and 2, 3 the inner residues which we aim to design. Specifically, we condition MDGEN on the trajectory of the flanking residues' backbone coordinates and generate the residue identities of the inner two residues. To carry out this design for a single tetrapeptide, we draw 1000 samples from MDGEN to estimate the mode of its joint distribution over the inner two residues.The conditioning information (the trajectories of the outer two residues' backbone coordinates) is different for the two evaluation settings of designing transitions with high flux or for designing arbitrary transitions. However, for both of them, the start and end frames are provided as conditioning information via the key frames. In the high flux setting, the conditioning information is obtained by sampling 1000 paths from the reference MD simulation of length 10 ps with 100 frames that start and end in the desired states. These states are determined as those with the maximum flux between them (see the paths about interpolation above for a description of flux). When designing residues that give rise to arbitrary random paths, the trajectories are randomly sampled from the reference simulation.After sampling 1000 pairs of residues for the inner two residues, we select the most frequently occurring pair as the final design. For this design, we report the sequence recovery (the fraction of residues that match the original sequence of the MD simulations from which the conditioning information was sampled).Inpainting Baselines. We aim to assess the benefit that is obtained by the trajectory-based inference of MDGEN over a baseline that only takes the start frame or the start and end frame as input for designing residues that transition between two states. Thus, we construct DYNMPNN and S-MPNN. These baselines use the same architecture as MDGEN in the inpainting setting, but DYNMPNN only obtains the start and end frames as key frames and via their roto-translation offsets for the first and last frames. S-MPNN is the analog with only the first frame.Notably, in the inpainting setting, MDGen and the baselines do not treat torsion angles, and all torsion angle entries of the SE(3)-invariant tokens are set to 0. Furthermore, the model does not take the amino acids of the flanking residues as input.We make this choice of withholding all information about amino acid identities since otherwise, the models overfit on the arbitrary identities of the flanking residues and do not generalize to the test set.Protein Simulations For training and evaluation on proteins, we use trajectories from the ATLAS dataset (Vander Meersche et al., 2024) , which includes 3 replicates of 100 ns explicit-solvent, all-atom simulations for each of 1390 non-membrane protein monomers. The proteins are chosen from the PDB as representatives of all available ECOD domains and are thus structurally non-redundant. We split the dataset into 1265 training, 39 validation, and 82 test proteins by PDB release date following Jing et al. (2024) . At training time, we randomly select a protein, select one of the three replicates, subsample every 40 frames, obtaining a training target with 250 frames. We train with random crops of up to 256 residues, but draw samples for the full protein at inference time. To compute statistical similarity of the MDGEN ensembles with the ground truth MD ensembles, we compare the 250 frames with 30k pooled frames from all three trajectories. Baseline metrics and runtimes for AlphaFlow and MSA subsampling are taken directly from Jing et al. (2024) . Analysis and visualization code for Table 4 and Figure 6 are provided courtesy of Jing et al. (2024) . Inpainting Design. We aim to sample trajectories conditioned on the dynamics of the two flanking residues of the tetrapeptide; in particular, the model determines the identities and dynamics of the two inner residues. We focus on dynamics scaffolding as one possible higher-level objective of inpainting: given the conformational transition of the observed residues, we hope to design peptides that support flux between the corresponding Markov states.Thus, for each test peptide, we select a 100-ps transition between the two most well-connected Markov states, mask out the inner residue identities and dynamics, and inpaint them with our model. To evaluate the designs, we compute the fraction of generated residue types that are identical to the tetrapeptide in which the target transition is known to occur. We compare MDGEN with a bespoke inverse folding baseline that is provided the two terminal states (i.e., two fully observed MD frames), and thus designs peptides that support the two modes (rather than additionally a partially-observed transition between them). We call this baseline DYNMPNN, and it otherwise has the same architecture and settings as MDGEN. We find (Table 3 ) that MDGEN recovers the ground-truth peptide substantially more often than DynMPNN when conditioned on a high-flux path or (as a sanity check) a random path from the reference simulatiom.Scaling to Long Trajectories. Although Section 3.1 showed that our model can emulate long trajectories, this was limited to rollouts of 1000 frames at a time with coarse 10 ps timesteps, potentially missing faster dynamics or disrupting slower dynamics. Thus, we investigate generating extremely long consistent trajectories that capture timescales spanning several orders of magnitude within a single model sample. To do so, we replace the time attention in our baseline SiT architecture with a non-causal Hyena operator (Poli et al., 2023) , which has O(N log N ) rather than O(N 2 ) overhead. We overfit on 100k-frame, 10-ns trajectories of the pentapeptide MDGEN and compare the torsional autocorrelation functions computed from a single generated trajectory with a single ground truth trajectory (Figure 5 ).Although not yet comparable to the main set of forward simulation experiments due to data availability and architectural expressivity reasons, these results demonstrate proof-of-concept for longer context lengths in future work.Protein Simulation. To demonstrate the applicability of our method for larger systems such as proteins, we train a model to emulate explicit-solvent, all-atom simulations of proteins from the ATLAS dataset (Vander Meersche et al., 2024) conditioned on the first frame (i.e., forward simulation). We follow the same splits as Jing et al. (2024) . Due to the much larger number of residues, we generate samples with 250 frames and 400 ps timestep, such that a single sample emulates the 100 ns ATLAS reference trajectory. The difficulty of running fully equilibrated trajectories for proteins prevents the construction of Markov state models used in our main evaluations. Instead, we compare statistical properties of forward simulation ensembles following Jing et al. (2024) . Our ensembles successfully emulate the ground-truth ensembles at a level of accuracy between AlphaFlow and MSA subsampling while being orders of magnitude faster per generated structure than either (Table 4 ).  
 
 Appendix:  


Decision: The paper presents a new approach to a timely topic -- modeling molecular dynamics -- with two accepted papers in NeurIPS last year. There are several contributions, lending ideas from CV and NLP (such as masking and impainting) to model a joint distribution of the simulation trajectory. As such, these new training strategies allow for the use of more data than previous approaches. The tasks the authors explore are in line with standard benchmarks. New prediction tasks are introduced, mostly of limited interest and practical use (e.g. trajectory upsampling). The major limitation and strength of the approach are its main innovations: 1) the molecular representation (tokenization), which locks in the models to be proteins or peptides, with no clear strategy to generalize to molecules more broadly, and 2) the trajectory representation, requiring end users to make hard decisions at training time, which may limit inference time usefulness. The underlying learning objective is not well justified from a physical perspective: e.g. a trajectory is learned as a joint distribution suggesting a time-dependence of the generator. There are settings where this is the case, however, in all the systems explored the paper the generating process is homogeneous and Markovian, and therefore at odds with the model and training setup. Nevertheless, the performance appears good in certain cases (small peptides), but poor on larger scale application, e.g. ATLAS protein dataset.

Reviews: review: paper_summary: paper_summary: The authors introduce MDGen as a novel approach for modeling MD trajectories. They demonstrate the capabilities of this method in tasks such as interpolation, upsampling, and inpainting of small peptides. The accuracy as well as speed of the new approach compared to the ground truth baseline is quantitatively evaluated. Initial experiments toward upscaling to small proteins are shown.
strength_weakness: strength_weakness: 1. Parts of Sections 3.1 and 3.2 are very condensed and hard to follow. A more detailed description in the SI would be helpful, where the most important aspects of the cited work is also repeated.
2. The suitability of the chosen representation for longer amino acid chains is questionable. This is also mentioned in the manuscript, but nonetheless, proteins are mentioned many times (more than 30) in the manuscript, while almost all experiments are actually performed on very small peptides. It should be stated in a more prominent place that upscaling to proteins is not trivial.
3. The representation limits the model to learn MD trajectories of natural amino acids, as no all-atom representation is used directly. This should be made clearer in the manuscript.

Minor points: A lot of figures have no proper axis labels (e.g. Fig 3, 4, 5, 6). This should be fixed. The best models in Table 4 should be indicated in bold.
questions: questions: 1. How often do clashes and other high-energy structures occur in the generated trajectories?
2. When comparing to other methods and approaches in the experimental section - do all of them use a similar reduced representation or do the other methods generate all-atom representations?
limitations: limitations: The approach is limited to peptides. Transfer to any other molecules is questionable due to a lack of suitable representation/tokenization.
paper_summary: paper_summary: The paper presents a novel generative model for molecular dynamics (MD) trajectories called MDGEN. This model aims to serve as a flexible surrogate for MD simulations by generating entire trajectories conditioned on initial frames. It addresses tasks such as forward simulation, transition path sampling, trajectory upsampling, and dynamics-conditioned molecular design. The model is evaluated on tetrapeptide simulations and demonstrates its capability to generate reasonable ensembles of protein monomers.
strength_weakness: strength_weakness: Complexity and Accessibility: The model’s complexity might pose challenges for reproducibility and accessibility for researchers who are not deeply familiar with both molecular dynamics and advanced generative modeling techniques.
Evaluation on Larger Systems: While the paper provides proof-of-concept evaluations on proteins, the primary focus remains on smaller tetrapeptides. The model's scalability and effectiveness on larger and more complex molecular systems need further exploration.
Dependence on Key Frames: The reliance on key frames for conditional generation limits the model’s ability to perform unconditional generation or inpainting of residue roto-translations, which could be a significant limitation in certain applications.
Computational Resources: The paper lacks detailed information on the computational resources required for training and inference, which is crucial for understanding the practical implications of using MDGEN in various research settings.
questions: questions: How can the model be adapted or improved to reduce its reliance on key frames?

Exploring techniques for unconditional generation or alternative ways to handle the roto-translations without predefined key frames could enhance the model's flexibility and applicability.
What architectural changes or enhancements could improve the model's performance on larger molecular systems such as proteins?

Investigating more scalable architectures or hybrid approaches that combine the current method with other techniques tailored for large systems could address this limitation.
How does the computational cost of training the model compare to traditional MD simulations, and what are the implications for its practical use?

Providing detailed information on computational requirements and potential optimizations could help in assessing the model's feasibility for widespread use.
What alternative tokenization strategies could be explored to extend the model's applicability to a wider range of molecular systems?

Research into tokenization methods that can handle diverse molecular structures and dynamics could broaden the model's utility.
How can additional conditioning types (e.g., textual descriptions, experimental data) be incorporated into the model, and what benefits might they provide?

Experimenting with and integrating various forms of conditioning could enhance the model's ability to generate more accurate and contextually relevant trajectories.
What are the potential impacts of data quality and availability on the model's performance, and how can these challenges be mitigated?

Addressing data-related challenges through techniques like data augmentation, transfer learning, or synthetic data generation could improve the model's robustness and applicability.
Can additional evaluation metrics be developed to provide a more comprehensive assessment of the generated trajectories' quality?

Identifying and implementing new evaluation criteria could offer deeper insights into the strengths and limitations of the model's output.
limitations: limitations: Reliance on Key Frames:

The model relies on key frames for parameterizing the roto-translations, which means it cannot perform unconditional generation or inpainting of residue roto-translations. This dependency might limit its applicability to scenarios where key frames are not easily obtainable or where full trajectories need to be generated from scratch.
Scalability to Larger Systems:

The architecture shows weaker performance on larger systems such as proteins compared to smaller peptides. This suggests that the current model and architecture might not be well-suited for handling the complex motions and larger size of protein structures without further modifications or enhancements.
Computational Resources:

While the paper mentions significant speedups compared to traditional MD simulations, the computational resources required for training the model (e.g., GPU hours) are not explicitly discussed. This information is crucial for understanding the practicality and scalability of the approach.

Generalization to Diverse Systems:

The current tokenization and modeling strategies are tailored to peptides and proteins. For more diverse molecular systems such as organic ligands, materials, or explicit solvent systems, alternative tokenization strategies might be necessary. This limits the immediate applicability of the model to a broader range of molecular simulations.
Limited Exploration of Additional Conditioning:

The paper primarily explores conditioning on initial frames and residue identities. Other types of conditioning, such as textual or experimental descriptors, are not explored but could open up further applications and improve the model's utility.
Data Availability and Quality:

The success of the model heavily depends on the availability of high-quality MD trajectory data. For many complex systems, obtaining such data can be challenging, which could limit the model's applicability and performance.
Evaluation Metrics:

While the paper uses several rigorous metrics for evaluation, the choice of metrics may not fully capture all aspects of the generated trajectories' quality. Additional metrics or more diverse evaluation criteria could provide a more comprehensive assessment.
paper_summary: paper_summary: In this work, the authors proposed MDGen, a new framework that aims to model molecular dynamics trajectories via generative modeling techniques. By properly encoding the Protein MD trajectories according to the characteristics of key frames, MDGen adopts flow matching techniques (both continuous and discrete flow matching) to generatively model MD trajectories. As a unified framework, MDGen is able to perform diverse tasks including forward simulation, interpolation, upsampling and inpaiting. Extensive experiments are conducted to demonstrate the effectiveness of MDGen.
strength_weakness: strength_weakness: 1. Lack of discussion on related works. This work does not discuss related works on the same topic. Some works are mentioned in the Introduction section, but I still recommend that there should be an independent Related Works section for comprehensive discussion.  Here are also several works that are worth discussing: (1) EGNO, which uses neural operator learning approach to also model the trajectory dynamics of molecules; (2) DiffMD, which uses diffusion models to simulate molecular dynamics. The quality of this work should be further improved if the authors could carefully discuss the differences between MDGen and these works and the strengths of MDGen compared to these works.

2. Lack of ablation studies. MDGen is composed of several parts, including the design of the backbone model, the design choices of flow matching framework, and the adoption of Hyena architecture for efficiency consideration. In addition to the aimed tasks, it would further improve the quality of this work if the authors could conduct ablation studies on these aspects to help readers know what the influence of each part of MDGen is.
questions: questions: N/A
limitations: limitations: The authors carefully discuss the limitations of this work.
paper_summary: paper_summary: The paper presents a new framework for generating trajectory of molecular geometries, ie, generative modeling for molecular dynamics. The paper proposes tokenization methods to tokenize the trajectory and learn flow models on the data. Experiments demonstrate the effectiveness of several tasks including forward sampling,  interpolation, and up sampling.
strength_weakness: strength_weakness: 1. Limited ML technical contribution, as all components exist in previous molecular generative models.

2. The experiment is comprehensive from a domain perspective. However, I feel the experiments lack some benchmarking comparison with state-of-the-art molecular generative models for related tasks. See my question below.
questions: questions: I think existing methods can also tackle several tasks. For example, for the forward sampling task, previous generative MD models like Timewarp (Klein et al., 2024) and ITO (Schreiner et al., 2024) can also be used for the task. A numerical comparison with these baselines can help to justify the effectiveness of the proposed method.
limitations: limitations: Limitation is nicely discussed in the paper.
paper_summary: paper_summary: The authors propose MDGen -- a generative model to sample molecular dynamics trajectory conditioned on key frames. This is a direct application of video generation techniques to solve domain challenges in protein modeling. Specifically, SiT and flow matching models are used to sample SE(3)-invariant representation of all-atom protein representations. This work demonstrates the effectiveness of MDGen primarily on tetrapeptide systems, where the authors showcase four downstream application tasks including forward simulation, interpolation, upsampling, and dynamics-conditioned inpainting. 

In general, I find this manuscript well-written and easy to follow. The model performance looks reasonable on tetrapeptides, yet the results are proof-of-concept in nature and generalization to larger proteins remain challenging. However, it is one of the pioneering work in AI protein modeling to directly emulate MD simulation trajectories using data-driven approaches. To that end, I think it would be beneficial for this work to gain visibility across the research community to inspire future studies.
strength_weakness: strength_weakness: - Benchmark and evaluation results on tetrapeptides, although comprehensive, are proof-of-concept in nature. It may not be sufficient to demonstrate transferability to general protein systems.
- Performance on ATLAS (i.e., larger proteins instead of short peptides) does not seem promising. MDGen performance is worse than AlphaFlow in Table 4. I wonder if the main bottleneck is training data quality/availability, or model architecture?
questions: questions: - L#141, when $K > 1$, how to ensure roto-translation prediction consistency across $K$ key frames and obtain a final $\hat{g}_j$?
- Table 2, with 100 ns being the ground truth, the non-zero JSD in the last column originates from subsampling the simulation trajectory?
- Figure 2F. My understanding is that sidechains exhibit faster dynamics while backbone motions are slower. The low correlation for backbone suggests that MDGen is not good at learning slower dynamics, which are typically more interesting to researchers?
- Temporal coherence between generated protein conformations is mainly evaluated using auto-correlation in this work. Is it possible to show other metrics to capture detailed structural quality and variation during time evolution?
- Why is MDGen more effective at sequence recovery than MPNN? More explanation and analysis would be helpful here.
- Would it be possible to emulate MD simulation trajectory of the 12 fast folding proteins from [Shaw 2009](https://dl.acm.org/doi/abs/10.1145/1654059.1654126)? They are smaller than ATLAS proteins and longer than tetrapeptides, with much longer simulation time and rich dynamics.
- It would be nice to see if MDGen could infer a trajectory given an [apo/holo pair](https://arxiv.org/abs/2304.02198).
limitations: limitations: See weaknesses and questions.
paper_summary: paper_summary: The paper suggests a flow-based generative framework on molecular trajectories, with various downstream tasks such as forward simulation and transition path sampling. Additionally, the model is trained in a transferable setting, across tetrapeptides.
strength_weakness: strength_weakness: 1. Experiment baselines

The baselines of experiments are mostly the Markov State Models. I think it would also be good if there were some comparison between other models, though I understand that many prior works targeted Alanine dipeptide not tetrapeptides.
- Forward simulation: ITO$^{[1]}$, Timewarp$^{[2]}$

- Interpolation (Transition path sampling): PIPS$^{[3]}$

2. (Minor) Necessity of additional tasks

The necessity of additional tasks relatively seems weak compared to tasks such as forward simulation, TPS, specially the inpainting design. Rather than additional tasks, ablations for stability might be a better? One can obviously see that scaling to long trajectories shows the stability against the time scale, and protein simulation shows the stability against space complexity.

**Minor typos, suggestions**

- Definition of S-MPNN only exists in the Appendix. It would great to point out that more details are in the appendix, in the first paragraph of section 4.4
- Figure 6 is not referenced in the main paper, only the appendix
- Figure 2F, reference of that blue indicates the side chains and orange indicates the backbones seems to be missing

[1] Implicit transfer operator learning: Multiple time-resolution surrogates for molecular dynamics, NIPS 2023

[2] Timewarp: transferable acceleration of molecular dynamics by learning time-coarsened dynamics, NIPS 2023

[3] Stochastic Optimal Control for Collective Variable Free Sampling of Molecular Transition Paths, NIPS 2023
questions: questions: 1. Difference between downstream tasks

Could the upsampling task be seen as a superset of interpolation? Since upsampling with two given frames would the same as interpolation.

2. Training on one tetrapeptide

Just curious, though the authors has presented a transferable setting, are there any results when the model is trained for a specific tetrapeptide and tested on downstream tasks?
limitations: limitations: 1. Unconditional generation

As the authors have mentioned, unconditional generation is impossible since the model relies on key frames.

Comments: title: Thanks for the clarifications
Response to Rebuttal
comment: Some of my concerns have been addressed. I raised my score.
Dear authors,

Thanks for your response. Most of my concerns have been properly addressed. I think this pioneering work on generative modeling for protein/peptide dynamics should be shared with the community to inspire follow-up work. I have raised my scores accordingly.
Thank you for your clarifications. Most of my concerns have been addressed. I choose to keep my positive rating.
Thank you for the clarifications. This answers all my questions and I encourage that all updates announced here will be integrated in the final version. As the changes do not change the main message and claims, I will keep my score as it is, so I continue to support acceptance of the paper.
I thank the authors for the detailed response and additional experiments. Additional experiments on forward simulation shows better performance of MDGen over Timewarp and ITO, where prior works fail in longer simulation time. All questions have been resolved, and I have raised my score accordingly.