UniSim: A Unified Simulator for Time-Coarsened Dynamics of Biomolecules
Authors: Ziyang Yu
Wenbing Huang
Yang Liu
Venue: ICML 2025 poster (2025)
Keywords: molecular dynamics
unified 3D molecular pretraining
forward simulation
cross-domain transferability
Primary area: applications->chemistry_physics_and_earth_sciences

Abstract: Molecular Dynamics (MD) simulations are essential for understanding the atomic-level behavior of molecular systems, giving insights into their transitions and interactions. However, classical MD techniques are limited by the trade-off between accuracy and efficiency, while recent deep learning-based improvements have mostly focused on single-domain molecules, lacking transferability to unfamiliar molecular systems. Therefore, we propose **Uni**fied **Sim**ulator (UniSim), which leverages cross-domain knowledge to enhance the understanding of atomic interactions. First, we employ a multi-head pretraining approach to learn a unified atomic representation model from a large and diverse set of molecular data. Then, based on the stochastic interpolant framework, we learn the state transition patterns over long timesteps from MD trajectories, and introduce a force guidance module for rapidly adapting to different chemical environments. Our experiments demonstrate that UniSim achieves highly competitive performance across small molecules, peptides, and proteins.

UniSim: A Unified Simulator for Time-Coarsened Dynamics of Biomolecules
Ziyang Yu Wenbing Huang? 23
Yang Liu
Abstract
small molecules
Molecular Dynamics (MD) simulations are essen- tial for understanding the atomic-level behavior molecular systems, giving insights into their tran- sitions and interactions. However, classical MD techniques are limited by the trade-off between curacy and efficiency, while recent deep learning- based improvements have mostly focused on single-domain molecules, lacking transferability to unfamiliar molccular systems. Therefore, we propose Unified Simulator (UniSim), which lever- ages cross-domain knowledge enhance the un- derstanding of atomic interactions. First, ploy multi-head pretraining approach to learn unified atomic representation model from large and diverse set of molecular data. Then, based on the stochastic interpolant framework, we learn the state transition patterns over long timesteps from MD trajectories, and introduce force guidance module for rapidly adapting different chemical environments. Our experiments demonstrate that UniSim achieves highly competitive performance across small molecules, peptides, and proteins.
timestep
peptides
UniSim
proteins
Figure I. UniSim enables time-coarsened dynamics simulations small molecules, peptides, and proteins over long timestep
2018). One of the MD's core objectives is to generate trajectories of molecular states that adhere to underlying physics constraints
a period of time, given the initial state of the molecular system and environment configura- tions (Allen et al., 2004). To achieve this end, classical MD methods (van Gunsteren & Oostenbrink, 2024) update the next step of motion by numerically integrating Newton equations or Langevin dynamics (Langevin, 1908), with the potential energy and atomic forces calculated based on the current molecular state. It should be noted that the stability of the numerical integration requires extremely small timestep At = 10 10-15g during MD simulations (Plimpton, 1995).
1. Introduction
Molecular Dynamics (MD) simulations, an in silico method used to comprehend the time evolution of molecular sys- tems in given environments, serve as a fundamental and essential tool in various fields like computational chem- istry, pharmacology, material design, and condensed mat- ter physics (Van Gunsteren & Berendsen, 1990; I Lindorff- Larsen et al., 2011 1; Hollingsworth & Dror, 2018; Lau et al.,
In terms of how energy is calculated, classical MD meth- ods can be divided into Quantum Mechanies (QM) meth- ods (Griffiths & Schroeter, 2019) and empirical force field methods (Pearlman et al., 1995; Vanommeslaeghe al., 2010). On the one hand, QM methods provide highly ac- curate energy calculations, while their great computational complexity makes it prohibitively expensive for accurate transition path sampling on a long time scale, like protein folding (Lindorff- Larsen al., 2011), On the other hand, empirical force field methods accelerate the sampling of conformation transition pathways, some MD methods based on reinforcement learning (Shin et al., 2019), adaptive sampling (Markwick et al., 2011: Botu & Ramprasad, 2015a), as well as enhanced sampling (Bal Neyts, 2015; M. Sultan & Pande. 2017; Wang et al., 2021; Célerse et al., 2022) have been proposed, yet achieving a
"Equal contribution 'Department of Computer Science and Technology, Tsinghua University, Beijing, China Gaolligg School Artificial intelligence, Renmin University of China, Beijing, China 'Enginerrin Research Center of text-Generation Intelli- gent Search and Recommendation, MOE "Institute for Al Industry Research (AIR), Tsinghua University. Beijing, China. Correspon- dence to: Wenbing Huang <hwenbing @ 126.com>, Yang Liu <liuyang201 tsinghua.cdu.cn>.
faster but less accurate. To
Proceedings of the ndnd International Conference Machine Learning. Vancouver, Canada. PMLR 267, 2025. Copyright 2025 by the author(s).

---

UniSim: A Unified Simulator for Time-Coarsened Dynamics of Biomolecules
qualitative improvement in efficiency while maintaining the accuracy remains challenging.
from MD trajectories, and introduce force guidance module for rapidly adapting to different chemical envi- ronments.
Recently, a surge of deep learning methods have been pro- posed to boost MD simulations from different aspects (Noé et al., 2019; Köhler et al., 2023; Lu et al., 2024; Klein et al., 2024; Schreiner et al., 2024; Wang et al., 2024a;b; Yu et al., 2024). Specifically, a series of methods categorized as time-coarsened dynamics aim to accelerate simulations by learning the push forward from X, to X,+> with much larger timestep T >> At. where X, denotes the molecular state the wall-clock time & (Klein et al., 2024; Schreiner et al., 2024; Yu et al., 2024). Although these methods can in principle achieve rapid long- ttmm sampling, several issues arise in practical applications: 1) Almost all previous works are restricted to a single molecular domain (e.g., peptides or proteins) within a fixed environment, lacking the transfer- ability across different scenarios. 2) Some models leverage hand-crafted representations specific domains (e-g.. " carbons in leucines), which significantly impair their ability to recognize unfamiliar molecules, such as proteins with unnatural amino acids (Link et al., 2003).
2. Related Work
3D Molecular Pretraining Given the success of the pre- training paradigm in the fields of NLP, some recent works have proposed pretraining methods based on 3D molecular structures (7aidi et al., 2023; Luo et al., 2022; Zhang et al., 2023; Gao et al., 2022; Jiao et al., 2023; Zhou et al., 2023; Feng et al., 2024; Jiao et al., 2024; Zhang et al., 2024). Considering the scarcity of labeled molecular data, the Noisy (Zaidi et al., 2023) method proposes the denois- ing regularization paradigm for self-supervised pretraining on equilibrium conformations. Correspondingly, machine learning forcefields (Botu & Ramprasad, 2015b; Chmicla et al., 2017: 2018; 2023) explore the off-equilibrium con- formation space by force-centric training. Furthermore, ET-OREO (Feng et al., 2024) proposes novel pretrain- ing method for learning unified representations encompass- ing both equilibrium and off-equilibrium conformations. EPT (Jiao et al., 2024) adopts the denoising pretraining on coarse-grained blocks and extends to learn representations of nomti- molecules. Additionally, DPA-2 (Zhang et al., 2024) adapts to diverse chemical and materials sys- tems by leveraging multi-task approach for different data sources. In this work, we propose novel pretraining ap- proach that is scalable to multi-domain biomolecules and adopts the multi-task approach to distinguish different force- fields and equilibrium states.
Taking all these into consideration, we propose a pre- trained model for Unified full-atom time-coarsened dynam- ies Simulation (UniSim), which is transferable small molecules, peptides as well as proteins, and can be easily adapted to various chemical environments by parameter- efficient training. An illustration for one-step simulation performed by UniSim is displayed in Figure 1. Firstly, ow- ing to the scarcity of MD trajectory datasets, we propose to pretrain a unified atomic representation model on multi- domain 3D molecular datasets under different chemical en- vironments and equilibrium states. Based on the pretrained model, we then leverage the stochastic interpolants gener- ative framework (Albergo et al., 2023) to learn the push forward from X, to X with a predefined long timestep T. Finally, in order to better adapt to specific chemical envi- ronments with different force conditions (temperatures, pres- sures, solvents etc.), we follow FBM (Yu et al., 2024) and introduce force guidance regulate the sampling process of molecular trajectories. In summary, our contributions are:
Cone--Carseeeed Dynamics To overcome the constraints of numerical integration stability of classical MD simula- tions, a series of deep Icarning-bac -bsed methods (Fu et al., 2023; Klein et al., 2024; Li et al., 2024; Jing et al., 2024; Schreiner et al., 2024; Hsu et al., 2024; Klein & Noe, 2024; Du et al., 2024; Yu et al., 2024) learn the push forward from X, to X.+r over long timesteps T to enable rapid state tran- sitions, which can be classified as time-coarsened dynamics. Various generative frameworks have been leveraged to learn the all-atom dynamics from data pairs of trajectories, in- cluding augmented normalizing flows (Klein et al., 2024), score-based diffusion (Hsu et al., 2024; Schreiner et al., 2024), flow matching (Klein & Noe, 2024) and stochastic interpolants (Du et al., 2024; Yu et al., 2024). On the con- trary. Folow (Li et al., 2024) and MDGen (Jing et al., 2024) learn dynamics of proteins with boarse-grained representa- tions. Notably, FBM (Yu et al., 2024) introduces the force guidance to comply with the underlying Boltzmann distri- bution, achieving state-of-theaa performance on peptides. According to FBM, we first learn a unified vector field based on the stochastic interpolants framework, then we leverage the force guidance for parameter-efficien fine-tuning to
To our best knowledge, UniSim is the first deep learning-based generative model that tailored for trans- ferable time-coarsened dynamics on cross-domain molecular systems.
We employ a multi-task pretraining approach to learn a unified atomic representation model from cross- domain molecular data, leveraging novel techniques to tackle unbalanced molecular scales and provide fine- grained atomic representations.
Based on the stochastic interpolant framework, we learn the state transition patterns over long-time steps
2

---

UniSim: A Unified Simulator for Time-Coarsened Dynamics of Biomolecules
diverse chemical environments and forcefields.
3.2. Unified Pretraining
In this section, we will introduce the techniques applied for pretraining UniSim on multi-domain 3D molecular datasets. Please note that, to comply with the symmetry properties of 3D molecules, we leverage the neural network TorchMD-NER (Pelaez et al., 2024) as the model architecture.
3. Method
We present the overall workflow of UniSim in this section with an illustration in Figure 2. In $ 3.1, we will define the task formulation with necessary notations. In $ 3.2, we will propose novel techniques for obtaining unified atomic representation model by pretraining on diverse datasets of multi-domain biomolecules. Afterwards, will introduce the regimen of learning time-coarsened dynamics based on the stochastic interpolant in $ 3.3. Finally, we will in- corporate the force guidance technique to adapt to diverse chemical environments and present the scheme of parameter- efficient fine-tuning in $ 3.4. All proofs of propositions are provided in $ B.
graph
Subgraph A crucial challenge in training unified representation models arises from the vast scale discrepancy between molecular systems: small molecules typically contain tens of atoms, while pro- teins often comprise hundreds or thousands atoms. Most of previous works that construct molecular graphs using KNN (Kong et al., 2023) or radius cutoff (Schreiner et al., 2024) neglect the scale discrepancy, which may inhibit the transferability across molecular domains. To prop- erly address the issue, we propose the so-called gradient- environment subgraph method to bridge the gap in scales of cross-domain molecules.
3.1. Task Formulation
We assume that there are m datasets (Di,... D con- sisting of 3D molecular conformations and datasets (TI.... Ti) consisting of MD trajectories. Each molecule represented as G = (Z, X). where Z E RN denotes the atomic types of N atoms, X E Nx3 denotes the atomic positions. If the molecular conformation is off-equilibrium, we denote E 1g- + IR as the MD potential function and F-
Specifically, for each macromolecule G with more than 1,000 atoms, we randomly select an atom e with Cartesian coordinates T. E RR. Given the predefined thresholds S 8min < 6max, the gradient subgraph Gy and the environment subgraph G. are constructed follows:
Ve X)= Nx3 as its MD forces, otherwise F 0. Our goals are listed below:
g3 Gاlا E g, ||ā - -Tall2 < 8min);
(2)
G. = Lili E 1|ā5 - -Eell2 < 8max); }.
(3)
1. We first obtain a unified atomic representation model @p pretrained on the datasets {Di,... .Dm). Formally, the pretrained model defined as:
It is easy to show that G, S G... Moreover, for each atom in gy, the distance to any other atom in g \ Ge should be at least 8 Omin- Therefore, as long as 6max ômii is set to be large enough, interactions between atoms outside of G. and those in Gg can be neglected, which is consistent with physics principles.
H, V = ,(T, (TT, X)),,
where E 13 is set for the compatibility to the generation framework, H E Nx1l denotes atomic representations of H channels, and V E Nx3*H denotes SO(3)- -qqiivarra vectors of all atoms.
A Afterwards, the environment subgraph Ge rather than the whole graph G will serve as the input to the model. Edges between atoms are constructed based on a predefined cutoff reut- It should be noted that only the atoms within G participate in the calculation of the training objective, since the contribution of the truncated atoms of the original graph G to the atoms within G. \ G, cannot be shielded.
2. Given the coarsened timestep 7 and training data pairs {(G., G+++)) randomly sampled from trajec- tory dataset T; (1 < sis l), we train a vector field model ø with the pretrained & serving as a graph en- coder, which learns the push forward from X, to X based on the stochastic interpolant framework.
A Aomic Embedding Expansion In addition to the scale discrepancy, the molecular specificity is another key factor hindering the development of a unified atomic representa- tion. In domains like proteins, atoms of the same type ex- hibit different but regular patterns (e.g., CA and CB), while wet lab experiments (Rossmann & Arnold, 2001) elucidate that atoms of the same pattern probably share consistent properties like bond lengths. Due to the hard constraints on bond lengths and angles, these patterns exhibit discrete char- acteristics, which 3D GNNs tailored for continuous features
3. For adapting to diverse chemical environments (e.g., solvation), we train a force guidance kernel Ç on the corresponding trajectory dataset T, (1 < j < l), which incorporates underlying physics principles into genera- tion (Yu et al., 2024). The parameters of networks yp, o are kept frozen, serving as reusable backbones.
3

