A trust layer for fast pedestal-stability surrogates: calibrated error bars, out-of-distribution alarms and few-shot recalibration

Amin Vakhshouri

cs.LGNov 4, 2026

A trust layer for fast pedestal-stability surrogates: calibrated error bars, out-of-distribution alarms and few-shot recalibration

Authors: Amin Vakhshouri

Organizations: dataquanta LLC, Delaware, USA

Abstract

Machine-learning surrogates of linear peeling-ballooning stability now run in milliseconds, but they are mostly deployed as point estimators: they do not say how wrong they may be, or when a query lies outside what they were trained on. For a fusion company designing a new device, every equilibrium of interest is outside the training data of existing surrogates, so this is the case that matters. We describe a model-agnostic trust layer for stability surrogates, made of conformal error bars, out-of-distribution alarms that route doubtful cases to the physics code, and recalibration from a few simulations of a new design, and we test it with pre-registered hypotheses. On a large public dataset of JET-like equilibria with ideal-MHD growth rates, the error bars hold their nominal coverage in distribution and fall short under a shift in plasma shape; a few dozen simulations of the new shape, or routing the cases the alarms flag, restore it. On resistive BOUT++ surrogates, the bars stay calibrated on held-out runs, the alarm flags tokamaks absent from training, and a scenario scan is screened in seconds instead of hours. Across equilibria from the SPARC library and the DIII-D coil set, a surrogate trained on our own graded BOUT++ labels beats the baseline on unseen equilibria, while a surrogate trained elsewhere ranks the same cases backwards and its alarm declines them. Building our own neural-operator surrogate, the harness exposed two input defects whose fix closed its gap to a simple network. For every surrogate tested, the trust layer keeps coverage near target, with a bar width set by that surrogate's own error. Not every pre-registered hypothesis held, and we report those results and the limits of all of them.

Figures & tables

Figure 1: Trust layer on KARHU (MLP ensemble, $K=5$ ). (a) Coverage in distribution and under a soft and a hard triangularity shift, with standard and adjusted calibration. (b) Few-shot recalibration on the new shape: mean and worst-5% coverage over 200 draws, and interval half-width. (c) Exploratory: coverage of the cases kept after routing what each alarm flags. (d) Coverage by true growth rate, one interval for all vs. per predicted- $\gamma$ range.

	Result	Verdict
H1	89.9% at 90%, $q=\pm 0.016$ , $R^{2}=0.933$	pass
H2	hard shift: 85.3% (criterion $<85\%$ )	not met
H3	soft shift: 89.0% $\to$ 91.1% weighted	pass (weak)
H4	$m=10/30/100$ : 91.5/90.3/90.4%, $\pm 0.055/0.026/0.024$	knee at 30
H5	AUROC: spread 0.63, distance 0.73	mixed
H6	$\gamma>0.2$ : 78.4% $\to$ 88.4% (Mondrian)	pass

Table 1: Pre-registered outcomes on KARHU (MLP ensemble).

Figure 2: The deployed architecture through the same trust layer on KARHU. (a) Mean absolute error in distribution (S0) and on the shifted shapes (S1 target). (b) Calibrated half-width and in-distribution coverage at a 90% target.

	MLP	v2	v3	v4
$R^{2}$ (S0)	0.9333	0.8802	0.8836	0.9325
MAE (S0)	0.0078	0.0165	0.0159	0.0116
Coverage (S0)	89.9%	90.2%	90.3%	90.3%
Half-width (S0)	0.016	0.037	0.036	0.025
MAE, S1 target	0.0106	0.0176	0.0180	0.0139
Coverage, S1 target	85.3%	88.0%	87.0%	86.1%

Table 2: Same trust-layer harness, four surrogates on KARHU (ensembles of five; S0 in distribution, S1 hard and S2 soft triangularity shift; p5: worst 5% of 200 draws of the 30 recalibration cases).

Figure 3: Onboarding the high-triangularity family (pre-registered; three repeats, bands show the range). (a) Error on the new family’s test set against the number of new simulations, random against guided selection; diamonds mark where each reaches the target halfway between the base model and the model trained on the whole pool. (b) Coverage of the conformal bar calibrated on 30 random simulations. (c) Its half-width.

Figure 4: The deployable surrogate on five real G-EQDSK equilibria. Columns: plasma boundary used; pressure profile against the KARHU training envelope (5–95%), with the model’s radial grid as ticks; global inputs against their training range (dark: 5–95%, light: 0.5–99.5%); verdict and reasons.

Equilibrium	In training	Spread / mean	Alarm
SPARC 1358	yes	1.0 $\times$	–
SPARC 1514	yes	1.1 $\times$	–
DIII-D 145419	yes	0.6 $\times$	–
JET 79692	no	5.9 $\times$	yes
NSTX 132588	no	9.9 $\times$	yes
ITER 15 MA	no	4.8 $\times$	yes

