Neural Network Reparametrization for Accelerated Optimization in Molecular Simulations

Authors: nima  dehmamy
jeet  mohapatra
subhro  das
tommi  jaakkola
mit  csail
Venue: NeurIPS_cc_2024_Conference
Published: 2024-05-15 07:19:41

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.

Main paper: 
 background: 
 Traditional optimization in physics-based models, like (MD), faces unique challenges due to the shallow nature of these models, where physical DOF are 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., 2011; Kingma & Ba, 2014) or 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 nondiagonal Hessians in physics problems, and CG can be restrictive and require intensive back-mapping and force-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 representation, which can simplify the optimization process. It 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 to reducing DOF but can also increase them when beneficial, offering an adaptive solution to the specific needs of a 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-matching and back-mapping, our neural reparametrization approach defines modes based on the spectrum of a canonical Hessian, directly incorporating these into the neural network'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 or costly to produce in molecular and material design. By not relying on training data, our approach provides a robust framework for tackling complex optimization problems, from molecular dynamics to protein folding, with improved efficiency and without the constraints of data availability.Let X ∈ X ≃ R n×d 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 a steep non-convex function, computations can be slow. Traditional coarse-graining (CG) maps X to a reduced space of CG variables, Z ≃ R k×d , where k ≪ n. Implementing dynamics using CG modes requires determining the inter-mode forces ("force-matching") and how to revert to X ("back-mapping").Force-matching. The fine-grained (FG) energy function,EQUATIONThe process of finding L CG is called force-matching, traditionally solved analytically but increasingly with machine learning for enhanced accuracy Jin et al. (2022) ; Majewski et al. (2023) .Back-mapping. The map Z ∼ R k×d is not unique, often resulting in multiple possible X for a given Z. Back-mapping typically involves sampling or optimization to find physically plausible X configurations, avoiding scenarios like overlapping atoms or high energies. This can be complex when many X map to the same Z, with current methods ranging from geometric reconstruction Lombardi et al. (2016) to refinement with molecular dynamics Badaczewska-Dawid et al. (2020); Roel-Touris & Bonvin (2020) and data-driven approaches Yang & Gómez-Bombarelli (2023) ; Wang et al. (2022) .Instead of traditional CG, which reduces DOF through a mapping to a reduced space, our approach reparametrizes the DOF X as a function of CG-like modes. This reparametrization, given by X = ρ(Z), where ρ : Z → X , offers a flexible, reversible mapping that inherently includes benefits such as direct access to fine-grained modes and elimination of force-matching and back-mapping needs:Reparametrization: X = ρ(Z), ρ : Z → X (2)1. Flexible parametrization: Leveraging neural overparametrization and architecture design.2. Direct access to fine-grained modes: X = ρ(Z) avoids the need for back-mapping.The success of optimization in molecular systems is 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 eigendecomposition of the Hessian H = ∇∇L . 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 a fundamental challenge: To maintain numerical stability, the learning rate must be small enough to handle the fastest modes, but this makes the slow modes evolve at a 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-graining 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 a 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. 
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, r ij ≡ X i -X j (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, v 2 = ∥v∥ η ≡ v T ηv, where η may be the Euclidean metric η = diag(1, 1, 1) or the Minkowski metric η = diag(-1, 1, 1, 1). For example, the Euclidean norm v T v in d dimensions is invariant under rotations v → gv, where g ∈ SO(d).Energy function structure. Let r denote the matrix of distances with r ij = ∥r ij ∥ η . Any function of r ij is invariant under symmetries that keep ∥ • ∥ η invariant. Assuming additivity, the energy function can be written as:EQUATIONwheref ij (z) = f ji (z).For example, the Coulomb potential between particles i and j with charges q i and q j respectively, is given byf ij (z) = kq i q j /z. The Lennard-Jones potential f ij (z) = A ij /z 1 2 - B ij /z 6 inmolecular systems is also of this form.The Hessian of potentials of the form equation 4 has the special property that it is 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 a graph Laplacian. Recall the Laplacian of an undirected graph with adjacency matrix A is defined as L = Lap(A) = D -A, where D is the degree matrix with elementsD ij = δ ij k A ik .The components of Laplacian can also be written asL ij = k A ik (δ ij -δ jk ).We show that the Hessian of L in equation 4 is a Laplacian. Let ∂ i ≡ ∂/∂X i and let r = ηr/r be the dual unit vector of r. First, observe that ∂ i r jk = rjk (δ ij -δ ik ) where rjk is the unit vector of r jk and δ ij is the Kronecker delta (1 if i = j, 0 otherwise). Let Hes[g] denote the Hessian of a function g. We find that (app. E)Hes[L ](X) ij = ∂ i ∂ j L (X) = k (δ ij -δ jk )H ik (X) = Lap(H) ij (5)where H ik (X) = Hes[f ik ](r ik ). Note that H has four indices, with components H µν ij , having two particle indices i, j and two spatial indices µ, ν. Thus, for every pair of spatial indices µ, ν, 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 use 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 A of a graph, let Â and L be the "unweighted" adjacency and Laplacian matrices, where Âij = 1 if A ij ̸ = 0 and zero otherwise. It follows that the null spaces of L and L are shared: Theorem 2.1 (Null Space of the Laplacian). Let Null[M ] denote the null space of a symmetric real matrix M . The null space of the unweighted Laplacian L is contained within the null space of the weighted Laplacian L, i.e., Null[ L] ⊆ Null [L] .Sketch of proof. For any vector, v ∈ R n , v T Lap(A)v = ij A ij (v i -v j ) 2 . Since Âij = 0 yields A ij = 0, but not necessarily vice versa, null vectors of Null[ L] ⊆ Null[L]. See appendix for full proof.Definition 2.1 (Slow manifold). Let L be a graph Laplacian (undirected, weighted or unweighted), with spectral expansion L = n i=1 λ i ψ i ψ T i . Let ε ≪ 1 and λ max = max{λ i } be the largest eigenvalue of L. We define the slow manifold asSlow ε [L] = Span ψ i |λ i | < ε 2 λ 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 L = Lap(A) and L = Lap( Â). Then Slow ε [L] overlaps with Slow ε [ L] up to O(ε 2 ) 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 = 1 2 CW C T and L = 1 2 CC T . For a random configuration X the edge weights W will be random, as they arising from derivatives of f ij (r ij ) in equation 20 (unless f ij is quadratic which makes W constant). Then, using the assumption of randomness on the weights W , we can show the slow modes of L are perturbations of order ε 2 on slow modes of L . See Appendix D for proof.Implications for Coarse-Graining. The identification of slow modes in the Hessian is crucial for coarse-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-Graining 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.The slow modes of the Hessian Hes[L ](X) = Lap(H(X)) can dynamically change during optimization. To ensure the robustness of these modes, we need a proxy for the unweighted adjacency matrix Â ≡ H. To this end, we aggregate Hessians from perturbed configurations Samples(X) = {X ′ = X + δX}: Lower means better for both energy and time. In Bond+LJ (left), a quadratic potential i (r ii+1 -1) 2 attracts nodes i and i + 1. A weak LJ potential attracts nodes i and i + 10 to form loops. In LJ loop (right) both the backbone i, i + 1 and the 10x weaker loop are LJ. 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, j ∈ Z n are the particle indices and the Frobenius norm ∥H ij ∥ 2 = µ,ν (H µν ij ) 2 sums over the feature indices (note that X µ i has a particle index i and a feature index µ ∈ {1, . . . d}). Then, we extract the slow modes of the backbone, by doing a spectral expansionBackbone: H ij = X ′ ∈Sample(X) ∥H ij (X ′ )∥ 2 (7)H = i λ i ψ i ψ T i and picking ψ i with |λ i | < ε 2 max j [λ j ],for some small ε < 1. The intuition behind equation 7 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, we 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 H(X ′ ) up to O(ε 2 ) errors. 
simplified energy computation:: 
 The energy for CG-like modes is L CG (Z) = L (ρ(Z)).While this method can be computationally intensive as L CG (Z) is computed using X, the efficiency gains in optimization speed and depth of energy minimization can offset the costs. A nxn = Linear Reparametrization = Z X d k n d Graph Adjacency 𝛹 T GNN Reparametrization Z X d h n = n d GCN 𝜎(AZW) GCN 𝜎(Ah 1 W) Feat Proj MLP 𝛹 T 𝛹 slow slow slow Input: Loss L, config X0, modes k Output: Optimized X /* 1. Compute Hessian backbone */ H ← ∇∇L(X0) Samples ← {X0 + δX | δX ∼ N (0, σ)} // Perturbed samples Hij ← X ′ µ,ν (H µν ij (X ′ )) 2 //← GNN(A, dims = [h, h, 3]) Z ∼ N (0, 1) n×h /* 4. Optimize */ while not converged do X ← gnn(Z) θ ← θ -η∇ θ L(X)// θ: Z and GNN weights end return X Neural Architectures for Reparametrization The reparametrization function ρ can range from simple linear projections to complex neural networks. Initially, we employ a linear projection onto identified slow modes:EQUATIONMore generally, ρ 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 use a "Hessian backbone" as a graph, which acts as a weighted graph adjacency matrix for a GNN. In our experiments, we observe this GNN to have significant advantages over the direct as well as 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 a large class of energy functions important in molecular systems.We apply our method to protein folding using classical MD forces. Settings: We use gradient descent to minimize L (X). All experiments (both CG and baseline) use the Adam optimizer with a learning rate 10 -2 and early stopping with |δL | = 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 baseline. 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 1UNC proteins, results in a 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. protein as it goes from an unfolded to a 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 2fs 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 × (#AminoAcids), 30%, 50%, and 70%. We use a two stage process. First, we use CG as in equation 3 X = ρ(Z) = Z T Ψ Slow and minimize L CG (Z) = L (ρ(Z)) over Z. After convergence to X 0 = ρ(Z 0 ), we add δX to X 0 and optimize the fine-grained δX, starting with δX = 0.GNN model: We use a GNN consisting of a graph convolution (GCN) layer with self-loops and one node-wise MLP layer, projecting the GNN output to 3D to get particle positions. The GCN takes Z h0 ∈ R n×h0 as input, with h 0 > 3 and has weights W G ∈ R h0×h1 . Then, GCN output gets a Tanh activation and is passed to the MLP layer to yield X. The CG parameters in this case are Z h , W G and the weights and biases of the MLP.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 3 shows many experiments using GD, CG, and GNN. In the quadratic Bond+LJ case, GNN yields a good GNN1 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 a worse (higher) energy than our GD and GNN runs, but it takes a fraction of the time (it has many efficiency tricks that our code doesn't have). The dashed line shows the energy achieved by OpenMM after 10k steps. As we see, some of our GNN models reach energies close to the 10k steps of openMM in a fraction of the steps. All experiments show the best energy among three runs. c shows the effect of initialization on the GD runs. We do find the protein converges to significantly different conformations based on the init. speedup, while CG yields better energies. The benefit of CG and GNN become more apparent in the harder pure LJ problem, where GD fails to 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 are not modeled and appear as hydrogen-bonding and hydrophobicity terms; app. A). In protein folding our energy function consists of five potential energies: bond length E bond , bond angles E angle , van der Waals E vdW , hydrophobic E hp and hydrogen bonding E H Ceci et al. (2007) . Figure 7 shows an example of these coupling matrices for the Enkephalin (1PLW) protein. To evaluate the effect of our CG model, we run experiments on four small proteins: Chignolin (5AWL), 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. (2017) , and it comprehensively models the following interactions:• Bond lengths E bond and bond angles E angle• Non-bonded interactions including van der Waals E vdW and electrostatic E elec forcesWe utilize the functional forms and parameters specified in the AMBER force field:E bond = bonds k bond (r -r 0 ) 2 E angle = angles k angle (θ -θ 0 ) 2 (8) E torsion = torsions V n [1 + cos(nω -γ)] E vdW = i<j A ij r 12 ij - B ij r 6 ij (9) E elec = i<j q i q j 4πϵ 0 ϵ r r ij (10)Here, r and θ represent the bond lengths and angles, respectively, with r 0 and θ 0 as their equilibrium values. The torsional term E torsion includes a sum over all torsion angles ω, with periodicity n, amplitude V n , and phase γ. The Lennard-Jones potential in E vdW is characterized by parameters A ij and B ij , and E elec is calculated using the Coulombic potential with partial charges q i , q j and the relative permittivity ϵ r .In this simulation, we exclude the modeling of solvent effects entirely, focusing solely on 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:L(X) = E bond + E angle + E torsion + E vdW + E elec (11)Figure 7 shows the interaction matrices for the Enkephalin (1PLW) 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 our approach.Protein results: Denoting the final energy and run time of the GNN model by E GN N and t GN N , and baseline by E 0 and t 0 , we compute the energy improvement factor δ Ê = E 0 /E GN N and speedup factor δ t = t 0 /t GN N , to plot different proteins together. Figure 4a shows the mean of δ Ê vs δ t over the 10 runs for GNN the model with hidden dimensions 300 (error bars are 1 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 use RMSD, comparing the final layouts to the PDB structure.Figure 5 shows the RMSD value evolution using different methods. While, in most cases, OpenMM reaches a deeper RMSD value, our models could serve as a good initializer for accelerating molecular dynamics. To evaluate the robustness of these results, we ran sweeps over the learning rate, varied the number of GNN layers (one or two layers), and varied the initialization.Figure 6 shows 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 in 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 back-mapping. However, more experiments are needed to compare it against traditional CG methods. In fact, using ML to learn force-matching might provide further advantage by removing the need to evaluate L CG (Z) = L (X) via the finegrained 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.In protein folding our energy function consists of five potential energies for: bond length E bond , bond angles E angle , Van der Waals E vdW , hydrophobic E hp and hydrogen bonding E H Ceci et al. (2007) . Note that we are ignoring the solvent (e.g. water) and writing using potentials, or force fields. To calculate the force field, we use distance, r, and angle-based, Θ, potentials. For each amino acid, we use the rdkit Landrum et al. (2020) package to acquire bond length, r 0 , and bond angle, θ 0 (every triplet of atoms defining the bond), information that we use to define quadratic energies E bond and E angle . We use Lennard-Jones (LJ) potentials, V p,q (r) = r -p -r -q , to approximate E vdW between all pairs of atoms, E H between atoms prone to form a hydrogen bond (certain H and O, in our case), E hp between atoms in hydrophobic residues, yieldingEQUATION)Here the coupling matrix [σ vdW ] ij = a i + a j where a i is the vdW radius of atom i. We note that our choices for ϵ H , ϵ vdW , ϵ hp and k bond , k angle , can be a source of error. Additionally, we "softened" the LJ potential to V p,q = 1/(r p + ζ) -1/(r q + ζ) with ζ = 0.65, which is large and significantly reduces the penalty for overlapping atoms and may reduce accuracy.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 we use in the energy optimization.Let X ∈ X ≃ R n×d be a set of degrees of freedom (e.g. particle positions, bond angles, etc.) and let L : X → R 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 its continuous variant, gradient flow (GF)EQUATION)where ε is the matrix of learning rates (LR). In simple GD where ε = cI is a single constant times identity, GD evolves at different rates in different directions, with some being much slower than others. At a given X, these "slow modes" are the eigenvectors of the Hessian H(X) = ∇∇L (X) with eigenvalues closest to zero, as we review below. We will first define fast and slow modes in the simple quadratic case and then generalize them to non-convex cases in the next section. Fast and slow modes for quadratic Loss. Consider the case where L (X) = 1 2 Tr X T HX . Here H is a Hermitian matrix and the Hessian of L , with a spectral expansion given by H = i λ i ψ i ψ T i , λ i ∈ R and ψ i ∈ R n . In this basis we have X(t) = i c i (t)ψ i with c i : R → R d . Projecting equation 13 onto one of the eigenmodes we getEQUATIONwhere we assumed dψ i /dt = 0. From equation 14 we see that the decay/growth rate along mode ψ i is |ελ i |. Hence, modes with λ i close to zero are the "slow modes", evolving very slowly, and large |λ i | defines the "fast modes". Since c i (t) = c i (0) exp[-t/τ i ] with time scale τ i = 1/(ελ i ), the fast modes evolve exponentially faster than slow modes. This disparity in the rates results in slow convergence, because the fast modes force us to choose smaller ε to avoid numerical instabilities. Two potential ways to fix the issue with disparity in time scales are: 1) make rates isotropic (second-order methods and adaptive gradients); 2) mode truncation or compression (CG). We will briefly review the former here.Adaptive gradient and second-order methods. Newton's method uses ε = ηH(X) -1 which makes GD isotropic along all modes, but it is expensive (O((3n) 3 ) in our case). Quasi-Newton methods, e.g. BFGS Fletcher (2013) , approximate H -1 iteratively, but are generally also slow. Another, more efficient approach is adaptive gradient methods, such as AdaGrad Duchi et al. (2011) and Adam Kingma & Ba (2014) which approximate H by g t g T t + η where g t = k i=1 γ i ∇L (X(t -i)) is some discounted average over past gradients and η a small constant. For efficiency, in practice we 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 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(t) and let δt be a small time interval. We are looking for modes which are almost stationary over δt. To identify these modes, we can for instance find perturbations δX which would have almost zero dynamics. concretely we find the dynamics of X + δX asd dt (X + δX) = -ε∇L (X + δX) ≈ -ε∇L (X) -εHδX + O(δX) 2 (15)meaning, a small δX adds εHδX to the dynamics.Thus if δX is a zero mode of the Hessian, HδX = 0, it won't change the dynamics of X. To define slow modes, we can slightly relax this and look for normalized modes ψ = δX/∥δX∥ whose associated time scale is much longer than a desired time scale δtτ = |εψ T Hψ| = |ελ| ≫ δt (16)which just means that we need to find the approximate zero modes of the Hessian H(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 ∆t ≫ δt. The essence of our algorithm is to ignore fast modes and project and evolve the system on the "slow manifold" spanned by the slow modes of the Hessian. However, the main challenge is 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.Invariant potentials. In systems of interacting particles in physics, most of the leading interactions are pairwise and involve relative features, r ij ≡ X i -X j (distance vector, relative angle, etc). Moreover, they are often invariant under certain global symmetries, such as Euclidean symmetries (translation and rotation) or Lorentz symmetry (relativistic particles). These symmetries keep some 2-norm of vectors, v 2 = ∥v∥ η ≡ v T ηv invariant. Here η may be the Euclidean metric η = diag(1, 1, 1) or the Minkowski metric η = diag(-1, 1, 1, 1) for relativistic problems, etc. For example, the Euclidean norm v T v in d dimensions is invariant under rotations v → gv, where g ∈ SO(d), and the Minkowski norm is invariant under the Lorentz group SO(1, d -1).Let r denote the matrix of distances with r ij = ∥r ij ∥ η . Any function of r ij is invariant under symmetries that keep ∥ • ∥ η invariant. A general invariant energy function can combine r ij for different i, j in arbitrary ways. Usually in physical systems each pair contributes an additive term in to the total energy. Assuming additivity, the energy has a formEQUATIONwhere f ij (z) = f ji (z) (symmetric under i ↔ j). For example, when particle i has electric charge q i , the Coulomb potential between i, j can be written as in equation 17 using f ij (z) = kq i q j /z. 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 van der Waals:EQUATION)Here σ ij = a i + a j , where a i 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 in finding its slow modes.The Hessian of potentials of the form equation 17 has the special property that it is 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 canonical slow modes.Hessian as a graph Laplacian. Let ∂ i ≡ ∂/∂X i and let r = ηr/r be the dual unit vector of r. First, observe that ∂ i r jk = rjk (δ ij -δ ik ) where rjk is the unit vector of r jk and δ ij is the Kronecker delta (1 if i = j, 0 otherwise). Let Hes[g] denote the Hessian of a function g. We find that (app. E)EQUATIONwhere H ik (X) = Hes[f ik ](r ik ) and given byEQUATIONNote that H has four indices, with components H µν ij , having two particle indices i, j and two spatial indices µ, ν. Recall the Laplacian of an undirected graph with adjacency matrix A is defined as L = Lap(A) = D -A, where D is the degree matrix with elements D ij = δ ij k A ik . The components of Laplacian can also be written asL ij = k A ik (δ ij -δ jk ). Thus, we see that the Hessian of L is indeed the Laplacian of H Hes[L ](X) ij = k (δ ij -δ jk )H ik = Lap(H) ij (21)where for every pair of spatial indices the Hessian is a 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.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 + δX with random δX ∼ N (0, T ) to get an ensemble of Hessians H = {H(X + δX)} and find the overlap of the slow modes of the Hessians in H. However, this is expensive, roughly O(mkn 2 ) for m = |H| and k slow modes. We cannot recompute the Hessian slow modes often. We also want a method which is 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(X ′ ). We then compute a backbone from these Hessians of the form Backbone:H ij = X ′ ∈Sample(X) ∥H ij (X ′ )∥ 2 (22)Here i, j ∈ Z n are the particle indices and the Frobenius norm ∥H ij ∥ 2 = µ,ν (H µν ij ) 2 sums over the feature indices (note that X µ i has a particle index i and a feature index µ ∈ {1, . . . d}). Then, we extract the slow modes of the backbone, by doing a spectral expansion H = i λ i ψ i ψ T i and pickingψ i with |λ i | < ε 2 max j [λ j ],for some small ε < 1.The intuition behind equation 22 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, 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 between the modes of the backbone H and the actual Hessians H(X ′ ). Importantly, entries in H(X ′ ) 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, as 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(ε 2 ) errors.Definition D.1 (weighted graph). Let Ĝ = (V, E) be a graph with vertices V = Z n , edges E ⊆ V × V. Let Â ∈ R n×n denote the adjacency matrix Âij = 1 if (i, j) ∈ E and 0 otherwise. We denote a weighted graph as G = (V, E, W) where W : E → R are the weights of the edges. Let A denote the adjacency matrix of G, where A ij = W(i, j) or zero if (i, j) ̸ ∈ 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 D ij = δ ik k A ik . Definition D.2 (Slow manifold). Let L be a graph Laplacian (undirected, weighted or unweighted), with spectral expansion L = n i=1 λ i ψ i ψ T i . Let ε ≪ 1 and λ max = max{λ i } be the largest eigenvalue of L. We define the slow manifold as 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 × E → {±1}, where for any edge e = (i → j) ∈ E, C i,e = -1 and C j,e = 1, 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 w e = A ij , for 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 = 1 2 CW C T (proof in app. E.1).EQUATIONBecause G and Ĝ share the same vertices and edges, their incidence matrix C is the same. From Lemma D.2, L = 1 2 CW C T and L = 1 2 CC T as Ĝ is unweighted. Using SVD, C = U SV T and defining R = U S/ √ 2 and Q = V T W V , we haveEQUATIONNote that for a random configuration X the edge weights W will be random, as they arising from derivatives of f ij (r ij ) in equation 20 (unless f ij 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 a Q = V T W V is somewhere between the distribution of W (for sparse graphs with |E| ∼ O(|V|)) and a Wigner Semi-circle (for dense graphs with |E| ∼ O(|V| 2 )). See appendix E.2 for more discussion. We also assume Q has no 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.1. For details, refer to appendix E.4. From the SVD, C = U SV T , the slow subspace isEQUATIONNormalize Ŝ = S/ max[S] and make them all positive (e.g. absorb their sign into U ). For some ε < 1 sort the SV such that Ŝ = diag(S ε , S 1 ) where the diagonal matrices S ε < ε and S 1 ≥ ε. Now, the problem of finding Slow ε [L] becomes finding eigenvectors of the matrix M = ŜQ ŜT with eigenvalues O(ε 2 ). Using S ε ∼ O(ε) and S 1 ∼ O(1), we can pull factors of ε out from M and write it asEQUATION)whereε2 ≡ ε 2 n A /n C is rescaled so that the random matrices Â ∈ R n A ×n A and C ∈ R n C ×n Chave a similar range of eigenvalues. Next, using a perturbative ansatz for eigenvectors ψ ′ = ψ + εδψ and eigenvaluesλ ′ = λ + εδλ, we solve M ψ ′ = λ ′ ψ ′ up to O(ε 2 ) corrections.To find slow modes for L we start from ψ ∈ Slow ε [ L] . Specifically, we start with an eigenvector ψ A of Â and concatenate it with zeros to get ψ = (ψ A , 0). We have M 0 ψ = λψ with λ = ε2 λ A . Using first-order perturbation theory, we find the corrections δλ to the eigenvalues and eigenvectors to beEQUATIONPutting all together we find the slow eigenvector ψ ′ = ψ + εδψ up to order O(ε 2 ) to beEQUATIONmeaning to first order in ε the corrections to eigenvalues of slow modes vanishes. This is desired because the slow mode eigenvalues are O(ε 2 ). We also observe that slow modes of L are mostly confined to Slow ε [ L] and only get O(ε) contributions from the fast subspace of L.As a 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{v|v ∈ R n , M v = 0} denote the null space of a matrix M ∈ R n×n . The null space of the Laplacian L (unweighted) is contained in the null space of Laplacian L (weighted), meaning Null[ L] ⊆ Null[L]. Lemma D.4. Null[ L] = Null[R T ] Proof. ∀v ∈ Null[ L], 0 = v T Lv = ∥R T v∥ 2 and ∀v ∈ Null[R T ], Lv = RR T v = 0. Proof of proposition D.3. ∀v ∈ Null[ L], Lv = RQR T v = 0 hence, Null[ L] ⊆ Null[L].Note that Null[ L] and ⊆ Null[L] 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 our method for coarse-graining using a set of canonical slow modes.We want to Compute the Hessian of equation 17, L (X) = ij (r ij ). Let r = ηr/r be the dual unit vector of r. First, note thatEQUATIONThen the gradient becomesEQUATIONwhere we used rjk = -r kj to show both terms in (δ ij -δ ik ) yield the same output. Finally, the Hessian becomesEQUATIONThis is because the components of Laplacian can be be writtenEQUATIONwill also have random Gaussian entries. When m = n, V defines the eigenbasis of Q and W ee 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 a distribution similar to W ee . For dense graphs, where m ∼ O(n 2 ), every entry of Q will involve a weighted sum over multiple W ee . 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-circle law. In both cases (sparse and dense graphs) the spectrum of Q has a finite variance and sits somewhere between a Gaussian and a semi-circle.We want to know how much the slow modes of weighted and unweighted graphs to overlaps. With the spectral expansion L = i λ i ψ i ψ T i Define the slow subspace as in equation 23EQUATIONwhere λ max ( L) = max{λ i } = max ψ [ψ T Lψ/∥ψ∥ 2 ] is the largest eigenvalue of L and ε ≪ 1. In terms of the singular values (SV) of the incidence matrix C = U SV T , the slow subspace becomesEQUATIONWe will show that the slow modes in weighted L = CW C T are perturbations to the slow modes of L. DefineM = SV T W V S T = SQS T (41) Normalize Ŝ = S/ max[S]. Break the space down to the slow and fast subspaces, based on whether Ŝii < ε or not. First, since L is positive semi-definite, we can make all S ii ≥ 0. Let Ŝ = S/ max S.We sort the dimensions in Ŝ to have the small SVs appear first. Denote the block in Ŝ where S ii < ε by S ε . We haveEQUATIONWe know the null space of L, where S ii = 0, is shared with L. First, we remove the null space from L and L, calling the remainder L 0 and L0 and the remaining SVs Ŝ. Then in this remainder subspace we need to find parts which are O(ε). We sort the dimensions in Ŝ to have the small SVs appear first. We denote the block in Ŝ where S 2 ii < ε max[S 2 ] by S ε . We haveEQUATIONBecause S ε is O(ε) and S 1 is O(1), we will factor out the factors of ε from blocks in M and writeEQUATIONHere A and C are random matrices built from their corresponding blocks in Q and sandwiched between S ε /ε (for A), and S 1 (for C), which have O(1) values. The spectrum of Q has a 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 asQ = ΨΛΨ T , Λ = diag(λ i ) n i=1 , Ψ = [ψ i ] n i=1 .(45) This is because when Q ij ∼ N (0, σ) we have (ignoring Bessel's correction for k ≫ 1).EQUATIONwhere we assumed Tr{Q}/n ≈ mean(Q) = 0. Since a block Q k of size k is k 2 entries sampled from the same distribution as Q, we expectEQUATION• 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 a 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 a 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 to 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 to 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 NeurIPS does not require releasing code, the conference does require all submissions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. (d) We recognize that reproducibility may be tricky in some cases, in which case authors are welcome to 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., to registered users), but it should be possible for other researchers to have some path to reproducing or verifying the results. 
open access to data and code: 
 Question: Does the paper provide open access to the data and code, with sufficient instructions 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. • Providing as much information as possible in supplemental material (appended to the paper) is recommended, but including URLs to data and code is permitted. 
experimental setting/details: 
 Question: Does the paper specify all the training and test details (e.g., data splits, hyperparameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?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 in the core of the paper to a 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, or as supplemental material. 
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 are accompanied by error bars, confidence 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. • It is OK to report 1-sigma error bars, but one should state it. The authors should preferably report a 2-sigma error bar than state that they have a 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 are out of range (e.g. negative error rates). • If error bars are reported in tables or plots, The authors should explain in the text how they were calculated and reference the corresponding figures or tables in the text. 
experiments compute resources: 
 Question: For each experiment, does the paper provide sufficient information on the computer resources (type of compute workers, memory, time of execution) needed to reproduce the experiments?Answer: [TODO]Justification: [TODO]Guidelines:• The answer NA means that the paper does not include experiments.• The paper should indicate the type of compute workers CPU or GPU, internal cluster, or cloud provider, including relevant memory and storage.• The paper should provide the amount of compute required for each of the individual experimental runs as well as estimate the total compute. Guidelines:• The answer NA means that there is no societal impact of the work performed.• If the authors answer NA or No, they should explain why their work has no societal impact or why the paper does not address societal impact. • Examples of negative societal impacts include potential malicious or 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 out 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 is 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 of models, providing defenses in addition to attacks, mechanisms for monitoring misuse, mechanisms to monitor how a system learns from feedback over time, improving the efficiency and accessibility of ML). 
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, or scraped datasets)? Answer: [NA] Justification: [TODO] Guidelines:• The answer NA means that the paper poses no such risks.• Released models that have a 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 is challenging, and many papers do not require this, but we encourage authors to take this into account and make a best faith effort.12. Licenses for existing assets Question: Are the creators or 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?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 or dataset.• The authors should state which version of the asset is used and, if possible, include a 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, paperswithcode.com/datasets has curated licenses for some datasets. Their licensing guide can help determine the license of a dataset. • For existing datasets that are 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 creators. 
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 was 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. 
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 as details about compensation (if any)?Answer: [NA] Justification: [TODO] Guidelines:• The answer NA means that the paper does not involve crowdsourcing nor research with human subjects. • Including this information in the supplemental material is fine, but if the main contribution of the paper involves human subjects, then as much detail as possible should be included in the main paper. • We recognize that the procedures for this may vary significantly between institutions and locations, and we 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 as the institution conducting the review.Scientific simulations, particularly in molecular dynamics (MD), face fundamental challenges in finding optimal configurations. The energy landscapes of 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 gradientbased 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 a hierarchy of energy scales that is challenging to optimize simultaneously. 
: 
 To address these challenges, coarse-graining (CG) methods have emerged as a 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-mapping (returning to the original DOF) and force-matching (finding the forces experienced by CG modes) (Jin et 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.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-grained (FG) modes as functions of CG modes. This reparametrization concept, similar to techniques such as Deep Image Priors (Ulyanov et 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 × d degrees of freedom (n particles in d dimensions), our neural representation can use a higher-dimensional 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 a structure informed by 'slow modes'-inherently stable collective modes identified through spectral analysis of the system's dynamics. We show that these modes arise naturally from the structure of physical Hessians, which are Laplacian matrices over particle indices for a 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-Jones 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 outperform 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. 
 
 Appendix:  


Decision: 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.

Reviews: review: paper_summary: paper_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.

 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.

The 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.
strength_weakness: strength_weakness: 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.  

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

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

Missing citations:
line 56-59
1. Traditional optimization in physics-based models, like (MD), faces unique challenges due to the
shallow nature of these models, where physical DOF are 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 method

The authors should cite relevant papers for the above paragraph.



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

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

[3]Bitzek, E., Koskinen, P., Gahler, F., Moseler, M., and Gumbsch, P. Structural relaxation made simple. Physical review letters, 97(17):170201, 2006.
questions: questions: Questions
1) How are the inputs to GCN network 'Z_h0' initialized? 
2) Could you clarify what is the loss function used to train the GNN network? details related to training shall be provided.
3) 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.  
4) 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?