---

UniSim: A Unified Simulator for Time-Coarsened Dynamics of Biomolecules
a. pretraining
b. vector field model
c. force guidance kernel
Velectty Head
Guidance Kernel up
4
Denoiser Head Oz
E [0,1]
froeen parameters
Unified Atomic Represontation Model
protrained
d. inference
=(t) R = b((x..))
- e(()) (2e())d
y(t) (t )
Denoising Head
Hood
Head K.
c=
Figure 2. Illustration for the overall workflow of UniSim. a. The unified atomic representation model yp is pretrained on multi-domain 3D molecules, where data from different chemical environments fed the corresponding output head. b. Based the stochastic interpolant framework, vector field models v, 77s are trained on MD trajectories to learn the push forward from X, to Xe+r with timestep To adapt different chemical environements, additional networks V, & are trained parameters frozen. d. Given an initial state, inference is performed by iteratively solving an SDE with the diffusion time t from 0 to 1.
intermediate forcefield Ve,, with other
may struggle capture.
the expanded embedding of atom i as follows:
This implies that an effective embedding approach must capture these patterns. Using only the periodic table as vo- cabulary would yield low-resolution representations, miss- ing those domain-speciff, regularities. Instead, we propose the atomic embedding expansion technique, extending el- ements of the periodic table to multiple discrete patterns that serve the expanded vocabulary. Given the molecular graph, the model automatically maps each atom to the most possible pattern of the element based on its neighbors, thus simplifying the understanding of complex but highly regular structures.
2, liin(Ali)..... w A A. e IR..
(6)
Then Z E RN is the concatenation of the expanded embeddings z; for all atoms i.
Unified Multi-Head Pretraining With the graph topol- ogy and the atomic representation well prepared, pro- pose the technique for unified pretraining multi-domain molecules. First of all, following Feng et al. (2024), all molecular data can be preliminarily classified into equilib- rium and off-equilibrium categories. For the off-equilibrium molecular conformations X with real MD forces F, we use the following objective for disciplinedd pretraining:
Specifically, We first define basic atomic vocabulary A, RAX well expanded atomic vocabulary A, E AxXH based on all possible element types that appear in the datasets, where A represents the number of element types, D represents the number of regular patterns for each element. Next, for any atom i of the constructed molecular graph & and its neighbors N,, we compute the expanded weight vector w; follows:
L. =
V x
Hout (i)
F
with
Houts Vout GVP(H, V),
(8)
where H, V. are obtained by the representation model (4) based on Eq. (1), GVP is the Geometric Vector Perceptron (GVP) layer first introduced by Jing et al. (2021) and further updated by Jiao et al. (2024).
Σ (bf() 3 Ablil E R JEN;
n, =
w. softmax (lin(A.[). n.)) [0,
(5)
For the equilibrium molecular conformations, we adopt the denoising pretraining method. Specifically, first add Gaussian noise the equilibrium conformation:
where represents the element-win multiplication opera- tor, lin denotes the linear layer, rbf denotes the radial basis function, Vili represents the atomic embedding of atom dy denotes the Euclidean distance between atom i and j. The vector w. is further considered as the probability that atom appears in one the D possible regular chemical environments. To allow for propaatoo,,,tion, we calculate
X = C( + oge), e~ (O,
(9)
where C(-) is the centering operator to ensure translational neutrality, da is the hyperparameter control the noise

---

UniSim: A Unified Simulator for Time-Coarsened Dynamics of Biomolecules
scale. Following Eqs. (1) and (8), we obtain out R by using& X' as the input. Then the training objective is given by:
locity v and denoiser 7/ are linked by an SDE:
<(t) 7(t)
ax, b(t, di
1/ (1, xx) dt + (2((() dB;
X
X
6. V X
H out
2
(10)
(15)
where b(t, X) rtt, + (t) X), c(t) is a prede- fined function with regard to 4, and B, denotes the standard Wiener process.
However, different datasets may have calculated MD forces under varying conditions (e.g., solvation), leading to mis- alignment between data distributions. To address the issue, DPA-2 (Zhang et al., 2024) proposes the multi-head pretrain- ing technique, introducing the unified descriptor for atomic representations and the fitting network that is updated exclu- sively with the specific pretraining dataset. Similarly, we use series of GVP heads in Eq. (8) with non-shared weights to process data with force labels calculated in different force- fields. For convenience, we denote the training objective in Eq. (7) calculated by the k-th (1 S k S K) head as (R) where K represents the number of forcefields across the pretraining datasets. The ultimate training objective for unified multi-head pretraining is given by:
Note that different choices of I(.) and 7()) can induce infi- nite SDEs. As a well-studied case, we follow the setting of Yu et al. (2024):
[tt, X x,) = 1X1 + (1- 7(1) - t)o,,
vi(1
(16)
where 0, R. is hyperparameter to control the perturba- tion strength. Substituting them into Eqs. (13) and (14), the training objectives are given by:
L. = E() - (X, xo)12+ (1,,(1, -
(17)
--= L. = 0) + ( F + 0),
(11)
where v and 7)= are implemented as neural networks, the expectation is taken over the diffusion time / following the uniform distribution on [o, 1]. training data pairs (Xo. X1)
where T(-) is the indicator function, l is the identifier for distinguishing different chemical environments.
and the intermediate state X, following Eq. (12).
Here we introduce the implementation details of networks and f/-- First, the pretrained network yp is initialized as graph encoder, which now takes the diffusion time e [0, 1J instead of [ E {0, 1) as input. The invariant and equivariant outputs of yp, H and V, are further fed to two GVP layers with non- non-shared weights as the output heads of and l/2, respectively. Considering different scales of labels, an additional Vector LayerNorm (VLN) is stacked before each GVP layer, defined as as:
3.3. Vector Field Model for Dynamics
In this section, we will present the generative framework of UniSim, named as the vector field model, to learn the push forward from X, to X of MD trajectories with a prede- fined long timestep r. For compatibility the generative framework, will denote the initial state as Xo and the terminal state as X for each training data pair.
Here we leverage the stochastic interpolant (Albergo et al., 2023) as our generative framework. Given two probability density functions go qı Rdd + R20. a stochastic inter- polant between go and q1 is a stochastic process X, defined as:
V std,
VLN(V) 0.
ER
(18)
where a learnable parameter, (.) represents calcu- lating the standard value along dimension i of the tensor.
X = 1(,, X,, + 7()
(12)
where the pair (X,, X2) is drawn from the probability mea- sure // that marginalizes on go and q1. Ž ~ (O, E [0, 1] denotes the diffusion time, T(.): [0, 1] x Rd x Rd -> Rd and 7(-): : [o, 1] + IR are mapping functions that should satisfy certain properties.
3.4. Force Guidance Kernel for Finetuning
In this section, we will introduce the network Ç, named as the force guidance kernel, which provides the mobility to different chemical environments by learning "virtual" force defined on X, and incorporating into generation. In- spired by Wang et al. (2024a); Yu et al. (2024), we hope the marginal distribution Pt generated by ( satisfies:
To obtain training objectives, we first define the velocity and the denoiser 7)= as:
v(4, x) = (1, X X X),
(13)
P-(-) x qe (-) exp(-ae so(-) = =1(-) =()), (19)
n (t, x = E X = X
(14)
where ge represents the marginal distribution generated by the vector field model o, a is the hyperparameter of guid- ance strength, and Ee(-) denotes an intermediate potential
where the expectation is taken over the data pairs (Xo, X)).
can be further proved that the stochastic process X,, ve-
5

---

UniSim: A Unified Simulator for Time-Coarsened Dynamics of Biomolecules
defined X,, which is continuous with the MD potential e at two endpoints. Therefore, any change of MD potentials owing to different chemical environments will reflect the change of generated marginal distributions based on S, serving as the guidance of underlying physics principles.
training objective of the force guidance kernel is given by:
C, Blllvo(, X.) + Foll + || (t, X,) + 11
g Xx (Xo X X+,)
+ +1,bbr,, X.)
11
aElg X Xo=
)g(( X X1)
For simplification, we assume that the stochastic interpolant that generates Pt shares the same form as in Eq. (16), im- plying that pr (X,, Xo. X1) = qr (X, X1). According to Yu et al. (2024), the closed-form of the intermediate forcefield Ve.(.) is given by:
(24)
where Fo, F1 correspond to MD forces of X X respec- tively, and the expectation in the denominator term is taken over all training data pairs of a mini-batch.
3.5. Inference
(X,) (X X (o, X ( X,,
X,) X aBlq: XX, Xo: X (X,, X1)]
During inference, the SDE of Eq. (15) will be discretized into T equidistant steps for generation, where T is a hy- perparameter. Given the initial state as Xo. a new state X, is generated through the 7-step discrete Markov pro- cess, which completes one inference iteration. Subsequently, the newly generated state serves the initial state for the next iteration, by which UniSim is able to autoregressively generate trajectories for any given chain length.
(20)
where g g(X,, X,) = (pp(-o(eXx)) + =(xx))) and
M ((x, X1, x,) = Vlogg((X; - V log 9: [X,|Xa Xo. X)).
According to Albergo et al. (2023), -7 (t) Xt) is unbiased estimation of Vlg (X:); and Vlog qe (X, X1) has a closed-form solution based on Eqs. (12) and (16), thus all terms of Eq. (20) can be calculated during training without any approximation.
Empirically, we find that performing inference without post- post-procsssing leads to unstable conformation generation. Therefore, after each iteration, we add a conformation re- finement step for peptides and proteins. Details and related analysis can be found in C.2 and ş E.2, respectively.
Further, Proposition 3.1 reveals the closed- form of SDE that
generates pe:
Proposition 3.1. Assume that marginals q: and Pt are gen- erated by b(t, X), (t, X) and (t, X), (t, X) based on Eq. (15), respectively. Given the probability measure v of data pairs satisfying v(Xo, xI) = v(X,, Xo). we have the
4. Experiments
following equalities hold:
4.1. Experimental Setup
8'(1,))= = b(t, X),
(21)
Firstly, the pretraining datasets D, (1 S i m) are listed as follows: 1) PCQM4Mx2 (Hu et al., 2021), a quantum chemistry dataset with around 3M small molecules of equilibrium conformations optimized by DFT. 2) ANI- 1x (Smith et al., 2020), small organic molecule dataset consisting of 5M DFT calculations for QM properties like energies, atomic forces, etc. 3) PepMD (Yu et al., 2024), a peptide dataset including peptides of 3-10 residues with the sequence identity threshold of 60%, where we perform MD simulations using OpenMM (Eastman et al., 2017) to gen- erate MD) trajectories of 283 peptides. We adopt the same test set split of 14 peptides in the original paper. 4) Pro- tein monomers processed by Jiao et al. (2024), subset of PDB (Berman et al., 2000) including protein monomer crys- tal structures. 5) ATLAS (Vander Meersche et al., 2024), a protein dataset gathering all-atom MD simulations of protein structures, which is chosen for structural diversity by ECOD domain classification (Schaeffer et al., 2017). Following Cheng al. (2024), We selected proteins from the dataset of 2024.09.21 21 with no more than 500 residues and a coil percentage not exceeding 50%, resulting in total of 834 data entries. We then apply the sequence clustering with the threshold of 30% by MMseq2 (Steinegger & Söding, 2017), obtaining 790/14 4 as the train/test splits. 6) Solvated Protein
Datasets
xt, + (X).
(22)
Finally, we introduce the implementation of the force guid- ance kernel Ç. Keeping the parameters of the representation model 9 and the vector field model = fv. 7=} frozen, we use another TOrChMD-NET & as the graph encoder of Ç. which is initialized with the same hyperparameters of yp. To leverage the unified atomic representation, we adopt the residual mechanism, where the invariant output H of will be added to that of V during both training and inference.
After the invariant and equivariant features are obtained from V, we construct the network &e fit the intermediate forcefield (-r(-) in the same interpolation form as in Yu et al. (2024):
((-) = (1 - t))o(-) + (xp1())++++))))))))) (-),
(23)
where networks db; (i E {0, (,) are implemented as the same architecture of v. and v/o, 2/1 are used to fit MD force- fields at two endpoints = 0, 1, respectively. Therefore, the
6

---

UniSim: A Unified Simulator for Time-Coarsened Dynamics of Biomolecules
Fragments (SPF) (Unke& & Meuwly, 2019), dataset prob- ing many-hody intermolecular interactions between protein fragments and water molecules. We randomly split the above pretraining datasets for training and validation by 4:1.
chain length of 10 The evaluation results on all 14 test peptides of PepMD are shown in Table 1, where UniSim/g refers to the version that performs inference using only the vector field model 6 without force guidance.
Secondly, we will validate the transferability of UniSim across three molecular domains: 1) MD17 (Chmiela et al., 2017) and MD22 (Chmiela et al.. 2023) used for training and evaluation respectively, serving as representatives for small molecules. 2) PepMD as introduced above for peptides and 3) ATLAS for proteins. More details of pretraining datasets, trajectory datasets and MD simulation setups are shown in
Based on Table 1, we first observe that UniSim outperforms the baselines on nearly all metrics, particularly when com- pared to FBM that uses a similar framework, demonstrating the effectiveness of our pretraining techniques. Moreover, by introducing the force guidance, UniSim shows a signif- icant improvement in validity while maintaining the same level of distribution similarity, which reveals deeper com- prehension of the underlying physics constraints.
$ D.
Baselines Considering computational resource limitations, we compare UniSim with baseline models on peptides and with MD trajectories on other molecular domains. Specifi- models
Table 1. Results on the test set of PepMD. Values of each metric are shown in mean/std of all 14 test peptides. The best result for each metric is shown in bold and the second best is underlined.
cally, we select the following deep learning- as our baselines: 1) FBM (Yu et al., 2024), the current state- of-the- model based on bridge matching, learning time- coarsened dynamics on peptides with steerable force guid- ance. 2) Timewarp (Klein et al., 2024), generative model leveraging the augmented normalizing flow and MCMC techniques, exhibiting transfcrability to small peptide sys- tems. 3) ITO (Schreiner et al., 2024), conditional diffusion model tailored for learning dynamics on varying time res- olutions. 4) Score Dynamics (SD) (Hsu al., 2024), score matching diffusion model that captures transitions of collective variables of interest.
MoneS
15 DISTANCE (1)
VAL-CA (f)
CUNTACT (1)
TIC-2D
FIMM
0.3000 134
0:0550 IDS
TIMEWARP
0 0.362/0095
0.386-0122 120
0.1140.110
0.74541061
0,19500051
ITO
0 0.367/0077
0.371 131
0.455 126
0.74881055
0.1600.
0.17481.99
SD
0.727.0.899
0.77600057
0.541/0.113
0,782
0,2680.266
0.6660.166
UniSiii/g
0.232/01135
0 161
0.1100.115
0,735/0.064
0,50500.112
0.163010766
UniSian
0.338/0149
0.336 159
0.516/.124
0.731/0074
0,57500.13
0.157.0.03
For more intuitive understanding, we provide the visualiza- tion of the metrics for the two test cases in Figure 3. UniSim exhibits close alignment with MD trajectories with regard to pairwise distances and residue contact rates. Further- more, UniSim shows good recovery of known metastable states, with samples located mainly in high-density regions. Though fully reproducing the free energy landscape may re- quire longer trajectory lengths, UniSim still demonstrates a basic understanding of the intrinsic Boltzmann distribution.
Metries We employ the same metrics in Yu et al. (2024) to comprehensively evaluate the distributional similarity, validity and flexibility of generated ensembles. Briefly, the metrics are listed below: 1) The proportion of conforma- tions with neither bond break nor bond clash, termed as VAL-CA (Lu et al., 2024). 2) The root mean square er- ror of contact maps between generated ensembles and MD trajectories, termed as CONTACT (Janson et al., 2023). 3) The Jensen-Shannon (IS) distance on projected feature spaces of Pairwise Distances (PWD), Radius-of-Gyration (RG), the slowest two Time-lagged Independent Compo- nents (TIC) Herrazzeer et al., 2013) as well as their joint distribution, termed as TIC-2D, The mean value of JS distances along each dimension is reported.
17702 C
"saoeii We
1ar8 0
4.2. Evaluation on Peptides
Based on the unified atomic representation model ip ob- tained through pretraining, we train on PepMD to derive unified vector field model p The rationale behind the choice is that peptides have moderate scales and exhibit high structural flexibility, making it suitable for transferring to other molecular domains subsequently. For fair compari- son, all baselines are trained from scratch on PepMD until convergence, and sample trajectories of each test peptide for
Figure 3. The visualization of comprchensive metrics peptide 1i7u_C (upper) and 1ar8_0 (lower). The left column shows the joint distribution of pairwise distances. The middle column demon- strates the residue contact map, where data in the lower and upper triangle are obtained from UniSim and MD, respectivee respectively. The right
column displays TIC-2D plots for the slowest two components, where contours indicate the kemee density estimated on MD tra- jectories and the generated conformations are shown in scatter.
7

