A Multi-Grained Symmetric Differential Equation Model for Learning Protein-Ligand Binding Dynamics

Authors: shengchao  liu
weitao  du
yanjing  li
zhuoxinran  li
vignesh  bhethanabotla
nakul  rampal
omar  yaghi
u c berkeley
anima  anandkumar
hongyu  guo
jennifer  chayes
Venue: ICLR_cc_2024_Workshop_AI4DiffEqtnsInSci
Published: 2024-02-11 06:14:14

Abstract: In drug discovery, molecular dynamics (MD) simulation for protein-ligand binding provides a powerful tool for predicting binding affinities, estimating transport properties, and exploring pocket sites. There has been a long history of improving the efficiency of MD simulations through better numerical methods and, more recently, by augmenting them with machine learning (ML) methods. Yet, challenges remain, such as accurate modeling of extended-timescale simulations. To address this issue, we propose NeuralMD, the first ML surrogate that can facilitate numerical MD and provide accurate simulations of protein-ligand binding dynamics. We propose a principled approach that incorporates a novel physics-informed multi-grained group symmetric framework. Specifically, we propose (1) a BindingNet model that satisfies group symmetry using vector frames and captures the multi-level protein-ligand interactions, and (2) an augmented neural ordinary differential equation solver that learns the trajectory under Newtonian mechanics. For the experiment, we design ten single-trajectory and three multi-trajectory binding simulation tasks. We show the efficiency and effectiveness of NeuralMD, with a 2000$\times$ speedup over standard numerical MD simulation and outperforming all other ML approaches by up to \textasciitilde 80\% under the stability metric. We further qualitatively show that NeuralMD reaches more stable binding predictions.

Main paper: 
 introduction: 
 The simulation of protein-ligand binding dynamics is one of the fundamental tasks in drug discovery (Kairys et al., 2019; Yang et al., 2020; Volkov et al., 2022) . Such simulations of binding dynamics are a key component of the drug discovery pipeline to select, refine, and tailor the chemical structures of potential drugs to enhance their efficacy and specificity. To simulate the protein-ligand binding dynamics, numerical molecular dynamics (MD) methods have been extensively developed. However, the numerical MD methods are computationally expensive due to the expensive force calculations on individual atoms in a large protein-ligand system.To alleviate this issue, machine learning (ML) surrogates have been proposed to either augment or replace numerical MD methods to estimate the MD trajectories. However, all prior ML approaches for MD are limited to single-system dynamics (e.g., small molecules or proteins) and not proteinligand binding dynamics. A primary reason is the lack of large-scale datasets. The first large-scale dataset with binding dynamics was released in May 2023 (Siebenmorgen et al., 2023) , and to our knowledge, we are now the first to explore it in this paper. Further, prior ML-based MD approaches limit to studying the MD dynamics on a small time interval (1e-15 seconds), while simulation on a longer time interval (e.g., 1e-9 seconds) is needed for specific tasks, such as detecting the transient and cryptic states (Vajda et al., 2018) in binding dynamics. However, such longer-time MD simulations are challenging due to the catastrophic buildup of errors over longer rollouts (Fu et al., 2022a) .Another critical aspect that needs to be integrated into ML-based modeling is the group symmetry present in protein-ligand geometry. Specifically, the geometric function should be equivariant to rotation and translation (i.e., SE(3)-equivariance). The principled approach to satisfy equivariance is to use vector frames, which have been previously explored for single molecules (Jumper et al., 2021) , but not yet for the protein-ligand binding complexes. The vector frame basis achieves SE(3)equivariance by projecting vectors (e.g., positions and accelerations) into the vector frame basis, and such a projection can maintain the equivariant property with efficient calculations (Liu et al., 2023) .Our Approach, NeuralMD. We propose NeuralMD, a multi-grained physics-informed approach designed to handle extended-timestep MD simulations for protein-ligand binding dynamics. Our multi-grained method explicitly decomposes the complexes into three granularities: the atoms in ligands, the backbone structures in proteins, and the residue-atom pairs in the complex, to obtain a scalable approach for modeling a large system. We achieve group symmetry in BindingNet through the incorporation of vector frames, and include three levels of vector frame bases for multi-grained modeling, from the atom and backbone level to the residue level for binding interactions.Further, our ML approach NeuralMD preserves the Newtonian mechanics. In MD, the movement of atoms is determined by Newton's second law, F = m • a, where F is the force, m is the mass, and a is the acceleration of each atom. By integrating acceleration and velocity w.r.t. time, we can obtain the velocities and positions, respectively. Thus in NeuralMD, we formulate the trajectory simulation as a second-order ordinary differential equation (ODE) or second-order stochastic differential equation (SDE) problem. Specifically, we augment derivative space by concurrently calculating the accelerations and velocities, allowing simultaneous integration of velocities and positions. 
