Efficient approximation of molecular kinetics using random Fourier features

The automated extraction of essential information about the thermodynamic and kinetic properties of molecular systems from large-scale computer simulations remains one of the major open problems in computational physics and chemistry to this day. The challenge consists of processing very long time series of high-dimensional data in such a way that relevant features of the molecular system are determined automatically, with limited user input or intervention. The precise definition of what is regarded as relevant depends on the context, but in many applications, the determination of slow relaxation processes and their associated metastable sets is among the most important features to be detected. This applies, for example, to protein folding, ligand binding, or protein–protein association.^1–3 

Let us highlight two milestones that have shaped the field in the last two decades: The first was the inception of Markov State Models (MSMs) in the late 1990s and early 2000s,^4–6 which allowed for the construction of easily interpretable discrete kinetic models from which a wealth of information could be extracted. The second milestone was the variational approach to conformational dynamics (VAC) for reversible systems^7,8 and its extension, the variational approach to Markov processes (VAMP),⁹ which provided a rigorously defined scalar objective function (called the VAMP score) to assess the quality of kinetic models. Combined, these developments have led to a standard workflow for the construction of MSMs, usually consisting of a linear dimensionality reduction using time-lagged independent component analysis (TICA),¹⁰ geometric clustering, as well as MSM estimation and validation.^11,12 It was also shown that MSMs and the VAC are closely related to the Koopman operator framework^13,14 and, in particular, the extended dynamic mode decomposition (EDMD) algorithm,^15,16 which have primarily emerged in the analysis of fluid dynamics and related fields. The Koopman framework is therefore a unifying theme for the construction of kinetic models. Extensions of the EDMD algorithm can also be used to approximate the generator of dynamical systems [e.g., the Kolmogorov backward operator for stochastic differential equations (SDEs)], as shown in Ref. 17.

Despite the remarkable success of MSMs, the construction process remains somewhat unsatisfactory, at least from a theoretical perspective. The initial TICA step presupposes that there is a kinetically informative linear subspace of the initial set of descriptors used for MSM building. In order to move beyond this requirement, algorithms based on deep neural networks (DNNs) were introduced in recent years,^18,19 allowing one to learn in one step the complete non-linear transformation from raw atomic coordinates to low-dimensional features that optimize the VAMP score. Both the MSM construction process (including state definition) and neural networks require a non-linear optimization step, which is typically solved by means of stochastic algorithms. A significant number of hyper-parameters need to be tuned in both cases, such as network architecture, network depth, and height for DNNs, or number of cluster centers, number of iterations, convergence threshold, etc. for clustering-based MSMs.

A different avenue that has been considered is to employ reproducing kernels, which have long been known in machine learning^20,21 as a powerful non-linear model class with rich theoretical properties. Kernel methods have also been demonstrated in many applications to be fairly robust if the available data are only small- to medium-sized. Kernel-based versions of EDMD were suggested in Refs. 22–24 and applied to high-dimensional molecular simulations in Ref. 25. Extracting slow timescales from kernel representations requires the solution of a large-scale generalized eigenvalue problem, where the size of the matrix depends on the number of snapshots. Even for moderate data sizes, the solution to this problem can be challenging. Moreover, although kernel methods usually require only a few hyper-parameters (such as the bandwidth of the hugely popular Gaussian kernel), these need to be tuned accurately, which requires solving the above-mentioned eigenvalue problem many times, which can be prohibitively time-consuming.

In this work, we show that kernel methods can be made tractable for the analysis of molecular kinetics by employing Random Fourier Features (RFFs),²⁶ a well-known low-rank approximation technique for kernel matrices, to convert the original large-scale eigenvalue problem into a much smaller one that can be easily solved. We derive the associated small-scale matrix problem both for the Koopman operator and its generator and show that this approach can be interpreted as applying EDMD with a specific randomly chosen basis set. We demonstrate that the VAMP score can be effectively used for hyper-parameter tuning of the kernel method. The combination of these ideas leads to an efficient approximation algorithm that is robust with respect to the data size and the stochastic Fourier feature selection. Therefore, our method presents a promising new direction for the automated analysis of molecular kinetics.

The remainder of this paper is structured as follows: In Sec. II, we review elements of the Koopman formalism, the EDMD approximation algorithm, reproducing kernel Hilbert spaces (RKHSs), and Koopman modeling using RKHSs. Then, in Sec. III, we present our main contribution, an approximation algorithm for the Koopman operator or generator based on Random Fourier Features. Further computational details, especially regarding model selection, are discussed in Sec. IV, while applications to benchmarking problems in molecular dynamics (MD) are shown in Sec. V.

We will briefly introduce the stochastic Koopman operator and its generator, as well as numerical methods to estimate these operators and their spectral properties from data.

