Learning quasistationarity diagnostics for accelerated molecular dynamics
New ideas in molecular dynamics (SCICADE 2026 minisymposium)
Mathematics for Materials Modeling, Institute of Mathematics, EPFL
June 29, 2026
Outline
% Blackboard Bold % Calligraphic % Positive definite symmetric matrices
% Script % Fraktur % Other Symbols % Alternative definition, safer % % % General Physics/Math % Geometry and Neighborhoods % Hessian Related % Harmonic Oscillators % Simplified/Indexed Notation % Other Specific Notation $$
\[ \definecolor{gris_C}{RGB}{96,96,96} \definecolor{blanc_C}{RGB}{255,255,255} \definecolor{pistache_C}{RGB}{176,204,78} \definecolor{pistache2_C}{RGB}{150,171,91} \definecolor{jaune_C}{RGB}{253,235,125} \definecolor{jaune2_C}{RGB}{247,208,92} \defineclor{darkblue}{RGB}{} \definecolor{orange_C}{RGB}{244,157,84} \definecolor{orange2_C}{RGB}{239,119,87} \definecolor{bleu_C}{RGB}{126,151,206} \definecolor{bleu2_C}{RGB}{90,113,180} \definecolor{vert_C}{RGB}{96,180,103} \definecolor{vert2_C}{RGB}{68,141,96} \definecolor{pistache_light_C}{RGB}{235,241,212} \definecolor{jaune_light_C}{RGB}{254,250,224} \definecolor{orange_light_C}{RGB}{251,230,214} \definecolor{bleu_light_C}{RGB}{222,229,241} \definecolor{vert_light_C}{RGB}{216,236,218} \]
- Motivation long-time simulation of metastable systems
- Theoretical results coarse-graining of killed Markov processes
- Application quasi-stationarity diagnostics with RNNs
Ongoing work with choices/extensions to discuss
Dynamical quantities in MD
Molecular Dynamics (MD) models atomic motion with a stochastic process \(X\) in a configuration space \(\mathcal{X}\) (allowed positions of nuclei)
Invariant distribution is the Gibbs measure for the interaction potential \(V:\mathcal{X}\to\mathbb{R}\) \[ \mu(\mathrm{d}x) = \frac{1}{Z_\beta}\mathrm{e}^{-\beta V(x)}\,\mathrm{d}x \in\mathcal{P}(\mathcal{X}),\quad \beta = 1/(k_{\mathrm{B}}T), \] giving access to thermodynamic/static quantities (via ergodic path averages)
dynamical/trajectorial/kinetic quantities (reaction rates, transport coefficients, time-correlations, reactive states…) are much more difficult (averages in path space)
Metastability and the timescale problem
Metastable trajectories on a free energy landscape
Dynamical properties typically require long trajectories because of metastability (timescale problem)
Multiple timescales(\(\mathrm{fs}\) simulation timestep/ \(\mathrm{ps}\) mesoscopic fluctuations/\(\mu\mathrm{s}-\mathrm{ks}\) macroscopic transitions)
Local description of metastability
- \(X\) enters a metastable state \(\Omega \subset \mathcal{X}\)
- \(X\) reaches a local equilibrium in \(\Omega\)
- \(X\) stays in local equilibrium for a long time before \(\tau_{\partial\Omega}=\)exit time
Local equilibria are quasi-stationary distributions
A probability measure \(\nu^X\in \mathcal{P}(\Omega)\) is a quasi-stationary distribution 1 (QSD) if for any \(A\subset\Omega\) and \(t>0\), \[ \mathbb{P}_{\nu^X}\left(X_t\in A \,|\, \tau_{\partial \Omega} > t\right) = \nu^X(A) \]
Under broad conditions on \(X\) and \(\Omega\), \(\exists!\) QSD \(\nu^X\) 2, and for each \(\eta_0^X \in \mathcal{P}(\Omega)\), \(\nu^X\) is the long-time limit in law conditioned on survival (Yaglom limit) \[ \underset{t\to+\infty}{\lim}\,\mathrm{d}_{\mathrm{TV}}(\eta_t^X,\nu^X) = 0,\qquad \eta_t^X = \mathrm{Law}\left(X_t\,\middle|\, \tau_{\partial\Omega} > t \right) \]
Accelerated dynamics
Accelerated MD methods exploit properties of the QSD to generate long approximate trajectories at lower cost
Tradeoff: short/accurate MD vs long/state-to-state AMD
Parallel replica method1
Sample QSD using long-time limit of Fleming-Viot system
Fleming-Viot process
Interacting particle system \(({\bf X}_t)_{t\geqslant 0}\) where \({\bf X}_t = \left(X_t^{(i)}\right)_{1\leqslant i\leqslant N}\in \mathcal{X}^N\)
Algorithm
- Each \(X^{(i)}\) is an independent replica \(X\)
- When \(X^{(i)}\) exits, resample independently uniformly from survivors
Particle approximation of \(\eta_t^X\) via propagation of chaos results \[ \widehat{\eta}^X_{t,N}\xrightarrow{N\to +\infty} \eta_t^X,\qquad \widehat{\eta}^X_{t,N} = \frac1N\sum_{i=1}^N \delta_{X_t^{(i)}}, \] uniformly in time (see2 for state of the art results)
Convergence diagnostics
Having a reliable and efficient algorithm requires knowing when \(\eta_t^X \approx \nu^X\)
Bottleneck develop sensitive/specific quasi-stationarity diagnostics \(\iff\) accurate estimates of decorrelation times
Gelman-Rubin diagnostic
\[ \widehat\eta_{t,N}^{X,(i)} = \frac1T\int_0^T \delta_{X_s^{(i)}}\,\mathrm{d}s,\qquad \widehat\eta_{t,N}^{X} = \frac1N\sum_{i=1}^N \widehat\eta_{t,N}^{X,(i)} \]
A Gelman-Rubin-type trace1 is defined, for some collective variable (CV) \(\xi:\mathcal{X}\to\mathbb{R}\), by \[ R_{t,N}(\xi) = \frac{\mathrm{Var}\left(\xi_\star\widehat\eta_{t,N}\right)}{\displaystyle\frac1N \sum_{i=1}^N \mathrm{Var}\left(\xi_\star \widehat\eta_{t,N}^{(i)}\right)}\xrightarrow{t\to +\infty} 1 \] GR decorrelation time \[ T_{\mathop{\mathrm{corr}}}^{GR}(\alpha,t_0) = \inf\left\{t\geqslant t_0,\middle|\,\underset{\xi\in \{\xi_1,\dots,\xi_\ell\}}{\max} R_{t,N}(\xi) < 1+\alpha\right\} \]
Possible approaches
- Physical heuristics 2
- Quantitative analysis in asymptotic regimes345
- MCMC (non)convergence diagnostics 6 on Fleming-Viot population(Gelman-Rubin)
- Our approach: effective dynamics (see7 8 for global constructions)
Only 3. and 4. are generic enough to apply to many systems (e.g. biomolecules)
- 3. assumes that each \(\xi(\mathbf{X})\) can be treated as a collection of independent Markov processes
- Requires knowledge of good \(\xi\)/neglects correlations between replicas
- Conservative rate of convergence
Coarse graining framework
Approximate the conditional law of \(\xi(X)\) with that of a Markovian approximation \(Z\)
Measure-valued flow
\(X,Z\) be an irreducible jump processes on finite state spaces \(\mathcal{X},\,\mathcal{Z}\) killed on \(\partial\mathcal{X},\partial\mathcal{Z}\). The law of \(Y\in\{X,Z\}\) conditioned on survival is a non-linear measure-valued flow \[ \eta_t^Y(f) = \frac{\eta_0^Y(Q_t^Y f)}{\eta_0^Y(Q_t^Y \mathbf{1}_{\mathcal{Y}})} \] where \(Q_t^Y\) is the killed semigroup \[ Q_t^Y f(y) = \mathbb{E}_y\left[f(Y_t)\mathbf{1}_{\tau_{\partial\mathcal{Y}}>t}\right] \]
Spectral decomposition
By Perron-Frobenius, spectral decomposition on \(Q^Y\), \((\lambda^Y, \nu^Y, h^Y)\) principal eigentriplet and a spectral gap \(\gamma^X>0\) \[ \|Q_t^Y\| \le C_\lambda^Y \mathrm{e}^{-\lambda^Y t},\qquad \|Q_t^Y P^Y\| \le C_\gamma^Y e^{-(\lambda^Y + \gamma^Y)t}\qquad P^Y f= f - \nu^Y(f)h^Y \]
Coarse-graining operators
Given a surjective CV \(\xi: \mathcal{X}\to \mathcal{Z}\) and a family \(\mathcal{F} =(\pi_z)_{z\in \mathcal{Z}}\), \(\pi_z\in \mathcal{P}(\xi^{-1}(z))\), define \[ \Pi^Z:\ell^\infty(\mathcal{X})\to\ell^\infty(\mathcal{Z}),\qquad \Pi^Z f(z) = \pi_z(f) \] and by pullback a projector onto fiber-constant observables \[ \Pi^X f = (\Pi^Z f)\circ\xi \]
Uniform bounds on the coarse-graining bias
Let \(\pi_z\) be the conditioning of \(\nu^X\) on \(\xi^{-1}(z)\), assume the timescale separation condition \[ \left\| e^{t(I-\Pi^X)\mathcal{L}^X(I-\Pi^X)} \right\| \le C_\zeta e^{-\zeta t} \]
Theorem 1 (in preparation)
Suppose \(\zeta > 2(\lambda^X + \gamma^X)\) and \(\eta_0^Z = \xi_\star \eta_0^X\)
There exists \(Z^*\in\mathrm{Markov}(\mathcal{Z})\) such that for all \(T > 0\), \[ \mathrm{d}_{\mathrm{TV}}\left(\xi_\star\eta_T^X,\, \eta_T^{Z^*}\right) \le \frac{1}{\zeta} \left( C e^{-\gamma^X T} + C' T e^{-(\gamma^{Z^*} \wedge \gamma^X) T} \right)\xrightarrow{T\to\infty}0 \] for explicit prefactors \(C,C'\) independent of \(T\)
We have quasimode estimates (obviously \(\nu^Z =\xi_\star \nu^X\)) \[ \|h^X - \Pi^X h^X\| \le \frac{C_\lambda^X C_\zeta}{\zeta - \lambda^X} \|\mathcal{L}^X \Pi^X - \Pi^X \mathcal{L}^X\|, \qquad \|\Pi^{Z^*} h^X - h^{Z^*}\| = \frac{C_\lambda^X C_\zeta C_\gamma^{Z^*}}{\gamma^{Z^*}(\zeta - \lambda^X)} \|\mathcal{L}^X \Pi^X - \Pi^X \mathcal{L}^X\|^2 \]
Proof approach by backward numerical analysis inspired by 2
Small-parameter illustration
Consider \(\mathcal{L}_\epsilon^X = J_{\text{mac}} - K + \epsilon^{-1} J_{\text{mic}}\), with \(J_{\text{mic}}=\bigoplus_{z\in\mathcal{Z}} A_z\) fiber-reducible (\(A_z\in \mathrm{Jump}(\xi^{-1}(z))\))
Supervised learning idea
- Teach a model to detect quasi-stationarity on an ensemble of processes \({Z}\) given trajectories of \(\mathbf{Z}\) (Fleming-Viot)
- Deploy the model given trajectories of \(\xi(\mathbf{X})\) (CV-Fleming-Viot population)
- Hope that scale-separation holds for \((\xi,X)\) and that the ensemble covers \(Z^*\)
Classification problem
Approximate the conditional probability \[ p_\theta(\mathscr{Z}) \approx h^*(\mathscr{Z}) := \mathbb{P}_{(z,t_i,t_f,Z)\sim \rho_{\mathrm{synth}}}\left(t_f>t_{\mathop{\mathrm{corr}},Z}(z)\middle|\,\Psi_{t_i,t_f}\left({\bf Z}^z\right)=\mathscr{Z}\right) \]
- Fleming-Viot started from \(z\): \({\bf Z}^z\)
- Honest decorrelation time \(t_{\mathop{\mathrm{corr}},Z}(z) = \inf\{t>0:\,\mathrm{d}_{\mathrm{TV}}\left(\eta_t^Z,\nu^Z\right) < \varepsilon\}, \eta_0^Z =\delta_z\)
- Featurization of \({\bf Z}^z\) on \([t_i,t_f]\): \(\Psi_{t_i,t_f}({\bf Z}^z)\)
- Synthetic distribution: \(\rho_{\mathrm{synth}}\in \mathcal{P}(\mathcal{Z}\times (0,\infty)^2\times \mathrm{Markov}(\mathcal{Z}))\)
Specific instantiation
One-dimensional CV \(\mathcal{Z}= (0,1)\) and \(\Psi_{t_i,t_f}({\bf Z}^z)\) is a sequence of time-pooled population histograms \[ \left(\frac{1}{\tau}\int_{t_i}^{t_i+\tau}\widehat{\eta}^Z_{s,N}\,\mathrm{d}s,\,\dots,\frac{1}{\tau}\int_{t_f-\tau}^{t_f}\widehat{\eta}^Z_{s,N}\,\mathrm{d}s\right) \] (recall \(\widehat{\eta}^Z_{s,N}\) is the empirical measure of \({\bf Z}_s^x\))
Randomize \((\tau,N)\) in training, and classify conditionally on this information
Note that \(h^*\) minimizes the cross-entropy functional \[ J[h] = \mathbb{E}_{\rho_{\mathrm{synth}}}\left[1_{t_f>t_{\mathop{\mathrm{corr}},Z}(z)}\log\,h\left(\Psi_{t_i,t_f}\left({\bf Z}^z\right)\right) - 1_{t_f\leqslant t_{\mathop{\mathrm{corr}},Z}(z)}\log\left\{1-h\left(\Psi_{t_i,t_f}\left({\bf Z}^z\right)\right)\right\}\right] \]
Synthetic ensemble
Synthetic reversible diffusion ensemble
Sample free energy/diffusion landscapes under \(\rho_{\mathrm{synth}}\)
Distribution of honest decorrelation times
Sample honest decorrelation times under \(\rho_{\mathrm{synth}}\)
Model and training
We train on subsequences of \(\Psi_{t_i,t_f}({\bf Z}^x) =(X_1,\dots,X_K)\)
Decorrelation times (hence \(K\)) spans several orders of magnitude natural choice of architecture is a recurrent neural network (RNN)
\[
\forall\,1\leqslant k\leqslant K,\quad e_k = E_\theta(X_k),\quad h_k = R_\theta(e_k,h_{k-1},\color{red}{N,\tau}),\quad p_\theta(X)_k = D_\theta(h_k),\quad h_0=0
\]
We minimize for \[ \theta^* \in \mathrm{Argmin}_\theta J[p_\theta] \] using Adam + batched Monte-Carlo estimators of \(J\) in practice
Implementation/training using Flux.jl, hyperparameter selection using Bayesian optimization framework NePS1
Test results: accuracy
We compare diagnoses on the original classification task in \(\rho_{\mathrm{synth}}\)
Accuracy of Gelman-Rubin diagnostic (th=\(\alpha\), stride=\(t_0\))
Accuracy of RNN diagnostic (decision threshold=\(\alpha\))
Test results: ROC curves
MD results: alanine dipeptide
Solvated alanine dipeptide and its backbone dihedrals \(\xi=(\phi,\psi)\)
Estimated \(\xi_\star\nu^X\), with initial samples for the Fleming-Viot process
MD results: time efficiency
RNN diagnostic
| Friction \ π₀ | Saddle 1 | Saddle 2 | Saddle 3 | Saddle 4 | Saddle 5 (minimum) | ||||||||||||||||||||||||||||||||||||||||||||||||||
|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|
| γ = 1 |
|
|
|
|
|
||||||||||||||||||||||||||||||||||||||||||||||||||
| γ = 2 |
|
|
|
|
|
||||||||||||||||||||||||||||||||||||||||||||||||||
| γ = 5 |
|
|
|
|
|
||||||||||||||||||||||||||||||||||||||||||||||||||
| γ = 10 |
|
|
|
|
|
||||||||||||||||||||||||||||||||||||||||||||||||||
| γ = 20 |
|
|
|
|
|
α ∈ {0.5, 0.6, 0.7, 0.8, 0.9}
Gelman-Rubin diagnostic
| Friction \ π₀ | Saddle 1 | Saddle 2 | Saddle 3 | Saddle 4 | Saddle 5 (minimum) | ||||||||||||||||||||||||||||||||||||||||||||||||||
|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|
| γ = 1 |
|
|
|
|
|
||||||||||||||||||||||||||||||||||||||||||||||||||
| γ = 2 |
|
|
|
|
|
||||||||||||||||||||||||||||||||||||||||||||||||||
| γ = 5 |
|
|
|
|
|
||||||||||||||||||||||||||||||||||||||||||||||||||
| γ = 10 |
|
|
|
|
|
||||||||||||||||||||||||||||||||||||||||||||||||||
| γ = 20 |
|
|
|
|
|
α ∈ {0.2, 0.1, 0.05, 0.02, 0.01}
Estimators of \(T_{\mathrm{corr}}(\alpha)/t_{\mathop{\mathrm{corr}}}\), where \(t_{\mathop{\mathrm{corr}}}\) is a decorrelation time for \(\xi_\star\eta_t^X\) (\(\xi\in\{\phi,\psi\}\)) estimated using many Fleming-Viot sample trajectories
MD results: bias
RNN diagnostic
| Friction \ π₀ | Saddle 1 | Saddle 2 | Saddle 3 | Saddle 4 | Saddle 5 (minimum) | ||||||||||||||||||||||||||||||||||||||||||||||||||
|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|
| γ = 1 |
|
|
|
|
|
||||||||||||||||||||||||||||||||||||||||||||||||||
| γ = 2 |
|
|
|
|
|
||||||||||||||||||||||||||||||||||||||||||||||||||
| γ = 5 |
|
|
|
|
|
||||||||||||||||||||||||||||||||||||||||||||||||||
| γ = 10 |
|
|
|
|
|
||||||||||||||||||||||||||||||||||||||||||||||||||
| γ = 20 |
|
|
|
|
|
α ∈ {0.5, 0.6, 0.7, 0.8, 0.9}
Gelman-Rubin diagnostic
| Friction \ π₀ | Saddle 1 | Saddle 2 | Saddle 3 | Saddle 4 | Saddle 5 (minimum) | ||||||||||||||||||||||||||||||||||||||||||||||||||
|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|
| γ = 1 |
|
|
|
|
|
||||||||||||||||||||||||||||||||||||||||||||||||||
| γ = 2 |
|
|
|
|
|
||||||||||||||||||||||||||||||||||||||||||||||||||
| γ = 5 |
|
|
|
|
|
||||||||||||||||||||||||||||||||||||||||||||||||||
| γ = 10 |
|
|
|
|
|
||||||||||||||||||||||||||||||||||||||||||||||||||
| γ = 20 |
|
|
|
|
|
α ∈ {0.2, 0.1, 0.05, 0.02, 0.01}
Thank you !
Future steps
- Generalize the theoretical bounds to (hypo)elliptic diffusions
- Train diagnostics for 2D/3D CVs
- Get serious about calibration of synthetic ensemble / benchmark on more serious systems
- Use transfer-operator learning (TICA, shallow learning a la1…) to parametrize \(\xi\) locally/on the fly
References