preliminiaries: 
 Ligands. In this work, we only consider binding complexes with small molecules as ligands. Small molecules can be treated as sets of atoms in the 3D Euclidean space, {f (l) , x (l) }, where f (l) and x (l) represent the atomic numbers and 3D Euclidean coordinates for atoms in each ligand, respectively.Proteins. Proteins are essentially chains of amino acids or residues, where there are 20 natural amino acids. Noticeably, amino acids are made up of three components: a basic amino group (-NH 2 ), an acidic carboxyl group (-COOH), and an organic R group (or side chain) that is unique to each amino acid. Additionally, the carbon that connects all three groups is called C α . (We refer to this Wiki page for more details.) In this work, due to the large volume of atoms in proteins, we will use coarse-grained modelings on proteins and binding complexes. With this regard, the backbone-level data structure for each protein is {f (p) , {x(p) N , x (p) Cα , x (p)C }}, for the residue type and the coordinates of N -C α -C in each residue, respectively. (We may ignore the superscript in the coordinates of backbone atoms for brevity since such backbone structures are unique for residues in proteins) In addition to the backbone level, for a coarser-grained data structure, we further consider residue-level modeling for binding interactions, {f (p) , x (p) }, where the coordinate of C α is treated as the residue-level coordinate, i.e., x (p) x (p) Cα . Molecular Dynamics Simulations. Generally, molecular dynamics (MD) describes how each atom in a molecular system moves over time, following Newton's second law of motion:EQUATION)where F is the force, m is the mass, a is the acceleration, x is the position, and t is the time. Then, an MD simulation will take Newtonian dynamics to get the trajectories. The Numerical molecular dynamics (MD) methods fall into two categories: Newtonian and Langevin dynamics. Newtonian dynamics is ideal for systems with minimal thermal effects or when deterministic trajectories are needed. Conversely, Langevin dynamics is preferred for systems where thermal effects are significant, particularly at finite temperatures. In this study, we introduce two versions: an ordinary differential equation (ODE) solver for Newtonian dynamics and a stochastic differential equation (SDE) solver for Langevin dynamics.Problem Setting: Protein-Ligand Binding Dynamics Simulation. In this work, we focus on simulating the protein-ligand binding dynamics in the semi-flexible setting Salmaso & Moro (2018) , i.e., proteins with rigid structures and ligands with flexible movements (Siebenmorgen et al., 2023) . Thus, the problem is formulated as follows: suppose we have a rigid protein structure {f (p) , {xEQUATIONC }} and a ligand with its initial structure and velocity, {f (l) , xEQUATION0 }. We want to predict the trajectories of ligands following the Newtonian dynamics, i.e., the movement of {x (l) t , ...} over time. We also want to clarify two critical points about this problem setting. (1) We consider trajectory prediction, i.e., positions as labels, and no explicit energy and force labels are considered. ML methods for energy prediction followed with numerical ODE/SDE solver may require smaller timescales (around 1e-15 seconds), while trajectory prediction, which directly predicts the positions, is agnostic to the magnitude of timescales. This is appealing for datasets like MISATO with larger timescales (1e-9 seconds). ( 2) Each trajectory is composed of a series of geometries of ligands, and such geometries are called snapshots. 