A critical modeling decision for the (g)EDMD approach is the choice of the dictionary {ϕ₁, …, ϕ_n}. Recent works^22–25 have shown that accurate results, even for large systems, can be obtained using methods based on reproducing kernel Hilbert spaces (RKHS) generated by a kernel function k. These approaches lead to powerful approximation spaces, yet they only depend on a few model parameters.

For many popular kernels, including the Gaussian RBF kernel, the associated RKHS is infinite-dimensional and densely embedded into the most relevant function spaces. Kernels, therefore, provide powerful approximation spaces, reducing the problem of dictionary selection to the choice of a single function, the kernel, and its hyper-parameters, such as the bandwidth for the Gaussian RBF kernel.

The linear problems (26) and (27) represent large-scale generalized eigenvalue problems. Solving these problems efficiently for large datasets is one of the central challenges common to most applications of kernel methods. In Sec. III, we present an approach based on a randomized low-rank representation obtained from the Fourier transformation of the kernel function.

In this section, we will present our main contribution: an approach to solving the generalized eigenvalue problems (26) and (27) using a randomized low-rank decomposition.

It is then sufficient to diagonalize these reduced matrices in order to solve (37). If the truncation rank of M is significantly smaller than p, additional computational savings can be harnessed this way. Algorithm 1 summarizes our solution method in compact form.

Let us break down the computational effort required for Algorithm 1: The SVD of M requires $O(mp2)$ operations. Assembling the reduced matrix costs $O(rpm)$ operations for the Koopman operator. In the case of generator approximation, we obtain an additional dependence on the dimension due to the required contraction of derivatives of the kernel function. We can upper-bound the resulting effort as $O(mp(d2+r))$ in the non-reversible case and by $O(m(rpd+r2d2))$ in the reversible case. Note that this bound assumes a dense and state-dependent diffusion field a; it reduces to linear dependence on d for diagonal diffusions, such as overdamped Langevin dynamics. Finally, the diagonalization of the reduced matrix amounts to $O(r3)$ floating point operations. Comparing this to the full kernel method, we see that the most drastic computational savings will result from being able to choose p small compared to m, as it reduces a factor $O(m3)$ to $O(mp2)$⁠.

In this section, we will discuss aspects of the practical application of our spectral RFF method and also introduce the systems studied in the results section below. For the numerical examples, we will use the Gaussian RBF kernel or its periodic counterpart for periodic domains.

Before we can apply Algorithm 1, we have to choose the number of features p, the truncation threshold for singular values, and as many hyper-parameters as required by the kernel function (e.g., the bandwidth for Gaussian RFF kernels). We select these parameters by cross-validating the VAMP score metric on randomly selected subsets of the data, as explained below. The only exception is the truncation threshold for singular values of M. Here, we adopt the rule that all singular values smaller than 10⁻⁴ times the largest singular value are discarded. This choice is based on prior experience and works well in all considered examples.

To account for the estimation error stemming from the use of finite data and also to avoid over-fitting, we employ a standard cross-validation approach and first compute the optimal functions ϕ₀, …, ϕ_K−1 for each model class on a random subsample of the data (training data). We then recompute the VAMP score for the K-dimensional space spanned by ϕ₀, …, ϕ_K−1 on the remaining data (test data). This process is repeated n_test times, where n_test is either 10 or 20 below, and we choose the model class that optimizes the test score on average. The ratio of training data to test data is always chosen as a 75%/25% split.

The solution of the eigenvalue problem for the reduced matrix R provides estimates of the leading eigenvalues $λ̂i(t)$ or $κ̂i$⁠, respectively, and their associated eigenfunctions $ψ̂i$⁠. The first point to note is that (37) is a complex eigenvalue problem, so the eigenvectors can be scaled by arbitrary complex numbers. Before analyzing the eigenvectors, we simply apply a grid search over complex phase factors to determine the one that minimizes the imaginary part of all eigenvector entries and use the rescaled eigenvectors for further analysis.

The second system is molecular dynamics simulation data of the alanine dipeptide in explicit water at a temperature of T = 300 K, employing the Langevin thermostat. We generated a total of 1 × 10⁶ data points at 1 ps time spacing. Reference implied timescales were computed by a Markov state model in the well-known two-dimensional space of backbone dihedral angles ϕ and ψ, using a 30 × 30 box discretization and all 1 × 10⁶ data points.

The third example is molecular dynamics simulation data of the deca alanine peptide in explicit water at a temperature of T = 300 K; see Ref. 37 for a description of the simulation setup. The data comprises a total of 4 × 10⁶ time steps at 1 ps time spacing. As reference values, we refer to the MSM analysis presented in Ref. 38, which is built on a 500-state k-means discretization after applying time-lagged independent component analysis (TICA)¹⁰ in the space of the peptide’s 16 backbone dihedral angles.

