Source-linked AI summary
Time-lagged autoencoders: Deep learning of slow collective variables for molecular kinetics
Christoph Wehmeyer, Frank Noé
TL;DR
Molecular dynamics requires low-dimensional representations that retain slow kinetic behavior, but linear reductions may be inadequate for nonlinear metastable states. The paper develops time-lagged autoencoders, establishes their linear equivalence to TCCA and, for reversible data, TICA, and finds strong performance on toy models and alanine dipeptide.
Problem
Linear low-dimensional subspaces may not resolve nonlinear metastable states, while manually selecting suitable feature transformations places a burden on users.
Method
The paper trains nonlinear time-lagged autoencoders to learn low-dimensional encodings and decodings that optimize time-lagged reconstruction.
Results
Time-lagged autoencoders recover slow-state structure and relaxation timescales more effectively than TICA or PCA across the tested toy models and alanine dipeptide data.
Takeaways & Limitations
Time-lagged autoencoders provide nonlinear embeddings suited to capturing slow dynamics in high-dimensional molecular data.
Abstract
from arXiv · showhide
Inspired by the success of deep learning techniques in the physical and chemical sciences, we apply a modification of an autoencoder type deep neural network to the task of dimension reduction of molecular dynamics data. We can show that our time-lagged autoencoder reliably finds low-dimensional embeddings for high-dimensional feature spaces which capture the slow dynamics of the underlying stochastic processes - beyond the capabilities of linear dimension reduction techniques.
I. THEORY
The theory defines time-lagged autoencoders as encoders and decoders that compress time-series data into a lower-dimensional representation while minimizing reconstruction error.
- A time-series with N dimensions and T time-steps is mapped to a d-dimensional embedding, with d < N, designed to compress time-lagged data.
- The encoding E maps R^N to R^d, while decoding D maps R^d back to R^N to approximately reconstruct the time-lagged signal.
- The formulation uses mean-free coordinates and covariance matrices, with whitening requiring care when C00 or Cττ is rank-deficient.
- The reconstruction objective minimizes the average error over a selected class of encoding and decoding functions.
A. Linear TAE performs TCCA
For linear transformations, the time-lagged autoencoder reduces to a rank-d approximation of the Koopman dynamics and is equivalent to time-lagged canonical correlation analysis.
- Linear encodings and decodings represent mean-free data through matrix multiplications, with d selected from the largest singular values and corresponding singular vectors.
- The encoding projects N-dimensional data into d dimensions, while the decoding lifts it back into the N-dimensional space.
- The reconstruction objective becomes a linear least-squares problem whose full-rank solution is obtained by regression.
- The optimal rank-d approximation of the Koopman matrix is obtained from its rank-d singular value decomposition.
- The resulting whitening-based encoding and decoding solution is equivalent to time-lagged canonical correlation analysis.
B. Time-reversible linear TAE performs TICA
When the time series is time-reversible, the linear time-lagged autoencoder becomes equivalent to TICA, with the placement of singular-value scaling determining the representation.
- If C0τ is symmetric, the Koopman matrix singular-value decomposition is equivalent to an eigenvector decomposition.
- For a stationary time series with C00 = Cττ, the resulting matrices use the usual TICA eigenvectors.
- Including Σd in the decoder yields TICA, whereas including it in the encoder yields a kinetic map.
- These theoretical equivalences motivate nonlinear time-lagged autoencoders that optimize the reconstruction objective.
II. EXPERIMENTS
The experiments test nonlinear time-lagged autoencoders on toy models and alanine dipeptide, comparing reconstruction, low-dimensional representations, and kinetic-model performance with TICA and PCA.
- Nonlinear time-lagged autoencoders are evaluated on two toy models and alanine dipeptide against TICA with kinetic-map scaling and PCA.
- Validation compares reconstruction errors, correlations with known essential variables, and implied-timescale convergence of MSMs built in the encoded space.
- The implementation uses PyTorch networks with one or two hidden layers, a latent layer of size d, leaky rectified linear units, and dropout.
- Training and MSM discretization are repeated 100 times, with ensemble medians and 68% one-standard-deviation percentiles reported.
- The evaluation uses independent or bootstrapped reference data, two-thirds training data, and one-third validation data.
- Encoded coordinates are whitened, compared with reference data using CCA, and used to construct MSMs for implied-timescale validation.
A. Two-state toy model
In the nonlinearly separable two-state HMM, the TAE learns a one-dimensional representation that separates the hidden states and better preserves slow kinetics than TICA or PCA.
- The TAE-encoded variable clearly separates both hidden states, whereas TICA is blurred and PCA does not separate them.
- The TAE outperforms TICA in validation reconstruction error, especially at small transformation lagtimes.
- The TAE encoding yields an MSM that approximates the slowest relaxation timescale even at short lagtimes.
- TICA converges too slowly toward the true timescale and reaches a numerically invalid range before getting close, while PCA is unsuitable for recovering kinetics.
B. Four-state swissroll toy model
For a four-state HMM embedded in a nonlinear swissroll geometry, the TAE outperforms linear methods in reconstruction and state correlation, while MSM timescales remain accurate even for a one-dimensional embedding.
- The TAE outperforms TICA in validation reconstruction error, with the largest difference at small transformation lagtimes.
- The two-dimensional TAE encoding is nearly perfectly correlated with the true hidden-state time series, whereas TICA and PCA perform significantly worse.
- All methods fail to obtain high correlation in one dimension, indicating that the system is not perfectly separable with a single coordinate.
- TAE-based MSMs perfectly recover the reference timescales at all lagtimes, including with the one-dimensional embedding.
- MSMs built on TICA or PCA are systematically underestimated and mostly show no sign of convergence.
C. Molecular dynamics data of alanine dipeptide
In alanine dipeptide simulations represented by 30 heavy-atom coordinates, TAE improves regression error and MSM kinetic estimates, although PCA correlates best with the reference dihedral plane.
- The analysis uses RMSD-aligned heavy-atom positions as a 30-dimensional input and treats φ/ψ dihedrals as the reference representation.
- The TAE outperforms TICA in validation regression error.
- PCA achieves the best correlation with the φ/ψ plane, while TAE and TICA show similar correlations.
- TAE-based MSMs perform best for recovering the first three relaxation timescales, while PCA performs poorly across all timescales.
- TICA converges to the first two relaxation timescales more slowly than TAE, and its third-timescale convergence is not practically useful.
III. CONCLUSION
The paper develops time-lagged autoencoders for low-dimensional nonlinear embeddings of dynamical data. Linear TAEs recover established time-lagged methods, while nonlinear TAEs automatically learn feature transformations that resolve slow processes in toy models and molecular-dynamics data.
- The study targets low-dimensional, nonlinear embeddings of dynamical data using time-lagged autoencoders.
- Linear TAEs are equivalent to TCCA and, for statistically time-reversible data, to TICA.
- Deep TAEs automatically optimize nonlinear feature transformations instead of requiring users to select them manually.
- Results on toy models and molecular-dynamics data indicate that low-dimensional representations of slow dynamics can be obtained with nonlinear TAEs.
- The approach is closely related to VAMPnets, and linear versions of these methods are connected to TCCA, TICA, Koopman models, or MSMs depending on inputs and reversibility.