Learning Markovian Dynamics with Spectral Maps

Understanding the long timescale dynamics of complex molecular systems constitutes a fundamental problem in chemical physics [1]. This dynamics is often governed by rare transitions between multiple long-lived metastable states that involve overcoming barriers much higher than the thermal energy ($\gg k_{\mathrm{B}}T$) [2, 3]. Although having many spatial and temporal scales, systems exhibiting metastable characteristics frequently display slow relaxation dynamics that can be captured by a few reaction coordinates, also called collective variables (CVs), to demarcate between physical states [4, 5]. The long timescale dynamics is often considered effectively Markovian and modeled as diffusion in a free-energy landscape spanned by CVs [6, 7, 8]. Under this view, the dynamics mainly depends on slowly varying variables, while fast degrees of freedom are treated as uncoupled thermal noise, leading to adiabatic timescale separation [9].

Relying merely on physical intuition or trial and error to identify slow CVs can be unsystematic and result in highly non-Markovian dynamics with long memory effects, which can hinder our understanding of the underlying physical process [10, 11]. The problem of timescale separation has been addressed in the Mori–Zwanzig projection operator formalism by constructing a generalized Langevin equation that describes the dynamics in a slow subspace. However, as noted by Zwanzig, “no a priori choice is specifically indicated by theory” [12] to select this slow subspace. Recently, there has been considerable interest in constructing reduced representations of complex systems using data-driven techniques [13, 14, 15, 16]. Many such techniques are developed to learn CVs by separating slow and fast kinetics, which is particularly useful for metastable systems with multiple temporal scales [17, 18, 19, 20, 21, 22, 23, 24, 25, 26, 27, 28, 29, 30, 31, 32, 33, 34].

In this Communication, we further improve spectral map [34], an unsupervised statistical learning technique inspired by diffusion map [35, 36, 37] and parametric dimensionality reduction [38, 39, 40, 41, 42]. Spectral map constructs slow CVs to describe Markovian dynamics by maximizing the spectral gap between slow and fast eigenvalues of a Markov transition matrix, increasing timescale separation, and preventing memory effects. We improve the learning process of spectral map by implementing an adaptive algorithm to estimate transition probabilities, enabling the representation of multiscale and heterogeneous free-energy landscapes with long-lived metastable states according to their physical characteristics. Finally, through a standard Markov state model analysis, we validate that CVs constructed by spectral map encode dominant slow relaxation timescales.

Consider a high-dimensional system described by $n$ configuration variables $\mathbf{x}=(x_{1},\dots,x_{n})$ whose dynamics at temperature $T$ and a potential energy function $U(\mathbf{x})$ is sampled from an unknown equilibrium distribution. If we represent the system by its microscopic coordinates, the dynamics follows a canonical equilibrium distribution given by the Boltzmann density $p(\mathbf{x})=\operatorname{e}^{-\beta U(\mathbf{x})}/Z$, where $\beta=1/(k_{\mathrm{B}}T)$ is the inverse temperature, $k_{\mathrm{B}}$ is the Boltzmann constant and $Z=\int\operatorname{d\mathbf{x}}\operatorname{e}^{-\beta U(\mathbf{x})}$ is the partition function of the system.

We simplify the high-dimensional configuration space by mapping it into a reduced space $\mathbf{z}=\left(z_{1},\dots,z_{d}\right)$ given by a set of $d$ functions of the configuration variables, commonly referred to as CVs, where $d\ll n$. We encapsulate these functions in a target mapping [41, 42, 16]:

$$\mathbf{z}=\xi_{w}(\mathbf{x})\equiv\big{\{}\xi_{k}(\mathbf{x},w)\big{\}}_{k=1}^{d},$$

(1)

where $w$ are learnable parameters ensuring that the target mapping describes slow dynamics and minimizes memory effects. By sampling the system in the CV space, its dynamics follows a marginal equilibrium density $p(\mathbf{z})\propto\operatorname{e}^{-\beta F(\mathbf{z})}$, where $F(\mathbf{z})$ is a free-energy landscape:

$$F(\mathbf{z})=-\frac{1}{\beta}\log\int\operatorname{d\mathbf{x}}\,\delta\left(\mathbf{z}-\xi_{w}(\mathbf{x})\right)\operatorname{e}^{-\beta U(\mathbf{x})}$$

(2)

defined modulo an irrelevant constant.

