Time-lagged autoencoders: Deep learning of slow collective variables for molecular kinetics

I Theory

To motivate our approach we first consider the case of linear transformations. We show that in a similar way that a linear autoencoder and PCA are equivalent up to orthogonalization, linear TAEs are equivalent to time-lagged canonical correlation analysis (TCCA), and in the time-reversible case equivalent to TICA. Then we move on to employ nonlinear TAEs for dimension reduction.

is small in some suitable norm. We introduce two conventions:

Note that whitening must take into account that $\mathbf{C}_{00}$ and $\mathbf{C}_{\tau\tau}$ are often not full rank matrices (Wu et al., 2017).

Now we must find an encoding and decoding which minimizes the reconstruction error

for a selected class of functions $E$ and $D$. The simplest choice, of course, are linear functions.

As we operate on mean free data, we can represent linear encodings and decodings by simple matrix multiplications:

Eq. (7) is a linear least squares problem. The full-rank solution is given by the regression

This form is often referred to as half-weighted Koopman matrix and it arises naturally when using whitened data.

The last step to the solution lies in choosing the optimal rank-$d$ approximation to the Koopman matrix which is given by the rank-$d$ singular value decomposition,

where $d$ indicates that we take the $d$ largest singular values and corresponding singular vectors. Thus, a possible choice for the encoding and decoding matrices for whitened data is

With this choice, we find for the mean-free but non-whitened data

the non-whitened Koopman matrix consistently with (Noé and Nüske, 2013; Williams, Kevrekidis, and Rowley, 2015; Wu and Noé, 2017; Horenko et al., 2007):

Likewise, we find the non-whitened encoding and decoding matrices

where the encoding consists of a whitening followed by the whitened encoding, while the decoding starts with the whitened decoding followed by unwhitening. This solution is equivalent with time-lagged canonical correlation analysis (TCCA) (Hotelling, 1936; Wu and Noé, 2017).

I.2 Time-reversible linear TAE performs TICA

If the covariance matrix $C_{0\tau}$ is symmetric, the singular value decomposition of the full-rank Koopman matrix is equivalent to an eigenvector decomposition:

If we further have a stationary time series, i.e., $\mathbf{C}_{00}=\mathbf{C}_{\tau\tau}$, the non-whitened encoding and decoding matrices are

where $\mathbf{U}_{d}^{\top}\mathbf{C}_{00}^{-\frac{1}{2}}$ contains the usual TICA eigenvectors and multiplication with $\mathbf{\Sigma}_{d}$ transforms to a kinetic map. Thus, if we include $\mathbf{\Sigma}_{d}$ in the decoder part, this solution is equivalent to TICA (Molgedey and Schuster, 1994; Pérez-Hernández et al., 2013; Schwantes and Pande, 2013), while if we include it in the encoder, it is equivalent to a kinetic map (Noé and Clementi, 2015).

Motivated by these theoretical results we will employ TAEs to learn nonlinear encodings and decodings that optimize Eq. (6).

II Experiments

We put the nonlinear time-lagged autoencoder to the test by applying it to two toy models of different degree of difficulty as well as molecular dynamics data for alanine dipeptide. In all three cases, we compare the performance of the autoencoder with that of TICA (with kinetic map scaling) and PCA by

comparing the reconstruction errors (6) of the validation sets,

comparing the low-dimensional representations found with the known essential variables of the respective system by employing canonical correlation analysis (CCA), and

examining the suitability of the encoded space for building MSMs via convergence of implied timescales.

The time-lagged autoencoders used in this study are implemented using the PyTorch framework (Paszke et al., 2017) and consist of an input layer with $N$ units, followed by one or two hidden layers with sizes $H_{1}$ and $H_{2}$ and the latent layer with size $d$ which concludes the encoding stage. The decoding part also adds one or two hidden layers of the same sizes as in the encoding part, followed by the output layer with size $N$. All hidden layers employ leaky rectified linear units (Maas, Hannun, and Ng, 2013) (leaky parameter $\alpha=0.001$) and a dropout layer (Srivastava et al., 2014) (dropout probability $p=0.5$). We train the networks using the Adam (Kingma and Ba, 2014) optimizer.