---

UniSim: A Unified Simulator for Time-Coarsened Dynamics of Biomolecules
4.3. Transferability to Small Molecules
guidance kernel corresponding to the protein domain is trained on ATLAS subsequently.
We further investigate the performance of UniSim on small molecules. Keeping parameters of the atomic representation model yp and the vector field model & frozen, we train MD17 to obtain a force guidance kernel and evaluate on 5 organic molecules of MD22. Here we use the relative dis- tances between all heavy atom pairs as the projected feature to compute TIC and TIC-2D metrics. The ablation results comparing UniSim with UniSim/g are shown in Table 2. Apparently, the force guidance kernel helps the model trans- fer to new chemical environment with better distributional similarity in overall.
Afterwards, the models are evaluated on the test set of AT- LAS and the results are displayed in Table 3. Evidently, UniSim significantly outperforms the baselines across all metries, especially in terms of validity, achieving an im- provement from 5% (ITO) to 8%. Meanwhile, the introduc- tion of foree guidance leads to modest improvements over most metrics well.
Table 3. Results on the test set of ATLAS. Values of each metric are shown mean/std of test protein monomers. The best result for each metric is shown in bold.
J8 DISTANCE
Table 2. Results on the test set of MD22. Values of each metric are shown in mean/std of all 5 test molecules. The best result for cach metric shown in bold.
MODELS
VAL-CA (t) CONTACT ())
PWD
RG
TIC
FBM
0,519/00233
0.597/0.121
0.621/0.152
0.012/0.007
0,252/0.039
ITO
0,588/1027
0.775/0.042
0.624/b1211
0.052/0.008
0.428/0.020
SD
0,604/9.200
0.762/0.060
0.605/0.128
0.001/0.000
0,235/0.033
JS DISTANCE ())
MODELS
UniSim/g
0.508/1.021
0.56.0. 146
0.545/0.11
0.071/b.029
0.171/0.011
TIC
TIC-2D
UniSim
0,506/0021
0,554/0.199
0,542/0.159
0,079/0.03
0.17M0.031
UniSim/g
0.408/0.111
0.791/0.044
UniSim 0.368/0.132
0.765/0.063
Moreover, to intuitively demonstrate the simulation effi- ciency of UniSim on large proteins, we present the TIC-2D visualizations of the first 200 generated conformations for two protein test cases, as shown in Figure 5. Notably, al- though long-term stable simulations for large proteins still proves to be challenging, UniSim is able to cross the energy barrier and explore distinct metastable states with only a few inference iterations, which shows the potential to take the place of traditional MD in terms of efficiency.
To better understand how the force guidance kernel works, we provide visualizations on Ac-Ala3-NHMe and DHA in Figure It shows that the force guidance greatly helps com- prehend the free energy landscape, enabling more accurate transitions between metastable states. Accordingly, samples generated by UniSim are more likely concentrated in high- density regions, which complies with physics constraints of the specific chemical environment.
UniSim
UniSim
TC
TIC
TIC 0
3bn0
4b6i_D
TC3
AC-Ala3-N-IMe
b. DHA
Figure 5. TIC-2D plots of the first 200 generated conformations for a. 3bn0_A and b. 4b6i_D, Contours indicate the kernel density estimated on MD trajectories, the generated conformations are shown in scatter, and the blue dashed arrows represent the order in which the conformations are generated.
Figure 4. TIC and TIC-2D plots of UniSim (left) and UniSim/g (right) on a. Ac-Ala3-NHMe and b. DHA. The first row displays the free energy projection on TIC 1, and the second row demon- strates TIC plots for the slowest two components.
4.5. Case Study: Alanine-Dipeptide
4.4. Exploration of Proteins
To further validate the stability and applicability of UniSim in long-timescale simulations, we conduct additional experi- ments on a well-studied molecular system, alanine- -iippiee (AD), consisting of only 22 atoms while exhibiting compre- hensive free energy landscapes.
In this section, we explore the transfcrability of UniSim to proteins. Considering the complexity and specificity of protein structures, we first finetune the vector field model on ATLAS with a learning rate of le-4. Next, the force
8

---

UniSim: A Unified Simulator for Time-Coarsened Dynamics of Biomolecules
Specifically, following the task setup in Timewarp (Klein et al., 2024), we attempt finetune UniSim trained on PepMD to AD before performing long-timescale simu- lations. Firstly, we obtain three independently sampled MD trajectories of AD with the simulation time of 250 ns from mdshare', which are assigned as the train- ing/validationntesst trajectories. The coarsened timestep T is set to 100 ps. with 200,000 data pairs randomly sampled for training and validation from corresponding trajectories, respectively. UniSim is then finetuned on the curated AD dataset with the learning rate of le-4 for 300 epochs.
root-mean-square deviation (RMSD) to each representative conformation from trajectories generated by UniSim, which are illustrated in Figure 7. It can be observed that UniSim consistently exhibits an excellent recovery of metastable states with negligible deviation.
a. cs (0.05A)
C7cq (0.03A)
c. sr (0.08A)
After we obtain the best checkpoint of UniSim evaluated on the validation set, we perform long-timesen timescale simulations for chain length of 100,000 to explore the metastable states of AD. We show the Ramachandran and TIC-2D plots of UniSim and the test MD trajectory in Figure 6. Building upon previous research (Wang et al., 2014), UniSim has demonstrated robust performance in long-timescale sim- ulations by effectively exploring key metastable states of AD, including C5, C7eq. a'g. d'g. as well as ag. Moreover, the relative weights of generated conformation ensembles across different metastable states show good agreement with MD, indicating that UniSim is basically capable of repro- ducing the free energy landscape.
og (0.03A)
e. dg (0.05A)
Figure 7. Comparison of representative conformations of AD be- tween UniSim (hlue) and MD (yellow), including: a. C5, b. C7ea. at., d. og and oR. The RMSD values of each representative pair over heavy atoms are shown in brackets.
MD
UniSim
5. Conclusion and Future Work
In this work, we present a novel architecture called UniSim, the first deep learning- based model tailored for performing time-coarsened dynamics on biomolecules of diverse do- mains. To accommodate molecular inputs of different sizes and types, we obtain a unified atomic representation model by pretraining on multi-domain 3D molecular datasets with novel techniques. Afterwards, we leverage the stochastic interpolant framework to construct vector field model that learns
PS1
Phi
4
UniSim
dynamics on MD trajectories. Con- sidering the impact of different chemical environments in applications, we further introduce the force guidance kernel, which adjusts the expectation of observables by incorporat- ing "virttal"' forces into the stochastic process. Experiments conducted on small molecules, peptides and proteins have fully verified the superiority of UniSim in distribution sim- ilarity compared to MD trajectories and transferability to -idt-i--iittitttion molecular domains.
TICO
TICO
Figure 6. Visualization of Ramachandran plots (the first row) and TIC-2D plots (the second row) of UniSim and MD alanine- dipeptide.
In addition, we believe this work can be advanced in the following aspects: 1) Influenced by cumulative prediction errors, the validity of the samples generated by UniSim is not fully reliable, especially for macromolecules like proteins. Efficient cross-domain structure optimization de- serves further exploration. 2) The generated trajectories for evaluation in our experiment are relatively short, which may hinder the model from discovering more possible metastable states. The dynamics pattern of biomolecules over longer time scales is worth investigating.
Furthermore, we provide a more intuitive illustration of the accuracy in generating conformations of metastable states. Based on Ramachandran plots, we apply the K-means (Mac- Queen, 1967) algorithm to obtain 5 clusters from MD trajec- tories, and select the centroid of each eluster as the represen- tative conformation of the corresponding metastable state. Subsequently, we identify the conformation with the lowest
http:://makkkooooddeelliiiiioedo Aiaaala dipeptide