method: bindingnet and neuralmd: 
 In this section, we introduce BindingNet, a multi-grained SE(3)-equivariant geometric model for protein-ligand binding. The input of BindingNet is the geometry of the rigid protein and the ligand at time t, while the output is the force on each atom in the ligand. Our architecture is SE(3) equivariant by extending the equivairant frames for small molecules in Du et al. (2022) to 1. Atom level: F ligand ; 2. Protein level: F protein ; 3. F complex . See appendix D for the formal definition.Atom-Level Ligand Modeling. We first generate the atom embedding using one-hot encoding and then aggregate each atom's embedding, z (l) , by aggregating all its neighbor's embedding within the cutoff distance c. Then, we obtain the atom's equivariant representation by aggregating its neighborhood's messages as (x(l) i -x (l) j ) • z (l)i . A subsequent scalarization is carried out based on the atom-level vector frame as h(l) ij = (h (l) i ⊕ h (l) j ) • F ligand ,where ⊕ is the concatenation. Finally, it is passed through several equivariant message-passing layers (MPNN). The outputs are atom representation and vector (h (l) and vec (l) ), and they are passed to the complex module.Backbone-Level Protein Modeling. For the coarse-grained modeling of proteins, we consider three backbone atoms in each residue. We first obtain the atom embedding on three atom types, and then we obtain each atom's representation z (p) by aggregating its neighbor's representation. Then, we obtain an equivariant atom representation by aggregating the edge information, (x(p) i -x (p) j ) • z (p)i , within cutoff distance c. Following which is the scalarization on the residue frame h(p) ij = (h (p) i ⊕ h (p) j ) • F protein .Recall that we also have the residue type, and with a type embedding z(p) , we can obtain the final residue-level representation using an MPNN layer ash (p) = z(p) + (h (p) N,Cα + h (p)Cα,C )/2. We leave the Residue-Level Complex Modeling in appendix F.As clarified in Section 2, molecular dynamics follows Newtonian dynamics, and we solve it as an ordinary differential equation (ODE) problem or a stochastic differential equation (SDE) problem. The BindingNet takes in the molecular system geometry (x (l) t , x (p) ) at arbitrary time t, and outputs the forces.To learn the MD trajectory following second-order ODE, we propose the following formulation of the second-order ODE within one integration call:EQUATION)where F is the output forces from BindingNet. This means we augment ODE derivative space by concurrently calculating the accelerations and velocities, allowing simultaneous integration of velocities and positions. Ultimately, following Newtonian mechanics, the coordinates at time t are integrated as:EQUATIONC ).(3)On the other hand, Langevin dynamics introduces a stochastic component for large molecular systems with thermal fluctuations. Concretely, Langevin dynamics is an extension of the standard Newtonian dynamics with the addition of damping and random noise terms:EQUATIONEQUATION)The objective is the mean absolute error between the predicted coordinates and ground-truth coordinates: L = E t | x(l) t -x (l) t | . An illustration of NeuralMD pipeline is in appendix Figure 2. 
