Learning on-the-fly quasi-stationarity diagnostics from low-dimensional surrogates

Noé Blassel

stat.MLOct 26, 2026

Learning on-the-fly quasi-stationarity diagnostics from low-dimensional surrogates

Authors: Noé Blassel

Organizations: Mathematics for materials modelling, Institute of Mathematics EPFL, Lausanne, Switzerland

Abstract

Accelerated dynamics methods, such as parallel replica dynamics, require estimates of the decorrelation time, after which a process trapped in one of its metastable states is distributed according to the quasi-stationary distribution (QSD). In a method proposed by Binder, Lelièvre and Simpson, decorrelation times are defined by coupling a particle approximation of the non-linear conditioned probability flow of the process with MCMC-inspired diagnostics. We explore an alternative to MCMC diagnostics, based on coarse-graining of non-linear probability flows and simulation-based inference. We first develop a coarse-graining framework for killed Markov jump processes conditioned on survival, proving uniform-in-time, non-asymptotic estimates approximating the non-linear probability flow of the true process by the non-linear flow of a surrogate process, defined with a collective variable. The framework does not assume reversibility of the true process, and also gives bounds which scale explicitly in an adiabatic regime. Motivated by these results, we then construct a quasi-stationarity diagnostic using simulation-based inference: a recurrent neural network (RNN), trained on particle approximations of non-linear effective flows, estimates posterior probabilities of decorrelation. The learned diagnostic is then transferred to collective-variable traces of high-dimensional particle approximations. We investigate the behavior and robustness of this diagnostic on a set of numerical experiments, comparing it to the MCMC diagnostic of Binder et al.\ on both synthetic examples and a molecular system, in and out of its training prior distribution. When its prior assumptions hold, the RNN diagnostic outperforms the MCMC diagnostic in ROC analysis, is less conservative, fails less often, and has an interpretable decision threshold.

Figures & tables

Figure 1: Various terms appearing in the prefactors

C,C^{\prime}

from ( 52 ), as a function of the timescale separation

\zeta

Figure 2: Supremum in time of the total variation bias as a function of the timescale separation

\zeta

. The theoretical scaling in

1/\zeta

expected from Theorem 5 is indicated with a dashed line.

Figure 3: Total variation bias (solid lines) and theoretical bound from Theorem 5 (dotted lines) for various values of

\varepsilon

and choices of initial conditions.

Figure 4: Free energy profiles $U$ (solid lines) and diffusion profiles $D$ (dotted lines, in logarithmic scale) for sixteen independent samples from the training prior. The value of the inverse temperature $\beta$ is indicated with a color gradient.

Figure 5: Density per decade ( 179 ) of the honest decorrelation times

t_{0}(\mathcal{L}^{Z},Z_{0})

(black), of the RNN decorrelation times (blue) and of the Gelman–Rubin decorrelation times (orange), for decision thresholds

\alpha=0.1

(solid lines) and

\alpha=0.05

(dashed lines).

Figure 6: ROC curves for the RNN, Gelman–Rubin and histogram TV ( 183 ) diagnostics, within the training prior. Samples are aggregated over a grid of design parameters

(N,\tau)

. The ROC-AUC score of each diagnostic is indicated in the legend.

Figure 7: ROC-AUC scores of the RNN, Gelman–Rubin and histogram TV ( 183 ) diagnostics for the sub-classification tasks associated to each design setting $(N,\tau)$ , under the training prior.

Figure 8: Distributions of the declared decorrelation times

t_{\rm corr}(\alpha)

in the training prior, in units of

t_{0}(\mathcal{L}^{Z},Z_{0})

, for the RNN (blue) and Gelman–Rubin (orange) diagnostics, for

\alpha=0.1

(solid lines) and

\alpha=0.05

(dashed lines). The proportion of failed trajectories for each diagnostic is indicated in the legend.

Figure 9: Empirical decorrelation frequency as a function of empirical RNN predictions under the training prior. Perfect calibration is indicated with a dotted line. Confidence bands are computed with a bootstrap estimator, treating different Fleming–Viot trajectories on the same landscape

