Definition
TS-DAR (Transition State identification via Dispersion and vAriational principle Regularized neural networks) is a deep learning framework designed to identify Transition States in molecular dynamics. It builds upon the VAMPnets architecture but adds geometric constraints to solve VAMPnets’ inability to clearly define boundaries between states.
The Problem
Standard VAMPnets are good at identifying metastable states but often result in unclear/fuzzy boundaries where transition states get lumped into the stable basins.
The Method
TS-DAR uses the standard VAMPnet architecture to learn the kinetics, but adds two specific innovations:
1. Hyperspherical Latent Space
The state probability vectors from the VAMPnet are normalized to length 1, projecting them onto a Hypersphere. The dimensionality of the sphere (K) equals the number of metastable states.
2. Dispersion Loss
While VAMP-2 maximizes kinetic correlation, Dispersion Loss forces geometric separation. It calculates the Prototypes (centroids) of each metastable state on the hypersphere and minimizes the Cosine Similarity between these prototypes. This pushes the state clusters as far apart as possible (e.g., to the poles of the sphere), creating distinct empty gaps between them.
Total Loss Equation:
Transition State Detection (OOD)
Once the model is trained and the metastable states are separated, the empty gaps on the hypersphere are analyzed using an Out-of-Distribution (OOD) Score. The OOD score measures the distance of a specific frame from the nearest cluster center. A High OOD score means the frame is far from all known stable states → Flagged as a Transition State.
Old description
Deep Learning framework for identifying transition states of molecular systems. Designed to solve VAMPnets shortcoming of identifying the transition states rather than just the metastable states (basins).
It works by taking 2 frames of molecular coordinate data at a time.
Frame 1:
Frame 2:
These frames are the set of coordinates from the Molecular Dynamics Simulation. (x, y, z). The frames are sent into 2 identical Neural Networks in parallel with shared weights. The output from these NNs are a probability distribution vector (after being passed through a softmax function) that gives the probability of what state that frame is currently in. (ex: (0.1, 0.7, 0.2) - 10% state #1, 70% state #2, 20% state #3). Then, using a batch of these frames, the Covariance Matrices are calculated. The VAMP-2 score is derived from these matrices to quantify how well the features at time t linearly predict the features at . The loss is then sent back, after being combined with the dispersion loss (explained later), through the network to train the model to maximize the covariance and minimize this loss.
The Neural Networks are effectively taking the chaotic and nonlinear molecular coordinates (nonlinear in predicting the next coordinates at ) and transforming them into linearly correlated probability vectors. This effectively emulates the Koopman Operator which maps nonlinear data into linear functions of higher dimensionality.
These state vectors are projected as embeddings onto a hyperspherical Latent Space. The hypersphere has dimension which is equal to the number of metastable states the model is finding. This results in each axis representing a state. Along with the VAMP-2 loss, a dispersion loss is calculated. This dispersion loss finds the cosine similarity between the Prototypes (Centers) of the states and aims to maximize the distance between these Centers on the hypersphere. This creates gaps in between the metastable states, which is exactly where the transition states will be. The dispersion loss is then combined with the Vamp loss to get the total loss, which is then Backpropogated through the network.
Finally, the the gaps in the hypersphere (areas in-between the clusters of metastable states) are analyzed by calculating the Out of Distribution (OOD) score for every frame. Frames with high scores, meaning they fall into the empty regions far from any cluster center, are flagged as Transition States.