To account for the stochastic nature of the high dimensional data, the autoencoder training process, and the discretization when building MSMs, all simulations have been repeated 100 times while shuffling training and validation sets. We show the ensemble median as well as a one-standard-deviation percentile (68%). The evaluation process always follows the pattern

Gather the high dimensional data and reference low-dimensional representation via independent simulation or bootstrapping.

Train the encoder/decoder for all techniques on two thirds (training set) of the high dimensional data.

Compute the reconstruction error for the remaining third of the data (validation set).

Obtain encoded coordinates and whiten (training + validation sets).

Perform CCA to compare the encoded space to the reference data (training + validation sets).

Build MSMs (Scherer et al., 2015) on the encoded space (training + validation sets) and validate using implied timescales tests (Swope, Pitera, and Suits, 2004).

The first toy model is based on a two-state hidden Markov model (HMM) which emits anisotropic Gaussian noise in the two-dimensional $x/y$-plane. To complicate matters we perform the operation

which leads to the distribution shown in Fig. 2a. We compare one-dimensional representations found by applying TICA and PCA to the time series, and a TAE employing one hidden layer of 50 units in the encoding and decoding part each, and a bottleneck size of $d=1$.

The TAE-encoded variable overlaps very well with the hidden state time series and can clearly separate both hidden states, while TICA gives a more blurred picture with no clear separation and PCA does not seem to separate the hidden states at all (Fig. 2b). These differences are quantified by the CCA score between the encoded and true hidden state signals (Fig. 2d). The time-lagged autoencoder outperforms TICA at all examined transformation lagtimes in terms of the reconstruction error; the difference is particularly strong for small lagtimes (Fig. 2c). Finally, the encoding found by the TAE is excellently suited to build an MSM that approximates the slowest relaxation timescale even at short lagtimes (Fig. 2e). In contrast, the MSM based on TICA converges towards the true timescale too slowly, and does not get close to it before reaching the numerically invalid range $\tau>t_{2}$. The MSM build on PCA seems to be completely unsuitable for recovering kinetics (Fig. 2e).

II.2 Four-state swissroll toy model

The second toy model is based on a four-state hidden Markov model (HMM), which emits isotropic Gaussian noise in the two-dimensional $x/y$-plane, with the means of the states located as shown in Fig. 3a. To create a nonlinearly separable system, we perform the operation

which produces a picture that is reminiscent of the swiss roll commonly used as a benchmark for nonlinear dimension reduction (Fig. 3b). For this toy model, we examine two- and one-dimensional encodings. The TAE with $d=2$ uses a single hidden layer with 100 units in the encoder and decoder part, while the TAE with $d=1$ uses two hidden layers with 200 and 100 units in the encoder part and 100 and 200 units in the decoder.

In both cases, the time-lagged autoencoder outperforms TICA in terms of reconstruction error (Fig. 3c). Again, the difference is larger for small transformation lagtimes. Indeed, the TAE encoding is nearly perfectly correlated with the true hidden states time series (Fig. 3d), while both TICA and PCA are significantly worse and nearly identical to each other. In the one-dimensional case, all methods fail at obtaining a high correlation, indicating that this system is not perfectly separable with a single coordinate, even if it is nonlinear.

MSMs constructed on the encoded space also indicate that the TAE perfectly recovers the reference timescales at all lagtimes (Fig. 3e) – surprisingly this is also true for the one-dimensional embedding, despite the fact that this embedding is not well correlated with the true hidden time series. MSMs build on either the TICA or PCA space are systematically underestimated and mostly show no sign of convergence.

II.3 Molecular dynamics data of alanine dipeptide