Table 3: Ensemble-spread alarm on real equilibria (BOUT++ surrogate). Threshold: largest spread on 42 in-distribution held-out runs.

Figure 5: Surrogate and trust layer on BOUT++ labels, equilibria held out (pre-registered). Left: held-out predictions with the 90% CV+ interval. Centre: error on held-out SPARC equilibria by model. Right: the KARHU-trained surrogate’s output on the same SPARC cases, outside its range.

Figure 6: Scenario scan on 42 held-out BOUT++ operating points at a growth-rate budget of 0.25. Left: decision map in $(\eta,\nu)$ . Right: prediction $\pm$ calibrated interval, sorted by prediction, with BOUT++ values as dashes.

Regression Uncertainty Estimation

See all 26 papers

AI Surrogate Modeling for Real-Time Tokamak Equilibrium Prediction: Benchmarking Neural Architectures and Validation on EXL-50U
Aug 24, 2026Guoyang Shi, Zitong Zhang, Siqi Ding +5Neural Surrogate Modeling PDE Surrogate Modeling
Beijing ENN Fusion Energy Science and Technology Co., Ltd., Beijing 101111, China · Beijing Key Laboratory of High Magnetic Field Spherical Torus Fusion Energy, Beijing 101111, China · Hebei Key Laboratory of Compact Fusion, Langfang 065000, China +3
Cycle-Consistent and Uncertainty-Aware Neural Surrogates for Tokamak Edge Plasmas
Jul 23, 2026Abdourahmane Diaw, Sebastian De Pascuale, Jae-Sun Park +3Neural Surrogate Modeling PDE Surrogate Modeling
Fusion Energy Division, Oak Ridge National Laboratory, Oak Ridge, TN 37932, USA · DIFFER - Dutch Institute for Fundamental Energy Research, De Zaale 20, 5612 AJ Eindhoven, the Netherlands
Probabilistic and Geometry Aware Neural Surrogate of Scrape Off Layer Plasma Simulations
Sep 23, 2026Gabriele Gianuzzo, Stefan Dasbach, Fleur Hendriks +2Neural Surrogate Modeling Surrogate Modeling
Eindhoven University of Technology · DIFFER

Aug 24, 2026J/K move · Enter open · S save

AI Surrogate Modeling for Real-Time Tokamak Equilibrium Prediction: Benchmarking Neural Architectures and Validation on EXL-50U

Guoyang Shi, Zitong Zhang, Siqi Ding, Jianguo Chen+4

Beijing ENN Fusion Energy Science and Technology Co., Ltd., Beijing 101111, China · Beijing Key Laboratory of High Magnetic Field Spherical Torus Fusion Energy, Beijing 101111, China · Hebei Key Laboratory of Compact Fusion, Langfang 065000, China+3

Neural Surrogate Modeling PDE Surrogate Modeling OOD Generalization

Open paper AskPDF arXiv

Fast and reliable plasma equilibrium prediction is essential for real-time tokamak operation and control, but conventional Grad-Shafranov (GS) solvers are often too costly for real-time deployment. We develop an AI surrogate framework and benchmark five architectures (MLP, CNN, FNO, Transformer, and KAN) on a numerical GS database with 100,000 IID and 10,000 OOD samples. Under a unified protocol, we evaluate accuracy, inference efficiency, model scaling, and robustness. We also establish device-level validation on the EXL-50U tokamak by linking numerical GS solutions, surrogate predictions, and the standard Shape Editor reference to assess simulation-to-device consistency. The surrogates achieve errors of

10^{-3}

10^{-2}

relative to GS solutions, while the GS-to-device discrepancy remains at

10^{-3}

. Transformer gives the best IID accuracy, whereas CNN offers the best balance of accuracy, robustness, and speed, reaching 0.7 ms TensorRT latency. On unseen plasma geometries and parameter regimes, CNN and FNO show the strongest extrapolation stability, with 4%-5% relative

L_2

error, while models with weaker inductive biases degrade more substantially. Scaling data and model capacity improves interpolation but not necessarily extrapolation, revealing a trade-off between capacity and OOD generalization. Overall, this work provides a systematic, device-consistent benchmark for AI-based GS prediction and practical guidance for selecting reliable surrogates for real-time plasma control and fusion applications.

Abstract

Regression Uncertainty Estimation

Explore similar work

AI Surrogate Modeling for Real-Time Tokamak Equilibrium Prediction: Benchmarking Neural Architectures and Validation on EXL-50U

Cycle-Consistent and Uncertainty-Aware Neural Surrogates for Tokamak Edge Plasmas

Probabilistic and Geometry Aware Neural Surrogate of Scrape Off Layer Plasma Simulations