Real-Time Plasma State Prediction via FPGA-Accelerated Quantized Recurrent Probabilistic Neural Networks

Daniel Gaytan-Villarreal; Aiken Xie; Tu Pham; Rohit Sonker; Chiara Amendola; Matteo Cremonesi; Cong Hao; Jeff Schneider

physics.plasm-phSep 19, 2026

Real-Time Plasma State Prediction via FPGA-Accelerated Quantized Recurrent Probabilistic Neural Networks

Authors: Daniel Gaytan-Villarreal, Aiken Xie, Tu Pham, Rohit Sonker, Chiara Amendola, Matteo Cremonesi, Cong Hao, Jeff Schneider

Organizations: Department of Physics, Carnegie Mellon University, Pittsburgh, PA 15213, USA · Department of Applied Physics and Applied Mathematics, Columbia University, New York, NY 10027, USA · Department of Electrical and Computer Engineering, Georgia Institute of Technology, Atlanta, GA 30332, USA · Machine Learning Department, Robotics Institute, Carnegie Mellon University, Pittsburgh, PA 15213, USA

Abstract

Real time plasma state estimation for control of Tokamak devices are challenging due to the stringent latency requirements of the plasma control system (PCS). We present an end-to-end workflow for deploying a recurrent probabilistic neural network (RPNN) on FPGA hardware. We combine architecture size reduction with quantization-aware training via QKeras. The model is then synthesized using hls4ml, targeting a Xilinx Alveo U50 device. We report a design that fits comfortably within all four resource budgets (DSP, LUT, FF, BRAM) at deterministic sub-10~ $μ$ s single-timestep latency, meeting the requirements for real-time inference inside a model-predictive-control-style plasma control loop.

Figures & tables

Group	Representation Type	Signal	Dimension
States	Scalar	$\beta_{N}$ , line-averaged density, internal inductance ( $l_{i}$ ), $q_{0}$ , $q_{95}$ , $v_{\rm loop}$ , WMHD	7
States	Profile	electron temperature, ion temperature, density, rotation (4 PCA components each); pressure, Safety factor $q$ (2 PCA components each)	20
Actuators	Beam	power injected, torque injected	2
	Shape	$a_{\rm minor}$ , triangularity (top/bottom), $\kappa$ , $R$ and $Z$ coordinates of magnetic axis	6
	Gas	gas A puffing	1
	Other	plasma current target, toroidal field (magnitude + sign)	3

Table 1: List of all state and actuator variables. The total dimension (40) counts each unique variable once; the model’s actual per-timestep input is 53-dimensional, since the 13 actuator channels are each supplied twice — once as their current absolute level and once as the commanded delta (change) into the following timestep. See Fig. 1 .

Component	Parameters	Share
Encoder MLP	$290{,}304$	$8.4\%$
BatchNorm layer	$2{,}048$	$0.1\%$
GRU	$591{,}360$	$17.1\%$
Decoder (residual MLP)	$2{,}560{,}640$	$74.2\%$
Mean / log-variance heads	$6{,}966$	$0.2\%$
Variance pinning	$54$	$\sim\!0.0\%$

Table 2: Trainable parameter count by component for the full-size RPNN architecture.

Norm type	MSE	EV
LayerNorm	$0.02728$	$0.4967$
BatchNorm	$0.02622$	$0.5166$
$\Delta$ (BatchNorm vs. LayerNorm)	$-3.9\%$	$+4.0\%$

Table 3: Normalization-type training comparison: both variants trained from scratch under an identical curriculum and evaluated on the same chronological validation split (

982

shots,

2{,}407

sequences).

Norm type	MSE	EV
LayerNorm	$0.02407$	$0.4038$
BatchNorm	$0.02375$	$0.4243$
$\Delta$ (BatchNorm vs. LayerNorm)	$-1.3\%$	$+5.1\%$

Table 4: Normalization-type evaluation of Keras checkpoints trained from scratch (same models as Table 3 ), evaluated on the held-out test split.

Configuration	$\Delta$ Params	$\Delta$ MSE	$\Delta$ EV
dec384	$-29.9\%$	$-0.8\%$	$-0.1\%$
blocks1	$-30.4\%$	$+2.9\%$	$-2.0\%$
dec256	$-52.3\%$	$-2.2\%$	$+3.0\%$
dec256_blocks1	$-59.9\%$	$-0.8\%$	$+2.2\%$
gru128_dec256	$-63.2\%$	$+2.0\%$	$-0.4\%$
dec256_b0	$-63.7\%$	$-1.1\%$	$+1.7\%$

Table 5: Architecture size ablation, training-matched comparison on the chronological validation split, relative to the full-size baseline (

3{,}451{,}532

params, MSE

=0.02622

, EV

=0.5166

). Naming: dec

N

sets decoder_hidden_dim

