A Generative Foundation Model for Multimodal Tokamak Tasks

Yingshuai Ji; Haodong Feng; Guohong Deng; Gelin Bian; Yue Wang

cs.AINov 5, 2026

A Generative Foundation Model for Multimodal Tokamak Tasks

Authors: Yingshuai Ji, Haodong Feng, Guohong Deng, Gelin Bian, Yue Wang

Organizations: Zhongguancun Academy · XinZhuAI

Abstract

Tokamak experiments record diverse measurements, including waveforms, profiles, fields, and scalars, which are used to reconstruct the plasma, forecast its evolution, recover associated control signals, and estimate disruption risk. Most models are built separately for one task and a fixed set of inputs and outputs, and each relies on a prescribed subset of diagnostics rather than structuring the full multimodal observation set to extract global, cross-signal information about the discharge. These limitations motivate a reusable foundation model that adapts to each task's input-output schema and models heterogeneous target signals jointly, rather than retraining a separate model for every configuration. In this work, we extend the TokaMark benchmark with two groups of tasks for actuator action inversion and disruption prediction, yielding 20 MAST tasks in six groups with explicit schemas and evaluation protocols. Then, we introduce GenToka, a generative tokamak foundation model that jointly generates schema-selected targets through a shared conditional diffusion Transformer, with signal-specific codecs preserving native time series, profiles, fields, and scalars; foundation pretraining and task-specific adaptation transfer cross-signal structure across all six groups. Evaluation on MAST measurements shows that GenToka reduces macro-average normalized root mean squared error relative to TokaMind by 4.7% on Groups 1-4 and by 17.7% on Group 5, and improves macro-average window-level F1 on Group 6 by 8.4% at a common validation-selected threshold. GenToka provides a reusable, generative, and multimodal pathway for experimental tokamaks, linking plasma-state reconstruction, actuator inference, and disruption warning from shared measurements to support more integrated, data-driven fusion plasma prediction and control research.

Figures & tables

Task family	Conditioning information	Requested output
Equilibrium reconstruction	Co-temporal magnetic and plasma observations	Global equilibrium quantities, plasma boundary, or poloidal-flux field
Magnetic dynamics	Recent magnetic and plasma state with task-provided actuator traces	Future magnetic and equilibrium time series
Profile and diagnostic dynamics	Recent profiles and diagnostics with operating context	Future kinetic profiles, global plasma quantities, or diagnostic responses
Long-horizon MHD dynamics	Extended multimodal history with operating context	Future plasma-current, equilibrium, or diagnostic time series
Actuator action inversion	State and control history with an observed same-discharge future response	Future recorded control-signal time series
Disruption prediction	Strictly pre-anchor diagnostic and, where configured, control-related history	A scalar time-to-disruption risk surrogate

Table 1: Tokamak task families considered in this work. Each family can contain multiple context and target schemas without changing its underlying objective.

Figure 1: Overview of GenToka. (a) Unified signal representation: multimodal measurements are time-chunked and encoded with signal-specific DCT3D; context and target adapters with schema metadata yield clean context tokens and noisy target tokens. (b) Joint conditional diffusion: the same conditional Transformer denoises all requested targets jointly at each noise level while context tokens remain clean. (c) Native-space reconstruction: denoised target codes are inverse-normalized and decoded to schema-selected time series, profiles, fields, or scalar Soft-TTD scores. (d) Pretraining and task adaptation: foundation training uses co-temporal reconstruction with random context masking; each downstream schema is adapted in two stages by aligning interfaces with the backbone frozen, then refining the backbone and target pathway with the context tokenizer frozen. (e) Benchmark scope: twenty MAST tasks in six families: equilibrium reconstruction, magnetic dynamics, profile and diagnostic dynamics, long-horizon MHD-relevant dynamics, actuator action inversion, and disruption prediction.

Group	Task	Metric	TokaMind	GenToka	Scratch
Group 1	1-1	NRMSE $\downarrow$	0.1363	0.0941	0.0991
	1-2	NRMSE $\downarrow$	0.0439	0.0412	0.0412
	1-3	NRMSE $\downarrow$	0.1200	0.1156	0.1184
	Group 1 avg	NRMSE $\downarrow$	0.1001	0.0836	0.0862
Group 2	2-1	NRMSE $\downarrow$	0.1289	0.0880	0.0931
	2-2	NRMSE $\downarrow$	0.0408	0.0370	0.0375

