Toward a Multimodal Foundation Model for Plasma State Reconstruction
Authors: Bin Xia, Chuanfei Dong, Boting Li, Yilan Qin, Zikang Xie, Sung Hyun Son, Adam J. Stanier, Jongsoo Yoo, Hantao Ji, Ahmed Diallo, Han Ding, John Wise
Organizations: School of Physics, Georgia Institute of Technology, Atlanta, GA 30332, USA · Department of Astronomy, Department of Electrical and Computer Engineering, Boston University, Boston, MA 02215, USA · Princeton Plasma Physics Laboratory, Princeton, NJ 08540, USA · Department of Electrical and Computer Engineering, Boston University, Boston, MA 02215, USA · Department of Astronomy, Boston University, Boston, MA 02215, USA · Department of Astrophysical Sciences, Princeton University, Princeton, NJ 08544, USA · Los Alamos National Laboratory, Los Alamos, NM 87545, USA · Department of Mechanical Engineering, Carnegie Mellon University, Pittsburgh, PA 15213, USA
Plasma experiments observe coupled physical states through sparse, heterogeneous, and incomplete diagnostics. We present a proof of concept toward a multimodal foundation model for plasma state reconstruction during magnetic reconnection, motivated by the Facility for Laboratory Reconnection Experiments (FLARE). Using 250 synthetic VPIC simulations at fixed reference electron beta βe0=0.2, with 225 training runs and 25 disjoint validation runs, we train a 5.5M-parameter residual 3D U-Net to reconstruct three magnetic-field components and electron density from incomplete observations. The model combines spatial-only pooling, which preserves all 24 temporal frames, with global spatiotemporal self-attention and rotary positional encoding. Independently masking magnetic and density modalities across random, grid, block, and temporal patterns enables sparse sensing, cross-modal completion, density super-resolution, and temporal completion with a single set of weights and no task-specific retraining. With density fully hidden, just 0.1% magnetic spatial visibility reduces median density NRMSE from approximately 0.9 to 0.1. Conversely, observed density improves current-density reconstruction when magnetic data are absent or extremely sparse. With no magnetic observations, roughly 1% density visibility approaches the density reconstruction accuracy achieved with fully observed magnetic fields, although current structure remains more dependent on direct magnetic measurements. Full magnetic conditioning also stabilizes completion of missing future density frames. These results demonstrate bidirectional cross-modal predictive coupling within this simulation regime. Generalization across plasma regimes or to experimental measurements remains untested.
Figures & tables
Task
Overall MSE
Density NRMSE
Jy NRMSE
Spatial random
0.0025 [0.0018, 0.0049]
0.0456 [0.0333, 0.0666]
0.3173 [0.2535, 0.4796]
Spatial grid
0.0033 [0.0023, 0.0062]
0.0581 [0.0460, 0.0785]
0.1931 [0.1607, 0.2630]
Spatial block
0.0287 [0.0117, 0.0589]
0.0933 [0.0526, 0.1426]
0.3724 [0.1888, 0.4926]
Temporal random
0.0044 [0.0028, 0.0101]
0.0405 [0.0288, 0.0613]
0.1916 [0.1628, 0.2666]
Temporal block
0.0217 [0.0066, 0.0576]
0.0677 [0.0376, 0.1462]
0.2385 [0.1702, 0.4286]
Table 1: Standardized approximately 50%-visible evaluation. Entries are median [16th, 84th percentile] across 25 held-out runs.
Figure 1: Spatial density reconstruction with the magnetic field fully observed and density approximately 50% masked. Rows show random sampling, a regular checkerboard, block inpainting, and block outpainting. Prediction panels are raw model outputs rather than hard-clamped fields.
Figure 2: Paired magnetic ablation with nested magnetic probe sets. Blue curves use fully hidden density; orange curves use fully visible density with the identical magnetic masks. Left: density NRMSE. Right: derived Jy NRMSE. Bands denote the 16th–84th percentile across held-out runs.
Figure 3: Density super-resolution as a function of exact density probe count, shown as spatial visibility. Fully observed magnetic fields (blue) strongly constrain density even without density probes; when B is absent (orange), about 1% density visibility is sufficient to approach the density error floor, while Jy remains more dependent on magnetic observations.
Figure 4: Conditional temporal completion within a 24-frame context. Colors denote the missing-suffix length; solid and dashed curves correspond to fully observed and fully hidden magnetic fields. Full magnetic context stabilizes future density reconstruction, while B-hidden Jy error increases with completion depth.
Appendix figures & tables4 assets
Supplementary material from the paper’s appendix.
Appendix
Quantity
Definition
Units / normalization
Values
βe0
Reference electron beta, 2μ0n0Te0/B02
dimensionless
0.2
ν
Multiplier on the Coulomb collision rate evaluated at initial n0,Te0
dimensionless
0, 1, 2, 3, 5, 6, 7, 8, 9, 10
Bz0
Uniform axial field added to the coil vacuum field; reconnecting-plane component, not a guide field
B0
0, 0.1, 0.15, 0.2, 0.3
τ
Flux-core-current drive time
ωci0−1
40, 70, 100, 150, 200
dt
Output cadence Δout
ωci0−1
≃2
Grid
Processed spatial grid
–
154×62
Appendix
Table 2: Physical meaning and coverage of the VPIC parameters used for training and evaluation. The metadata field dt is an output cadence, not the PIC integration timestep.
Item
Version / setting
System
NERSC Perlmutter, A100-SXM4-80GB GPU nodes
NERSC module
pytorch/2.13.0
Python
3.12.13
PyTorch
2.13.0+cu130
CUDA build
13.0
NCCL
2.29.7
Appendix
Table 3: Software and distributed-training environment used for the final run.
Table 8
Figure 5: Sliding-window extension to 48-frame sequences with 1,000 fixed density probes. Left: median frame RMSE across 18 held-out runs for different slide steps and magnetic conditions. Middle/right: predictions and residuals for a representative B-hidden trajectory.
Reconstructing spatially resolved plasma dynamics from few sensors is essential for diagnostics, reduced-order modelling and control, yet remains difficult because the sparse measurements incompletely constrain multiscale, regime-dependent degrees of freedom. The Shallow Recurrent Decoder (SHRED) partially addresses spatial sparsity by using measurement histories; however, its fully connected decoder provides no explicit mechanism for resolving spatial structure across scales or explicit parametric dependency. We introduce the Recurrent Multiscale Affine-modulated Inference Network (ReMAIN), which preserves SHRED's recurrent temporal encoding but replaces its decoder with a U-Net whose feature hierarchy is conditioned by the recurrent state through feature-wise linear modulation. The temporal representation supplies both a dense prior and scale-specific modulation throughout the U-Net. A parametric extension jointly embeds the operating condition and sensor history, enabling reconstruction to adapt as the governing dynamics change with operating regime. ReMAIN is first benchmarked against SHRED on six one-dimensional nonlinear PDEs representing diverse dynamics. Across all benchmarks, it reduces reconstruction errors on unseen trajectories and more faithfully resolves sharp transitions, localized extrema and fine-scale variations. The parameter-conditioned model is then demonstrated on a collisionless E×B plasma subject to perpendicular axial electric and radial magnetic fields, with the electric-field strength serving as the operating parameter. ReMAIN reconstructs the high-dimensional, multiscale plasma state and recovers its regime-dependent spatiotemporal dynamics at electric-field strengths withheld from training. Together, ReMAIN improves sparse-sensor full-state reconstruction and, through parameter conditioning, generalizes across plasma operating regimes.
Maryam Reza, Farbod Faraji
Independent Researcher · Department of Computing, Huxley Building, Imperial College London, London SW7 2RH, United Kingdom
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
This work presents a new bidirectional autoregressive latent diffusion approach for predicting the evolution of multiple fields (mass density, pressure, velocity, and magnetic field components) for magnetohydrodynamics. We show that this bidirectional flow can be used as a self-supervised consistency metric for uncertainty and error estimation, which enables the model to estimate test-time uncertainty and error without access to ground truth, by comparing how closely flowing forwards and backwards in time returns to the same predicted fields. We also demonstrate this methods's potential to serve as a non-invasive plasma diagnostic, and show how adaptive feedback can be used to make the model more robust based on sparse diagnostics or limited views/measurements.
Alexander Scheinker
Los Alamos National Laboratory, Los Alamos, NM, 87545.