FusionMAE: large-scale pretrained model to optimize and simplify diagnostic and control of fusion plasma

Abstract: In magnetically confined fusion device, the complex, multiscale, and nonlinear dynamics of plasmas necessitate the integration of extensive diagnostic systems to effectively monitor and control plasma behaviour. The complexity and uncertainty arising from these extensive systems and their tangled interrelations has long posed a significant obstacle to the acceleration of fusion energy development. In this work, a large-scale model, fusion masked auto-encoder (FusionMAE) is pre-trained to compress the information from 88 diagnostic signals into a concrete embedding, to provide a unified interface between diagnostic systems and control actuators. Two mechanisms are proposed to ensure a meaningful embedding: compression-reduction and missing-signal reconstruction. Upon completion of pre-training, the model acquires the capability for 'virtual backup diagnosis', enabling the inference of missing diagnostic data with 96.7% reliability. Furthermore, the model demonstrates three emergent capabilities: automatic data analysis, universal control-diagnosis interface, and enhancement of control performance on multiple tasks. This work pioneers large-scale AI model integration in fusion energy, demonstrating how pre-trained embeddings can simplify the system interface, reducing necessary diagnostic systems and optimize operation performance for future fusion reactors.

• The paper introduces FusionMAE, a Transformer-based masked autoencoder that unifies 88 tokamak plasma diagnostics through self-supervised training and random masking.
• It achieves near-lossless signal reconstruction with a mean PCC of 98.6% and robust imputation on missing diagnostic channels.
• The unified 256-dimensional plasma state embedding enhances downstream tasks like disruption prediction, equilibrium reconstruction, and plasma evolution.

FusionMAE: Large-Scale Pretrained Model for Fusion Plasma Diagnostics and Control

Motivation and Context

The engineering complexity of magnetic confinement fusion reactors, particularly tokamaks, is driven by the need to integrate a multitude of diagnostic systems and control actuators to monitor and regulate the nonlinear, multiscale plasma dynamics. Existing approaches rely on a fragmented ecosystem of diagnostic subsystems, simulation modules, and control loops, resulting in operational fragility and high system overhead. While AI has demonstrated efficacy in specific fusion tasks (e.g., disruption prediction, equilibrium reconstruction), these solutions are typically narrow in scope and do not address the overarching challenge of system integration and robustness. FusionMAE is introduced as a domain-specific, large-scale foundational model that unifies diagnostic data into a compact, self-consistent plasma state embedding, enabling robust, simplified interfaces for both diagnostics and control.

Model Architecture and Training

FusionMAE is a Transformer-based masked autoencoder designed to process 88 diagnostic signals from the HL-3 tokamak, sampled at 1 kHz over a 10 ms window (88×10 input array). Each channel is projected into a 64-dimensional latent space via an MLP, with positional encoding added to preserve channel identity. The encoder comprises multiple Transformer blocks (multi-head self-attention, MLPs, batch normalization, residual connections), outputting a 5632-dimensional vector, which is compressed to a 256-dimensional plasma status embedding. The decoder mirrors the encoder structure, reconstructing the original signals from the embedding.

A key innovation is the use of random masking: 25% of input channels are masked during training, forcing the model to learn inter-channel correlations and reconstruct missing data. The loss function is a weighted MSE, prioritizing reconstruction of intentionally masked channels while preserving valid data. Training employs AdamW with scheduled learning rate and weight decay, and Gaussian noise augmentation to enhance robustness.

Emergent Capabilities

High-Fidelity Compression and Imputation

FusionMAE achieves near-lossless compression, with a mean Pearson correlation coefficient (PCC) of 98.6% between reconstructed and original signals. For masked channels, the model attains a PCC of 96.7%, demonstrating reliable imputation of missing diagnostics. The plasma status embedding remains stable under partial input dropout (5–20% missing), with only minor oscillations, indicating effective redundancy exploitation and resilience to diagnostic failure.

Unified Plasma State Embedding

The 256-dimensional embedding encapsulates the intrinsic physical state of the plasma, analogous to word embeddings in NLP. UMAP projections reveal clustering of embeddings along key physical parameters (Bt, Ip, ne, BN), confirming that the model autonomously organizes latent space according to underlying plasma physics. The embedding is robust to reductions in input diagnostics, supporting the notion of a universal, all-purpose vector for downstream tasks.

Automated Data Analysis

FusionMAE can reconstruct secondary parameters (e.g., plasma shape, position) that are typically derived from analysis codes such as EFIT, even when all such channels are masked. This demonstrates the model's capacity to infer physical relationships and approximate traditional analysis tools, streamlining data processing pipelines.

Downstream Task Performance

Three representative downstream tasks are evaluated:

• Disruption Prediction (DPR): Using the plasma status embedding as input yields AUC scores equal to or exceeding those obtained with raw data. When missing signals are imputed by FusionMAE, performance is further improved.
• Equilibrium Reconstruction (EFIT-NN): Models trained on embeddings or completed data outperform those using raw diagnostics, as measured by PCC.
• Plasma Evolution Prediction (PPR): Embedding-based models achieve higher PCC, and robustness to missing data is markedly superior compared to baseline approaches.

Performance degradation with increasing masked input is significantly slower for FusionMAE-supported algorithms, establishing a robust operational framework for future reactors.

Quantitative Analysis of Diagnostic Redundancy

FusionMAE provides a principled method to quantify the contribution of each diagnostic channel to plasma state estimation. By measuring the change in embedding similarity when individual channels are masked, the model identifies critical diagnostics (e.g., Bp, Te, ne, Bt) and quantifies redundancy. This enables empirical guidelines for minimum viable diagnostic configurations, balancing measurement robustness against system complexity.

Implications and Future Directions

FusionMAE demonstrates that large-scale, self-supervised pretraining architectures (inspired by BERT, GPT, and MAE) are highly effective in fusion data domains, transcending the limitations of small-model, supervised approaches. The unified plasma state embedding simplifies system interfaces, enhances downstream control and analysis, and provides resilience against diagnostic failure—a critical requirement for future fusion power plants.

Theoretical implications include the validation of latent manifold representations for complex physical systems and the feasibility of foundational models for scientific domains beyond NLP and vision. Practically, FusionMAE paves the way for reduced diagnostic system overhead, improved control precision, and accelerated fusion reactor development.

Future research should explore:

• Integration of reinforcement learning from human feedback for adaptive control.
• Mixture-of-experts architectures to scale model capacity and specialization.
• Chain-of-thought reasoning for interpretability and physical insight.
• Transfer learning across devices and campaigns to generalize foundational models.

Conclusion

FusionMAE establishes a new paradigm for fusion plasma diagnostics and control by leveraging large-scale, self-supervised Transformer architectures to unify, compress, and robustly reconstruct multi-source diagnostic data. Its emergent capabilities—high-fidelity imputation, universal plasma state embedding, automated data analysis, and superior downstream task performance—address critical challenges in fusion reactor engineering. The model's quantitative framework for diagnostic redundancy and its demonstrated resilience to missing data provide actionable insights for future reactor design. FusionMAE exemplifies the potential of foundational AI models to accelerate progress toward practical fusion energy solutions, with broad implications for scientific AI integration.