Nonlinear Discovery of Slow Molecular Modes using State-Free Reversible VAMPnets

TL;DR

This paper introduces SRV, a deep learning method that discovers nonlinear slow collective variables from molecular simulation data, aiding in understanding and enhancing sampling of complex molecular dynamics.

Contribution

The paper presents state-free reversible VAMPnets (SRV), a novel deep learning architecture that learns explicit, nonlinear, and orthogonal slow eigenfunctions for molecular systems.

Findings

SRV accurately captures slow eigenfunctions in toy models.

SRV identifies meaningful CVs in alanine dipeptide simulations.

SRV effectively models complex protein dynamics.

Abstract

The success of enhanced sampling molecular simulations that accelerate along collective variables (CVs) is predicated on the availability of variables coincident with the slow collective motions governing the long-time conformational dynamics of a system. It is challenging to intuit these slow CVs for all but the simplest molecular systems, and their data-driven discovery directly from molecular simulation trajectories has been a central focus of the molecular simulation community to both unveil the important physical mechanisms and to drive enhanced sampling. In this work, we introduce state-free reversible VAMPnets (SRV) as a deep learning architecture that learns nonlinear CV approximants to the leading slow eigenfunctions of the spectral decomposition of the transfer operator that evolves equilibrium-scaled probability distributions through time. Orthogonality of the learned CVs is naturally imposed within network training without added regularization. The CVs are inherently explicit and differentiable functions of the input coordinates making them well-suited to use in enhanced sampling calculations. We demonstrate the utility of SRVs in capturing parsimonious nonlinear representations of complex system dynamics in applications to 1D and 2D toy systems where the true eigenfunctions are exactly calculable and to molecular dynamics simulations of alanine dipeptide and the WW domain protein.

Equations

$$u_{t+\tau }(x)={\mathcal{T}}_t\circ u_t(x)=\frac{1}{\pi (x)}\int dy{p}_{\tau }^t(y,x)u_t(y)\pi (y)$$

$$u_{t+\tau }(x)=\mathcal{T}\circ u_t(x)=\frac{1}{\pi (x)}\int dy{p}_{\tau }(y,x)u_t(y)\pi (y).$$

$$\pi (x)p_{\tau }(x,y)=\pi (y)p_{\tau }(y,x).$$

$$⟨u_1(x)|\mathcal{T}\circ u_2(x){⟩}_{\pi }$$

$$=$$

$$⟨a|b{⟩}_{\pi }=\int a(x)b(x)\pi (x)dx.$$

$$\mathcal{T}\circ {\psi }_i(x)={\lambda }_i{\psi }_i(x).$$

$$⟨{\psi }_i|{\psi }_j{⟩}_{\pi }={\delta }_{ij}.$$

$${\chi }_t(x)=\sum _i⟨{\psi }_i|{\chi }_t{⟩}_{\pi }{\psi }_i(x).$$

$${\chi }_{t+k\tau }(x)={\mathcal{T}}^k\circ {\chi }_t(x)$$

$$=\sum _i⟨{\psi }_i|{\chi }_t{⟩}_{\pi }{\lambda }_i^k{\psi }_i(x)$$

$$=\sum _i⟨{\psi }_i|{\chi }_t{⟩}_{\pi }\exp \left(-\frac{k\tau }{t_i}\right){\psi }_i(x),$$

$$t_i=-\frac{\tau }{\log {\lambda }_i}.$$

$$u(x)=\sum _{i\ge 1}⟨{\psi }_i|u{⟩}_{\pi }{\psi }_i(x)=\sum _{i\ge 1}{c}_i{\psi }_i(x)$$

$$\tilde{\lambda }=\frac{⟨u\mid \mathcal{T}\circ u{⟩}_{\pi }}{⟨u\mid u{⟩}_{\pi }}=\frac{{\sum }_{i\ge 1}{c}_i^2{\lambda }_i}{{\sum }_{i\ge 1}{c}_i^2}\le \frac{{\sum }_{i\ge 1}{c}_i^2{\lambda }_1}{{\sum }_{i\ge 1}{c}_i^2}={\lambda }_1.$$

$${\tilde{\lambda }}_i=\frac{⟨{\tilde{\psi }}_i\mid \mathcal{T}\circ {\tilde{\psi }}_i{⟩}_{\pi }}{⟨{\tilde{\psi }}_i\mid {\tilde{\psi }}_i{⟩}_{\pi }},$$

$$⟨{\tilde{\psi }}_k\mid {\tilde{\psi }}_i{⟩}_{\pi }=0,0\le k<i.$$

$${\tilde{\psi }}_i=\sum _j{s}_{ij}{\zeta }_j.$$

$$C{s}_i={\tilde{\lambda }}_{i}Q{s}_i,$$

$$C_{jk}=⟨{\zeta }_j(x)|\mathcal{T}\circ {\zeta }_k(x){⟩}_{\pi },$$

$$Q_{jk}=⟨{\zeta }_j(x)|{\zeta }_k(x){⟩}_{\pi }.$$

$$⟨{\tilde{\psi }}_i\mid \mathcal{T}\circ {\tilde{\psi }}_i{⟩}_{\pi }$$

$$=$$

$$\approx$$

$$⟨{\tilde{\psi }}_i\mid {\tilde{\psi }}_i{⟩}_{\pi }$$

$$=$$

$$\approx$$

$${\tilde{\lambda }}_i=\frac{⟨{\tilde{\psi }}_i\mid \mathcal{T}\circ {\tilde{\psi }}_i{⟩}_{\pi }}{⟨{\tilde{\psi }}_i\mid {\tilde{\psi }}_i{⟩}_{\pi }}\approx \frac{\mathrm{E}\left[{\tilde{\psi }}_i(x_t){\tilde{\psi }}_i(x_{t+\tau })\right]}{\mathrm{E}\left[{\tilde{\psi }}_i(x_t){\tilde{\psi }}_i(x_t)\right]}.$$

$$⟨{\tilde{\psi }}_k\mid {\tilde{\psi }}_i{⟩}_{\pi }\approx \mathrm{E}\left[{\tilde{\psi }}_k(x_t){\tilde{\psi }}_i(x_t)\right]=0,0\le k<i.$$

$$C_{jk}=⟨{\zeta }_j(x)|\mathcal{T}\circ {\zeta }_k(x){⟩}_{\pi }\approx \mathrm{E}\left[{\zeta }_j(x_t){\zeta }_k(x_{t+\tau })\right],$$

Full text

I Introduction

Molecular dynamics (MD) simulations have long been an important tool for studying molecular systems by providing atomistic insight into physicochemical processes that cannot be easily obtained through experimentation. A key step in extracting kinetic information from molecular simulation is the recovery of the slow dynamical modes that govern the long-time evolution of system coordinates within a low-dimensional latent space. The variational approach to conformational dynamics (VAC) Noé and Nuske (2013); Nüske et al. (2014) has been successful in providing a mathematical framework through which the eigenfunctions of the underlying transfer operator can be estimated Noé and Clementi (2017); Klus et al. (2018).

