A Pre-trained Variational Autoencoder for Gyrokinetic Plasma Turbulence Surrogate Modeling

Minglei Yang; Marshall Nicholson; Diego Del-Castillo-Negrete; David Hatch; Guannan Zhang

physics.plasm-phSep 29, 2026

A Pre-trained Variational Autoencoder for Gyrokinetic Plasma Turbulence Surrogate Modeling

Authors: Minglei Yang, Marshall Nicholson, Diego Del-Castillo-Negrete, David Hatch, Guannan Zhang

Organizations: Oak Ridge National Laboratory, Oak Ridge, TN, 37831, USA · Grand Valley State University, Allendale, MI, 49401, USA · University of Texas at Austin, Austin, TX, 78712, USA

Abstract

Machine learning surrogate models offer a promising path toward accelerating plasma turbulence simulations. We present PreVAE-Turb, a surrogate modeling framework that leverages pre-trained variational autoencoders (VAEs) from the Stable Diffusion image generation model for efficient spatial compression of turbulence fields. The pre-trained VAE is fine-tuned on turbulence data using a physics-informed loss function that includes a spectral loss operating in Fourier space to enforce spectral accuracy across scales. The VAE is combined with convolutional long short-term memory (ConvLSTM) networks to learn temporal dynamics in latent space, with a manifold consistency error metric that monitors encode--decode consistency during autoregressive rollouts. We validate the framework on two-dimensional Hasegawa-Wakatani drift-wave turbulence and extend it to gyrokinetic turbulence from the GENE code, where a four-channel adaptation simultaneously predicts electrostatic potential, density, and parallel/perpendicular temperature fluctuations without requiring architecture redesign. Once trained, inference generates thousands of time steps in seconds on a single GPU, providing substantial computational acceleration compared to direct numerical simulation. The pre-trained approach offers a transferable methodology broadly applicable to various turbulence simulation codes.

Figures & tables

Figure 1: Overview of the PreVAE-Turb framework. (a) Input: multi-field 2D turbulence snapshots; the diagram labels the spatial dimensions as $H=N_{x}$ and $W=N_{y}$ . The VAE backbone is initialized from Stable Diffusion v1.5 [ 24 ] , pre-trained on natural images. (b) Two-stage architecture: Stage 1 fine-tunes the pre-trained VAE on turbulence data using a physics-informed loss combining reconstruction, gradient, spectral (upper right), and KL divergence terms. Stage 2 trains a ConvLSTM to predict temporal dynamics in latent space while preserving spatial correlations. The manifold consistency error (lower right) monitors encode–decode consistency of predicted latents. (c) Output: autoregressive prediction generates turbulence fields beyond the training horizon.

Figure 2: Training loss comparison for HW turbulence: pre-trained VAE (AutoencoderKL from Stable Diffusion v1.5, solid lines) vs. CNN-VAE trained from scratch (4-layer encoder/decoder, latent dim 64, dashed lines). (a) Total loss evolution. (b) Individual loss components: reconstruction (density + potential), gradient, and KL divergence. The KL loss shown is the mean per latent dimension. The pre-trained model achieves lower loss across all components, demonstrating greater representational capacity. For HW turbulence,

w_{\text{spec}}=0

Model	Field	MSE $\downarrow$	PSNR (dB) $\uparrow$	SSIM $\uparrow$	Correlation $\uparrow$
CNN-VAE	Density	0.256 $\pm$ 0.035	23.7 $\pm$ 1.1	0.955 $\pm$ 0.007	0.969 $\pm$ 0.005
CNN-VAE	Potential	0.137 $\pm$ 0.026	26.2 $\pm$ 1.3	0.972 $\pm$ 0.006	0.981 $\pm$ 0.004
Pre-trained	Density	0.041 $\pm$ 0.065	32.5 $\pm$ 2.7	0.993 $\pm$ 0.009	0.995 $\pm$ 0.008
Pre-trained	Potential	0.019 $\pm$ 0.050	35.7 $\pm$ 2.7	0.996 $\pm$ 0.007	0.997 $\pm$ 0.007

Table 1: Reconstruction quality metrics (mean

\pm

std) over 8000 HW turbulence snapshots. The pre-trained VAE substantially outperforms the CNN-VAE across all metrics for both density and potential fields.

