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

Tobia Boschi; Andrea Loreti; Nicola C. Amorisco; Rodrigo H. Ordonez-Hurtado; Cécile Rousseau; George K. Holt; Eszter Székely; Alexander Whittle

physics.plasm-phFeb 16, 2026

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,

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

Abstract

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

\mathcal{I}

and actuators

\mathcal{A}

are chunked and embedded by signal-specific codecs

E_{g}

to produce token embeddings

z_{i}\in\mathbb{R}^{K_{g(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

\hat{y}_{o}\in\mathbb{R}^{K_{o}}

; a target-availability mask

m_{o}

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	✓	Token Encoder / Output Adapter reinitialized (shape mismatch); other blocks reused
Change chunk length / max_chunks / stride	✓	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	$\Delta$	NRMSE	$\Delta$
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.

\Delta

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
		$\bm{K_{g}}$	dim/ $K_{g}$	EE	$\bm{K_{g}}$	dim/ $K_{g}$
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

\mathrm{dim}=HWT

. DCT3D columns report the source codecs used by warm-start models: embedding dimension

K_{g}

, compression ratio

\mathrm{dim}/K_{g}

, and explained energy (EE). VAE columns report the latent embedding dimension

K_{g}

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
			$\bm{(H,W,T)}$	$\bm{K_{g}}$	dim/ $K_{g}$	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

K_{g}

, compression ratio

\mathrm{dim}/K_{g}

with

\mathrm{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.

Multimodal Foundation Models

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

See all 349 papers

TokaMind for Power Grid: Cross-Domain Transfer from Fusion Plasma
May 10, 2026JC Wu, Norton Lee, Kai Siang ChenTime Series Classification Critical Slowing Down
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
Benchmarking Sensor Robustness in Plasma Diagnostic Models: A Systematic Evaluation on TokaMark
Jul 5, 2026Neerav GuptaIncomplete Data Imputation
Independent Researcher
IGNITE Tokamak World Model Architecture
Oct 1, 2026Peter Steiner, Azarakhsh Jalalvand, Nathaniel Chen +4World Model Learning Physics-Based Simulation
Department of Mechanical and Aerospace Engineering Princeton University Princeton, NJ · Princeton Plasma Physics Laboratory Princeton, NJ · Princeton University, Princeton Plasma Physics Laboratory Princeton, NJ

May 10, 2026J/K move · Enter open · S save

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

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

Open paper AskPDF arXiv

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.

Abstract

Multimodal Foundation Models

Explore similar work

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

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

IGNITE Tokamak World Model Architecture