---

UniSim: A Unified Simulator for Time-Coarsened Dynamics of Biomolecules
Acknowledgements
Cheng, K., Liu, C., Su, Q.. Wang. J., Zhang, L., Tang, Y., Yao, Y., Zhu, S., and Qi, Y. 4d diffusion for dynamic protein structure prediction with reference guided motion alignment. arXiv preprint arXiv:2408. 12419, 2024.
This work is jointly supported by the National Key R&D Program of China (No.2022ZD0100500) the National Nat- ural Science Foundation of China (No. 61925601, No. 62376276, No. 62276152), and Beijing Nova Program (20230484278).
Chmiela, S., Tkatchenko, A., Sauceda, H. E.,, Poltavsky, I., Schüit, K. T., and Müller, K.-R. Machine learning of ac- curate energy-conserving molecular force fields. Science advances, 3(5)xe1603015, 2017.
Impact Statement
This paper presents work whose goal to advance the field of molecular dynamics through developing a novel deep learning-based generative model, UniSim, which proposes solution to the key limitation in current methods of unified and efficient simulation for cross-domain molecules. Our work has the potential to impact both scientific research and practical applications in real-world scenarios, such as drug discovery and molecular property prediction. We wish our work could inspire future research this field.
Chmiela, S., Sauceda, H. E., Müller, K.-R., and Tkatchenko, A. Towards exact molecular dynamics simulations with machine-leamee foree fields. Nature communications, (1):3887, 2018.
Chmicla, S., Vassilev-Galindo, V., Unke, O. T., Kabylda, A., Sauceda, H. E., Tkatchenko, A., and Müller, K.- R. Accurate global machine learning force fields for molecules with hundreds of atoms. Science Advances, 9 (2)eadd08873, 2023.
References
Du, Y., Plainer, M., Brekelmans, R., Duan, C., Noe, F., Gomes, C. P., Aspuru-Guzik, A., and Neklyudov, K. Doob's lagrangian: A sample-efficient variational ap- proach to transition path sampling. The Thirty-eighth Annual Conference on Neural Information Processing Systems, 2024.
Albergo, M. S., Boffi, N. M., and Vanden Eijnden, E. Stochastic interpolants: A unifying framework for flows and diffusions. arXiv preprint arXiv:2003.08797, 2023.
Allen, M. P. et Introduction to molecular dynamics simulation. Computational soft matter: from synthetic polymers proteins, 23(1):1-28, 2004.
Eastman, P., Swails, J., Chodera, D., McGibbon, R. T., Zhao, Y., Beauchamp, K. A., Wang, L.-P., Simmonett, A. C., Harrigan, M. P., Stern, C. D., et al. Openmm 7: Rapid development of high performance algorithms for molecular dynamics. PLOS computational biology, 13(7): e e1005659, 2017.
Bal, K. M. and Neyts, E. C. Merging metadynamics into hyperdynamics: accelerated molecular simulations reach- ing time scales from microseconds seconds. Journal of chemical theory and computation, 11(10)4555 4554, 2015.
Feng, R., Zhu, Q., Tran, H., Chen, B., Toland, A., Ram- prasad, R., and Zhang, C. May the force be with you: Unified force-centric pre-training for 3d molecular con- formations. Advances in Neural Information Processing Systems, 36, 2024.
Berman, H. M., Westbrook, J., Feng, Z., Gilliland, G., Bhat, T. N., Weissig, H., Shindyalov, N., and Bourne, E. The protein data bank. Nucleic acids research, 28(1): 235- 035-222,22 2000.
Botu, V. and Ramprasad, R. Adaptive machine learning framework accelerate ab initio molecular dynamics. In- ternational journal of quantum chemistry, 115(16):1074- 1083, 2015a.
Fu, X., Xie, T., Rebello, N. J., Olsen, B., and Jaakkola, T. S. Simulate time-integrated coarse-grained molecular dynamics with multi-scale graph networks. Transactions on Machine Learning Research, 2023.
Botu, V. and Ramprasad, R. Learning scheme predict atomic forces and accelerate materials simulations. Phys- ical Review B, 92(9))09433 2015b.
Gao, X., Gao, W., Xiao, W., Wang, 7.,, Wang, C., and Xiang, L. Supervised pretraining for molecular force fields and properties prediction. arXiv preprint 2rXiv:2211.14429, 2022.
Célerse, F., Inizan, T. J., Lagardère, L., Adjoua, O., Mon- marché, P., Miao, Y., Derat, E., and Piquemal, J.-P. An Sfficientgaussinn-acceeee gaussian-accelerated molecular dynamics (gamd) efficient
Griffiths, D. J. and Schroeter, D. Introduction to quantum
multilevel enhanced sampling strategy: Application to polarizable force fields simulations of large biological systems. Journal of Chemical Theory and Computation, 18(2):968997) 2022.
mechanics. Cambridge university press, 2019.
Hollingsworth, A. and Dror, R. O. Molecular dynamics simulation for all. Neuron, 99(6): 129-1 1143, 2018.
10

---

UniSim: A Unified Simulator for Time-Coarsened Dynamics of Biomolecules
Hsu, T., Sadigh, B., Bulatov, V., and Zhou, F. Score dy- namics: Scaling molecular dynamics with picoseconds time steps via conditional diffusion model. Journal of Chemical Theory and Computation, 20(6))2335-2345) 2024.
Langevin, P. Sur la théorie du mouvement brownien. Compl. Rendus, 146:530-533, 1908.
Lau, D., Jian, W., Yu, Z., and Hui, D. Nano-engineering of construction materials using molecular dynamies sim- ulations: Prospects and challenges. Composites Part B: Engineering, 143:282-299 291, 2018.
Hu, W., Fey, M.. Ren, H., Nakata, M., Dong, Y., and Leskovec, J. Ogb-lsc: A large-scale challenge for ma- chine learning on graphs. In Thirty-fifth Conference on Neural Information Processing Systems Datasets and Benchmarks Track (Round 2), 2021.
Li, S., Wang, Y., Li, M., Zhang, J., Shao, B., Zheng, N., and Tang. J. Flow: Frame-to Srame-to-freee to-frame coarse- saarseegrinedd molec- ular dynamics with se (3) guided flow matching. arXiv preprint arXiv:2405.00755 2024.
Janson, G., Valdes- Gariea, G., Heo, L., and Feig, M. Di- rect generation of protein conformational ensembles via machine learning. Nature Communications, 14(1):774, 2023.
Lindorff-Larsen. K., Piana, S., Dror, R. O., and Shaw, D. E. How fast-folding proteins fold. Science, 334(6055):517- 520, 2011.
Jiao, R., Han. J., Huang, W., Rong, Y., and Liu, Y. Energy- motivated equivariant pretraining for 3d molecular graphs. In Proceedings of the AAAI Conference on Artificial In- telligence, volume 37, pp. 8096-8100 2023.
Link, A. J., Mock, M. L., and Tirrell, D. A. Non-canonical amino acids in protein engineering. Current opinion in biotechnology, 14(6):603-60) 2003.
Lu, J., Zhong, B., Zhang, Z., and Tang. J. Str2str: A score- based framework for zero-shot protein conformation sam- pling. In The Twelfth International Conference on Learn- ing Representations, 2024.
Jiao, R., Kong, X., Yu, Z., Huang, W., and Liu, Y. Equivari- ant pretrained transformer for unified geometric learning on multi-domain 3d molecules. In ICLR 2024 Work- shop on Generative and Experimental Perspectives for Biomolecular Design, 2024.
Luo, S., Chen, T., Xu, Y., Zheng, S., Liu, T.-Y., Wang, L., and He, D. One transformer understand both 2d & 3d molecular data. In The Eleventh International Conference on Learning Representations, 2022.
Jing, B., Eismann, S., Soni, P. N., and Dror, R. O. Equivari- ant graph neural networks for 3d macromolecular struc- ture. arXiv preprint arXiv:2106.0388
2021.
M. Sultan, M. and Pande, V. S. tica-metadynamies: ac- celerating metadynamics by using kinetically selected collective variables. Journal of chemical theory and com- putation, 13(6):2440-244 2447, 2017.
Jing, B., Stark, H., Jaakkola, T., and Berger, B. Generative modeling of molecular dynamics trajectories. In The eiirty-eighth Annual Conference on Neural Information Processing Systems, 2024.
MacQueen, J. Some methods for classification and analy- sis of multivariate observations. In Proceedings of the Fifth Berkeley Symposium on Mathematical Statistics and Probability, Volume 1: Statistics, volume 5, pp. 28 281-298. University of California press, 1967.
Klein, L. and Noe, F. Transferable boltzmann generators. In The Thirty-eighth Annual Conference on Neural Infor- mation Processing Systems, 2024.
Klein, L., Foong, A., Fjelde, T., Mlodozeniec, B., Brockschmidt, M., Nowozin, S., Noé, F., and Tomioka, R. Timewarp: Transferable acceleration of molecular dy- namics by learning time-coarsened dynamics. Advances in Neural Information Processing Systems, 36, 2024.
Markwick, R., Pierce, L. C., Goodin, D. B., and McCam- mon, J. A. Adaptive accelerated molecular dynamics (ad-amd) revealing the molecular plasticity of p450cam. The journal of physical chemistry letters, 2(3):158-14) 2011.
Köhler, J., Invernizzi, M., De Haan, P., and Noé, F. Rigid body flows for sampling molecular crystal structures. In International Conference on Machine Learning, pp. 17301-17326. PMLR, 2023.
Noé, F., Olsson, S., Köhler, J., and Wu, H. Boltzmann gener- ators: Sampling equilibrium states of many-body systems with deep learning. Science, (6457))) 147, 2019.
Kong, X., Huang, W., and Liu, Y. Conditional antibody de- sign as 3d equivariant graph translation. In The Eleventh International Conference on Learning Representations, 2023.
Onufriev, A., Bashford, D., and Case, D. A. Exploring pro- tein native states and large-scale conformational changes with a modified generalized born model. Proteins: Struc- ture, Function, and Bioinformatics, 55(2))383-39) 2004.

---

