Neural Transport Nested Sampling

David Yallup; Will Handley

cs.LGSep 24, 2026

Neural Transport Nested Sampling

Authors: David Yallup, Will Handley

Organizations: Kavli Institute for Cosmology Cambridge Institute of Astronomy, University of Cambridge, Madingley Road, Cambridge, CB3 0HA, UK

Abstract

Sampling from Boltzmann distributions of molecular systems is an inference problem that has seen significant recent developments fuelled by advances in neural density estimation. We develop a novel sampling algorithm, Neural Transport Nested Sampling (NTNS), which combines the classical strengths of nested sampling with modern neural flow-based methods. NTNS uses a flow matching velocity as the drift in a Metropolis--Hastings corrected Langevin kernel inside a nested sampling outer loop, requiring only evaluations of the target energy function and providing scalable estimation of the full partition function of high-dimensional particle systems. We benchmark NTNS on challenging molecular sampling benchmarks, scaling up to Lennard--Jones clusters of 55 interacting particles, where it reduces both interatomic distance and energy Wasserstein errors to reference MCMC by over an order of magnitude relative to the strongest neural baselines at lower wall-clock cost. To our knowledge, NTNS is also the first neural sampler to return a calibrated, temperature resolved partition function estimate at this scale, recovering the phase structure across temperature from a single run.

Figures & tables

	DW-4		LJ-13		LJ-55
Method	$r$ -W 2 $\downarrow$	$E$ -W 2 $\downarrow$	$r$ -W 2 $\downarrow$	$E$ -W 2 $\downarrow$	$r$ -W 2 $\downarrow$	$E$ -W 2 $\downarrow$
iDEM	$0.296\,{\color[rgb]{0.5,0.5,0.5}\scriptstyle\pm\,0.043}$	$\mathbf{0.226}\,{\color[rgb]{0.5,0.5,0.5}\scriptstyle\pm\,0.072}$	$0.102\,{\color[rgb]{0.5,0.5,0.5}\scriptstyle\pm\,0.007}$	$55.7\,{\color[rgb]{0.5,0.5,0.5}\scriptstyle\pm\,37.0}$	$0.135\,{\color[rgb]{0.5,0.5,0.5}\scriptstyle\pm\,0.001}$	$48.4\,{\color[rgb]{0.5,0.5,0.5}\scriptstyle\pm\,0.42}$
AS	$0.212\,{\color[rgb]{0.5,0.5,0.5}\scriptstyle\pm\,0.033}$	$0.445\,{\color[rgb]{0.5,0.5,0.5}\scriptstyle\pm\,0.086}$	$0.044\,{\color[rgb]{0.5,0.5,0.5}\scriptstyle\pm\,0.005}$	$1.39\,{\color[rgb]{0.5,0.5,0.5}\scriptstyle\pm\,0.24}$	$0.060\,{\color[rgb]{0.5,0.5,0.5}\scriptstyle\pm\,0.002}$	$36.4\,{\color[rgb]{0.5,0.5,0.5}\scriptstyle\pm\,0.40}$
ASBS	$0.172\,{\color[rgb]{0.5,0.5,0.5}\scriptstyle\pm\,0.044}$	$0.304\,{\color[rgb]{0.5,0.5,0.5}\scriptstyle\pm\,0.116}$	$0.029\,{\color[rgb]{0.5,0.5,0.5}\scriptstyle\pm\,0.005}$	$2.24\,{\color[rgb]{0.5,0.5,0.5}\scriptstyle\pm\,0.22}$	$0.082\,{\color[rgb]{0.5,0.5,0.5}\scriptstyle\pm\,0.001}$	$40.2\,{\color[rgb]{0.5,0.5,0.5}\scriptstyle\pm\,0.50}$
NTNS	$\mathbf{0.154}\,{\color[rgb]{0.5,0.5,0.5}\scriptstyle\pm\,0.034}$	$0.286\,{\color[rgb]{0.5,0.5,0.5}\scriptstyle\pm\,0.086}$	$\mathbf{0.018}\,{\color[rgb]{0.5,0.5,0.5}\scriptstyle\pm\,0.008}$	$\mathbf{0.58}\,{\color[rgb]{0.5,0.5,0.5}\scriptstyle\pm\,0.15}$	$\mathbf{0.008}\,{\color[rgb]{0.5,0.5,0.5}\scriptstyle\pm\,0.001}$	$\mathbf{2.68}\,{\color[rgb]{0.5,0.5,0.5}\scriptstyle\pm\,0.17}$

Table 1 : Wasserstein-2 distances to MCMC reference samples for interatomic distances (

r

-W 2 ) and energy (

E

-W 2 ) for NTNS (ours) and other neural samplers. Bold text indicates the best performance for each metric. For LJ-55, the NTNS and iDEM entries are median-seed runs drawn from the five-seed sweep detailed in appendix G .

Appendix figures & tables5 assets

