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

TL;DR

TokaMind is an open-source multi-modal transformer model for fusion plasma modeling that effectively integrates heterogeneous diagnostics, outperforming benchmarks and enabling efficient task adaptation in tokamak research.

Contribution

The paper introduces TokaMind, a novel multi-modal transformer framework for tokamak plasma data, with new multi-modal embedding techniques and superior performance on the MAST benchmark.

Findings

Fine-tuned TokaMind outperforms baseline on most tasks.

Lightweight fine-tuning can surpass training from scratch.

Multi-modal pretraining benefits tokamak plasma modeling.

Abstract

We present TokaMind, an open-source foundation model framework for fusion plasma modeling, based on a Multi-Modal Transformer (MMT) and trained on heterogeneous tokamak 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 training-free Discrete Cosine Transform embedding (DCT3D) and provide a clean interface for alternative embeddings (e.g., Variational Autoencoders - VAEs). We evaluate TokaMind on the recently introduced MAST benchmark TokaMark, comparing training and embedding strategies. Our results show that fine-tuned TokaMind outperforms the benchmark baseline on all but one task, and that, for several tasks, lightweight fine-tuning yields better performance than training the same architecture from scratch under a matched epoch budget. These findings highlight the benefits of multi-modal pretraining for tokamak plasma dynamics and provide a practical, extensible foundation for future fusion modeling tasks. Training code and model weights will be made publicly available.