Neural Network Reparametrization for Accelerated Optimization in Molecular Simulations
Authors: Nima Dehmamy
Csaba Both
Jeet Mohapatra
Subhro Das
Tommi Jaakkola
Venue: NeurIPS 2024 poster (2024)
Keywords: coarse-graining
molecular dynamics
protein-folding
repametrization
hessian
graph neural networks
Primary area: machine_learning_for_physical_sciences
TL;DR: We introduces a neural net reparametrization technique that significantly enhances the efficiency and accuracy of energy minimization in molecular simulations

Abstract: We propose a novel approach to molecular simulations using neural network reparametrization, which offers a flexible alternative to traditional coarse-graining methods. 
Unlike conventional techniques that strictly reduce degrees of freedom, the complexity of the system can be adjusted in our model, sometimes increasing it to simplify the optimization process. 
Our approach also maintains continuous access to fine-grained modes and eliminates the need for force-matching, enhancing both the efficiency and accuracy of energy minimization.
Importantly, our framework allows for the use of potentially arbitrary neural networks (e.g., Graph Neural Networks (GNN)) to perform the reparametrization, incorporating CG modes as needed. 
In fact, our experiments using very weak molecular forces (Lennard-Jones potential) the GNN-based model is the sole model to find the correct configuration. 
Similarly, in protein-folding scenarios, our GNN-based CG method consistently outperforms traditional optimization methods. 
It not only recovers the target structures more accurately but also achieves faster convergence to the deepest energy states.
This work demonstrates significant advancements in molecular simulations by optimizing energy minimization and convergence speeds, offering a new, efficient framework for simulating complex molecular systems.

Neural Network Reparametrization for Accelerated Optimization in Molecular Simulations
Nima Dehmamy IBM Research
Csaba Both Northeastern University both, ctnortheastern.c
Jeet Mohapatra MIT CSAIL jeetnolnit.cdt
Subhro Das IBM Research
Nina Dohmany@i bn. con
subbro. subbro.dar@iaae www.imm.ccmmmmmm con
Tommi Jaakkola
MIT CSAIL
touni monnicssil..mmmmmm nit.edu
Abstract
We propose a novel approach to molecular simulations using neural network reparametrization, which offers flexible alternative to traditional coarse-gran soars-graiinng methods. Unlike conventional techniques that strictly reduce degrees freedom, the complexity of the system can be adjusted in our model, sometimes increasing it to simplify the optimization process. Our approach also maintains continuous access to fine-grained modes and eliminates the need for force- matching, enhanc- ing both the efficiency and accuracy of energy minimization. Importantly, our framework allows for the use of potentially arbitrary neural networks (e.g., Graph Neural Networks (GNN)) to perform the reparametrization, incorporating CG modes as needed. In fact, our experiments using very weak molecular forces (Lennard-Jones potential) the GNN-based model the sole model to find the correct configuration. Similarly. in pootein-folding scenarios, GNN-based CG method consistently outperforms traditional optimization methods. not only recovers the target structures more accurately but also achieves faster convergence to the deepest energy states. This work demonstrates significant advancements molecular simulations by optimizing energy minimization and convergence speeds. offering a new, efficient framework for simulating complex molecular systems.
Scientific simulations, particularly in molecular dynamics (MD), face fundamental challenges in finding optimal configurations. The energy landscapes these systems are characterized by numerous saddle points and local minima, making it difficult for traditional optimization methods to discover the most stable states. This complexity stems from the interplay between different scales of interactions, from strong covalent bonds to weak van der Waals forces, leading to slow convergence in gradient- based methods and often suboptimal results. For instance, in protein folding, the strong peptide bonds create steep energy barriers while weak hydrophobic interactions guide the overall folding process, creating hierarchy of energy scales that is challenging to optimize simultaneously.
To address these challenges, yoarse-graining (CG) methods have emerged as popular approach, reducing computational complexity by decreasing the number of degrees of freedom (DOF) and clustering them into collective modes. While these methods have shown success (Pak & Voth 2018 Hollingsworth Dror 2018), they face significant limitations. Traditional CG approaches require cumbersome procedures such as back- mack-mapping (returning to the original DOF) and force-matching (finding the forces experienced by CG modes) (Jin al. 2022), which can limit their efficiency and scalability. Moreover, the strict reduction of DOF in conventional CG can sometimes oversimplify the system, losing important fine-grained details necessary for accurate energy minimization.
The code can be found http:///:/.tttccc..oaaaaaaaaaaaa
38th Conference on Neural Information Processing Systems (NcurIPS 2024).

---

In this paper, we introduce an innovative alternative that overcomes these limitations through neural reparametrization. Instead of strictly reducing DOF as in conventional CG. our approach leverages an overparametrized neural ansatz. to represent fine- fine-grained (FG) modes as functions of CG modes. This reparametrization concept, similar to techniques such as Deep Image Priors (Ulyanov al. 2018). enables the neural network to dynamically represent the FG system while maintaining continuous access to FG modes and eliminating the need for force-matching. The overparametrization provides additional flexibility in navigating the energy landscape while the physical system has n x d degrees of freedom (n particles in d dimensions), our neural representation can use a higher-dimensionall latent space to find paths around energy barriers that might be difficult to traverse in the original space.
A key innovation in our approach is the incorporation of Graph Neural Networks (GNN) with structure informed by 'slow modes' -inhereln stable collective modes identified through spectral analysis of the system's 's dynamics. We show that these modes arise naturally from the structure of physical Hessians, which are Laplacian matrices over particle indices for broad class of potential energies. By focusing on these slow modes, which typically cause convergence bottlenecks in traditional optimization, we can significantly accelerate the learning process. The GNN architecture allows us to safely increase learning rates without stability issues, resulting in both faster dynamics progression and the discovery of lower energy states compared to direct optimization methods.
The effectiveness of our approach is demonstrated through experiments on both synthetic systems and real molecular structures. In particular, for protein folding with weak Lennard-Jone interactions. where traditional methods often struggle with the shallow energy landscape, our GNN-based model consistently finds deeper energy minima. This success can be attributed to two key factors: the ability of the overparametrized representation to explore the energy landscape more effectively, and the incorporation of physically meaningful slow modes into the neural architecture, which helps guide the optimization toward stable configurations.
The main contributions of this work are:
1. CG via reparametrization: A new paradigm that circumvents traditional challenges like force-matching and back-mapping.
2. Robust slow modes: Effective identification and utilization of stable modes across various systems.
3. MD simulations: Demonstrated improvements in efficiency and depth of energy exploration in protein dynamics.
4. Overparametrization benefits:
Evidence that an overparametrized framework can outper- form traditional DOF reduction in terms of convergence speed and energy minimization.
5. Data-free optimization: Our method modifies the optimization landscape without the need for training data, enhancing its applicability and efficiency.
Background
Traditional optimization in physics- based models, like (MD), faces unique challenges due to the shallow nature of these models, where physical DOF the trainable weights. Additionally, the interactions occur at multiple scales, from strong covalent bonds to weak van der Waals forces, leading to slow convergence in gradient-based methods.
To address these challenges, conventional strategies include preconditioning with methods like adaptive gradient (Duchi et al., 2211 Kingma & Bal 2014) quasi-Newton (Fletcher 2013). and CG. which simplifies the system by truncating DOF to focus on collective modes. However, both approaches have limitations: preconditioning methods struggle with cost and inefficacy due to non- diagonal Hessians in physics problems, and CG can be restrictive and require intensive back-mapping and force-maten matching steps (Jin et al. 2022).
In contrast, our approach utilizes neural network reparametrization to dynamically adjust system complexity, which may involve overparametrization. This method allows for flexible system represen- tation, which can simplify the optimization process. can help avoid local minima and accelerates convergence by exploring the configuration space more efficiently.

---