Supplementary material from the paper’s appendix.

Appendix

Nested sampling
Parameter	DW-4	LJ-13	LJ-55
Live points $m$	2048	2048	2048
Deletion / chains $k$	1024	1024	1024
Inner steps $T$	8	39	165
Initial prior oversample factor	100	100	100
Retrain interval	every iteration	every iteration	every iteration

Appendix

Table 2 : Hyperparameters for NTNS experiments.

Ablation	Setting	$r$ - $W_{2}$	$E$ - $W_{2}$	Runtime (s)	Evals (M)
Baseline	-	$0.026\,{\color[rgb]{0.5,0.5,0.5}\scriptstyle\pm\,0.006}$	$0.605\,{\color[rgb]{0.5,0.5,0.5}\scriptstyle\pm\,0.129}$	583	10.15
Inner steps $T$	$19$	$0.031\,{\color[rgb]{0.5,0.5,0.5}\scriptstyle\pm\,0.006}$	$0.476\,{\color[rgb]{0.5,0.5,0.5}\scriptstyle\pm\,0.077}$	598	5.05
	$78$	$0.021\,{\color[rgb]{0.5,0.5,0.5}\scriptstyle\pm\,0.008}$	$0.541\,{\color[rgb]{0.5,0.5,0.5}\scriptstyle\pm\,0.185}$	612	20.17
Drift time $t^{*}$	$0.1$	$0.031\,{\color[rgb]{0.5,0.5,0.5}\scriptstyle\pm\,0.006}$	$0.570\,{\color[rgb]{0.5,0.5,0.5}\scriptstyle\pm\,0.141}$	607	10.19
	$0.25$	$0.027\,{\color[rgb]{0.5,0.5,0.5}\scriptstyle\pm\,0.005}$	$0.895\,{\color[rgb]{0.5,0.5,0.5}\scriptstyle\pm\,0.168}$	617	10.15
	$0.5$	$0.044\,{\color[rgb]{0.5,0.5,0.5}\scriptstyle\pm\,0.009}$	$1.078\,{\color[rgb]{0.5,0.5,0.5}\scriptstyle\pm\,0.232}$	606	10.19

Appendix

Table 3 : Ablation results on LJ-13. Single-axis perturbations to NTNS on the

d{=}39

Lennard-Jones target; the top row is the baseline, each block varies one hyperparameter.

r

W_{2}

E

W_{2}

are mean

\pm

std of the

1

D Wasserstein-2 over

10

batches of

1000

NS-weighted draws vs. permuted MCMC reference batches; runtime is wall time on one GH200.

Method	$m/k$	Runs	$r$ - $W_{2}$ $\downarrow$	$E$ - $W_{2}$ $\downarrow$	$\log\mathcal{Z}(1)$	Time (s)	Evals
NTNS	$2048/1024$	$5$	$\mathbf{0.017}\,{\color[rgb]{0.5,0.5,0.5}\scriptstyle\pm\,0.020}$	$\mathbf{3.97}\,{\color[rgb]{0.5,0.5,0.5}\scriptstyle\pm\,2.03}$	$233.85\,{\color[rgb]{0.5,0.5,0.5}\scriptstyle\pm\,0.83}$	$19{,}541$	$8.0{\times}10^{7}$
NSS	$2048/1024$	$10$	$0.066\,{\color[rgb]{0.5,0.5,0.5}\scriptstyle\pm\,0.052}$	$25.74\,{\color[rgb]{0.5,0.5,0.5}\scriptstyle\pm\,19.37}$	$230.7\,{\color[rgb]{0.5,0.5,0.5}\scriptstyle\pm\,11.6}$	$144\,{\color[rgb]{0.5,0.5,0.5}\scriptstyle\pm\,15}$	$4.82{\times}10^{8}$

Appendix

Table 4 : Matched LJ-55 comparison between NTNS and classical NSS. Values are mean

\pm

standard deviation across independent runs. Lower

W_{2}

is better. Runtime is wall time on one GH200; evaluations count target-energy calls.

	DW-4		LJ-13		LJ-55
Method	Time (s)	Evals	Time (s)	Evals	Time (s)	Evals
iDEM	9,633	$1{\times}10^{8}$	19,877	$1{\times}10^{8}$	$43{,}200^{*}$	$6{\times}10^{7}$
AS	7,260	$4{\times}10^{6}$	$43{,}200^{*}$	$4{\times}10^{6}$	$43{,}200^{*}$	$\mathbf{3{\times}10^{5}}$
ASBS	6,960	$4{\times}10^{6}$	34,680	$\mathbf{3{\times}10^{6}}$	$43{,}200^{*}$	$\mathbf{3{\times}10^{5}}$
NTNS	233	$\mathbf{2{\times}10^{6}}$	583	$1{\times}10^{7}$	19,541	$8{\times}10^{7}$

Appendix