A special case of VAC which estimates linearly optimal slow modes from mean-free input coordinates is known as time-lagged independent component analysis (TICA) Pérez-Hernández et al. (2013); Noé and Nuske (2013); Nüske et al. (2014); Noé and Clementi (2015); Noé et al. (2016); Pérez-Hernández and Noé (2016); Schwantes and Pande (2013); Klus et al. (2018). TICA is a widely-used approach that has become a standard step in the Markov state modeling pipeline Husic and Pande (2018). However, it is restricted to form linear combinations of the input coordinates and is unable to learn nonlinear transformations that are typically required to recover high resolution kinetic models of all but the simplest molecular systems. Schwantes et al. address this limitation by applying the kernel trick with TICA to learn non-linear functions of the input coordinates Schwantes and Pande (2015). A special case of a radial basis function kernels was realized by Noé and Nuske in the direct application of VAC using Gaussian functions Noé and Nuske (2013). Kernel TICA (kTICA), however, suffers from a number of drawbacks that have precluded its broad adoption. First, its implementation requires $O(N^{2})$ memory usage and $O(N^{3})$ computation time for a dataset of size $N$, which becomes intractable for large datasets. In practice, this issue can be adequately alleviated by selecting a small number of landmark points to which kernel TICA is applied, and then constructing interpolative projections of the remaining points using the Nyström extension Harrigan and Pande (2017). Second, kTICA is typically highly sensitive to the kernel hyperparameters Harrigan and Pande (2017), and extensive hyperparameter tuning is typically required to obtain acceptable results. Third, the use of the kernel trick and Nyström extension compromises differentiability of the latent space projection. Exact computation of the gradient of any new point requires the expensive calculation of kernel function $K(x,y)$ of new data $x$ with respect to all points $y$ in the training set, and gradient estimations based only on the landmarks are inherently approximate M. Sultan and Pande (2017). Due to their high cost and/or instability of these gradient estimates, the slow modes estimated through kTICA are impractical for use as collective variables (CVs) for enhanced sampling or reduced-dimensional modeling approaches that require exact gradients of the CVs with respect to the input coordinates.

Deep learning offers an attractive alternative to kTICA as means to solve these challenges. Artificial neural networks are capable of learning nonlinear functions of arbitrary complexity Hassoun (1995); Chen and Chen (1995), are generically scalable to large datasets with training scaling linearly with the size of the training set, the network predictions are relatively robust to the choice of network architecture and activation functions, and exact expressions for the derivatives of the learned CVs with respect to the input coordinates are available by automatic differentiation Bastien et al. (2012); Baydin et al. (2018); Paszke et al. (2017); Schmidhuber (2015). A number of approaches utilizing artificial neural networks to approximate eigenfunctions of the dynamical operator have been proposed. Time-lagged autoencoders Wehmeyer and Noé (2018) utilize auto-associative neural networks to reconstruct a time-lagged signal, with suitable collective variables extracted from the bottleneck layer. Variational dynamics encoders (VDEs) Hernández et al. (2018) combine time-lagged reconstruction loss and autocorrelation maximization within a variational autoencoder. While the exact relationship between the regression approach employed in time-lagged autoencoders and the VAC framework is not yet formalized Wehmeyer and Noé (2018), variational autoencoders (VAEs) have already been studied as estimators of the dynamical propagator Hernández et al. (2018) and in furnishing collective variables for enhanced sampling Sultan et al. (2018). The most pressing limitation of VAE approaches to date is their restriction to the estimation of the single leading eigenfunctions of the dynamical propagator. The absence of the full spectrum of slow modes fails to expose the full richness of the underlying dynamics, and limits enhanced sampling calculations in the learned CVs to acceleration along a single coordinate that may be insufficient to drive all relevant conformational interconversions.

In this work, we propose a deep-learning based method to estimate the slow dynamical modes that we term state-free reversible VAMPnets (SRVs). SRVs take advantage of the VAC framework using a neural network architecture and loss function to recover the leading modes of the spectral hierarchy of eigenfunctions of the transfer operator that evolves equilibrium-scaled probability distributions through time. In a nutshell, the SRV discovers a nonlinear featurization of the input basis to pass to the VAC framework for estimation of the leading eigenfunctions of the transfer operator.

This approach shares much technical similarity with, and was in large part inspired by, the elegant VAMPnets approach developed by Noé and coworkers Mardt et al. (2018) and deep canonical correlation analysis (DCCA) approach developed by Livescu and coworkers Andrew et al. (2013). Both approaches employ Siamese neural networks to discover nonlinear featurizations of an input basis that are optimized using a VAMP score. VAMPnets differ from DCCA in optimizing a VAMP-2 rather than VAMP-1 score, making it better suited to applications to time series data due to its theoretical grounding in the Koopman approximation error Mardt et al. (2018). VAMPnets seek to replace the entire MSM construction pipeline of featurization, dimensionality reduction, clustering, and construction of a kinetic model. The objective of SRVs is not to perform direct state space partitioning but rather to learn continuous nonlinear functions of the input data to generate a nonlinear basis set with which to approximate the eigenfunctions of the transfer operator. The design of SRVs differs from that of VAMPnets in two important ways to support this goal. Indeed the name given to our approach is intended to both indicate its heritage with VAMPnets and these distinguishing features. First, the SRV optimizes the VAC as a variational principle for stationary and reversible processes Noé and Nuske (2013); Nüske et al. (2014), whereas VAMPnets employ the more general VAMP principle that applies to non-stationary and non-reversible processes Mardt et al. (2018). As such, SRVs are designed for applications to molecular systems obeying detailed balance where the VAC permits us to take molecular trajectories that may not strictly obey detailed balance and make a biased estimation of the slow eigenfunctions of the reversible dynamics rather than a less biased estimator of the possibly non-reversible dynamics contained in the finite data. Second, VAMPnets employ softmax activations within their output layers to generate $k$-dimensional output vectors that can be interpreted as probability assignment to each of $k$ metastable states. This network architecture achieves nonlinear featurization of the input basis and soft clustering into metastable states. The output vectors are subsequently optimized using the VAMP principle to furnish a kinetic model over these soft/fuzzy states by approximating the eigenfunctions of the transfer operator over the states. Importantly, even though the primary objective of VAMPnets is clustering, the soft state assignments can be used to approximate the transfer operator eigenfunctions using a reweighting procedure. However, since the soft state assignments produced by softmax activations of the output layer are constrained to sum to unity, there is a linear dependence that requires $(k+1)$ output components to identify the $k$ leading eigenfunctions. The second distinguishing feature of SRVs is thus to employ linear or nonlinear activation functions in the output layer of the network. By eschewing any clustering, the SRV is better suited to directly approximating the transfer operator eigenfunctions as its primary objective, although clustering can also be performed in the space of these eigenfunctions in a post-processing step. This seemingly small change has a large impact in the success rate of the network in successfully recovering the transfer operator eigenfunctions. Training neural networks is inherently stochastic and it is standard practice to train multiple networks with different initial network parameters and select the best. Numerical experiments on the small biomolecule alanine dipeptide (see Section III.C) in which we trained 100 VAMPnets and 100 SRVs for 100 epochs employing optimal learning rates showed that both VAMPnets and SRVs were able to accurately recover the leading eigenfunctions and eigenvalues in quantitative agreement with one another, but that VAMPnets exhibited a 29% success rate in doing so compared to 70% for SRVs. Accordingly, the architecture of SRVs is better suited to direct estimation of the transfer operator eigenfunctions and may be preferred when the objective is to estimate these functions as continuous, explicit, and differentiable functions of the input coordinates that can be used to infer the mechanisms of molecular conformational transitions, employed directly in enhanced sampling calculations, and passed to standard MSM construction pipelines to perform microstate clustering and estimation of discrete kinetic models.