Neural Reparametrization in Practice Our neural reparametrization approach is not limited reducing DOF but can also increase them when beneficial, offering an adaptive solution the specific needs of simulation. This flexibility is crucial for addressing the hierarchy of interactions in molecular systems, where different forces operate at vastly different scales.
Implementation and Comparison to CG While CG methods focus on predefined collective modes and often involve laborious optimization steps like force-matee matching and back-mapping, our neural reparametrization approach defines modes based on the spectrum of canonical Hessian, directly incorporating these into the neural network's 's architecture. This not only bypasses the need for traditional CG steps but also enhances the adaptability and speed of the optimization process.
Advantages Over Traditional Methods Our method diverges from traditional data-driven ML approaches that require extensive datasets, which are often unavailable costly produce in molecular and material design. By not relying on training data, our approach provides robust framework for tackling complex optimization problems, from molecular dynamics to protein folding. with improved efficiency and without the constraints of data availability.
1.1 Traditional Coarse-graining
Raxd represent the degrees of freedom (DOF), such as particle positions or bond angles, and let L: X + R denote the energy or potential function. The objective is to simulate the dynamics of the system or to find high-likelihood configurations X* that represent deep local minima of L. Given that n is typically large and L is steep non-convex function, computations can be slow. Traditional coarse- goansegraining (CG) maps X to a reduced space of CG variables, where k < n. Implementing dynamics using CG modes requires determining the inter-mode forces ("force-matching") and how to revert to X bbback-mapping"
Let e E
The fine-grained (FG) energy function, Lrg X R. needs approximate potential Leg : > IR such that for E X,
Force-matching.
Φ X Z, CG:
Lca(o(X)) = Lx)(X).
(1)
The process of finding Leg called force-matchinn traditionally solved analytically but increasingly with machine learning for enhanced accuracy Jin et al. (2022); Majewski et al. (2023).
Back- bacck-mapping. The map Z Rixd is not unique, often resulting in multiple possible X for a given Z. Back-mapping typically involves sampling optimization to find physically plausible X configurations, avoiding scenarios like overlapping atoms or high energies. This can be complex when many X map the same Z, with current methods ranging from geometric reconstruction Lombardi et al. (2016) to refinement with molecular dynamics [Badaczewska-Die Dawid et al. (2020);
Rocl-Touris & Bonvin (2020) and data-driven approaches Yang GómezzBombarelli (2023 Wang
et al. (2022).
1.2 Neural Reparametrization as an Alternative to Coarse-graining
Instead of traditional CG, which reduces DOF through mapping a reduced space, our approach reparametrizes the DOF X as a function of CG-like modes. This reparametrization, given by X p(Z), where p: Z - X, offers flexible, reversible mapping that inherently includes benefits such as direct access to fine-gres ine-grained modes and elimination of porce-matching and back- mapping needs:
Reparametrization: X p(Z), p: Z
1. Flexible parametrization: Leveraging neural overparametrization and architecture design.
2. Direct access to fine-grained modes: X p(7) avoids the need for back-mapping.
3. Simplified energy computation: The energy for CG-like modes is Leg (Z) (p(())).
While this method can be computationally intensive as Lcg (Z) is computed using X, the efficiency gains in optimization speed and depth of energy minimization can offset the costs.

---

Linear
Graph Adjacency
GNN Reparametrization
Reparametrization
d
d
R
X
A =
GCN d(AZ))
GCN d(Ah'W)
Feat Proj MLP
=
Compute
Input: Loss L. config Xo, modes
Output: Optimized X
/+ 1. Compute Hessian backbone
H - VVL(Xo)
Samples fx 6X (V), // Perturbed samples
Compute Backbone
H +- Ex E (H (X')² // Hessian backbone (xyz norm)
/+ 2. Extract slow modes
+/
L f Lap(H)
// Laplacian of backbone graph
A, V eigvecs(L)
Extmet Slow Modes
slow &- V argsoro R]]
// Pick bottom eigenvectors
eigeestlapp(ll
/* 3. Construct GMN
// use slow mode projection as adjacency gnn + GNN(A, dims = [b, 3J)
A + Vston VT
4. Neural Reparametrizaion
Z ~ NO, 1)mxh
X GNNNZ W)
/ 4. Optimize
*/
while not converged do
X gnn(2n 2)
5. Optimize min LAUNNG Wh
- VoL(()
11 8: and GNN weights
end return
Figure 1: Overview of the neural reparametrization method. Top: Architectures used for reparametrization. In linear reparametrization, X = zT LT' V slow modes to construct graph with adjacency A = V stow V slow and use it in GCN layers to obtain X (Z). Left: Flowchart showing the key steps the method. Right: Detailed algorithm for implementation.
the GNN case, we use the
Neural Architectures for Reparametrization The reparametrization function p can range from simple linear projections to complex neural networks. Initially, we employ a linear projection onto identified slow modes:
== p(Z) = 2T Stow Slow
Σ 7T
iESlow
More generally, p may be a deep neural network (DNN), similar to the approach taken in prior work suggesting neural priors (e.g. Deep Image Priors Ulyanov et al. (2018)).
Graph Neural Networks (GNN) for Dynamic Reparametrization: Extending beyond linear models, we explore the use of GNNs, inspired by recent advancements in graph-based optimizations Both et al. (2023). Here, the GNN reparametrizes node states and was shown to find both lower energy states and exhibit faster convergence. Our idea is to which acts as weighted graph adjacency matrix for a GNN. In our experiments, we observe this the direct well linear reparametrization equation 3 The details of our GNN architecture are discussed in Section 3 Next, we derive the properties of the slow modes for large class of energy functions important in molecular systems.
"Hessian backbone" as graph,
GNN have significant advantages
1.3
The role of the Hessian
The success of optimization in molecular systems fundamentally limited by the disparity in evolution rates along different modes of the system. Near any configuration X, the dynamics of

---

gradient-based optimization can be understood through the eigendecomposttion of the Hessian = VV... The eigenvectors of H define the natural modes of the system, with their eigenvalues determining how quickly these modes evolve under gradient descent. Modes with large eigenvalues (fast modes) evolve rapidly but constrain the learning rate to ensure stability, while modes with eigenvalues close to zero (slow modes) evolve orders of magnitude more slowly, leading to extremely slow convergence, particularly near saddle points.
This disparity presents fundamental challenge: To maintain numerical stability, the learning rate must be small enough handle the fastest modes, but this makes the slow modes evolve at glacial pace. Traditional approaches like adaptive gradient methods attempt to address this by approximating a diagonal preconditioner, but they struggle with the strongly coupled nature of physical systems where the Hessian is far from diagonal. While conventional coarse- graiiiig partially addresses this by eliminating fast modes, it introduces other challenges such as force-matching and back-mapping.
Our approach takes a different perspective: instead of eliminating modes, we seek to identify and directly incorporate slow modes into our optimization process. However, this raises two key challenges. First, as the system evolves, the Hessian changes, potentially altering which modes are slow. Second, even if we can identify slow modes, we need way to modify the optimization to preferentially explore these directions. The next section addresses the first challenge by proving that slow modes of physical Hessians are remarkably robust, arising from fundamental symmetries of the underlying interactions. We then show how these robust slow modes can be effectively utilized through neural reparametrization.
2 Properties of Physical Hessians
We will now show that the Hessian of potential energies important in physics and molecular systems enjoy certain properties that lead to the robustness of slow modes. In short, if we find a stable backbone for Hessians of different configurations X, then the slow modes of the Hessian at X are close to the slow modes derived from the backbone.
Invariant potentials. In systems of interacting particles in physics, leading interactions are often pairwise and involve relative features, riy = X; X, (distance vector, relative angle, etc). These interactions are invariant under global symmetries, such as Euclidean symmetries (translation and rotation) or Lorentz symmetry (relativistic particles). These symmetries maintain the invariance of certain norms, 2² || x² 1y,, sy,, where 7) may be the Euclidean metric 1) diag(1, 1, or the Minkowski metric 17 = diag( 1,, For example, the Euclidean norm va v in d dimensions is invariant under rotations e + gu, where g E SO(d).
Energy function structure. Let denote the matrix of distances with rig = || Ily. Any function of rty is invariant under symmetries that keep || lly invariant. Assuming additivity, the energy function can be written as:
((x) (a(
where (1)(=) = fz((2). For example, the Coulomb potential between particles and j with charges y: and q respectively, is given by fig(2) kg:g; /z. The Lennard-Jones potential fz)(2) = A B,y/ 26 in molecular systems is also of this form.
2.1 Hessian of invariant potentials
The Hessian of potentials of the form equation ( has the special property that it the graph Laplacian of a weighted graph which depends on X, as we show now (see appendix E for details). This will play a crucial role in our argument about the robustness of the slow modes.
Hessian as graph Laplacian. Recall the Laplacian of an undirected graph with adjacency matrix is defined as = Lap(A) D A, where D the degree matrix with elements Dij d.p Σx Aig- The components of Laplacian can also be written as Lij = >x (āx 5, We show that the Hessian of L in equation 4 is a Laplacian. Let O = a/ax, and let f = qpr be the dual unit vector of r. First, observe that a,r,k (ox di) where fyke is the unit vector
5

---

of rjk and ō.z is the Kronecker delta (1 if = j. 0 otherwise). Let Hes[g] denote the Hessian of function g. We find that (app. E
- Hes = = E(os (6.
Ha = Lap(H)is
(5)
where (X) = Hes[ Note that H has four indices, with components H having two particle indices i and two spatial indices u.v. Thus, for every pair of spatial indices N,M, the Hessian H is a Laplacian over particle indices. The Hessian being Laplacian has an important effect on its null eigenvectors. To show this we make of the incidence matrix.
We are interested in the eigenvalues and eigenvectors of H, as these characterize the slow and fast modes of the system. First, given a weighted adjacency matrix of a graph. let and L be the "unweighted" adjacency and Laplacian matrices, where A.j = Asj * 0 and zero otherwise. follows that the null spaces of I. and are shared:
Theorem 2.1 (Null Space the Laplacian), Ler Null(M] denote the null space of a symmetric real matrix M. The rull space of the unweighted Laplacian is contained within the null space of the weighted Laplacian L, i.e., [u C [u
Sketch of proof. For any vector, IR", v² Lap(A)u Eis Ag((: v2)². Since A.j = 0 yields
Ag = 0, but not necessarily vice versa, null vectors of Null[L] C Null[L]. See appendix for full proof.
Definition 2.1 (Slow manifold). Let I be graph Laplacian (undirected, weighted or unweighted),
with spectral expansion = Σi=1 Apbro? ". Let < and Amax = -max(A] be the largest eigenvalue of L. We define the slow manifold
Slow:[L] = Span( (b/l1aal < <>²>max)
(6)
Theorem 2.2 (Slow modes of weighted Laplacians). Let A be the adjacency matrix of a weighted graph and À be its unweighted counterpart. Let T. = Tap(A) and f. = ())(A). Then [T. overlaps with Slow, [ up O(E²) corrections from the rest of the modes.
The sketch of the proof relies on relating the spectra of the weighted and unweighted Laplacians using the incidence matrix C. as L = CCCC and L = |CCTT. For a random configuration X the edge weights W will be random, as they arising from derivatives of (re()))) in equation (unless faz is quadratic which makes W constant). Then, using the assumption randomness on the weights W, we can show the slow modes of L are perturbations of order e² on slow modes of LL See Appendix for proof.
Implications for Coarse-Graining. The identification of slow modes the Hessian is crucial for coarse-grainge granse-graining... as these modes capture the essential dynamics of the system at larger scales. By focusing on these slow modes, we can develop reduced models that retain the key physical properties while being computationally more efficient.
Coarse- GrarssnGraining via Slow Modes. The identification of slow modes in the Hessian enables an effective coarse-graining approach, where fast dynamics are averaged out, retaining only the slow, relevant dynamics. This method is particularly advantageous in reducing computational complexity while preserving critical structural information.
2.2 Hessian Backbone and Robust Slow Modes
The slow modes of the Hessian Hes (X) = Lap(H(X)) can dynamically change during optimization. To ensure the robustness of these modes, need proxy for the unweighted adjacency matrix Å H. To this end, aggregate Hessians from perturbed configurations Samples(X) = [X' XXX.
Backbone: H
Σ
X'emmmlle(X)

---

Bond+LJ loop
Pure LJ loop
GD
CG Rop
GNN
GD
CG Rep
GNN
Figure 2: Synthetic loop experiments. Example runs of the synthetic loop experiments with n 400 nodes. On the left (Bond+LJ), the potential is the sum of a quadratic bond potential Erond and weak LJ (12,6) ELJ- The bonds form a line graph Asond connecting node to + 1, and a 10 weaker Atoop connecting node to i 10 via the LJ potential. To the right (Pure LJ) where the interactions are all LJ, but with a coupling matrix A = Abond + 0.1Atoop. In Bond+LJ, GD already finds good energies and the configuration is reasonably close to a loop, though flattened. Both linear CG reparametrization (CG Rep) and GNN also find a good layout. The pure LJ case is much more tricky. But in most runs, GD almost gets the layout, but some nodes remain far away. The CG Rep fails to bring all the pieces together. Only GNN succeeds in finding the correct layout.
Time Bondl loup n=2000
Enorgy vs Time
0.35
-8016
(redet
model naae
0.38-
600K
00
GYM
0.05
200
110
6002
3302
6600
0.22
-0.5
0.00
-0.30
Time
Time
Figure 3: Synthetic loop folding (n 1000). Lower means better for both energy and time. In Bond+LJ (left), quadratic potential >. (ras+) - 1) attracts nodes i and i + 1. A weak LJ potential attracts nodes and i + 10 to form loops. In IJ loop (right) both the backbone i, + and the 10x weaker loop are IJ Orange crosses denote the baseline GD, green is GNN and blue is CG. The dots are different hyperparameter settings (LR, Nr. CG modes, stopping criteria, etc.) with error bars over 5 runs. In Bond+LJ, CG yields slightly better energies but takes longer, while GNN can converge faster to GD energies. In pure LJ, using CG and GNN can yield significantly better energies.
This aggregation helps identify consistently significant components across configurations, aiding in the extraction of reliable slow modes that remain effective over extended periods of optimization, In equation 7 i E Z., are the particle indices and the Frobenius norm ||H Em (H
sums over the feature indices (note that X has a particle index and a feature index μ E (1,...d d)).
Then, we extract the slow modes of the backbone, by doing a spectral expansion H E, Aphae o
and picking v with |^,] < 8² max, [²,]. for some small E < 1. The intuition behind equation is to identify the components in the sampled Hessians which have consistently high magnitudes. If we had taken a simple mean we could get very small values, because the components can fluctuate randomly. Also, if we had taken the variance instead of the norm, would get zero for quadratic L, where H is constant and has no variance, As we discussed above, the slow modes of the backbone H approximate the slow modes of sampled (X') up to OEE2) errors.
3 Experiments
We apply our method to protein folding using classical MD forces.
Settings: We use gradient descent to minimize {X). All experiments (both CG and baseline) use the Adam optimizer with a learning rate 10-2 and early stopping with 18.21 10-6 tolerance and 5 steps patience. We ran each experiment four times.
Baseline: we use gradient descent (GD) with Adam optimizer on the MD energy as bascline.
7

---

1.04-
3631
GNN
30-
1,92
SAMIL
2,002
30211
IPEW
INGD
3.00
6903
6997
10
1.5
2.9 25 30
0.872
0.575
000
5.005
0.032
0.555
Figure 4: Protein folding simulations Figure (a) shows the energy improvement factor (FG energy GNN energy) in the function of the speedup factor (FG time / GNN time) for the six selected proteins marked with different colors (c). In all cases, the GNN parameterization leads to speed improvement while it converges higher energy. (b) However, the higher energy in some cases, 2JOF and UNC proteins, results in slightly lower RMSD value, which measures how close the final layout is to the PDB layout. The data points are averaged over ten simulations per protein.
Opennm
GNN
Unfolded
GD
GNN 100
GNN 500
RMSD (backoone)
GNN
OpenMM
200
400
600
800
Time
Figure 2JOF (Trp-Cage) protein folding. Figure (a) shows the RMSD value evolution of the 2JOF protein as it goes from an unfolded to folded stage. At every step, we calculated the RMSD of the current layout compared to the PDB layout. We ran the OpenMM simulations at 298K temperature with 2I's timestep for 800000 steps, while the GNN and GD simulations were performed for 400000 steps with various hidden dimensions (10, 100, 300, 500). The black curves show the stochastic nature of protein folding using OpenMM. (b) The first figure shows the PDB (red) and unfolded (blue) layout; the second one is the GNN 500 final layout (blue), while the third is one of the OpeNMM layouts, corresponding to the black curve.
CG model: We use four different choices for the fraction of the eigenvectors to use in CG equation 3 3 x (#Amini) AminoAcids)) 30%, 50%. and 70%. We use a two stage process. First, we use CG as in equation 3 = p p(Z)= 2d V Slow and minimize Loo (Z) L (p(2)) over Z. After convergence to X. = Zo). we add 6 to X, and optimize the fine-grained 6XX,, starting with oX 0.
GNN model: We GNN consisting of graph convolution (GCN) layer with self-loops and one node-wise MI.P layer, projecting the GNN output to 3D to get particle positions. The GCN takes Zno = input, with ho and has weights We h Then, GCN output gets Tanh activation and is passed to the MLP layer to yield X. The CG parameters in this case are Zn. We and the weights and biases of the MLP.
Synthetic coil: We use quadratic and LJ potentials to make synthetic systems whose minimum energy state should be a coil (looping every 10 nodes), inspired by MD potentials. Figure shows many experiments using GD, CG, and GNN. In the quadratic Bond+LJ case, GNN yields good

---

Robustness of GNN results
b
1830
-1225
-1370
130
1.50
CAN 1eco
530
2001 Lerts
11230
-359
-2200
-197
10
-133
15
33
Figure 6: Learning rate and initialization in protein folding for pdb 2JOF: We conducted a sweep of the learning rate to see how robust the advantage of GNN over direct GD the energy achieved by GI and GNN vs the number of iterations and wallclock time. GNNI and GNN2 use one and two GCN layers, respectively. We used early stopping which generally stopped the runs after 3-5k steps. The grey star shows the OpenMM results after 5k steps, which has worse (higher) energy than our GD and GNN runs, but it takes a fraction the time (it has many efficiency tricks that our code doesn't 'tt have). The dashed line shows the energy achieved by OpenMM after 10k steps. As we see, some our GNN models reach energies close to the 10k steps of openMM in fraction of the steps. All experiments show the best energy among three runs. shows the effect of initialization on the GD runs. We do find the protein converges to significantly different conformations based the init.
and we show
speedup, while CG yields better energies. The benefit of CG and GNN become more apparent in the harder pure IJ problem, where GD fails find good energies, while CG finds much deeper energies, followed by GNN (Fig. 2
Protein folding with classical MD: We implement a simplified force-field with implicit solvent (i.e. water molecules not modeled and appear as hydrogen- bonginnnnninggggg and hydrophobicity terms; Aj. In protein folding our energy function consists of five potential energies: bond length Esond. bond angles Eangle- der Waals Evaw, hydrophobic Exp and hydrogen bonding E Ceci et al. (2007). Figure shows an example of these coupling matrices for the Enkephalin (1PLW) protein. To experiments four small proteins: Chignolin (5AWL),
app.
evaluate the effect of our CG model, Trp-Cage (2JOF), Cyclotide (2MGO) and Enkephalin (1PLW).
Protein Folding with Classical MD Using AMBER Force Field In the updated simulation approach, we incorporate the AMBER force field, known for its accurate representation of molecular interactions, particularly in proteins. This force field is implemented using the parameters from OpenMM Eastman et al. (017)) and comprehensively models the following interactions:
Bond lengths Esond and bond angles Eangle
Torsional angles Etorsion
Non-bonded interactions including van der Waals Evaw and electrostatic Esteo forces
We utilize the functional forms and parameters specified in the AMBER force field:
Ebond
Σ khonalr - ro)2
Eangle =
Σ
Kangte (0 0o)
(8)
bands
angles
Etorsion
Σ torsions
Vn
cos(nw - -
Evaw
Σ
Aij
Bij
(9)
12
6
xzj
9:9p
Felec
(10)
Here, r and represent the bond lengths and angles, respectively, with To and 0 as their equilibrium values. The torsional term Etorsion includes over all torsion angles with periodicity

---

amplitude Vn, and phase y, The Lennard-Joner potential in Evaw is characterized by parameters Aij and B,j. and Eeloc calculated using the Coulombic potential with partial charges d: g and the relative permitrivity
In this simulation, exclude the modeling of solvent effects entirely, focusing solely the protein in vacuum. This approach simplifies the computational model while emphasizing the direct interactions within the protein.
The overall energy of the system is then given by:
((X)= = Eboud + Eangle + Etorsion + EvdW + Eelec
(11)
Figure shows the interaction matrices for the Enkephalin (IPLW) protein. Our framework has been extended to efficiently compute these energies and gradients, facilitating the simulation of protein folding dynamics in our coarse-grained model. We test our model on several small proteins including Chignolin (5AWL), Trp-Cage (2JOF), Cyclotide (2MGO), and Enkephalin (1PLW) to evaluate the effectiveness of approach.
Protein results: Denoting the final energy and run time of the GNN model by EGNN and LGNN, and baseline by Eg and fo. we compute the energy improvement factor & Fo/ EGNN N and speedup factor öf = fo/tGNN to plot different proteins together. Figure 4/a shows the mean of &E vs dł over the 10 runs for GNN the model with hidden dimensions 300 (error bars are STD). Overall, we find that all GNN models outperform the baseline in terms of run time and, eventually, also with energy improvement. To measure the folding quality, we RMSDD, comparing the final layouts the PDB structure.
Figure 5 shows the RMSD value evolution using different methods. While, in most cases, OpenMM reaches deeper RMSD value, our models could serve as a good initializer for accelerating molecular dynamics. To evaluate the robustness of these results, the number of GNN layers (one or two layers), and varied the initialization.
sweeps over the learning rate, varied
Figure s/owws the results of these tests for the protein 2JOF. We used early stopping for switching from CG to FG in our GNN models and GD, resulting 3-5k iteration steps. Compared with 5k steps of OpenMM simulations, both our GNN models and GD with Adam reach significantly deeper energies with fewer steps (a), with the lowest energies being all GNN. However, OpenMM takes less wall-clock time (b). Nevertheless, the depth of the energies achieved by GNN at 3-5k steps is close to 10k steps with OpenMM. More efficient implementations of our GNN may further improve these results.
Discussion
We showed preliminary evidence that CG through reparametrization can yield some improvements over non-CG baseline in protein folding, both in terms of run time as well as energy. This method has the advantage that it does not require force-matching or mapping However, more experiments are needed compare it against traditional CG methods. In fact, using ML learn force-matching might provide further advantage by removing the need to evaluate LCg (2) ((X) via the fine- grained modes X. Also, while our canonical slow modes are derived for physical Hessians, the reparametrization approach to CG is general and could be applied to other ML problems.
Acknowledgment
JM was partly supported by a grant from the MIT-IBM Watson AI Lab. CB's work was done partly during his internship at the MIT-IBM Watson AI Lab.
References
Aleksandra E Badaczewska-David. Andrzej Kolinski, and Sebastian Kmiecik. Computational reconstruction of atomistic protein structures from coarse-grained models. Computational and structural biotechnology journal, 18:162-176, 2020.
10

---

Csaba Both, Nima Dehmamy, Rose Yu, and Albert-László Barabási. Accelerating network layouts using graph neural networks. Nature Communications, 14(1):1560, 2023.
Ceci, A Mucherino, M D' A'Auzzo, Daniela Di Serafino, S Costantini, Facchiano, and Colonna. Computational methods for protein fold prediction: ab-initio topological approach. Data Mining in Biomedicine, pp. 391 029114292220 2007.
John Duchi, Elad Hazan, and Yoram Singer. Adaptive subgradient methods for online learning and stochastic optimization. Journal of machine learning research, 12(7), 2011.
Peter Eastman, Jason Swails, John D Chodera. Robert T McGibbon, Yutong Zhao, Kyle Beauchamp. Lee-Ping Wang, Andrew C Simmonett, Matthew P Harrigan, Chaya D Stern, et al. Openmm 7: Rapid development high performance algorithms for molecular dynamies. PLoS computational biology, 13(7);e 1005659, 2017.
Roger Fletcher. Practical methods of optimization. John Wiley & Sons, 2013.
Vineet Gupta, Tomer Koren, and Yoram Singer. Shampoo: Preconditioned stochastic tensor optimiza- tion. In International Conference on Machine Learning, pp. 1842-1850. PMLR, 2018
Scott A Hollingsworth and Ron O Dror. Molecular dynamics simulation for all. Neuron, 99(6): 1129-1143, 2018.
Jachyeok Jin, Alexander Pak, Aleksander EP Durumeric, Timothy D Loose, and Gregory Voth. Bottom-up coarse-graining: Principles and perspectives. Journal of Chemical Theory and Computation, 18(10):5759555 57919 2022.
Diederik P Kingma and Jimmy Ba. Adam: method for stochastic optimization. arXiv preprint 14Xiv:1412.69800 2014.
Greg Landrum, Paolo Tosco, Brian Kelley, sriniker, gedeck, NadineSchneider, Riccardo Vianello, Ric, Andrew Dalke, Brian Cole, AlexanderSavelyev, Matt Swain, Samo Turk, Dan N, Alain Vaucher, Eisuke Kawashima, Macicj Wójcikowski, Daniel Probst, guillaume godin, David Cosgrove, Axel Pahl, JP, Francois Berenger, strets 133, JLVarj,, Noel O' O'Boyle, Patrick Fuller, Jan Holst Jensen, Gianluca Sforna, and DoviathGavid. rdkit/rdkit: 2020 03 1 (q1 2020) release, March 2020. URL
/ttcp:://doi.or//00. org/10 52w1//enooo. 3732262
Leandro E Lombardi, Marcelo A Martí, and Luciana Capece. Cg2aa: backmapping protein coarse- grained structures. Bioinformatics, 1235-11235 2016.
Maciej Majewski, Adrià Pérez. Philipp Thölke, Stefan Doerr, Nicholas E Charron, Toni Giorgino, Brooke E Husic, Cecilia Clementi, Frank Noé, and Gianni De Fabritiis. Machine learning coarse- grained potentials of protein thermodynamics. Nature Communications, 14(1):5739, 2023.
James Martens and Roger Grosse. Optimizing neural networks with kronecker- factored approximate curvature. In International conference machine learning, pp. 2408-2441.. PMLR, 2015.
Alexander J Pak and Gregory A Voth. Advances in coarse-grained modeling of macromolecular complexes. Current opinion in structural biology, 52:119-126 2018.
Jorge Roel-Touris and Alexandre MJJ Bonvin. Coarse-grained (hybrid) integrative modeling of
biomolecular interactions. Computational and structural biotechnology journal, 18:1182-1190, 2020.
Dmitry Ulyanov, Andrea Vedaldi, and Victor Lempitsky. Deep image prior. In Proceedings of the IEEE conference on computer vision and pattern recognition, pp. 9446-9454, 2018.
Wujie Wang. Minkai Xu, Chen Cai, Benjamin Kurt Miller, Tess Smidt, Yusu Wang, Jian Tang, and Rafael Gomez- Bombarelli. Generative coarse-graining of molecular conformations. arXiv preprint arXiv:2201.12176. 2022.
Soojung Yang and Rafael Gómez-Bombarelli. Chemically transferable generative backmapping of coarse-grained proteins. arXiv preprint arX1v:2303.01569, 2023.

---

A Protein folding with classical MD
In protein folding our energy function consists of five potential energies for: bond length Ebond. bond angles Eangle. Van der Waals Evaw. hydrophobic Enp and hydrogen bonding Æø Ceci et al. (2007). Note that are ignoring the solvent (e.g. water) and writing using potentials, or Iorce helds. To calculate the force field, we use distance, and ingle-based, θ, potentials. For each amino acid, we use the rdkit Landrum al. (2020) package to acquire bond length, Np. and bond angle, 8g (every triplet of atoms deñnnnn the bond), information that we use to define quadratic energies Ebond and Eangle- We use Lennard-Jones (IJ) potentials, Vn.g(r) (r) = rP - rr approximate Evaw between all pairs of atoms, Er between atoms prone to form a hydrogen bond (certain H. and O in our case), Phyp between atoms in hydrophobic residues, yielding
(() = Ebbond + Eangle + EvdW + EH + Enp
=Kbona(r ro)2 2 Kangie 00)
(
+ EudW V12,6
+ + Vo,
Ehp Vo,4
(12)
OndW
0H
Tbp
Here the coupling matrix [ardwlis aj + aj where a, is the vdW radius of atom i. For atoms which form H- H-bonds, [onli = (b; 1.5A (hydrogen bonding radius) with b; if forms an H-bond, and b; 0 otherwise. [ohpli q + Cj where & 2Å if atom is in hydrophobic residue and c, = 0 otherwise.
We note that our choices for cH, CydW, Chp and khord, Kangle, can be source of error. Additionally, we "softened" the LJ potential to Vp.g 1/(r" + +) 1/(r" + {) with Ç (= 0.65, which is large and significantly reduces the penalty for overlapping atoms and may reduce accuracy.
B Additional Figures
Myltoee bund masic (An0)
Hystrshebe be matte
60
Figure 7: Enkephalin (1PLW). a) The peptide chain is built by stacking amino acids on each other using the peptide bond length from the literature, 1.32 Å. b) Van der Waals, hydrogen bond, and hydrophobic interaction matrix, that in the energy optimization.
C Energy minimization
Let X E ~ Raxd be a degrees of freedom (e.g. particle positions, bond angles, etc.) and let L X + IR. be the energy (loss) function. We are interested in finding configurations X which are local minima of L. We can find such X* using a gradient descent (GD), or continuous variant, gradient flow (GF)
dX dl
-
(13)
where is the matrix of learning rates (LR). In simple GD where : = cf is single constant times identity, GD evolves at different rates different directions, with some being much slower than others. At a given X, these "slow modes" are the eigenvectors of the Hessian HXX) VV (X) with eigenvalues closest to zero, as we review below. We will first define fast and slow modes the simple guadratic case and then generalize them non-convex cases in the next section.
12

---

2J0F
2MGO
-0.0100
*
MD 2JOF
0.026
MD ZMGO
-0. 01215
2J0F
-0.028
2MG0
0.0155
00.300
-0.032
0.0200
x
-0.034
-0.02 00.022
-0.036
0. 0.0250
0.02275
0..000
300
350
400
450
500
550
600
340
360
380
400
420
440
460
480
time
time
1PLW
5AWL
ND 1PLW
-0.020
*
5AWL SAWIL
x
-0.015
CG 1PLW
-0.025
0.020
-0.035
0.025
0.040
00045
-0.030
-0.050
300
325
350
375
400
425
450
475
500
350
400
450
500
time (s)
time
Figure 8:
Comparison performance of CG Hessian versus baseline MD. Point sizes correspond to the number of CG modes used.
Fast and slow modes for quadratic Loss. Consider the case where = = - Here H is Hermitian matrix and the Hessian of L,, with a spectral expansion given by H
E, S X; E R and di R" In this basis we have X(t) Σ, (t) with & Rd.
Projecting equation 13 onto one of the eigenmodes we get
de; dt
GT dX
- x di
(14)
where we assumed /dt/dtt = O. From equation 14 we see that the decay/growth rate along mode d/ | Hence, modes with A close zero are the "slow modes", evolving very slowly, and large [Ag] defines the "fast modes", Since (t) c, (0) exp[-t/r. TT. with time scale T; 1/(=X;), the fast modes evolve exponentially faster than slow modes. This disparity convergence, because the fast modes force us to choose smaller E to avoid numerical instabilities. Two potential ways to fix the issue with disparity in time scales are: make rates isotropic (second-order methods and adaptive gradients); 2) mode truncation or compression (CG). We will briefly review the former here.
the rates results in slow
Adaptive gradient and second-order methods. Newton's method uses E = which makes GD isotropic along all modes, but expensive (O((3))t in our case). Quasi-Newton meth- ods, e.g. BFGS Fletcher (2013), approximate H H-1 iteratively. but are generally also slow. Another, more efficient approach is adaptive gradient methods, such AdaGrad Duchi al. (2011) and Adam
Kingma & Ba (2014) which approximate H by / 9t + where 9: Σ. Y V.X(((t -)) is some discounted average over past gradients and 17 small constant. For efficiency, in practice only use the diagonal part of this matrix to approximate H-1 As we will see in experiments, this approximation, while being far superior to GD with constant LR, is still very slow for MD tasks.
Most second-order methods are designed to work for generic problem and don't make strong assumptions about the spectrum of the Hessian. Recent second-order methods such as K-FAC Martens & Grosse (2015) and Shampoo Gupta et al. (2018) work with block diagonal approximations of the Hessian (or the Fisher information matrix), which usually emerges in deep learning models
13

---

MD
60
Nalive
86
142
198
Figure 9: The folded structures of the 2JOF protein by using the CG and baseline method. The numbers in front of the rows are the numbers of eigenvectors used in the CG reparametrization. Dashed frames show the minimum energy embedding in each case, while the thick line frame highlights the absolute minimum layout.
due to model architecture. Instead, we will exploit the spectral properties of the Hessian in physics problems. Fast and slow modes generally arise in physics due to vastly different strengths forces (e.g. weak van der Waals vs strong chemical bonds).
C.1 Generalized fast and slow modes
The notion of fast and slow modes is helpful for the analysis of any time slice of the dynamics during
which the Hessian is not changing dramatically. Consider a configuration X (1) and let 5t be a small time interval. We are looking for modes which are almost stationary over 8t. To identify these modes, we can for instance find perturbations 6X which would have almost zero dynamics. concretely we find the dynamics of X 8XX as
d
x 6X) -X- - dt
=H6X 0(xx)²
(15)
meaning, a small 8XX adds eH6 to the dynamics.
Thus if 8XX is a zero mode of the Hessian, H8X = 0, it won't change the dynamics of X. To define slow modes, we slightly relax this and look for normalized modes 6x//|6XX] whose associated time scale is much longer than a desired time scale 8t.
T H
|=A| >
8t
(16)
which just means that we need to find the approximate zero modes of the Hessian (X).
CG by projecting to the slow manifold. Because the dynamics of the modes above is very slow over öt, we can safely increase the time scale and run their dynamics over much longer periods At >> St. The essence of our algorithm is to ignore fast modes and project and evolve the system on "slow manifold" spanned by the slow modes of the Hessian. However, the main challenge how to deal with the fact that the Hessian is not constant and depends on the configuration X. We address this point next. We show that for a large class of physical potentials one can find a reliable set of approximate slow modes.
14

---

D Properties of Physical Hessians
Invariant potentials. In systems interacting particles in physics, most of the leading interactions pairwise and involve relative features, Tij X; Xj (distance vector, relative angle, etc). Moreover, they are often invariant under certain global symmetries, such as Euclidean symmetries (translation and rotation)
Lorentz symmetry (relativistic particles). These symmetries keep some 2-norm of vectors, x² = |,||, = 1' invariant. Here may be the Euclidean metric 7) diag(1, 1, 1) or the Minkowski metric 1) diag(- 1, 1) for relativistic problems, etc. For example. the Euclidean norm v² v in d dimensions is invariant under rotations v + go, where g E (d),,, and the Minkowski norm is invariant under the Lorentz group SO 1).
Let r denote the matrix of distances with riy || lyy Any function of Tij is invariant under symmetries that keep || !l, invariant. A general invariant energy function can combine rij for different j in arbitrary ways. Usually in physical systems each pair contributes an additive term in the total energy. Assuming additivity, the energy has form
L(X) Σ (rg)
(17)
where = fz(z) (symmetric under ++ j). For example, when particle has electric charge q. the Coulomb potential between i.j can be written as in equation 17 using fz)(2) = kg.93/ ///2.
Similarly, weak van der Waals (vdW) forces in molecular systems, which are modeled as Lennard- Jones potential, are also of the form in equation 17 with
der Waals:
fey(Tig) Vp.
Tij Jij
(r)((r)
1 T2
(18)
r
Here oij = 04 aj. where a, is the vdW radius of particle i, and vdW uses p = 12, q 6. Next, we show that the Hessian of equation 17 has an important property which aids finding its slow modes.
D.1 Hessian of invariant potentials
The Hessian of potentials of the form equation 17 has the special property that is the graph Laplacian of a weighted graph which depends on X, as we show now (see appendix E details). This will play crucial role in our argument about canonical slow modes.
Hessian as a graph Laplacian. Let 0 a/ax, and let f /r/r be the dual unit vector of First, observe that 0:rik = Fy (ō.y - 5 where f5 is the unit vector of ryk and d,z is the Kronecker delta (1 if j. 0 otherwise). Let Hes[gl denote the Hessian of a function g. We find that (app. E
Hes = 20,
Σ (6.g )33) (x)
(19)
where (X) (es(ful(rak) and given by
Hx(()= = (Fi((() (il(()
(v)
0 2 0 +
(20)
e=
Note that H has four indices, with components H having two particle indices i,j and two spatial indices Jr, Recall the Laplacian of an undirected graph with adjacency matrix A is defined as (A) = D A, where D is the degree matrix with elements Dzg 8zz Σ Ask- The
components of Laplacian can also be written as L,j = Σ* An 6zx). Thus, we see that the Hessian of L is indeed the Laplacian of H
Hes (es{_ = = - 6,)) = Lap(H)
(21)
where for every pair of spatial indices the Hessian is Laplacian over particle indices. The Hessian being Laplacian has an important effect on its null eigenvectors. To show this we make use of the incidence matrix.
15

---

D.2 Canonical backbone for the Hessian
As the Hessian depends on X. it is not clear whether slow modes found at a given X would be applicable to other X. We need some guarantee that a set of modes exist which are approximately slow modes for the Hessian at a range of different X. We could use multiple perturbed configurations
X + 6X with random 5X ~ (00 T) to get ensemble of Hessians H = [ + 8X)} and find the overlap of the slow modes of the Hessians in H. However, this expensive, roughly O(mkn²) for m = [H] and & slow modes. We cannot recompute the Hessian slow modes often. We also want method which more efficient than quasi-Newton methods such as BFGS. Our solution is to find a backbone for the sampled Hessians whose slow modes are guaranteed to be approximate slow modes of the actual Hessians. The key observation is that the Hessian in equation 21] is a Laplacian of a weighted graph. We show that the slow modes of weighted Laplacians overlap significantly with their unweighted counterparts.
We want to extract a set of slow modes from the sampled Hessians H(''). We then compute a backbone from these Hessians of the form
Backbone:
H
E
[X')('))
(22)
(ample(e(
Here i,j E Za are the particle indices and the Frobenius norm |H 112 = E (H sums over the feature indices (note that X has a particle index and feature index μ E {1, ...dd.dd. Then, we extract the slow modes of the backbone, by doing a spectral expansion H =; Antro?? and picking d/z with |A:] < 8² max, [A,], for some small E < 1.
The intuition behind equation 22 is to identify the components in the sampled Hessians which have consistently high magnitudes. If we had taken simple mean we could get very small values, because the components can fluctuate randomly. Also, if we had taken the variance instead of the norm, we would get zero for quadratic L, where H is constant and has no variance. However, these intuitions do not show that there would be any connection hetween the modes of the backbone H and the actual Hessians H (X'). Importantly, entries in H have signs, which affects the spectrum, whereas all entries in H are positive. So why should the spectra of H and H be related? This is where the structure of L comes into play. Indeed, we show below, for many physical L, the slow modes of the backbone H approximate the slow modes of sampled H (X') up to O(=²) errors.
Definition D.1 (weighted graph). Let 9 = (V,E) be a graph with vertices V = Zn. edges & V.
Let Å E Raxx denote the adjacency matrix A.j 1 if (i.j) E E and 0 otherwise. We denote weighted graph as G = E W) where W : - R are the weights of the edges. Let A denote the adjacency matrix of G. where Aij = j) or zero if (i, & E. The Laplacian L Lap( (A) of an undirected weighted graph is defined analogous to the unweighted graph as L == D - A with degree matrix elements Dig 8.k Σx Aik.
Definition D.2 (Slow manifold). Let T. be a graph Laplacian (undirected, weighted or unweighted), & 1 and Xmax = maxx (x+] be the largest
with spectral expansion I, = Σi= Axbhaa? w. Let
eigenvalue of L. We define the slow manifold as
Slow, [T.] = Span W//lAal < <E²2 EdXmax))
(23)
Theorem D.1 (Slow modes of weighted Laplacians). Let A be the adjacency matrix of a weighted graph and Å be its unweighted counterpart. Let L Lap(A) and L = Lap(A). Then Slow, [L]
overlaps with Slow, [i. up O(E²) corrections from the rest of the modes.
To prove this we will make use of the incidence matrix representation of the Laplacian. Definition D.3 (Incidence matrix). Given a weighted graph g = &V,E, W), define its incidence matrix as C: V & + [=1,, where for any edge = j) E E, C2,e -1 and Cz,- = and zero for other components.
Lemma D.2 (Laplacian in terms of the incidence matrix). Let w vec(W(E)) be the vector of all weights indexed in the same order as the columns of C. with we = Aur, e = (i.j) and let W be a diagonal matrix with w on its diagonal. Then, the Laplacian L Lap(A) can be written as L = CWC CT (proof in app. E.1
Because G and & share the same vertices and edges, their incidence matrix C' is the same. From Lemma D.2 L JCWCT and L = CCCT as & is unweighted. Using SVD, C USVT and
16

---

defining R= = US/V2 and Q = WV, we have
L= RQRQ..
(24)
Note that for a random configuration X the edge weights W will be random, as they arising from derivatives of fay (raz)) in equation 20 (unless fúj is quadratic which makes W constant). Therefore, we will assume Q has a uniform Gaussian distribution. Assuming W is also Gaussian, the spectrum of such Q V WV is somewhere between the distribution of W (for sparse graphs with [E] ~ O(VI)) ) and a Wigner Semi-circle (for dense graphs with |&] O(V)))), See appendix E.2 for more discussion. We also assume Q has particular block structure and that the spectrum of any diagonal block of Q should also follows a distribution similar to all of Q.
Slow subspace. We now sketch the proof for Theorem D.I For details, refer to appendix E.4 the SVD, = USVT, the slow subspace is
From
Slow, [ = S < max[S]]
(25)
Normalize $ = S/ max [S] and make them all positive (e.g. absorb their sign into U). For some
E sort the SV such that $ = diag(S,, S1) where the diagonal matrices Sy < ε and S 2 e. Now, the problem of finding Slow, [L]] becomes finding eigenvectors of the matrix M = ŠQST with cigenvalues O(=²). Using S. O()) and S1 ~ O(),, we can pull factors of E out from M and write
it as
E 0 C)
0 BT
B 0)
M = Mo + 86M,
Mo =
8M =
(26)
where 82 82 /n / rescaled so that the random matrices Å E RRAXXA Xna and E RR xnc
have a similar range of eigenvalues. Next, using perturbative ansatz for eigenvectors v/ = 1/ + Edr and eigenvalues X' X + ε6A, we solve M = Xabl up to Ogg²) corrections.
To find slow modes for J. we start from v/ E Slow; [E] Specifically, we start with eigenvector A of Á and concatenate it with zeros get (v) 0). We have Mar X with À 2² AA Using first-order perturbation theory, we find the corrections 5λ to the eigenvalues and eigenvectors to be
du (Moo - (( (()-))-
8A 1 8M
(27)
Putting all together we find the slow eigenvector v/ = vi + Edr up to order O(s2) to be
Slow [L] e 8(C-EEAA BTT ) = + + 0(E))
vi
(28)
meaning to first order in é the corrections to eigenvalues of slow modes vanishes. This is desired because the slow mode eigenvalues are O(E2). We also observe that slow modes of are mostly confined to [L] and only get O()) contributions from the fast subspace of L.
As side, it follows that all weighted graphs share the null space of the unweighted Laplacian. Proposition D.3 (Shared null space). Let Null[M] Span E RM,Mu 0} denote the null space of a matrix M E R The null space of the Laplacian L (unweighted) is contained in the null space of Laplacian L (weighted), meaning Null C Null[E].
Lemma D.4. Null(f.] = Nu]](R]]
Proof. Vu e Null o - ||RTu v12 and Ve e Null[R"] Lo RRT v = 0.
Proof of proposition D.3 Ve e Null Lo= RQR u hence, Null[ C Null[L].
Note that Null / and C [ull are not necessarily the same because weights can be zero, which could make the null space of the weighted graph larger than the unweighted one. Next, we present method for coarsiggrainigg using set of canonical slow modes.
17

---

E Invariant additive dyadic potentials
We want to Compute the Hessian of equation vector of r. First, note that
L(x Eij(rey) Let f = or/r be the dual unit
Öryk 0, 1111X, - Xx1l7
Or = 0x,
rjk (ōeg 6(k) f/ 8sk)
"1 ryh
(29)
Then the gradient becomes
Öryh jk
0.2(x) Σ f(altrs)
Ox:
Σ Maltissa 8cx)
j.k
= 2 E Plg(ru)ing.
(30)
where used îyN -fky to show both terms in (ōis 8ex) yield the same output. Finally, the Hessian becomes
-0,0,,=0;0,211)) 20, Σ (ra
-2 (ia(ri @ Pax + (
Σ (85e 65x) (rax 2 AM
Sl (raa) (1
Ozi (n
Pun 0,rah) rik
rik
( 71
Pur @ Askk
Σ (53z 876) fil(ra))Fr @ Aikk + Sike (rac)
risk
rit
(1 (n-080) (i(())6 08+ Pla(v)
(8.j - 654)
v
v=ria
(SHLOV)) [1((v)
8 8 i + Sia(()
(8is - 6z4))
Σ Ha(()) (Rig - 85x) = LLa((H)))))
(31)
This is because the components of Laplacian can be be written
Lig = ITap(A);; = (D - A)is
8ū Σ Art A,j = Σ Aux 85k)
(32)
18

---

E.1 Incidence matrix
The Laplacian D A of an undirected graph with adjacency be written as CWC /2 using the incidence matrix C and the edge weights W. This can be shown as follows
[owcT Σ CF
Σ
c( A A
(da - )u)) Autlsi - 0gh)
Cdand -11091 - Supde + AM
k
Σ
(ōuōς) 51155))A
8,y 2A15 2(D - A)ij = 22g
(33)
where we assumed A&ı- Ask (undirected graph).
So the same derivation of the backbone also holds for this case. The idea is that using the incidence matrix C and edge weights W (as diagonal matrix), any Laplacian L be decomposed as = CWC Then, doing SVD C = USVT we have
L = USV.SSTWW WV = UMUT
(34)
Where the matrix M = WVVT has interesting property, namely that its null space includes the null space of the unweighted Laplacian Lo = CC To see this note that Lo - US St which means columns Ug are the eigenvectors of Lo with eigenvalues S. The null eigenspace of Lo are the U; for which S; = 0. This subspace will also be null subspace for L, because that block is also zero in M, because M.j Σ. S. Va W RR Sy. So, whenever S, 0 or Sj = 0, Mis O, meaning that whole block in MM is zero and MU, = 0 (wite it better).
Example: power law. Hessian
Let f(r) = rp. We have f =pr and p" plp 1)re-2 yielding the
H = VVf(r) = p [(p² - (2p) f pn]
(35)
Bi Awt 2-2 [(p² - 2p) Puk f.x - pr) ik
(36)
Example: Lennard-Jones,
This potential has the form
f(r)= 4: (f)) (
(37)
where for classic van-der Waals potential p 24 12. The Hessian for this potential given by
vV)r)))))) = (((
p+2 [(p 2p) 8 f ml] (7) 9+2
[(9²
2g) f 69 f qm]
(38)
and Bi = H
E.2 Structure and spectrum of of VVVVV
To consider only the relevant subspace SVD, we have U, E Raxn and V E RR with n |
and Eg). For connected undirected graph m 2 2(n - 1) and V is full rank (VTV= I,).
Note the edge weights W come from the forces fúz (rty in equation 20 which for an arbitrary X
will be random. Assuming Gaussian distribution Wee (O o)) for all edges e, the matrix Q
19

---

will also have random Gaussian entries. When m V defines the eigenbasis of Q and W are the eigenvalues of Q. Similarly, in sparse graphs. where m O(n). V is approximately the eigenbasis and the spectrum of Q should have distribution similar to Wee. For dense graphs, where m O(n²), every entry of Q will involve a weighted sum over multiple Wee Then, from central limit theorem, entries of Q will asymptotically have a Gaussian distribution. From random matrix theory, we know that such Q will have a spectrum which follows the Wigner-semi senner--mm--iirlee law. In both cases (sparse and dense graphs) the spectrum of Q has finite variance and sits somewhere between a Gaussian and semi-circle.
E.3 Generalization to nonzero but small SV
We want to know how much the slow modes of weighted and unweighted graphs to overlaps. With the spectral expansion L = Ei Aeprei? Define the slow subspace as in equation 23
Slow [L] = Span ch: ||A:] < E² Amax(()
(39)
where Amax (L) max(A;) = max L ||1//|111111111111)))1 is the largest eigenvalue of and & < 1. In terms of the singular values (SV) of the incidence matrix C USV", the slow subspace becomes
Slow; [4] = {il Sa. < E max[S]}
(40)
We will show that the slow modes in weighted L CWC are perturbations to the slow modes of L. Define
M = S = SQST
(41)
Normalize $ = S/ max Break the space down to the slow and fast subspaces, based on whether Sü not. First, since L is positive semi demindefinite, can make all S 2 0. Let $ = S/ max S.
We sort the dimensions in S to have the small SVs appear first. Denote the block in S where S < by Sg. We have
2 0
S2
i
(42)
SF
We know the null space of i. where S4 = O, is shared with L. First, we remove the null space from L and L, calling the remainder Lo and Lo and the remaining SVs S. Then in this remainder subspace we need to find parts which are O(E). We sort the dimensions in S to have the small SVs appear first. We denote the block in S where S < max by Sg. We have
Se S.Qa1S)) (Mez MM
Me1
M = max (exx STRSSSTT
SIQNS SQnS,
Mn
(43)
Because Sg is O(s) and S₁ is 0(1), we will factor out the factors of e from blocks in M and write
EE2A =B)
M = max
(44)
EB"
C
Here A and C are random matrices built from their corresponding blocks in Q and sandwiched between S (for A), and S, (for C), which have (11) values. The spectrum of Q has distribution between a Gaussian with mean zero and a Wigner semi-circle, also centered around zero. We expect spectra of A and C to be similar to Q. Denote the spectral expansion of Q as
Q V
A = (A)
V = [
(45)
This is because when Qi N(, (7) we have (ignoring Bessel's 's correction for k > 1).
1 101 1 > x Var(A)
0² = Var(Qu)) ~
(46)
where we assumed Tr(Q)/n ~ mean(() = 0. Since block Qs of size ls is b22 entries sampled
from the same distribution as Q, we expect
| 1l
11Q:11
1
1 Var(Q.)
Var(Qu) k
(47)
h22
12
Thus, rescaling A E 2 X and C' E xae we get
A
C
Ō=
Var(A) Var(O)
(48)
nA
VTc
20

---

E.4 Approximate slow modes of L
If M did not have the off-diagonal blocks B, then Slow, and Slow. [Ż] would coincide, as the
Sg block and the S block would not mix when B Define Mo as the block matrix of M with
B 0.
A 0 0 C
Mo =
(49)
Using spectral expansions
= WAAAWA
C= AcAcV'
(50)
the eigenvectors of Mo consist of
V/ 0
0 (b:z (boz)= Aci (DO()
Mo
5² AAi
Mo
(51)
Since we are looking for slow modes, we must also consider the magnitudes of Aai and λci. Since A and C' entries are random samples from Q, we expect them have semi-cirele or Gaussian distribution similar to Q. Thus, we can use the variances of eigenvalues of and C' as proxy for the how the magnitudes of AAs and Acs compare. From equation 48 we have
1
E[A2] RS
Var(A) E[AG]
(52)
nA
Based on this we define rescaled é such that =²\Ai still has a smaller magnitude than dçi on average, meaning we want
"E[A3> < EEA
sAnA < nC
T
1
(53)
TC
We choose E such that the condition in equation 53 is satisfied. We can express M in terms of é by rescaling A and B to E2 = A and é BB. Now eigenvalues of have the same variance as eigenvalues of C. For brevity, denote M = max
max{s)] [S] M. We have
8 8B c).
M
(54)
T
To find how slow modes of AM = SQST/ max[S] differ from slow modes of SST we break M into block diagonal part and an OEE) off-diagonal perturbation
(62 0
-)
0 B
B)
B
M Mo + E8M,
Mo
SM
(55)
As in equation 51 eigenvectors of A n and C' are eigenvectors of Mo. Now want to find eigenvectors of MM with small OE=2) eigenvalues up order é corrections by treating 6M as perturbation.
(Mo + ( + (ō0) = (A + (8A) + EdO)
Mor + &6 + M,d+) + 0 = Auh + E(AAp) + A61)) + O
5M Mn 6Aph Ao
(56)
We only need the components of ف④) orthogonal to th, we can assume SUs v 0. From this we have
82 = Th 8M + T Mody = &8M
(57)
where we used T Mndg = ARy dzb = 0. Plugging equation 57 into equation 56 we can solve for dyb by inverting the matrices
(Mo A)ōy = (6) -8M)
= 8y6 (Mo A in) 5M)
(58)
21

---

where we added small yy to make the matrix Mo À invertible, as À is one of its eigenvalues. To find slow modes, start from slow modes of Mo which in the A subspace. Let 4 be eigenvector of A with A AAt'A. Concatenating v with zeros in the C' subspace have
Mod Aaa
(59)
Using this to compute aλ in equation 57 we have
82 v 8M = ( 0)
BT (A =0
(60)
meaning to first order in é the corrections to eigenvalues of slow modes vanishes. This is desired because the slow mode eigenvalues are O and we find that with this & ansatz the corrections will get are also at least O(E²). Next, we compute the corrections Ou to the eigenvectors. Plugging e into equation 58 with E² AA and Œλ we have
(Mo + in)
(E À X in)
(C (C-
0
81/
(Mu + -ntiinn-- -
( (C- S
0
(61)
where we dropped in in the lower block because E? XA is unlikely to be also eigenvalue of C, as A and C are random matrices.
Using the relation e = =B with the original and putting all together we find the eigenvector vl/ th + E0 up order (s²) to be
VAA (6(C 82 -
V
(62)
M = Axx' + +0E2) = -(E?)
(63)
22

---

NeurIPS Paper Checklist
The checklist is designed encourage best practices for responsible machine learning research, addressing issues of reproducibility, transparency, research ethics, and societal impact. Do not remove the checklist: The papers not including the checklist will be desk rejected. The checklist should follow the references and follow the (optional) supplemental material. The checklist does NOT count towards the page limit.
Please read the checklist guidelines carefully for information on how to answer these questions. For each question in the checklist:
You should answer [Yes], [No]. or [NA].
[NA] means either that the question Not Applicable for that particular paper or the relevant information is Not Available.
Please provide short (1-2 sentence) justification right after your answer (even for NA).
The checklist answers are an integral part of your paper submission. They are visible to the reviewers, area chairs, senior area chairs, and ethics reviewers. You will be asked to also include (after eventual revisions) with the final version of your paper, and its final version will published with the paper.
The reviewers of your paper will be asked use the checklist as one of the factors their evaluation. While "[Yes] is generally preferable to "[No]" ", it is perfectly acceptable to answer "[No] providedd proper justification is given (e.g.. "error bars are not reported because it would be too computationally expensive" or "we were unable to find the license for the dataset we used"). In general, answering binary way, we
"[No] or "[NA]" is not grounds for rejection. While the questions are phrased acknowledge that the true answer is often more nuanced, please just use your best judgment and write justification to elaborate. All supporting evidence can appear either in the main paper or the supplemental material, provided in appendix. If you answer [Yes] to a question, in the justification please point to the section(s) where related material for the question can be found.
IMPORTANT, please:
Delete this instruction block, but keep the section heading "NeurIPS paper checklist",
Keep the checklist subsection headings, questions/answers and guidelines below.
Do not modify the questions and only use the provided macros for your answers.
1. Claims
Question:
Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?
Answer: [Yes]
Justification: [TODO]
Guidelines:
The answer NA means that the abstract and introduction do not include the claims made in the paper.
The abstract and/or introduction should clearly state the claims made, including the contributions made in the paper and important assumptions and limitations. A No or
NA answer to this question will not perceived well by the reviewers.
The claims made should match theoretical and experimental results, and reflect how much the results be expected to generalize other settings.
is fine include aspirational goals as motivation long is clear that these goals are not attained by the paper.
2. Limitations
Question:
Does the paper discuss the limitations of the work performed by the authors?
Answer: [Yes]
Justification: [TODO]
23

---

Guidelines:
The answer NA means that the paper has no limitation while the answer No means that the paper has limitations, but those are not discussed in the paper.
The authors are encouraged to create a separate "Limitations" section their paper.
The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (e.g., independence assumptions, noiseless settings, model well-specification, asymptotic approximations only holding locally). The authors should reflect on how these assumptions might be violated in practice and what the implications would be.
The authors should reflect on the scope of the claims made, e.g., if the approach was only tested on a few datasets or with few runs. In general, empirical results often depend on implicit assumptions, which should be articulated.
The authors should reflect on the factors that influence the performance of the approach. For example, facial recognition algorithm may perform poorly when image resolution low or images are taken in low lighting. Or speech-to-tex system might not be used reliably to provide closed captions for online lectures because it fails to handle technical jargon.
The authors should discuss the computational efficiency of the proposed algorithms and how they scale with dataset size.
If applicable, the authors should discuss possible limitations of their approach to address problems of privacy and fairesss.
While the authors might fear that complete honesty about limitations might be used by reviewers as grounds for rejection, a worse outcome might be that reviewers discover limitations that aren't acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play an impor- tant role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed to not penalize honesty concerning limitations.
3. Theory Assumptions and Proofs
Question: For each theoretical result, does the paper provide the full set of assumptions and complete (and correct) proof?
Answer: [Yes]
Justification: [TODO]
Guidelines:
The answer NA means that the paper does not include theoretical results.
All the theorems, formulas, and proofs in the paper should be numbered and cross- referenced.
All assumptions should be clearly stated or referenced the statement of any theorems.
The proofs can either appear in the main paper or the supplemental material, but if they appear in the supplemental material, the authors are encouraged to provide short proof sketch to provide intuition.
Inversely, any informal proof provided in the core of the paper should be complemented by formal proofs provided in appendix or supplemental material.
Theorems and Lemmas that the proof relies upon should be properly referenced.
4. Experimental Result Reproducibility
Question: Does the paper fully disclose all the information needed to reproduce the main ex- perimental results of the paper to the extent that affects the main claims and/or conclusions of the paper (regardless of whether the code and data are provided or not)?
Answer: [Yes]
Justification:
Most experiments are uploaded in supplemental. Rest will be provided upon
request.
Guidelines:
The answer NA means that the paper does not include experiments.

---

If the paper includes experiments, a No answer to this question will not be perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data are provided or not.
If the contribution is dataset and/or model, the authors should describe the steps taken to make their results reproducible or verifiable.
Depending on the contribution, reproducibility can be accomplished in various ways. For example, if the contribution is a novel architecture, describing the architecture fully might suffice, or if the contribution is specific model and empirical evaluation, it may be necessary to either make it possible for others to replicate the model with the same dataset, or provide access the model. In general. releasing code and data is often one good way to accomplish this, but reproducibility can also be provided via detailed instructions for how to replicate the results, access a hosted model (e.g., in the case of a large language model), releasing of a model checkpoint, or other means that are appropriate to the research performed.
While NeurlPS does not require releasing code, the conference does require all submis- sions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example
(a) If the contribution is primarily new algorithm, the paper should make it clear how to reproduce that algorithm. (b) the contribution is primarily a new model architecture, the paper should describe the architecture clearly and fully.
(c) If the contribution is new model (e.g., a large language model), then there should either be way to access this model for reproducing the results or way to reproduce the model (e.g., with an open-source dataset instructions for how to construct the dataset).
(d) We recognize that reproducibility may be tricky in some cases, in which case authors are welcome describe the particular way they provide for reproducibility. In the case of closed-source models, it may be that access to the model is limited in some way (e.g., registered users), but it should be possible for other researchers to have some path to reproducing or verifying the results.
5. Open access to data and code
Question: Does the paper provide open access to the data and code, with sufficient instruc- tions to faithfully reproduce the main experimental results, as described in supplemental material?
Answer: [Yes]
Justification:
Most code in supplement.
Guidelines:
The answer NA means that paper does not include experiments requiring code.
Please see the NeurIPS code and data submission guidelines (nttps://nips.co
for more details.
While we encourage the release of code and data, we understand that this might not be possible, so "No" an acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (e.g., for new open-source benchmark).
The instructions should contain the exact command and environment needed to run to reproduce the results. See the NeurlPS code and data submission guidelines (https
/ /ttppp.cc/pbllccguideeee http://oeeubiiiiiioooooooooooooo
for more details.
The authors should provide instructions on data access and preparation, including how to access the raw data, preprocessed data, intermediate data, and generated data, etc.
The authors should provide scripts to reproduce all experimental results for the new proposed method and baselines. only subset of experiments are reproducible, they should state which ones are omitted from the script and why.
At submission time, to preserve anonymity, the authors should release anonymized versions (if applicable).
25

---

Providing as much information as possible in supplemental material (appended to the paper) recommended, but including URLs to data and code permitted.
6. Experimental Setting/Details
Does the paper specify all the training and test details (e.g., data splits, hyper- parameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?
Question:
Answer: [Yes]
Justification: in code and text.
Guidelines:
The answer NA means that the paper does not include experiments.
The experimental setting should be presented the core the paper to level of detail that is necessary to appreciate the results and make sense of them.
The full details can be provided either with the code, in appendix, supplemental material.
7. Experiment Statistical Significance
Question: Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?
Answer: [Yes]
Justification: In most cases, multiple runs and their statistics are provided in figures and text.
Guidelines:
The answer NA means that the paper does not include experiments.
The authors should answer "Yes" if the results accompanied by error bars, confi- dence intervals, or statistical significance tests, at least for the experiments that support the main claims of the paper.
The factors of variability that the error bars are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, or overall run with given experimental conditions).
The method for calculating the error bars should be explained (closed form formula, call to a library function, bootstrap, etc.)
The assumptions made should be given (e.g.. Normally distributed errors).
It should be clear whether the error bar is the standard deviation or the standard error of the mean.
OK to report 1-sigma error bars, but one should state The authors should preferably report a 2-sigma error bar than state that they have 96% CI, if the hypothesis of Normality of errors is not verified.
For asymmetric distributions, the authors should be careful not to show in tables or figures symmetric error bars that would yield results that out of range (e.g. negative error rates).
If error bars reported in tables or plots, The authors should explain in the text how they were calculated and reference the corresponding figures or tables the text.
8. Experiments Compute Resources
For each experiment, does the paper provide sufficient information on the com- puter resources (type of compute workers, memory, time of execution) needed to reproduce the experiments?
Question:
Answer: [TODO]
Justification: [TODO]
Guidelines:
The answer NA means that the paper does include experiments.
The paper should indicate the type of compute workers CPU or GPU, internal cluster, or cloud provider, including relevant memory and storage.
26

---

The paper should provide the amount of compute required for each of the individual experimental runs as well estimate the total compute.
The paper should disclose whether the full research project required more compute than the experiments reported in the paper (e.g., preliminary or failed experiments that didn't make into the paper).
9. Code Of Ethies
Question: Does the research conducted in the paper conform, in every respect, with the NeurIPS Code of Ethics htatpp///eerriss..cpppbbbe/iiiit
Answer: [Yes]
Justification: [TODO]
Guidelines:
The answer NA means that the authors have not reviewed the NeurIPS Code of Ethics.
If the authors answer No, they should explain the special circumstances that require a deviation from the Code of Ethics.
The authors should make sure to preserve anonymity (e.g.. if there special consid- eration due to laws or regulations in their jurisdiction).
10. Broader Impacts
Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?
Answer: [NA]
Justification: [TODO]
Guidelines:
The answer NA means that there no societal impact of the work performed.
If the authors answer NA or No, they should explain why their work has no societal impact why the paper does not address societal impact.
Examples of negative societal impacts include potential malicious unintended uses (e.g., disinformation, generating fake profiles, surveillance), fairness considerations (e.g., deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations.
The conference expects that many papers will be foundational research and not tied to particular applications, let alone deployments. However, if there is a direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that an improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point that a generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster.
The authors should consider possible harms that could arise when the technology is being used as intended and functioning correctly, harms that could arise when the technology being used as intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology.
If there are negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release models, providing defenses in addition attacks, mechanisms for monitoring misuse, mechanisms to monitor how system learns from feedback over time, improving the efficiency and accessibility ML).
11. Safeguards
Question: Does the paper describe safeguards that have been put in place for responsible release of data or models that have a high risk for misuse (e.g., pretrained language models, image generators, scraped datasets)?
Answer: [NA]
Justification: [TODO]
Guidelines:
The answer NA means that the paper poses such risks.
27

---

Released models that have high risk for misuse or dual use should be released with necessary safeguards to allow for controlled use of the model. for example by requiring that users adhere to usage guidelines or restrictions to access the model or implementing safety filters.
Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images.
We recognize that providing effective safeguards challenging, and many papers do not require this, but we encourage authors to take this into account and make best faith effort.
12. Licenses for existing assets
Are the creators original owners of assets (e.g., code, data, models), used in the paper, properly credited and are the license and terms of use explicitly mentioned and properly respected?
Question:
Answer: [NA]
Justification: [TODO]
Guidelines:
The answer NA means that the paper does not use existing assets.
The authors should cite the original paper that produced the code package dataset.
The authors should state which version of the asset used and, if possible, include URL.
The name of the license (e.g., CC-BY 4.0) should be included for each asset.
For scraped data from a particular source (e.g., website), the copyright and terms of service of that source should be provided.
If assets are released, the license, copyright information, and terms of use in the package should be provided. For popular datasets, paperswi thcode. com datasets has curated licenses for some datasets. Their licensing guide can help determine the license of dataset.
For existing datasets that are re- re-packaged, both the original license and the license of the derived asset (if it has changed) should be provided.
If this information is not available online, the authors are encouraged to reach out to the asset's 's creators.
13. New Assets
Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?
Answer: [NA]
Justification: [TODO]
Guidelines:
The answer NA means that the paper does not release new assets.
Researchers should communicate the details of the dataset/code/model as part of their submissions via structured templates. This includes details about training, license, limitations, etc.
The paper should discuss whether and how consent obtained from people whose asset is used.
At submission time, remember to anonymize your assets (if applicable). You can either create an anonymized URL or include an anonymized zip file.
14. Crowdsourcing and Research with Human Subjects
Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given to participants and screenshots, if applicable, as well details about compensation (if any)?
Answer: [NA]
Justification: [TODO]
28

---

Guidelines:
The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.
Including this information the supplemental material is fine, but if the main contribu- tion of the paper involves human subjects, then as much detail as possible should be included the main paper.
According to the NeurIPS Code Ethics, workers involved in data collection, curation, other labor should be paid at least the minimum wage in the country of the data collector.
15. Institutional Review Board (IRB) Approvals or Equivalent for Research with Human Subjects
Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed the subjects, and whether Institutional Review Board (IRB) approvals (or an equivalent approval/review based on the requirements of your country or institution) were obtained?
Answer: [NA]
Justification: [TODO]
Guidelines:
The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.
Depending on the country in which research is conducted, IRB approval (or equivalent)
may be required for any human subjects research. If you obtained IRB approval, you should clearly state this in the paper.
We recognize that the procedures for this may vary significantly between institutions and locations, and expect authors to adhere to the NeurIPS Code of Ethics and the guidelines for their institution.
For initial submissions, do not include any information that would break anonymity (if applicable), such the institution conducting the review.
29


Reviews:
[{"review_id": "1exh1CNw8z", "reviewer": "Reviewer_fqNB", "summary": "This paper proposes a novel approach for molecular simulations using neural network reparametrization. The authors first motivate the need for this work, specifically the traditional coarse-graining (CG) methods reduce the number of degrees of freedom (DOF) to improve computational efficiency. However, they require back-mapping and force-matching steps, which can be cumbersome.\n\n The major contribution is the framework of hessian backbone that allows calculation of hessian using weighted graph Laplacian(although restricted to pairwise invariant potentials). With GNN reparameterization the experiments on synthetic coil shows good improvement in speedup, but not necessarily in reaching lower energies.  The proposed system instead of reducing DOF, allows for a flexible representation of the system. Their  neural reparametrization approach is not limited to reduce DOF but can also increase them when required.\n\nThe paper showcases the effectiveness of the method on LJ systems and protein folding simulations. Results suggest the reparametrization approach, especially using GNNs, can achieve ower energy states compared to traditional CG methods and also faster convergence.", "questions": "Questions\n1) How are the inputs to GCN network 'Z_h0' initialized? \n2) Could you clarify what is the loss function used to train the GNN network? details related to training shall be provided.\n3) Can the approach be used for disordered glassy systems, e.g. Kob-Anderson Binary LJ Model Glass which is known to have slow dynamics. Does it gets stuck in higher energy minima? It will be good to show the coarse grained approach on glassy systems.  \n4) The authors mention they use Adam optimizer with a learning rate 10−2. Can the authors show impact of learning rate, lower and high Learning rates? Does the conclusion remain same?\n\nMinor Comments\n\n1) Typo in line 143: Power of z in repulsive term\n2) Typo in line 176: spelling of 'weighted'\n3) Line 131 typo: mdoes", "limitations": "Yes", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 2, "contribution": 3, "strengths": "1) Proposed a hessian backbone approach to get slow modes using graph Laplacian\n2) Eliminates the need for force matching and back mapping by reparameterization using slow modes and GNN  \n3) Average over perturbed configurations taken to address dynamic variation in Hessian\n4) Data-free optimization: Doesn't require extensive training data, unlike traditional machine learning approaches.\n5) Code is also provided.", "weaknesses": "1) Different choices for fraction of eigenvectors in CG equations are mentioned 3x (#AminoAcids), 30%, 50%, and 70%, but the results corresponding to them are not shown, it is important as it is related to the 'epsilon' in slow mode calculations.  \n\n2) Literature review is weak, several works both recent and classical on optimization in molecular systems are missing e.g. learned optimizers[1], graph reinforcement learning [2], FIRE [3] \n\n4) Presentation issues:  what is n and d in line 85? The paper can be made more reader friendly if there is a table which has symbol/variable names and their meaning.\n\nMissing citations:\nline 56-59\n1. Traditional optimization in physics-based models, like (MD), faces unique challenges due to the\nshallow nature of these models, where physical DOF are the trainable weights. Additionally, the\ninteractions occur at multiple scales, from strong covalent bonds to weak van der Waals forces,\nleading to slow convergence in gradient-based method\n\nThe authors should cite relevant papers for the above paragraph.\n\n\n\n[1]Merchant, A., Metz, L., Schoenholz, S.S. &amp; Cubuk, E.D.. (2021). Learn2Hop: Learned Optimization on Rough Landscapes. <i>Proceedings of the 38th International Conference on Machine Learning</i>, in <i>Proceedings of Machine Learning Research</i> 139:7643-7653 Available from https://proceedings.mlr.press/v139/merchant21a.html.\n\n[2]Bihani, V., Manchanda, S., Sastry, S., Ranu, S. &amp; Krishnan, N.M.A.. (2023). StriderNet: A Graph Reinforcement Learning Approach to Optimize Atomic Structures on Rough Energy Landscapes. <i>Proceedings of the 40th International Conference on Machine Learning</i>, in <i>Proceedings of Machine Learning Research</i> 202:2431-2451 Available from https://proceedings.mlr.press/v202/bihani23a.html.\n\n[3]Bitzek, E., Koskinen, P., Gahler, F., Moseler, M., and Gumbsch, P. Structural relaxation made simple. Physical review letters, 97(17):170201, 2006.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you, we appreciate your pertinent comments. Please also check our general response above about improvements to the presentation. \n## Weaknesses: \n1. The plot was mistakenly omitted. We are adding it to the appendix. \n2. Thanks, adding the citations. \n3. n is the number of particles, d is the embedded dim, e.g. d=3 for 3D coordinates x,y,z.\n4. Missing citations: we are adding the citations for that paragraph\n\n## Questions: \n1. 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. \n2. 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. \nOur 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.  \nOur 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.\n3. 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.\n4. 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.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "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": "hB48OjftBL", "reviewer": "Reviewer_rxS5", "summary": "This paper presents a novel approach to molecular simulations using neural network reparametrization as an alternative to traditional coarse-graining methods. The key idea is to reparametrize fine-grained modes as functions of coarse-grained modes through a neural network, maintaining continuous access to fine-grained modes and eliminating the need for force-matching. The authors demonstrate improved performance on Lennard-Jones potentials and protein folding simulations compared to conventional methods.", "questions": "* What is $n_0$ in line 235?\n* Does the design of your GNN need to consider equivariance such as SE(3)? i.e. if you rotate your system, the Hessian and the Laplacian would be transformed in a predicitable way.\n* How well do you expect this method to generalize to more complex force fields or larger biomolecular systems?", "limitations": "See above.", "rating": 5, "confidence": 3, "soundness": 2, "presentation": 1, "contribution": 3, "strengths": "* Theoretical foundation: The paper provides a solid theoretical analysis of the properties of physical Hessians and how they relate to slow modes in the system.\n* Experimental results: The approach shows promising results on both synthetic systems (Lennard-Jones potentials) and protein folding simulations, demonstrating faster convergence and lower energy states in many cases.", "weaknesses": "* Unclear presentation: The paper lacks a clear problem definition and objective function. A flowchart or algorithm illustrating the coarse-graining process, GNN structure, and how it drives molecular dynamics would significantly improve clarity.\n* Insufficient explanation of DOF experiments: The paper would benefit from an algorithm or graphical illustration explaining the process of finding effective degrees of freedom.\n* Inadequate explanation of GNN graph structure: The authors could have done a better job explaining the graph used in the GNN, particularly the force constant matrix and its physical meaning.\n* The paper would benefit from a more extensive comparison to state-of-the-art coarse-graining and optimization methods. Additionally, a more thorough analysis of the computational costs and scalability of the approach would strengthen the paper.\n\nMinor Issues:\n* Typos on Line 131, 176", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "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. \n\n## Weaknesses: \n1. __Unclear presentation:__ We agree. Please see above and the attached pdf for flow-chart and algorithm. \n2. __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]\\|$.  \n3. __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. \n4. __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. \n\n## Questions: \n1. __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.\n2. __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)). \n3. __Scaling:__ \nOur 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.  \nThe 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.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "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": "WLYhP6lH5z", "reviewer": "Reviewer_ti8n", "summary": "This paper proposes an efficient approach for finding the optimal conformation in terms of the energy function. By identifying the slow modes, the proposed method can reduce the computational complexity while extracting the core movement for simulating the dynamics. Specifically, the authors observe a connection between the Laplacian and the Hessian of the targeted potential functions, leading to the coarse-graining by identifying the slow modes.", "questions": "* Is it possible to find the globally optimal structures? How much does the initial state affect the final predicted structures?\n* For the large-scale protein, how much can the proposed method reduce the computational cost?\n* Does the proposed method can significantly reduce the memory cost?\n\nTypos:\n\nline 131: mdoes -> modes\n\nline 176: is relies -> relies\n\nline 176: eighted -> weighted", "limitations": "As the authors noted, more experimental results are needed to show the effectiveness of the proposed method. Especially, I recommend comparing the computational cost and the quality of the predicted structures with the state-of-the-art structure prediction models on large-scale proteins.", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 2, "contribution": 2, "strengths": "* The proposed method for a proxy of the unweighted adjacency matrix is reasonable and sound.\n* The connection between the graph laplacian and the hessian of energy function is interesting and general.\n* The proposed method is efficient compared to the gradient descent method which is widely used for simulating the protein dynamics.", "weaknesses": "* The efficiency of this work is largely dependent on the complexity of the targeted potential function as it still requires computing the gradient from the potential function. In other words, if we have a large-scale AI model that can more efficiently predict the most stable structures in one shot such as the AlphaFold series, the proposed method might not be helpful to find the optimal conformation.\n* The optimization process needs to be clearly described. I suggest the authors provide an algorithm for how to optimize the structures including the intervention of the potential functions and the model architecture.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "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. \n\n## Weaknesses: \n1. __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. \n2. __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.\n\n## Questions: \n1. __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. \n2. __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. \n3. __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.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "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": "ZEDVH7bPrK", "reviewer": "Reviewer_qFP8", "summary": "The paper proposes a novel neural network reparameterization approach which provides a flexible alternative to traditional coarse-graining methods for molecular simulations. Unlike CG methods that strictly reduce degrees of freedom, the proposed model can dynamically adjust system complexity, potentially increasing it to simplify optimization. Fitting a neural network which maps the reduced space to the fine-grained space, the approach maintains continuous access to fine-grained modes and eliminates the need for force-matching, thereby enhancing both the efficiency and accuracy of energy minimization. The framework incorporates arbitrary neural networks like GNNs, to perform reparametrization, achieving significant improvements in molecular simulations by optimizing energy minimization and convergence speeds.", "questions": "1. The proposed approach has an edge over full-atom simulations and CG methods. Meanwhile, various approaches such as metadynamics with collective variables are commonly used to learn and simulate the dynamics of the system in a reduced space. Have the authors compared the proposed neural network reparametrization method with those approaches?\n\n1. While slow modes represent directions in the configuration space along which the energy landscape is relatively flat so that they capture low-frequency motions, it is not necessarily clear that these are the most relevant directions for the system's evolution that we care about. In other words, the slow modes do not necessarily include more \"science\". Thus, I'm wondering whether relying solely on slow modes may overlook critical aspects of the system's behavior.\n\n1. Is the reparameterized neural network generalizable to different molecular systems or molecular structures under different conditions (e.g. temperature, pressure)?\n\n1. Empirical potentials are used in this work. However, empirical potentials cannot achieve high accuracy. Have the authors tried machine learning potentials for MD simulations?", "limitations": "The authors discussed the limitation of lack of experiments for comparison against traditional CG methods. No potential negative societal impact is involved.", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 2, "contribution": 3, "strengths": "1. The neural network reparameterization approach gets rid of the force-matching and the non-unique back-mapping in traditional CG methods and shows promising potential in enhancing molecular simulations. \n\n1. The Hessian matrix provides a precise mathematical framework to identify slow modes that correspond to the most stable and significant collective motions in the system. Such an innovative strategy does not need to use an encoder-decoder for FG-CG mapping. \n\n1. Dynamically adjusting the effective DOF allows the model to increase complexity when needed, capturing essential details without a fixed reduction in resolution, thus resulting in more accurate representations of the system. By focusing on slow modes, the method can more effectively explore the energy landscape, avoiding local minima and achieving better convergence to global minima.\n\n1. The authors provide solid mathematical derivation to support their proposed approach.", "weaknesses": "1. The manuscript needs improvement: it is difficult to follow the paper as the contents are not well-organized and some crucial details are missing. \n    - The introduction and background could be reorganized. It's better to first introduce the key challenge and discuss the approach and key contributions at a high level, while the details of CG and the proposed approach could be moved to Methods. This can also remove the overlapping between the Introduction and Background.\n    - Due to the lack of a high-level introduction to the overall workflow (usually the last paragraph and Figure 1), the discussion on NN reparameterization and Hessian of potential energies is confusing at first glance because readers are not clear about their roles in the big picture. Similarly, the subsections in the Hessian part propose many theorems and corollaries but do not explicitly address their contribution to the proposed workflow. \n    - The Experiment part should stress what the experiment settings are, where we can find the results, and what conclusions we can draw. The information is not clearly stressed even though I can find some of it in the paragraphs. \n    - Some crucial information is missing, for example, the GNN training details and the loss function. Without such information, it's hard to assess the soundness of the method. \n    - Figure captions also needs improvement. For example, the caption of Figure 1 consists of many short sentences, but the key takeaway are not strengthened.\n\n1. The major contribution on machine learning of this paper is the use of a neural network to map slow modes to FG configurations. However, this neural network model is not novel and does not introduce new ML techniques or architectures. Therefore, I'm afraid that the work may not be fully aligned with the scope of the conference. Additionally, it is unclear whether a neural network is necessary for this mapping, as other simpler or more traditional methods might suffice for this task. A more thorough justification for the use of a neural network over other potential approaches is needed.\n\n1. The manuscript involves biology background knowledge especially in the experiments. However, they are not well explained, making the experimental session more difficult to follow.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "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.\n\n## Weaknesses \n__1.__ We agree that intro and background will benefit from significant reorganization. \n1. 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.\n2. We are making a new Fig. 1 (see attached pdf) to outline the methodology. \n3. 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.\n4. 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$ . \n5. We are improving figure captions. \n\n__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.\n\nFurthermore, 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.\n\n__3.__ We agree with the need for clearer bio background and will make the following short additions to aid readers:\n\n1. **Overview of Experiments**: We'll introduce two main experimental setups:\n   - **Energy Minimization on Synthetic Systems**: Forming a coil using LJ potentials to mimic molecular forces.\n   - **MD for Protein Folding**: Using the AMBER force field.\n\n2. **Details on Synthetic Coil Experiments**:\n   - **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.\n   - **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.\n\n3. **Protein Folding Using MD**:\n   - 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.\n   - 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.\n\n## Questions \n1.  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. \n2. 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. \n\n3. 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. \n4. 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.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "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:
Thank you all for very constructive comments. Some issues were raised by multiple reviewer. Here we will address the shared concerns. Below, we respond point-by-point to each review. 

## Presentation
We agree that the presentation of the paper needs significant improvement and reorganization, especially intro and background. We are doing the following: 
1. __Problem statement:__ 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. __Flow-chart and algorithm:__ We are making a new Fig. 1 (see attached pdf) to outline the methodology. We are also adding the step-by=step of the algorithm (detailed next). 
3. __Motivation for using Hessian slow-modes:__ 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. __GNN details:__ We add GNN parameter details in a table in the appendix. Our GNN consists of GCN layers with self-loops and residual connections. For the experiments in Fig 1, the GNN hidden dims are [20,10, 3], i.e. starting from $Z_h$ with 20 dim embedding, and one GCN layer with 10 hidden dims and a projection layer down to 3D. For the protein experiments we had dims [100,100,3]. We used $n/3$ slow modes $\Psi$ to get adjacency matrix $A= \Psi\Psi^T$ used in the GNN layers, which are GCN 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$ .   Here $h^l \in \mathbb{R}^{n\times d_l}$ ($n$ particles in $d_l$ hidden dims). We concatenate outputs along feature dims of all GCN layers into $h = [Z_h| h^1...]$ (dense residual connections) and pass them through a final projection along feature dims to get them to $d=3$ dimensions. 
5. __Figures:__ We are improving figure captions. We are also including missing figures on using different number of Hessian modes (30%-70%) for proteins. 
6. __More experiments:__ We are adding new figures comparing our final energies with OpenMM’s `simulation.minimizeEnergy` which is the most head-to-head comparison with our method. However, we OpenMM also uses many tricks for efficiency, such as distance cutoff for forces and may use Barnes-Hut or other space partitioning, which we haven't implemented. Thus the only fair comparison we can make right now is against direct gradient descent (no reparam). We run a __learning rate sweep__ for some protein simulations. We find the lowest loss and closest to OpenMM minimum to be consistently GNN (see attached pdf). We also find that at a similar number of iteration steps (5k) OpenMM was at much higher energies compared to the GNN. 
## Flow-chart and algorithm
Most of the reviewers asked for a clear algorithm or flow-chart of our methodology. We are including the flow-chart in the attached pdf and will add the algorithm steps as follows:
1. Compute functional Hessian of the energy function $H = \nabla \nabla \mathscr{L}$ w.r.t. Particle positions $X$. Because $X \in \mathbb{R}^{n\times d}$ ($n$ particles with $d$ features), this Hessian will have four indices as $H_{i j}^{\mu\nu} = \partial^2 \mathscr{L}/\partial X_i^\mu \partial X_j^\nu$
2. Evaluate the Hessian over a small ensemble of perturbed positions $\mathbf{Samples}(X) = \{X’ = X+\delta X\}$ and compute the __Backbone__  $\mathbf{H} \_{ij}$ $= \sum_{X’\in \mathbf{Samples}(X)}\sum_{\mu\nu} H_{ij}^{\mu\nu}(X’)^2$(eq 7). 
3. Compute or approximate $k$ eigenvectors $\Psi$ of $\mathbf{H}$ with smallest magnitude eigenvalues. 
4. Perform neural reparametrization: express $X$ using a neural ansatz $X = \mathrm{NN}(\Psi,\theta)$ where $\mathrm{NN}$ uses the Hessian modes $\Psi$ and has trainable parameters $\theta$. 
5. Optimize the same los $\mathscr{L}(X)= \mathscr{L}(\mathrm{NN}(\Psi,\theta))$ over NN parameters $\theta$ instead of the original $X$.

The procedure is quite simple and could be summarized in the following pythonic pseudocode using pytorch (actual implementation slightly different for efficiency): 

```python
# Hessian backbone
H = functional.hessian(Loss)
H_samples = tensor([H(x) for x in samples])
H_backbone = H_sampels.norm((2,4)).sum(0)
H_lap = Laplacian(H_backbone)

# slow modes
eig_vals, Psi = eigh(H_lap)
Psi_slow = Psi[:, argsort(abs(eig_val))[:k]]

# reparametrization
A = Psi_slow @ Psi_slow.T
gnn = GNN(A, hidden_dims = [h,h,3])
Z = random.normal(n,h).requires_grad()
X = gnn(Z)

# optimization 
optimizer = Adam(parameters = [Z]+list(gnn.parameters()) 

for i in range(steps):
	optimizer.zero_grads()
	loss = Loss(X)
	loss.backward()
	optimizer.step()

```

Meta-review: Accelerating molecular simulations is an important problem in several scientific domains. In this work, authors propose a novel approach to address this challenge through neural reparametrization. All the reviewers have appreciated the proposed framework and its relevance. However, they have also raised concerns regarding the clarity of presentation, setting up the problem, improved comparison with the state of the art and more extensive literature review. Authors have satisfactorily addressed these comments. It is recommended that all these changes, along with the modifications to reflect the review responses, are implemented in the camera ready version of the manuscript.
Decision: Accept (poster)