VAMPnets: Deep learning of molecular kinetics

Andreas Mardt; Luca Pasquali; Hao Wu; Frank Noé

Source-linked AI summary

VAMPnets: Deep learning of molecular kinetics

Andreas Mardt, Luca Pasquali, Hao Wu, Frank Noé

arXiv:1710.06012v2Oct 16, 2017stat.MLphysics.bio-phphysics.chem-phphysics.comp-ph

TL;DR

Molecular-kinetics modeling traditionally relies on an expert-designed sequence of feature engineering, dimension reduction, clustering, and Markov-model estimation. This paper introduces VAMPnets, neural networks trained with a VAMP variational score to replace that pipeline end to end. The method performs equally or better than state-of-the-art Markov modeling methods and produces interpretable few-state kinetic models.

Problem
Existing molecular-kinetics workflows require substantial modeling expertise because poor choices in feature construction, dimension reduction, clustering, or model estimation can cause large errors.
Method
VAMPnets use neural networks trained by maximizing a VAMP variational score to transform time-lagged molecular configurations into kinetic states.
Results
VAMPnets perform equally or better than state-of-the-art Markov modeling methods and directly produce accurate interpretable models with few states.
Takeaways & Limitations
The framework combines featurization, dimension reduction, and fuzzy clustering into a single end-to-end molecular-kinetics model.

Abstract

from arXiv · show

There is an increasing demand for computing the relevant structures, equilibria and long-timescale kinetics of biomolecular processes, such as protein-drug binding, from high-throughput molecular dynamics simulations. Current methods employ transformation of simulated coordinates into structural features, dimension reduction, clustering the dimension-reduced data, and estimation of a Markov state model or related model of the interconversion rates between molecular structures. This handcrafted approach demands a substantial amount of modeling expertise, as poor decisions at any step will lead to large modeling errors. Here we employ the variational approach for Markov processes (VAMP) to develop a deep learning framework for molecular kinetics using neural networks, dubbed VAMPnets. A VAMPnet encodes the entire mapping from molecular coordinates to Markov states, thus combining the whole data processing pipeline in a single end-to-end framework. Our method performs equally or better than state-of-the art Markov modeling methods and provides easily interpretable few-state kinetic models.

INTRODUCTION

Molecular-kinetics analysis traditionally uses a multistep pipeline whose choices depend heavily on expert judgment. VAMP provides a variational basis for replacing that handcrafted pipeline with a learned structure optimized for Markov-process dynamics.

MD simulations generate extensive data from which stationary, kinetic, and mechanistic properties must be extracted automatically.
The conventional workflow transforms coordinates into features, reduces dimension, clusters the data, and estimates a transition or rate model before coarse-graining it.
Suboptimal decisions at any pipeline step can substantially deteriorate the kinetic model, with feature selection failures causing large modeling errors.
VAMP defines variational scores that measure the optimality of models for arbitrary Markov processes, including nonreversible and non-stationary dynamics.
Maximizing the VAMP-2 score selects feature transformations whose subspaces align with dominant Koopman singular functions.

B. Learning the feature transformation using VAMPnets

VAMPnets use paired time-lagged configurations in two neural-network lobes to learn feature transformations by maximizing a VAMP-2 score. Validation data assess learned features and hyperparameters while limiting overfitting.

Neural networks approximate the basis functions χ0 and χ1 that VAMP optimizes for Markov-process kinetics.
Two parallel lobes receive xt and xt+τ, learn χ0(xt) and χ1(xt+τ), and compute a training VAMP-2 score from transformed transition pairs.
The network can include a constant basis function and use mean-free covariance matrices of feature-transformed coordinates in the score.
Score gradients are back-propagated through both lobes; cloned lobes can be trained using the combined gradient from both outputs.
VAMP-2 validation scores from held-out data are used to assess learned features and select hyperparameters such as network size.

C. Dynamical model and validation

The learned transformations produce Koopman or Markov models that are evaluated through reversibility choices, Chapman–Kolmogorov consistency, and implied-timescale behavior across lag times.

The time-lagged covariance estimate is generally nonsymmetric, so the resulting Koopman or Markov model is typically not time-reversible.
Reversibility can be enforced by reweighting, but the reported results avoid classical PCCA+ because the VAMPnet structure directly supplies the model.
Because the model is Markovian, its transition matrices should satisfy the Chapman–Kolmogorov equation, although approximate MD models can satisfy it only within statistical uncertainty.
Model selection uses eigenvalue decompositions of Koopman matrices and implied timescales evaluated as functions of lag time.
A suitable lag time is chosen where implied timescales remain approximately constant, after which Chapman–Kolmogorov consistency is tested.

D. Network architecture and training

The chosen VAMPnet architecture uses cloned fully connected lobes, dimensionality reduction, and Softmax outputs that represent fuzzy memberships and yield a fuzzy Markov state model.

The applied networks use identical cloned fully connected lobes and generally reduce dimensionality from more input nodes to fewer output nodes.
Softmax outputs are nonnegative and sum to one, allowing each activation to represent the probability of membership in a state.
The architecture combines featurization, dimension reduction, and fuzzy clustering into a fuzzy Markov state model.
Training uses paired configurations (xt, xt+τ) and the Adam stochastic-gradient-descent method, with repeated training and randomized 90%/10% training-validation splits.

E. Asymmetric double well potential