The structure of this paper is as follows. We first derive the theoretical foundations of the SRV as a special case of VAC, and then demonstrate its efficacy against kTICA and state-of-the-art TICA-based MSMs in applications to 1D and 2D toy systems where the true eigenfunctions are known and in molecular simulations of alanine dipeptide and WW domain.

II Methods

We first recapitulate transfer operator theory and the variational approach to conformational dynamics (VAC) Noé and Nuske (2013); Nüske et al. (2014); Schütte et al. (2001); Prinz et al. (2011); Schwantes and Pande (2015), choosing to borrow the notational convention from Ref. Prinz et al. (2011). We then demonstrate how the VAC specializes to TICA, kTICA, and SRVs.

II.1 Transfer operator theory

Given the probability distribution of a system configuration $p_{t}(x)$ at time $t$ and the equilibrium probability distribution $\pi(x)$, we define $u_{t}(x)=p_{t}(x)/\pi(x)$ and the transfer operator $\mathscr{T}_{t}=\mathscr{T}_{t}(\tau)$, known formally as the Perron-Frobenius operator or propagator with respect to the equilibrium density Klus et al. (2018), such that

[TABLE]

where $p^{t}_{\tau}(y,x)=\mathbb{P}(x_{t+\tau}=x|x_{t}=y)$ is a transition density describing the probability that a system at $y$ at time $t$ evolves to $x$ after a lag time $\tau$. In general, $p^{t}_{\tau}(y,x)$ depends on not only current state $y$ at time $t$, but also previous history, and is therefore time dependent. Under the Markovian assumption, which becomes an increasingly good approximation at larger lag times $\tau$, $p^{t}_{\tau}(y,x)$ becomes a time homogeneous transition density $p_{\tau}(y,x)$ independent of $t$ and the transfer operator $\mathscr{T}_{t}$ can be written as $\mathscr{T}$, where

[TABLE]

If the system is at equilibrium, then it additionally obeys detailed balance such that

[TABLE]

Given any two state functions $u_{1}(x)$ and $u_{2}(x)$, we appeal to Eq. 2 and 3 to write

[TABLE]

which demonstrates that $\mathscr{T}$ is self-adjoint with respect to the inner product

[TABLE]

Let $\{\psi_{i}(x)\}$ be eigenfunctions of $\mathscr{T}$ corresponding to the eigenvalues $\{\lambda_{i}\}$ in non-ascending order

[TABLE]

The self-adjoint nature of $\mathscr{T}$ implies that it possesses real eigenvalues $\{\lambda_{i}(x)\}$ and its eigenvectors $\{\psi_{i}(x)\}$ form a complete orthonormal basis Wu and Noé (2017); Nüske et al. (2014); Noé and Nuske (2013), with orthonormality relations

[TABLE]

Normalization of the transition density $\int dxp_{\tau}(y,x)=1$ together with the assumption of ergodicity implies that the eigenvalue spectrum is bounded from above by a unique unit eigenvalue such that $1=\lambda_{0}>\lambda_{1}\geq\lambda_{2}\geq\ldots$ Nüske et al. (2014); Prinz et al. (2011). Any state $\chi_{t}(x)$ at a specific time $t$ can be written as a linear expansion in this basis of $\{\psi_{i}\}$

[TABLE]

The evolution of $\chi_{t}(x)$ to $\chi_{t+k\tau}(x)$ after a time period $k\tau$ can be written as

[TABLE]

where $t_{i}$ is the implied timescale corresponding to eigenfunction $\psi_{i}$ given by

[TABLE]

This development makes clear that the eigenvalue associated with an eigenfunction characterizes its temporal autocorrelation. The pair $\{\psi_{0},\lambda_{0}=1\}$ corresponds to the equilibrium distribution. Smaller positive eigenvalues in the range $1>\lambda>0$ decay increasingly faster in time. The self-adjointness of $\mathscr{T}$ assure all eigenvalues are real but does not prohibit negative eigenvalues, which are mathematically admissible but unphysical on the grounds that they are measures of autocorrelation. Accordingly, negative eigenvalues are rarely observed for well trained models for which sufficient training data is available at sufficiently high temporal resolution, and their appearance is a numerical indication that the “slow modes” identified by the approach cannot be adequately resolved from the data and/or model at hand.

II.2 Variational approach to conformational dynamics (VAC)

Under the VAC, we seek an orthonormal set $\{\tilde{\psi}_{i}\}$ to approximate $\{\psi_{i}\}$ under the orthogonality conditions given by Eq. 6. Typically we are interested not in the full $\{\psi_{i}\}$ but only the leading eigenfunctions corresponding to the largest eigenvalues and thus longest implied timescales.

We first observe that $\psi_{0}(x)=1$ is a trivial eigenfunction of $\mathscr{T}$ with eigenvalue $\lambda_{0}=1$ corresponding to the equilibrium distribution at $t\rightarrow\infty$. This follows from Eq. 2 and 3 whereby $\mathscr{T}\circ\psi_{0}(x)=\mathscr{T}\circ 1=\frac{1}{\pi(x)}\int dy\ p_{\tau}(y,x)\pi(y)=\frac{1}{\pi(x)}\int dy\ p_{\tau}(x,y)\pi(x)=1=\psi_{0}(x)$.

To learn $\psi_{1}(x)$, we note that any state function $u(x)$ which is orthogonal to $\psi_{0}(x)$ can be expressed as

[TABLE]

where $c_{i}=\langle\psi_{i}|u\rangle_{\pi}$ are expansion coefficients, and

[TABLE]

Since $\tilde{\lambda}$ is bounded from above by $\lambda_{1}$ by the variational principle Nüske et al. (2014); Schwantes and Pande (2015), we can exploit this fact to approximate the first non-trivial eigenfunction $\psi_{1}$ by searching for a $u$ that maximizes $\tilde{\lambda}$ subject to $\langle u\mid\psi_{0}\rangle_{\pi}=0$. The learned $u$ is an approximation to first non-trivial eigenfunction $\tilde{\psi}_{1}$.

We can continue this procedure to approximate higher order eigenfunctions. In general we approximate $\tilde{\psi}_{i}(x)$ by maximizing

[TABLE]

under the orthogonality constraints

[TABLE]

In essence, the VAC procedure combines a variational principle Nüske et al. (2014) with a linear variational approach perhaps most familiar from quantum mechanics Szabo and Ostlund (2012). Given an arbitrary input basis $\{\zeta_{j}(x)\}$, the eigenfunction approximations may be written as linear expansions

