UniSim: A Unified Simulator for Time-Coarsened Dynamics of Biomolecules

• Q1: The first claim & its transferability

Firstly, we strongly agree that the term "reusable backbone" is more suitable to describe UniSim, which will be used in our revised manuscript.

Secondly, 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.

Lastly, we are unable to apply identical experimental setups to baselines for validating transferability probably due to:

1. Architectural Constraints: Methods like FBM incorporate protein-specific atomic representations (e.g., residue index) that are fundamentally incompatible with small molecules.
2. Pretraining Deficiency: None of them provides pretraining interfaces, preventing knowledge transfer to novel atomic species after training solely on peptide/protein data.

• Q2: The term "multi-task pretraining"

We 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.

• Q3: The practical usefulness in longer timescales

To 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), with its stability been verified based on RMSD to the initial state (Figure 2).

• Q4: The energy minimization technique & efficiency

Firstly, 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.

Regarding the computational efficiency, we provide Table 5, which demonstrates that UniSim maintains a 1–2 order-of-magnitude advantage in ESS/s (used in Timewarp) over MD, with refinement step included.

• Q5: The energy normalization trick

Firstly, 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.

• Q6: The visualization of TIC plots

Given 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) for long-timescale simulations.

• Q7: Hyperparameter sensitivity

In 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.

• Q8: Intuition of the atomic embedding expansion

Here we will elaborate on the rationale behind our design of atomic embedding expansion.

Since 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.

Specifically, we validate the effectiveness of atomic embedding expansion with the ablation on PepMD shown in Table 6.

• Q9: Intuition of using the "gradient-environment subgraph" approach

Thank 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.

• 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.

Thank 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.

Specifically, 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, 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.

After 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. 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.

Furthermore, 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. 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.

Reference

[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.

[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.

• 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.

Thanks for the question. We will elaborate on the design rationale of the gradient-environment subgraph as clearly as possible.

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.

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.

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.

Accordingly, our hyperparameters should satisfy:

• Scale Matching: $G_g$ contains $10^1$-$10^2$ atoms (aligning with small molecule sizes)
• Physical Consistency: $\delta$_max - $\delta$_min ≥ 10 Å ensures force computation completeness

Moreover, 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.

• 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.

Thanks for the question. Regarding the impact of $\tau$ on the model performance, we provide the following explanations:

• 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.
• 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.
• 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.

In conclusion, we believe that a moderate $\tau$ is a reasonable choice.

Reference

[1] Kong, X., Huang, W., & Liu, Y. Generalist Equivariant Transformer Towards 3D Molecular Interaction Learning. In Forty-first International Conference on Machine Learning.

• Q1: The validity metrics for protein simulations remain unsatisfactory.

We claim that the unsatisfactory validity metric of ATLAS stems from two primary factors:

1. We pursue a unified modeling framework across small molecules and proteins, which necessitates certain compromises in protein-specific modeling (e.g., MSA).

2. 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 and Table 2, 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.

• Q2: Some evaluation criteria are not thoroughly explored.

We sincerely appreciate your valuable suggestions regarding the evaluation criteria. We address your concerns through the following systematic investigations:

1. 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. From the table, we can summarize two key observations:

   • 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.

   • As $T$ increases, most metrics show degrading trends. This is likely because excessive discretization of SDE leads to greater error accumulation.

2. 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).

• Q3: The lack of rigorous statistical analysis.

We have now updated the statistical analysis for protein validity metrics (Table 4), which provides rigorous statistical validation of our performance claims.

• Q4: a brief discussion of methods that integrate reinforcement learning/enhanced sampling/adaptive sampling.

In response to your valuable feedback, we will further enhance the Introduction section with a brief discussion on those methods.

• Q5: Regarding the choice of the coarse timestep $\tau$.

Thanks for the instructful question. We provide the following responses:

1. 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.
2. 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.
3. Finally, we briefly discuss the impact of different $\tau$ on the performance.
   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.
   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.
   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.

In conclusion, we believe that a moderate $\tau$ is a reasonable choice.

• Q6: Regarding the training Data Balance.

We employed several engineering tricks to ensure the pretraining dataset remains as balanced as possible:

• During training, we require molecules of each batch originate from the same dataset.
• We ensure that the total number of atoms per molecule in a batch is similar using dynamic batch size.
• We impose the same upper limit on the number of batches for each dataset in a single epoch.

This approach ensures that the total number of atoms contributed by each dataset during a single epoch remains approximately balanced.