Variational deep learning of equilibrium transition path ensembles

TL;DR

This paper introduces a neural network-based variational method to analyze equilibrium transition paths, accurately estimate reaction rates, and elucidate mechanisms with minimal data in stochastic systems.

Contribution

It develops a novel variational approach that combines neural networks and path sampling to efficiently learn reaction mechanisms and rate contributions in equilibrium processes.

Findings

Accurately estimates reaction rates with minimal trajectories.

Provides detailed insight into reactive mode contributions.

Demonstrates effectiveness on model systems and alanine dipeptide.

Abstract

We present a time dependent variational method to learn the mechanisms of equilibrium reactive processes and efficiently evaluate their rates within a transition path ensemble. This approach builds off variational path sampling methodology by approximating the time dependent commitment probability within a neural network ansatz. The reaction mechanisms inferred through this approach are elucidated by a novel decomposition of the rate in terms of the components of a stochastic path action conditioned on a transition. This decomposition affords an ability to resolve the typical contribution of each reactive mode and their couplings to the rare event. The associated rate evaluation is variational and systematically improvable through the development of a cumulant expansion. We demonstrate this method in both over- and under-damped stochastic equations of motion, in low-dimensional model systems and the isomerization of solvated alanine dipeptide. In all examples, we find that we can obtain quantitatively accurate estimates of the rates of the reactive events with minimal trajectory statistics, and gain unique insight into the transitions through the analysis of their commitment probability.

Tables

Table 1. Table 1: Contribution to rate in implicit solvent

Label	Type	Description	Contribution
$d_4$	Dihedral	4 - 6 - 8 - 14	0.416
$d_3$	Dihedral	4 - 6 - 8 - 10	0.160
$d_8$	Dihedral	10 - 8 - 14 - 16	0.079
$a_3$	Angle	5 - 4 - 6	0.048
$d_{11}$	Dihedral	1 - 6 - 4 - 5	0.043
$d_7$	Dihedral	10 - 8 - 14 - 15	0.042
$a_1$	Angle	1 - 4 - 5	0.025
$a_8$	Angle	8 - 14 - 15	0.023
$a_{10}$	Angle	15 - 14 - 16	0.022
$a_{11}$	Angle	14 - 16 - 18	0.021
$a_9$	Angle	8 - 14 - 16	0.021
$a_6$	Angle	6 - 8 - 14	0.019
$a_2$	Angle	1 - 4 - 6	0.018
$d_{12}$	Dihedral	8 - 16 - 14 - 15	0.014
$a_4$	Angle	4 - 6 - 8	0.011
$a_7$	Angle	10 - 8 - 14	0.009
$d_9$	Dihedral	8 - 14 - 16 - 18	0.004
$d_{10}$	Dihedral	15 - 14 - 16 - 18	0.002
$a_5$	Angle	6 - 8 - 10	0.002
$d_2$	Dihedral	5 - 4 - 6 - 8	0.001
$d_1$	Dihedral	1 - 4 - 6 - 8	-0.009
$d_6$	Dihedral	6 - 8 - 14 - 16	-0.010
$d_5$	Dihedral	6 - 8 - 14 - 15	-0.012

Table 2. Table 2: Contribution to rate in explicit solvent

Label	Type	Description	Contribution
$a_{13}$	Angle	0 - 1 - 3	0.035
$a_{41}$	Angle	17 - 16 - 18	0.035
$d_4$	Dihedral	4 - 6 - 8 - 14	0.034
$a_{23}$	Angle	7 - 6 - 8	0.032
$a_{33}$	Angle	11 - 10 - 12	0.032
$a_{12}$	Angle	0 - 1 - 2	0.031
$a_{15}$	Angle	2 - 1 - 3	0.030
$a_{34}$	Angle	11 - 10 - 13	0.029
$d_{21}$	Dihedral	5 - 4 - 6 - 7	0.026
$a_{47}$	Angle	20 - 18 - 21	0.025
$d_7$	Dihedral	10 - 8 - 14 - 15	0.024
$d_3$	Dihedral	4 - 6 - 8 - 10	0.023
$d_{11}$	Dihedral	1 - 6 - 4 - 5	0.023
$a_{27}$	Angle	9 - 8 - 10	0.021
$d_{23}$	Dihedral	4 - 6 - 8 - 9	0.019
$a_{35}$	Angle	12 - 10 - 13	0.018
$a_{44}$	Angle	16 - 18 - 21	0.017
$d_{33}$	Dihedral	9 - 8 - 10 - 12	0.015
$a_{42}$	Angle	16 - 18 - 19	0.015
$d_{17}$	Dihedral	3 - 1 - 4 - 5	0.015

Equations

$${\gamma }_i{\dot{\mathbf{r}}}_i(t)={\mathbf{F}}_i\left({\mathbf{r}}^N\right)+{\mathbf{\eta }}_i(t)$$