To estimate the separation between effective timescales characteristic of the system, we model its reduced dynamics as a Markov chain using kernel functions. To measure the similarity between CV samples $\mathbf{z}_{k}$ and $\mathbf{z}_{l}$, we introduce a Gaussian kernel for the reduced representation:

$$G(\mathbf{z}_{k},\mathbf{z}_{l})=\exp\left(-\frac{1}{\varepsilon}{\lVert\mathbf{z}_{k}-\mathbf{z}_{l}\rVert^{2}}\right)$$

(3)

where $\lVert\mathbf{z}_{k}-\mathbf{z}_{l}\rVert$ denotes pairwise Euclidean distances between every pair of CV samples $k,l=1,\dots,N$ and $N$ is the number of samples. The Gaussian kernel exhibits a notion of locality by defining a neighborhood around each sample of radius $\varepsilon$ [17].

In our previous proof-of-concept work that proposed spectral map [34], $\varepsilon$ was kept constant throughout the learning process. However, metastable states often have multimodal characteristics, resulting in spatially heterogeneous free-energy landscapes. Therefore, to improve the ability of spectral map to adjust to metastable states, we estimate the sample-dependent scale factors by adaptively balancing local and global scales as:

$$\varepsilon_{kl}(r)=\lVert\mathbf{z}_{k}-\eta_{r}(\mathbf{z}_{k})\rVert\cdot\lVert\mathbf{z}_{l}-\eta_{r}(\mathbf{z}_{l})\rVert,$$

(4)

where each term defines a ball centered at $\mathbf{z}$ of radius $\eta_{r}(\mathbf{z})>0$. For convenience, we define the radius by the fraction of the neighborhood size $r\in[0,1]$, allowing us to decide which scale is more relevant. Specifically, the Gaussian kernel describes a local neighborhood around each sample for values $r$ close to 0 (i.e., the nearest neighbors), which correspond to deep and narrow states. For values of $r$ around 1 (the farthest neighbors), it considers more global information, corresponding to shallow and wide states. Intermediate values of $r$ maintain a balance between spatial scales. For additional details about estimating $\varepsilon_{kl}$, see the Supplementary Material (Sec. S2).

As the marginal equilibrium density in the CV space is often far from uniform for dynamical systems with complex free-energy landscapes, we need a density-preserving kernel for data sampled from any underlying probability distribution. For this, we employ an anisotropic diffusion kernel as introduced for diffusion maps [35]:

$$K(\mathbf{z}_{k},\mathbf{z}_{l})=\frac{G(\mathbf{z}_{k},\mathbf{z}_{l})}{\sqrt{\varrho(\mathbf{z}_{k})\varrho(\mathbf{z}_{l})}},$$

(5)

where $\varrho(\mathbf{z}_{k})=\sum_{l}G(\mathbf{z}_{k},\mathbf{z}_{l})$ is a kernel density estimate. Next, we build a Markov transition matrix by row-normalizing $K$:

$$m_{kl}\sim M(\mathbf{z}_{k},\mathbf{z}_{l})=\frac{K(\mathbf{z}_{k},\mathbf{z}_{l})}{\sum_{n}K(\mathbf{z}_{k},\mathbf{z}_{n})}$$

(6)

which models a discrete Markov chain in the CV space $m_{kl}=\mathrm{Pr}\{\mathbf{z}_{\tau+1}=\mathbf{z}_{l}\;|\;\mathbf{z}_{\tau}=\mathbf{z}_{k}\}$ expressing a probability of transition between CV samples $\mathbf{z}_{k}$ and $\mathbf{z}_{l}$ in an auxiliary (non-physical) time $\tau$. The Markov chain approximates the long-time asymptotics of the system by describing the dynamics by the Fokker–Planck anisotropic diffusion [35].

Finally, to estimate the dominant timescales encoded in the system, we perform a spectral decomposition of the Markov transition matrix:

$$M\Psi_{k}=\lambda_{k}\Psi_{k},$$

(7)

where $\Psi_{k}$ and $\lambda_{k}$ are the $k$-th right eigenfunctions and eigenvalues of $M$, respectively. The real-valued eigenvalues of $M$ are (sorted in non-ascending order):

$$\lambda_{0}=1>\lambda_{1}\dots\geq\lambda_{N},$$

(8)

where the eigenvalue $\lambda_{0}$ corresponds to the equilibrium distribution of the Markov chain (Eq. 6) given by the eigenfunction $\Psi_{0}$. The dominant eigenvalues related to the slowest relaxation timescales in the system can be found by associating each eigenvalue with an effective timescale [45]:

