TokaMind

TokaMind is a multi-modal, token-based Transformer foundation model for tokamak plasma dynamics, trained on data from the MAST (Mega Ampere Spherical Tokamak) facility.

This repository contains the following pretrained checkpoints:

Variant	Parameters	Folder
TokaMind-Base	9.32M	tokamind-base-v1/
TokaMind-Tiny	5.29M	tokamind-tiny-v1/

Model Description

TokaMind implements a schema-flexible tokenization pipeline and a modular multi-modal Transformer with per-output adapters. Windowed multimodal inputs and actuators are chunked and compressed by signal-specific codecs into tokens. Tokens are projected to a shared model dimension, processed by a Transformer backbone, and mapped to targets via modality heads and per-output adapters.

The model is introduced in the paper TokaMind: A Multi-Modal Transformer Foundation Model for Tokamak Plasma Dynamics (submitted to the 32nd SIGKDD Conference on Knowledge Discovery and Data Mining, 2026), and evaluated against the TokaMark benchmark.

Intended Use

• Warmstarting finetunes on TokaMark benchmark tasks
• Transfer learning for tokamak plasma signal prediction
• Research into multi-modal Transformer architectures for scientific time-series

Training Data

Pretrained on MAST tokamak data, accessed via the TokaMark dataset. The dataset covers multimodal plasma signals including magnetics, profiles, and actuators across thousands of shots.

How to Use

Each checkpoint folder contains:

• Model weights (checkpoints/best)
• Embedding artifacts (embeddings/dct3d.yaml, embeddings/dct3d_indices/)
• Config snapshot used for pretraining

Download a variant and place it under runs/ in your TokaMind installation:

runs/
└── tokamind_base/
    ├── tokamind_base.yaml
    ├── checkpoints/
    │   └── best
    └── embeddings/
        ├── dct3d.yaml
        └── dct3d_indices/

Then warmstart a finetune:

python scripts_mast/run_finetune.py \
  --task task_2-1 \
  --init warmstart \
  --model tokamind_base \
  --emb_profile dct3d \
  --tag exp1

See the TokaMind repository for full installation and usage instructions.

Companion Resources

Resource	Link
TokaMind paper	arXiv:2602.15084
TokaMark paper	arXiv:2602.10132
TokaMind repository	UKAEA-IBM-STFC-Fusion-FMs/tokamind
TokaMark repository	UKAEA-IBM-STFC-Fusion-FMs/tokamark

Citing TokaMind

If you use TokaMind, please cite:

@article{boschi2026tokamind,
  title={TokaMind: A Multi-Modal Transformer Foundation Model for Tokamak Plasma Dynamics},
  author={
    Boschi, Tobia and Loreti, Andrea and Amorisco, Nicola C and Ordonez-Hurtado, Rodrigo H and
    Rousseau, C{\'e}cile and Holt, George K and Sz{\'e}kely, Eszter and Whittle, Alexander and
    Jackson, Samuel and Agnello, Adriano and Pamela, Stanislas and Pascale, Alessandra and
    Akers, Robert and Bernabe Moreno, Juan and Thorne, Sue and Zayats, Mykhaylo
  },
  journal={arXiv preprint arXiv:2602.15084},
  year={2026}
}

If you also use the TokaMark benchmark, please cite:

@article{rousseau2026tokamark,
  title={TokaMark: A Comprehensive Benchmark for MAST Tokamak Plasma Models},
  author={
    Rousseau, C{\'e}cile and Jackson, Samuel and Ordonez-Hurtado, Rodrigo H. and
    Amorisco, Nicola C. and Boschi, Tobia and Holt, George K and Loreti, Andrea and
    Sz{\'e}kely, Eszter and Whittle, Alexander and Agnello, Adriano and Pamela, Stanislas and
    Pascale, Alessandra and Akers, Robert and Bernabe Moreno, Juan and Thorne, Sue and
    Zayats, Mykhaylo
  },
  journal={arXiv preprint arXiv:2602.10132},
  year={2026}
}