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

Liyun Zhang; Naoya Mamada; Kentaro Sakai; Takeo Hoshi; Toru Aonishi

stat.COAug 19, 2026

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

Authors: Liyun Zhang, Naoya Mamada, Kentaro Sakai, Takeo Hoshi, Toru Aonishi

Organizations: GSFS, The University of Tokyo Japan

Abstract

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.

Figures & tables

Figure 1. Comparative overview of MCMC inference in extreme-noise plasma diagnostics. Top: The forward model

F(\theta)

maps latent VDFs to noisy spectra. Markers: green star (ground truth), yellow circle (prediction), and blue ellipse (posterior uncertainty). A: Physical MCMC yields theoretically exact posteriors but is computationally prohibitive and highly noise-sensitive (slender ellipse). B: Standard NN Surrogate (homoscedastic MSE loss) is fast but yields structurally biased and overconfident posteriors (circle far from star). C: Proposed Heteroscedastic Surrogate (GNLL loss) predicts data-dependent log-variance. This learned uncertainty acts as a topological smoother, penalizing overconfident proposals to enable fast and stable MCMC inference (ellipse’s appropriate shape, nearly aligned circle and star).

Figure 2. Workflow of the proposed heteroscedastic neural surrogate and its MCMC integration. (1) Neural Surrogate: A backbone processes abstract physical parameters $\theta$ (latent VDF) via a dual-head network, predicting expected spectrum $\mu_{\mathrm{GNLL}}(\theta)$ and channel-wise log-variance $s_{\mathrm{GNLL}}(\theta)$ . (2) GNLL Optimization: The network is trained via Gaussian Negative Log-Likelihood (GNLL) loss, where learned aleatoric uncertainty attenuates residual penalties in high-noise regions. (3) MCMC Inference: The frozen surrogate evaluates the heteroscedastic log-likelihood. The learned variance regularizes the Metropolis-Hastings acceptance ratio, yielding a robust, physically consistent posterior.

Figure 3. Qualitative comparison of 2D Posterior Probability Density Functions (a: best seeds with minimal absolute $T_{e}$ error, marked $\boldsymbol{\dagger}$ in Table 1 ) and parameter trace plots across 20 independent seeds (b, c). Method A: yields precise ‘best-case’ estimates, though 95% credible intervals remain widely dispersed due to extreme noise sensitivity. Method B: collapses into a biased and overconfident geometry (almost vanishing error bars) by overfitting localized shot noise. Method C: Our proposed GNLL Surrogate acts as a topological smoother, bypassing noise-induced local minima. Its 95% credible intervals consistently encompass ground-truth values, ensuring the estimated uncertainty aligns with actual estimation errors under extreme noise.

Test ID	Method A (Physical MCMC)		Method B (Standard MSE)		Method C (Proposed GNLL)
Test ID	$T_{e}$ (eV)	$n_{e}(\times 10^{18}\text{ m}^{-3})$	$T_{e}$ (eV)	$n_{e}(\times 10^{18}\text{ m}^{-3})$	$T_{e}$ (eV)	$n_{e}(\times 10^{18}\text{ m}^{-3})$
Seed 01	$6000.18^{\boldsymbol{\dagger}}$	$2.008^{\boldsymbol{\dagger}}$	$5975.22$	$2.004$	$5948.87$	$1.998$
Seed 02	$5895.83$	$1.979$	$6002.32^{\boldsymbol{\dagger}}$	$1.994^{\boldsymbol{\dagger}}$	$6017.00$	$1.997$
Seed 03	$5984.37$	$2.007$	$6051.74$	$2.019$	$6139.09$	$2.032$
Seed 04	$5894.42$	$1.987$	$6002.89$	$2.001$	$6060.64$	$2.013$
Seed 05	$6015.62$	$2.003$	$6022.20$	$2.002$	$6020.44$	$2.002$

Table 1. Comprehensive MCMC inference results across 20 independent extreme-noise blind tests (Target Truth:

T_{e}=6000\text{ eV},n_{e}=2.000\times 10^{18}\text{ m}^{-3}

). The summary metrics—Ensemble Std, RMSE to Truth, and Inference Time per posterior PDF—quantitatively demonstrate that the proposed GNLL surrogate (Method C) achieves

\sim 1500\times

acceleration over the physical MCMC (Method A) while surpassing the standard MSE (Method B) in both accuracy (RMSE) and stability (Std). Values marked with a dagger (

\boldsymbol{\dagger}

) denote seeds visualized in Figure 3 (a), chosen for minimum absolute