Finally, we re-analyze the 39-residue protein NTL9, which was simulated on the Anton supercomputer by D. E. Shaw Research; see Ref. 39 for details. The data comprises around 1.2 × 10⁶ frames, corresponding to a time spacing of Δ_t = 2 ns. As a reference, we apply linear TICA to the set of 666 minimal heavy atom inter-residue distances. We also rank these distances according to the fraction of simulation time where a contact is formed, that is, the distance does not exceed 0.35 nm; see Ref. 38, and calculate TICA models using only the first [20, 100, 200, 300, 400, 500, 666] of these distances.

We start by analyzing the Lemon-Slice potential (48), shown in Fig. 1(a), which features four metastable states corresponding to the wells of the potential energy. We use this system as a first test case for approximating the generator $L$ due to the availability of all relevant quantities in analytical form.

Figure 1(b) shows the VAMP score as a function of the bandwidth for a selection of data sizes m and feature numbers p. We see that while m = 1000 data points are not sufficient, for m = 5000 data points, the score stabilizes for a range of bandwidths σ ∈ [0.3, 0.8] and already attains optimal values for very small p. We confirm this finding by plotting the first three non-trivial eigenvalues of the generator, for m = 5000 and p = 50, as a function of the bandwidth in Fig. 1(c). Indeed, accurate estimates for eigenvalues can be obtained in the suggested regime of bandwidths. After calculating the eigenvectors of the reduced matrix for σ = 0.4 and applying the PCCA analysis, we see that the metastable structure is perfectly recovered; see Fig. 1(d). Thus, the VAMP score reliably guides the hyper-parameter search toward a very efficient approximation of the leading spectral components of this system.

For alanine dipeptide, we approximate the Koopman operator with lag time t, using the periodic Gaussian kernel on the well-known two-dimensional reaction coordinate space given by the dihedral angles ϕ and ψ. Figure 2(a) shows the familiar free energy landscape of the system in the reaction coordinate space. For all RFF-based approximations, we downsample the dataset to m = 20 000 points, corresponding to a time spacing of 50 ps. We apply the same model selection protocol as before and find that a range of kernel bandwidths σ ∈ [0.4, 1.0] can be stably identified as optimal in terms of the VAMP score. We observe in Fig. 2(b) that again, a small number of Fourier features, p = 50, is sufficient to arrive at an optimal VAMP score. We also confirm that the optimal parameter regime is stable across different lag times. Note that the VAMP score necessarily decreases with increasing lag time, as all eigenvalues but λ₀(t) decay exponentially with t. As shown in Fig. 2(c), the optimal range of bandwidths indeed allows for accurate estimation of the three leading implied timescales of the system. Further, we also verify that the metastable decomposition of the dihedral space into four states is correctly recovered using the RFF model; see Fig. 2(d). The dramatic rank reduction for the kernel matrix achieved by using a small number of Fourier features allows us to easily compute eigenvalues and test scores for a broad range of parameters using m = 20 000 data points in real space. Solving the same problem by means of standard methods would be significantly more costly if the same data size was used.

We use the deca alanine example to demonstrate that, by means of low-rank kernel methods, generator models can be efficiently learned on a reaction coordinate space that is more than just one- or two-dimensional and that dynamical hypotheses can be tested this way. We project the system onto the space of its 16 interior backbone torsion angles. The generator of full molecular dynamics is not readily available in closed form. It is well known, however, that if observed only in a subset of position space, many systems driven by thermostatically controlled molecular dynamics behave like reversible overdamped Langevin dynamics at long timescales.⁴¹ This can be rigorously shown for underdamped Langevin dynamics by using the inverse friction constant as a re-scaling of time, leading to overdamped dynamics in position space in the large friction limit.

We find in Fig. 3(a) that a kernel bandwidth of σ ∈ [2.0, 8.0] leads to an optimal VAMP score for the projected generator. Using about m = 1000 data points in real space and p = 300 Fourier features is sufficient to obtain stable results, rendering the calculation highly efficient. The optimality of this regime of bandwidths is confirmed by plotting the first three non-trivial eigenvalues of the generator, as shown in Fig. 3(b). We note that these eigenvalues differ from those of the reference MSM by a constant factor. This is due to the re-scaling of time inherent in assuming overdamped Langevin dynamics in position space, rendering the effective dynamics too fast. However, we also see in Fig. 3 that all slow eigenvalues are indeed re-scaled by the same factor. Moreover, the PCCA analysis of the slowest eigenvectors obtained for the optimal RFF model also recovers the metastable states of the deca alanine system, corresponding to the folded and unfolded states as well as two intermediates. The effective overdamped Langevin model learned by the kernel method thus retains all major features of the original molecular dynamics, up to the re-scaling of time.