Our third example involves real MD data from three independent simulations of 250 ns each (Nüske et al., 2017; Harvey, Giupponi, and Fabritiis, 2009) from which we repeatedly bootstrap five sub-trajectories of length 100 ns. The features on which we apply the encoding are RMSD-aligned heavy atom positions which yield an $N=30$-dimensional input space. Although we do not know the optimal two-dimensional representation, we assume that the commonly used ($\phi$,$\psi$) backbone dihedrals contain all the relevant long-time behavior of the system (except for methyl rotations which do not affect the heavy atoms (Zheng et al., 2013)), and we thus use these dihedral angles as a reference to compare our encoding spaces to. Fig. 4a and b show the free energy surface for the reference representation and the assignment of ($\phi$,$\psi$)-points to the four most slowly-interconverting metastable states.

The TAE outperforms TICA in terms of the regression error (Fig. 4c). While all three methods find a two-dimensional space that correlates relatively well with the ($\phi$,$\psi$)-plane, PCA achieves, surprisingly the best correlation (Fig. 4d-e), while the TAE and TICA are similar. This result is put into perspective by the performances of MSMs built upon the encoding space (Fig. 4f-h). Here, TAE clearly performs best. The TICA MSM does converge to the first two relaxation timescales, although slower than the TAE in the first relaxation timescale, while its convergence of the third relaxation timescale is too slow to be practically useful. PCA performs poorly for all relaxation timescales.

III Conclusion

We have investigated the performance of a special type of deep neural network, the time-lagged autoencoder, to the task of finding low-dimensional, nonlinear embeddings of dynamical data. We have first shown that a linear time-lagged autoencoder is equivalent to time-lagged canonical correlation analysis, and for the special case of statistically time-reversible data equivalent to the time-lagged independent component analysis commonly used in the analysis of MD data. However, in many datasets, the metastable states are not linearly separable and there is thus no low-dimensional linear subspace that will resolve the slow processes, resulting in large approximation errors of MSMs and other estimators of kinetics or thermodynamics. In these cases, the traditional variational approach puts the workload on the user who can mitigate this problem by finding suitable feature transformations of the MD coordinates, e.g., to contact maps, distances, angles or other other nonlinear functions in which the metastable states may be linearly separable. In a deep TAE, instead, we take the perspective that the nonlinear feature transformation should be found automatically by an optimization algorithm. Our results on toy models and MD data indicate that this is indeed possible and low-dimensional representations can be found that outperform those found by naive TICA and PCA.

Our approach is closely related to the previously proposed VAMPnet approach that performs a simultaneous dimension reduction and MSM estimation by employing the variational approach of Markov processes (Mardt et al., sion). By combining the theoretical results from this paper with those of (Noé and Nüske, 2013; Wu and Noé, 2017), it is clear that in the linear case all these methods are equivalent with TCCA, TICA, Koopman models or MSMs, depending on the type of inputs used, and whether the data are reversible or nonreversible. We believe that there is also a deeper mathematical relationship between these methods in the nonlinear case, e.g., when deep neural networks are employed to learn the feature transformation, but this relationship is still elusive. Both the present approach, that minimizes the TAE regression error in the input space, as well as the variational approach, that maximizes a variational score in the feature space (Noé and Nüske, 2013; Wu and Noé, 2017), are suitable to conduct hyper-parameter search (McGibbon and Pande, 2015). The present error model (6) is based on least square regression, or in other words, on the assumption of additive noise in the configuration space, while VAMPnets do not have this restriction. Also, VAMPnets can incorporate the MSM estimation in a single end-to-end learning framework. On the other hand, the autoencoder approach has the advantage that, in addition to the feature encoding, a feature decoding back to the full configuration space is learned, too. Future studies will investigate the strengths and weaknesses of both approaches in greater detail.

Acknowledgments

We are grateful for insightful discussions with Steve Brunton, Nathan Kutz, Andreas Mardt, Luca Pasquali, and Simon Olsson. We gratefully acknowledge funding by European Commission (ERC StG 307494 “pcCell”) and Deutsche Forschungsgemeinschaft (SFB 1114/A04).

References