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; Yapeng Zhang; Jiayi Zhi; Hanyue Zhao; Tianyuan Liu

Source-linked AI summary

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, Yapeng Zhang, Jiayi Zhi, Hanyue Zhao, Tianyuan Liu

arXiv:2608.23217v1Aug 24, 2026physics.plasm-phcs.AI

TL;DR

Real-time tokamak control requires fast equilibrium prediction, but conventional GS computation can be too costly. The paper benchmarks five neural surrogate architectures under unified IID/OOD, efficiency, scaling, and robustness evaluations, and validates them against EXL-50U references. Transformer is most accurate on IID data, while CNN offers the best overall balance and CNN/FNO are most robust OOD, with scaling exposing a capacity–generalization trade-off.

Problem
Real-time tokamak control requires millisecond or sub-millisecond equilibrium computation, while prior GS learning studies provide limited systematic IID/OOD, cross-architecture, and device-relevant evidence.
Method
The study generates common GS reference solutions and benchmarks MLP, CNN, FNO, Transformer, and KAN within a unified surrogate framework, including scaling, efficiency, and EXL-50U validation.
Results
Transformer achieves the highest IID accuracy, CNN reaches 0.7 ms TensorRT latency with a favorable accuracy–robustness balance, and CNN/FNO show the strongest OOD robustness.
Takeaways & Limitations
The benchmark provides practical guidance for selecting AI surrogates for real-time tokamak equilibrium prediction and plasma-control applications.
Takeaways & Limitations
Existing evidence remains comparatively limited for complex nonlinear profiles, IID-to-OOD generalization, and explicit device-relevant validation; the common targets also derive from numerical GS solutions.

Abstract

from arXiv · show

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.

1. Introduction

Tokamak equilibrium prediction must balance numerical accuracy with millisecond-scale computational demands. Although learning-based GS approaches are promising, systematic cross-architecture, IID/OOD, and device-relevant evaluations remain limited.

Tokamak operation depends on accurate MHD-equilibrium characterization and control, which supports stability assessment, transport modeling, and related analyses.
Millisecond or sub-millisecond equilibrium computation is required for real-time control, while repeated GS solves can constrain integrated discharge modeling and scenario optimization.
Data-driven surrogates learn solver-generated input–output mappings to provide rapid inference for computationally intensive scientific tasks, including plasma equilibrium reconstruction.
Prior GS studies leave complex nonlinear profiles, IID-to-OOD generalization, and fair unified benchmarking insufficiently characterized.
The paper compares MLP, CNN, FNO, Transformer, and KAN for mapping plasma parameters and geometric descriptors to GS equilibrium solutions.
The proposed framework establishes architecture benchmarking, simulation-to-device validation, and analysis of scaling and generalization characteristics.

2. Tokamak Equilibrium Prediction Problem

The study formulates tokamak equilibrium prediction through the Grad–Shafranov equation, parameterizing plasma profiles and fixed boundaries before generating reference solutions and evaluation datasets.

Problem formulation: The GS equation describes axisymmetric tokamak equilibria through the poloidal flux function ψ(R,Z), whose contours define magnetic flux surfaces.The formulation follows from ideal MHD force balance and assumes independence from the toroidal angle.
Profile and boundary parameterization: Normalized flux sets ψ̂=0 at the magnetic axis and ψ̂=1 at the plasma boundary, enabling equilibria with different currents and dimensions to share a numerical scale.Pressure-gradient and poloidal-current source terms jointly determine the toroidal current-density distribution and internal flux-surface structure.
Profile and boundary parameterization: The pressure and poloidal-current profiles are controlled by np, mp, nf, mf, λ, and β0, with β0 typically set to 0.8.The indices shape the pressure and current profiles, while λ is associated with the total plasma current Ip.
Profile and boundary parameterization: The fixed plasma boundary is parameterized by major radius R0, minor radius a, elongation κ, and triangularity δ, with ψ=0 imposed on the boundary.These geometric parameters define the boundary used in the fixed-boundary equilibrium problem.
Reference solutions: The finite-difference GS solver uses successive over-relaxation on a 101×161 rectangular grid and requires 10.2 s per case on the MATLAB CPU baseline.Interior flux values are updated from neighboring grid points and nonlinear source terms, while boundary nodes remain fixed at ψ=0.
Datasets: The evaluation includes 100,000 IID samples, 10,000 OOD samples with shifted parameter ranges, and an EXL-50U validation dataset comparing predictions with GS and Shape Editor references.IID testing measures interpolation within training ranges, whereas OOD testing measures extrapolation beyond them.

