Neural Surrogate HMC: On Using Neural Likelihoods for Hamiltonian Monte Carlo in Simulation-Based Inference

TL;DR

The paper introduces Neural Surrogate HMC, combining neural likelihood estimation with Hamiltonian Monte Carlo to accelerate and stabilize Bayesian inference when simulators are expensive or unstable. It provides a formal three-step pipeline, practical guidelines for training data and priors, and empirical convergence checks. The method is demonstrated on the Parker equation for galactic cosmic ray transport, achieving state-of-the-art posterior constraints while delivering substantial computational speedups. These contributions offer a scalable approach for inference in complex physical simulators and other SBI contexts requiring repeated, differentiable likelihood evaluations.

Abstract

Bayesian inference methods such as Markov Chain Monte Carlo (MCMC) typically require repeated computations of the likelihood function, but in some scenarios this is infeasible and alternative methods are needed. Simulation-based inference (SBI) methods address this problem by using machine learning to amortize computations. In this work, we highlight a particular synergy between the SBI method of neural likelihood estimation and the classic MCMC method of Hamiltonian Monte Carlo. We show that approximating the likelihood function with a neural network model can provide three distinct advantages: (1) amortizing the computations for MCMC; (2) providing gradients for Hamiltonian Monte Carlo, and (3) smoothing over noisy simulations resulting from numerical instabilities. We provide practical guidelines for defining a prior, sampling a training set, and evaluating convergence. The method is demonstrated in an application modeling the heliospheric transport of galactic cosmic rays, where it enables efficient inference of latent parameters in the Parker equation.

Figures (14)

• Figure 1: Our three-step approach to sampling from the posterior $p(\theta|x)$: (1) produce a dataset $(\theta_i, x_i)_{i=1 \ldots M}$ by drawing parameters from the prior $p(\theta)$ and running simulator $\mathcal{S}$ forward. (2) Train an NLE $q(x | \theta ; \phi)$ on the simulated dataset to approximate the simulator. (3) Sample from the posterior $p(\theta_i | x_i)$ for new observations $x_i$ with HMC, using the NLE for all likelihood evaluations. Adapted from deistler2025simulation.
• Figure 2: The NLE (left) takes as input eight latent parameters of the heliosphere and predicts the GCR flux near Earth for 32 rigidity steps. These steps are uniformly distributed in logspace between 0.2 and 200 GV. The NLE is trained with targets from a simulator $S$ that suffers from instabilities (the jagged portion of simulation curve).
• Figure 3: Mean absolute error (left) and mean squared error (right) with numerical model solutions on the test dataset, shown for five models trained on different bootstrap-sampled datasets.
• Figure 4: Posterior samples for AMS-02 interval 24 January to 19 February 2015 across the five sampled parameters. Plotted histograms show the sample distribution from five data replicates of full train size (NLE trained on different bootstrap-sampled training datasets). The Gelman-Rubin $\hat{R}$ statistic for each parameter is shown. Histograms are scaled to the same height across parameters.
• Figure 5: Gelman-Rubin $\hat{R}$ statistic for AMS-02 interval 24 January to 19 February 2015 plotted for the five sampled parameters against training set size. Lines show the statistic across five HMC replicates (blue), model replicates (orange), and data replicates (green). The dashed red line marks the convergence threshold of $\hat{R} = 1.1$.