Table 5 : Computational cost comparison; all runs use a single NVIDIA GH200 GPU. NTNS achieves substantially the lowest walltime across all baselines, while AS / ASBS use fewer energy queries by design—their adjoint-matching objective only requires

\nabla E

at trajectory endpoints. Every iDEM, AS and ASBS evaluation is differentiated to obtain

\nabla E

; NTNS evaluations are of

E

alone. Entries marked ∗ reached the 12 h training time limit before convergence.

Method	$r$ - $W_{2}$ $\downarrow$	$E$ - $W_{2}$ $\downarrow$	$\log\mathcal{Z}(1)$
iDEM	$0.136\,{\color[rgb]{0.5,0.5,0.5}\scriptstyle\pm\,0.093}$	$78.6\,{\color[rgb]{0.5,0.5,0.5}\scriptstyle\pm\,75.3}$	$273.2\,{\color[rgb]{0.5,0.5,0.5}\scriptstyle\pm\,22.2}^{*}$
NTNS ( $\alpha=0.5$ )	$0.041\,{\color[rgb]{0.5,0.5,0.5}\scriptstyle\pm\,0.024}$	$8.4\,{\color[rgb]{0.5,0.5,0.5}\scriptstyle\pm\,3.5}$	$229.5\,{\color[rgb]{0.5,0.5,0.5}\scriptstyle\pm\,2.2}$
NTNS ( $\alpha=0.3$ )	$\mathbf{0.017}\,{\color[rgb]{0.5,0.5,0.5}\scriptstyle\pm\,0.020}$	$\mathbf{3.97}\,{\color[rgb]{0.5,0.5,0.5}\scriptstyle\pm\,2.03}$	$233.85\,{\color[rgb]{0.5,0.5,0.5}\scriptstyle\pm\,0.83}$

Appendix

Table 6 : Wasserstein-2 distances to MCMC reference and partition function estimate for the LJ-55 seeded runs in fig. 5 , reported as mean

\pm

std across the five seeds at each setting. Lower

W_{2}

is better. The two NTNS rows differ only in the target acceptance

\alpha

( section F.1 );

\alpha=0.3

is our canonical setting. The single-run LJ-55 entries in the main table 1 for both NTNS and iDEM are median-seed runs drawn from this table: NTNS (

r

W_{2}=0.008

E

W_{2}=2.68

) and iDEM (

r

W_{2}=0.135

E

W_{2}=48.4

). The same NTNS run is shown in Fig. 2(b) and fig. 3 . One of the iDEM trainings diverged so is omitted. ∗ Reference value taken from Akhound-Sadegh et al. [13] .

Explore similar work

Quenched Ensemble Sampling
Sep 14, 2026David YallupPhase Transitions Effective Field Theory
Kavli Institute for Cosmology Cambridge Institute of Astronomy, University of Cambridge, Madingley Road, Cambridge, CB3 0HA, UK
Generative Nested Sampling of Atomistic Thermodynamic Landscapes
Sep 2, 2026Alessandro Coretti, Nico Unglert, Sebastian Falkner +2Molecular Dynamics General Relativity
Faculty of Physics, University of Vienna, 1090 Vienna, Austria
MetaDNS: Enhancing Exploration in Discrete Neural Samplers via Well-Tempered Metadynamics
May 20, 2026Xiaochen Du, Juno Nam, Jaemoo Choi +7Accelerated Sampling Molecular Dynamics
Massachusetts Institute of Technology · Georgia Institute of Technology

Sep 14, 2026, stat.MLJ/K move · Enter open · S save

Quenched Ensemble Sampling

David Yallup

Kavli Institute for Cosmology Cambridge Institute of Astronomy, University of Cambridge, Madingley Road, Cambridge, CB3 0HA, UK

Phase Transitions Effective Field Theory

Open paper AskPDF arXiv

Some of the sharpest challenges in sampling from the energy functions of physical systems arise at phase transitions, where the density of states changes abruptly and many sampling algorithms stall. Nested sampling is a particle method that traverses the density of states under a hard energy constraint and is known to be robust to such transitions, but its application in high dimension is limited by the difficulty of sampling under that constraint. In this work we introduce Quenched Ensemble Sampling, which generalises the hard constraint to a family of repulsive potentials at the energy boundary. This preserves the quenched path of monotonically decreasing energy while making the constrained target amenable to scalable gradient-based kernels. We demonstrate on synthetic models of phase transitions that our method estimates the marginal likelihood and draws posterior samples across a first-order transition where popular alternatives such as tempering fail. We apply the procedure to marginal likelihood estimation in Bayesian neural networks, enabling model comparison between network architectures. Finally, in a high-dimensional continuous lattice field theory, we show that this method traverses a first-order transition and estimates the partition function.

Abstract

Explore similar work

Quenched Ensemble Sampling

Generative Nested Sampling of Atomistic Thermodynamic Landscapes

MetaDNS: Enhancing Exploration in Discrete Neural Samplers via Well-Tempered Metadynamics