A Pre-trained Variational Autoencoder-based Surrogate Model for Plasma Turbulence

Plasma turbulence represents a critical challenge in controlled nuclear fusion, causing energy losses and reactor component damage that hinder fusion efficiency. Traditional numerical simulations of plasma dynamics are computationally prohibitive, limiting real-time prediction capabilities essential for fusion reactor control. We present PreVAE-Turb, an advanced machine learning surrogate model using pre-trained variational autoencoders (VAEs) for plasma turbulence prediction. Our approach builds upon the successful Generative Artificial Intelligence Turbulence(GAIT) model [1,2] by implementing a pre-training and fine-tuning strategy using the AutoencoderKL from Stable Diffusion v1.5 [3], coupled with Long Short-Term Memory (LSTM) networks for temporal evolution modeling. The key innovation lies in repurposing a pre-trained convolutional variational autoencoder, originally trained on large-scale image datasets, to encode multi-channel turbulence data into structured latent representations. This approach demonstrates superior reconstruction quality and faster convergence compared to training convolutional neural networks (CNNs) from scratch, while maintaining the physical interpretability of plasma dynamics. We validate our model on two distinct turbulence systems: the Hasegawa-Wakatani equations and data from the gyrokinetic codes GENE and GENE-X, representing a substantial advancement toward realistic plasma turbulence modeling. The model achieves significant speedup while preserving critical statistical properties including Fourier power spectra, proper orthogonal decomposition, and time auto-correlation. The pre-trained VAE offers notable advantages in efficiency and performance. This work establishes a foundation for scalable surrogate models enabling real-time plasma control and long-term transport predictions essential for fusion reactor optimization.