[TABLE]

Adopting this basis, the VAC can be shown to lead to a generalized eigenvalue problem analogous to the quantum mechanical Roothaan-Hall equations Nüske et al. (2014); Szabo and Ostlund (2012)

[TABLE]

where

[TABLE]

[TABLE]

Here $s_{i}$ is the (eigen)vector of linear expansion coefficients for the approximate eigenfunction $\tilde{\psi}_{i}$, and $\tilde{\lambda}_{i}$ is the associated eigenvalue. The spectrum of solutions of Eq. 15 yield the best linear estimations of the eigenfunctions of the transfer operator $\mathscr{T}$ within the basis $\{\zeta_{j}(x)\}$. The generalized eigenvalue problem can be solved by standard techniques Watkins (2007).

Eq. 12 and 13 serve as the central equations for TICA, kTICA and SRVs. We first show how these equations can be estimated from simulated data, and then how these three methods emerge as specializations of VAC under particular choices for the input basis.

II.3 Estimation of VAC equations from trajectory data

Here we show how Eq. 12 and 13 can be estimated from empirical trajectory data Noé and Nuske (2013); Nüske et al. (2014); Wu and Noé (2017). The numerator of Eq. 12 becomes

[TABLE]

where $\mathrm{E}\left[\tilde{\psi}_{i}(x_{t})\tilde{\psi}_{i}(x_{t+\tau})\right]$ can be estimated from a trajectory $\{x_{t}\}$. The denominator follows similarly as

[TABLE]

The full expression for Eq. 12 becomes

[TABLE]

Similarly, Eq. 13 becomes

[TABLE]

Using the same reasoning, the components (Eq. 16 and 17) of the generalized eigenvalue problem (Eq. 15) are estimated as

[TABLE]

[TABLE]

II.4 Time-lagged independent component analysis (TICA)

In TICA, we represent $\tilde{\psi}_{i}(x)$ as a linear combination of molecular coordinates $x$, where $a_{i}$ is a vector of linear expansion coefficients and $C$ is an additive constant

[TABLE]

The orthogonality condition Eq. 13 of $\tilde{\psi}_{i}$ relative to $\tilde{\psi}_{0}(x)=1$ becomes

[TABLE]

It follows that

[TABLE]

and therefore Eq. 24 can be written as

[TABLE]

where $\delta x=x-\mathrm{E}\left[x\right]$ is a mean-free coordinate. Under this specification for $\tilde{\psi}_{i}(x)$, Eq. 12 and Eq. 13 become

[TABLE]

[TABLE]

which are exactly the objective function and orthogonality constraints of TICA Schwantes and Pande (2013); Pérez-Hernández et al. (2013).

II.5 Kernel TICA (kTICA)

One way of generalizing TICA to learn nonlinear features is though feature engineering. Specifically, if we can find a nonlinear mapping $\phi$ that maps configurations $x$ to appropriate features $\phi(x)$, we can apply TICA on these features. However, designing good nonlinear features typically requires expert knowledge or expensive data preprocessing techniques. Therefore, instead of finding an explicit mapping $\phi$, an alternative approach is to apply the kernel trick using a kernel function $K(x,y)=\phi(x)\cdot\phi(y)$ that defines an inner product between $\phi(x)$ and $\phi(y)$ as a similarity measure in the feature space that does not require explicit definition of $\phi$.

To apply the kernel trick to TICA, we need to reformulate TICA in terms of this kernel function. It can be shown that in Eq. II.4, the coefficient $a_{i}$ is linear combination of $\{\delta x_{t}\}\cup\{\delta x_{t+\tau}\}$ (see Supplementary Information of Ref. Schwantes and Pande (2015)) and may therefore be written as

[TABLE]

Under this definition for $a_{i}$ Eq. II.4 and 30 become,

[TABLE]

[TABLE]

Now the objective function Eq. II.5 and constraints Eq. II.5 only depend on inner products between any pair of elements in $\{\delta x_{t}\}\cup\{\delta x_{t+\tau}\}$.

To obtain a nonlinear transformation, we replace the linear similarity measure, which is the inner product $\delta x\cdot\delta y$ of two vectors, with a symmetric nonlinear kernel function $K(x,y)$. This transforms Eq. II.5 and II.5 to

[TABLE]

[TABLE]

which define the objective function and orthogonality constraints of kTICA. As detailed in Ref. Schwantes and Pande (2015), Eq. II.5 and II.5 can be simplified to a generalized eigenvalue problem that admits efficient solution by standard techniques Watkins (2007).

Although kTICA enables recovery of nonlinear eigenfunctions, it does have some significant drawbacks. First, it has high time and space complexity. The Gram matrix $K=[K(x,y)]_{N\times N}$ takes $O(N^{2})$ time to compute and requires $O(N^{2})$ memory, and the generalized eigenvalue problem takes $O(N^{3})$ time to solve, which severely limits its application to large datasets. Second, results can be sensitive to the choice of kernel and there exist no rigorous guidelines for the choice of an appropriate kernel for a particular application. This limits the generalizability of the approach. Third, the kernels are typically highly sensitive to the choice of hyperparameters. For example, the Gaussian (radial basis function) kernel is sensitive to noise for small kernel widths $\sigma$, which leads to overfitting and overestimation the implied timescales. A large $\sigma$ on the other hand typically approaches linear TICA results which undermines its capacity to learn nonlinear transformations Harrigan and Pande (2017); Sultan et al. (2018). This hyperparameter sensitivity typically requires signifiant expenditure of computational effort to tune these values in order to obtain satisfactory results Schwantes and Pande (2015); Harrigan and Pande (2017). Fourth, kTICA does not furnish an explicit expression for the mapping $\phi:x\rightarrow\phi(x)$ projecting configurations into the nonlinear feature space Hernández et al. (2018). Accordingly, it is not straightforward to apply kernel TICA within enhanced sampling protocols that require the learned latent variables to be explicit and differentiable functions of the input coordinates.

One way to ameliorate the first deficiency by reducing memory usage and computation time is to employ a variant of kTICA known as landmark kTICA. This approach selects $m<<N$ landmarks from the dataset, computes an $m$-by-$m$ Gram matrix, and then uses the Nyström approximation to estimate the original $N$-by-$N$ Gram matrix Harrigan and Pande (2017).

II.6 State-free reversible VAMPnets (SRVs)

In general, the eigenfunction approximations need not be a linear function within either the input coordinate space or the feature space. The most general form of Eq. 12 and 13 is

[TABLE]

[TABLE]

We now introduce the SRV approach that employs a neural network $f$ to learn nonlinear approximations to $\{\tilde{\psi}_{i}\}$ directly without requiring the kernel trick. The neural network $f$ maps configuration $x$ to a $n$-dimensional output $f_{i}(x)\ (i=1,...,n)$, where $n$ is the number of slow modes we want to learn. Then a linear variational method is applied to obtain the corresponding $\tilde{\psi}_{i}(x)$ such that $\{\tilde{\psi}_{i}(x)\}$ form an orthonormal basis set that minimizes the network loss function.