(U,D,\beta)

as independent samples.

Figure 10: Empirical frequency of decorrelation as a function of the probability predicted by the RNN diagnostic, under each prior shift (blue lines) and under the training prior (grey lines). Perfect calibration is indicated with a dotted line.

Figure 11: Configurations of the solvated dimer in each of its three metastable conformations. The dimer atoms are shown in black, and the solvent atoms within

1.5r_{c}

, contributing the bulk of

n_{\rm coord}

in green.

Figure 12: (a) Free energy of the dimer length

r_{12}

(solid line) and unsolvated free energy (dashed line). The state

\Omega

is shaded, and the initial points are indicated in blue. (b) Free energy as a function of

\overline{\xi}

\Omega

, with draws

\beta U

from the training prior (grey lines).

Figure 13: Histogram distances

\widehat{D}_{f}(t)

( 198 ) to their quasi-stationary state as a function of time, for

f=\overline{\xi},n_{\rm coord}

and

(\overline{\xi},n_{\rm coord})

. Columns correspond to different initial conditions, and rows to different values of the parameter

D_{d}

. Reference decorrelation times

t_{0}(f)

defined in ( 199 ) are indicated with vertical lines (solid gray:

f=\overline{\xi}

, dashed pink:

f=(\overline{\xi},n_{\rm coord}))

. The tolerance threshold

\varepsilon_{\rm corr}

is indicated with dotted horizontal lines.

Figure 14: Distributions of the declared decorrelation times

t_{\rm corr}(\alpha)

, in units of

t_{0}(\overline{\xi})

, for the RNN (blue) and Gelman–Rubin (orange) diagnostics on the solvated dimer, for

\alpha=0.1

(solid lines) and

\alpha=0.05

(dashed lines). The proportion of trajectories on which each diagnostic fails is indicated in each panel.

Figure 15: Exit distribution from diagnosed samples, for $D_{d}=0.05$ (top) and $D_{d}=1$ (bottom). The left panels show empirical survival functions for the exit time, with the short-time behavior represented in the insets. Right panels show the empirical probability of exiting through $\overline{\xi}=0$ for each initial point, with $95\%$ confidence intervals, treating Fleming–Viot particles as independent. The reference values and its confidence interval are indicated by the black line and grey band, and the heuristic error level $\pm\varepsilon_{\rm corr}/2$ by dotted lines.

Figure 16: Empirical survival functions of the exit time conditioned on the exit point, through

\overline{\xi}=0

(red) or

\overline{\xi}=1

(blue), starting from a Fleming–Viot population. Left panel: reference decorrelated population. Middle panel: population decorrelated according to the RNN. Right panel: population decorrelated according to the Gelman–Rubin diagnostics.

Figure 17: Behavior of the two-variable diagnostics for

\alpha=0.1

. Top left: distribution of decorrelation times in units of

t_{0}(\overline{\xi},n_{\rm coord})

(see Figure 14 ). Top right: exit-side distribution per initial point (see right panels of Figure 15 ). Bottom: survival functions conditioned on the exit side, for the RNN (left) and Gelman–Rubin (right, see Figure 16 ).

Table 1: Summary statistics for the RNN (a) and Gelman–Rubin (b) diagnostics in the experiments of Section 4 . The column

t_{\rm corr}/t_{0}

gives

50\%\,(25\%,75\%)

quantiles of

t_{\rm corr}(\alpha)/t_{0}

over all trajectories. Failed trajectories with

t_{\rm corr}(\alpha)>10t_{0}

are reported as

{>}10

. The columns “early” and “failed” give the empirical false decorrelation risk (

t_{\rm corr}(\alpha)<t_{0}

) and failure risk (

t_{\rm corr}(\alpha)>10t_{0}

), respectively. The reference

t_{0}