3. AI Surrogate Modeling Framework

The framework casts equilibrium prediction as a common parameter-to-field mapping, allowing different AI surrogate architectures to be compared under the same task definition.

Problem formulation: The objective is to predict two-dimensional GS poloidal-flux fields from varying plasma parameters and boundary geometries.The input includes a parameter vector and a computational grid in the poloidal plane.
Surrogate mapping: An AI surrogate maps the input parameters and grid to a predicted poloidal-flux field, which is compared with the high-fidelity numerical GS output.The formulation uses trainable surrogate parameters while keeping the prediction target fixed across architectures.
Unified comparison: The common parameter-to-field formulation supports fair comparison even when candidate models process inputs differently or are not all neural operators.The shared task separates architecture differences from differences in the prediction objective.

(1) High-fidelity dataset generation

High-fidelity supervised targets are generated by sampling equilibrium parameters and solving the GS equation numerically for each sampled input.

Dataset generation: For each sampled parameter vector fi, the numerical GS solver computes a corresponding reference poloidal-flux field ui.The resulting input–field pairs form the offline dataset used for surrogate modeling.
Dataset generation: The dataset therefore represents a mapping from sampled equilibrium parameters to solver-computed two-dimensional magnetic-flux fields.This construction supplies supervised examples for training, validation, and testing.
Common targets: The numerical solutions provide common reference targets for all five surrogate architectures.Using identical targets ensures that performance differences primarily reflect architecture and training behavior rather than data generation.

(2) AI surrogate training

Each surrogate is trained to approximate the same parameter-and-coordinate-to-flux mapping under a unified protocol, with data misfit forming the core loss component.

Training objective: Each candidate architecture learns the mapping from input parameters and spatial coordinates to the corresponding poloidal-flux field.The architectures may process inputs differently, but they share the same prediction target.
Unified protocol: The models use common training, validation, and testing partitions and are optimized under a unified protocol.This design supports controlled comparison across the candidate architectures.
Training objective: The total loss includes a data-misfit term comparing predicted and true fields.The supplied formulation identifies the data-misfit component as the core part of the training objective.

(3) Model validation

The study validates five neural surrogate architectures under a unified protocol for predicting the two-dimensional poloidal-flux field from physical, geometric, and spatial inputs. The architectures differ in their treatment of spatial structure, spectral behavior, global interactions, and nonlinear parameter dependence.

CNN: The CNN uses a U-Net-style encoder–decoder with multiscale convolutions, upsampling, and skip connections to reconstruct the flux field on the original grid.The design captures local gradients and boundary features while expanding the effective receptive field for global equilibrium structure.
FNO: The FNO represents equilibrium mappings with Fourier-space global convolution, frequency-mode transformation, and inverse transformation back to the spatial domain.Its discretization-invariant formulation supports inference across mesh resolutions without retraining.
KAN: KAN replaces fixed MLP activations with learnable univariate functions, hierarchically composing them for expressive and interpretable nonlinear mappings.The univariate functions are typically constructed from B-spline bases with trainable coefficients.
Architecture benchmark: Five architectures—MLP, CNN, FNO, Transformer, and KAN—predict the same poloidal-flux field from shared physical, geometric, and spatial inputs.Their identical evaluation enables comparison of accuracy, inference efficiency, model capacity, and OOD robustness.
Transformer: The Transformer formulates GS prediction as sequence-to-field mapping, using jointly encoded geometry and coordinates with self-attention for long-range interactions.The five models therefore provide distinct inductive strategies for spatial, spectral, global, and parameter-dependent structure.

4. Results and Analysis

The results show rapid convergence and strong IID prediction across the surrogate architectures, with Transformer achieving the highest accuracy. Error distributions and representative contour comparisons further support the quality of the predictions.

Training convergence: All architectures reduce validation loss rapidly within 50 epochs and converge after approximately 150 epochs at final error levels near 10^-7.Transformer, KAN, FNO, and CNN converge smoothly, while MLP shows mild late-epoch oscillations.
IID accuracy: 0.201% and 0.242% are the Transformer’s test L1 and L2 errors, respectively, the best accuracy among the compared models.KAN, FNO, and CNN follow closely, while MLP has relatively larger prediction errors.
IID accuracy: Below 0.02 percentage points separates training and test errors for every model, indicating limited generalization gaps under IID conditions.The near-identical training and test errors are consistent with the reported absence of overfitting.
Error distributions: 0.0002 is the MAE achieved by Transformer, KAN, FNO, and CNN, whose error distributions are concentrated near zero.Transformer has the sharpest peak, while MLP has MAE = 0.0003.
Qualitative validation: Predicted magnetic-flux contours closely match reference solutions for circular and D-shaped plasma configurations across the evaluated models.The qualitative comparisons capture the reported agreement in complex topological features.

