Neural Network Reparametrization for Accelerated Optimization in Molecular Simulations

Thank you for your valuable feedback. We agree that the manuscript requires substantial revisions for clarity and flow. Below, we aim to address each of your concerns, and to resolve the criticisms thoroughly.

Weaknesses

1. We agree that intro and background will benefit from significant reorganization.

1. Our revised intro will start with an overview of challenges in scientific simulations, such as the proliferation of saddle points and local minima leading to suboptimal results. While conventional dim reduction methods like CG offer partial solutions, they encounter issues like back-mapping and force-matching. Instead, we propose an innovative approach using an overparametrized neural ansatz. We demonstrate that CG reparam or a well-designed GNN ansatz, incorporating Hessian slow modes, achieve significantly lower energy states compared to direct optimization.
2. We are making a new Fig. 1 (see attached pdf) to outline the methodology.
3. In Sec. 2 on the Hessian, we'll clarify the motivation for using slow modes, rooted in the difference in fast vs slow mode evolution rate, which causes slow convergence at saddle points. Our goal is to adapt the optimization process to grant direct access to slow modes, hoping that it helps escape such saddle points. However, this approach faces challenges: 1) Changes in the Hessian may alter the slow modes during optimization. 2) The need to modify the optimization to favor slow modes. We address these by showing the stability and robustness of slow modes and by proposing linear CG and GNN reparam. Our experiments show superior efficacy of the GNN approach.
4. We add GNN parameter details in a table in the appendix. For the experiments in Fig 1, the GNN hidden dims are [20,10, 3]. We used $n/3$ slow modes $\Psi$ to get adjacency matrix $A= \Psi\Psi^T$ used in the GNN layers with output $h^l = \sigma(Ah^{l-1}W +W_s \odot h^{l-1} + b)$ with self-loop weights $W_s$ and biases $b$ .
5. We are improving figure captions.

2. We appreciate your comments and recognize the need for a clear justification of our approach. While our NN model is not a new architecture, using it to map Hessian data to FG modes is a novel approach within physics simulations. This use of NN diverges from traditional ML, mainly centered around supervised and unsupervised learning. We use NN as an ansatz in scientific optimization. This opens up new potential uses of NN in the realm of AI for science. It should also be useful in other ML tasks where saddle points or flat minima are problematic.

Furthermore, our results clearly demonstrate that simpler methods like linear slow mode reparam often fall short compared to GNN. For instance, in Fig. 1 "Pure LJ Loop," where the Lennard-Jones (LJ) potential complicates optimization due to its flatness and shallow minima, traditional methods like GD and linear CG reparam achieve low energies but fail to accurately model the complete coil formation, which the GNN effectively accomplishes. This is evidence for the efficacy of our neural reparam in complex energy landscapes, making our work highly relevant for ML-focused venues.

3. We agree with the need for clearer bio background and will make the following short additions to aid readers:

1. Overview of Experiments: We'll introduce two main experimental setups:

   • Energy Minimization on Synthetic Systems: Forming a coil using LJ potentials to mimic molecular forces.
   • MD for Protein Folding: Using the AMBER force field.

2. Details on Synthetic Coil Experiments:

   • Quadratic Bonds + LJ: We'll describe the energy function $E = E_{bond} + E_{LJ}$ where $E_{bond}$ and $E_{LJ}$ are the quadratic bonds and LJ interactions, designed to form a coil.
   • Pure LJ: Focuses on the challenge of optimizing in a flat energy landscape of LJ, highlighting the difficulty in achieving the lowest energy state due to extremely small gradients.

3. Protein Folding Using MD:

   • We'll provide a primer on MD's role in studying proteins, emphasizing the use of the AMBER force field to model essential interactions like bond stiffness and angle constraints.
   • We simplify MD by omitting solvents and focus on energy minimization using AMBER params from OpenMM. This aligns with our synthetic experiments, facilitating interfacing with ML frameworks (pytorch) to implement our approach.

Questions

1. As we understand it, metadynamics is mostly concerned with exploring and mapping the free energy landscape, so the goal is different. Nevertheless, our method can be combined with metadynamics by replacing the collective variables with our reparam.

