PaperlayerView exact revision on arXiv

Source-linked AI summary

Neural Transport Nested Sampling

David Yallup, Will Handley

arXiv:2609.29413v1cs.LGstat.ML

TL;DR

Molecular Boltzmann sampling calls for ensemble samples and partition functions, but existing neural samplers have limitations in energy-gradient use, accuracy guarantees, and calibrated temperature-resolved estimates. NTNS puts an online-trained flow-matching drift inside a Metropolis-corrected Langevin kernel within nested sampling. On LJ-55 it improves both reported Wasserstein errors over the strongest prior neural baseline at lower wall-clock training cost, while one run estimates temperature-resolved partition functions.

  • Problem

    Molecular sampling requires Boltzmann-ensemble samples and their normalising partition function for downstream free-energy estimation, conformational analysis, and phase-behaviour studies.

    Source: §1 · p. 1

  • Method

    NTNS accelerates nested sampling’s constrained-prior replacement with an online-trained flow-matching drift in a Metropolis-adjusted Langevin kernel using energy evaluations.

    Source: §1 · p. 2

  • Results

    On LJ-55, NTNS reduced both interatomic distance and energy Wasserstein errors by over an order of magnitude versus the strongest prior neural baseline, at substantially lower wall-clock training cost.

    Source: §1 · p. 2

  • Takeaways & Limitations

    A single NTNS run returns a calibrated, temperature-resolved partition function estimate and supports post-hoc sampling across temperatures, recovering the LJ-55 phase diagram without additional simulations.

    Source: §1 · p. 2

    Evaluation cost is currently driven in part by scaling inner-kernel steps with problem dimension, though the authors describe this allocation as tunable.

    Source: §5 · p. 9

Abstract

from arXiv · show

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.

1 Introduction

Molecular sampling must produce Boltzmann-ensemble samples and partition functions for downstream tasks such as free-energy estimation and phase-behaviour analysis. NTNS combines nested sampling with a learned, Metropolis-corrected flow drift, achieving strong Lennard–Jones benchmark accuracy while also estimating temperature-resolved partition functions from one run.

1 Introduction

  • Molecular sampling requires Boltzmann-ensemble samples and their normalising partition function for downstream free-energy estimation, conformational analysis, and phase behaviour.The setting can be sample-free: energy evaluations are available, but target samples and the partition function are not.Source: §1 · p. 1
  • Existing neural diffusion samplers scale to challenging problems, but differentiate the energy, remain tied to learned finite-time dynamics, and do not intrinsically return calibrated temperature-resolved partition functions.Transport-augmented particle methods provide normalising-constant estimates, but their molecular-system efficiency is limited by learned proposal quality.Source: §1 · p. 2
  • On LJ-55, NTNS reduced both interatomic distance and energy Wasserstein errors by over an order of magnitude versus the strongest prior neural baseline, at substantially lower wall-clock training cost.On DW-4 and LJ-13, it was competitive with the best baselines on every metric and best on most.Source: §1 · p. 2
  • A single NTNS run returns a calibrated, temperature-resolved partition function estimate and supports post-hoc sampling across temperatures, recovering the LJ-55 phase diagram without additional simulations.The paper describes this as the first such neural sampling result at scale.Source: §1 · p. 2
  • NTNS uses an online-trained flow-matching drift in a Metropolis-adjusted Langevin kernel to accelerate nested-sampling constrained-prior replacement.The flow is trained by regression on live points, and the kernel uses only energy evaluations.Source: §1 · p. 2

2 Background

This background introduces nested sampling as a temperature-independent route to partition functions, flow matching as a way to learn transport, and prior neural and particle methods that motivate NTNS’s corrected local-mutation design.

2 Background

  • The background frames NTNS at the intersection of particle methods for normalising constants and neural transport for sample generation in a sample-free Boltzmann setting.The section identifies nested sampling, flow matching, and neural sampler baselines as the necessary background.Source: §2 · pp. 2–3

2.1 Nested Sampling

  • Nested sampling estimates the partition function by transforming an integral into one over prior volume, with live points sampling successively constrained priors and a temperature-independent trajectory yielding Z(β).The hard energy constraint makes replacement increasingly expensive as the constrained region develops multiple modes, curved ridges, or thin shells.Source: §2.1 · p. 3

2.2 Flow Matching and Density Estimation

  • Flow matching trains a velocity field by regression between base and target samples; NTNS uses it as a proposal inside an MH-corrected kernel rather than directly sampling from the smooth flow.The figure depicts live points defining the constrained prior and the flow proposal used within the corrected kernel.Source: §2.2 · pp. 3–4, Figure 1 · p. 4