The method proceeds as follows. Given a neural network $f$ with $n$-dimensional outputs, we feed a training set $X=\{x\}$ and train the network with loss function $L$

[TABLE]

where $g$ is a monotonically decreasing function. Minimizing $L$ is equivalent to maximizing the sum over $\tilde{\lambda}_{i}$ and therefore maximizing the sum over $t_{i}$ (Eq. 9). The $\{\tilde{\psi}_{i}\}$ correspond to linear combinations of the neural network outputs $\{f_{i}(x)\}$

[TABLE]

computed by applying the linear VAC to the neural network outputs Noé and Nuske (2013); Nüske et al. (2014). The linear VAC is equivalent to the generalized eigenvalue problem in Eq. 15 where $\zeta_{j}(x)=f_{j}(x)$.

Minimization of the loss function by gradient descent requires the derivative of $L$ with respect to neural network parameters $\theta$

[TABLE]

The first two partial derivatives are straightforwardly calculated by automatic differentiation once a choice for $g$ has been made. The third partial derivative requires a little more careful consideration. We first expand this final derivative using Eq. 15 to make explicit the dependence of $\tilde{\lambda}_{i}$ on the matrices $C$ and $Q$

[TABLE]

To our best knowledge, no existing computational graph frameworks provide gradients for generalized eigenvalue problems. Accordingly, we rewrite Eq. 15 as follows. We first apply a Cholesky decomposition to $Q$, such that

[TABLE]

where $L$ is a lower triangular matrix. We then left multiply both sides by $L^{-1}$ to obtain

[TABLE]

Defining $\tilde{C}=L^{-1}C(L^{T})^{-1}$ and $\tilde{s}_{i}=L^{T}s_{i}$ we convert the generalized eigenvalue problem into a standard eigenvalue with a symmetric matrix $\tilde{C}$

[TABLE]

where the Cholesky decomposition assures numerical stability. Now, Eq. 41 becomes

[TABLE]

where all terms are computable using automatic differentiation: $\frac{\partial\tilde{\lambda}_{i}}{\partial\tilde{C}}$ from routines for a symmetric matrix eigenvalue problem, $\frac{\partial\tilde{C}}{\partial C}$ and $\frac{\partial\tilde{C}}{\partial Q}$ from those for Cholesky decomposition, matrix inversion, and matrix multiplication, and $\frac{\partial C}{\partial\theta}$ and $\frac{\partial Q}{\partial\theta}$ by applying the chain rule to Eq. 16 and 17 with $\zeta_{j}(x)=f_{j}(x)$ and computing the derivatives $\frac{\partial f}{\partial\theta}$ through the neural network.

Training is performed by passing $\{x_{t},x_{t+\tau}\}$ pairs to the SRV and updating the network parameters using mini-batch gradient descent using Adam Kingma and Ba (2014) and employing the automatic differentiation expression for $\partial L/\partial\theta$ to minimize the loss function. To prevent overfitting, we shuffle all $\{x_{t},x_{t+\tau}\}$ pairs and reserve a small portion as validation set with which to implement early stopping. Training is terminated when validation loss no longer decreases for a pre-specified number of epochs. A schematic diagram of the SRV is shown in Fig. 1.

We note that we do not learn $\{\tilde{\psi}_{i}\}$ directly in our neural network, but obtain it as a weighted linear combination of $\{f_{i}\}$. Specifically, during training we learn not only the weights of neural network, but also the linear expansion coefficients $\{s_{ij}\}$ that yield $\{\tilde{\psi}_{i}\}$ from $\{f_{i}\}$. Conceptually, one can consider the neural network as learning an optimal nonlinear featurization of the input basis to pass through the VAC framework. After training is complete, a new out-of-sample configuration $x$ can be passed through the neural network to produce $f$, which is then transformed via a linear operator $\{s_{ij}\}$ to get the eigenfunctions estimates $\{\tilde{\psi}_{i}\}$. Since the neural network is fully differentiable and the final transformation is linear the SRV mapping $x\rightarrow\{\tilde{\psi}_{i}\}$ is explicit and differentiable, making it well suited to applications in enhanced sampling.

An important choice in our network design is the function $g(\tilde{\lambda})$ within our loss function. In theory it can be any monotonically decreasing function, but motivated by Refs. Mardt et al. (2018) and Wu and Noé (2017), we find that choosing $g(\tilde{\lambda})=-\tilde{\lambda}^{2}$ such that the loss function corresponds to the VAMP-2 score yields good performance and possesses strong theoretical grounding. Specifically, the VAMP-2 score may be interpreted as the cumulative kinetic variance Noé and Clementi (2015); Husic and Pande (2018) analogous to the cumulative explained variance in principal component analysis Pearson (1901), but where the VAMP-2 score measures kinetic rather than conformational information. The VAMP-2 score can also be considered to measure the closeness of the approximate eigenfunctions to the true eigenfunctions of the transfer operator, with this score achieving its maximum value when the approximations become exact Mardt et al. (2018). This choice may also be generalized to the VAMP-r score $g(\tilde{\lambda})=-\tilde{\lambda}^{r}$ Wu and Noé (2017). We have also observed good performance using $g(\tilde{\lambda})=1/\log(\tilde{\lambda}$), which corresponds to maximizing the sum of implied timescales $\sum_{i}t_{i}$ (Eq. 9).

We also note that it is possible that the transfer operator may possess degenerate eigenvalues corresponding to associated eigenfunctions with identical implied timescales. True degeneracy is expected to be rare within the leading eigenfunctions but may arise from symmetries in the system dynamics; approximate degeneracy may be encountered wherein timescales become indistinguishable within the training data available to the model. In principle, either of these situations could lead to numerical difficulties wherein degenerate eigenfunctions cannot be stably resolved due to the existence of infinitely many equivalent linear combinations. In such instances, all eigenfunctions associated with repeated eigenvalues should be retained and mutually aligned within successive rounds of training by determining the optimal rotation matrix within their linear subspace. In practice, we suspect that minor differences in the implied timescales induced by the finite nature of the training data will break the degeneracy of the eigenvalues and allow for stable numerical recovery. A post hoc analysis of any groups of close eigenvalues can then be performed to attempt to resolve the root of any suspected underlying symmetry.

The SRV learning protocol is quite simple and efficient. It only requires $O(N)$ memory and $O(N)$ computation time, which makes it ideal for large datasets. It also does not require selection of a kernel function to achieve appropriate nonlinear embeddings Sultan et al. (2018); Hernández et al. (2018) Instead we appeal to the universal approximation theorem Hassoun (1995); Chen and Chen (1995), which loosely states that a neural network with more than one hidden layer and enough number of hidden nodes can approximate any nonlinear continuous function without need to impose a kernel. Lifting the requirement for kernel selection and its attendant hyperparameter tuning is a significant advantage of the present approach over kernel-based techniques. Training such a simple neural network is possible using standard techniques such as stochastic gradient descent and is largely insensitive to our choice of network architecture, hyperparameters, and activation functions. Using a default learning rate, large batch size, and sufficiently large network gives excellent results as we demonstrate below. Furthermore, we use the same number of hidden layers, number of nodes in each hidden layer, and hidden layer activation functions for all four applications in this work, demonstrating the simplicity, robustness, and generalizability of the SRV approach.

