GraphVAMPNet, using graph neural networks and variational approach to markov processes for dynamical modeling of biomolecules

TL;DR

This work addresses learning accurate, high-resolution kinetic models from long molecular dynamics trajectories. It introduces GraphVAMPNet, an end-to-end framework that combines VAMPNet with graph neural networks (SchNet) to produce graph-based molecular embeddings, enabling a more detailed metastable-state decomposition and faster convergence of implied timescales. Across Trp-Cage, Villin, and NTL9, GraphVAMPNet achieves higher VAMP-2 scores and passes Chapman-Kolmogorov tests, with attention revealing residue-level contributions to state classification. The approach reduces reliance on hand-crafted features and provides interpretable, scalable models for biomolecular kinetics, with potential extensions to incorporating reversibility constraints and transfer learning.

Abstract

Finding low dimensional representation of data from long-timescale trajectories of biomolecular processes such as protein-folding or ligand-receptor binding is of fundamental importance and kinetic models such as Markov modeling have proven useful in describing the kinetics of these systems. Recently, an unsupervised machine learning technique called VAMPNet was introduced to learn the low dimensional representation and linear dynamical model in an end-to-end manner. VAMPNet is based on variational approach to Markov processes (VAMP) and relies on neural networks to learn the coarse-grained dynamics. In this contribution, we combine VAMPNet and graph neural networks to generate an end-to-end framework to efficiently learn high-level dynamics and metastable states from the long-timescale molecular dynamics trajectories. This method bears the advantages of graph representation learning and uses graph message passing operations to generate an embedding for each datapoint which is used in the VAMPNet to generate a coarse-grained representation. This type of molecular representation results in a higher resolution and more interpretable Markov model than the standard VAMPNet enabling a more detailed kinetic study of the biomolecular processes. Our GraphVAMPNet approach is also enhanced with an attention mechanism to find the important residues for classification into different metastable states.

Figures (7)

• Figure 1: Overview of the architecture of GraphVAMPNet method. Given a molecular structure at time $t$ and a lagtime later $t+\tau$ , molecular graphs are built using the nearest neighbor of the chosen atoms. Several graph convolution operations are performed resulting in representation for each node. A hierarchical pooling is done to find a latent representation of the full graph which is concatenated between time $t$ and $t+\tau$. The full network is then optimized by maximizing a VAMP-2 score.
• Figure 2: TrpCage system A) Implied timescale (ITS) plot for SchNet as feature transformation in VAMPNet (errors are taken from 95 % confidence interval from 10 different trainings) B) CK-test for SchNet using lagtime of 20 ns C) Free energy landscape (FEL) in a 2d graph embedding F) state assignment of the 2d graph embedding using 0.95 cutoff.
• Figure 3: A) Representative structure of each metastable state in TrpCage with their probabilities B) average attention score between $C_\alpha$ atoms for each cluster C) averaged attention score for each residue of TrpCage in each cluster which is the scaled sum of rows.
• Figure 4: Villin system A) Implied timescale (ITS) plot for SchNet as feature transformation in VAMPNet (errors are taken from 95 % confidence interval from 10 different trainings) B) CK-test for SchNet using lagtime of 20 ns C) Free energy landscape (FEL) in a 2d graph embedding F) state assignment of the 2d graph embedding using 0.75 cutoff.
• Figure 5: Representative structure of each metastable state in Villin with their probabilities B) average attention score between $C_\alpha$ atoms for each cluster C) averaged attention score for each residue of Villin in each cluster which is the scaled sum of rows.