=N

; blocks1 / b

K

set decoder_num_res_blocks

=1

K

(default 3); gru

N

sets gru_hidden_dim

=N

(default 256); hid

N

sets the shared backbone width hidden_dim

=N

(default 512).

Configuration	$\Delta$ Params	$\Delta$ MSE	$\Delta$ EV
dec256	$-52.3\%$	$-1.5\%$	$+4.2\%$
gru128_dec256	$-63.2\%$	$-1.3\%$	$+2.0\%$
gru128_dec256_b1	$-70.8\%$	$-0.7\%$	$-0.1\%$
hid256_dec256	$-66.0\%$	$-0.2\%$	$+3.0\%$
hid256_dec256_b1	$-73.6\%$	$-0.2\%$	$+1.6\%$
hid128_gru128_dec256_b1	$-85.7\%$	$+0.2\%$	$-4.7\%$

Table 6: The same architecture size ablation configurations as Table 5 , evaluated on the held-out chronological test split (

990

most-recent shots,

3{,}000

sequences, never touched during training or model selection), relative to the full-size baseline on this split (MSE

=0.02375

, EV

=0.4243

). This is the population used for all deployment-accuracy comparisons in this paper. Rows are ordered by performance (best to worst

\Delta

MSE), not by parameter count as in Table 5 .

Panel of Figure 1 — Figure 1: RPNN architecture, drawn with the layer widths of the deployed configuration ( hid128_gru64_dec128_b1 , Sec. III.1 ); the full-size baseline shares the identical topology with wider layers (backbone $512$ , GRU $256$ , decoder $512$ with three residual blocks), and each block’s width reduction relative to it is noted below the block (e.g. $512\!\to\!128$ ). States (27) and actuators (13, supplied both as their current absolute level and as the commanded delta into the following timestep, for a 53-dimensional encoder input) are normalized and encoded, then passed through a BatchNorm layer before a single-layer GRU; the GRU output is concatenated with the encoder’s own output (bypassing the GRU) and passed to a residual-MLP decoder, whose output feeds two linear heads producing the predicted mean and log-variance, the latter passing through a variance-pinning nonlinearity.

Component	Parameters	Share
Encoder MLP	$23{,}424$	$13.3\%$
BatchNorm layer	$512$	$0.3\%$
GRU	$37{,}248$	$21.2\%$
Decoder (residual MLP)	$107{,}264$	$61.1\%$
Mean / log-variance heads	$6{,}966$	$4.0\%$
Variance pinning	$54$	$\sim\!0.0\%$

Table 7: Parameter count by pipeline component for the deployed architecture, hid128_gru64_dec128_b1 (

175{,}628

total params,

-94.9\%

relative to the full-size architecture of Table 2 ), using the same convention as Table 2 .

Precision	MSE	$\Delta$ MSE	EV	$\Delta$ EV
ap_fixed<16,6>	$0.02487$	$+2.8\%$	$0.3932$	$-4.8\%$

Table 8: 16-bit QAT on the chosen deployment architecture, final accuracy after fine-tuning, relative to the FP32 reference for this architecture (MSE

=0.02420

, EV

=0.4128

) on the held-out test split.

Figure 2: Fully autoregressive ensemble replay of held-out test shot

193369

(

225

steps), FP32 baseline (blue, all 15 bootstrap members, 45 rollouts) vs. the deployed 16-bit hid128_gru64_dec128_b1 model (green, same convention), against ground truth (dashed black). Top row: all six profile quantities at the final timestep in physical units, vs. minor radius index (the safety-factor profile is stored, and shown, as

1/q

). Remaining rows: four representative scalar quantities (

\beta_{N}

l_{i}

q_{0}

q_{95}

) over the full shot, plotted against absolute time into the discharge (

320

–

4{,}820\,

ms, so the ramp-up is included); RMSE for these and the three not plotted here ( dssdenest , vloop , wmhd_EFIT01 ) is reported in full in Table 9 . Faint lines are individual rollout samples; bold lines are the ensemble mean. Vertical axes are clipped to the range of the ground truth and ensemble means, so a few individual sampled trajectories extend beyond the plotted range.

Quantity	FP32	16-bit	$\Delta$
betan_EFIT01	$0.4068$	$0.3775$	$-7.2\%$
dssdenest	$0.3010$	$0.3125$	$+3.8\%$
li_EFIT01	$0.5686$	$0.6422$	$+12.9\%$
q0_EFIT01	$0.5772$	$0.4221$	$-26.9\%$
q95_EFIT01	$0.6557$	$0.5706$	$-13.0\%$
vloop	$0.4857$	$0.5031$	$+3.6\%$

Table 9: Scalar-quantity RMSE (normalized units) over the full autoregressive rollout of shot

193369

(16-bit vs. floating-point 32-bit precision).