III Results

III.1 1D 4-well potential

In our first example, we consider a simple 1D 4-well potential defined in Ref. Schwantes and Pande (2015) and illustrated in Fig. 2. The eigenfunctions of the transfer operator for this system are exactly calculable, and this provides a simple initial demonstration of SRVs. We construct a transition matrix as a discrete approximation for the transfer operator by dividing the interval $[-1,1]$ into 100 evenly-spaced bins and computing the transition probability $p_{ij}$ of moving from bin $i$ to bin $j$ as,

[TABLE]

where $V_{j}$ and $V_{i}$ are the potential energies at the centers of bins $j$ and $i$ and $C_{i}$ is the normalization factor for bin $i$ such that the total transition probability associated with bin $i$ sums to unity. We then define a unit-time transition matrix $P(1)=[p_{ij}]_{100\times 100}$, and a transition matrix of lag time $\tau$ as $P(\tau)=P(1)^{\tau}$. In this model, we use a lag time $\tau=100$ for both theoretical analysis and model learning.

By computing the eigenvectors $\{v_{i}\}$ of the transition matrix $P(\tau=100)$, we recover the equilibrium distribution $\pi=v_{0}$. The corresponding eigenfunctions of transfer operator $\mathscr{T}(\tau=100)$ are given by $\psi_{i}=v_{i}/\pi$. The top four eigenfunctions are shown in Fig. 3, where first eigenfunction $\psi_{0}$ corresponds to trivial stationary transition, and next three to transitions over the three energy barriers.

Using the calculated transition matrix $P(1)$, we generate a 5,000,000-step trajectory over the 1D landscape by initializing the system within a randomly selected bin and then propagating its dynamics forward through time under the action of $P(1)$ Schwantes and Pande (2015). The state of the system at any time $t$ is represented by the 1D $x$-coordinate of the particle $x(t)\in\mathbb{R}^{1}$. We then adopt a fixed lag time of $\tau=100$ and learn the top three non-trivial eigenfunctions using both kTICA and SRV. We note that for kTICA, we cannot compute the full 5,000,000-by-5,000,000 Gram matrix, so instead we select 200 landmarks by K-means clustering and use the landmark approximation. We select a Gaussian kernel with $\sigma=0.05$. In contrast, the SRV has no difficulty processing all data points. We use two hidden layers with 100 neurons each, giving a final architecture of [1, 100, 100, 3]. The activation function for all layers are selected to be $\tanh(x)$, and we employ a VAMP-2 loss function. The SRV network is constructed and trained within Keras Chollet .

The results for kTICA and SRV are shown in Fig. 4. We find that both methods are in excellent agreement with the theoretical eigenfunctions. The small deviation between the estimated timescales for both methods and the theoretical timescales is a result of noise in the simulated data. This result demonstrates the capacity of SRVs to recover the eigenfunctions of a simple 1D system in quantitative agreement with the theoretical results and excellent correspondence with kTICA.

III.2 Ring potential

We now consider the more complicated example of a 2D modified ring potential $V(r,\theta)$. This potential contains a narrow ring potential valley of 0 $k_{B}T$ and four barriers of heights 1.0 $k_{B}T$, 1.3 $k_{B}T$, 0.5 $k_{B}T$, and 8.0 $k_{B}T$. The expression of the potential is given by Eq. 47, and it is plotted in Fig. 5

[TABLE]

We use the same procedure outlined in the previous example to generate the theoretical eigenfunctions of transfer operator and simulate a trajectory using a Markov state model. The region of interest, $[-1,1]\times[-1,1]$, is discretized into 50-by-50 bins. In this model, we use a lag time of $\tau=100$ for both theoretical analysis and model learning. The transition probability $p_{ij}$ of moving from bin $i$ to bin $j$ is given by Eq. 48, where $C_{i}$ is the normalization factor for bin $i$ such that the total transition probability associated with bin $i$ sums to unity.

[TABLE]

Once again we compute the first three non-stationary theoretical eigenfunctions of $\mathscr{T}(\tau=100)$ from the transition matrix and illustrate these in Fig. 6. We then numerically simulate a 5,000,000-step trajectory over the 2D landscape under the action of the transition matrix, and pass these data to a 1000-landmark kTICA model employing a Gaussian kernel and an SRV with the same hidden layer architecture and loss function as the previous example. The state of the system at any time $t$ is defined by the 2D ($x$,$y$)-coordinate pair representing the particle location $(x(t),y(t))\in\mathbb{R}^{2}$. Again small deviations between the estimated and theoretical timescales should be expected due to noise and finite length of the simulated data.

The kTICA results employing the optimal bandwidth $\sigma$ of the Gaussian kernel are shown in Fig. 6. Although it gives a reasonable approximation of the eigenfunctions within the ring where data are densely populated, the agreement outside of the ring is much worse. This is due to the intrinsic limitation of a Gaussian kernel function: a small $\sigma$ leads to an accurate representation near landmarks but poor predictions for regions far away, while a large $\sigma$ produces better predictions far away at the expense of local accuracy. Moreover, the kTICA results depend sensitively on both the number of landmarks and the bandwidth of the Gaussian kernel. In Fig. 7 we report the test loss given by Eq. 38 on a dynamical trajectory of length 1,000,000 for kTICA models with different kernel bandwidths $\sigma$ and numbers of landmarks selected by K-means clustering. The approximations of the leading non-trivial eigenfunctions are reported in Fig. 8. We note that only when we use a large number of landmarks can we achieve reasonable results, which leads to expensive computations in both landmark selection and calculation of the Gram matrix. Moreover, even with a large number of landmarks, the range of $\sigma$ values for which satisfactory results are obtained is still quite small and requires substantial tuning.

In contrast, the SRV with architecture of [2, 100, 100, 3] shows excellent agreement with the theoretical eigenfunctions without any tuning of network architecture, activation functions, or loss function (Fig. 6). The SRV eigenvalues closely match those extracted by kTICA, and the SRV eigenfunctions show superior agreement with the theoretical results compared to kTICA. This result demonstrates the capacity of SRVs in leveraging the flexibility and robustness of neural networks in approximating arbitrary continuous function over compact sets.

III.3 Alanine dipeptide

Having demonstrated SRVs on toy models, we now consider their application to alanine dipeptide in water as a simple but realistic application to molecular data. Alanine dipeptide (N-acetyl-L-alanine-N′-methylamide) is a simple 22-atom peptide that stands as a standard test system for new biomolecular simulation methods. The molecular structure of alanine dipeptide annotated with the four backbone dihedral angles that largely dictate its configurational state is presented in Fig. 9. A 200 ns simulation of alanine dipeptide in TIP3P water and modeled using the Amber99sb-ILDN forcefield was conducted at T = 300 K and P = 1 bar using the OpenMM 7.3 simulation suite Eastman et al. (2012, 2017). Lennard-Jones interactions were switched smoothly to zero at a 1.4 nm cutoff, and electrostatics treated using particle-mesh Ewald with a real space cutoff of 1.4 nm and a reciprocal space grid spacing of 0.12 nm. Configurations were saved every 2 ps to produce generate a trajectory comprising 1,000,000 configurations. The instantaneous state of the peptide is represented by the Cartesian coordinates of the 22 atoms $x(t)\in\mathbb{R}^{66}$, where the influence of the solvent molecules is treated implicitly through their influence on the peptide configuration. In this case the theoretical eigenfunctions of the transfer operator are unavailable, and we instead compare the SRV results against those of kTICA.

