TokaMind: A Multi-Modal Transformer Foundation Model for Tokamak Plasma Dynamics
Authors: Tobia Boschi, Andrea Loreti, Nicola C. Amorisco, Rodrigo H. Ordonez-Hurtado, Cécile Rousseau, George K. Holt, Eszter Székely, Alexander Whittle, Samuel Jackson, Adriano Agnello, Stanislas Pamela, Alessandra Pascale, Robert Akers, Juan Bernabe Moreno, Vassil Alexandrov, Mykhaylo Zayats
Organizations: IBM Research Dublin, Ireland · UK Atomic Energy Authority Abingdon, UK · STFC Hartree Centre Daresbury, UK
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.
Figures & tables
Figure 1. TokaMind tokenization and model architecture. Windowed multi-modal inputs I and actuators A are chunked and embedded by signal-specific codecs Eg to produce token embeddings zi∈RKg(i) (outputs are embedded at window level to form targets in the same space). A Token Encoder projects each token to the shared model dimension d and adds learned metadata embeddings (signal, role, modality, relative position). A Transformer Backbone processes the variable-length token set using an attention mask for missing/padded tokens and outputs the [CLS] (classification) token embedding. Modality-specific heads (TS: time series, P: Profile, V: Video) and per-output adapters predict embedded targets y^o∈RKo ; a target-availability mask mo excludes missing outputs from the supervised loss.
Change between runs
Model adaptation
Remove some inputs / actuators
✓
Fewer tokens; Transformer Backbone unchanged
Add new input / actuator
✓
New Token Encoder projection initialized; other blocks reused
Remove outputs
✓
Corresponding Output Adapters absent; loss ignores them via masks
Add new outputs
✓
New Output Adapters initialized; existing adapters reused
Change embedding dim / encoder for an existing signal
Different token count; relative-position encoding preserves a consistent notion of recency
Table 1. List of possible changes in the inputs and outputs between model runs and the corresponding adaptations.
Baselines
TokaMind
Task
Mean
Pers.
CNN
C+LSTM
Scratch
FT-Tiny
FT-Base
1-1
0.757
—
0.209
0.222
0.140
0.153
0.150
1-2
0.998
—
0.047
0.051
0.042
0.044
0.043
1-3
0.927
—
0.148
0.159
0.129
0.122
0.122
Group 1
0.894
—
0.135
0.144
0.104
0.106
0.105
2-1
0.817
—
0.275
0.220
0.180
0.192
0.181
Table 2. Test NRMSE on TokaMark . C+LSTM denotes the CNN+LSTM baseline and Pers. the Persistence baseline. Dashes indicate tasks for which Persistence is not defined. Best and second-best values per row are shown in bold and underlined, respectively. Group and overall averages are unweighted means over tasks.
Figure 2. Task-level relative NRMSE improvement over the strongest available benchmark baseline, averaged within each group. Positive values indicate lower NRMSE; negative values indicate worse performance.
VAE FT-Base
VAE Scratch
Task
NRMSE
Δ
NRMSE
Δ
1-1
0.165
+0.015
0.153
+0.013
1-2
0.041
-0.002
0.046
+0.004
1-3
0.122
+0.000
0.130
+0.001
Group 1
0.109
+0.004
0.110
+0.006
2-1
0.198
+0.017
0.189
+0.009
Table 3. Group 1 and Group 2 test NRMSE for VAE-based embeddings under FT-Base and Scratch regimes. Δ reports the NRMSE difference relative to the corresponding DCT3D result under the same training regime.
Appendix figures & tables2 assets
Supplementary material from the paper’s appendix.
Appendix
Signal
(H, W, T)
DCT3D
VAE
Kg
dim/ Kg
EE
Kg
dim/ Kg
Inputs shared by Groups 1 and 2
magnetics-flux_loop_flux
(15,1,25)
44
8.52
1.000
12
31.25
magnetics-b_field_pol_probe_ccbv_field
(40,1,25)
54
18.52
0.999
30
33.33
magnetics-b_field_pol_probe_obr_field
(18,1,25)
57
7.89
1.000
15
30.00
magnetics-b_field_pol_probe_obv_field
(18,1,25)
45
10.00
1.000
15
30.00
Appendix
Table A.1. Group 1 and Group 2 input and actuator signals. (H,W,T) denotes the native tensor shape for a single 5 ms chunk, with dim=HWT . DCT3D columns report the source codecs used by warm-start models: embedding dimension Kg , compression ratio dim/Kg , and explained energy (EE). VAE columns report the latent embedding dimension Kg and the corresponding compression ratio. Inputs are shared by Groups 1 and 2, while actuators are used only by Group 2.
Group
Task
Signal
DCT3D
Neural Baselines
TokaMind
(H,W,T)
Kg
dim/ Kg
EE
CNN
C+LSTM
Scratch
FT-Tiny
FT-Base
Group 1
1-1
equilibrium-beta_normal
(1,1,1)
1
1.00
1.000
0.222
0.243
0.157
0.170
0.172
equilibrium-beta_pol
(1,1,1)
1
1.00
1.000
0.224
0.230
0.153
0.166
0.170
equilibrium-beta_tor
(1,1,1)
1
1.00
1.000
0.225
0.267
0.157
0.175
0.173
equilibrium-bphi_rmag
(1,1,1)
1
1.00
1.000
0.285
0.286
0.179
0.199
0.195
equilibrium-bvac_rmag
(1,1,1)
1
1.00
1.000
0.271
0.278
0.154
0.180
0.169
Appendix
Table A.2. Detailed per-signal test NRMSE in native space (expanded view of Table 2 ). DCT3D columns summarize the task-specific output codecs used by the warm-start models: native tensor shape (H,W,T) , embedding dimension Kg , compression ratio dim/Kg with dim=HWT , and explained energy (EE; rounded to 3 decimals). Neural-baseline columns report CNN and CNN+LSTM performance. Best and second-best NRMSE values per row are shown in bold and underlined, respectively.
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 F1=0.837±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
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.
We introduce IGNITE, a generative world foundation model for fusion plasma behavior simulation trained in a self-supervised manner from over a decade of unlabeled experimental data at the DIII-D National Fusion Facility. The core of IGNITE is a dynamics model that can simulate DIII-D discharges from a given set of actuator trajectories. These trajectories can be supplied or generated on-the-fly from a textual prompt or from desired experimental outcomes. The model architecture consists of several spatio-temporal tokenizers that embed the different input modalities, including time-series like spatio-temporal measurement data, image sequences, and high-resolution spectrograms, each of which collected at vastly different time scales. The backbone is composed of an auto-regressive dynamics model that has the capacity to predict entire DIII-D discharges given initial latent plasma states and actuator trajectories over a theoretical infinite horizon. IGNITE paves the way towards efficient AI-driven experimental planning and world modeling for nuclear fusion.
Peter Steiner, Azarakhsh Jalalvand, Nathaniel Chen +4
Department of Mechanical and Aerospace Engineering Princeton University Princeton, NJ · Princeton Plasma Physics Laboratory Princeton, NJ · Princeton University, Princeton Plasma Physics Laboratory Princeton, NJ