T_{e}

error.

Variant	Variance Modeling	MCMC Likelihood	Std ( $T_{e}$ )	Std ( $n_{e}$ )	RMSE ( $T_{e}$ )	RMSE ( $n_{e}$ )
Standard MSE	None	Uniform	$103.43$	$0.0179$	$105.19$	$0.0179$
Global GNLL	Global Scalar	Uniform	$90.51$	$0.0145$	$88.29$	$0.0146$
Decoupled GNLL	Channel Vector	Uniform	$87.48$	$0.0143$	$85.84$	$\mathbf{0.0144}$
Proposed GNLL	Channel Vector	Data-Dependent	$\mathbf{78.88}$	$\mathbf{0.0141}$	$\mathbf{83.76}$	$\mathbf{0.0144}$

Table 2. Ablation study decoupling the architectural and inferential components of our framework. We evaluate the necessity of channel-wise variance modeling and its explicit integration into the MCMC likelihood evaluation. Metrics isolate the inference precision (unbiased Std and RMSE), with units being eV for

T_{e}

and

\times 10^{18}\text{ m}^{-3}

for

n_{e}

Explore similar work

Accelerating Bayesian inverse design in computational fluid dynamics using neural operators
May 25, 2026Bipin Tiwari, Omer SanNeural Surrogate Modeling Bayesian Inverse Problems
Department of Mechanical and Aerospace Engineering, University of Tennessee, Knoxville,2026 TN 37996, USA
Fast Bayesian equipment condition monitoring via simulation based inference: applications to heat exchanger health
Apr 22, 2026Peter Collett, Alexander Johannes Stasik, Simone Casolo +1Simulation-Based Inference Efficient Inference
Cognite AS, Oslo, Norway · Department of Data Science, Norwegian University of Life Sciences, Ås, Norway · Department for Mathematics and Cybernetics, SINTEF Digital, Oslo, Norway
Continually learning neural-operator surrogate for three-dimensional airborne electromagnetic Bayesian inversion
Aug 26, 2026Jaehong Chung, Andrew Lockwood, Jef CaersPDE Surrogate Modeling Bayesian Inverse Problems
Mineral-X, Department of Earth and Planetary Sciences, Stanford University, USA · Xcalibur Smart Mapping, Perth, Australia

May 25, 2026J/K move · Enter open · S save

Accelerating Bayesian inverse design in computational fluid dynamics using neural operators

Bipin Tiwari, Omer San

Department of Mechanical and Aerospace Engineering, University of Tennessee, Knoxville,2026 TN 37996, USA

Neural Surrogate Modeling Bayesian Inverse Problems Posterior Sampling Inverse Design Neural Operators Computational Fluid Dynamics

Open paper AskPDF arXiv

Bayesian inverse design provides a principled framework for inferring aerodynamic geometries from sparse flow observations while quantifying uncertainty. However, its practical use in computational fluid dynamics (CFD) is severely limited by the cost of repeated high-fidelity simulations required for gradient-based Markov chain Monte Carlo (MCMC) sampling. While surrogate models are commonly proposed to reduce this cost, their effect on posterior geometry and uncertainty, especially for shock-dominated flows, remains poorly understood. In this work, we demonstrate that neural operator surrogates can be embedded directly within the MCMC inference loop while preserving posterior structure. Using a fully Bayesian inverse formulation of quasi-one-dimensional nozzle flow, we demonstrate that geometry parameterization plays a decisive role in identifiability and posterior conditioning, with cubic B-splines yielding stable and physically meaningful uncertainty estimates. Building on this formulation, a Deep Operator Network trained on CFD-generated data is substituted for the CFD solver within a No-U-Turn Sampler, while keeping the likelihood model, priors, and sampling configuration unchanged. Across sparse to fully observed regimes, surrogate-based inference reproduces the posterior geometry and uncertainty trends of the CFD reference. As a result of surrogate integration, total inference time is reduced to under one second, corresponding to a speedup exceeding three orders of magnitude. In addition, a direct inverse neural operator is examined as a deterministic alternative for inverse design, enabling single-shot geometry reconstruction without posterior sampling. These results demonstrate that neural operator-accelerated Bayesian inference enables practical, uncertainty-aware inverse design workflows for aerodynamic applications.

Abstract

Explore similar work

Accelerating Bayesian inverse design in computational fluid dynamics using neural operators

Fast Bayesian equipment condition monitoring via simulation based inference: applications to heat exchanger health

Continually learning neural-operator surrogate for three-dimensional airborne electromagnetic Bayesian inversion