Table 2: GenToka model results across all twenty tasks. Groups 1–5 report native-space NRMSE; Group 6 reports window-level AUROC and F1, with

\tau^{\star}=0.4

for F1. Means are unweighted over tasks. Scratch uses the GenToka architecture with random initialization and the same downstream recipe for each task. Bold identifies the best displayed value, including ties.

Group	Task	Metric	LSTM	TokaMark-CNN
Group 1	1-1	NRMSE $\downarrow$	0.1853	0.1491
	1-2	NRMSE $\downarrow$	0.0488	0.0408
	1-3	NRMSE $\downarrow$	0.1302	0.1354
	Group 1 avg	NRMSE $\downarrow$	0.1215	0.1084
Group 2	2-1	NRMSE $\downarrow$	0.1633	0.1376
	2-2	NRMSE $\downarrow$	0.0461	0.0482

Table 3: Benchmark reference results for all twenty tasks. LSTM and TokaMark-CNN are trained separately for each task from random initialization and evaluated with the common native-space metrics. Groups 1–4 contain the fourteen TokaMark forward tasks; Groups 5 and 6 contain the six newly created tasks (actuator action inversion and disruption prediction). Means give equal weight to tasks. Disruption F1 uses the frozen threshold

\tau^{\star}=0.4

. The two reference implementations use their respective representations and training recipes (Appendix A.4 ).

Figure 2: On a co-temporal reconstruction example, GenToka tracks the measured electron temperature and density profiles more closely than TokaMind, matching the lower aggregate equilibrium error summarized above.

Figure 3: During magnetic-dynamics forecasting, GenToka follows the separatrix in $(R,Z)$ more tightly than TokaMind on this discharge, echoing the largest single-task NRMSE gain in the family.

Figure 4: When predicting kinetic profiles ahead in time, GenToka produces smoother, better-aligned

T_{e}

and

n_{e}

than TokaMind, consistent with the stronger reductions on the most profile-heavy targets.

Figure 5: At a 100 ms horizon the two models diverge on magnetic-axis and separatrix geometry on the same shot, explaining why long-horizon averages can tie even when individual targets favor different models.

Figure 6: Given an observed plasma response in the same discharge, GenToka reconstructs the recorded PF-coil voltages more faithfully than TokaMind, in line with the large gains on magnetic actuation inversion.

Figure 7: Relative test-NRMSE improvement of GenToka over TokaMind by task family (positive is better): GenToka leads on equilibrium, magnetic, profile, and actuator-inversion families, ties on long-horizon MHD, and Scratch without pretraining is consistently weaker.

Figure 8: Near a disruption-warning window with multivariate history, GenToka assigns clearer separation at the shared risk threshold than TokaMind, matching the improved F1 despite similar AUROC.

Appendix figures & tables4 assets

Supplementary material from the paper’s appendix.

Appendix

Task	Observed history	Observed future response	Prediction target	Window and held-out cohort
5-1	Magnetic diagnostics; PF-coil, solenoid and plasma currents; past PF voltages; optional NBI power.	P2/P4/P5 coil currents, solenoid current and $I_{p}$ ; optional LCFS $(R,Z)$ contour.	Four PF voltage traces: P1, P2, P4, P5.	History: 5 ms; response/target: 25 ms; 1 ms cut grid; 1,103 test shots.
5-2	Magnetic diagnostics; PF-coil, solenoid and plasma currents; line-integrated density; Thomson $T_{e},n_{e}$ ; past current/density requests, NBI power, cumulative gas and gas-valve voltages.	Thomson electron-density profile.	Six gas-valve voltage traces: G1, G2, G3, G4, G11, G12.	History: 150 ms (Thomson: 5 ms); response/target: 50 ms; 1 ms cut grid; 681 test shots.
5-3	Thomson $T_{e},n_{e}$ ; D-alpha and upper/lower soft-X-ray camera voltages; past current/density requests, NBI power, cumulative gas and gas-valve voltages.	Three-channel D-alpha voltage time series.	Six gas-valve voltage traces: G1, G2, G3, G4, G11, G12.	History: 5 ms; response/target: 50 ms; 1 ms cut grid; 803 test shots.
6-1	Nine signals: $I_{p}$ , line-integrated density, radiated power, D-alpha voltage, lower soft-X-ray camera voltage, $q_{95}$ , internal inductance $l_{i}$ , diamagnetic energy and minor radius.	—	Soft time-to-disruption risk.	History: 100 ms; 10 ms cut grid; 933 test shots / 39,967 windows.
6-2	Five signals: $I_{p}$ , line-integrated density, D-alpha voltage, lower soft-X-ray camera voltage and minor radius.	—	Soft time-to-disruption risk.	Same history, grid, shots, window coordinates and targets as 6-1.
6-3	Plasma current $I_{p}$ .	—	Soft time-to-disruption risk.	Same history, grid, shots, window coordinates and targets as 6-1.