A VAMPnet learns a fuzzy discretization of an asymmetric bistable process from time-lagged coordinates, placing states to resolve the slow transition and reduce Koopman approximation error.

The network receives current and time-lagged mean-free x coordinates and maps them to five Softmax states.The VAMP-2 score uses the four largest singular values.
The learned output states concentrate around the transition region, which improves representation of the dominant Koopman eigenfunction.
The resulting model provides a fuzzy discretization of the asymmetric doublewell dynamics.

F. Protein folding model

In a five-dimensional folding model, VAMPnets learn the nonlinear radial reaction coordinate from Cartesian coordinates and accurately represent the folded-to-unfolded kinetics.

The five-dimensional system is bistable along the reaction coordinate r = |x| because energetic and entropic effects favor different radial regions.
A VAMPnet with two output nodes identifies folded and unfolded states while assigning intermediate memberships to the transition region.
The network learns the nonlinear mapping r = |x| using only the variational score.
The implied timescales and Chapman-Kolmogorov test indicate that the model predicts long-time kinetics almost perfectly.

G. Alanine dipeptide

For alanine dipeptide, VAMPnets learn kinetic coordinates and state discretization from heavy-atom Cartesian coordinates, while the resulting model reproduces the transition kinetics well.

The network must learn both the transformation to φ and ψ torsion-angle space and an optimal metastable-state discretization from Cartesian coordinates.
The six-state model represents transition probabilities between metastable states, omitting probabilities below 0.5%.
The model performs excellently in the Chapman-Kolmogorov test.

H. Choice of lag time, network depth and number of output states

Model performance depends on lag time, depth, and output-state count: successful training occurs only over a limited lag-time range, deeper networks work best at four to seven layers, and VAMPnets outperform MSMs with few states.

Choice of lag time: Training succeeds with significant probability for lag times from 4 to 32 picoseconds, but the success rate remains below 40%.Many failed runs do not resolve the rarely occurring third-slowest process.
Network depth: A high optimization success rate is found for networks with four to seven layers.Deeper networks reduce dimensionality less radically per layer but are more difficult to train.
Number of output states: Two output states separate negative and positive φ regions, while three states additionally distinguish α and β regions of the Ramachandran plot.
Number of output states: Higher output-state counts produce finer discretizations and smaller interconversion timescales until transition regions begin to be discretized.
Number of output states: The best scores using four singular values are achieved with at least six output states, which separate six metastable states.
Comparison with MSMs: MSM VAMP-2 scores are significantly worse than VAMPnet scores below 20 states, whereas VAMPnets directly produce accurate few-state models.

I. VAMPnets learn to transform Cartesian to torsion coordinates

VAMPnets learn internal representations from Cartesian coordinates that recover backbone torsional structure, including the organization of Ramachandran regions.

Pearson correlations of 0.95 and 0.92 link two bottleneck activations to unseen φ and ψ torsion angles, respectively.The network was trained with a two-node bottleneck before the five-node output layer.
The bottleneck representation reproduces the Ramachandran plot’s contiguous clusters, connectivity, and separation of small- and large-φ regions.Four minima at small φ values are grouped with correct connectivity, while large-φ states are separated; two large-φ substates remain unresolved.

J. NTL9 Protein folding dynamics

The NTL9 experiment applies VAMPnets to all-atom folding and unfolding trajectories, producing a compact kinetic model that matches a larger MSM and passes long-time validation.

A five-layer VAMPnet was trained at τ = 10 ns on 111,000 uniformly sampled steps from a 1.11 ms NTL9 trajectory.Internal coordinates were used because folding and unfolding provide no unique reference structure for Cartesian alignment.
Five-state VAMPnet relaxation timescales are on par with those from a 40-state MSM using state-of-the-art estimation methods.The five-state representation is described as easier to interpret and analyze.
The Chapman–Kolmogorov test shows excellent agreement between long-time predictions and direct estimates.
VAMPnets replace featurization, dimension reduction, MSM estimation, and optional coarse-graining with a VAMP-optimized deep neural network.The network can learn feature transformations, dimension reduction, and mappings from MD time steps to fuzzy clusters.
VAMPnets remain limited in supporting multiensemble and augmented Markov models, MSM statistical error estimators, and robust optimization against suboptimal local maxima.These extensions are identified as challenges for future studies.

METHODS

The methods define neural-network architecture and training choices around output dimension, depth, validation, regularization, and lag time, while Figure 6 specifies the NTL9 evaluation panels.

Neural network structure: The network performs dimension reduction when its input dimension exceeds its output dimension, with output size tied to the number of Koopman singular functions approximated.
Neural network structure: Network depth d excludes the input layer, and tested depths ranged from 2 to 11.The structure is fixed by output dimension and depth.
Training and hyperparameters: Dropout is used during training, while hyperparameters are selected with the VAMP-2 validation score to avoid overfitting.The hyperparameters include regularization, dropout, batch size, and Adam learning rate.
NTL9 evaluation: Figure 6 combines state decompositions, relaxation timescales, and a Chapman–Kolmogorov test for NTL9 folding kinetics.The figure compares VAMPnet and TICA-plus-k-means kinetics and reports 95% confidence intervals over 100 training runs.
Training and hyperparameters: A batch size of 4000 was found to be a good compromise between slower training for large batches and stronger score fluctuations for small batches.
Lag-time selection: Larger lag times better isolate slow processes and are therefore more suitable for networks with fewer output nodes.Lag time is selected after choosing the output-node count and hidden layers.