We use the NTL9 dataset to demonstrate once again the ability of the proposed method to efficiently and robustly analyze large datasets. We downsample the available trajectory data to around m = 15 000 data points at 200 ns time spacing. As mentioned earlier, m = 15 000 is still a significant data size for kernel methods. We project the MD simulations onto the 666-dimensional space of minimal heavy atom inter-residue distances. It is known from previous publications that these coordinates capture the folding process of NTL9 very well.^38,42 Indeed, just using linear TICA on the top-ranked 300 distances allows for an accurate estimation of the slowest implied timescale, t₁, associated with the folding process; see the top green line in Fig. 4(b).

Figure 4(a) shows that using Gaussian bandwidth parameters σ ∈ [10, 50] can be identified as optimal in terms of the VAMP score, again across different lag times. In Fig. 4(b), we first compare the slowest timescale t₁ obtained for σ = 15 as a function of the lag time for different feature numbers p (blue lines). We see that already for p ≥ 300, satisfactory convergence of t₁ with t can be obtained, reducing the computational cost by about a factor of two thousand compared to the full kernel eigenvalue problem. The folded and unfolded states of the protein can also be identified correctly using these parameters, see Fig. 5. Interestingly, timescale convergence improves with p in about the same way that the performance of TICA improves as more and more of the ranked distances are added. However, we show that the kernel method is more robust with respect to the selection of input coordinates by randomly selecting d = 50 out of the 666 distances and recomputing both the VAMP score and implied timescales for different bandwidths, keeping the number of Fourier features fixed at p = 300. We see in Figs. 4(a) and 4(b) that both the optimal value of the VAMP score as well as the timescale estimates remain stable across random selections of the input distances [magenta lines in panel (a), red lines in panel (b)]. The convergence of timescales is, therefore, mainly dependent on the number of Fourier features p, which can be automatically tuned using the VAMP score, and not on the selection of input coordinates. Moreover, we also find that under random distance selection, the optimal regime of kernel bandwidths changes as the dimension of the input coordinate space becomes significantly smaller. This is reliably reflected by the behavior of the VAMP score; see again Fig. 4(a).

We have introduced a new approach to extracting the dominant eigenvalues of dynamical operators associated with stochastic processes, corresponding to the slow timescale dynamics in molecular kinetics. Building on kernel representations of the Koopman operator and generator, our main novelty is a dual low-rank representation of the kernel matrices using random Fourier features. We have derived the corresponding reduced eigenvalue problem and provided an interpretation of the method in terms of a random dictionary of plane waves. Using four examples of increasing complexity, we have shown that hyper-parameters of the method can be very effectively tuned using the VAMP score metric. Moreover, we have shown that only a few hundred Fourier features were sufficient to obtain accurate and robust results for all systems considered. This finding allowed us to rapidly scan kernel hyper-parameters for data sizes in the ten thousand on coordinate spaces of dimensions up to several hundred. The computational savings compared to the full eigenvalue problem, measured in floating point operations, were on the order of three to four orders of magnitude.

The main promise of our method is the ability to analyze large datasets using a very generic and, in fact, randomly generated basis set. The success of the method depends on the number of required Fourier features. Future research will need to test if this number indeed remains small even for larger systems, which would make the method scalable beyond the benchmarking examples studied here. Combining our approach with re-weighting techniques would allow the application of the method to enhanced sampling simulation data. It will also be interesting to see if kernel functions with more specific properties can be used, for instance, kernels incorporating more detailed symmetries than only translational invariance, kernels acting on mixed types of domain (periodic and non-periodic), or kernels associated with Sobolev spaces, which have rich analytical properties. In addition, concerning the theoretical analysis, it is known that the timescale separation between slow and fast dynamics expected for metastable systems induces an approximate low-rank structure of the Koopman operator. Our results suggest that this structure can be captured by a small number of random features. Providing a more detailed analysis of this heuristic argument presents another future research direction.

The authors thank D. E. Shaw Research for providing the molecular dynamics simulation data for NTL9.

The authors have no conflicts to disclose.

Feliks Nüske: Conceptualization (lead); Data curation (lead); Formal analysis (equal); Investigation (equal); Methodology (lead); Software (lead); Validation (lead); Visualization (lead); Writing – original draft (lead); Writing – review & editing (lead). Stefan Klus: Conceptualization (supporting); Formal analysis (equal); Investigation (equal); Methodology (supporting); Validation (supporting); Visualization (supporting); Writing – original draft (supporting); Writing – review & editing (supporting).

The methods presented in this paper are available as part of the KoopmanLib library at https://github.com/fnueske/KoopmanLib. The complete data and code to reproduce the examples shown above are available from a public repository at https://doi.org/10.5281/zenodo.8036525, with the exception of the raw molecular dynamics simulation data, which can be obtained directly from the authors or, in the case of NTL9, from the copyright owners. PCCA decompositions and all Markov state model results were computed using the deeptime library.⁴³