Appendix

Table 4: Schemas of the six extension tasks. Windows are relative to the cut time

t_{c}

and are left-closed, right-open. Task 5 conditions on a paired observed future plasma response to reconstruct the corresponding actuator time series. Task 6 uses causal history to predict a continuous soft time-to-disruption target.

Task	Target component	TokaMind	GenToka	Reduction (%)
5-1	XDC_PF_F_P1	0.1545	0.0706	54.3
5-1	XDC_PF_F_P2	0.0710	0.0360	49.3
5-1	XDC_PF_F_P4	0.1437	0.0825	42.6
5-1	XDC_PF_F_P5	0.1507	0.0881	41.5
5-2	G1	0.2930	0.2842	3.0
5-2	G2	0.2874	0.2761	3.9

Appendix

Table 5: Decoded native-space component NRMSE on the frozen Group 5 test cohorts. Each component uses the shared signal-role standard deviation and equal weighting across accepted test shots. Relative reduction is measured against TokaMind. Component averages differ from task NRMSE because task NRMSE pools output components within each shot before averaging across shots.

Setting	LSTM	TokaMark-CNN
Training epochs	50	50
Optimizer	AdamW	Adam
Learning rate	$6\times 10^{-4}$	$3\times 10^{-4}$
Weight decay	$10^{-2}$	0
Global batch size	1024	2048
Supervision / selection	Masked coefficient MSE	Masked native-space MSE

Appendix

Table 6: Training settings for the two benchmark reference implementations. Both use seed 54, eight GPUs, a 2% linear learning-rate warmup, and 3000 fixed shot-stratified validation windows for checkpoint selection. The reported result for each task uses its best validation checkpoint.

$\tau$	0.1	0.2	0.3	0.4	0.5	0.6	0.7	0.8	0.9
Mean F1	0.4105	0.4884	0.5391	0.5570	0.5226	0.4675	0.3932	0.2983	0.1724

Appendix

Table 7: Validation macro-average F1 for the common disruption-risk threshold. The average covers TokaMind and GenToka on Tasks 6-1–6-3; no test labels are used. Bold entries identify the selected threshold and its validation score.

Multimodal Foundation Models

Multimodal foundation models learn shared capabilities across text, images, audio, video, and other signals for understanding, generation, and action.

See all 357 papers

Feb 16, 2026

TokaMind: A Multi-Modal Transformer Foundation Model for Tokamak Plasma Dynamics

We present TokaMind, to our knowledge the first open-source foundation model for tokamak plasma dynamics, based on a Multi-Modal Transformer (MMT) and pretrained on heterogeneous diagnostics from the publicly available MAST dataset. TokaMind supports multiple data modalities (time-series, 2D profiles, and videos) with different sampling rates, robust missing-signal handling, and efficient task adaptation via selectively loading and freezing four model components. To represent multi-modal signals, we use a lightweight fixed-basis Discrete Cosine Transform embedding (DCT3D) and provide a clean interface for alternative embeddings (e.g., Variational Autoencoders). We evaluate TokaMind on the recently introduced MAST benchmark TokaMark, which comprises 14 tasks with heterogeneous reconstruction and forecasting objectives. Our results show that fine-tuned TokaMind outperforms the strongest benchmark baseline on all but one task. Compared with training the same architecture from scratch under a matched epoch budget, warm-start adaptation is most beneficial on demanding downstream settings, including long-horizon forecasting and high-dimensional equilibrium objectives. These findings highlight the value of multi-modal pretraining for tokamak plasma dynamics and provide a practical, extensible foundation for future fusion modeling tasks. Training code and model weights are publicly available at github.com/UKAEA-IBM-STFC-Fusion-FMs/tokamind and huggingface.co/UKAEA-IBM-STFC, respectively.

