SpectralTF-MAE: A Pre-trained Spectral Foundation Model for Tokamak Fusion Plasma Diagnostics

Xiao Wang; Zhijie Lu; Duoqing Yang; Jiangtao Yu; Jin Tang; Zichao Lin; Hongming Zhang; Guosheng Xu

cs.CVOct 28, 2026

SpectralTF-MAE: A Pre-trained Spectral Foundation Model for Tokamak Fusion Plasma Diagnostics

Authors: Xiao Wang, Zhijie Lu, Duoqing Yang, Jiangtao Yu, Jin Tang, Zichao Lin, Hongming Zhang, Guosheng Xu,

Organizations: School of Computer Science and Technology, Anhui University, Hefei 230601, China · Institute of Plasma Physics, Hefei Institutes of Physical Science, Chinese Academy of Sciences, Hefei 230031, China · New Laboratory of Pattern Recognition, Institute of Automation, Chinese Academy of Sciences, Beijing, China

Abstract

Spectroscopic diagnostics provide rich information for characterizing plasma states in tokamak fusion experiments, while existing learning-based approaches are often designed for individual tasks and rely heavily on labeled data. In this work, we propose SpectralTF-MAE, a self-supervised spectral foundation model for tokamak plasma diagnostics. The model is pretrained on paired charge-exchange recombination spectroscopy (CXRS) and X-ray crystal spectroscopy (XCS) measurements through masked spectral reconstruction. To better capture diagnostically important spectral structures, we introduce a peak-aware masking strategy together with a peak-enhanced reconstruction objective, encouraging the model to focus on characteristic spectral regions. A shared Transformer encoder is employed to learn transferable representations across different diagnostic modalities. The pretrained model is further adapted to several downstream tasks, including plasma-parameter prediction, cross-diagnostic spectral generation, and spectral peak analysis. Experiments on EAST diagnostic data demonstrate that the proposed framework achieves competitive performance across different tasks and consistently benefits from self-supervised spectral pretraining. These results indicate the potential of SpectralTF-MAE as a general and transferable representation learning framework for plasma spectroscopic diagnostics.

Figures & tables

Panel of Figure 1 — Fig. 1: (a-b). Demonstration of different principles for spectral - data acquisition; (c-f). Performance comparison of downstream tasks based on our newly proposed spectral foundation model.

Fig. 2: An overview of our proposed pre-trained spectral foundation model (upper panel) and downstream tasks (lower panel) for tokamak fusion plasma diagnostics.

Method	XCS $\rightarrow$ CXRS			CXRS $\rightarrow$ XCS
Method	MAE $\downarrow$	RMSE $\downarrow$	$R^{2}$ $\uparrow$	MAE $\downarrow$	RMSE $\downarrow$	$R^{2}$ $\uparrow$
ResNet-50 [ 37 ]	0.00683	0.01398	96.90	0.02866	0.04759	90.04
ResNet-101 [ 37 ]	0.00677	0.01403	96.88	0.02880	0.04912	89.39
ViT-Small [ 12 ]	0.01540	0.02625	89.08	0.06026	0.11271	44.13
ViT-Base [ 12 ]	0.01552	0.02655	88.83	0.06012	0.11262	44.21
ViT-Large [ 12 ]	0.01522	0.02603	89.26	0.05993	0.11254	44.29

TABLE I: Comparison of methods on bidirectional cross-diagnostic spectral generation. For XCS-to-CXRS generation, MAE and RMSE are calculated in the normalized CXRS characteristic peak region; for CXRS-to-XCS generation, they are calculated in the normalized XCS characteristic peak region. All

R^{2}

values are multiplied by 100.

Method	$n_{e1}$			$n_{e2}$			$T_{i}$			$T_{e}$
Method	MAE $\downarrow$	RMSE $\downarrow$	$R^{2}$ $\uparrow$	MAE $\downarrow$	RMSE $\downarrow$	$R^{2}$ $\uparrow$	MAE $\downarrow$	RMSE $\downarrow$	$R^{2}$ $\uparrow$	MAE $\downarrow$	RMSE $\downarrow$	$R^{2}$ $\uparrow$
ResNet-50 [ 37 ]	0.29697	0.41943	77.18	0.16697	0.26488	89.15	44.38147	66.45265	97.18	98.92274	165.28667	97.86
ResNet-101 [ 37 ]	0.29142	0.42682	76.37	0.18265	0.30272	85.83	48.54834	73.08930	96.59	110.23382	185.75727	97.30
ViT-Small [ 12 ]	0.26455	0.37971	81.30	0.15322	0.25030	90.31	42.46420	65.57747	97.26	86.27301	142.69225	98.41
ViT-Base [ 12 ]	0.24610	0.34939	84.17	0.12798	0.20237	93.67	39.81028	61.63069	97.58	76.46207	125.69500	98.77
ViT-Large [ 12 ]	0.22526	0.34130	84.89	0.12143	0.20529	93.48	38.71456	59.66490	97.73	73.02069	130.84304	98.66

TABLE II: Comparison of methods on plasma-parameter prediction from XCS measurements. The

T_{i}

and

T_{e}

metrics are calculated by aggregating all test samples and all 32 radial positions. All

R^{2}

values are multiplied by 100.

Method	Exact- $K$ $\uparrow$	F1 $\uparrow$	Weak Recall $\uparrow$	Center MAE $\downarrow$	Recon. $R^{2}$ $\uparrow$	NRMSE $\downarrow$
Task-specific Baseline	98.88	77.81	22.22	0.9631	98.54	0.02579
SpectralTF-MAE-Base (Random)	98.92	80.67	19.05	0.9646	99.37	0.01584
SpectralTF-MAE-Base (Pretrained)	98.92	83.64	23.81	0.8174	99.22	0.01767