UniSim: A Unified Simulator for Time-Coarsened Dynamics of Biomolecules
Pearlman, D. A., Case, D. A., Caldwell, J. W., Ross, W. S., Cheatham III, T. E.,, DeBolt, S., Ferguson, D., Seibel, G., and Kollman, P. Amber, package of computer programs for applying molecular mechanics, normal mode analysis, molecular dynamics and free energy calculations to simu- late the structural and energetic properties of molecules. Computer Physics Communications, 91(1-3):1-41 1995.
Van Gunsteren, W. F. and Berendsen, H. J. Computer simu- lation of molecular dynamics: methodology. applications, and perspectives in chemistry. Angewandte Chemie Inter- national Edition in English, 29(9):99--1023, 1990.
van Gunsteren, W. F. and Oostenbrink, C. Methods for dlassiall-eehanne molecular simulation in chemistry: Achievements, limitations, perspectives. Journal of Chemical Information and Modeling, (4(16)6221163000 2024.
Pelaez, R. P., Simeon, G., Galvelis, R., Mirarchi, A., East- man, P., Doerr, S., Thölke, P., Markland, T. E., and Fabri- tiis, G. D. Torchmd-net 2.0: Fast neural network poten- tials for molecular simulations, 2024.
Vander Meersche, Y., Cretin, G., Gheeraert, A., Gelly, J.-C., and Galochkina, T. Atlas: protein flexibility description from atomistic molecular dynamics simulations. Nucleic acids research, 52(DI):D384 D392, 2024.
Pérez-Hernández, G., Paul, F., Giorgino, T., De Fabritiis, G., and Noé, F. Identification of slow molecular order parameters for markov model construction. The Journal of chemical physics, (1)(1), 2013.
Vanommeslacghe, K., Hatcher, E., Acharya, C., Kundu, S., Zhong, S., Shim, J., Darian, E., Guvench, O., Lopes, P., Vorobyov, I., et al. Charmm general force field: A force field for drug-like molecules compatible with the charmm all-atom additive biological force fields. Journal of computational chemistry, 31 31(4):671-690, 2010.
Plimpton, S. Computational limits of classical molecular dynamics simulations. Computational Materials Science, 4(4):361-366 1995.
Rossmann, M. G. and Arnold, E. International tables for crystallography volume F: crystallography of biological macromolecules. Wiley Online Library, 2001.
Wang. H., Schütte, C., Ciccotti, G., and Delle Site, L. Ex- ploring the conformational dynamics of alanine dipep- tide in solution subjected to an external electric field: A nonequilibrium molecular dynamics simulation. Journal of Chemical Theory and Computation, 10(4):1376-1386. 2014.
Schaeffer, R. D., Liao, Y., Cheng, H., and Grishin, N. V. Ecod: new developments in the evolutionary classifica- tion of domains. Nucleic acids research, 45(D1):D296. D302, 2017.
Schreiner, M. Winther, O., and Olsson, S. Implicit transfer operator learning: Multiple time-resolution models for molecular dynamics. Advances in Neural Information Processing Systems, 36, 2024.
Wang, J., Arantes, P. R., Bhattarai, A., Hsu, R. V., Pawnikar, S., Huang, Y.-m. M., Palermo, G., and Miao, Y. Gaussian accelerated molecular dynamics: Principles and applica- tions. Wiley Interdisciplinary Reviews: Computational Molecular Science, 11(5)=e1211, 2021.
Shin, K., Tran, D. P., Takemura, K., Kitao, A., Terayama, K., and Tsuda, K. Enhancing biomolecular sampling with reinforcement learning: A tree search molecular dynamics simulation method. ACS omega, 1(9):13853- 13862, 2019.
Wang, I.,, Shen, Y., Wang, Y., Yuan, H., Wu, Y., Gu, Q., et al. Protein conformation generation via force-guided se (3) diffusion models. In Forty-first International Conference on Machine Learning. 2024a.
Smith, S., Zubatyuk, R., Nebgen, B., Lubbers, N., Barros, K., Roitberg, A. E., Isayev, O., and Tretiak, S. The ani- 1cex and ani- 1x data sets, coupled-cluster and density functional theory properties for molecules. Scientific data, 7(1::134, 2020.
Wang, T., He, X., Li, M., Li, Y., Bi, R., Wang, Y., Cheng, C., Shen, X., Meng, J., Zhang, H., et al. Ab initio char- acterization of protein molecular dynamics with ai2bmd. Nature, pp. 1-9, 2024b.
Steinegger, M. and Söding, J. Mmseqs2 enables sensi- tive protein sequence searching for the analysis of mas- sive data sets. Nature hidtebomloog.33(11:0000---33 35011)000261100 2017.
Yu, Z., Huang, W., and Liu, Y. Force-guided bridge match- ing for full-atom time-coarsened dynamics of peptides. arXiv preprint arXiv:2408.1 15126, 2024.
Zaidi, S., Schaarschmidt, M., Martens, J., Kim, H., Teh, Y. W., Sanchez-Gonzan A., Battaglia, P., Pascanu, R., and Godwin, J. Pre-training via denoising for molecu- lar property prediction. In The Eleventh International Conference on Learning Representations, 2023.
Unke, O. Т. and Meuwly, M. Physnet: A neural network for predicting energies, forces, dipole moments, and partial charges. Journal of chemical theory and computation, 15 (6):3678-3693 2019.
12

---

UniSim: A Unified Simulator for Time-Coarsened Dynamics of Biomolecules
Zhang, D., Liu, X., Zhang, X., Zhang, C., Cai, C., Bi, H., Du, Y., Qin, X., Peng, A., Huang, J., al. Dpa-2: large atomic model as multi-task learner. npj Computational Materials, 10(1):293, 2024.
Zhang, Z.,, Xu, M., Jamasb, A. R., Chenthamarakshan, V., Lozano, A., Das, P., and Tang, J. Protein representa- tion learning by geometrie structure pretraining. In The Eleventh International Conference on Learning Repre- sentations, 2023.
Zhou, G., Gao, Z., Ding, Q., Zheng, H., Xu, H., Wei, Z., Zhang, L., and Ke, G. Uni-mol: A universal 3d molecu- lar representation learning framework. In The Eleventh International Conference on Learning Representations, 2023.
13

---

UniSim: A Unified Simulator for Time-Coarsened Dynamics of Biomolecules
A. Reproducibility
Our code is available at httttp/://githuub..m////lleeeeee
B. Proofs of Propositions
Proof of Proposition 3.1. Firstly, denote the marginal score of the vector field model Φ as s(t, X), which is given by:
s(t, x) = E qe
X
X [ = X
(25)
// V log
4: Xo. X (X) dx dx,
(26)
Vq.l
x,) dx, ox
(27)
1
v // xXXXx (x X1) dx, ox,
(28)
q (X)
2 (X) log qt (X),
(29)
qo(X)
where in the second equality we use the Bayesian rule, and the third equality justified by assuming the integrands satisfy the regularity conditions of the Leibniz Rule. q(Xo, X1) denotes the joint distribution of training data pairs.
Further, according to Albergo et al. (2023), we have the following inference:
s(t, x) = - (t,
(30)
Based on the assumption of Eq. (19), we have:
V loggt (X) V Vlogo( (X).
(31)
According to the above derivation, by replacing V log [), V Vo with -? --- (t)r (t)) X), respec- tively, we show that (t, X) = (t, X) + VE, X) holds.
Secondly, we denote the following two terms on the probability path g::
E X
uf(1, x) E
X.
X), ur (1. X)
X].
(32)
Similarly, we define u', (t, X) and u³, (t, x) on the probability path pt in the same form. Based on Eqs. (13), (14) and (16),
we have:
u(t, x) = E(a,1(t, Xt, X, X = x = E| - X
(33)
(t x) E x X 1x, (1 x)
XX X]
(34)
Vi(1- t)o
Combining the connection of i(t,, X) vlt, X) + (t, X), we have:
b(t,X) = E[A, - -xo + (1-2t)o,
X tx, (1 x)xo
X X
(35)
2/1(1-1 11)
VI[1 1)0,
EE (1 2t) 2x)x + :x11 - (1 x)
X, =
(36)
2 2/1(11 1)
(u X) ur u1,11 X)).
(37)
Similarly, (t,,, x) = (ut (t, x) + of X)). According to Yu et al. (2024), it has been proven that u/ X) - (t, ()
(t, X) - (t, (t, x) given the probability measure v of training data pairs satisfying v v(X,, X,) = vXx,, Xo). Therefore, we can derive that (1(1, = b(t, X).
14

---

UniSim: A Unified Simulator for Time-Coarsened Dynamics of Biomolecules
C. Training and Inference Details
C.1. Normalization of Potential and Force Labels
Previous works Wang et al. (2024a); Yu et al. (2024) have mentioned the numerical instability of potential energies across different molecular systems. To ensure the stability of the training process, the potential labels are normalized during pre-processing. Specifically, For each trajectory of one molecular system, we denote the concatenation of unnormalized potential labels of all conformations as E E R.M, where M represents the trajectory length. The normalized potentials given by:
max E)
É
E E1-1,0
(38)
max E) min E)
where max(-) and min(-) denote the maximum and minimum element of the vector. Based on the technique, all potential energy labels are normalized to 0] without manipulating the energy distribution of each ensemble.
For atomic force labels which are comparably more stable, we unify them into MJ/ (mol-nm), whose scale is close to standard Gaussian distributions empirically.
C.2. Conformation Refinement
In practice, we found that small prediction errors the autoregressive generation process would accumulate over time, resulting in unreasonable conformations especially for proteins. Following Wang al. (2024a); Yu al. (2024), after generating a new conformation, we introduce an energy minimization step using OpenMK (Eastman et al., 2017) with harmonic constraints for local structure refinement. Specifically, independent harmonic constraints first added all heavy atoms with spring constant of 10 kcal/mol-A A in case of altering the overall conformation. The minimization tolerance is set to be 2.39 kcal/mol-A A without maximal step limits. Afterwards, the energy minimization step will be applied for peptides or proteins, while small molecules will not undergo any additional treatment.
C.3. Hyperparameters
The hyperparameters of UniSim for model constructing, training and inference are shown in Table 4.
D. Experimental Details
D.1. Dataset Details
Dataset Curation All peptides of PepMD are simulated by OpenMM (Eastman et al., 2017) with the implicit solvation model GB-OBC I (Onufriev al., 2004), where show MD simulation setups in Table 5. Moreover, use the same forcefield to obtain MD potential and force labels of protein conformations in ATLAS. The pretraining dataset details shown in Table 6, and we summarize the trajectory datasets in Table 7.
Test Set Molecules of each test set are listed below, where peptides or protein monomers are named as "(PDB.io. (chain.idi"
MD22: AT-AT-CG-CG, AT-AT, DHA, stachyose, Ac-Ala3-NHMe.
PepMD: 1hhg C, P, 1k83 M, 1bz9.C, C, li7u.C, Igxc B, 1ar8 2xa7 P. 1e28 C, Igy3.F, F, 1n73.I, 1fpr B, laze B, 1qj6_1.
ATLAS: lchd A, 2nnu A, 2bjq-A A, 2ov9.C, CC A, 7mm7_A A, 4adn B, A, 3eye 1h2e A, 1jq5 A, 4b6i Aznr_A, I.uf_A.
E. Additional Experimental Results
E. Ablation Study
Hy perparameeer Sensitivity In order to investigate the hyperparameter sensitivity, we conduct ablation studies of SDE steps T and the guidance strength a on PepMD test set, where the ablation results are shown Table From the table, we
15

---

UniSim: A Unified Simulator for Time-Coarsened Dynamics of Biomolecules
Table 4. Hyperparameters of UniSim.
Hyperparameters
Values
Model
Hidden dimension H
256
FFN dimension
512
RBF dimension
Expand embed dimension D
32
# attention heads
8
# layers of atomic representation model
4
# layers of force guidance kernel
4
Gradient subgraph threshold 8min
8A
Environment subgraph threshold 8max Cutoff threshold Teut Pretraining noise scale de
20A
5
0.04
SDE perturbation strength d
0.2
Training
Learning rate Optimizer
5e-4
Adam
Warm up steps
1,000
Warm up scheduler
LamdaLR
Training scheduler
ReduceL P patience=5, min (r=l-r))))7)
Inference
SDE steps
[15,25,50]
Guidance strength
[0.05,0.06,0.07]
Table 5. MD simulation setups using OpenMM.
Property
Value
Forcefield
AMBER-14 14
Integrator
LangevimMiddlccnneeaaatr
Integration time step
Ifs
Frame spacing
Ips
Friction coefficient
0.5ps
Temperature
300K
Solvation model
GB-OBCI
Electrostatics
CutoffNonPeriodie
Cutoff
2.0nm
Constraints
HBonds
Table 6. Pretraining dataset details.
Domain
Dataset
# Items
Equilibrium
Off-equilibrimm
Traj.
Forcefields
PCQM4Mv2
3M
small molecule
ANI-Ix
55
x DFT
peptide
PepMD
1M
AMBER-14/GB-OBG
PDB
200K
X
x
protein
ATLAS
500K
AMBER-I4/GB-OBCI
SPF
~2
X
DeIPBB-DDXB/)/E222T2VV X
16

---

UniSim: A Unified Simulator for Time-Coarsened Dynamics of Biomolecules
Table 7. Details of the trajectory datasets.
Properties
MD17&22
PepMD
ATLAS
Frame spacing
0.5fs
Ips
10ps
Simulation time per traj
not fixed
100ns
10ns
Coarsened time
100ps
10ps
100ps
Training trajs
8
269
790
# Training pairs per traj
5,000
5,000
500
# Valid pairs per traj
500
500
100
Test trajs
5
14
14
can summarize two key observations:
1. Within the tested range, the validity metric improves significantly as increases, while other metrics generally exhibit deteriorating trends. This suggests that the force guidance kernel enhances the comprehension of physical priors, thereby facilitating higher-quality sample generation, but may simultaneously constrain the exploration breadth of the state space to some extent.
2. As T' increases, most metrics show varying degrees degradation. This is likely because excessive discretization of the SDE process leads greater error accumulation, indicating that a small T sufficient for balancing accuracy and efficiency within the bridge matching framework.
Table Ablation results SDE steps T and the guidance strength a on PepMD test set. Values of cach metric are first averaged over independent runs for each peptide and then shown in mean/std of all 14 test peptides.
JS DISTANCE (L)
Hyperparameters
VAL-CA (1)
CONTACT (4)
PWD
RG
TIC
TIC-2D
T = 15, a 0.05
0.328/0.149
0.330/0.189
0.510/0.124
0.731/0.074
0.575/0.139
0.157/0.088
T = 15, 0.06
0.340/0. 143
0.372/0. 187
0.51
0.740/0.059
0.594/0.100
0.167/0.099
T 15, 0.07
0.349/0.144
0.384/0.206
0.523/0.132
0.736/0.074
0.607/0.138
0.195/0.091
T 25, 0.05
0.391/0.14]
0,474/0.
0.507/0.134
0.738/0.066
0.441/0.117
0.231/0.087
T 25, 0.06
0.404/0.171
0.445/0.222
0.505/0.142
0.734/0.078
0.468/0.136
0.244/0.110
T 25, 0.07
0.409/0.159
0.488/0.218
0.516/0.129
0.742/0.078
0.496/0.147
0.239/0.114
Atomic Embedding Expansion To investigate the contribution of atomic embedding expansion to pretraining, we conduct additional ablation experiments. Specifically, we remove the expanded embedding A, from the model architecture while keeping all other parameters the same as in Table 4. The model is then pretrained and finetuned on the PepMD dataset following the same procedure as described the main text, and the ablation results on PepMD test set are shown in Table 9, evident that removing Ag leads to a decline in nearly all evaluation metrics to varying degrees, demonstrating the effectiveness and necessity of our atomic embedding expansion technique cross-domain scenarios.
Table Ablation Results atomic embedding expansion on PepMD test set. Values each metric are shown in mean/std all test peptides. The best result for each metric is shown in bold.
JS DISTANCE (L)
MODELS
VAL-CA (1)
CONTACT (U)
PWD
RG
TIC
TIC-2D
UniSim/g
0.332/0. 135 0.389/0.175
0.332/0.161
0.510/0. 0.516/0.135
0.738/0.064
0.505/0.112
0.162/0.076
UniSim/g w/o A
0.453/0.233
0.732/0.053
0.397/0.132
0.228/0.119
17