$$h_X[{\mathbf{r}}^N(t)]=\{\begin{array}{l}1{\mathbf{r}}^N(t)\in X\\ 0{\mathbf{r}}^N(t)\notin X\end{array}$$

$$k=\frac{d}{dt}\frac{⟨h_A(0)h_B(t)⟩}{⟨h_A⟩}=\frac{d}{dt}⟨h_{B|A}(t)⟩$$

$$k{t}_f=\int \mathcal{D}[\mathbf{X}]h_{B|A}(t_f)P[\mathbf{X}]$$

$${\gamma }_i{\dot{\mathbf{r}}}_i={\mathbf{F}}_i({\mathbf{r}}^N)+{\mathbf{\lambda }}_i({\mathbf{r}}^N,t)+{\mathbf{\eta }}_i(t)$$

$$k_{\lambda }{t}_f=\int \mathcal{D}[\mathbf{X}]h_{B|A}(t)P_{\lambda }[\mathbf{X}]$$

$$\ln k_{\lambda }{t}_f$$

$$=\ln k{t}_f+\ln {⟨e^{\Delta U_{\lambda }}⟩}_{B|A}$$

$$\displaystyle{\color[rgb]{0,0,0}\definecolor[named]{pgfstrokecolor}{rgb}{0,0,0}\pgfsys@color@gray@stroke{0}\pgfsys@color@gray@fill{0}\Delta U_{\lambda}[\mathbf{X}]=-\sum_{i=1}^{N}\frac{1}{4\gamma_{i}k_{\mathrm{B}}T}\int_{0}^{t_{f}}dt\,|\boldsymbol{\lambda}_{i}|^{2}-2\boldsymbol{\lambda}_{i}\cdot(\mathbf{\gamma_{i}\dot{r}_{i}}-\mathbf{F_{i}})}$$

$$\ln k{t}_f$$

$$=-⟨\Delta U_{\lambda }{⟩}_{B|A}-\sum _{n=2}^{\infty }\frac{1}{n!}{\mathcal{C}}_{B|A}^{(n)}(\Delta U_{\lambda })$$

$$\displaystyle{\color[rgb]{0,0,0}\definecolor[named]{pgfstrokecolor}{rgb}{0,0,0}\pgfsys@color@gray@stroke{0}\pgfsys@color@gray@fill{0}\ln kt_{f}}$$

$$\displaystyle{\color[rgb]{0,0,0}\definecolor[named]{pgfstrokecolor}{rgb}{0,0,0}\pgfsys@color@gray@stroke{0}\pgfsys@color@gray@fill{0}\implies\ln kt_{f}}$$

$${\partial }_{t}q({\mathbf{r}}^N,t)=-\sum _{i=1}^N\frac{{\mathbf{F}}_i\left({\mathbf{r}}^N\right)}{{\gamma }_i}\cdot {\nabla }_{i}q({\mathbf{r}}^N,t)-\frac{k_{\mathrm{B}}T}{{\gamma }_i}{\nabla }_i^{2}q({\mathbf{r}}^N,t)$$

$${\mathbf{\lambda }}^{\ast }({\mathbf{r}}^N,t)=2k_{\mathrm{B}}T\nabla \ln q({\mathbf{r}}^N,t)$$

$$\displaystyle{\color[rgb]{0,0,0}\definecolor[named]{pgfstrokecolor}{rgb}{0,0,0}\pgfsys@color@gray@stroke{0}\pgfsys@color@gray@fill{0}\mathcal{L}_{\lambda}=\bigg{\langle}\sum_{n=0}^{t_{f}/\Delta t}\sum_{i=1}^{N}\frac{\Delta t}{4\gamma_{i}k_{\mathrm{B}}T}\big{(}-|\boldsymbol{\lambda}_{i}(n\Delta t)|^{2}}$$

$$\displaystyle{\color[rgb]{0,0,0}\definecolor[named]{pgfstrokecolor}{rgb}{0,0,0}\pgfsys@color@gray@stroke{0}\pgfsys@color@gray@fill{0}+2\boldsymbol{\lambda}_{i}(n\Delta t)\cdot\boldsymbol{\eta}_{i}(n\Delta t)\big{)}\bigg{\rangle}_{B|A}}$$

$$2k_{\mathrm{B}}T\ln q({\mathbf{r}}^N,t)=\sum _{n=1}^{n_{\mathrm{b}}}{c}_n{\phi }_n({\mathbf{r}}^N,t)$$

$$\mathbf{c^{*}}=\left[\bigg{\langle}\int_{0}^{t_{f}}dt\;\nabla\mathbf{\Phi}\otimes\nabla\mathbf{\Phi}\bigg{\rangle}_{B|A}\right]^{-1}\bigg{\langle}\int_{0}^{t_{f}}dt\;\boldsymbol{\eta}\cdot\nabla\mathbf{\Phi}\bigg{\rangle}_{B|A}$$

$$V(x,y)/k_{\mathrm{B}}T=2[6+4x^4-6y^2+3y^4+10x^2(y^2-1)]$$

$${\color[rgb]{0,0,0}\definecolor[named]{pgfstrokecolor}{rgb}{0,0,0}\pgfsys@color@gray@stroke{0}\pgfsys@color@gray@fill{0}h_{A}(t)=\Theta(-x(t)+0.85)\qquad h_{B}(t)=\Theta(x(t)-0.85)}$$

$${\phi }_n(x,y,t)=e^{-a_x(x-x_n{)}^2}{e}^{-a_y(y-y_n{)}^2}{e}^{-a_t(t-t_n{)}^2}$$

$$k{t}_f=\exp \left[-\sum _{i=1}^{Nd}⟨\Delta U_{{\lambda }^{\ast }}^i{⟩}_{B|A}\right]$$

$$\Delta U_{{\lambda }^{\ast }}^i$$

$$P[\Delta U_{{\lambda }^{\ast }}^{\alpha }]={⟨\delta (\Delta U_{{\lambda }^{\ast }}^{\alpha }-\Delta U_{{\lambda }^{\ast }}^{\alpha }[\mathbf{X}])⟩}_{B|A}$$

$${\mathbf{J}}_{\mathbf{r}}(\tilde{\mathbf{r}})=\left[\begin{array}{ccc}{\nabla }_{\mathbf{r}}{\tilde{\mathbf{r}}}_1& \cdots & {\nabla }_{\mathbf{r}}{\tilde{\mathbf{r}}}_{\tilde{n}}\end{array}\right]$$

$$\mathbf{\lambda }({\mathbf{r}}^N,t)={\mathbf{J}}_{\mathbf{r}}^T(\tilde{\mathbf{r}})\cdot \tilde{\mathbf{\lambda }}(\tilde{\mathbf{r}},t)$$

$$\Delta U_{{\lambda }^{\ast }}$$

$$\Delta {\tilde{U}}_{{\lambda }^{\ast }}^{jk}={\int }_0^{t_f}dt\frac{{\tilde{\lambda }}_j^{\ast }{\tilde{\lambda }}_k^{\ast }}{4k_{\mathrm{B}}T}{\Gamma }_{jk}^{-1}$$

$$m{\dot{\mathbf{v}}}_i=-\gamma {\mathbf{v}}_i+{\mathbf{F}}_i({\mathbf{r}}^N)+{\mathbf{\eta }}_i{\dot{\mathbf{r}}}_i={\mathbf{v}}_i$$

$${\partial }_{t}q=-\sum _{i=1}{\mathbf{v}}_i\cdot {\nabla }_{{\mathbf{r}}_i}q-\frac{\gamma {\mathbf{v}}_i}{m}\cdot {\nabla }_{{\mathbf{v}}_i}q+\frac{{\mathbf{F}}_i}{m}\cdot {\nabla }_{{\mathbf{v}}_i}q+\frac{2\gamma k_{\mathrm{B}}T}{m^2}{\nabla }_{{\mathbf{v}}_i}^{2}q$$

$${\mathbf{\lambda }}_i^{\ast }({\mathbf{r}}^N,{\mathbf{v}}^N,t)=\frac{2\gamma k_{\mathrm{B}}T}{m}{\nabla }_{{\mathbf{v}}_i}\ln q({\mathbf{r}}^N,{\mathbf{v}}^N,t)$$

Full text

Introduction

In complex systems, understanding the mechanism of transitions between long-lived metastable states is hampered by the general collective nature of the dynamics and the difficulty of observing these rare but important events.1 While methods like transition path sampling2 exist to harvest rare events computationally, their distillation into mechanistic descriptions is cumbersome, and the conversion of that description into quantitative statements of their rate is challenging.3; 4 Here, we present a method that uses a neural-network ansatz with a variational optimization procedure to compute the time dependent commitment probability from a reactive trajectory ensemble. The method involves learning a unique policy, in the form of an optimal external control force, that reweights a reactive conditioned path ensemble to an unconditioned ensemble that reacts autonomously. The optimal force is simply related to the commitment probability,5; 6 and serves as an ideal descriptor of the reaction. The reweighting principle developed within the framework of variational path sampling7 is expressed in terms of the stochastic action, which allows us to decompose the rate into additive contributions from different degrees of freedom, including collective coordinates that describe molecular transitions. This decomposition provides a means of identifying relevant order parameters without making a-priori assumptions. The combination of the mechanistic insight afforded by an interpretable representation of the reaction and the validation through a variational evaluation of the rate, provides a robust method for distilling features of equilibrium transition path ensembles.

The investigation of reactive events requires access to timescales that are considerably longer than the local relaxation time of the system. The canonical approach to investigate these processes has leveraged physically intuitive low-rank descriptions of the system to infer mechanistic insight, and bridge the timescales through reactive flux calculations or importance sampling.8; 9; 10; 11; 12 The notion of an ideal reaction coordinate capable of providing a complete description of the reactive event dates back to Onsager13, and was formalized within the context of chemical physics as the committor– a map between the phase space position of a system and the likelihood of it reacting3; 14; 15; 16. Learning this high dimensional function has attracted interest from a diversity of fields, and significant advances has been made through methods that employ importance sampling and machine learning.17; 18; 16; 19; 20; 21; 22; 23; 24; 25; 26; 27; 28; 23; 29; 30; 31; 32; 33; 34; 35Some notable approaches have leveraged the confinement of the transition region to compute it using string methods17; 16; 36, coarse-grained the phase-space to approximate it through diffusion maps19; 28; 37; 38, and parameterized neural-networks by either fitting the committor directly18; 21; 34 or solving the variational form of the steady-state backward Kolmogorov equation22 by combining it with importance sampling methods23; 24; 25. While the learning procedures applied previously have been successful in fitting high dimensional representations of the reaction coordinate or committors, their nonlinearity has largely resulted in a difficulty in interpreting the relative importance of physically distinct descriptors and converting those descriptors into a robust measure of the rate. Earlier developments of methods based on likelihood maximization 18; 39; 40 have offered linear ways to make this analysis tractable to complex processes41; 42; 43; 44; 45. However, these approaches have overwhelmingly relied on physical intuition to express likelihood functions.44

The method that we present builds off of variational path sampling7; 46; 6; 47; 48 that has provided an alternative approach for sampling rare events. These methods and related ones employ ideas from stochastic optimal control49; 50; 51; 52, and are most useful in investigating nonequilibrium steady states as they do not invoke detailed balance. Of particular interest is a recent method6 that has detailed how to express a low-rank ansatz for an optimal control force to drive rare events and estimate their rates. Our work exploits the fact that the optimization of this control force, or policy, is related to the time dependent committor. We find that in equilibrium systems, where path sampling methods afford a way to generate a reference reactive trajectory ensemble, the optimization of this committor becomes straightforward, and allows the use of a neural-network (NN) ansatz to solve the time-dependent backward Kolmogorov equation53, providing a time dependent and probabilistic representation of the reaction. While the method computes a nonlinear function, the form of the optimized loss is given by the difference in stochastic actions that quantifies the distance between a conditioned and a reference trajectory ensemble. For systems in which we saturate the variational bound, this quantity is unique and linearly decomposable on a per-coordinate basis, and can be understood as a measure of the importance of each coordinate to conditioning a trajectory to be reactive. This metric is purely based on the intrinsic mechanism of the reaction, and can be extended to collective coordinates, allowing us to identify the relevant reaction descriptors without making a-priori assumptions.

This paper is organized as follows. First, we review the variational path sampling formalism to discuss the theory behind this method. Next, we validate this method by applying it to a couple of low dimensional systems where numerically exact results are possible. We probe the sensitivity of this method to limited statistics as well as the applicability to systems integrated with underdamped equations of motion. Then, we illustrate how the per-coordinate stochastic action encodes the relevance of a coordinate to the reaction. Finally, we apply this method to study the isomerization of alanine dipeptide in implicit and explicit solvent. In both of these cases, we show how the method can be used to infer a mechanistic picture of the reaction, and identify important reaction descriptors among a redundant set of internal coordinates.

I Variational path sampling formalism

For simplicity, we consider a system evolving under an overdamped Langevin equation of the form,

[TABLE]

where $\dot{\mathbf{r}}_{i}$ is the rate of change of the $i$th particle’s position at time $t$ in $d$ dimensions, $\gamma_{i}$ is the friction coefficient, and $\boldsymbol{\eta}_{i}(t)$ denotes a Gaussian random force with mean $\langle\boldsymbol{\eta}_{i}(t)\rangle=0$ and variance $\langle\boldsymbol{\eta}_{i}(t)\otimes\boldsymbol{\eta}_{j}(t^{\prime})\rangle=2\gamma_{i}k_{\mathrm{B}}T\delta_{ij}\mathbf{1_{d}}\delta(t-t^{\prime})$ where $\otimes$ denotes the cross-product operator, $\mathbf{1_{d}}$ is an identity-matrix of size $d\times d$ and $k_{\mathrm{B}}T$ is Boltzmann’s constant times the temperature. The conservative force $\mathbf{F}_{i}\left(\mathbf{r}^{N}\right)=-\nabla_{i}V\left(\mathbf{r}^{N}\right)$ is given by the gradient of the potential $V\left(\mathbf{r}^{N}\right)$ with $\mathbf{r}^{N}$ the full $N$-particle configuration. We are interested in investigating reactive events, so we consider potentials that exhibit metastability.

We consider transitions between two metastable states, $A$ and $B$, which in general are collections of configurations defined through the indicator functions $h_{A}[\mathbf{r}^{N}(t)]$ and $h_{B}[\mathbf{r}^{N}(t)]$, where

[TABLE]

for $X=\{A,B\}$. For the rest of the the paper, indicator functions are going to written down simply as functions of time in favor of brevity. The rate for the $A\rightarrow B$ transition can be defined by the time derivative of the side-side correlation function,9

[TABLE]

where $\langle\cdots\rangle$ denotes an average computed over a stationary distribution and $h_{B|A}$ is the conditional probability of starting in $A$ and ending in $B$ at $t$. Provided a separation of timescales between the local relaxation time within a state, $\tau_{\mathrm{mol}}$, and $1/k$, the rate is given by the path integral

[TABLE]

where when $t_{f}$ is in the range $\tau_{\textrm{mol}}<t_{f}\ll 1/k$, the probability to transition grows linearly with time. The path integral sums over all trajectories $\mathbf{X}=\{\mathbf{r}^{N}(0),\dots,\mathbf{r}^{N}(t_{f})\}$, or the timeseries of the state of the system evolved for time $t_{f}$, weighted by the likelihood of observing a trajectory $P[\mathbf{X}]$. This path integral is a trajectory partition function associated with reactive paths,54 and equal to the transition probability between $A$ to $B$ in time $t_{f}$.

Variational path sampling uses the path partition function representation of the rate together with a dynamical reweighting approach55 to extract reactive paths effectively,46 evaluate rates accurately,6 and we show here, provide detailed mechanistic information concerning the rare event. Variational path sampling does this by considering the system as before, but under the action of an additional time-dependent drift $\boldsymbol{\lambda}_{i}(\mathbf{r}^{N},t)$, which enters the equation of motion as

[TABLE]

where the conservative force, noise and friction are the same as the reference system without $\boldsymbol{\lambda}_{i}(\mathbf{r}^{N},t)$. For this driven system, the rate $k_{\lambda}$ between the same two metastable states $A$ and $B$ is given by an analogous relation as in the reference system

[TABLE]

where $P_{\lambda}[\mathbf{X}]$ denotes the probability of observing a trajectory $\mathbf{X}$ integrated using Eq. 5. By virtue of the Girsanov transformation, these two rate expressions can be related to each other. Specifically, using the Radon-Nikodym derivative to define the change in stochastic action, $\Delta U_{\lambda}[\mathbf{X}]=\ln P_{\lambda}[\mathbf{X}]/P[\mathbf{X}]$, the rate in the driven system can be rewritten as56

[TABLE]

where we have employed $\langle\dots\rangle_{B|A}=\langle h_{B}(t_{f})h_{A}(0)\dots\rangle/\langle h_{A}(0)\rangle$ as a conditional average over a reference reactive ensemble to relate the two rates. For the case of the overdamped Langevin equation, the change in stochastic action is given by a difference of Onsager-Machlup actions57

[TABLE]

where $|\dots|$ denotes the norm of the vector, and $\cdot$ denotes the dot product operator. The discretized form of the OM action can be simplified further by noting that the difference between the the time derivative of the positions and the conservative force is simply given by the noises.

The relationship between rates in Eq. I is exact for any time-dependent drift $\boldsymbol{\lambda}_{i}(\mathbf{r}^{N},t)$. It is distinct from that employed previously,6 which related the reference and driven rates to an expectation value in the driven system. In variational path sampling, we consider a class of $\boldsymbol{\lambda}_{i}(\mathbf{r}^{N},t)$ which enforce the transition to occur with probability 1. In such a case, provided access to a reactive path ensemble in which to evaluate the expectation values, the rate in the reference system can be obtained directly as an exponential average,

[TABLE]

or a cumulant expansion, where $\mathcal{C}_{B|A}^{(n)}(\Delta U_{\lambda})$ denotes the $n$’th cumulant of $\Delta U_{\lambda}$ averaged in the reactive ensemble. We will refer to these different estimators as $k^{(\exp)}$ for the exponential average and $k^{(n)}$ for the cumulant expansion where $n$ will denote where the sum was truncated.

Truncation of the cumulant expansion for $n=1$ provides a variational bound of the rate. This is seen by applying Jensen’s inequality to Eq. I,

[TABLE]

where we have used conservation of probability $\ln k_{\lambda}t_{f}\leq 0$ in the second step to eliminate the rate of the driven process from the inequality. Hence, the rate in the reference system is just the Kullback-Leibler (KL) divergence between the driven and reference path ensembles, or equivalently the mean change in action. In equilibrium systems, this relation is similar to the variational structure of transition state theory, which also provides an upper bound to the rate.9 However, this expression is also closely related to the reversible work theorem in equilibrium thermodynamics,58 as it relates the smallest change required to transform one ensemble to another.59; 60 In this case, the transformation is between an unconditioned path ensemble and a reactive path ensemble. Just as the minimum amount of work done on a physical system is given by its reversible limit, which reflects the way in which a system would naturally transform, so too we find the minimum driving force to ensure a reaction is related to the way in which a system would naturally react.61 This is shown by noting that the force that saturates this bound in Eq.  11 is the Doob force, denoted $\boldsymbol{\lambda}^{*}(\mathbf{r}^{N},t)$, and is related to the solution of the backward Kolmogorov equation 62; 63; 5; 6. For an overdamped Langevin dynamics this is,

[TABLE]

with boundary conditions $q(\mathbf{r}^{N},t_{f})=h_{B}(t_{f})$ and $q(\mathbf{r}^{N},0)=h_{A}(0)$. The function that solves this expression, $q(\mathbf{r}^{N},t)$, is the time-dependent committor function,5; 6 or the probability of reaching state $B$ at $t_{f}$ given a position $\mathbf{r}^{N}$ at time $t$. In the stationary limit, where the separation of timescales prohibits multiple transitions, $q(\mathbf{r}^{N},t)$ reduces to the time independent committor function of transition path theory.64; 65 The explicit relation between the Doob force $\boldsymbol{\lambda}^{*}(\mathbf{r}^{N},t)$ and $q(\mathbf{r}^{N},t)$ is,

[TABLE]

where by construction this force makes all trajectories reactive, and the reactions occur as they would in the original system. This force uniquely saturates the inequality in Eq. 11, thus providing a unique description of the reaction in a complex system.6

This formalism allows us to compute both the time-dependent committor $q(\mathbf{r}^{N},t)$ and the rate $k$ from a reactive trajectory ensemble by parameterizing the external force $\boldsymbol{\lambda}$ and optimizing it by maximizing the expectation value of the change in action averaged within the reactive trajectory ensemble. In this work, we will consider parameterizing $\boldsymbol{\lambda}$ with both linear functional forms as well as a non-linear form provided by a neural network. The optimization of either is done by defining a loss function, $\mathcal{L}_{\lambda}$, as

[TABLE]

where the sum is over each particle that the noises act on, and $\boldsymbol{\lambda}_{i}$ is the component of the driving force on the degree of freedom associated with the noise. This loss function is just the change in stochastic action in the discretized form, so in optimizing $\boldsymbol{\lambda}_{i}$ we are simultaneously optimizing our estimate of the rate in the reference system. This optimization occurs over $n_{\mathrm{int}}$ iterations, and requires averages within the reactive trajectory ensemble, which we will generate with standard path sampling tools like transition path sampling. Since this method requires the positions and the noises at each time step, the method of generation of transition path ensemble requires some care. For complex systems where we use transition path sampling, we store the random number seeds for each runs, and then save the new trajectories only when they are accepted by rerunning them with the same seed. We find this method to be minimally slower than running standard transition path sampling. The option to save and change the seeding is available in most molecular dynamics simulation packages.

II Choice of Ansatz and Convergence

The accuracy of the rate estimate, and the mechanistic information afforded by the evaluation of $q(\mathbf{r}^{N},t)$, depends on the fidelity with which the function can be represented. This depends on the ansatz used to expand it, and in particular, its expressibility. It also depends on the ease by which the function is learned, as inevitably the reactive path ensemble needed to train $q(\mathbf{r}^{N},t)$ will be computationally expensive to generate. In this section, we consider the relative merits of expanding the driving force in both linear and nonlinear bases, and assess their accuracy and data efficiency.

II.1 Linear Function Ansatz

We first consider the case of linear function approximations. A linear functional for $\boldsymbol{\lambda}(\mathbf{r}^{N},t)=2k_{\mathrm{B}}T\nabla\ln q(\mathbf{r}^{N},t)$ can generically be expressed as its associated potential,

[TABLE]

where $c_{n}$ and $\varphi_{n}(\mathbf{r}^{N},t)$ denote the $n$’th coefficient and basis function, and $n_{\mathrm{b}}$ denotes the total number of basis functions. This can be written compactly as $\boldsymbol{\lambda}(\mathbf{r}^{N},t)=\nabla[\mathbf{c}\cdot\mathbf{\Phi}(\mathbf{r}^{N},t)]$, where $\mathbf{c}$ is the $n_{\mathrm{b}}$ length vector of coefficients and $\mathbf{\Phi}(\mathbf{r}^{N},t)$ is the vector of basis functions. For a linear functional expansion, the optimal set of coefficients $\mathbf{c}^{*}$ has a closed form that can be computed by taking the derivative of the loss function in Eq. 14 and setting it to 0. Computing the coefficients reduces to solving a $n_{\mathrm{b}}\times n_{\mathrm{b}}$ set of linear equations, whose solution is

[TABLE]

where $\otimes$ denotes an outer product. For an orthonormal basis, the optimal coefficients are simply related to the average noise-weighted basis function66, but in general, the functions are not expected to be orthonormal. Because of this simplicity in training, linear bases are particularly efficient to employ. In cases where the reaction coordinate can be described well by a limited set of coordinates or order parameters, they can also be accurate7; 46; 6; 33; 31.

To understand the utility of a linear functional approximation, we consider a particle evolving in a two-dimensional external potential with two reactive channels visualized in Fig. 1 (A). The potential $V(x,y)$ is

[TABLE]

where $x$ and $y$ are dimensionless coordinates and we have worked in a reduced unit system determined by $k_{\mathrm{B}}T=\gamma_{x}=\gamma_{y}=1$, and employed a first order Euler integrator with timestep equal to 0.004 $t^{*}$ with $t^{*}=\gamma_{x}/k_{\mathrm{B}}T$ as our reduced time unit. We considered transitions defined by the indicator functions

[TABLE]

where $\Theta$ denotes the Heaviside step function. A reactive path ensemble was generated by running brute force trajectories in order to sample 400 reactions, and the rate was evaluated by computing the side-side correlation function. We found that $t_{f}/t^{*}=2$ was a sufficient observation time to be in the linear growth regime for the transition probability with $\ln kt_{f}=-6.1\pm 0.1$.

The linear approximation used were localized Gaussian basis functions of the form

[TABLE]

where the Gaussian centers $\{x_{n},y_{n},t_{n}\}$ were equally spaced on a grid within the range of $x=[-1.5,1.3]$, $y=[-1.6,1.6]$ and $t=[0,2]$ and Gaussian widths were choosen such that $\{a_{x}=1.4/(n_{\mathrm{b}}^{1/3}-1),a_{y}=1.6/(n_{\mathrm{b}}^{1/3}-1),a_{t}=1/(n_{\mathrm{b}}^{1/3}-1)\}$. The expansion coefficients, $\mathbf{c^{*}}$, were computed using Eq. 16 averaged over the path ensemble consisting of the 400 reactive trajectories. The optimization was done in one step, $n_{\mathrm{int}}=1$ by solving the linear equation in Eq. 16, where we found the loss function immediately converged to the brute force estimate of the rate for $n_{\mathrm{b}}=12^{3}$ as shown in Fig. 1 (B). The dependence of the rate estimate with the size of the basis is shown in Fig. 1 (C), where the loss decays slowly, obtaining a value of the rate statistically indistinguishable from the brute force estimation of the rate for $n_{\mathrm{b}}=8^{3}$. This slow decay could be mitigated somewhat with fine tuning the basis, but we do not explore that here.

II.2 Neural Network Function Ansatz

Since the form of the force is rapidly varying and nonlinear, saturation of the inequality in Eq. 11 requires a large number of basis functions. If we express the linear ansatz in the full configuration space, the number of basis set coefficients grows exponentially with the degrees of freedom, making it intractable to converge the loss to the rate for complex systems. In order to circumvent the exponential scaling of the number of basis sets with the dimensionality of the system, we consider employing a neural network (NN) ansatz to compute the time dependent committor, associated Doob force through automatic differentiation, and evaluate the rate through optimization. The input comprises the features selected for expressing the force and is connected to two hidden layers. For the two hidden layers, the Swish activation function67 is used as its derivatives are free from discontinuities, while also being exempt from the weight decay problem.68 The penultimate layer only contains a single unit, with a sigmoid activation function. The output of this layer is the model’s estimate of the time-dependent committor, $q(\mathbf{r}^{N},t)$. The final layer is a lambda layer, which simply computes the log of the committor. The output of this layer represents the many-body potential, $\ln q(\mathbf{r}^{N},t)$, and the forces can be computed by taking a derivative of the output with respect to the input coordinates via autodifferentiation. While one can simply parametrize the forces instead of the committor, this architecture automatically enforces the conservativeness of the potential and offers a simple way to obtain the committor without the need to perform a multidimensional integration.

For the NN ansatz for the same 2D system above, $x$, $y$ and $t$ were used as the input features and optimization was performed using the RMSprop optimizer69 on 200 reactive trajectories. The learning rate was choosen to be 0.001, and for each iteration the loss function and associated gradients were evaluated over half of the trajectories drawn randomly from the ensemble. The training curve plotted in Fig. 1 (B) shows that the loss function plateaus to $\ln kt_{f}$ within $n_{\mathrm{int}}=20$ indicating that this ansatz was successful in learning the exact time-dependent committor quickly. While the training required multiple iterations, the number of parameters used to converge to the brute force rate was around 500 without specific optimization, fewer than required in the naive linear function approximation. The flexibility of the NN ansatz and the relatively swift training suggest it as a viable means of approximating the time dependent committor. As a consequence, in the remainder of the manuscript, we consider only the performance of the NN ansatz.

II.3 Convergence with Limited Statistics

To illustrate the efficiency of this method, we tested the convergence of the NN ansatz with the statistics used to compute the rate and time dependent committor within the previously introduced model two-dimensional potential. Specifically, we tested the convergence of the NN ansatz with the number of reactive trajectories used in training, as well as the time lag between configurations along a reactive trajectory. For both cases, we use two estimators, one which probed how closely the restricted trajectory ensemble is to the full trajectory ensemble, and a second which indicated how well a model trained on an approximated trajectory ensemble performs on the original trajectory ensemble. All of these estimators are based on comparisons between the OM action and $\ln kt_{f}$, as the agreement between the two signifies the success of the model in learning the true committor. Hence, the difference between the two offers a natural metric to probe the error between the parameterized and the true committor. We denote the error from each of these approximate estimates as $\Delta\mathcal{L}$.

For the first case, we vary the number of trajectories $N_{\mathrm{t}}$ used for training the model. The model trained on this limited trajectory ensemble, with force denoted by $\lambda^{N_{\mathrm{t}}}$ is then used to compute the first cumulant for the original trajectory ensemble comprised of the full 200 trajectories, $\langle\Delta U_{\lambda^{N_{\mathrm{t}}}}\rangle_{B|A}$. This is compared to the first cumulant obtained by training the model on the original trajectory ensemble with an optimal estimate of the rate. The difference of these two values, plotted in Fig. 2 (A), is an indicator of how close the committor trained on the restricted ensemble is to the actual committor. This plot shows that the estimator converges quickly with $N_{\mathrm{t}}$, and suggests that about 50 trajectories are sufficient to learn the time-dependent committor for this specific system. Another way of probing how the restricted trajectory ensemble compares to the original trajectory ensemble is to perform both training and averaging in the restricted trajectory ensemble, and compare that estimate to the true rate. This difference between the action averaged in a restricted ensemble $\langle\Delta U_{\lambda^{N_{\mathrm{t}}}}\rangle_{B|A,N_{\mathrm{t}}}$, also shown in Fig 2 (A), is observed to be negative for $N_{\mathrm{t}}=10$, indicating overfitting of the model to the restricted trajectory ensemble. However, this error vanishes quickly, and plateaus to 0 for $N_{\mathrm{t}}=50$. This is a reflection of the transition path ensemble and the similarity of different reactive trajectories. The error bars for all of these cases are obtained by training 5 different models on $N_{t}$ randomly selected trajectories from the trajectory ensemble.

For the second case, we approximate the trajectory ensemble by storing only every configurations after a timelag of $d_{c}\Delta t$ where $\Delta t$ is the timestep used to integrate the trajectory. The original reactive trajectory ensemble comprises 200 trajectories with 500 discrete timesteps, and the number of configurations used per trajectory are obtained by dividing 500 by $d_{c}$. We train the model in this approximated trajectory ensemble and compute the same error estimates. The first estimate compares the first cumulant averaged in the original trajectory ensemble ($d_{c}=1$) with the model trained on the approximated trajectory ensemble to the loss computed by performing both averaging and training in the original trajectory ensemble, $\langle\Delta U_{\lambda_{d_{c}}}\rangle_{B|A}$. In this case, this estimate probes how well the the NN ansatz is able to extrapolate the forces for timesteps that it has not been trained on. The difference plotted in Fig. 2 (B) indicates that this extrapolation fails quickly. The second estimate probes the loss obtained by performing both averaging and training in the approximated trajectory ensemble whose loss is $\langle\Delta U_{\lambda_{d_{c}}}\rangle_{B|A,d_{c}}$. To get this estimate, the variance in Eq. 14 had to be scaled by a factor of $d_{c}^{-1}$ to account for the change in the effective timestep $\Delta t$. This difference plotted in Fig. 2 (B) shows that this approximation only works well for $d_{c}\leq 5$. The poor scaling with $d_{c}$ reflects the fact that stochastic diffusions with different variances have no overlap in the continuum limit.70

III Rate decomposition and feature selection

From an information theoretic point of view, the rate is a ratio of a conditioned and an unconditioned trajectory partition function2; 54; 71; 72. Our optimization directly minimizes the KL-divergence between a trajectory ensemble driven with force $\boldsymbol{\lambda}$ and the undriven reactive trajectory ensemble. As the KL-divergence is expressible by the change in stochastic action along a trajectory, it involves a sum over all the degrees of freedom that the noises act on. For a suboptimal force, the rate is given by the average of the exponential of this quantity, coupling the noises from different degrees of freedom. However, when the variational bound is saturated and the rate is given by a simple mean, the accompanying change of action is linearly decomposable. This decomposition provides mechanistic insight, and affords a means of optimizing the features that form the representation of $\boldsymbol{\lambda}$. We generally find a NN ansatz to saturate the bound in Eq. 11, which allows us in this section to explore a variety of featurizations and their corresponding contributions to the rate. Specifically, we consider networks with Cartesian and collective coordinates, as well as those integrated with underdamped equations of motion.

III.1 Cartesian coordinates

Using an NN ansatz allows us to compute the exact time dependent committor and associated Doob force. When Eq. 11 is saturated, the rate is given by the first cumulant of the change in action. This allows us to decompose the rate into independent contributions,

[TABLE]

where

[TABLE]

is the contribution to the rate per stochastic degree of freedom. The change in action, $\Delta U^{i}_{\lambda^{*}}$, is strictly positive, indicative of the transition probability being less than 1, and results from functional minimization of Eq.  14. The stochastic action for the Langevin equation is a sum of Gaussian random variables for each degree of freedom at each timeslice, and a change in stochastic action is a difference of Gaussian random variables. Given this, and recognizing the quadratic dependence on $\lambda^{*}_{i}(t)$ for the change in action, we observe that $\lambda^{*}_{i}(t)$ is essentially fitting the bias in the Gaussian noises generated when conditioning the stochastic process to react. Therefore, only degrees of freedom that require activation, or a rare sequence of noises, will accumulate a significant change in stochastic action or contribute significantly to the rate. Degrees of freedom that are uncorrelated with the reaction will not contribute to the rate, as their noises will remain unbiased.

To illustrate how this decomposition can be used to identify the relevance of coordinates, we consider the same 2D system visualized in Fig. 1 (A) and perform a decomposition of the rate. The two stochastic coordinates, $x$ and $y$, are fed into the neural network ansatz and optimized. The resultant distributions for the individual components of the stochastic action, $P[\Delta U_{\lambda^{*}}^{\alpha}]$, for $\alpha=\{x,y\}$, defined as

[TABLE]

are shown in Fig. 3 (A). Neither of the two distributions show a complete overlap with the distribution of the total rate, $P[\Delta U_{\lambda^{*}}]$, indicating that both $x$ and $y$ are important in describing the reaction coordinate. However $P[\Delta U_{\lambda^{*}}^{x}]$ is shifted towards larger values, and the expectation value of $\langle\Delta U_{\lambda^{*}}^{x}\rangle_{B|A}$ is found to be be larger than $\langle\Delta U_{\lambda^{*}}^{y}\rangle_{B|A}$, allowing us to quantitatively assert that the coordinate $x$ is more important to the reaction than $y$, as it encodes more information of the conditioned path ensemble. However $y$ is still relevant, in agreement with intuition from the geometry of the potential.

III.2 Collective coordinates

While the decomposition above can quantify the relevance of an coordinate to a reactive event, they are expressed in the bare Cartesian coordinates that enter into the equation of motion. As such, their utility is diminished in many-particle systems which are translationally and rotationally invariant, and for which the number of degrees of freedom is large. A canonical approach in the study of rare events in complex systems is to employ collective coordinates, which are nonlinear combinations of the original Cartesian coordinates and may encode the expected symmetries of the system. In order to extend the formalism into this regime, we consider the transformation between Cartesian and collective coordinates, $\mathbf{r}\rightarrow\tilde{\mathbf{r}}$, and its subsequent impact on the rate decomposition. The Jacobian of the transformation $\mathbf{J}_{\boldsymbol{r}}(\tilde{\boldsymbol{r}})$ is,

[TABLE]

which is a matrix of $Nd\times\tilde{N}$ partial derivatives where $\tilde{N}$ is the size of the collective variable function space. Under this transformation, the original forces $\boldsymbol{\lambda}(\mathbf{r}^{N},t)$ and the transformed forces $\boldsymbol{\tilde{\lambda}}(\mathbf{\tilde{r}},t)$ are related by

[TABLE]

where the force acting on the original coordinate $\mathbf{r}_{i}$ due to the force $\boldsymbol{\tilde{\lambda}}_{j}$ which depends on the collective coordinate $\tilde{\mathbf{r}}_{j}$ is given by a product of $\tilde{\boldsymbol{\lambda}}_{j}$ and the Jacobian element $J_{ij}$. Inserting this into the expression for the optimal stochastic action, in Eq. 20, we obtain

[TABLE]

with

[TABLE]

where the initially linearly independent factors from each Cartesian coordinate, indexed by $i$, are expressed as a pair of contributions from the collective coordinates, indexed by $j$ and $k$. From this form it is evident that the contribution to the rate incurred from the transformed coordinates $\mathbf{\tilde{r}}$ are not necessarily independent of each other or bipartite72. Zero coupling between $\tilde{r}_{j}$ and $\tilde{r}_{k}$ is obtained when the effective friction $\Gamma^{-1}_{jk}=\sum_{i=1}J_{ij}J_{ik}/\gamma_{i}=\delta_{jk}/\gamma$, with $\gamma_{i}=\gamma$, a condition that requires the friction weighted transformed coordinates to be orthogonal.

As an illustration of the decomposition under a change of coordinates, we consider the same 2D system as before, but rather than parameterizing $\boldsymbol{\lambda}_{i}$ on $x$ and $y$, we transform into polar coordinates $(x,y)\rightarrow(r,\theta)$, where $r=x^{2}+y^{2}$ and $\tan\theta=y/x$. We quantify the contributions to the rate from the polar coordinates, by training the NN ansatz on the polar coordinates. The partials are prepared ahead of time and are passed into the loss function along with the noises. The relative action distributions in the transformed coordinates, $P[\Delta\tilde{U}_{\lambda^{*}}^{\alpha,\alpha^{\prime}}]$ for $\alpha,\alpha^{\prime}=\{r,\theta\}$ computed using Eq. 24 are shown in Fig. 3 (B). Since polar coordinates are orthogonal and $\gamma_{x}=\gamma_{y}$, the coupling term $\Delta\tilde{U}_{\lambda^{*}}^{\alpha,\alpha^{\prime}}=0$ for $\alpha\neq\alpha^{\prime}$. We observe that the distribution corresponding to the coordinate $\theta$ almost perfectly overlaps with the total action distribution, indicating that $\theta$ is an excellent descriptor of the reaction coordinate. The distribution for $r$ is centered around 0 and narrow, illustrating it is unbiased by conditioning on a reaction and thus contributes little to the rate. This decomposition of the rate in collective coordinates provides a simple metric to identify the relevance of physically meaningful descriptors to a reactive process, without making any a-priori assumptions about the reaction. The form of this metric is purely based on the physical mechanism of the reaction, as it quantifies how conditioning a trajectory ensemble to be reactive shifts the noise distributions per-coordinate. This allows us to do hypothesis testing for the relevance of collective coordinates, and discover the coordinates that are gating the rare event, and those that are uncorrelated with barrier crossing, through the size of their contribution to the rate. This hypothesis testing requires the saturation of the variational bound, which if not achieved points to the lack of relevant features in the NN ansatz.

III.3 Importance of velocity

In a general molecular system, motion is not overdamped and as a consequence the full phase space spanned by both configurational coordinates as well as their conjugate velocities are required to specify a reactive trajectory. In order to understand the importance of including velocity degrees of freedom in a parameterization of $\boldsymbol{\lambda}$, we consider formally when it can be neglected. For concreteness, we consider an the underdamped Langevin equation of the form

[TABLE]

where $\mathbf{v}_{i}$ is the velocity of particle $i$ and the rest of the quantities are defined in the same way as in Eq. 1. For simplicity we take the mass, $m$, and friction $\gamma$ to be independent of particle index, though generalizations are straightforward. We start by noting that the backward Kolmogorov equation takes the form,

[TABLE]

which when solved with the same boundary conditions as Eq. 12 yields the time dependent committor function $q(\mathbf{r}^{N},\mathbf{v}^{N},t)$ whose arguments we suppress above for ease of notation. Since the noise acts only on the velocities, the Doob force is given by the gradient of the committor with respect to the velocities rather than the positions,

[TABLE]

thus naively it would seem that parameterizing a velocity dependence is crucial whenever an underdamped equation is used. However, in the limit that $\gamma^{-1}\rightarrow 0$, we find that the velocity dependence can be safely ignored.

This can be understood via application of perturbation theory, where $q(\mathbf{r}^{N},\mathbf{v}^{N},t)$ is expanded in orders of $\gamma^{-1}$73; 36; 74. To first order, $q(\mathbf{r}^{N},\mathbf{v}^{N},t)$ becomes

[TABLE]

where $q_{0}$ is independent of the velocity. Substituting the approximated form of $q$ into the underdamped backward Kolmogorov equation, we find,

[TABLE]

which when averaged over the Maxwell-Boltzmann distribution, yields

[TABLE]

which to first order in $1/\gamma$ is identical to Eq. 12, the overdamped backward Kolmogorov equation, with $q(\mathbf{r}^{N},\mathbf{v}^{N},t)\approx q_{0}(\mathbf{r}^{N},t)$. As a consequence, the committor in the overdamped limit becomes a function solely of $\mathbf{r}^{N}$ and the Doob force is given by a gradient with respect to position. In Appendix A, we show that to $\mathcal{O}(\gamma^{-2})$ this approximation also saturates the variational inequality for the rate expression.

In order to gain intuition for when the higher order terms in Eq. 28 become negligible, we consider the reaction of a particle in a simple double well potential of the form

[TABLE]

where $x$ is a dimensionless coordinate and we take $k_{\mathrm{B}}T=m=1$, which determines a dimensionless time unit $t^{*}=\sqrt{m/k_{\mathrm{B}}T}$. We considered transitions between states defined by the indicator functions

[TABLE]

and obtain 400 reactive trajectories each of length $t_{f}/t^{*}=5$ using a timestep of 0.01 $t^{*}$ using first order integrator. We studied this system over a range of $\gamma/\gamma^{*}$ between 0.1 and 1 with $\gamma^{*}=m/t^{*}$. We trained a NN ansatz only on the positions and time, and compared the optimized value of the loss function to the brute-force rates evaluated from a direct mean first passage time calculation. Figure 4 (A) shows the difference between the two estimates, along with the brute-force rate as a function of $\gamma/\gamma^{*}$. The reactive rates show a Kramers’ turnover10 at $\gamma/\gamma^{*}\approx 0.3$, and the optimized loss is consistently off by a factor of 1.5 for $\gamma/\gamma^{*}<0.3$. After the turnover, the error in rate estimate decreases monotonically, until it completely vanishes for $\gamma/\gamma^{*}=1$. It is surprising that this relatively small friction is already consistent with the overdamped limit.

To further understand the importance of velocity in the time dependent committor, we train a model to optimize for the velocity-dependent commitor for $\gamma/\gamma^{*}=0.1$ and $\gamma/\gamma^{*}=1$. For both of these cases, the optimized value of the loss was within a standard error of the true rate. The plot of the optimized committor evaluated at time $t=t_{f}/2$ as a function of velocity and position is shown in Figs. 4 (B) and (C). The functional dependence on time is not strong away from $t=0$ and $t=t_{f}$. For $\gamma/\gamma^{*}=1.0$, $q(x,v,t)$ depends weakly on $v$, with the dependence being captured by a linear shift along $x$ to an otherwise simple sigmodal dependence on $x$. This is precisely the dependence expected from the expansion in Eq. 28. However, for $\gamma/\gamma^{*}=0.1$ the committor is strongly sensitive to the velocity. For $v=0$ the large $x$ behavior of $q(x,v,t)$ slowly converges to 1 reflecting the potential for the particle even at large values of $x$ to fail to react. For negative velocities, the inflection point of $q(x,v,t)$ is shifted to positive values of $x$, consistent with corresponding values of the potential that are low enough below the barrier that the particle is trapped. Correspondingly, for positive velocities, $q(x,v,t)$ is shifted to negative values of $x$, reflecting the high likelihood of reacting even for positions not quite to the top of the barrier. This behavior is not reproducible by scaling a spatially dependent committor by a simple constant. Hence, featurization of the velocity, or expansion of the committor to higher orders in $\gamma^{-1}$ is required to accurately encode the time-dependent committor for this low-friction regime.

IV Application to Alanine Dipeptide

To examine the efficacy of this method for a complex molecular system, we investigate the isomerization of alanine dipeptide. Alanine dipeptide has two metastable conformations. It can transition between these two states via the rotation of the Ramachandran angles $\phi$ and $\psi$. A multitude of path sampling methods have focused on this model due to the collective nature of this transition in the gas phase and in solution. While the transition can be tracked using the $\phi$ and $\psi$, they serve only as order parameters and are not sufficient in describing the complete reaction coordinate or committor. 14; 18 Significant advancements in methods to parameterize the time independent committor have been made by resolving this model along physically motivated, predetermined order parameters28; 75; 76; 77; 23; 78; 79; 80; 21; 81. As we show, choosing among a large number of internal coordinates without consideration of their correlation or coupling risks neglecting important aspects of the transition path ensemble. This is because internal coordinates do not form an orthogonal set of coordinates, and collective motions such as the rotations of a single dihedral angle can be coupled with the motions of angles and other dihedrals. Below we first consider isomerization of alanine dipeptide in implicit solvent, and then in explicit solution. For both we parameterize $\ln q(\mathbf{r}^{N},t)$ using the NN ansatz.

IV.1 Isomerization in implicit solvent

In implicit solvent we consider isomerization of alanine dipeptide between its {\color[rgb]{0,0,0}\definecolor[named]{pgfstrokecolor}{rgb}{0,0,0}\pgfsys@color@gray@stroke{0}\pgfsys@color@gray@fill{0}C7}_{\mathrm{eq}} and {\color[rgb]{0,0,0}\definecolor[named]{pgfstrokecolor}{rgb}{0,0,0}\pgfsys@color@gray@stroke{0}\pgfsys@color@gray@fill{0}C7}_{\mathrm{ax}} conformations, as visualized in the inset in Fig. 5. To investigate this reaction, we first generated a reactive trajectory ensemble. Simulations were performed in OpenMM82 and the AMBER ff14SB forcefield83 was used for parametrizing the dipeptide interactions. A Langevin thermostat with the leap-frog discretization was used as the integrator84. The timestep was chosen to be 1 fs, $\gamma$ was set to $10$ ps*-1* and the transition path length $t_{f}$ was set as 1 ps. The indicator functions identifying the metastable wells {\color[rgb]{0,0,0}\definecolor[named]{pgfstrokecolor}{rgb}{0,0,0}\pgfsys@color@gray@stroke{0}\pgfsys@color@gray@fill{0}C7}_{\mathrm{ax}} and {\color[rgb]{0,0,0}\definecolor[named]{pgfstrokecolor}{rgb}{0,0,0}\pgfsys@color@gray@stroke{0}\pgfsys@color@gray@fill{0}C7}_{\mathrm{eq}} were defined using the Ramachandran angle $\phi$,

[TABLE]

and the first trajectory was generated by running forward and backward simulations from the top of the saddle point along the dihedral $\phi=0$. Transition path sampling2 was used to obtain a reactive trajectory ensemble and the shooting from the top method85 was used to generate new trajectories. This method offers a way to decrease correlations between the trajectories as well as increase the acceptance rate for new trajectories by performing shooting moves within a restricted region near the saddle point, which for this case was chosen as $-\pi/6\leq\phi\leq\pi/6$. A total of 1000 trial trajectories were generated, and the acceptance rate came out to be approximately 0.4. Every 5th trajectory in the ensemble was saved and used for analysis for a total of 200 trajectories. For the choice of reaction descriptors, we used all the internal coordinates that did not involve the hydrogen atoms. This set consists of 9 bonds, 11 angles and 12 dihedrals, the latter two of which contain 9 total redundant coordinates. These internal coordinates along with the Jacobians matrices are computed for the 200 saved trajectories, and saved to be used for training.

We use the underdamped approximation discussed in section Eq. 28, which exempts us from including the velocities as a part of the feature set. The loss function is modified accordingly for the action of the Langevin leap-frog integrator86 implemented in OpenMM, as shown in Appendix B. The RMSProp optimizer with a learning rate of 0.001 was used to train the model, and training was performed for 8000 steps. A 90-10 training validation split was used for optimization, and the splits were randomized every 100 epochs. Figure 5 shows the value of the training set loss, along with the reaction rate obtained by computing the mean first passage time of 400 reactive trajectories generated independently. The loss function plateaus around the 4000th step, with the value being within one standard error of the true rate.

To gain mechanistic insight into the reaction, we performed the decomposition of the relative action in terms of the internal coordinates. Following Eq. 24, we computed the $\langle\tilde{U}_{\lambda^{*}}^{jk}\rangle_{B|A}$ matrix and visualize it in Fig. 6 (A). The angles and dihedrals are represented using the letters $a$ and $d$ respectively, and the numbers in the subscripts are defined in Appendix C. We observe that the matrix of contributions to the rate is sparse, with only a few select coordinates and their couplings obtaining a significant value. The contributions from the distances have been removed from the plot as their combined value was calculated to be statistically indistinguishable from zero. The effective decoupling of the bond vibrations from the angles and dihedrals provides evidence that the NN-based ansatz is not overfitting redundant features from the limited input dataset, consistent with physical intuition for stiff bonds.

We also observe a strong coupling among the dihedrals and the angles. Internal coordinates do not form an orthogonal set of coordinates, and the off-diagonal terms in the matrix indicate that the reaction is mediated by the coupling between these internal degrees of freedom. Moreover, the change in action is delocalized between sets of internal coordinates that have been ignored in previous studies. We note that the off-diagonal elements of the matrix are negative, while the diagonal terms are positive. This prevents us breaking down the rate in terms of additive contributions from different degrees of freedom. However, the matrix is symmetric so we can sum over the rows of the matrix and define the contribution from a single collective coordinate $j$ as

[TABLE]

where $\Delta\tilde{U}^{jk}_{\lambda^{*}}$ is defined the same way as in Eq. 24. The reason for using this measure is based on the sum rule for the OM action along transformed coordinates, which from using calculus of variation follows

[TABLE]

which is distinct from the case of Cartesian coordinates, where the relation is given by $[\lambda^{*}_{i}]^{2}=\lambda_{i}^{*}\eta_{i}$. This difference reflects the fact that the noises along pairs of transformed coordinates are not necessarily independent of each other and cannot be assumed to follow bipartite dynamics72. Summing over the rows of the matrix can be understood as a marginalization of the OM action over all the coupled coordinates Plotted in Fig. 6 (B), this decomposition is found to be positive for almost all the internal coordinates except for two. These negative values are within the standard error. This allows us to extract the leading contributors to the {\color[rgb]{0,0,0}\definecolor[named]{pgfstrokecolor}{rgb}{0,0,0}\pgfsys@color@gray@stroke{0}\pgfsys@color@gray@fill{0}C7}_{\mathrm{ax}}\rightarrow{\color[rgb]{0,0,0}\definecolor[named]{pgfstrokecolor}{rgb}{0,0,0}\pgfsys@color@gray@stroke{0}\pgfsys@color@gray@fill{0}C7}_{\mathrm{eq}} reaction. We observe that the Ramachandran angle $\phi$ ($d_{4}$) is found to incur the largest contribution. This is a remarkable result as no a-priori information of the reaction coordinate was passed into the model for training. While the indicator functions that were used to define the boundaries of the metastable wells were defined using $\phi$, the optimization scheme itself did not require any description of the indicator functions. Yet, this method automatically finds the Ramachandran angle $\phi$ to contribute the most to isomerization, out of 32 internal coordinates, 9 of which are redundant.

This decomposition reveals other leading contributors to the reaction and highlights other order parameters that are activated. The C-N-Cα-Cβ ($d_{3}$) and the Cβ-Cα-C-N ($d_{8}$) torsions are found to be the next two leading contributors, suggesting that rotation of the Ramchandran angle $\phi$ is strongly coupled to the orientation of the alkyl bond. Some other important internal coordinates that are selected by this method include the O-C-N angle ($a_{3}$), the C-O-C-N improper torsion ($d_{11}$) and the Cβ-Cα-C-O ($d_{7}$) torsion. These internal coordinates also emphasize the importance of the relative orientation of the O-C bond and the methyl-bond. Our final observation is that the contribution from other Ramachandran angle $\psi$ ($d_{6}$) is found to be effectively zero. This is another significant result as $\psi$ has long been used as the 2nd order parameter to explore the isomerization of alanine dipeptide due to the topology of the free energy surface. As this has been mentioned previously,14 future studies should consider coarse-graining along other dihedral angles for performing committor analysis of this reaction.

IV.2 Isomerization in Explicit Solvent

Finally to demonstrate the ability to tackle very high dimensional systems, we explore the conformational isomerization of alanine dipeptide in explicit solvent. As the potential energy landscape of along the Ramachandran angles is modified due to solvent interactions14; 87, we consider the isomerization between the {\color[rgb]{0,0,0}\definecolor[named]{pgfstrokecolor}{rgb}{0,0,0}\pgfsys@color@gray@stroke{0}\pgfsys@color@gray@fill{0}C5} and $\alpha_{L}$ states, visualized in Fig. 7. The equations of motion and forcefields for the peptide are the same as in the implicit solvent study, and the TIP3P forcefield88 is used for parameterizing the water molecules. Lorentz-Berthelot mixing rules are used for the peptide-water interactions. A periodic box of volume 27 nm3 is used with 862 water molecules, and the Ewald Summation is used for computing the long ranged interactions. The basin definitions for $\alpha_{L}$ and {\color[rgb]{0,0,0}\definecolor[named]{pgfstrokecolor}{rgb}{0,0,0}\pgfsys@color@gray@stroke{0}\pgfsys@color@gray@fill{0}C5} are the same as that of the {\color[rgb]{0,0,0}\definecolor[named]{pgfstrokecolor}{rgb}{0,0,0}\pgfsys@color@gray@stroke{0}\pgfsys@color@gray@fill{0}C7}_{\mathrm{ax}} and {\color[rgb]{0,0,0}\definecolor[named]{pgfstrokecolor}{rgb}{0,0,0}\pgfsys@color@gray@stroke{0}\pgfsys@color@gray@fill{0}C7}_{\mathrm{eq}} states, respectively. The same method as before is used for obtaining a reactive trajectory ensemble, with an ensemble of 200 reactive trajectories used for learning the time dependent committor.

For this reaction, we restrict the input feature to only parameterize the internal coordinates of the peptide. Noting that bond vibrations are decoupled from rotations of the dihedrals and that water interactions are mediated through hydrogen bonds, our input feature set comprises all the 36 angles and 45 dihedrals and contains 42 redundancies. While the solvent degrees of freedom can be parameterized using symmetry functions89; 90; 91, our goal is to illustrate how our method can leverage a quantitative insight into the reaction mechanism even when it does not have access to the full phase space.

Optimization of the NN ansatz is performed using the RMSProp Optimizer for 2500 steps, and the results are shown in Fig. 7 (B). The loss function plateaus to a value of 2 higher than $\ln kt_{f}$, which was computed through evaluation of the first mean passage time from 400 independently run trajectories. This discrepancy between the first cumulant and the true rate is expected, as the feature set excludes the relevant solvent degrees of freedom. Regardless, we are able to obtain the correct rate from this method by computing the second cumulant and exponential average, the form of which is given in Eqs. 9 and 10. Both these estimators are plotted, and are observed to converge to the rate computed independently. The agreement between the exponential estimator and second cumulant is only expected when the loss function is perturbatively close to the true value, otherwise the additional cumulants would be needed. Note that even in the case where the Doob force is not fully optimal, driven trajectories are almost surely reactive.

Since the variational bound is not saturated, the sum given in Eq. 20 does not equal the rate. This means that the description of the time dependent committor is not exact. However, because we are perturbatively close we can still use the method to extract the relative importance of degrees of freedom as before. First, we perform the same decomposition as Eq. 34. While not visualized, the $\Delta\tilde{U}^{jk}_{\lambda}$ matrix is found to be less sparse for the reaction in solvent, due to the renormalization of the solvent effects into the peptide degrees of freedom. To gain quantitative insight, we sum of the rows of the $\Delta\tilde{U}^{jk}_{\lambda}$ matrix as before, and plot the contributions from the 20 leading features. Plotted in Fig. 8 (A), this decomposition shows that the rotations of angles involving hydrogen atoms become more important than the internal rotations of the peptide dihedrals. Most of the leading contributors are angles that involve one or two hydrogen atoms, reemphasizing the effect of solvent in mediating this reaction. This is in accord with the findings of previous papers on the isomerization of solvated alanine dipeptide. 18; 92; 21; 93 However what is striking is that no single mode is dominant, with no internal coordinate accounting for more than 5% of the rate.

To confirm the role of the hydrogen atoms, we plot a decomposition in terms of the individual atoms in Fig. 8 (B) using the action expressed in the bare coordinates. The plot reveals that the methyl carbon contributes the most to the rate, followed by the acetyl carbonyl oxygen atom. However, the combined importance of the hydrogens far outweighs both. This finding also illuminates why the addition of solvent transforms the reactive mechanism. Both these atoms strongly interact with water molecules via hydrophilic and hydrophobic effects that are mediated through hydrogen-bonding and volume exclusion, respectively.18; 92; 21; 93 We find that this method is able to provide a rate estimate and quantify the renormalized contributions from different degrees of freedom even when it does not have access to the full phase space. This feature can be particularly useful for more complex systems, where a complete description of the system is not tractable due to computational or memory bottlenecks.

Conclusion

We have detailed a novel method that can be used to evaluate the time-dependent committor and the rate from a reactive trajectory ensemble. The method employs an ansatz for parameterizing a many-body potential that is related to the time-dependent committor, and can be optimized by variationally solving the backward Kolmogorov equation, as expressed through a trajectory reweighting theory used within variational path sampling. For reactive processes in equilibrium, where the cost of obtaining a reactive trajectory ensemble is independent of the rarity of the reaction, this method provides a simple procedure to compute the rate and distill mechanistic information.

Combining this optimization scheme with a neural network ansatz for the time dependent committor allows us to saturate the variational rate bound, and gives us a complete description of the transition path ensemble. Specifically, we have described how to decompose the rate in terms of additive contributions from different degrees of freedom. This procedure of quantifying contributions can be applied to collective coordinates and order parameters that are used for characterizing reactions of complex molecular systems. We showcase this decomposition by investigating the reaction of Brownian particles in simple potentials in underdamped and overdamped regimes. We have shown how to apply this procedure to conformational changes in solution, leveraging insightful information about the reactive event even when the full phase space is not provided as training data. In cases where the variational bound is not saturated, the rate can still be computed using other estimators. This decomposition could lend insight into the design of models to accurately recover kinetic information.94

The formalism employed, casts the time dependent committor as an optimal control force naturally making this model generative. Specifically, when the variational bound is saturated, a time dependent control force is produced that generates reactive trajectories in an unbiased manner. While not used as such here, this procedure can be employed to glean higher order statistics of the reactions over and above the rate.6 When the variational bound is not saturated the control force can still be applied to generate unbiased transition path statistics through ensemble reweighting.55 One could envision an iterative procedure in cases where path sampling is difficult, for example in cases of long diffusive trajectories, where initial control forces are gradually optimized through alternative cycles of training and reactive ensemble generation.

As the method is based on ensembles of trajectories and path reweighting, there is no formal restriction to equilibrium systems. Indeed, variational path sampling has been initially applied to systems whose dynamics break detailed balance. As such the procedures developed here for NN based function approximations and rate decompositions transfer over directly to rare transitions in nonequilibrium steady-states. However, traditional path sampling techniques that render the generation of a path ensemble simple in equilibrium are not typically as effective away from equilibrium. For those systems, one would have to consider using path sampling methods that do not invoke detailed balance.95; 96; 97 The iterative procedure alluded to above is likely a robust means of extending this methodology to study phase transitions in active matter and driven assembly.

Acknowledgements

We would like to thank Dr. Avishek Das and Dr. Jorge Rosa-Raíces for illuminating discussions regarding variational path sampling methods. We would also like to thank Dr. Clay Batton for useful comments regarding the manuscript. This material is based upon work supported by the U.S. Department of Energy, Office of Science, Office of Advanced Scientific Computing Research and Office of Basic Energy Sciences, via the Scientific Discovery through Advanced Computing (SciDAC) program.

Data Availability

The data for the plots and the illustration of this method on the 2D potential can be found within the Github repository https://github.com/ansingh1214/Deep-VPS.

Appendix A Saturation of variational bound

In Section III.3 we detailed how the time dependent committor can be approximated for a formally underdamped system evolving in an overdamped regime. Here we demonstrate that the approximate form of the time dependent committor saturates the variational rate bound up to order $\mathcal{O}(\gamma^{-2})$. Using the approximation $q(\mathbf{r}^{N},\mathbf{v}^{N},t)=q_{0}(\mathbf{r}^{N},t)+m\mathbf{v}\cdot\nabla_{\mathbf{r}}q_{0}(\mathbf{r}^{N},t)/\gamma+\mathcal{O}(\gamma^{-2})$ we consider the log transform, $Q=\ln q$, which to equivalent order in perturbation theory is

[TABLE]

where we will use $Q_{0}=\ln q_{0}$. For an underdamped equation of motion the relative action, $\Delta U_{\lambda}$ is given by

[TABLE]

just as in the overdamped case. Substituting the underdamped Doob force,

[TABLE]

into $\Delta U_{\lambda}[\mathbf{X}]$ we find,

[TABLE]

The first term can be resolved using Ito’s Lemma

[TABLE]

Substituting this back to the relative action, we get:

[TABLE]

Finally, using the perturbative approximation of $Q$ in Eq. A, and substituting the approximated form of the backward Kolmogorov equation in Eq. 29 yields

[TABLE]

Hence, $\Delta U_{\lambda^{*}}$ quantifies the transition probability between the states $A$ and $B$ over time $t_{f}$ when averaged over an initial distribution in $A$.

Appendix B Relative action for Langevin leap-frog integrator

The equations of motion for the Langevin leap-frog integrator is given by86; 82

[TABLE]

where the definitions of $\mathbf{v}_{i}$, $\mathbf{r}_{i}$ and $\mathbf{F}_{i}$ are the same as in Eq. 25, $m_{i}$ is the mass of particle $i$, $\gamma$ is friction coefficient, $\Delta t$ is the timestep, and $\alpha=\exp[-\gamma\Delta t]$. The noise, $\boldsymbol{\eta_{i}}$ is a Gaussian random variable with mean $\langle\boldsymbol{\eta}_{i}(t)\rangle=0$ and variance $\langle\boldsymbol{\eta}_{i}(t)\otimes\boldsymbol{\eta}_{j}(t^{\prime})\rangle=k_{\mathrm{B}}T(1-\alpha^{2})m_{i}^{-1}\delta_{ij}\mathbf{1}\delta(t-t^{\prime})$. For this discretization, the relative stochastic action is

[TABLE]

which is the same general form as in the overdamped case.

Appendix C Internal coordinates for alanine dipeptide

In the studies on alanine dipeptide we parameterized our NN ansatz for the time dependent committor based on a set of internal coordinates. In Tables I and II we define each of the angles and dihedrals referred to in the main text based on the atom numbering Fig. 9.