Quantity	FP32	16-bit	$\Delta$
temp	$0.0864$	$0.0593$	$-31.4\%$
itemp	$0.1664$	$0.3034$	$+82.3\%$
dens	$0.5862$	$0.6106$	$+4.2\%$
rotation	$0.3173$	$0.4214$	$+32.8\%$
pres_EFIT01	$0.0788$	$0.0287$	$-63.6\%$
q_EFIT01	$0.0203$	$0.0805$	$+296.7\%$

Table 10: Profile RMSE (normalized units, 33-point radial reconstruction) at the final timestep of the same rollout (ensemble mean vs. ground truth).

Figure 3: Full-state explained variance (Eq. 1 with

y=s_{s+t}

and

\hat{y}

the model’s accumulated prediction) as a function of autoregressive rollout horizon over steps

0

–

100

, pooled over every starting position in every held-out test shot (

112{,}436

(shot, start) combinations,

30

sampled continuations each). Top row: the average across all

27

output channels and five representative scalars (

\beta_{N}

l_{i}

q_{0}

q_{95}

, line-averaged density); bottom row: the first PCA component of all six profile quantities. Curves compare the full-size FP32 baseline (blue: mean mode; crimson: sample mode) against the deployed 16-bit hid128_gru64_dec128_b1 model (green: mean; orange: sample), each averaged over four members from the ensemble, with shaded bands showing the standard error across members. Curves are smoothed for display with a centered

15

-step rolling average; the underlying per-step values are unsmoothed.

Figure 4: Uncertainty metrics over rollout horizon for the three predictive-distribution methods (blue: Sample-TS1; green: Sample-TSInf; red: Mean-TS1, the reference paper’s palette), comparing the full-size FP32 baseline ensemble (dark, solid; members $\{1,2,3,7\}$ ) against the deployed 16-bit hid128_gru64_dec128_b1 ensemble (light, dashed; members $\{1,4,7,12\}$ ) within each method’s hue. Panels, left to right: empirical coverage of the $90\%$ prediction interval (Eq. 2 , dotted line marks the $0.90$ target); miscalibration area (Eq. 3 , lower is better); EV averaged over all $27$ output channels (axis clipped below $0$ ). Shaded bands are $90\%$ shot-level cluster-bootstrap confidence intervals ( $B=200$ ) of each mean curve, computed from identical resamples for every curve so visible gaps are not resampling noise. The horizon extends to the data ceiling $t=277$ , but the active population collapses past the $200$ -step mark of Fig. 3 — $n=3{,}508$ combinations at $t=200$ , $1{,}213$ at $t=225$ , $283$ at $t=250$ , $7$ at $t=275$ — so only the longest shots remain and the tail beyond $t\approx 250$ is qualitative (the brief EV spike near $t\approx 267$ is the $n=7$ population, not a recovery).

Resource	Used	Budget	Util.
LUT	$262{,}381$	$744{,}638$	$35.24\%$
FF (REG)	$280{,}955$	$1{,}579{,}376$	$17.79\%$
DSP	$3{,}868$	$5{,}948$	$65.03\%$
BRAM	$562$	$1{,}163$	$48.32\%$
URAM	$0$	$636$	$0.0\%$

Table 11: Post-route resource utilization of the deployed stateful hid128_gru64_dec128_b1 kernel (16-bit QAT,

\mathrm{RF}=64

, GRU hidden state as explicit kernel I/O) on the Alveo U50, relative to the user-accessible resource budget.

	Baseline (3.45M)		Ablated (176K)
Batch	CPU	A100	CPU	A100
$1$	$1{,}011$	$1{,}249$	$608$	$1{,}085$
$16$	$106$	$81.0$	$45.0$	$67.3$
$128$	$34.2$	$10.1$	$9.9$	$8.7$
$1{,}024$	$24.9$	$1.45$	$2.83$	$1.08$
FPGA, end-to-end latency	$2.64\,\mu$ s

Table 12: Software-vs-FPGA inference cost for one control-loop tick. Software rows: median wall-clock per call divided by batch size (in

\mu

s per sample), for the FP32 TensorFlow implementations of the full-size baseline (

3.45

M parameters) and the deployed ablated architecture (

176

K parameters), on an 8-core AMD EPYC 7763 allocation and a dedicated NVIDIA A100 80 GB (PCIe, exclusive use), timed as a compiled tf.function with device synchronization included. Batch size

1

is the operative regime for closed-loop control, where each tick depends on the previous prediction and cannot be batched. FPGA rows: the deployed 16-bit stateful kernel of Table 13 .