Figure 3: Statistical validation of PreVAE-Turb on HW turbulence for density $n$ (top row) and potential $\phi$ (bottom row). From left to right: spatial spectrum $S_{f}(k)$ , temporal spectrum $\Omega_{f}(\omega)$ , POD spectrum $\sigma_{\ell}$ , and temporal autocorrelation $C_{f}(\tau)$ . Black solid: HW ground truth. Blue dashed: PreVAE-Turb prediction $t\in[8000,13000]$ . Red dotted: PreVAE-Turb prediction $t\in[18000,23000]$ . PreVAE-Turb accurately reproduces the statistical properties across both fields and different prediction windows.

Figure 4: Cross-field correlations for HW turbulence. Left: $\phi$ - $n$ cross-correlation $C_{\phi n}(\tau)$ showing agreement between HW model (black) and PreVAE-Turb predictions at different time windows (blue dashed: $t\in[8000,13000]$ , red dotted: $t\in[18000,23000])$ . Middle: joint probability distribution $P(\phi,n)$ for HW ground truth. Right: joint probability distribution $P(\phi,n)$ for PreVAE-Turb prediction. PreVAE-Turb preserves the characteristic tilted elliptical structure indicating the phase relationship between potential and density fluctuations.

Figure 5: VAE training loss comparison for GENE turbulence: pre-trained VAE (solid lines) vs. CNN-VAE trained from scratch (dashed lines). Left: total loss evolution. Right: individual loss components (reconstruction, gradient, spectral, and KL divergence). The pre-trained model achieves lower loss across all reconstruction-related terms while maintaining comparable KL regularization.

Figure 6: Snapshot comparison for GENE turbulence fields. Top row: GENE ground truth at $t=1200$ . Middle row: PreVAE-Turb prediction at $t=1200$ . Bottom row: PreVAE-Turb prediction at $t=2500$ (long-horizon extrapolation). Columns from left to right: electrostatic potential $\phi$ , density $n$ , parallel temperature $T_{\parallel}$ , and perpendicular temperature $T_{\perp}$ . PreVAE-Turb accurately reproduces coherent structures at $t=1200$ and maintains stable turbulence patterns at $t=2500$ , demonstrating long-term prediction stability.

Figure 7: Validation metrics for PreVAE-Turb across all four fields (

\phi

n

T_{\parallel}

T_{\perp}

). Columns: spatial spectrum

S_{f}(k)

, temporal spectrum

\Omega_{f}(\omega)

, POD spectrum

\sigma_{\ell}

, and autocorrelation

C_{f}(\tau)

. Black: GENE ground truth,

t\in[400,900]

. Blue: PreVAE-Turb,

t\in[400,900]

. Red: PreVAE-Turb,

t\in[2400,2900]

Figure 8: Ablation study: validation metrics for PreVAE-Turb trained without spectral loss (

w_{\text{spec}}=0

). Compared to Figure 7 , differences are visible in POD spectra and autocorrelation, while the spatial spectra remain broadly similar.

Figure 9: Manifold consistency error

E_{\text{mani}}

(Eq. 6 ) during autoregressive rollout, averaged over latent elements. The error remains around

0.006

, indicating stable encode–decode consistency over the plotted interval.

Figure 10: Cross-correlation functions between GENE field pairs. From left to right:

\phi

n

(particle transport),

\phi

T_{\perp}

(perpendicular heat transport),

n

T_{\perp}

(density-temperature coupling), and

T_{\parallel}

T_{\perp}

(temperature anisotropy). Black lines: GENE ground truth. Blue lines: PreVAE-Turb prediction. The agreement in cross-correlations demonstrates that PreVAE-Turb preserves the multi-field coupling structure essential for transport predictions.

Figure 11: Temporal evolution of normalized midplane transport proxies. Left: particle flux

\Gamma=\langle\tilde{n}\,v_{r}\rangle

. Middle: parallel heat flux

Q_{\parallel}=\langle\tilde{T}_{\parallel}\,v_{r}\rangle

. Right: perpendicular heat flux

Q_{\perp}=\langle\tilde{T}_{\perp}\,v_{r}\rangle

. Black: GENE ground truth (

t=0

–

1400

). Blue: PreVAE-Turb prediction (

t=0

–

4000

). Faint lines show raw data; bold lines show smoothed (moving average) values. The red dashed line at

t=400

marks the start of autoregressive prediction; before this, the outputs are VAE reconstructions of ground truth latents (warmup period).

Figure 12: Normalized midplane transport-proxy probability density functions (top row) and temporal autocorrelations (bottom row) for particle flux

\Gamma

(left), parallel heat flux

Q_{\parallel}

(middle), and perpendicular heat flux

