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, Bo Lyu, Qiang Liu
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
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
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 → CXRS
CXRS → XCS
MAE ↓
RMSE ↓
R2↑
MAE ↓
RMSE ↓
R2↑
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 R2 values are multiplied by 100.
Method
ne1
ne2
Ti
Te
MAE ↓
RMSE ↓
R2↑
MAE ↓
RMSE ↓
R2↑
MAE ↓
RMSE ↓
R2↑
MAE ↓
RMSE ↓
R2↑
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 Ti and Te metrics are calculated by aggregating all test samples and all 32 radial positions. All R2 values are multiplied by 100.
Method
Exact- K↑
F1 ↑
Weak Recall ↑
Center MAE ↓
Recon. R2↑
NRMSE ↓
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 R2 values are multiplied by 100.
Table 6
Labeled Data
XCS → CXRS R2
Parameter Mean R2
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 → CXRS R2↑
CXRS → XCS R2↑
Parameter Mean R2↑
5×10
97.20
94.92
92.88
7×8
97.11
95.05
92.53
7×10
96.55
94.82
93.24
5×14
97.16
95.30
94.22
7×14
96.89
95.15
92.99
TABLE VII: Effect of XCS spectral patch size on representative downstream tasks. All R2 values are multiplied by 100.
CXRS Patch
XCS → CXRS R2↑
CXRS → XCS R2↑
Parameter Mean R2↑
6×8
96.80
95.21
93.27
3×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 → CXRS R2↑
Parameter Mean R2↑
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 → CXRS
CXRS → XCS
MAE ↓
RMSE ↓
R2↑
MAE ↓
RMSE ↓
R2↑
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 R2 values are multiplied by 100.
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.
Fig. 6: Relative-error distributions of predicted plasma temperatures: (a) electron temperature Te and (b) ion temperature Ti . Each subfigure shows the error distribution as a function of the normalized radius ρ , together with the within- ±10% coverage curve and the marginal error distribution.
Fig. 7: Parity plots for electron-density prediction: (a) ne1 and (b) ne2 . The solid markers represent model predictions against the corresponding ground-truth values, and the dashed diagonal denotes ideal agreement.
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
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.
Auroral spectrographs such as the Auroral Spectrograph In Skibotn (ASIS) record hundreds of thousands of emission spectra, but only a few hundred can be labelled by an expert. To exploit the rest, we pretrain a 1D Vision Transformer with a masked autoencoder on 223,000 unlabelled spectra. Without labels, its representation recovers the emission-line intensity ratios that physicists use to diagnose the precipitating particles (R^2 0.91 vs. 0.77 for an untrained control) and, under one linear probe, classifies as well as 13 features designed by experts. Fine-tuned, the model outperforms the previous supervised auroral classifier on its own benchmark (macro-AP 88.5 vs. 77.8), reaches 0.870 mAP, and exceeds the same architecture trained from scratch by +0.159 with 10% of the labels; attribution shows that it uses both N2+ bands. Could an existing pretrained model replace it? Two astronomical spectral foundation models and a time-series model transfer according to their spectral window: SpectraFM, trained in the infrared, falls below the untrained control, whereas SpecFormer, trained in the optical, approaches in-domain pretraining without reaching it.
Matthieu Le Lain, Gaël Cessateur, Sébastien Lefèvre
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