2.3 Related Prior Work

  • Neural diffusion samplers use learned drifts or scores but lack a correction guarantee, while learned-proposal particle methods typically use global proposals whose acceptance degrades with dimension.NTNS instead uses learned dynamics as a guiding force for energy-based Metropolis correction and local mutation.Source: §2.3 · p. 4

3 Method

NTNS keeps the nested-sampling outer loop and replaces its constrained-prior mutation step with a learned, Metropolis-corrected kernel. A flow-matching velocity supplies a local Langevin drift, while finite mutation budgets and adaptive training choices shape practical accuracy and cost.

3 Method

  • NTNS retains nested sampling’s reweight, resample, mutate, and replace sequence, changing only the constrained-prior mutation kernel.The outer loop fixes an energy threshold and records removed live points as dead with temperature-dependent weights.Source: §3 · p. 5

3.1 Neural drift Monte Carlo kernels

  • The flow-matching velocity acts as a local MALA drift, avoiding flow-density evaluation and its backward ODE while retaining an MH correction.The proposal density is Gaussian; computing it uses velocity-network evaluations rather than evaluating the flow density.Source: §3 · p. 5, §3.1 · pp. 5–6

3.1.1 Learned kernels within nested sampling

  • The MH-corrected mutation kernel preserves the constrained prior regardless of flow quality, but finite inner moves can still leave finite-mixing error.The number of inner steps T is treated as a numerical-accuracy parameter and assessed by varying it.Source: §3.1.1 · p. 6

3.2 Training and adaptation

  • NTNS trains each flow on surviving live points, warm-starts from the previous iteration, adapts the step size, and restarts training if acceptance collapses.The live set serves as the sample buffer, and flow training dominates runtime in the described implementation.Source: §3.2 · pp. 6–7

4 Experiments

Across molecular benchmarks, NTNS performs strongly against neural samplers, with its clearest accuracy and convergence results on LJ-55. A single LJ-55 run also supports temperature reweighting, phase-structure recovery, and validation against an intermediate-temperature reference.

4 Experiments

  • Benchmarks cover DW-4, LJ-13, and LJ-55 against neural samplers and a nested-slice-sampling control, using distance metrics computed against long-run parallel MCMC chains.The targets range from an 8D double well to 39D and 165D Lennard–Jones clusters.Source: §4 · p. 7
  • On LJ-55, NTNS reduces both r-W2 and E-W2 by over an order of magnitude, converges consistently in about 5.5 hours, and closely matches MCMC distance and energy summaries.The comparison table reports Wasserstein-2 distances to MCMC references; Figure 2 displays the distance and energy distributions.Source: §4 · p. 7, Table 1 · p. 7, Figure 2 · p. 8

4.1 Partition function calculation

  • A single LJ-55 run reweights samples across β to expose phase structure; at β = 0.6, NUTS validates recovery, and thermodynamic integration gives 148.8 ± 0.4 versus NTNS’s 147.5 ± 0.4.The run produced 486,400 weighted samples, with effective sample sizes of 54,705 at β = 0.6 and 59,282 at β = 1.0.Source: §4.1 · pp. 7–8

5 Discussion

NTNS combines a local learned drift with population-based exploration, and a matched control attributes its gain over classical nested sampling to the mutation kernel. The discussion also identifies energy-evaluation trade-offs and computational choices that may constrain efficiency or merit further study.

5 Discussion

  • Replacing global flow proposals with a local learned drift while retaining population exploration places NTNS between corrected MCMC and global search; a matched control attributes its gain over classical nested sampling to the kernel.The control holds the outer schedule fixed while replacing the learned mutation kernel.Source: §5 · p. 8
  • NTNS avoids force evaluations and target clipping or smoothing, but adjoint methods use fewer energy evaluations, a relevant trade-off for expensive neural-network energy functions.The current benchmarks use cheap target energies, making the evaluation-count difference negligible in practice.Source: §5 · pp. 8–9
  • Current computational cost depends on flow-training allocation, a basic warm-start heuristic, and inner-kernel steps that scale with dimension; the authors describe these choices as potentially conservative.They identify particle-history use, online inner-step tuning, and more direct reuse of repeated steps as future directions.Source: §5 · p. 9
  • The authors propose extending the approach to an interacting SMC system, amortizing transports across systems, and using samples to estimate free-energy differences.The terminating flow can also be simulated as a neural sampler.Source: §5 · p. 9