The 45 pairwise distances between the 10 heavy atoms were used as features with which to perform kTICA employing employing a Gaussian kernel, 5000 landmarks obtained from K-means clustering, and a lag time of $\tau$ = 20 ps. The intramolecular pairwise distances or contact matrix are favored within biomolecular simulations as an internal coordinate frame representation that is invariant to translational and rotational motions Sittel et al. (2014). The leading three eigenfunctions $\tilde{\psi}_{i}\ (i=1,2,3)$ discovered by kTICA employing a manually tuned kernel bandwidth are shown in Fig. 10 superposed upon the Ramachandran plot in the backbone $\phi$ and $\psi$ torsional angles that are known to be good discriminators of the molecular metastable states Ferguson et al. (2011); Ren et al. (2005); Laio and Parrinello (2002); Barducci et al. (2008); Bolhuis et al. (2000); Valsson and Parrinello (2014); Stamati et al. (2010). The timescales of the 4th and higher order modes lie below the $\tau$ = 20 ps lag time so cannot be resolved by this model. Accordingly, we select three leading slow modes for analysis. From Fig. 10 it is apparent that the first slow mode captures transitions along $\phi$ torsion, the second characterizes the transitions between $\alpha$ and $(\beta,P_{//})$ basins, and the third motions between the $\alpha_{L}$ and $\gamma$ basins.

The trajectory is then analyzed at the same lag time using a [45, 100, 100, 3] SRV employing the same hidden layer architecture and loss function as the previous examples. The SRV eigenfunctions illustrated in Fig. 10 are in excellent agreement with those learned by kTICA, but importantly the implied timescale of the SRV leading mode is 17% slower than that extracted by kTICA. What is the origin of this discrepancy?

The current state-of-the-art methodology to approximate the eigenfunctions of the transfer operator is to construct a Markov state model (MSM) over a microstate decomposition furnished by TICA Pande et al. (2010); Schwantes and Pande (2013); Pande et al. (2010); Husic and Pande (2018); Trendelkamp-Schroer et al. (2015); Sultan and Pande (2017); Mittal and Shukla (2018); Harrigan et al. (2017); Wehmeyer et al. (2018); Scherer et al. (2018) Applying TICA to the 45 pairwise distances between the heavy atoms, we construct a MSM using PyEMMA Scherer et al. (2015), and present the results in Fig. 10. We see that the MSM eigenfunctions are in excellent accord with those learned by SRVs and kTICA, however while SRVs nearly exactly match the implied timescales of the MSM and results reported in Ref. Trendelkamp-Schroer et al. (2015), kTICA substantially underestimates the timescale of the slowest mode.

The underlying reason for this failure is the that spatial resolution in feature space is limited by number of landmarks, and if the Euclidean distance in feature space does not closely correspond to kinetic distances then it requires much finer spatial resolution to resolve the correct slow modes. For alanine dipeptide, pairwise distances between heavy atoms are not well correlated with the slowest dynamical modes – here, rotations around backbone dihedrals – and therefore landmark kTICA models built on these features have difficulty resolving the slowest mode. This issue may be alleviated by employing more landmarks, but it quickly becomes computationally intractable on commodity hardware to use significantly more than approximately 5000 landmarks. Another option is to use features that are better correlated with the slow modes. For alanine dipeptide, it is known that the backbone dihedrals are good features, and if we perform kTICA using these input features we do achieve much better estimates of the implied timescale of the leading mode ($t_{1}$ = 1602 ps). In general, however, a good feature set is not known a priori, and for poor choices it is typically not possible to obtain good performance even for large numbers of landmarks.

Our numerical investigations also show the implied timescales extracted by SRVs to be very robust to the particular choice of lag time. The reliable inference of implied timescales from MSMs requires that they be converged with respect to lag time, and slow convergence presents an impediment to the realization of high-time resolution MSMs. The yellow bars in Fig. 11 present the implied time scales of the leading three eigenfunctions computed from the mean over five SRVs constructed with lag times of $\tau$ = 10, 40, and 100 ps. It is clear that the implied timescales are robust to the choice of lag time over a relatively large range.

Given the implied timescales evaluated at four different lag times – $\tau$ = 10, 20, 40, and 100 ps – this provides an opportunity to assess the dynamical robustness of SRVs by subjecting them to a variant of the Chapman-Kolmogorov test employed in the construction and validation of MSMs Mardt et al. (2018). As discussed in Section. II.1, the VAC framework is founded on the Markovian assumption that the transfer operator is homogeneous in time. In previous two toy examples, this assumption holds by construction due to the way the trajectory data were generated. However, for a real system like alanine dipeptide, there is no such guarantee. Here we test a necessary condition for the Markovian assumption to hold. In general we have

[TABLE]

If the Markovian assumption holds, then the transfer operators are independent of $t$, such that

[TABLE]

The corresponding eigenvalue and eigenfunctions are

[TABLE]

and

[TABLE]

where $\{\tilde{\lambda}_{i,\tau}\}$ and $\{\tilde{\psi}_{i,\tau}(x)\}$ are the estimated eigenvalues and eigenfunctions of $\mathscr{T}(\tau)$, and $\{\tilde{\lambda}_{i,k\tau}\}$ and $\{\tilde{\psi}_{i,k\tau}(x)\}$ those of $\mathscr{T}(k\tau)$. Appealing to Eq. 50, it follows that

[TABLE]

providing a means to compare the consistency of SRVs constructed at different choices of $\tau$. In particular, the implied timescales for the eigenfunctions of $\mathscr{T}(k\tau)$ estimated from an SRV constructed at a lag time $k\tau$

[TABLE]

should be well approximated by those estimated from an SRV constructed at a lag time $\tau$

[TABLE]

If this is not the case, then the assumption of Markovianity is invalidated for this choice of lag time.

We present in Fig. 11 the predicted implied timescales over the range of lag times $\tau$ = 2–120 ps calculated from an SRV constructed at a lag time of $\tau$ = 20 ps. These predictions are in excellent accord to the implied timescales directly computed from SRVs constructed at lag times of $\tau$ = 10, 40, and 100 ps, demonstrating satisfaction of the Chapman-Kolmogorov test and a demonstration of the Markovian property of the system at lag times $\tau\gtrsim 10$ ps Mardt et al. (2018); Wehmeyer et al. (2018); Chodera et al. (2006); Pande et al. (2010); Husic and Pande (2018).

III.4 WW domain

Our final example considers a 1137 $\mu s$ simulation of the folding dynamics of the 35-residue WW domain protein performed in Ref. Lindorff-Larsen et al. (2011). We use 595 pairwise distances of all Cα atoms to train a TICA-based MSM and a [595, 100, 100, 2] SRV with the same hidden layer architecture and loss function as all previous examples. We use lag time of $\tau$ = 400 ns (2000 steps) for both models, and focus on the two leading slowest modes. The implied timescales of higher-order modes lie close to or below the lag time and so cannot be reliably resolved by this model. The slow modes discovered by the TICA-based MSM and the SRV are shown in Fig. 12a projected onto the two leading TICA eigenfunctions (tIC1, tIC2) to furnish a consistent basis for comparison. The MSM and SRV eigenfunctions are in excellent agreement, exhibiting Pearson correlation coefficients of $\rho$ = 0.99, 0.98 for the top two modes respectively. The implied timescales inferred by the two methods are also in good agreement.

An important aspect of model quality is the convergence rate of the implied timescales with respect to the lag time $\tau$. The lag time must be selected to be sufficiently large such that the state decomposition is Markovian whereby dynamical mixing with states is faster than the lag time and interconversion between states is slower, but it is desirous that the lag time be as short as possible to produce a model with high time resolution Pande et al. (2010). Better approximations for the leading eigenfunctions of the transfer operator typically lead to convergence of the implied timescales at shorter lag times. We construct 10 independent SRV models and 10 independent TICA-based MSMs over the WW domain trajectory data and report in Fig. 12b the mean and 95% confidence intervals of the implied timescale convergence with lag time. The SRV exhibits substantially faster convergence than the TICA-based MSM, particularly in the second eigenfunction. This suggests that the eigenfunctions identified by the SRV, although strongly linearly correlated with those identified by the TICA-based MSM, provide a better approximation to the leading slow modes of the transfer operator and produce a superior state decomposition. We attribute this observation to the ability of the SRV to capture complex nonlinear relationships in the data within a continuous state representation, whereas the MSM is dependent upon the TICA coordinates which are founded on a linear variational approximation to the transfer operator eigenfunctions that subsequently inform the construction of an inherently discretized microstate transition matrix from which we compute the MSM eigenvectors. This result demonstrates the viability of SRVs to furnish a high-resolution model of the slow system dynamics without the need to perform any system discretization and at a higher time resolution (i.e., lower lag time) than is possible with the current state-of-the-art TICA-based MSM protocol.

IV Conclusions

In this work we proposed a new framework that we term state-free reversible VAMPnets (SRV) for the discovery of a hierarchy of nonlinear slow modes of a dynamical system from trajectory data. The framework is built on top of transfer operator theory that uses a flexible neural network to learn an optimal nonlinear basis from the input representation of the system. Compared to kernel TICA and variational dynamics encoders, our SRV framework has many advantages. It is capable of simultaneously learning an arbitrary number of eigenfunctions while guaranteeing orthogonality. It also requires $O(N)$ memory and $O(N)$ computation time, which makes it amenable to large data sets such as those commonly encountered in biomlecular simulations. The neural network architecture does not require the selection of a kernel function or adjustment of hyperparameters that can strongly affect the quality of the results and be tedious and challenging to tune Sultan et al. (2018); Hernández et al. (2018). In fact, we find that training such a simple fully-connected feed-forward neural network is simple, cheap, and insensitive to batch size, learning rate, and architecture. Finally, the SRV is a parametric model, which provides an explicit and differentiable mapping from configuration $x$ to the learned approximations of the leading eigenvectors of the transfer operator $\{\tilde{\psi}_{i}\}$. These slow collective variables are then ideally suited to be utilized in collective variable-based enhanced sampling methods where the differentiability of the SRV collective variables enable their seamless incorporation with powerful biased sampling techniques such as metadynamics Laio and Parrinello (2002).

The SRV framework possesses a close connection with a number of existing methodologies. In the one-dimensional limit, SRVs are formally equivalent to variational dynamics encoders (VDEs) with an exclusive autocorrelation loss, subject to Gaussian noise Hernández et al. (2018). VDEs however, cannot currently generalize to multiple dimensions due to the lack of an orthogonality constraint on the learned eigenfunctions. By using the more general VAMP principle for non-reversible processes and targeting membership state probabilities rather than learning continuous functions, VAMPnets are obtained Mardt et al. (2018).

In regards to the analysis of molecular simulation trajectories, we anticipate that the flexibility and high time-resolution of SRV models will be of use in helping resolve and understand the important long-time conformational changes governing biomolecular folding and function. Moreover, it is straightforward to replace TICA-based MSMs with SRV-based MSMs, maintaining the large body of theoretical and practical understanding of MSM construction while delivering the advantages of SRVs in improved approximation of the slow modes and superior microstate decomposition. In regards to enhanced sampling in molecular simulation, the differentiable nature of SRV coordinates naturally enables biasing along the SRV collective variables (CVs) using well-established accelerated sampling techniques such as umbrella sampling, metadynamics, and adaptive biasing force. The efficiency of these techniques depends crucially on the choice of “good” CVs coincident with the important underlying dynamical modes governing the ling-time evolution of the system. A number of recent works have employed neural networks to learn nonlinear CVs describing the directions of highest variance within the data Chen and Ferguson (2018); Chen et al. (2018); Ribeiro et al. (2018); Ribeiro and Tiwary (2018). However, the high variance directions are not guaranteed to also correspond to the slow directions of the dynamics. Only variational dynamics encoders have been used to learn and bias sampling in a slow CV Sultan et al. (2018), but, as observed above, the VDE is limited to approximate only the leading eigenfunction of the transfer operator. SRVs open the door to performing accelerated sampling within the full spectrum of all relevant eigenfunctions of the transfer operator. In a similar vein, SRVs may also be profitably incorporated into adaptive sampling approaches that do not apply artificial biasing force, but rather smartly initialize short unbiased simulations on the edge of the explored domain Shamsi et al. (2018); Preto and Clementi (2014); Weber and Pande (2011); Zimmerman and Bowman (2015); Bowman et al. (2010); Hinrichs and Pande (2007); Doerr and De Fabritiis (2014). The dynamically meaningful projections of the data into the SRV collective variables is anticipated to better resolve the dynamical frontier than dimensionality reduction approaches based on maximal preservation of variance, and therefore better direct sampling along the slow conformational pathways. In sum, we expect SRVs to have a variety of applications not just in the context of molecular simulations, but also more broadly within the analysis of dynamical systems.

Acknowledgments

This material is based upon work supported by the National Science Foundation under Grant No. CHE-1841805. H.S. acknowledges support from the Molecular Software Sciences Institute (MolSSI) Software Fellows program (NSF grant ACI-1547580) Krylov et al. (2018); Wilkins-Diehr and Crawford (2018). We are grateful to D.E. Shaw Research for sharing the WW domain simulation trajectories.

Software availability

Software to perform slow modes discovery along with an API compatible with scikit-learn and Keras, Jupyter notebooks to reproduce the results presented in this paper, documentation, and examples have been freely available under MIT license at https://github.com/hsidky/srv.