---

UniSim: A Unified Simulator for Time-Coarsened Dynamics of Biomolecules
E.2. Inference Efficiency
In this section, we compare the inference efficiency of UniSim with that of MD performed using OpenMM. Following Timewarp (Klein et al., 2024), use the sfffetive-sammpl-siize per second of wall-clock time (ESS/s) as the evaluation metric. Figure illustrates the statistical results of ESS/s between UniSim and MD PepMD test set, demonstrating that UniSim achieves, average, approximately 25 times higher efficiency compared to MD.
Meanwhile, since conformation refinement using OpenMM is performed at each inference iteration (i.e, generating a new state) for peptide or protein generation, we provide the statistics of the number optimization steps required per iteration as follows: (1) mean: 69.3, (2) median: 55, and (3) maximum: 2,075, as illustrated in Figure 8. Additionally, for each iteration, the average inference time 0.120 s and the average optimization time is 0.152 Therefore, the computational overhead remains within the same order of magnitude with the refinement step.
60
200
175
50
Opt. steps / Iter
150
40
125
ESS/s
30
100
75
20
50
10
25
UniSim
UniSim Method
Method
Figure Statistical results for evaluating model's efficiency. Left. The sareiive-ssmll--iiize per second of wall-clock time (ESS/s) on PepMD test set. For the convenience of comparison, the values were converted into ratios relative to the median of results from MD (the blue dashed line). Right. The optimization steps performed by OpenMM per iteration on PepMD test set.
F. Computing Infrastructure
UniSim was trained on NVIDIA GeForce RTX 3090 GPUs within week. The inference procedures were performed on one NVIDIA GeForce RTX 3090 GPU,
18