$$t_{k}=-\frac{1}{\log\lambda_{k}}.$$

(9)

The largest gap between neighboring eigenvalues is called the spectral gap and determines the degree of the timescale separation between the slow and fast processes:

$$\sigma=\lambda_{k-1}-\lambda_{k},$$

(10)

where $k>0$ indicates the number of metastable states in the CV space [46].

The theory of spectral characterization of metastable states (see works by Gaveau and Schulman [47, 46] and references therein) explains that the spectral gap and degree of degeneracy in eigenvalue spectrum are related to the timescale separation and Markovian dynamics. If the eigenvalue of the Markov transition matrix is nearly degenerate $k+1$ times, it indicates that the equilibrium distribution breaks into $k$ metastable states with infrequent transitions between them. The converse is also true: if the equilibrium density breaks into metastable states separated by a free-energy barrier much larger than the thermal energy $k_{\mathrm{B}}T$, there is eigenvalue degeneracy.

To achieve the reduced dynamics that is effectively Markovian, it is crucial to have a spectral gap between neighboring eigenvalues, along with the near degeneracy of the dominant eigenvalue. This condition is essential for the maximal spectral gap at $k$ to lead to the separation into $k$ metastable states, which can help reduce memory effects in the dynamics. Therefore, we consider the spectral gap as a scoring function in spectral map.

Many techniques for learning slow CVs rely partly on identifying lag times. In contrast, spectral map learns CVs describing the slowest modes and dominant relaxation times without requiring a specific lag time. This is possible as spectral map estimates transition probabilities between samples, unlike other techniques that rely on counting configurations. Selecting an incorrect lag time can lead to non-Markovian dynamics and result in non-negligible memory effects, which can mix slow and fast timescales, causing significant errors in estimated kinetics [57, 58]. However, it is important to note that spectral map exchanges temporal parameters (lag time) for spatial parameters (neighborhood size). Estimating scale factors is also known to affect the reduced representation [59, 60]. Although methods for identifying slow CVs using a lag time [22, 26, 61] are well established and commonly used, spectral map is a recently developed technique. Therefore, it is yet to be determined if estimating scale factors is more practical than selecting lag times.

Spectral map employs the anisotropic diffusion kernel to estimate transition probabilities and maximize the timescale separation. This kernel has been first introduced for diffusion map [17], and therefore, it is important to compare these two techniques. In diffusion maps, the reduced space is approximated using eigenfunctions of a Fokker–Planck operator (diffusion coordinates) constructed from samples in the configuration space. Because of that, the validity of the eigenfunctions cannot be self-consistently verified, and their quality cannot be improved. Additionally, from a technical perspective, eigendecompositions of large transition matrices can be computationally expensive. In contrast to diffusion map, spectral map performs the spectral decomposition of a Markov transition matrix in the reduced space, which is perhaps the main difference between the two techniques. Spectral map does not rely on eigenfunctions as approximations of the reduced representation and only requires computing eigenvalues. These eigenvalues are then used to maximize the spectral gap to iteratively improve CVs corresponding to slow kinetics through the learning process. Initially, the reduced dynamics can have non-Markovian characteristics. As the spectral gap is maximized, memory effects are reduced. Moreover, the eigendecompositions are computed from data batches, which reduces computational costs.

In this Communication, we have further improved spectral map, an unsupervised machine-learning technique for learning slow CVs describing Markovian dynamics of complex metastable systems. By implementing an adaptive algorithm for estimating transition probabilities, we have added to spectral map the ability to represent heterogeneous and multiscale free-energy landscapes. We have verified with a standard Markov state model analysis that spectral map can be successfully used to construct CVs corresponding to the slowest modes. Our technique is applicable to a variety of physical processes with multiple temporal scales. Although spectral map is still in its preliminary version, it has shown potential in computing slow reaction coordinates and deserves further development.

Computational details and results, including information about simulation datasets, learning, and Markov state models are available in the Supplementary Material.

J. R. acknowledges funding from the Polish Science Foundation (START), the Ministry of Science and Higher Education in Poland, and the Japan Society for the Promotion of Science (JSPS). J. R. and T. G. are supported by the National Science Center in Poland (Sonata 2021/43/D/ST4/00920, “Statistical Learning of Slow Collective Variables from Atomistic Simulations”). D. E. Shaw Research is acknowledged for providing the CLN025 dataset.

Data and a reference implementation of spectral map are available on PLUMED-NEST [49] (www.plumed-nest.org), the public repository of the PLUMED consortium, as plumID:24.005.