Optimization of metastable timescales for accelerated sampling of trajectories

Molecular dynamics

Molecular dynamics (MD) probes the behavior of atomistic system at the level of classical motion of nuclei.

The timescale problem in MD

Dynamical setting

Molecular trajectories are solutions to a stochastic differential equation: \[ \mathrm{d}X_t = -a(X_t)\nabla V(X_t)\,\mathrm{d}t + \frac1\beta\mathrm{div}\,a(X_t)\,\mathrm{d}t + \sqrt{\frac{2}{\beta}}a(X_t)^{1/2}\,\mathrm{d}W_t, \]

• \(V:{\mathbb{{R}}}^d\to {\mathbb{{R}}}\) potential energy function, \(a:{\mathbb{{R}}}^{d\times d}\) s.p.d. diffusion coefficient, \(W\) standard Brownian motion in \({\mathbb{{R}}}^d\).

• Ergodic and reversible with respect to the standard Gibbs measure at temperature \((k\beta)^{-1}\) \[ \pi(\mathrm{d}x) = \mathrm{e}^{-\beta V(x)}\,\mathrm{d}x/Z_\beta. \]

• Continuous-time preconditioned Langevin Monte-Carlo / diffusion on the Riemannian manifold \(\mathcal M=({\mathbb{{R}}}^d,a^{-1})\)

• Evolution semigroup generator is symmetric on \(L^2({\mathbb{{R}}}^d,\pi)\): \[ \mathcal{L}_\beta \varphi = \frac1\beta\mathrm{e}^{\beta V}\mathrm{div}\,\left(\mathrm{e}^{-\beta V}a\nabla\varphi\right). \]

• We assume everywhere that \(V,a\in\mathcal W^{2,\infty}_{\mathrm{loc}}\), and \(a\) is locally elliptic: \(a^{-1}\in L^\infty_{\mathrm{loc}}\).

Local approach to metastability