corresponds to ( 180 ) for the synthetic experiments, and ( 199 ) for the molecular systems (“with

n_{\rm coord}

” corresponds to

f=(\overline{\xi},n_{\rm coord})

, or else

f=\overline{\xi}

only).

Appendix figures & tables3 assets

Supplementary material from the paper’s appendix.

Appendix

Panel of Figure 18 — Figure 18: Clustering of $\sigma(-T(\varepsilon))$ (blue dots, see ( 202 )). Black dots represent $\sigma(-T(0))=\sigma(\mathcal{L}_{\rm mic})$ . The right panel details the structure of the spectrum inside $\Gamma^{\rm s}$ , which comes from the probabilistic nature of $T^{(1)}$ (see ( 213 ) and the Perron–Frobenius argument above).

prior	$[\beta_{\min},\beta_{\max}]$	$[s_{\min},s_{\max}]$	$\delta_{\max}$
training	$[1,3]$	$[-4,-1]$	$0.2$
high $D$	$[1,3]$	$[-1,1]$	$0.2$
low $T$	$[4,6]$	$[-4,-1]$	$0.2$
low $T$ , high $D$	$[4,6]$	$[-1,1]$	$0.2$
low $T$ , one well	$[4,6]$	$[-4,-1]$	$0.05$

Appendix

Table 2: Parameter ranges for the training prior and the misspecified priors considered in Section 4.3 .

layer type(/non-linearity)	output dimension	number of parameters
input : $H_{k}\in\mathbb{R}^{d_{\rm hist}}$ , $d_{\rm hist}=64$ , $m(N,\tau)=\sqrt{N\tau}\in\mathbb{R}$
featurizer $f_{\theta}:\mathbb{R}^{d_{\rm hist}}\to\mathbb{R}^{d_{\rm cnn}}$
Conv1D $(5,\,1\to 16)$ / MaxPool $\,{}_{2}\circ$ leakyReLU	$32\times 16$	$96$
Conv1D $(5,\,16\to 32)$ / MaxPool $\,{}_{2}\circ$ leakyReLU	$16\times 32$	$2\,592$
Conv1D $(5,\,32\to 64)$ / MaxPool $\,{}_{2}\circ$ leakyReLU	$8\times 64$	$10\,304$
reshape layer $f_{\theta}(H_{k})\star m(N,\tau)$

Appendix

Table 3: Architecture of the deployed RNN diagnostic. Convolutional layers use a kernel width of

5

and a zero-padding size of

2

Markov Chains

See all 31 papers

Oct 12, 2026

$\texttt{NFDiag}$ : High-dimensional MCMC Convergence Diagnosis with Normalizing Flows

Accurately assessing the convergence of Monte Carlo Markov Chains is crucial for ensuring their validity in Bayesian inference problems. However, convergence assessment in high dimensions is known to be difficult, due to most convergence diagnostics becoming insensitive beyond a few dozen parameters. Conventional convergence statistics analyze marginal distributions which in high dimensions neglects increasing levels of information from the joint and conditional probability distributions. Thus we propose

\texttt{NFDiag}

, which taps into the information of the full multivariate joint posterior with a normalizing flow (NF). Our convergence diagnostic thereby also analyzes information from conditional distributions. In other words, by working in the full multivariate space, it is sensitive with respect to the geometry of the chain in high dimensions. Our diagnostic utilizes that NFs assume independent and identically distributed (iid) training samples, which is a testable assumption that will be violated by a non-converged or strongly correlated chain. Our diagnostic trains NFs on separate parts of a chain, and then assesses whether these correspond to the same multivariate target distribution. We test our diagnostic in up to 250 dimensions, and demonstrate it is universal and unbiased, and detects problematic chain quality also when Gelman-Rubin

\hat{R}

and the multivarite effective sample size do not.

Kutay Nazlı, Elena Sellentin

Leiden Observatory, Leiden University, Einsteinweg 55 2333 CC Leiden, Netherlands · Mathematical Institute, Leiden University, Einsteinweg 55 2333 CC Leiden, Netherlands