Minor Comments

1) Typo in line 143: Power of z in repulsive term
2) Typo in line 176: spelling of 'weighted'
3) Line 131 typo: mdoes
limitations: limitations: Yes
paper_summary: paper_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.
strength_weakness: strength_weakness: * 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.
* Insufficient explanation of DOF experiments: The paper would benefit from an algorithm or graphical illustration explaining the process of finding effective degrees of freedom.
* 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.
* 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.

Minor Issues:
* Typos on Line 131, 176
questions: questions: * What is $n_0$ in line 235?
* 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.
* How well do you expect this method to generalize to more complex force fields or larger biomolecular systems?
limitations: limitations: See above.
paper_summary: paper_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.
strength_weakness: strength_weakness: * 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.
* 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.
questions: questions: * Is it possible to find the globally optimal structures? How much does the initial state affect the final predicted structures?
* For the large-scale protein, how much can the proposed method reduce the computational cost?
* Does the proposed method can significantly reduce the memory cost?

Typos:

line 131: mdoes -> modes

line 176: is relies -> relies

line 176: eighted -> weighted
limitations: 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.
paper_summary: paper_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.
strength_weakness: strength_weakness: 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. 
    - 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.
    - 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. 
    - 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. 
    - 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. 
    - Figure captions also needs improvement. For example, the caption of Figure 1 consists of many short sentences, but the key takeaway are not strengthened.

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

