Descriptors-free Collective Variables From Geometric Graph Neural Networks

Enhanced sampling simulations make the computational study of rare events feasible. A large family of such methods crucially depends on the definition of some collective variables (CVs) that could provide a low-dimensional representation of the relevant physics of the process. Recently, many methods have been proposed to semi-automatize the CV design by using machine learning tools to learn the variables directly from the simulation data. However, most methods are based on feed-forward neural networks and require as input some user-defined physical descriptors. Here, we propose to bypass this step using a graph neural network to directly use the atomic coordinates as input for the CV model. This way, we achieve a fully automatic approach to CV determination that provides variables invariant under the relevant symmetries, especially the permutational one. Furthermore, we provide different analysis tools to favor the physical interpretation of the final CV. We prove the robustness of our approach using different methods from the literature for the optimization of the CV, and we prove its efficacy on several systems, including a small peptide, an ion dissociation in explicit solvent, and a simple chemical reaction.

💡 Research Summary

The paper introduces a fully automated framework for constructing collective variables (CVs) for enhanced‑sampling simulations by leveraging geometric graph neural networks (GVP‑GNNs) that take raw atomic coordinates as input. Traditional ML‑based CV design relies on a set of hand‑crafted descriptors (distances, angles, symmetry functions, etc.), which requires expert knowledge and often fails to capture all relevant degrees of freedom in complex systems. By representing a molecular system as a graph—atoms as nodes, spatially proximate atom pairs as edges—and assigning both scalar (e.g., one‑hot atom types, radial basis expansions of distances) and vector (normalized direction vectors) features to nodes and edges, the authors eliminate the need for any pre‑defined descriptors.

The core architecture is the Geometric Vector Perceptron (GVP), an equivariant neural block that processes scalar and vector channels simultaneously. In each message‑passing layer, node features are concatenated with edge features, passed through several GVP layers to generate messages, and then aggregated (averaged) back into the node. The GVP’s design—linear transformations of scalars, L2‑norm of vectors, gating mechanisms, and separate nonlinearities—ensures invariance to global translations, rotations, reflections (E(3) symmetry) and, because atom types are encoded as one‑hot vectors, invariance to permutations of identical atoms. After a configurable number of message‑passing steps, a global pooling operation over the final scalar node features yields a single scalar CV.

Two distinct learning objectives are explored. Deep Targeted Discriminant Analysis (DeepTDA) treats the CV as a discriminant function that maximizes separation between labeled metastable states while minimizing intra‑state variance. Deep Time‑lagged Independent Component Analysis (DeepTICA) is unsupervised; it maximizes the autocorrelation of the CV at a chosen lag time, thereby extracting the slowest dynamical mode. Both objectives are implemented as differentiable loss functions, allowing end‑to‑end training of the GVP‑GNN.

The methodology is validated on three benchmark systems:

1. Alanine dipeptide in vacuum – a classic two‑state test case defined by the φ/ψ dihedral angles. The GVP‑GNN learns a CV that reproduces the free‑energy surface obtained from the conventional φ/ψ coordinates. DeepTDA yields a discriminative CV that cleanly separates the two basins, while DeepTICA captures the slowest rotational mode, confirming that the network can learn both classification‑ and kinetics‑oriented CVs.

2. NaCl dissociation in explicit water – a noisy, high‑dimensional problem where the ion pair interacts with a fluctuating solvent shell. Despite the noisy training data, the GVP‑GNN discovers a low‑dimensional CV that combines ion‑ion distance with solvent coordination features, enabling clear discrimination between bound and dissociated states. The resulting free‑energy profile matches reference calculations that required extensive manual descriptor engineering.

3. Methyl migration in 2‑fluoro‑2,3‑dimethyl‑butane (FDMB) cation – a reaction where permutation symmetry of identical methyl groups is crucial. Because the GVP‑GNN is permutation‑invariant, the learned CV respects this symmetry and outperforms a comparable feed‑forward neural network CV in both barrier height estimation and transition‑path identification. This demonstrates the practical advantage of built‑in symmetry handling.

To aid interpretability, the authors propose post‑hoc analysis tools: (i) per‑node scalar contribution maps that highlight which atoms dominate the CV, and (ii) edge‑importance visualizations that reveal key inter‑atomic interactions. These visualizations bridge the “black‑box” nature of deep learning with chemical intuition, allowing researchers to extract physically meaningful insights from the learned CV.

From a computational standpoint, the GVP‑GNN requires far fewer trainable parameters than traditional descriptor‑based pipelines (e.g., SOAP, ACSF, PIP) and scales naturally with system size because the graph construction is O(N) in the number of atoms. Training times are comparable to, or faster than, feed‑forward networks using large descriptor vectors, and inference is efficient enough for on‑the‑fly biasing in metadynamics or umbrella sampling.

The authors also emphasize modularity: their codebase abstracts the graph construction, message‑passing, and loss functions, making it straightforward to swap the GVP block for other equivariant architectures (e.g., SE(3)‑Transformer, MACE) or to incorporate multi‑task objectives such as committor learning.

In summary, the paper makes four major contributions: (1) a descriptor‑free, fully automated pipeline for CV generation, (2) rigorous enforcement of E(3) and permutation symmetries via the GVP architecture, (3) demonstration of flexibility across supervised (DeepTDA) and unsupervised (DeepTICA) learning objectives, and (4) provision of interpretability tools that connect the learned CV to underlying atomic motions. This work paves the way for scalable, black‑box‑free enhanced‑sampling simulations of complex biomolecular and materials systems, reducing reliance on expert intuition and accelerating the discovery of accurate free‑energy landscapes.