Tobia Boschi, Andrea Loreti, Nicola C. Amorisco +13

IBM Research Dublin, Ireland · UK Atomic Energy Authority Abingdon, UK · STFC Hartree Centre Daresbury, UK

Multimodal Foundation Models Multimodal Pretraining Time Series Forecasting

Ask PDF

May 10, 2026

TokaMind for Power Grid: Cross-Domain Transfer from Fusion Plasma

TokaMind is a multi-modal transformer (MMT) foundation model pre-trained on tokamak plasma diagnostics data from MAST, where it was shown to outperform CNN-based approaches on fusion benchmarks. We investigate whether its learned representations generalize to physically distinct but structurally analogous domains. Through systematic experimentation across four domains-industrial bearing degradation, NASA CMAPSS turbofan degradation, and two independent power grid PMU datasets-we identify four transfer-favoring characteristics that help explain where TokaMind's pretrained representations are most effective. Power grid synchrophasor data matches this target-domain profile most directly, while industrial degradation datasets demonstrate that TokaMind can still yield useful performance under partial alignment, especially when task design and feature construction expose physically meaningful degradation structure. On the GESL/PNNL 500-event benchmark with provider-aware evaluation, TokaMind achieves test

\text{F1} = 0.837 \pm 0.040

(3seeds) for severe event classification. Our central finding, however, is not the aggregate score: classification difficulty is structurally determined by provider-level grid topology, not model capacity. In the single-window early-warning regime, TokaMind outperforms a CNN baseline (F10.889 vs.~0.878)--a reversal that disappears as more event windows are provided. Furthermore, Critical Slowing Down (CSD) indicators, used as a confidence gate rather than a classification label, improve F1 from 0.696 to 0.750 at 63% coverage-outperforming the CNN baseline (0.636) at any coverage level. These results establish the first cross-domain validation of TokaMind outside nuclear fusion and propose a transferability framework and revised evaluation protocol for multi-source PMU datasets.

JC Wu, Norton Lee, Kai Siang Chen

TaiScience Research Group, affiliated with Fu Jen Catholic University, Taiwan · Center for Geometry and Physics, Institute for Basic Science (IBS), South Korea · Fu Jen Catholic University, Taiwan

Time Series Classification Critical Slowing Down Cross-Domain Transfer Learning Time-Series Foundation Models Cross-Dataset Generalization Power Systems

Ask PDF

Jul 5, 2026

Benchmarking Sensor Robustness in Plasma Diagnostic Models: A Systematic Evaluation on TokaMark

Plasma diagnostic models for tokamak fusion devices are almost universally evaluated on clean, complete sensor data. In practice, fusion diagnostics fail regularly: acquisition systems start late, individual sensors die, and signal dropouts cluster precisely when a plasma disruption is approaching. We present the first systematic robustness benchmark for plasma diagnostic ML using the TokaMark dataset of 11,573 MAST shots, evaluating XGBoost, LSTM, Transformer, and the TokaMark CNN baseline across six physically-grounded failure scenarios and three imputation strategies. We introduce the Robustness Score (RS) for standardized cross-architecture comparison. Our central finding is that disruption-proximate sensor failure (corruption injected in the final window timesteps) collapses sequence model performance (LSTM +212% NRMSE) while a statistical feature model remains comparatively stable (XGBoost +37%). Forward-fill imputation eliminates nearly all degradation from random dropout for sequence models (LSTM +57% to ~0%), but offers little help when the end of the window is corrupted. Shot-level alarm evaluation using ground-truth disruption timestamps reveals that LSTM alarm detection collapses to TPR=0.00 under proximate sensor failure, while mean-fill imputation recovers it to TPR=1.00, a reversal of the pattern observed in NRMSE. Plasma current emerges as the single most critical diagnostic across all architectures (+73% to +140% upon removal). Code, data, and trained checkpoints are available at https://github.com/Neerav-Gupta/tokamark-robustness.

Neerav Gupta

Independent Researcher

Incomplete Data Imputation

Ask PDF

Abstract

Multimodal Foundation Models

Explore similar work

TokaMind: A Multi-Modal Transformer Foundation Model for Tokamak Plasma Dynamics

TokaMind for Power Grid: Cross-Domain Transfer from Fusion Plasma

Benchmarking Sensor Robustness in Plasma Diagnostic Models: A Systematic Evaluation on TokaMark