1. The manuscript involves biology background knowledge especially in the experiments. However, they are not well explained, making the experimental session more difficult to follow.
questions: 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?

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

1. Is the reparameterized neural network generalizable to different molecular systems or molecular structures under different conditions (e.g. temperature, pressure)?

1. 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: limitations: The authors discussed the limitation of lack of experiments for comparison against traditional CG methods. No potential negative societal impact is involved.

Comments: title: Request to reviewers to respond to rebuttal comments
Thanks
Replies to Rebuttal
comment: Thank you, we appreciate your time. We wanted to add a couple more points: 
* __Scaling for biomolecules:__  We forgot to mention, the method is already very scalable, as follows. We have observed that the Hessian slow modes overlap strongly with the slow modes of the Laplacian of the _molecule graph_. This is likely because the molecule enters the Hessian via quadratic bond energies. We found that just using the full molecule graph inside the GNN, instead of the Hessian slow modes, also yielded good results. Since the molecule graph is quite sparse, and the number of edges scale linearly with the nodes (average degree 2-3), passing through the GNN becomes linear in $n$, making it quite scalable. 
* __Equivariance:__ The slow modes can also be made equivariant to other spatial group by using the correct invariant metric. Ex: The Lorentz group $SO(1,3)$ from special relativity preserves the Minkowski metric $\eta = \mathrm{diag}(-1,1,1,1)$, meaning $g^T\eta g = \eta$ for $g\in SO(1,3)$. If instead of the $L_2$ norm on spatial indices we contract them with this metric  $\mathbf{H} \_{ij} = \sum_{\mu\nu\rho\sigma} H_{ij}^{\mu\nu}H_{ij}^{\rho\sigma} \eta_{\mu\rho}\eta_{\nu\sigma}$, we get equivariance under the Lorentz group.  