TABLE III: Comparison on spectral peak analysis. Exact-

K

, F1, and weak-peak recall are reported in percentage. Reconstruction

R^{2}

values are multiplied by 100.

Table 6

Labeled Data	XCS $\rightarrow$ CXRS $R^{2}$		Parameter Mean $R^{2}$
Labeled Data	Scratch	Pretrained	Scratch	Pretrained
10%	84.88	95.94	58.72	69.60
25%	94.33	96.90	70.95	81.38
50%	95.56	97.03	79.59	89.23
100%	95.70	97.16	86.03	94.22

TABLE VI: Data efficiency of spectral pretraining on representative downstream tasks.

XCS Patch	XCS $\rightarrow$ CXRS $R^{2}$ $\uparrow$	CXRS $\rightarrow$ XCS $R^{2}$ $\uparrow$	Parameter Mean $R^{2}$ $\uparrow$
$5\times 10$	97.20	94.92	92.88
$7\times 8$	97.11	95.05	92.53
$7\times 10$	96.55	94.82	93.24
$5\times 14$	97.16	95.30	94.22
$7\times 14$	96.89	95.15	92.99

TABLE VII: Effect of XCS spectral patch size on representative downstream tasks. All

R^{2}

values are multiplied by 100.

CXRS Patch	XCS $\rightarrow$ CXRS $R^{2}$ $\uparrow$	CXRS $\rightarrow$ XCS $R^{2}$ $\uparrow$	Parameter Mean $R^{2}$ $\uparrow$
$6\times 8$	96.80	95.21	93.27
$3\times 16$	97.16	95.30	94.22

TABLE VIII: Effect of CXRS patch size on downstream transfer. The XCS patch size is fixed to the configuration selected from the XCS patch study.

Mask Ratio	XCS $\rightarrow$ CXRS $R^{2}$ $\uparrow$	Parameter Mean $R^{2}$ $\uparrow$
60%	97.04	91.72
70%	97.18	92.59
80%	97.16	94.22
90%	97.12	92.28

TABLE IX: Sensitivity to the global masking ratio during spectral pretraining.

Pretraining Aug.	XCS $\rightarrow$ CXRS			CXRS $\rightarrow$ XCS
Pretraining Aug.	MAE $\downarrow$	RMSE $\downarrow$	$R^{2}$ $\uparrow$	MAE $\downarrow$	RMSE $\downarrow$	$R^{2}$ $\uparrow$
No	0.00809	0.01615	95.86	0.02080	0.03417	94.86
Yes	0.00632	0.01339	97.16	0.01796	0.03269	95.30

TABLE X: Effect of data augmentation during self-supervised pretraining on cross-diagnostic spectral generation. MAE and RMSE are calculated in the normalized target characteristic region, and

R^{2}

values are multiplied by 100.

Panel of Figure 12 — Fig. 3: Visualization of bidirectional spectral generation: (a) XCS-to-CXRS generation and (b) CXRS-to-XCS generation. Each subfigure compares the ground-truth target spectrum, the generated spectrum, the absolute-error map, and representative spectral profiles.

Fig. 4: Visualization of spectral peak decomposition for a representative sample. The observed spectrum and SpectralTF-MAE reconstruction are shown together with the ground-truth peak components and their corresponding predicted components.

Fig. 5: Peak-center localization error comparison between the baseline method and SpectralTF-MAE. Each point denotes a paired sample with the correct peak count in both methods, and the dashed diagonal indicates equal MAE.

Panel of Figure 15 — Fig. 6: Relative-error distributions of predicted plasma temperatures: (a) electron temperature $T_{e}$ and (b) ion temperature $T_{i}$ . Each subfigure shows the error distribution as a function of the normalized radius $\rho$ , together with the within- $\pm 10\%$ coverage curve and the marginal error distribution.

Panel of Figure 16 — Fig. 7: Parity plots for electron-density prediction: (a) $n_{e1}$ and (b) $n_{e2}$ . The solid markers represent model predictions against the corresponding ground-truth values, and the dashed diagonal denotes ideal agreement.

Representation Learning

Representation learning develops data transformations that expose useful structure for prediction, transfer, generation, and control.

See all 813 papers

TokaMind: A Multi-Modal Transformer Foundation Model for Tokamak Plasma Dynamics
Feb 16, 2026Tobia Boschi, Andrea Loreti, Nicola C. Amorisco +13Multimodal Foundation Models Multimodal Pretraining
IBM Research Dublin, Ireland · UK Atomic Energy Authority Abingdon, UK · STFC Hartree Centre Daresbury, UK
Benchmarking Sensor Robustness in Plasma Diagnostic Models: A Systematic Evaluation on TokaMark
Jul 5, 2026Neerav GuptaIncomplete Data Imputation
Independent Researcher
Self-Supervised Representation Learning: From Spectral Foundation Models to Auroral Emission Spectra
Sep 25, 2026Matthieu Le Lain, Gaël Cessateur, Sébastien LefèvreRepresentation Learning Self-Supervised Pre-Training
IRISA, UMR CNRS 6074, Université Bretagne Sud, Vannes, France · Royal Belgian Institute for Space Aeronomy, Brussels, Belgium · UiT - The Arctic University of Norway, Tromsø, Norway

Feb 16, 2026J/K move · Enter open · S save

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

Tobia Boschi, Andrea Loreti, Nicola C. Amorisco, Rodrigo H. Ordonez-Hurtado+12

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

Multimodal Foundation Models Multimodal Pretraining Time Series Forecasting

Open paper AskPDF arXiv

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.

Abstract

Representation Learning

Explore similar work

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

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

Self-Supervised Representation Learning: From Spectral Foundation Models to Auroral Emission Spectra