• For \(\Omega\subset{\mathbb{{R}}}^d\), local equilibrium is given by a quasi-stationary distribution (QSD) in \(\Omega\) i.e. a probability \(\nu_\Omega\in \mathcal M_1(\Omega)\) for which \[ \forall\,A\subset \Omega,\qquad \mathbb P_{\nu_\Omega}\left(X_t \in A\,\middle|\,t<T_{\mathrm{exit}}(\Omega)\right) = \nu_\Omega(A). \]
• Since \(V\), \(\Omega\) are bounded and \(a\) is uniformly elliptic on \(\Omega\), the QSD is unique, and related to the principal Dirichlet eigenvector: \[ \nu_\Omega(\mathrm{d}x) = \frac{u_{1,\Omega}(x)\mathrm{e}^{-\beta V(x)}\,\mathrm{d}x}{\int_\Omega u_{1,\Omega}\mathrm{e}^{-\beta V}},\qquad -\mathcal{L}_\beta u_{1,\Omega}(x) = \lambda_{1,\Omega}u_{1,\Omega}(x), \]
• Here, we consider the generator \(\mathcal{L}_\beta\) with absorbing Dirichlet boundary conditions on \(\partial \Omega\): \[ \mathcal D(\mathcal{L}_\beta) = H_0^1(\Omega) \cap \{u\in L^2(\Omega):\,\mathcal{L}_\beta u\in L^2(\Omega)\subset\mathcal D'({\mathbb{{R}}}^d)\}. \] It is self-adjoint with compact resolvent on \(L^2_\beta(\Omega):=L^2(\Omega,\mathrm{e}^{-\beta V(x)}\,\mathrm{d}x)\) for bounded \(\Omega\).

Spectral characterization of timescales

Denote \(\mu_t = \mathrm{Law}\left(X_t \middle| t<T_{\mathrm{exit}}(\Omega)\right)\) the law of the process conditioned on remaining trapped.

• Theorem (Le Bris et al. (2012))

  \[ \mu_0=\nu_\Omega \implies T_{\mathrm{exit}}(\Omega)\sim \mathcal E(\color{blue}{\lambda_{1,\Omega}}),\qquad T_{\mathrm{exit}}(\Omega) \perp\!\!\!\!\perp X_{T_{\mathrm{exit}(\Omega)}}. \]

• This result justifies the use of the Parallel Replica algorithm.

• The exit rate is given by the principal Dirichlet eigenvalue \(\lambda_{1,\Omega} = 1/\mathbb E_{\nu_\Omega}[T_{\mathrm{exit}}(\Omega)]\).

• Theorem (Le Bris et al. (2012))

  Convergence to the QSD and bias on the exit event are exponentially decaying. There exist \(C_1,C_2>0\) depending on \(\mu_0\) such that \[ \left\|\mu_t-\nu_\Omega\right\|_{\mathrm{TV}} \leqslant C_1 \mathrm{e}^{-(\color{red}{\lambda_{2,\Omega}-\lambda_{1,\Omega}})t}, \] \[ \left\|\mathrm{Law}_{\mu_t}(T_{\mathrm{exit}(\Omega)}-t,X_{T_{\mathrm{exit}}(\Omega)}) - \mathrm{Law}_{\nu_\Omega}(T_{\mathrm{exit}(\Omega)},X_{T_{\mathrm{exit}}(\Omega)})\right\|_{\mathrm{TV}}, \leqslant C_2 \mathrm{e}^{-(\color{red}{\lambda_{2,\Omega}-\lambda_{1,\Omega}})t}. \]

• The decorrelation rate is given by the spectral gap of the Dirichlet generator \(\lambda_{2,\Omega}-\lambda_{1,\Omega}\).

• Set \(T_{\mathrm{corr}}(\Omega) = -\log\,\varepsilon_{\mathrm{corr}}/(\lambda_{2,\Omega}-\lambda_{1,\Omega})\) in ParRep, for some fixed tolerance parameter \(0<\varepsilon_{\mathrm{corr}}<1\).

• In practice, raises the need for quantitative or statistical estimates of \(\lambda_{2,\Omega}-\lambda_{1,\Omega}\) (and also \(C_1,C_2\)).

A spectral shape optimization problem

• The goal becomes to maximize the efficiency of ParRep by maximizing the separation of timescales: \[ \underset{\Omega \subset {\mathbb{{R}}}^d\text{ open, bounded, connected}}{\max}\, \left(\color{red}{\lambda_{2,\Omega}-\lambda_{1,\Omega}}\right)/\color{blue}{\lambda_{1,\Omega}}. \]
• Make the exit rate as small as possible compared to the decorrelation rate. Quantitative measure of local metastability / timescale separation inside \(\Omega\).
• Standard choice takes \(\Omega=\mathcal A(z_0)\), the basin of attraction of some local minimum \(z_0\) for \(V\) under the steepest-descent dynamics \(\dot X(t) = -a(X(t))\nabla V(X(t))\).
• The case of the Dirichlet Laplacian \(a=\mathrm{Id}\), \(V=0\) is a famous question in spectral geometry (Payne-Polya-Weinberger conjecture).
• Question has been identified by MD practitioners, e.g. in Perez, Uberuaga, and Voter (2015), but not studied in depth.
• Problem is ill-posed as a global optimization problem, e.g., for \(V(x) = |x|^2/2 + \delta V(x)\), \(\delta V\in \mathcal C^\infty_c({\mathbb{{R}}}^d)\): \[ \lambda_{2,B(0,R)}/\lambda_{1,B(0,R)}\xrightarrow{R\to +\infty} +\infty. \]
• Typical situation in spectral optimization: introduce constraints, or focus on local optimization.

Local optimization with steepest ascent method

Explicit formulas for shape perturbations

Galerkin objective for high-dimensional systems

For systems of interest, \(d\gg 1\). In MD, common to use a collective variable (CV) \(\xi:{\mathbb{{R}}}^d\to{\mathbb{{R}}}^m\) with \(m\ll d\) for interpretation and biasing.

Example: solvated Alanine dipeptide

Consider alanine-dipeptide in solution \((d=1857)\).

Shape-optimization results

Shape-gradient ascent, implemented with FreeFem++.

Shape-optimization results

Degeneracy-aware steepest ascent method implemented in FreeFem++.

Quantifying gains in the original objective

Measure gain in timescale separation with respect to reference state definition, in a practical dynamical setting.

Inference of decorrelation timescale

For the various ensembles of initial conditions, infer decorrelation rate at the level of histogram convergence.

Comparison of free-energy basin and optimized state

Asymptotic optimization methodology in the low-temperature limit

• Introduce a parametric family of domains \((\Omega_{\beta,\alpha})_{\alpha\in\mathcal X,\beta>0}\), jointly parametrized by the asymptotic parameter \(\beta\) and the shape parameter \(\alpha\).
• Compute quantitative asymptotics of the objective: \[ \frac{\lambda_2(\Omega_{\beta,\alpha})}{\lambda_1(\Omega_{\beta,\alpha})} = \mathcal F(\beta,\alpha)(1+\mathrm{o}(1)) \] for some explicit \(\mathcal F\) in the limit \(\beta\to+\infty\).
• Optimize \(\mathcal F\) with respect to \(\alpha\) (the optimal parameter may itself depend on \(\beta\)).
• Main difficulty is to compute spectral asymptotics for domains which change with the temperature.
• Equivalent to the semiclassical analysis of the Witten Laplacian \[ -\mathrm{e}^{-\beta V/2}\mathcal{L}_\beta\mathrm{e}^{\beta V/2} = -\frac1\beta \Delta + \frac{\beta}{4}|\nabla V|^2 - \frac12\Delta V \] with \(\hbar:=1/\beta\)-dependent Dirichlet boundary conditions.

Main geometric assumption

• We restrict to \(a=\mathrm{Id}\), assume \(V\) is a Morse function (non-degenerate Hessians at critical points), and \(\Omega_\beta \subset \mathcal K\) with \(\mathrm{Crit} V|_{\mathcal K}=\{z_0,\dots,z_{N-1}\}\).
• Assume the following limit is well-defined in \((-\infty, +\infty]\) for each \(0\leqslant i < N\). \[ \alpha_i = \underset{\beta\to\infty}{\lim}\, \sqrt\beta\sigma_{\Omega_\beta}(z_i),\qquad \sigma_{\Omega_\beta}(x) = \begin{cases} d(x,\partial\Omega_\beta),& x\in \Omega_\beta,\\ -d(x,\partial\Omega_\beta)&x\in\mathcal X\setminus \Omega_\beta \end{cases}. \]
• If \(\alpha_i = +\infty\), \(z_i\) is far from the boundary, and if not, it is close.
• Assume there exists a function \(\delta :{\mathbb{{R}}}_+\to (0,\varepsilon(V))\) such that \[ \sqrt\beta\delta(\beta)\xrightarrow{\beta\to\infty}\,+\infty,\qquad\mathcal{O}_{i}^{-}(\beta) \subseteq B(z_i,\delta(\beta))\cap\Omega_\beta \subseteq \mathcal{O}_{i}^{+}(\beta), \] where the sets \(\mathcal{O}_{i}^{\pm}\) constrain the geometry of the boundary near the \(z_i\) in a specific way.

Main geometric assumption

First-order behavior for the full spectrum

Further geometric assumptions

• We now restrict the setting to the case of only one minimum far from the boundary \(z_0\) (plus arbitrary other critical points).
• Harmonic approximation then implies \(\lambda_{1,\Omega_\beta}\to 0\), \(\lambda_{2,\Omega_\beta}\to \lambda_{2,\alpha}^{\mathrm{H}}>0\).
• Define \(V^* = \min\left\{V(z),\,z\in\partial \left[{\mathbb{{R}}}^d\setminus \mathcal A(z_0)\right]\right\}\), where \(\mathcal A(z_0)=\) basin of attraction of \(z_0\) for \(X'(t)=-\nabla V(X(t))\).
• The minimizers are index one saddle points \(\left\{z_i,i\in I_{\min}\right\}\).
• Assume also the energetic condition \(\left(\mathcal A(z_0)\cap \{V<V^*+c\delta(\beta)^2\} \setminus \cup_{i\in I_{\min}} B(z_i,\delta(\beta))\right) \subset \Omega_\beta\) for all \(\beta\) and some \(c>0\).
• This condition avoids the presence of generalized saddle points with energy \(\leqslant V^*\), ensuring that the likeliest exit pathway is intrinsic to \(V\), and not an artefact of the domain geometry.

Modified Eyring-Kramers formula

Illustration on a 1D energy well

Semiclassical approximation of shape-optimization landscape