We hope you also find that the flow-chart and GNN structure plots address your concern and that you still consider the work for a higher score.
Thank you very much, we really appreciate your time. You comments really made a difference in our paper's structure and presentation. We think the reparametrization idea is an underexplored aspect of neural nets and that it has a lot of potential, especially as a new class of ansatze for scientific problems. We are also working on generalizability and transferrability of the modes to larger molecules. Fianlly, we are applying this methodology to other systems with metastable and slow dynamics such as glassy systems. We observe in some regimes the method may provide more advantage than in other regimes.
Thank you very much for your consideration. Yes, we agree with you about the limitations and are working to address that. One observation we had regarding protein dynamics was that, due to the strong quadratic forces from chemical bonds, the slow modes of the Hessian overlapped very highly with the slow modes of the Laplacian of _molecular graph_ alone (we will share the figure soon). This is a big deal, because the molecular graph is very sparse, with average degree between 2-3. Hence, instead of the expensive Hessian computation, we found using the sparse molecule graph led to faster performance. Our final results for protein folding all use the molecule graph as a proxy instead of the Hessian. This makes our method scalable for even large molecules. Additionally, our method could be applied to coarse-grained dynamics and aid CG models find deeper energies. 

We are also actively working on the DOF from Hessian problem and have found further strong results, suggesting one can explore the phase space of molecules efficiently using these DOF. We also agree that it merits its own paper.
Thank you, we will also produce the loss curve plots in the coming days. We will also try to include our preliminary results on the Kob-Anderson system. We are definitely planning to apply this method to glassy systems. It would be interesting to investigate if in certain hard regimes this methodology would provide more advantage than in other regimes.
Thanks for the authors' clarification. The authors clearly describe their methods by providing additional figures.

In my understanding, there are technical pros and cons for this work. For pros, it is data-efficient and requires less memory cost especially compared to the large structure prediction model such as AlphaFold. However, it is disadvantageous in terms of the computation cost of the Hessian which hinders increasing the scalability, and is dependent on the initial state as it can fall into the local minima conformations.

On the other hand, in the aspect of the methodology, the connection between the Hessian and Laplacian and identifying effective DOF are promising, which can lead to other future works that are related to protein dynamics.

By considering these points, I raise my score to 6
Dear Reviewers,

We are into the last 24 hours of author-reviewer discussions. As you might have noticed, authors have actively tried to engage in the discussions. Hence, I request all the reviewers to please respond to the rebuttal at the earliest so that the authors get a fair chance to represent themselves.

Thank you, 

AC
Thank the author for the response. The rebuttal has addressed most of my concerns. The methods is actually interesting and refreshing for ML research for scientific applications. With that, I will increase my score.
I thank the authors for their response. After these clarifications, I increase my score.

I agree with other reviewers, the presentation needs to be improved. Also, the new figure with learning rate sweep could be a continuous plot(instead of the discrete version), showing change in energy.
I acknowledge that I have read the rebuttal. I ll maintain my score.