4.2 Robustness Evaluation under IID and OOD Plasma Configurations

All five architectures perform strongly on IID equilibria, but extrapolation exposes large architecture-dependent differences. CNN and FNO are the most robust under unseen geometries and parameters, while KAN degrades severely.

CNN and FNO achieve the strongest OOD performance, with L2 errors of 4.3% and 4.0% under extreme geometric configurations.
Overall, CNN and FNO provide the strongest extrapolation behavior, indicating that structural inductive biases support distribution-shift robustness.
All models maintain IID L1 relative errors below 5×10^-3, with Transformer achieving the lowest median L1 error at 1.95×10^-3.
CNN records the best median extrapolation L1 error at 0.027, while FNO achieves a comparable 0.031.
Transformer retains intermediate extrapolation performance with a median L1 error of 0.038 despite superior interpolation accuracy.
KAN’s median extrapolation L1 error reaches 0.72, reflecting severe degradation from its interpolation performance.

4.3 Scaling Behavior of AI Surrogate Models

Scaling improves interpolation accuracy, but its effect on extrapolation depends on architecture, dataset coverage, and model capacity. CNN scales most consistently, whereas excessive FNO capacity and additional Transformer data can reduce OOD gains.

Dataset-size scaling: Transformer reaches an interpolation L1 error of 2.01×10^-3 with 60,000 samples, improving 50% from 4.23×10^-3 at 10,000 samples.
Dataset-size scaling: FNO’s extrapolation L1 error decreases 35%, from 6.48×10^-2 with 10,000 samples to 4.18×10^-2 with 60,000 samples.
Dataset-size scaling: CNN’s extrapolation L1 error falls 51%, from 8.73×10^-2 to 4.26×10^-2, as training coverage increases.
Dataset-size scaling: Transformer extrapolation varies non-monotonically with dataset size, unlike its consistently strong interpolation performance.
Model-capacity scaling: FNO extrapolation L1 error increases from 0.0391 to 0.0763 as model size grows from 0.01M to 1M parameters.
Overall scaling behavior: Larger datasets and models generally improve IID accuracy but do not necessarily improve OOD robustness, creating a capacity-generalization trade-off.

4.5 Device-Relevant Validation on the EXL-50U Tokamak

EXL-50U validation links numerical GS solutions, CNN predictions, and Shape Editor references to test simulation-to-device consistency. The CNN reproduces device-level flux geometry with relative discrepancies mainly in the 10^-2–10^-3 range.

Validation uses representative EXL-50U discharges #6460 and #7533, with GS solutions as ground truth and Shape Editor as the operational baseline.
Figure 14 compares poloidal-flux contours and relative-error distributions for GS, Shape Editor, and CNN predictions across both discharges.
CNN predictions closely match GS magnetic-axis positions, contour shapes, flux-surface geometry, and nested topology on both device cases.
CNN relative error predominantly lies in the 10^-2–10^-3 range over the plasma cross-section.
The dual-validation pathway connects GS accuracy with operational Shape Editor consistency, supporting device-level physical reliability of the surrogate approach.

5. Conclusion

The study establishes a unified benchmark for five AI architectures in fixed-boundary GS prediction, combining accuracy, efficiency, scaling, robustness, and device validation. Results favor Transformer for IID accuracy and CNN/FNO for OOD robustness, while highlighting a trade-off between interpolation gains and extrapolation reliability.

The framework benchmarks MLP, CNN, FNO, KAN, and Transformer using high-fidelity IID/OOD data and a unified evaluation protocol.
Transformer achieves the highest IID accuracy, with relative L1 and L2 errors of 0.201% and 0.242%, respectively.
CNN balances accuracy, OOD robustness, and efficiency, reaching 0.7 ms TensorRT latency.
GS-to-Shape Editor discrepancies remain predominantly on the order of 10^-3, while CNN-to-GS discrepancies are 10^-3–10^-2 on EXL-50U.
CNN and FNO show the strongest OOD robustness, whereas KAN severely degrades despite competitive IID accuracy.
Increasing dataset size and model capacity generally improves IID accuracy but does not necessarily enhance OOD robustness.
The benchmark provides practical guidance for selecting AI surrogates for fast tokamak equilibrium prediction and future real-time plasma-control applications.