2. We recognize the concern. But note that, while slow modes capture low-freq motions and help navigate flat energies, our approach integrates all modes both in the final relaxation phase as well as via residual connections in the GNN, ensuring no critical dynamics are missed. As in lines 29-33, our motivation for using slow modes is that they converge slowly and that including fast modes forces us to use smaller learning rates for numerical stability. However, the flexibility of our GNN method allows the model to adjustment the weight of fast and slow modes to decrease the energy faster.

3. Potentially yes, if the slow eigenvectors are transferable. The GNN graph comes from the Hessian backbone. In some settings such as a protein complex the total Hessian may be approximately block-diagonal, each components forming a block. The slow modes are then mostly localized on individual components. The learned GNN weights (encoding the actual atom locations, for example) are more challenging to transfer, but the loss should be invariant to symmetries of the system, such as SE(3) for proteins.

4. We have not, but that is a good suggestion. Our method can be used with any differentiable potential. We used the empirical AMBER force-field because of its popularity in MD.

We thank the reviewer for the time taken to review our work. Please check our shared rebuttal above for discussion on improving presentation and new experimental results.

Weaknesses:

1. Efficiency: We agree that models like AlphaFold take a different approach. However, we think our approach is crucial for domains where we do not have such huge database as proteins. Our approach is inherently "data-free," offering advantages in scenarios where access to molecular dynamics data is limited or unavailable, yet there is a need to simulate the entire system. Drug discovery or material design would benefit from faster ab initio simulations. This will allow us to go beyond learning patterns in existing molecules and experiment with entirely new ones.
2. Clarity of method: Please see our shared response at the top of the page. We have included an enhanced figure in the attached pdf that illustrates the process of extracting coarse-grained modes and their integration into the Graph Neural Network (GNN) reparametrization framework. We also sketch out the GNN architecture in there. This addition aims to clarify the methodology and strengthen the overall presentation of our approach.

Questions:

1. global minima and init: The energy landscape of proteins is inherently "rugged," (i.e. has many local minima and saddle points) making the search for a global minimum particularly challenging. Consequently, our objective is to identify the most favorable local minima. Even for relatively small proteins, varying initial conditions can result in different local minima. To illustrate how sensitive our methodology is to these initial conditions, we have included a figure in the pdf (see figure c and its caption). For the small protein 2JOF we find starting from three different initially unfolded positions, the final energy obtained by direct gradient descent come out different.
2. Large proteins: If the challenges to computing the Hessian for large proteins can be overcome (See our response to rev rxS5, Q3: scaling) our method could save computational costs because it generally requires a fraction of the number of iterations to reach a given energy. (see figure a in pdf). However, each step takes more compute due to forward and backward pass through the NN. With efficient hardware such as GPU, this overhead could reduce significantly.
3. memory cost: We think our model actually requires more memory due to the extra neural network. The extra memory cost can be linear or superlinear in the number of particles, depending on how the number of slow modes is chosen to scale with the system size. Our argument is that in the era of large compute, this memory overhead is ok, as the overparametrization can yield significant benefits in terms of convergence to deeper minima and avoiding saddle points.

Thank you for your comments. We try to address them below. Please also read our general rebuttal above, which details our plan for improving presentation and more.

Weaknesses:

1. Unclear presentation: We agree. Please see above and the attached pdf for flow-chart and algorithm.
2. Effective DoF: we will add more details, but the gist of it is using Corollary 2.3.1 and 2.3.2 as a loss function to find the symmetry operators $L$. We can approximately satisfy 2.3.1 when $\|LH\|< \epsilon$, which happens when $L$ lives in the slow mode subspace. Next, 2.3.2 can be satisfied when $[L,H]=0$, so we can solve an optimization problem $L_0 = \mathrm{argmin}_L \|[L,H]\|$.
3. GNN graph: please see attached pdf and general comment above. The adjacency of the graph (nxn with n being number of particles) is $A = \Psi\Psi^T$ with $\Psi$ being the slow modes.
4. Comparison with SOTA CG: True, but unfortunately existing MD frameworks are very opaque and difficult to modify. We tried and failed to interface our method with OpenMM. This seems to be general problem with the field of MD. However, we recently learned about an endeavor to implement MD in Jax. We will try to migrate to that. However, currently, we are unable to offer better comparisons with CG. We decided to present the approach inspite of lacking CG comparison because it is a flexible alternative to CG and can even be combined with CG.

