Toward a Multimodal Foundation Model for Plasma State Reconstruction

Bin Xia; Chuanfei Dong; Boting Li; Yilan Qin; Zikang Xie; Sung Hyun Son; Adam J. Stanier; Jongsoo Yoo

cs.LGOct 10, 2026

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,

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 ·

Abstract

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	$J_{y}$ 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 $J_{y}$ 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 $\mathbf{B}$ is absent (orange), about 1% density visibility is sufficient to approach the density error floor, while $J_{y}$ 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 $J_{y}$ error increases with completion depth.

Appendix figures & tables4 assets

Supplementary material from the paper’s appendix.

Appendix

Quantity	Definition	Units / normalization	Values
$\beta_{e0}$	Reference electron beta, $2\mu_{0}n_{0}T_{e0}/B_{0}^{2}$	dimensionless	0.2
$\nu$	Multiplier on the Coulomb collision rate evaluated at initial $n_{0},T_{e0}$	dimensionless	0, 1, 2, 3, 5, 6, 7, 8, 9, 10
$B_{z0}$	Uniform axial field added to the coil vacuum field; reconnecting-plane component, not a guide field	$B_{0}$	0, 0.1, 0.15, 0.2, 0.3
$\tau$	Flux-core-current drive time	$\omega_{ci0}^{-1}$	40, 70, 100, 150, 200
dt	Output cadence $\Delta_{\rm out}$	$\omega_{ci0}^{-1}$	$\simeq 2$
Grid	Processed spatial grid	–	$154\times 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.

Multimodal Foundation Models

Multimodal foundation models learn shared capabilities across text, images, audio, video, and other signals for understanding, generation, and action.

See all 353 papers

Reconstruction of Multiscale Plasma Dynamics Across Operating Regimes
Oct 7, 2026Maryam Reza, Farbod FarajiRecurrent Neural Networks Multiscale Time-Series Modeling
Independent Researcher · Department of Computing, Huxley Building, Imperial College London, London SW7 2RH, United Kingdom
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
Bidirectional Autoregressive Latent Diffusion for Forward and Inverse Magnetohydrodynamics
Jun 28, 2026Alexander ScheinkerUncertainty Quantification Autoregressive Diffusion
Los Alamos National Laboratory, Los Alamos, NM, 87545.

Oct 7, 2026J/K move · Enter open · S save

Reconstruction of Multiscale Plasma Dynamics Across Operating Regimes

Maryam Reza, Farbod Faraji

Independent Researcher · Department of Computing, Huxley Building, Imperial College London, London SW7 2RH, United Kingdom

Recurrent Neural Networks Multiscale Time-Series Modeling Multi-Scale Feature Fusion

Open paper AskPDF arXiv

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 \times 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.

Abstract

Multimodal Foundation Models

Explore similar work

Reconstruction of Multiscale Plasma Dynamics Across Operating Regimes

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

Bidirectional Autoregressive Latent Diffusion for Forward and Inverse Magnetohydrodynamics