experiments: 
 Datasets. We consider MISATO in our work (Siebenmorgen et al., 2023) . It is built on 16,972 experimental protein-ligand complexes extracted from the protein data bank (PDB) (Berman et al., 2000) . For each protein-ligand complex, the trajectory comprises 100 snapshots in 8 nanoseconds under the fixed temperature and pressure. We want to highlight that MD trajectories allow the analysis of small-range structural fluctuations of the protein-ligand complex. See appendix E for the basic statistics.Experiments Settings. We consider two experiment settings. The first type of experiment is the single-trajectory prediction, where both the training and test data are snapshots from the same trajectory, and they are divided temporally. The second type of experiment is the multi-trajectory prediction, where each data point is the sequence of all the snapshots from one trajectory, and the training and test data correspond to different sets of trajectories.Evaluation Metrics. For both experiment settings, the trajectory recovery is the most straightforward evaluation metric. To evaluate this, we take both the mean absolute error (MAE) and mean squared error (MSE) between the predicted coordinates and ground-truth coordinates. Stability is also an important metric for evaluating the predicted MD trajectory. The intuition is that the prediction on long-time MD trajectory can enter a pathological state (e.g., bond breaking), and stability is the measure to quantify such observation. It is defined as P i,j ∥x i -x j ∥ -b i,j > ∆, where b i,j is the pair distance at the last snapshot (the most equilibrium state), and we take ∆ = 0.5 Å. Another metric considered is frames per second (FPS) on a single Nvidia-V100 GPU card, and it measures the MD efficiency.Results on single-trajectory prediction We provide the visualization of some examples in Figure 1 . Quantitative results are in Table 2 and baselines are introduced in appendix G. The first observation is that the baseline VertletMD has a clear performance gap compared to the other methods. This verifies that using ML models to predict the energy (or force) at each snapshot, and then using a numerical integration algorithm can fail in the long-time simulations (Fu et al., 2022a) . Additionally, we can observe that the recovery error of trajectory (MAE and MSE) occasionally fails to offer a discernible distinction among methods (e.g., for protein-ligand complex 3EOV, 1KT1, and 4G3E), though NeuralMD is slightly better. However, the stability (%) can be a distinctive factor in method comparisons, where we observe NeuralMD outperform on all 10 tasks up to ˜80%. More detailed analysis is also provided in appendix G.Numerical methods for MD can be classified into classical MD and ab-initio MD, where ab-initio MD calculates the forces using a quantum-mechanics-based method, such as DFT, while classical MD uses an approximated function fit to ab-initio calculations to predict the atomic forces. More recently, the machine learning methods for MD have opened a new perspective by utilizing the group symmetric tools for geometric representation and the neural ODE Chen et al. ( 2018).On one hand, many works have studied the protein-ligand binding problem in the equilibrium state (Stepniewska-Dziubinska et al., 2018; Jiménez et al., 2018; Jones et al., 2021; Yang et al., 2023) , but not the MD simulation for binding dynamics. On the other hand, existing machine learning (ML) methods have studied molecular simulation (Zhang et al., 2018; Doerr et al., 2020; Musaelian et al., 2023; Fu et al., 2022b) , but they are mainly studying small molecules or proteins.In this work, we propose an ML framework as an MD simulation to learn the protein-ligand binding dynamics. In this article, a 3D molecular graph is represented by a collection of 3D point clouds. The corresponding symmetry group is SE(3), which consists of translations and rotations. Recall that we define the notion of equivariance functions in R 3 in the main text through group actions. Formally, the group SE( 3) is said to act on R 3 if there is a mapping ϕ : SE(3) × R 3 → R 3 satisfying the following two conditions:1. if e ∈ SE(3) is the identity element, then ϕ(e, r) = r for ∀r ∈ R 3 .2. if g 1 , g 2 ∈ SE(3), then ϕ(g 1 , ϕ(g 2 , r)) = ϕ(g 1 g 2 , r) for ∀r ∈ R 3 .Then, there is a natural SE(3) action on vectors r in R 3 by translating r and rotating r for multiple times. For g ∈ SE(3) and r ∈ R 3 , we denote this action by gr. Once the notion of group action is defined, we say a function f : R 3 → R 3 that transforms r ∈ R 3 is equivariant if:f (gr) = gf (r), for ∀ r ∈ R 3 .On the other hand, f : R 3 → R 1 is invariant, if f is independent of the group actions:f (gr) = f (r), for ∀ r ∈ R 3 .For some scenarios, our problem is chiral sensitive. That is, after mirror reflecting a 3D molecule, the properties of the molecule may change dramatically. In these cases, it's crucial to include reflection transformations into consideration. More precisely, we say an SE(3) equivariant function f is reflection anti-symmetric, if:EQUATIONfor reflection ρ ∈ E(3). Frame is a popular terminology in science areas. In physics, the frame is equivalent to a coordinate system. For example, we may assign a frame to all observers, although different observers may collect different data under different frames, the underlying physics law should be the same. In other words, denote the physics law by f , then f should be an equivariant function.There are certain ways to choose the frame basis, and below we introduce two main types: the orthogonal basis and the protein backbone basis. The orthogonal basis can be built for flexible 3D point clouds such as atoms, while the protein backbone basis is specifically proposed to capture the protein backbone.Since there are three orthogonal directions in R 3 , one natural frame in R 3 can be a frame consisting of three orthogonal vectors: F = (e 1 , e 2 , e 3 ).Once equipped with a frame (coordinate system), we can project all geometric quantities to this frame. For example, an abstract vector x ∈ R 3 can be written as x = (r 1 , r 2 , r 3 ) under the frame F , if: x = r 1 e 1 + r 2 e 2 + r 3 e 3 . An equivariant frame further requires the three orthonormal vectors in (e 1 , e 2 , e 3 ) to be equivariant. Intuitively, an equivariant frame will transform according to the global rotation or translation of the whole system. Once equipped with an equivariant frame, we can project equivariant vectors into this frame:x = r1 e 1 + r2 e 2 + r3 e 3 .We call the process of x → r := (r 1 , r2 , r3 ) the scalarization or projection operation. Since ri = e i • x is expressed as an inner product between equivariant vectors, we know that r consists of scalars.In this article, we assign an equivariant frame to each node/edge, therefore we call them the local frames. Given two atoms with 3D positions (x i , x j ), we can find the atom (denoted by x k ) that is nearest to the center of (x i , x j ) by KNN algorithms. Then the equivariant frame is defined by:EQUATIONThe Gram-Schmidt orthogonalization makes sure that the Vector-Frame(x i , x j ) is orthonormal.Reflection Antisymmetric Since we implement the cross product × for building the local frames, the third vector in the frame is a pseudo-vector. Then, the projection operation is not invariant under reflections (the inner product between a vector and a pseudo-vector change signs under reflection). Therefore, our model can discriminate two 3D geometries with different chiralities.Our local frames also enable us to output equivariant vectors by multiplying scalars (v 1 , v 2 , v 3 ) with the frame:v = v 1 • e 1 + v 2 • e 2 + v 3 • e 3 .Equivariance w.r.t. cross-product The goal is to prove that the cross-product is equivariant to the SE(3)-group, i.e.:gx × gy = g(x × y), g ∈ SE(3)-Group (9)Geometric proof. From intuition, with rotation matrix g, we are transforming the whole basis, thus the direction of gx × gy changes equivalently with g. And for the value/length of gx × gy, because |gx × gy| = ∥gx∥ • ∥gy∥ • sin θ = ∥x∥ • ∥y∥ • sin θ = |x × y|. So the length stays the same, and the direction changes equivalently. Intuitively, this interpretation is quite straightforward.Analytical proof. A more rigorous proof can be found below: First, we have that for the rotation matrix g:EQUATIONwhereg i , x, y ∈ R 3×1 . Because A T C • B T D -A T D • B T C = (A × B) T (C × D), so we can have:EQUATIONThen because:det(g) = (g2 × g3) T g1 = g T 1 g1 = 1 =⇒(g2 × g3) T g1g -1 1 = g T 1 g1g -1 1 =⇒(g2 × g3) T = g T 1 .(12)Thus, we can haveEQUATIONVector-Frame(xi+ δx, xj + δx) = Gram-Schmidt{xi -xj, xi -x k , (xi -xj) × (xi -x k )}. (18)Because the basis is based on the difference of coordinates, it is straightforward to observe thatGram-Schmidt{v 1 + t, v 2 + t, v 3 + t} = Gram-Schmidt{v 1 , v 2 , v 3 }.So the frame operation is transition equivariant.Vector-Frame(xi, xj) ̸ = Vector-Frame(-xi, -xj).From intuition, this makes sense because the cross-product is anti-symmetric.A simple counter-example is the original geometry R and the reflected geometry by the original point -R. Thus the two bases before and after the reflection group is the following:EQUATION)The bases between v 1 , v 2 , v 3 and {-v 1 , -v 2 , v 3 }} are different, thus such frame construction is reflection anti-symmetric.Once we have the three vectors as the vector frame basis, the next thing is to do modeling. Suppose the frame is F = (e 1 , e 2 , e 3 ), then for an equivariant vector (tensor) h, the scalarization is:EQUATIONProteins are essentially macromolecules composed of thousands of residues (amino acids), where each residue is a small molecule. Thus, it is infeasible to model all the atoms in proteins due to the large volume of the system, and such an issue also holds for the protein-ligand complex. To address this issue, we propose BindingNet, a multi-grained SE(3)-equivariant model. The vector frame basis ensures SE(3)-equivariance, and the multi-granularity is achieved by considering frames at three levels.Vector Frame Basis for SE(3)-Equivariance. Recall that the geometric representation of the whole molecular system needs to follow the physical properties of the equivariance w.r.t. rotation and translation. Such a group symmetric property is called SE(3)-equivariance. We also want to point out that the reflection or chirality property is equivariant for properties like energy, yet it is not for the ligand modeling with rigid protein structures (i.e., antisymmetric to the reflection). The vector frame basis can handle this naturally, and we leave a more detailed discussion in Appendix D, along with the proof on group symmetry of vector frame basis. In the following, we introduce three levels of vector frames for multi-grained modeling.Atom-Level Vector Frame for Ligands. For small molecule ligands, we first extract atom pairs (i, j) within the distance cutoff c, and the vector frame basis is constructed using the Gram-Schmidt as:EQUATION)where × is the cross product. Note that both x (l) i and x(l) jare for geometries at time thenceforth, we omit the subscript t for brevity. Such an atom-level vector frame allows us to do SE(3)-equivariant message passing to get the atom-level representation.Backbone-Level Vector Frame for Proteins. Proteins can be treated as chains of residues, where each residue possesses a backbone structure. The backbone structure comprises an amino group, a carboxyl group, and an alpha carbon, delegated as N -C α -C. Such a structure serves as a natural way to build the vector frame. For each residue in the protein, the coordinates are x N , x Cα , and x C , then the backbone-level vector frame for this residue is:EQUATION)This is built for each residue, providing a residue-level representation.Residue-Level Vector Frame for Protein-Ligand Complexes. It is essential to model the protein-ligand interaction to better capture the binding dynamics. We achieve this by introducing the residue-level vector frame. More concretely, proteins are sequences of residues, marked as {(fEQUATIONi+1 , ...}. Here, we use a cutoff threshold c to determine the interactions between ligands and proteins, and the interactive regions on proteins are called pockets. We construct the following vector frame for residues in the pockets sequentially:F complex = ( x (p) i -x (p) i+1 x (p) i -x (p) i+1 , x (p) i × x (p) i+1 x (p) i × x (p) i+1 , x (p) i -x (p) i+1 x (p) i -x (p) i+1 × x (p) i × x (p) i+1 x (p) i × x (p) i+1).Through this complex-level vector frame, the message passing enables the exchange of information between atoms from ligands and residues from the pockets. The illustration of the above three levels of vector frames can be found in Figure 3 . Once we build up such three vector frames, we then conduct a scalarization operation (Du et al., 2022) , which transforms the equivariant variables (e.g., coordinates) to invariant variables by projecting them to the three vector bases in the vector frame.In this section, we provide more details on the MISATO dataset (Siebenmorgen et al., 2023) .For small molecule ligands, we ignore the Hydrogen atoms. We also plot the distribution of the energy gap between each time step and the initial snapshot, i.e., E t -E 0 . The distribution is in Figure 6 . We can observe that as the time processes, the mean of the energy stays almost the same, yet the variance gets higher. Residue-Level Complex Modeling. Once we obtain the atom-level representation and vector (h (l) , vec (l) ) from ligands, and backbone-level representation (h (p) ) from proteins, the next step is to learn the protein-ligand interaction. We first extract the residue-atom pair (i, j) with a cutoff c, based on which we obtain an equivariant interaction edge representation h ij = (h(l) i +h (p) j )•(x (l) i -x (p) j ). After scalarization, we can obtain invariant interaction edge representation h ij = h ij • F complex . Finally, we adopt an equivariant MPNN layer to get the atom-level force as:EQUATIONIn the last equation, the ultimate force predictions can be viewed as two parts: the internal force from the molecule vec(l)i and the external force from the protein-ligand interaction vec (pl) ij .In this section, we provide more details on the model architecture in Figure 7 , and hyperparameter details in Table 1 . To verify the effectiveness and efficiency of NeuralMD, we design ten single-trajectory and three multi-trajectory binding simulation tasks. For evaluation, we adopt the recovery and stability metrics (Fu et al., 2022a) . NeuralMD achieves 2000× speedup compared to the numerical methods. We observe that NeuralMD outperforms all other ML methods (Zhang et al., 2018; Musaelian et al., 2023; Fu et al., 2022b; Wu & Li, 2023; Arts et al., 2023) on 12 tasks using recovery metric, and NeuralMD is consistently better by a large gap using the stability metric (up to ˜80%). Qualitatively, we illustrate that NeuralMD realizes more stable binding dynamics predictions in three case studies. They are three protein-ligand binding complexes from Protein Data Bank (PDB), as shown in Figure 1 .Baselines. In this work, we mainly focus on machine learning methods for trajectory prediction, i.e., no energy or force labels are considered. GNN-MD is to apply geometric graph neural networks (GNNs) to predict the trajectories in an auto-regressive manner (Siebenmorgen et al., 2023; Fu et al., 2022b) . More concretely, GNN-MD takes as inputs the geometries at time t and predicts the geometries at time t + 1. DenoisingLD (denoising diffusion for Langevin dynamics) (Arts et al., 2023; Wu & Li, 2023; Fu et al., 2022b ) is a baseline method that models the trajectory prediction as denoising diffusion task (Song et al., 2020) , and the inference for trajectory generation essentially becomes the Langevin dynamics. CG-MD learns a dynamic GNN and a score GNN (Fu et al., 2022b) , which are essentially the hybrid of GNN-MD and DenoisingLD. Here, to make the comparison more explicit, we compare these two methods (GNN-MD and DenoisingLD) separately. Additionally, we consider VerletMD, an energy prediction research line (including DeePMD (Zhang et al., 2018) , TorchMD (Doerr et al., 2020) , and Allegro-LAMMPS (Musaelian et al., 2023) ), where the role of ML models is to predict the energy, and the MD trajectory is obtained by the velocity Verlet algorithm, a numerical integration method for Newtonian mechanics. We keep the same backbone model (BindingNet) for energy or force prediction for all the baselines.The main results in shown in Table 2 .One main benefit of using NeuralMD for binding simulation is its efficiency. To show this, we list the computational time in Table 3 . We further approximate the wall time of the numerical method for MD simulation (PDB 5WIJ). Concretely, we can get an estimated speed of 1 nanosecond of dynamics every 0.28 hours. This is running the simulation with GROMACS (Van Der Spoel et al., 2005 ) on 1 GPU with 16 CPU cores and a moderately sized water box at the all-atom level (with 2 femtosecond timesteps). This equivalently shows that NeuralMD is ˜2000× faster than numerical methods. A more challenging task is to test the generalization ability of NeuralMD among different trajectories. The MISATO dataset includes 13,765 protein-ligand complexes, and we first create two small datasets by randomly sampling 100 and 1k complexes, respectively. Then, we take 80%-10%-10% for training, validation, and testing. We also consider the whole MISATO dataset, where the data split has already been provided. After removing the peptide ligands, we have 13,066, 1,357, and 1,357 complexes for training, validation, and testing, respectively.The quantitative results are in Table 4 . First, we can observe that VerletMD has worse performance on all three datasets, and the performance gap with other methods is even larger compared to the single-trajectory prediction. The other two baselines, GNN-MD and DenoisingLD, show similar performance, while NeuralMD outperforms in all datasets. Notice that stability (%) remains more distinguishable than the two trajectory recovery metrics (MAE and MSE).Recall that, in the main paper, we have discussed using the semi-flexible binding setting, i.e., proteins with rigid structures while small molecule ligands with flexible structures, and the goal is to predict the trajectories of the ligands. If we want to take both proteins and ligands with flexible structures, one limitation is the GPU memory cost. However, we would like to mention that it is possible to do NeuralMD on protein-ligand with small volume, and we take an ablation study to test them as below.Problem Setup. Both the proteins and ligands are flexible, and we want to predict their trajectories simultaneously. In the main paper, we consider three levels of vector frames. Here in the flexible setting, due to the large volume of atoms in the protein-ligand complex, we are only able to consider two levels, i.e., the atom-level and residue-level. Thus, the backbone model (BindingNet) also changes accordingly. The performance is shown in Table 5 , and we can see that NeuralMD is consistently better than the GNN-MD on all three metrics and all 10 single trajectories. We omit the multi-trajectory experiments due to the memory limitation.  
 
 Appendix:  



Reviews: review: title: title: Authors propose NeuralMD which is a ML surrogate model, multi-grained symmetric differential equation model for learning protein-ligand binding dynamics.
main_review: main_review: Authors propose NeuralMD which is a ML surrogate model, multi-grained symmetric differential equation model for learning protein-ligand binding dynamics.
title: title: Review for NeuralMD
main_review: main_review: The authors have introduced a ML surrogate model for learning the Protein-ligand binding dynamics. The manuscript is well written with all the relevant details explicitly mentioned. 

I have few overall concerns in this topic, please address them below

1. ML surrogates are highly dependent on the dataset they are trained for. What exactly is the motivation here to develop ML surrogates here ?
2. Is there any real time computation of MD simulation that is really needed for warranting the development of ML surrogates 
2. ML surrogates are not extrapolative in nature and perform poorly on cases outside the dataset considered for train and testing. In such as case how do we account for such outliers in online computation.