6 Conclusion

NTNS accelerates nested sampling’s constrained-prior replacement step with a learned drift in a MALA kernel, and the paper reports improved accuracy at scale with economical energy evaluations and runtime.

6 Conclusion

  • NTNS combines nested sampling with a learned drift in a MALA kernel and scales to 55 interacting particles with improved accuracy and economical energy evaluations and runtime.The conclusion identifies larger particle systems and more complex, costly energy functions as open challenges.Source: §6 · pp. 9–10

Code and Data Availability

  • The algorithm code and example scripts for reproducing the experiments are available on GitHub.The paper gives the repository as https://github.com/yallup/ntns.Source: Code and Data Availability · p. 10

AI and LLM tool use

  • The authors report using GPT-5.4 to refine portions of the draft and Claude Opus 4.7 to refine the code, while taking responsibility for the final content.Source: AI and LLM tool use · p. 10

Broader Impact

  • The authors envisage downstream impact in drug discovery and novel-materials research, while noting that their work uses synthetic energy functions.Source: Broader Impact · p. 10

A Target Distributions

  • The particle targets use a standard Gaussian prior restricted to the zero-centre-of-mass subspace, with translation symmetry removed and O(d) × S_N invariance.Energy depends only on pairwise distances; the stated ambient dimensions refer to R^Nd.Source: A Target Distributions · p. 16

A.1 Double-Well Potential (DW-4)

  • DW-4 has four particles in R^2, giving eight dimensions, and its target has a bimodal interatomic-distance distribution.Its pairwise energy uses a standard Gaussian prior on the zero-centre-of-mass subspace.Source: §A.1 · p. 16

A.2 Lennard-Jones Cluster (LJ-13)

  • LJ-13 is a 13-particle Lennard–Jones system in R^3 with 39 dimensions and reduced units ε = σ = 1.Its prior is a standard Gaussian restricted to the zero-centre-of-mass subspace.Source: §A.2 · p. 16

A.3 Lennard-Jones Cluster (LJ-55)

  • LJ-55 extends the Lennard–Jones system to 55 particles in R^3 and presents a solid–liquid phase transition benchmark.The system has 165 dimensions and uses a standard Gaussian prior restricted to the zero-centre-of-mass subspace.Source: §A.3 · p. 16

B Reference MCMC Data

The paper evaluates generated samples against independent reference configurations from long MCMC runs at each target temperature.

B Reference MCMC Data

  • Reference MCMC samples from Klein et al. [6], redistributed in the iDEM codebase [13], provide diagnostic comparisons.The references are long runs at the target temperature and validate posterior samples using pairwise-distance and energy histograms.Source: B Reference MCMC Data · p. 16
  • DW-4 and LJ-13 each have 100,000 training, 10,000 validation, and 10,000 test configurations; LJ-55 has 10,000 in each split.Configuration shapes are (4, 2), (13, 3), and (55, 3), respectively.Source: B Reference MCMC Data · p. 16
  • The reference configurations serve only as independent ground truth for comparison; NTNS does not use them for training.Source: B Reference MCMC Data · p. 16

C Metrics

This material defines the sample-comparison metrics and documents the NTNS implementation, including its learned-drift kernel and the correctness argument for constrained-prior sampling.

C Metrics

  • The reported 2-Wasserstein summaries are r-W2 for pooled pairwise distances and E-W2 for target energies, rather than raw-coordinate W2.Raw-coordinate W2 is omitted because particle-permutation symmetry can make its value dominated by solver inefficiency and finite-sample noise; variance estimates use ten disjoint reference batches.Source: C Metrics · p. 17

D.1 EGNN Architecture

  • The particle velocity network is an E(n)-equivariant graph neural network whose stacked layers update particle positions and features through edge messages.Its output is an accumulated, mean-subtracted displacement, producing an E(n)-equivariant zero-mean velocity field.Source: §D.1 · p. 17

D.2 Hyperparameters

  • Table 2 lists NTNS hyperparameters held fixed across benchmarks except where the section F ablations specify otherwise.Source: §D.2 · p. 17, Table 2 · p. 18

D.3 Algorithm Pseudocode

  • Algorithm 1 performs the nested-sampling outer iteration, training the flow on surviving live points before replacing dead points with mutated samples.The pseudocode describes the drift MALA kernel as the replacement step’s mutation mechanism.Source: §D.3 · pp. 17–18
  • The kernel adapts its Langevin step size between outer iterations using pooled empirical acceptance, and uses T = d Langevin steps per replacement.For particle systems, projected noise and a mean-subtracted drift keep the chain on the zero-centre-of-mass subspace.Source: §D.3 · pp. 18–19

