Source-linked AI summary
Coarse Graining Molecular Dynamics with Graph Neural Networks
Brooke E. Husic, Nicholas E. Charron, Dominik Lemm, Jiang Wang, Adrià Pérez, Maciej Majewski, Andreas Krämer, Yaoyi Chen, Simon Olsson, Gianni de Fabritiis, Frank Noé, Cecilia Clementi
TL;DR
Atomistic simulations remain limited at long timescales, motivating coarse-grained models that preserve thermodynamic conclusions. The paper introduces CGSchNet, which learns molecular features with SchNet inside supervised force matching, and reports accurate thermodynamics for small biomolecular systems with improved free-energy surfaces, robustness, and reduced regularization. Its learned representations support transferability across system size and sequence, while transfer across thermodynamic states remains limited.
Problem
Atomistic MD cannot reach the longest biomolecular timescales, while existing supervised coarse-graining approaches require hand-selected molecular features that limit transferability.
Method
CGSchNet integrates SchNet’s continuous-filter graph neural network into CGnet to learn molecular features for supervised force matching.
Results
CGSchNet reproduces thermodynamics for small biomolecular systems and yields more accurate free energy surfaces, greater hyperparameter robustness, and less regularization than CGnet.
Takeaways & Limitations
The learned molecular representations make the architecture transferable across molecular system size and sequence, supporting development of transferable coarse-grained force fields.
Takeaways & Limitations
CGSchNet is not readily transferable across thermodynamic states, and the optimal coarse-grained model resolution remains an unresolved challenge.
Abstract
from arXiv · showhide
Coarse graining enables the investigation of molecular dynamics for larger systems and at longer timescales than is possible at atomic resolution. However, a coarse graining model must be formulated such that the conclusions we draw from it are consistent with the conclusions we would draw from a model at a finer level of detail. It has been proven that a force matching scheme defines a thermodynamically consistent coarse-grained model for an atomistic system in the variational limit. Wang et al. [ACS Cent. Sci. 5, 755 (2019)] demonstrated that the existence of such a variational limit enables the use of a supervised machine learning framework to generate a coarse-grained force field, which can then be used for simulation in the coarse-grained space. Their framework, however, requires the manual input of molecular features upon which to machine learn the force field. In the present contribution, we build upon the advance of Wang et al.and introduce a hybrid architecture for the machine learning of coarse-grained force fields that learns their own features via a subnetwork that leverages continuous filter convolutions on a graph neural network architecture. We demonstrate that this framework succeeds at reproducing the thermodynamics for small biomolecular systems. Since the learned molecular representations are inherently transferable, the architecture presented here sets the stage for the development of machine-learned, coarse-grained force fields that are transferable across molecular systems.
I. INTRODUCTION
Coarse graining extends biomolecular simulation beyond atomistic length and time scales by reducing structural degrees of freedom while targeting preserved physical properties. This work replaces hand-selected molecular features in supervised force matching with learnable graph-based features, producing a transferable architecture.
- Motivation: Coarse graining enables more efficient simulation and analysis of biomolecular systems beyond the length and time scales accessible to all-atom MD.It can also elucidate physical components important to molecular processes.
- Coarse-graining framework: Coarse-grained models combine a mapping from atomistic entities to interaction-site beads with a physical model for those beads.Machine learning has been used to learn both mapping and model components in bottom-up coarse graining.
- Approach: The paper formulates bottom-up force matching as supervised learning of a coarse-graining model to reproduce the thermodynamics of small biomolecular systems.The approach focuses on learning the coarse-grained force field rather than the mapping.
- Methodological contribution: CGSchNet learns molecular features through graph convolutional neural networks instead of requiring hand-selected features.It leverages SchNet’s transferable representation-learning architecture within CGnet.
- Reported advantages: The reported protocol produces more accurate free energy surfaces, greater robustness to hyperparameter choices, and less regularization than hand-selected features.These properties support transferability across molecular systems.
A. Force matching
Force matching approximates a thermodynamically consistent coarse-grained potential by matching projected atomistic forces with forces generated by a reduced-coordinate energy function. The exact potential is generally intractable because it requires exhaustive sampling of atomistic configurations.
- Atomistic force matching: Force matching trains a trial energy function by comparing reference atomistic forces with forces obtained from the negative derivative of the trial energy.The reference dataset contains atomistic coordinates and corresponding forces, while the trial function has parameters Θ.
- MS–CG adaptation: MS–CG maps N atomistic coordinates to n < N coarse-grained beads, then matches projected atomistic forces with predicted coarse-grained forces.The mapping matrix Ξ clusters atoms into interaction sites and links all-atom and coarse-grained representations.
- Force projection: The instantaneous coarse-grained force is the projection of the atomistic force into coarse-grained space, with ΞF = (ΞΞ⊤)−1Ξ given as one general projection.Other choices of force mapping are possible.
- Potential of mean force: The thermodynamically consistent coarse-grained energy is the potential of mean force, defined through atomistic configurations that map to a given coarse-grained structure.Multiple atomistic configurations can correspond to the same coarse-grained configuration after degrees of freedom are removed.
- Approximation: Force matching is needed because the exact potential of mean force is intractable when the full atomistic configuration space cannot be exhaustively sampled.The method therefore approximates the coarse-grained energy by minimizing a force-matching objective.
B. Coarse graining as a supervised machine learning problem
A variational formulation makes force matching a supervised learning problem for thermodynamically consistent coarse-grained force fields. CGSchNet extends CGnet by learning transferable molecular representations with SchNet rather than relying on hand-selected structural features.
- Variational foundation: The variational principle identifies a scalar loss whose minimum corresponds to the thermodynamically consistent coarse-grained force field.This enables supervised or semi-supervised machine learning formulations of coarse graining.
- CGnet: CGnet is a neural network architecture designed to minimize the force-matching loss while enforcing known system properties such as symmetries and equivariances.A trained CGnet can serve as a force field for new coarse-grained data points.
- Motivation: The original CGnet required regularization and system-specific tuning, whereas CGSchNet was introduced to address its feature-size and transferability limitations.The baseline had to be retrained when system geometry changed because its feature size was fixed.
- SchNet: SchNet uses continuous filter convolutions in a graph neural network to learn representations of particles embedded in three-dimensional space.Its filters depend on interparticle distances, preserving expected invariances.
- CGSchNet: CGSchNet incorporates SchNet into CGnet to replace hand-selected structural features with learned molecular features.The resulting hybrid framework converts CGnet into a transferable architecture across molecular systems.
A. Original CGnet architecture
The original CGnet learns a coarse-grained energy correction from hand-selected structural features and differentiates its scalar output to obtain conservative forces. Its system-specific feature representation makes transferability difficult and can require retraining for new geometries.
- Data preparation: CGnet prepares coarse-grained structures and mapped atomistic-force labels from atomistic configurations using a mapping matrix Ξ.The mapping produces n-bead representations from N-atom data, with n < N.
- Feature representation: Original CGnet uses distances, planar angles, and torsional angles as manually selected structural features.The present architecture replaces these fixed features with learned representations.
- Regularization: Regularized CGnets can include harmonic priors for selected distances and angles and repulsion priors for nonbonded distances.Prior coefficients can be selected through cross-validated hyperparameter optimization.
- Energy and forces: A scalar final network output makes the learned coarse-grained force field curl-free and therefore energy-conserving.Autodifferentiation is used to differentiate the learned energy with respect to coordinates.
- Transferability: System-specific hand-selected features make CGnet sensitive to hyperparameters and generally require retraining when the molecular geometry changes.The fixed feature size is tied to each system’s geometry.
1. Learning molecular representations with SchNet
SchNet learns molecular representations on coarse-grained bead graphs using distance-dependent continuous-filter convolutions. These learned, symmetry-preserving features feed a beadwise CGnet energy model while retaining permutation-invariant learned energies, subject to the priors used here.
- Graph representation: SchNet represents particles as three-dimensional graph nodes and uses interparticle-distance-dependent filters to preserve expected molecular invariances.In CGSchNet, the nodes are coarse-grained beads rather than atomistic particles.
- Learned embeddings: Trainable bead embeddings encode bead types and allow SchNet to learn molecular representations instead of using fixed heuristic features.Applications use nuclear charges or amino acid identities as bead types.
- Interaction updates: Continuous convolutions transform radial-basis-featurized inter-bead distances into filters that update neighboring bead representations.Neighborhoods may include all beads or only beads within a cutoff.
- Interaction updates: Residual interaction blocks refine bead representations, using nonlinear trainable functions with a hyperbolic-tangent nonlinearity.The residual update is intended to prevent gradient annihilation in deep networks.
- Energy output: A beadwise CGnet network maps final SchNet representations to scalar bead energies, which are summed into a permutation-invariant learnable energy.The priors used in this study are not permutation invariant and therefore break that invariance in the full model.
- Hybrid architecture: CGSchNet uses SchNet interaction layers to learn the input representation for a beadwise CGnet output network.In this study, the architecture is trained on individual molecular systems rather than used transferably.
2. CGSchNet: a transferable architecture for coarse graining
CGSchNet combines CGnet force matching with SchNet-learned molecular features to address CGnet’s reliance on system-specific structural descriptors. The resulting representation is more flexible and is designed to support transfer across molecular systems, while coarse-grained free-energy labels remain unavailable for direct energy matching.
- Motivation: CGnet’s reliance on system-specific structural features prevents it from learning a transferable coarse-grained force field.The paper addresses this by incorporating SchNet into CGnet.
- Architecture: CGSchNet replaces predetermined distances, angles, and torsions with features learned by a SchNet subnetwork.The learned representation scales better with system size and is amenable to transferability.
- Architecture: SchNet requires coordinates and bead types, allowing chemical or physical intuition to enter through bead identities rather than fixed structural features.Examples include amino acid identities for protein systems.
- Transferability: Shared bead-type embeddings can apply across systems with the same components, supporting transfer across size and sequence space.The paper presents this as a benefit of the SchNet-based representation.
- Training objective: CGSchNet is trained with a coarse-grained force-matching loss comparing predicted forces with known training-set forces.A separate energy-matching term cannot be incorporated straightforwardly because coarse-grained free energies require an intractable microscopic integral.
C. Coarse-grained simulations
A trained CGSchNet serves as a coarse-grained force field whose predicted forces are integrated with Langevin dynamics. The capped-alanine setup preserves six atoms and uses φ and ψ to characterize the central alanine’s heavy-atom dynamics.
- Simulation protocol: A trained CGSchNet predicts forces on coarse-grained coordinates, which an integrator uses to propagate simulations forward in time.Repeating force prediction and integration generates a dataset of the desired duration.
- Simulation protocol: Langevin dynamics propagate the coarse-grained coordinates using bead masses, a collision rate, and stochastic noise.The implementation uses a BAOAB Langevin integrator.
- Simulation protocol: Overdamped Langevin dynamics omit inertia and use a diffusion matrix D = M^-1k_BT/γ.This formulation does not require bead masses separately because the product Mγ can be used without separating its factors.
- Capped alanine: The capped-alanine coarse-grained representation preserves six shaded atoms, while φ and ψ completely describe the central alanine’s heavy-atom dynamics.Capped alanine is used as a benchmark for molecular-dynamics method development.
- Capped alanine: Two 100-ns CGSchNet trajectories show multiple transitions among basins in both φ and ψ reaction coordinates.The plotted darker curves are smoothed with a moving average of 250 simulation frames.
A. Capped alanine
CGSchNet learns a six-bead capped-alanine force field from atomistic coordinates and forces, then generates coarse-grained simulations whose free-energy surfaces reproduce the baseline’s major basins and barriers. Noise-calibrated comparisons indicate that its distributional error is approximately comparable to 0.1 Å coordinate perturbations, while retaining greater visual fidelity.
- Mapping and training: A six-bead mapping retains capped alanine’s backbone heavy atoms and Cβ for CGSchNet force-field training on atomistic coordinates and forces.The selected beads are C–[N–Cα–C]Ala–N plus alanine Cβ.
- Simulation protocol: 100 parallel 200 ns coarse-grained simulations sample capped alanine from varied Ramachandran-space starting positions.The trajectories include transitions among the major conformational basins.
- Free-energy reproduction: CGSchNet captures the same major free-energy basins as the all-atom baseline and approximates the one-dimensional free-energy barriers.The comparison uses φ and ψ Ramachandran coordinates and one-dimensional landscapes along each reaction coordinate.
- Noise calibration: Increasing Cartesian noise smears free-energy surfaces and reduces barriers, with perturbations evaluated across σ = 0.1 Å, 0.2 Å, and 0.3 Å.The perturbed surfaces are constructed by recalculating φ and ψ after adding Gaussian coordinate noise.
- Distributional comparison: Both KL divergence and MSE increase with noise, while CGSchNet’s error is approximately comparable to the perturbed dataset at σ = 0.1 Å.Despite comparable metric values, CGSchNet more closely resembles the baseline surface visually than the broadened 0.1 Å perturbation.
B. Chignolin
CGSchNet is evaluated on chignolin using a ten-α-carbon coarse-grained representation and compared with baseline, noisy-baseline, and CGnet simulations. It reproduces major thermodynamic features, while the study explicitly excludes kinetic analysis because force matching is limited to thermodynamic consistency.
- Representation and model: The chignolin CGSchNet model retains only the ten α-carbons as coarse-grained beads and assigns amino-acid-specific SchNet environments.Training data come from an explicit-solvent atomistic simulation with stored forces.
- Thermodynamic analysis: MSM analysis featurizes each system with 45 α-carbon pairwise distances and uses TICA to obtain four slow reaction coordinates.The same analysis framework is applied to baseline, CGSchNet, CGnet, and noisy-baseline datasets.
- Free-energy reproduction: CGSchNet reasonably approximates free-energy barriers along the TICA reaction coordinates.The coordinates are reweighted using each MSM’s stationary distribution before constructing one-dimensional free-energy surfaces.
- Comparison with CGnet: For both divergence metrics, CGSchNet corresponds to a lesser added-noise value than the CGnet simulation.KL divergence and MSE both increase monotonically with the magnitude of baseline-coordinate noise.
- Conformational ensembles: CGSchNet captures chignolin’s three major atomistic basins corresponding to folded, unfolded, and misfolded ensembles.The basins are identified from MSM-derived free-energy surfaces.
- Scope boundary: The study does not attempt kinetic analysis because force matching is limited to thermodynamic consistency.Matching dynamics in addition to thermodynamics remains an open challenge in the authors’ stated scope.
V. DISCUSSION
CGSchNet combines variational force matching with learnable graph-based molecular features to create a transferable coarse-grained force-field architecture. On capped alanine and chignolin, it improves free-energy-surface accuracy and robustness relative to CGnet, while requiring less regularization, though larger proteins remain an open challenge.
- Architecture and transferability: The variational force-matching framework provides the theoretical basis for supervised learning of thermodynamically consistent coarse-grained force fields.The framework connects force matching to a supervised machine-learning formulation through a variational limit.
- Architecture and transferability: CGSchNet replaces CGnet’s hand-selected molecular features with a learnable SchNet representation based on continuous-filter graph convolutions.This design is intended to support transfer across molecular systems of different sizes and sequences.
- Performance: CGSchNet produced more accurate chignolin free-energy surfaces than CGnet and was more robust to network hyperparameter choices.The comparison concerns simulations of capped alanine and the 10-amino-acid miniprotein chignolin.
- Performance: CGSchNet uses less regularization than CGnet and does not require the additional Lipschitz constraint applied to CGnet weight matrices.The authors identify reduced regularization as a third performance advantage over the predecessor architecture.
- Scope and outlook: The current protocol has demonstrated success for a capped monopeptide and a 10-amino-acid miniprotein, while accurate force fields for larger proteins remain an open challenge.The authors suggest that larger systems may require new sampling strategies, structural priors, or architectural changes.
- Scope and outlook: CGSchNet is transferable across system size and sequence but is not readily transferable across thermodynamic states.Related work is noted as addressing multiple thermodynamic states using the same variational principle in a semi-supervised setting.
- Scope and outlook: The accompanying software package is presented as a framework for advancing transferable coarse-grained force fields and future work on coarse-grained dynamics and model resolution.The paper identifies dynamics learning and resolution optimization as outstanding challenges.
Appendix A: Software
The cgnet software package implements the paper’s methods and provides notebook-based examples, with core dependencies for numerical computing, machine learning, trajectory analysis, and visualization.
- Software: The cgnet package is available under the BSD-3-Clause license and requires NumPy, SciPy, and PyTorch.Optional functionality depends on pandas, MDTraj, and Scikit-learn.
- Software: Examples are provided in Jupyter notebooks, which also require Matplotlib.The Langevin-dynamics simulation code is adapted from OpenMM, while the SchNet implementation is inspired by SchNetPack.