Component	Cycles	Time
AXI in: enqueue/demux (117 words)	$120$	$1.146\,\mu$ s
Encoder L1 (Dense $53\!\to\!128$ , 27-bit in)	$70$ – $71$	$0.669$ – $0.678\,\mu$ s
Encoder ReLU	$0$	$0$
Encoder L2 (Dense $128\!\to\!128$ )	$65$ – $66$	$0.621$ – $0.630\,\mu$ s
BatchNorm layer	$63$	$0.602\,\mu$ s
GRU (64-unit, stateful, $h$ explicit I/O)	$70$ – $71$	$0.669$ – $0.678\,\mu$ s

Table 13: Per-component latency breakdown for the deployed stateful hid128_gru64_dec128_b1 kernel (Alveo U50, io_stream /Resource strategy,

\mathrm{ReuseFactor}=64

, GRU hidden state as explicit kernel I/O). Cycle counts are Vitis HLS 2022.1 C-synthesis per-instance figures; wall-clock times use the real post-route achieved kernel clock of

104.7

MHz (

9.55

ns/cycle). The AXI wrapper rows are the physical interface cost of streaming the

117

input and

118

output words at one word per cycle; see the text for how to read the per-layer rows.

FPGA-Based Neural Network Acceleration

Field-programmable gate array (FPGA)-based neural network acceleration maps models to reconfigurable hardware for fast, energy-efficient inference.

See all 90 papers

An Embedded RISC-V Evaluation of Kolmogorov--Arnold Networks in Hard-Constrained Recurrent Physics-Informed Models
Aug 1, 2026Enzo Nicolas Spotorno, Josafat Leal FilhoEfficient Neural Network Inference Kolmogorov-Arnold Networks
Department of Informatics and Statistics - Graduate Program in Computer Science (PPGCC)
Real-time virtual circuits for plasma shape control via neural network emulators: integration and testing in the MAST-U PCS
Aug 26, 2026Matthew J. Marshall, Edward Jones, Graham J. McArdle +10AI Control Feedback Control
Hartree Centre, STFC, Daresbury Laboratory, Daresbury, WA4 4AD, United Kingdom · UKAEA, Culham Campus, Abingdon, OX14 3DB, United Kingdom
PQuantML: A Tool for End-to-End Hardware-aware Model Compression
Mar 27, 2026Roope Niemi, Anastasiia Petrovych, Arghya Ranjan Das +9Model Compression Neural Network Compression
European Center for Nuclear Research (CERN), CH-1211 Geneva, Switzerland · Purdue University, West Lafayette, IN 47907, USA · California Institute of Technology, Pasadena, CA 91125, United States +4

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

An Embedded RISC-V Evaluation of Kolmogorov--Arnold Networks in Hard-Constrained Recurrent Physics-Informed Models

Enzo Nicolas Spotorno, Josafat Leal Filho

Department of Informatics and Statistics - Graduate Program in Computer Science (PPGCC)

Efficient Neural Network Inference Kolmogorov-Arnold Networks Physics-Informed ML On-Device Inference

Open paper AskPDF arXiv

Hard-constrained recurrent physics-informed networks (HRPINNs) embed known dynamics inside a recurrent numerical integrator and restrict a neural branch to learning only the residual dynamics that the first-principles model does not capture. Kolmogorov--Arnold Networks (KANs) have been proposed as parameter-efficient replacements for multilayer perceptrons (MLPs) in such residual branches, but their learnable B-spline activations follow a markedly different execution profile. Building on prior work that characterized when a vanilla B-spline KAN matches or underperforms an MLP as an HRPINN residual branch in discovery accuracy, this paper asks whether that parameter efficiency survives deployment. Using identical trained weights, we measured execution latency, energy per integration step, and dependability under post-training quantization in the closed recurrent loop on a RISC-V RV64GC platform without vector extensions (StarFive VisionFive~2, SiFive U74). For the two accuracy-comparable pairs, the KAN residual branch executed

13.5\times

and

8.0\times

slower and consumed

11.3\times

and

5.6\times

more energy per integration step (3.7,

μ

J against 0.33,

μ

J for the smallest pair); across all four parameter-matched size tiers the ranges are

4.7\times

14.5\times

and

4.7\times

18.7\times

. Under INT8 quantization, KAN trajectories diverged up to

43\times

earlier than matched MLPs; the damage traces to weight quantization, not to input-side knot-interval misassignment. These results indicate that the parameter efficiency reported for KANs does not transfer to deployment cost on scalar embedded cores, and that an MLP residual branch is the more dependable default for embedded HRPINN deployment unless specific quantization co-design is used.

Abstract

FPGA-Based Neural Network Acceleration

Explore similar work

An Embedded RISC-V Evaluation of Kolmogorov--Arnold Networks in Hard-Constrained Recurrent Physics-Informed Models

Real-time virtual circuits for plasma shape control via neural network emulators: integration and testing in the MAST-U PCS

PQuantML: A Tool for End-to-End Hardware-aware Model Compression