E Proof of Correctness

  • The Metropolis–Hastings correction preserves the constrained prior as the stationary distribution, regardless of the learned drift, under the stated mild conditions.Proposals outside the energy constraint are rejected; the proof relies on a valid proposal density, not on the flow parameters.Source: E Proof of Correctness · pp. 19–20

F Ablations

On LJ-55, the matched control attributes NTNS’s lower distance and energy errors to its learned mutation dynamics rather than the nested-sampling outer loop. Ablations show how acceptance targets and computational settings affect fidelity, runtime, and energy evaluations.

F Ablations

  • Across four LJ-13 ablation axes, sample quality remains robust while runtime and evaluation costs respond predictably to their relevant settings.Halving T to 19 reduces target evaluations to approximately 5×10^6, with r-W2 of 0.031 versus the baseline 0.026; increasing m to 4096 alone raises runtime to 2071 s.Source: F Ablations · p. 20, Table 3 · p. 20
  • Flow epochs are the dominant runtime handle, while inner-kernel steps T dominate target-evaluation count.Changing flow epochs from 150 to 75 reduces runtime to 464 s, while 300 epochs increases it to 859 s; doubling both m and k to (4096, 2048) recoups most of the larger-live-set runtime cost at 851 s.Source: F Ablations · pp. 20–21
  • Increasing m to 4096 while holding k = 1024 slows compression and yields a more accurate log Z, while doubling both m and k preserves the baseline compression rate.The authors note that both larger configurations increase per-iteration target evaluations because replacement-chain count scales with k.Source: F Ablations · p. 20

F.1 Inner-kernel target acceptance

  • At α = 0.3, LJ-55 fidelity improves on both r- and E-W2, and cross-seed scatter in log Z(1) is about 3× tighter than at α = 0.5.The TI cross-check moves from 3.4σ to 1.5σ of the NUTS reference; the authors attribute the departure from the canonical MALA target to the constrained prior’s hard edge.Source: §F.1 · p. 21
  • Across the tested acceptance targets, α = 0.3 adapts to a larger mean step size than α = 0.5 and is adopted for LJ-55 experiments.Mean ϵ is approximately 2.1×10^-4 at α = 0.3 versus 9.7×10^-5 at α = 0.5.Source: §F.1 · p. 21

F.2 Classical nested sampling control

  • NTNS achieves substantially lower distance and energy errors than matched NSS, which uses approximately six times more target evaluations.NSS is faster in wall time because the analytic Lennard–Jones energy is exceptionally cheap, while NTNS runtime is dominated by flow training.Source: Table 4 · p. 21, §F.2 · p. 21

G Convergence and Reproducibility

NTNS uses a partition-function-based stopping rule, while the diffusion baselines run to fixed time limits. Its LJ-55 results are consistent across seeds, and the reported cost comparison distinguishes energy-only evaluations from gradient-based evaluations.

G Convergence and Reproducibility

  • NTNS stops when the remaining live-set contribution to the partition function is negligible, whereas iDEM, AS, and ASBS use fixed wall-clock budgets and timed out on LJ-55.The typical stopping threshold is log Z_live − log Z < −3, which targets β = 1 and can be tuned for other temperatures or runtime budgets.Source: G Convergence and Reproducibility · p. 21
  • Across the five-seed sweep, NTNS is consistently the best-performing sampler, while continued diffusion-baseline training does not yield a stable answer.Figure 5 shows NTNS reproducing the MCMC reference across seeds, unlike iDEM.Source: G Convergence and Reproducibility · p. 22
  • At α = 0.5, acceptance can collapse to zero and trigger an inner-kernel restart, whereas this did not occur across five seeds at α = 0.3.At the canonical α = 0.3, occasional non-finite flow-training gradients are filtered online without an across-seed sample-quality difference.Source: G Convergence and Reproducibility · p. 22
  • NTNS evaluates energy alone, while iDEM, AS, and ASBS differentiate each evaluation; AS and ASBS use fewer energy queries by design, and NTNS has the lowest reported wall time.Entries marked with an asterisk reached the 12-hour training limit before convergence.Source: Table 5 · p. 22
  • NTNS’s median-seed LJ-55 run has r-W2 = 0.008 and E-W2 = 2.68, versus iDEM’s r-W2 = 0.135 and E-W2 = 48.4.The table reports five-seed results for the two NTNS acceptance settings and iDEM; one iDEM training that diverged is omitted.Source: Table 6 · p. 22