Markov Chain Monte Carlo Bayesian Inference

Ask PDF

Aug 19, 2026

Heteroscedastic Neural Surrogate Modeling for Robust and Rapid Bayesian Inference in Fusion Plasma Diagnostics

Bayesian inference via Markov Chain Monte Carlo (MCMC) provides effective parameter estimation, but its real-time application in complex physical systems is hindered by heavy computational bottlenecks and extreme sensitivity to statistical noise. We address this by proposing a neural-network-based probabilistic surrogate framework for rapid and robust MCMC inference. Using fusion plasma Thomson scattering diagnostics as a challenging, noise-dominated testbed, our approach employs a dual-head architecture to simultaneously estimate the expected physical emission spectrum and the channel-wise intrinsic measurement noise variance. By optimizing a Gaussian Negative Log-Likelihood (GNLL) objective, the learned aleatoric uncertainty dynamically buffers the sampler against pathological shot noise. Evaluations demonstrate that this surrogate framework achieves > 1500x acceleration over exact physical forward models, while simultaneously reducing inference error (RMSE) by >20% compared to standard homoscedastic neural baselines, offering a highly promising paradigm for real-time physical analysis.

Liyun Zhang, Naoya Mamada, Kentaro Sakai +2

GSFS, The University of Tokyo Japan

Neural Surrogate Modeling Bayesian Inference Markov Chain Monte Carlo Heteroscedastic Regression

Ask PDF

Jun 8, 2026

Backward Coherence and Hidden-State Stability in Recurrent Neural Networks: A Quasi-Reverse-Martingale Theory

Recurrent neural networks maintain a hidden state

h_t

, but its probabilistic meaning is often unclear. We study hidden-state stability through backward coherence: the extent to which

h_t

can be reconstructed from

h_{t+1}

by a learned backward projector

g_φ

. Under contraction and summable backward drift, the hidden-state sequence forms a quasi-reverse-martingale. This yields almost-sure convergence, rates under mixing, an interpretable limiting representation, finite pathwise stopping times, and a theoretical framework for time-uniform confidence sequences. Simulations support the theory. Backward-coherence regularisation reduces the empirical quasi-martingale total

\hat Q

43

58%

, reaches stability

28

44%

earlier than an unregularised RNN, and gives tracking-error recovery consistent with geometric bounds. Additional tests confirm echo-state forgetting rates bounded by

ρ

and verify the increment-sum tube

R_t

with

100%

simultaneous coverage, although

R_t

is conservative; in practice, the defect-tail proxy

\hat Q_t

is the more useful monitor. The backward-coherence loss is also equivalent to minimising a Kullback--Leibler divergence in a Gaussian backward model, linking the method to variational inference. Extensions cover

φ

-mixing inputs, change-point tracking, and finite-sample concentration. Three real-data studies further validate the approach. On PhysioNet 2012 ICU data, the Reverse Martingale RNN (RMRNN) matches RNN mortality-prediction AUC while reaching stable representations 13 hours earlier. On FRED-MD, it reduces one-month-ahead forecast error by about fourfold under concept drift. On UCI Human Activity Recognition, it maintains lower post-transition tracking error with geometric decay. The guarantees apply under the stated assumptions; universality is not claimed.

Yuan-chin Ivan Chang

Institute of Statistical Science, Academia Sinica · Academia Road, Section 2, Nankang, Taipei 115, Taiwan

Recurrent Neural Networks Neural Network Stability

Ask PDF

Abstract

Markov Chains

Explore similar work

NFDiag\texttt{NFDiag}NFDiag: High-dimensional MCMC Convergence Diagnosis with Normalizing Flows

Heteroscedastic Neural Surrogate Modeling for Robust and Rapid Bayesian Inference in Fusion Plasma Diagnostics

Backward Coherence and Hidden-State Stability in Recurrent Neural Networks: A Quasi-Reverse-Martingale Theory

$\texttt{NFDiag}$ : High-dimensional MCMC Convergence Diagnosis with Normalizing Flows