Q_{\perp}

(right). Black: GENE ground truth (

t\in[400,1400]

). Blue: PreVAE-Turb prediction (

t\in[400,1400]

). Red: PreVAE-Turb long-horizon prediction (

t\in[1400,4000]

). PreVAE-Turb reproduces the flux distributions and decorrelation timescales, with visible shifts in the long-horizon PDF peaks.

Model	Stage	Hardware	Wall Time
GENE Solver	Simulation	4 $\times$ NVIDIA A100 (Perlmutter)	235 min
PreVAE-Turb	VAE fine-tuning	RTX A5000	220 min
	ConvLSTM training	RTX A5000	53 min
	Inference (2000 steps)	RTX A5000	13.7 s
	Inference (4000 steps)	RTX A5000	19.4 s

Table 2: Computational efficiency comparison for GENE turbulence (

128\times 48

grid, 4 fields). GENE evolves the full five-dimensional distribution function, whereas the surrogate predicts four two-dimensional midplane fields. Training costs are one-time (1000 epochs each); inference costs are per trajectory.

Explore similar work

A Shortcut to Statistically Steady-State Turbulence with Flow Matching
Jul 14, 2026Gianluca Galletti, Gerald Gutenbrunner, William Hornsby +5Turbulent Flows Generative Models
Institute for Machine Learning, JKU Linz · United Kingdom Atomic Energy Authority, Culham campus · Mistral AI
Surrogate Modeling of 3D Rayleigh-Benard Convection with Equivariant Autoencoders
Date pendingFynn Fromme, Hans Harder, Christine Allen-Blanchette +1Surrogate Models Conditional Variational Autoencoder
Lamarr Institute for Machine Learning and Artificial Intelligence and Karlsruhe Institute of Technology · Lamarr Institute for Machine Learning and Artificial Intelligence and Department of Computer Science at Paderborn University · Department of Mechanical and Aerospace Engineering Center for Statistics and Machine Learning Princeton University +1
Physics-Preserving Latent Compression for Zero-Shot Resolution Transfer in 3D Turbulence
Jun 19, 2026Yilong Dai, Yiming Sun, Yiheng Chen +4Turbulent Flows Data Compression Methods
The University of Alabama · University of Pittsburgh · University of Maryland, College Park

Jul 14, 2026, physics.plasm-phJ/K move · Enter open · S save

A Shortcut to Statistically Steady-State Turbulence with Flow Matching

Gianluca Galletti, Gerald Gutenbrunner, William Hornsby, Lorenzo Zanisi+4

Institute for Machine Learning, JKU Linz · United Kingdom Atomic Energy Authority, Culham campus · Mistral AI

Turbulent Flows Generative Models Statistical Inference Shortcut Learning

Open paper AskPDF arXiv

Many nonlinear physical systems exhibit an initial transient phase in which perturbations grow before nonlinear interactions lead to a statistically steady state. While this saturated regime is of primary interest, direct numerical simulations must resolve the full transient dynamics before reaching it, incurring significant computational cost. In Computational Fluid Dynamics, reduced-order approaches such as Large Eddy Simulation mitigate computational cost by modeling small-scale dynamics, enabling tractable approximations of turbulent flows. In contrast, for systems such as gyrokinetics, comparably effective closures for the full dynamics are not generally available, and high-fidelity simulations remain necessary. Existing surrogate modeling approaches for these systems are autoregressive, hence they suffer from accumulating error. We instead propose to bypass explicit time evolution by directly modeling the distribution of saturated states under an ergodicity assumption, stating that ensemble averages over samples are equivalent to time averages of a single long simulation. We introduce GyroFlow, a latent generative model that directly estimates steady-state statistics of gyrokinetic turbulence in 5D phase space, without resolving the transient phase. GyroFlow generates saturated snapshots from noise, conditioned on dimensionless operating parameters and outperforms autoregressive, reduced-order, and other generative approaches, while providing substantial speedup. To evaluate generation quality we propose FGyD, a distributional metric computed in the latent space of a pretrained gyrokinetic model, and show that it correlates with downstream flux accuracy and solver convergence. Finally, GyroFlow can be used to warm-start the numerical code used to produce the data.

Abstract

Explore similar work

A Shortcut to Statistically Steady-State Turbulence with Flow Matching

Surrogate Modeling of 3D Rayleigh-Benard Convection with Equivariant Autoencoders

Physics-Preserving Latent Compression for Zero-Shot Resolution Transfer in 3D Turbulence