Reviews:
[{"review_id": "9tNMsRv0a9", "reviewer": "Reviewer_mRvK", "summary": "The paper presents UniSim, a unified simulator for time‐coarsened dynamics of biomolecules. It proposes a multi-task pretraining method to learn a unified atomic representation across small molecules, peptides, and proteins. The simulator uses a stochastic interpolant framework to learn long-timestep state transitions and introduces a force guidance module to adapt the generated trajectories to different chemical environments. Experiments are reported on peptides, small molecules, and proteins, with comparisons against baselines such as FBM, Timewarp, ITO, and Score Dynamics.", "questions": "", "limitations": "", "rating": 3, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "- **Q1**: The validity metrics for protein simulations remain unsatisfactory.\n\nWe claim that the unsatisfactory validity metric of ATLAS stems from two primary factors: \n\n1. We pursue a unified modeling framework across small molecules and proteins, which necessitates certain compromises in protein-specific modeling (e.g., MSA).\n\n2. The number of parameters matters: we conduct the ablation analysis by increasing the hidden dimension $H$ from 128 to 256. As demonstrated in [Table 1](https://anonymous.4open.science/r/uni-sim-3BBB/rebuttal/ablation_hidden_dim_pepmd.png) and [Table 2](https://anonymous.4open.science/r/uni-sim-3BBB/rebuttal/ablation_hidden_dim_atlas.png), the modification yields substantial improvements in validity metrics across both the PepMD and ATLAS test sets, suggesting that parameter scaling can effectively address certain performance deficiencies.\n\n- **Q2:** Some evaluation criteria are not thoroughly explored.\n\nWe sincerely appreciate your valuable suggestions regarding the evaluation criteria. We address your concerns through the following systematic investigations:\n\n1. **Hyperparameter Sensitivity**: we have now conducted the ablation study for the inference step parameter $T$ and the guidance strength $\\alpha$ on the PepMD test set, which are shown in [Table 3](https://anonymous.4open.science/r/uni-sim-3BBB/rebuttal/ablation_T_alpha_pepmd.png). From the table, we can summarize two key observations:\n\n   - The validity metric improves as $\\alpha$ increases, while other metrics generally exhibit deteriorating trends. This suggests that the force guidance kernel enhances the comprehension of physical priors, but may constrain the exploration of the state space to some extent.\n\n   - As $T$ increases, most metrics show degrading trends. This is likely because excessive discretization of SDE leads to greater error accumulation.\n\n2. **Error accumulation over long rollouts**: we have performed long-timescale simulations (100,000 rollouts) on a well-studied system, alanine-dipeptide, where the detailed setup is elaborated in our response to **Q1,** **Reviewer KWXR**.  The stability of UniSim in long-timescale simulations has been verified by the alignment of RMSD relative to the initial state ([Figure 2](https://anonymous.4open.science/r/uni-sim-3BBB/rebuttal/ad_rmsd_over_time.png)).\n\n- **Q3:** The lack of rigorous statistical analysis.\n\nWe have now updated the statistical analysis for protein validity metrics ([Table 4](https://anonymous.4open.science/r/uni-sim-3BBB/rebuttal/statistical_analysis.png)), which provides rigorous statistical validation of our performance claims.\n\n- **Q4**: a brief discussion of methods that integrate reinforcement learning/enhanced sampling/adaptive sampling.\n\nIn response to your valuable feedback, we will further enhance the Introduction section with a brief discussion on those methods.\n\n- **Q5**: Regarding the choice of the coarse timestep $\\tau$.\n\nThanks for the instructful question. We provide the following responses:\n\n1. First, we did not conduct ablation experiments on the coarsened timestep $\\tau$. That's because selecting different $\\tau$ values would require retraining the model on entirely new datasets, which is not worthwhile.\n2. Second, we explain our methodology for selecting $\\tau$ across different datasets. According to Timewarp [1] where $\\tau$ was set to 500 ps, we opted for values of $\\tau$ in the same order of magnitude. We then select an appropriate $\\tau$ to ensure the variance of the distance between training data pairs was close to 1, thereby avoiding numerical instability during training.\n3. Finally, we briefly discuss the impact of different $\\tau$ on the performance.\n   1. If $\\tau$ is extremely small, the conformations of $(x_0,x_1)$ would be very similar. While this may intuitively improve model accuracy, the sampling efficiency may become the bottleneck instead.\n   2. If $\\tau$ falls within a reasonable range, increasing $\\tau$ reduces the time correlation between $(x_0,x_1)$, making the learning task more challenging but probably enhances the exploration of the state space.\n   3. If $\\tau$ exceeds a certain threshold, $x_0$ and $x_1$ can be considered as independent samples from the Boltzmann distribution, rendering the model incapable of learning dynamic transition features between states.\n\nIn conclusion, we believe that a moderate $\\tau$ is a reasonable choice.\n\n- **Q6**: Regarding the training Data Balance.\n\nWe employed several engineering tricks to ensure the pretraining dataset remains as balanced as possible:\n\n- During training, we require molecules of each batch originate from the same dataset.\n- We ensure that the total number of atoms per molecule in a batch is similar using dynamic batch size.\n- We impose the same upper limit on the number of batches for each dataset in a single epoch.\n\nThis approach ensures that the total number of atoms contributed by each dataset during a single epoch remains approximately balanced.", "claims_and_evidence": "The paper claims that UniSim is the first unified model for time-coarsened dynamics across diverse biomolecular domains and that its multi-task pretraining and force guidance improve distributional similarity and validity. While the theoretical formulations are supported by derivations, the experimental evidence is less convincing. In particular, the validity metrics for protein simulations remain unsatisfactory, and the improvements over existing methods are modest. These factors weaken the claims regarding robust performance and cross-domain transfer.", "methods_and_evaluation_criteria": "The proposed methods include:\n- A unified atomic representation using an SO(3)-equivariant GNN.\n- A stochastic interpolant-based generative framework to bridge long simulation timesteps.\n- A force guidance module to enforce Boltzmann-like behavior during sampling.\n\nThe evaluation criteria (JS distances, VAL-CA, CONTACT, etc.) are standard and appropriate for MD simulations. However, some aspects (e.g., sensitivity to hyperparameters and error accumulation over long rollouts) are not thoroughly explored.", "theoretical_claims": "The paper presents a proof (Proposition 3.1) linking the stochastic interpolant framework with the modified dynamics when force guidance is applied. A cursory check suggests the derivations are consistent. No major errors were detected in the presented theoretical claims, though a detailed line-by-line verification is challenging within the review context.", "experimental_designs_or_analyses": "The experiments span multiple molecular domains with comparisons against several baselines. The design is standard, and the benchmarks are appropriate. However, the results—particularly for proteins—do not show significant improvements in key validity metrics. The lack of rigorous statistical analysis (e.g., confidence intervals) and limited ablation studies on the effect of force guidance and timestep choices reduce confidence in the claimed advantages.", "relation_to_broader_scientific_literature": "The work is positioned well within the recent literature on deep learning for MD simulations. It builds on approaches such as FBM, Timewarp, ITO, and Score Dynamics, aiming to integrate aspects from each—namely, unified representations, time-coarsened simulation, and force-guided sampling. The paper accurately cites recent advances and situates its contributions relative to them.", "essential_references_not_discussed": "While the paper covers most key works, it would benefit from a brief discussion of methods that integrate reinforcement learning or adaptive sampling techniques for MD acceleration, as well as approaches from enhanced sampling (e.g., metadynamics) that address similar problems in different ways.", "other_strengths_and_weaknesses": "N/A", "other_comments_or_suggestions": "N/A", "questions_for_authors": "1. Statistical Significance: Can the authors provide statistical analysis (e.g., confidence intervals or p-values) for the reported improvements, especially in the protein validity metrics?\n2. Hyperparameter Sensitivity: How does the choice of the coarse timestep ($\\tau$) affect the simulation accuracy and validity? Have experiments been conducted to assess sensitivity to this parameter?\n3. Training Data Balance: Was the unified model trained on a balanced dataset across small molecules, peptides, and proteins? If not, could the imbalance affect the reported transferability claims?", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "84QQzdf2OG", "reviewer": "Reviewer_Cnw4", "summary": "The paper introduces UniSim, a deep learning-based unified simulator for time-coarsened MD simulation. The framework aims to improve the transferability and efficiency of long-timescale molecular simulations across different biomolecular domains (small molecules, peptides, and proteins), and consistes of a pretraining module with unified atomic representation, a state transition module using stochastic interpolant , and optional force guidance kernel. Evaluations on small molecules, peptides, and proteins show competitive results w.r.t. several baselines.", "questions": "", "limitations": "", "rating": 3, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "- **Q1:** Writing clarity in the methods needs improvement. For example, the gradient-environment subgraph part was a bit confusing and needs some more intuition on the design.\n\nThanks for the question. We will elaborate on the design rationale of the gradient-environment subgraph as clearly as possible.\n\n**Scale Invariance:** A crucial challenge in training unified representation models arises from the vast scale discrepancy between molecular systems: small molecules typically contain ~10^1 atoms, while proteins often comprise ~10^3 or more atoms. Conducting direct full-atom representation learning without proper processing could lead to non-transferrable representations across different molecular domains. \n\n**Physical Faithfulness:** We notice that, from Newtonian mechanics, atomic motion is governed by the forces acting on each atom. Crucially, long-range intermolecular forces (e.g., van der Waals interactions) decay exponentially with distance, becoming negligible beyond ~10 Å empirically. This implies that the force acting on any atom predominantly originates from its local environment within a 10 Å radius sphere. This physical insight motivates our design of the environment subgraph paradigm, where atomic force computation can be localized to such spherical subgraphs. This approach makes it feasible to decompose large biomolecules into manageable subgraphs for training.\n\n**Computational Efficiency:** Furthermore, during training, we hope that the number of atoms that contribute to loss computation should be the same order of magnitude as that of small molecules or peptides (10^1~10^2 atoms). We denote those atoms as $G_g$. From the above, the complete force-determining environment for $G_g$ is the union of the spherical subgraphs centered on each atom in $G_g$ with radius of around 10 Å, denoted as $G_e$. By using $G_e$ as the training input, we effectively constrain the number of atoms involved in gradient computation to $|G_g|$, which aligns with our intention. Specifically, for the sake of convenience, we adopt Eqs. (2-3) for constructing training data, thus avoiding computing the union of hundrends of point sets.\n\nAccordingly, our hyperparameters should satisfy:\n\n- Scale Matching: $G_g$ contains $10^1$-$10^2$ atoms (aligning with small molecule sizes)\n- Physical Consistency: $\\delta$\\_max - $\\delta$\\_min ≥ 10 Å ensures force computation completeness\n\nMoreover, regarding the impact of the selection of different $\\delta$\\_min and $\\delta$\\_max, we argue that the choice of $\\delta$\\_min can be arbitrary as long as the spherical region with radius delta_min contains an approximately appropriate number of atoms. In contrast, the selection of $\\delta$\\_max must be grounded in physical priors. Crucially, if $\\delta$\\_max - $\\delta$\\_min is so small that it shields some strong interactions, it will inevitably degrade model performance. On the contrary, once $\\delta$_max - $\\delta$\\_min exceeds a certain threshold (e.g., 10 Å), further increasing $\\delta$\\_max yields negligible benefits.\n\n- **Q2:** Effect of predefined timestep τ on model accuracy is unexplored. Interesting to discuss how τ would affect model performance. No major experiment is required, but would like at least some discussion on the intuition.\n\nThanks for the question. Regarding the impact of $\\tau$ on the model performance, we provide the following explanations:\n\n- If $\\tau$ is extremely small (comparable to the MD integral timestep), the conformations of $(x_0,x_1)$ would be very similar. While this may intuitively improve model accuracy, the efficiency could become comparable to or even worse than classical MD, contradicting our goal of accelerating MD simulations.\n- If $\\tau$ falls within a reasonable range, increasing $\\tau$ reduces the time correlation between $(x_0,x_1)$, making the learning task more challenging but enabling the model to explore more state space with a shorter simulation.\n- If $\\tau$ exceeds a certain threshold, $x_0$ and $x_1$ can be considered as independent samples from the Boltzmann distribution, rendering the model incapable of learning dynamic transition features between states.\n\nIn conclusion, we believe that a moderate $\\tau$ is a reasonable choice.\n\n**Reference**\n\n> [1] Kong, X., Huang, W., & Liu, Y. Generalist Equivariant Transformer Towards 3D Molecular Interaction Learning. In *Forty-first International Conference on Machine Learning*.", "claims_and_evidence": "Mostly yes, the novelty  of combining atomistic pretraining with stochastic-interpolant-based dynamics modeling is backed by previous literature. The effectiveness of the force guidance kernel is supported by ablation, and cross-domain transferability is shown on datasets such as MD22.", "methods_and_evaluation_criteria": "Yes, see above.", "theoretical_claims": "Proofs are included in the supplement which I didn't check.", "experimental_designs_or_analyses": "The designs are largely consistent with the claims, see above.", "relation_to_broader_scientific_literature": "The work enables a broader, cross-domain application of time-coarsen MD simulation and also promotes the combination of pretraining with dynamic modelling.", "essential_references_not_discussed": "Not that I found.", "other_strengths_and_weaknesses": "Strengths\n1. The pretraining + stochastic interpolant framework is novel.\n2. Force guidance kernel further improves flexibility in different chemical environments.\n3. Results showing good cross-domain transferability.\n\n Weaknesses:\n1. Writing clarity in the methods needs improvement. For example, the gradient-environment subgraph part was a bit confusing and needs some more intuition on the design.\n2. Effect of predefined timestep τ on model accuracy is unexplored. Interesting to discuss how τ would affect model performance. No major experiment is required, but would like at least some discussion on the intuition.", "other_comments_or_suggestions": "See above.", "questions_for_authors": "See above.", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "3tV2u0NJE0", "reviewer": "Reviewer_KWXR", "summary": "The paper introduces UniSim, a deep learning framework designed to simulate biomolecular dynamics over coarse-grained time steps. The method unifies the treatment of small molecules, peptides, and proteins by first learning a unified atomic representation through multi-task pretraining on a diverse set of 3D molecular datasets. Building on this representation, the authors employ a stochastic interpolant framework to model the transition probability of molecular states over long timesteps, thus enabling time-coarsened dynamics simulation. A force guidance kernel is further introduced to adapt the model to various chemical environments and force field conditions. The approach is evaluated on several datasets (e.g., PepMD, MD17/MD22, ATLAS) and compared against SOTA methods, with experiments demonstrating competitive performance in capturing both the structural and distributional properties of MD trajectories.", "questions": "", "limitations": "", "rating": 4, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "- **Q1**: Potential Issues: Although the experiments are comprehensive, the generated trajectories are relatively short, which may limit insights into long-term dynamics. So adding experiments about the exploration of with conformational space, such as the experiments in F3low, can further illustrate the point.\n\nThank you for your instructful suggestion! To further validate the stability and applicability of UniSim in long-timescale simulations, we conducted additional experiments on a well-studied molecular system, alanine dipeptide (AD), consisting of only 22 atoms while exhibiting comprehensive free energy landscapes.\n\nSpecifically, following the task setup in Timewarp [1], we attempt to finetune our vector field model trained on PepMD (i.e., UniSim/g) to AD before performing long-timescale simulations. Firstly, we obtained three **independently** sampled MD trajectories of alanine dipeptide (AD) with the simulation time of 250 ns from [mdshare](https://markovmodel.github.io/mdshare/ALA2/#alanine-dipeptide), which were assigned as the training/validation/test set. The coarsened timestep $\\tau$ was set to 100 ps, and 200,000 data pairs were sampled from the training and validation set, respectively. Then UniSim/g was finetuned on AD with the learning rate of 1e-4 for at most 300 epochs, after which an force guidance kernel was trained based on the finetuned model with the learning rate of 5e-4.\n\nAfter we obtain the best checkpoint of UniSim evaluated on the validation set, we perform long-timescale simulations for a chain length of 100,000 to explore the metastable states of AD. We show the Ramachandran and TIC-2D plots of UniSim and the test MD trajectory in [Figure 1](https://anonymous.4open.science/r/uni-sim-3BBB/rebuttal/ad_ramachandran.png). Building upon previous research [2], UniSim has demonstrated robust performance in long-timescale simulations by effectively exploring key metastable states of AD, including C$5$, C$7_{eq}$, $\\alpha_{R}'$, $\\alpha_{R}$ as well as $\\alpha_{L}$. Moreover, the relative weights of generated conformation ensembles across different metastable states show good agreement with MD, confirming that UniSim accurately reproduces AD's free energy landscape in long-timescale simulations.\n\nFurthermore, to investigate the stability of UniSim in long-timescale simulations, we selected the root-mean-square deviation (RMSD) of heavy atoms relative to the initial state as the validation metric. The plots of RMSD over time compared to the MD trajectory are presented in [Figure 2](https://anonymous.4open.science/r/uni-sim-3BBB/rebuttal/ad_rmsd_over_time.png). As evident from the figure, the RMSD variations and ranges of the trajectory generated by UniSim exhibit strong consistency with that of MD, demonstrating the model's stability in long-timescale simulations.\n\n**Reference**\n> [1] Klein, L., Foong, A., Fjelde, T., Mlodozeniec, B., Brockschmidt, M., Nowozin, S., ... & Tomioka, R. (2023). Timewarp: Transferable acceleration of molecular dynamics by learning time-coarsened dynamics. Advances in Neural Information Processing Systems, 36, 52863-52883.\n\n> [2] Wang, H., Schütte, C., Ciccotti, G., & Delle Site, L. (2014). Exploring the conformational dynamics of alanine dipeptide in solution subjected to an external electric field: A nonequilibrium molecular dynamics simulation. Journal of Chemical Theory and Computation, 10(4), 1376-1386.", "claims_and_evidence": "- Main Claims:\nThe authors claim that UniSim is the first simulator capable of transferably modeling time-coarsened dynamics across diverse biomolecular domains. In addition, they assert that the multi-task pretraining yields a robust unified atomic representation, and that the incorporation of a force guidance module significantly enhances sample validity and adherence to underlying physical constraints.\n\n- Evidence:\nThe claims are supported by extensive experimental results. Quantitative metrics such as Jensen–Shannon (JS) distances are reported on multiple datasets. Comparative analyses and ablation studies (e.g., UniSim with and without force guidance) further corroborate the improvements brought by the proposed components. While the evidence is comprehensive, the paper also acknowledges limitations (e.g., relatively short generated trajectories and suboptimal validity in protein simulations), which helps contextualize the claims.", "methods_and_evaluation_criteria": "- Methodology: The approach comprises three main stages:\n1. Unified Pretraining: A multi-task pretraining strategy is used to learn an atomic representation from diverse 3D molecular datasets. An SO(3)-equivariant GNN is employed along with novel techniques like gradient-environment subgraph construction and atomic embedding expansion.\n2 .Vector Field Model: A stochastic interpolant framework is used to fit the probability transition kernel of MD trajectories over a coarse time step. This involves training neural networks to predict both a drift term and a noise term, thereby defining an SDE that governs state evolution.\n3. Force Guidance Module: To adapt to different chemical environments, a force guidance kernel is introduced. This module adjusts the learned transition dynamics using “virtual” forces derived from MD potential labels.\n- Evaluation Criteria:\nThe method is evaluated using metrics that assess both the physical plausibility (e.g., VAL-CA, contact map errors) and the distributional similarity (JS distances on various projections) of generated trajectories relative to ground-truth MD data. These criteria are well chosen for the problem of capturing the complex free energy landscapes and metastable states in biomolecular dynamics.", "theoretical_claims": "The paper states theoretical propositions (e.g., Proposition 3.1) that connect the learned vector field model with the force guidance adjustments. The proofs—provided in the appendix—are built on established frameworks from stochastic interpolants and score matching techniques. They derive closed-form relations between the drift and noise components of the SDE under both the baseline and force-guided settings. Although this is not novel, giving proof enhances the completeness of the paper.", "experimental_designs_or_analyses": "- Design:\nThe experimental design is extensive and multifaceted. The authors conduct evaluations across three molecular domains (small molecules, peptides, proteins) using multiple datasets. They perform ablation studies comparing the full UniSim model with a variant that omits the force guidance module (UniSim/g) to isolate its effect.\n\n- Analyses:\nQuantitative comparisons against competitive baseline methods (FBM, Timewarp, ITO, Score Dynamics) are provided via tables and figures that report JS distances, validity scores, and contact errors. Visualizations (e.g., TIC plots, contact maps) illustrate the capability of UniSim to capture the free energy landscape and the transitions between metastable states.\n\n- Potential Issues:\nAlthough the experiments are comprehensive, the generated trajectories are relatively short, which may limit insights into long-term dynamics. So adding experiments about the exploration of with conformational space, such as the experiments in F$^3$low, can further illustrate the point.", "relation_to_broader_scientific_literature": "This paper combines pre-training methods, stochastic interpolation methods in generative modeling, and the classifer-guided approach that is popular in generative models.", "essential_references_not_discussed": "This paper is a good summary of this field.", "other_strengths_and_weaknesses": "I think the concepts presented in this paper are not novel, but it is a very solid piece of work that blends the appropriate concepts more seamlessly and does to SOTA on the various task metrics, which is very helpful to the research community", "other_comments_or_suggestions": "No other comments or suggestions.", "questions_for_authors": "I don't have a specific question.", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "ljmFJTImH4", "reviewer": "Reviewer_mTdA", "summary": "This paper introduces UniSim, a model for simulating time-coarsened molecular dynamics of small molecules, peptides, and proteins. UniSim employs a multi-stage approach: First, a unified atomic representation model is pretrained on a diverse collection of molecular datasets (both equilibrium and off-equilibrium conformations) using a multi-task approach. Then, a vector field model is trained using the stochastic interpolants framework to predict the probability transition kernel for large timesteps (time-coarsening). Finally, a force guidance kernel is trained with other model parameters frozen to adapt (fine-tune) the model to specific chemical environments and force fields. The authors evaluate UniSim on benchmark datasets for each of the three molecular domains and compare to existing methods for time-coarsened molecular dynamics.\n\n## update after rebuttal\n\nPlease see comments below.", "questions": "", "limitations": "", "rating": 2, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "- **Q1:** The first claim & its transferability\n\nFirstly, we strongly agree that the term \"reusable backbone\" is more suitable to describe UniSim, which will be used in our revised manuscript.\n\nSecondly, we argue that the force kernel is not strictly necessary for cross-domain generalization. Figure 4 shows that UniSim/g reproduces TIC-2D distributions for MD22 small molecules without any fine-tuning procedure. This zero-shot transfer capability confirms the learned physical representations transcend specific training domains.\n\nLastly, we are **unable** to apply identical experimental setups to baselines for validating transferability probably due to:\n\n1. Architectural Constraints: Methods like FBM incorporate protein-specific atomic representations (e.g., residue index) that are fundamentally incompatible with small molecules.\n2. Pretraining Deficiency: None of them provides pretraining interfaces, preventing knowledge transfer to novel atomic species after training solely on peptide/protein data.\n\n- **Q2:** The term \"multi-task pretraining\"\n\nWe would like to clarify that \"multi-task pretraining\" (line 220) does **not** refer to different loss formulations here. In particular, different \"tasks\" is defined as different force fields where the force labels come from (line 220-229). By using such multi-task pretraining, we address label inconsistency issues owing to varying force field parameters across different dataset. We hope the misunderstanding can be resolved.\n\n- **Q3:** The practical usefulness in longer timescales\n\nTo specifically address your concerns and validate UniSim’s practical usefulness, we have performed long-timescale simulations on a well-studied system, alanine-dipeptide. Details of the experimental setup can be found in our response to **Q1,** **Reviewer KWXR**, where we show UniSim successfully reproduces the free energy landscape ([Figure 1](https://anonymous.4open.science/r/uni-sim-3BBB/rebuttal/ad_ramachandran.png)), with its stability been verified based on RMSD to the initial state ([Figure 2](https://anonymous.4open.science/r/uni-sim-3BBB/rebuttal/ad_rmsd_over_time.png)).\n\n- **Q4:** The energy minimization technique & efficiency\n\nFirstly, we emphasize that the energy minimization is **almost only** applied to hydrogen atoms with heavy atoms constrained by predefined harmonic potentials (line 682). For this reason, we placed the details in the appendix instead of the main text.\n\nRegarding the computational efficiency, we provide [Table 5](https://anonymous.4open.science/r/uni-sim-3BBB/rebuttal/ess_pepmd.png), which demonstrates that UniSim maintains a 1–2 order-of-magnitude advantage in ESS/s (used in Timewarp) over MD, with refinement step included.\n\n- **Q5:** The energy normalization trick\n\nFirstly, we claim that UniSim focuses on capturing relative energy differences rather than absolute values, as the weight of the Boltzmann constraint can be further modulated through the guidance strength $\\alpha$. Secondly, we confirm that all trajectories in PepMD and ATLAS datasets were sampled at 300K, while MD17 data was collected at around 500K. Thus the temperature consistency within each dataset is guaranteed.\n\n- **Q6:** The visualization of TIC plots\n\nGiven that UniSim's trajectories are relatively short (1,000 frames), contour plots failed to demonstrate clear distributions and yielded suboptimal visualization. By the way, Ramachandran plots in the same format as MD are shown in ([Figure 1](https://anonymous.4open.science/r/uni-sim-3BBB/rebuttal/ad_ramachandran.png)) for long-timescale simulations.\n\n- **Q7:** Hyperparameter sensitivity\n\nIn this experiment, the primary hyperparameters we tuned were: (1) the SDE inference steps $T$ and (2) the guidance strength $\\alpha$. Due to space limitations, please refer to **Q2, reviewer mRvK** for our ablation studies and further discussions.\n\n- **Q8:** Intuition of the atomic embedding expansion\n\nHere we will elaborate on the rationale behind our design of atomic embedding expansion.\n\nSince we aim for cross-domain representation learning, using fine-grained vocabulary specific to proteins would conflict with small-molecule representation. While using the periodic table is the simplest alternative, it will result in significant information loss on **regularly occurring patterns**. Therefore, we adopt the periodic table as the base vocabulary $A_b$ and extend each element with a supplementary vector of length $D$, thereby modeling recurring chemical environments while maintaining a unified representation.\n\nSpecifically, we validate the effectiveness of atomic embedding expansion with the ablation on PepMD shown in [**Table 6**](https://anonymous.4open.science/r/uni-sim-3BBB/rebuttal/ablation_Ae_pepmd.png).\n\n- **Q9:** Intuition of using the \"gradient-environment subgraph\" approach\n\nThank you for your valuable question. Due to space limitations, We kindly direct the reviewer to our response to **Q1, Reviewer Cnw4** for a detailed discussion.", "claims_and_evidence": "The main claims are:\n\n1. UniSim is the first time-coarsened dynamics model exhibiting transferability across diverse molecular domains (small molecules, peptides, and proteins)\n\n2. A multi-task pretraining approach learns a unified atomic representation\n\n3. The stochastic interpolant framework and force guidance allow for rapid fine-tuning to different chemical environments, learning the state transition patterns over long time steps\n\nRegarding the first claim, even though the authors write that their model is transferable, in practice, they train separate force guidance kernels for specific tasks. In my understanding, a truly transferable model should work \"out-of-the-box\" without any need for fine-tuning or output heads trained from scratch. It would thus be more accurate to say that they train a \"reusable backbone\", or \"model that can be efficiently finetuned\" (or similar formulations). Whether UniSim is the first model to achieve this level of transferability is also unclear: The authors should compare to existing models that are finetuned in a similar manner to UniSim, and then compare the performance of different finetuned backbones.\n\nFor the second claim, the \"multi-task pretraining\" is actually a single task (force matching) pre-training in disguise. The off-equilibrium loss L_o (Eq.7) is the standard squared error force loss that is also used to train machine-learned force fields: H_out[i] can be regarded as the atomic energy contribution of atom i to the total predicted energy sum_i H_out[i]. Then, grad_X sum_i H_out[i] just computes the gradient of this predicted energy, i.e., the associated forces. By minimizing L_o, the predicted forces are matched to the reference forces. It turns out that the equilibrium loss L_e (Eq.10) is actually the same objective in disguise (i.e., predicted forces are matched to a reference force). To see this, recall that any potential energy surface is well-approximated by a harmonic potential close to an equilibrium (take the Taylor expansion and truncate after the second-order term, the first-order term vanishes because the forces are zero at equilibrium structures, which are minima of the potential). Then, the term -(X'-X)/sigma_e is just the force of a harmonic potential given by V(X') = 1/(2 sigma_e) (X'-X)^2. The authors should make this connection more transparent and consider changing  \"multi-task pretraining\" to \"force matching pretraining\" (or a similar term), and clearly explain that L_e is just a clever trick to generate synthetic data with estimated force labels from equilibrium structures.\n\nFor the last claim, I think the shown data does not support the statement that the state transition patterns for long-time steps are learned. In fact, the TIC0/TIC1 plots shown in several figures clearly show that the sampled distributions do **not** closely approximate the probability distributions sampled during MD. Perhaps UniSim is competitive with other models that attempt time-coarsened dynamics, and I acknowledge that it is a step in the right direction, but the task is definitely not solved in a satisfactory manner.\n\nFinally, there is another claim that is hidden in the conclusion. The authors write: \"Experiments conducted on small molecules, peptides and proteins have fully verified the superiority of UniSim in distribution similarity compared to MD trajectories and transferability to out-of-distribution molecular domains.\" However, as I already stated above, the agreement to distributions sampled in MD trajectories looks actually quite poor. The metrics the authors look at (e.g., distance histograms, contact maps, etc.) are certainly informative, but they are not sufficient. The ultimate measure of quality must be how well actual thermodynamic observables are reproduced by the time-coarsened dynamics. After all, accelerated dynamics are useless if they give the wrong results, and the experiments that are done currently are not sufficient to \"fully verify the superiority of UniSim\", in particular with respect to regular MD.", "methods_and_evaluation_criteria": "The proposed methods are appropriate for the problem. However, the evaluation criteria are not sufficient to show that the method works well and is practically useful. As mentioned above, to truly measure the quality of a time-coarsened dynamics, it is necessary to check how well it reproduces thermodynamic observables extracted from regular MD simulations. I think the authors should at least try a very simple example, e.g., reproducing the free energy surface of alanine-dipeptide (a well-studied small model system) with UniSim, so readers are better able to judge whether the presented model is useful for practical studies or not.", "theoretical_claims": "The paper includes theoretical justifications, particularly for the stochastic interpolant framework (Section 3.3 and Appendix B). I have not checked the presented proofs carefully, but at a first glance, they seem sound to me.", "experimental_designs_or_analyses": "The experiments that were performed and the analyses that were done seem sound to me. However, there is an aspect of the work, which should be discussed more prominently: In Appendix C.2., the authors mention that every prediction of their model is followed by an energy refinement with OpenMM to prevent error accumulation. This is a very important detail that should be discussed in the main text, and it is also necessary to state how many optimization steps need to be performed on average. Every evaluation of OpenMM used for the minimization is roughly equivalent to a single regular MD time step, so this greatly affects the actual speedup offered by UniSim. For example, if UniSim can leverage a time step that is, say, 1000 times larger than a regular MD time step, but then on average 1000 optimization steps with OpenMM are necessary to stabilize the structure, there is no net gain in terms of efficiency (on the contrary, this would be less efficient than just running regular MD).", "relation_to_broader_scientific_literature": "The paper appears to cite the most relevant prior work. However, when discussing classical MD methods, quantum mechanics, and empirical force fields in the introduction, the provided references, while not outright wrong, are somewhat questionable. It almost makes it appear as if the cited works introduced the corresponding concepts. Perhaps it would be more appropriate to cite earlier seminal work here.", "essential_references_not_discussed": "I did not spot any essential omissions.", "other_strengths_and_weaknesses": "**Strengths**\n\n+ The core idea of a unified simulator for diverse biomolecules is valid.\n+ The paper is well-written and easy to follow.\n\n**Weaknesses**\n\n- Lack of a direct comparison of computational efficiency with baselines and classical MD (especially considering that minimization with regular force fields seems to be necessary to stabilize trajectories).\n- Lack of experiments that investigate the prediction of thermodynamic observables with the proposed framework, assessing the practical usefulness of the approach\n- Some aspects of the work are unnecessarily \"obfuscated\" (e.g., the multi-task pretraining is really a single task in disguise), or hidden in the appendix (e.g., the need for minimization with a classical force field for stabilization).\n- Some design choices are questionable, for example, the energy normalization described in Appendix C.1. effectively introdces a dataset-dependent unit conversion. This may not cause a large issue if all datasets are sampled at roughly comparable temperatures, but it seems unjustified from a theoretical standpoint.", "other_comments_or_suggestions": "I think the TIC0/TIC1 plots could be improved: Instead of showing a scatter of the time-coarsened sampling over a contour of the true probability distribution, it would be better to do a side-by-side comparison of the true probability distribution and the distribution sampled by UniSim's time-coarsened dynamics (both as contours). This would allow a more direct comparison and allow readers to better judge the agreement of time-coarsened dynamics with the ground truth.", "questions_for_authors": "1. Can you provide a quantitative comparison of the computational cost (e.g., wall-clock time or speedup factor) of UniSim (including the time needed for structure refinement via minimization with OpenMM) compared to classical MD?\n2. Have you performed any experiments to assess the stability and accuracy of UniSim on longer timescales? How do the cumulative prediction errors affect the results over longer simulations?\n3. How sensitive is the performance of UniSim to the choice of hyperparameters? Did you perform any hyperparameter optimization, and if so, what were the key findings?\n4. Could you elaborate on the motivation behind the atomic embedding expansion technique? How does it specifically help in distinguishing different chemical environments compared to using just the periodic table? Did you do ablation studies?\n5. Could you elaborate a bit more on the intuition and advantage of using the \"gradient-environment subgraph\" approach? What is the expected impact of different choices of delta_min and delta_max?", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Meta-review: This paper proposes a unified simulator for time‐coarsened dynamics of biomolecules, based on stochastic interpolant and force guided bridge matching to adapt the generated trajectories to different chemical environments. 

The reviewers generally appreciate the proposed idea and the experiments, Reviewer mRvK and mTdA expressed concerns on metrics and evaluation procedures of the method, i.e., it is not clear whether if the proposed method actually solves the considered problem better than classic MD simulations or prior work like Wang et al., 2019. Reviewer mTdA also expressed some concerns on presentation of the method. 

However, the reviewers generally agree that this work provides a solid improvement for coarsed-graining of MD for accelerated simulation. The authors also provide good effort to address the reviewer comments. I recommend acceptance for the paper, but hope the authors further improve their method through further experiments and incorporation of the reviewer comments.
Decision: Accept (poster)