H Phase-Diagram Analysis

An independent thermodynamic-integration check agrees with NTNS’s LJ-55 partition-function change across β. Heat-capacity peaks along the Gaussian-prior tempering path are interpreted as a softened cluster-formation transition, with the LJ-55 peak below the converted literature melting point.

H Phase-Diagram Analysis

  • Thermodynamic integration gives log Z(1.0) − log Z(0.6) = 148.8 ± 0.4, compared with NTNS’s 147.54 ± 0.44.Across five NTNS seeds, the mean difference is 147.75 ± 0.73, and log Z(1.0) is 233.85 ± 0.83.Source: H Phase-Diagram Analysis · p. 22
  • The per-atom heat-capacity peaks are 2.692 ± 0.089 at βc = 1.169 ± 0.053 for LJ-13 and 2.715 ± 0.038 at βc = 0.875 ± 0.032 for LJ-55.The curves are reweighted from a single NTNS dead-point trajectory under the Gaussian-prior Boltzmann convention, with variability across simulated prior-volume sequences.Source: H Phase-Diagram Analysis · pp. 22–23
  • The LJ-55 peak at βc = 0.875 is below the converted literature melting point near β ≈ 1.7, consistent with a softened transition along the Gaussian-prior tempering path.The energy convention differs by a factor of two from the standard literature convention, and Gaussian rather than hard-sphere confinement broadens and shifts the transition.Source: H Phase-Diagram Analysis · p. 23

I IMH Results and Scaling Analysis

The IMH flow-density proposal gives MCMC-comparable LJ-13 sample quality but scales poorly to LJ-55 because backward ODE density evaluation is costly. The supplementary discussion contrasts nested sampling’s energy-ordered path with annealing and shows additional sampler diagnostics.

I IMH Results and Scaling Analysis

  • IMH proposes by integrating the learned flow forward and applies an importance ratio plus the energy-constraint indicator in its Metropolis–Hastings acceptance rule.The kernel preserves the constrained prior for any flow Qϕ, so flow quality affects mixing efficiency; unlike MALA, IMH also integrates backward to evaluate the flow density.Source: I IMH Results and Scaling Analysis · pp. 23–25
  • At α = 0.5, LJ-55 MALA acceptance occasionally reaches zero and recovers after a restart, while at α = 0.3 it remains above zero without a restart.The acceptance panels also show per-iteration flow-training loss and adapted Langevin step size, with a dotted line marking the target α.Source: Figure 6 · p. 24
  • The LJ-13 and LJ-55 heat-capacity curves peak near β ≈ 1, which the authors interpret as cluster formation along the likelihood-tempering path.The plotted traces also show variability across simulated prior-volume sequences.Source: Figure 7 · p. 24
  • On LJ-13, IMH gives r-W2 = 0.011 ± 0.006 and E-W2 = 0.77 ± 0.19, comparable to MALA and within statistical error of the MCMC-vs-MCMC control.IMH maintains 15–37% acceptance over 51 nested-sampling iterations and completes in 4.4 hours at 3.3 points/second.Source: I IMH Results and Scaling Analysis · p. 25
  • On LJ-55 at 165 dimensions, IMH acceptance falls below 1% with prohibitive runtime, making density-corrected IMH practical only at moderate dimensions (d ≲ 50).Backward ODE density evaluation requires O(d) vector-Jacobian products per integration step; MALA avoids this density calculation by using flow velocity as a drift.Source: I IMH Results and Scaling Analysis · p. 25
  • For LJ-13, IMH has sample quality comparable to MALA but substantially higher runtime.The diagnostics compare pooled interatomic-distance and energy histograms against MCMC and show acceptance rates across nested-sampling iterations.Source: Figure 8 · p. 25

J Nested Sampling vs SMC Paths

  • Nested sampling follows an energy-sorted sequence of constrained distributions; the authors describe this path as advantageous near phase transitions compared with annealing.The passage notes annealing can face weight collapse near a critical temperature or slow down when its step size is adapted.Source: J Nested Sampling vs SMC Paths · pp. 25–26

K Supplementary Figures

  • Figure 9 compares NTNS, iDEM, AS, and ASBS with MCMC for the DW-4 interatomic-distance and energy distributions.The caption identifies the distributions shown but states no outcome of the comparison.Source: Figure 9 · p. 26

Related papers

Loading 2609.29413v1…