Questions:

1. n0 In the GNN reparam, n0​ denotes the node feature dimension of $Z_h$, which is the input to the GNN, thus n0 is also the input dim of GNN. Recall, the atom positions $X$ are reparametrized as $X = \rho(Z)= GNN_\theta (Z_h)$ (as in equation 2), where $Z=(Z_h, \theta)$ and $\theta$ are the GNN weight, with n0​ being the input hidden dimension.
2. Equivariance: very good question! Yes! The formula for the Hessian backbone (eq 7) takes a norm over spatial indices and is therefore invariant under SE(3). We will emphasize this point in the paper. Thus, the Laplacian is also SE(3) invariant and the slow modes do not have a spatial index (scalars under SE(3)).
3. Scaling: Our methodology is compatible with any differentiable potential energy function. The primary limitation for large biomolecular systems is computing the Hessian matrix to get the slow modes. In biomolecules and MD the loss is generally the sum of pair-wise or triple interactions and there fore the Hessian can also be written as the sum of such easy-to-compute Hessians. Thus we don't have to compute the full Hessian matrix until we want the spectrum. We have relied on existing methods such as Lanczos and Davison to get the slow spectrum. Whether the complexity of the force field becomes a burden for scaling depends on the computation graph of the Hessian from autograd, but our guess is that it won't be a problem.
   The GNN can be scaled up by sparsifying the adjacency or other methods people have developed, so we don't expect the reparamterized part to be an issue.

Thank you, we appreciate your pertinent comments. Please also check our general response above about improvements to the presentation.

Weaknesses:

1. The plot was mistakenly omitted. We are adding it to the appendix.
2. Thanks, adding the citations.
3. n is the number of particles, d is the embedded dim, e.g. d=3 for 3D coordinates x,y,z.
4. Missing citations: we are adding the citations for that paragraph

Questions:

1. For $Z_{h0}$ initialization, we used two methods: 1) random; 2) started from random but optimized together with GNN weights to match a given initial configuration of atoms. Method 2 introduces a minimal overhead: for protein 2JOF with 284 atoms it converges after 900 steps, taking 0.67 seconds.
2. The GNN loss is the energy function $\mathscr{L} = E$! This is the crucial point about reparametrization. We still deal with $E(X)$ but now $X$ is $X = \rho(Z)$ (as in eq 2) where $Z=(Z_h, \theta)$ and $\theta$ are the GNN weights. So the optimization problem changes from $\mathrm{argmin}_X E(X)$ to $\mathrm{argmin}_Z E(\rho(Z))$ for the GNN or other reparametrizations. Our framework is unsupervised. The training of the GNN is the goal: The positions of atoms are encoded in GNN weights and finding the final atom positions means minimizing the loss function which is the potential energy function. After the training, there is no inference step.
   Our GNN setup consists of two GCN layers with residual connections, followed by a projection to 3D. $Z_h$ and each GCN layer comprises 100 hidden dimensions. During training, once the energy curve begins to plateau (a minimum change in energy of 0.1 and 20 patience steps), we stop the GNN reparam. We then continue to optimize using full FG modes again using the same loss (energy) function as GNN.
3. This is a great suggestion. We believe that the Kob-Andersen Binary Lennard-Jones (LJ) glass model (KAM) could also be a potential application. KAM closely resembles our synthetic pure LJ loop simulations, with the key difference being that, in KAM, particles do not have a fixed underlying graph during optimization. Consequently, in KAM, we need to construct a spatial proximity graph and update it dynamically during training to leverage the GNN framework. Due to time constraints, we haven't conducted extensive experiments, but we have run a few, both a gradient descent (GD) and a GNN reparam.
4. Yes, we have run sweeps of the learning rate and we find the lowest energy is consistently achieved by a GNN, though direct GD beats some GNN that have large LR.