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
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=Nx and W=Ny . 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, wspec=0 .
Model
Field
MSE ↓
PSNR (dB) ↑
SSIM ↑
Correlation ↑
CNN-VAE
Density
0.256 ± 0.035
23.7 ± 1.1
0.955 ± 0.007
0.969 ± 0.005
Potential
0.137 ± 0.026
26.2 ± 1.3
0.972 ± 0.006
0.981 ± 0.004
Pre-trained
Density
0.041 ± 0.065
32.5 ± 2.7
0.993 ± 0.009
0.995 ± 0.008
Potential
0.019 ± 0.050
35.7 ± 2.7
0.996 ± 0.007
0.997 ± 0.007
Table 1: Reconstruction quality metrics (mean ± 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 ϕ (bottom row). From left to right: spatial spectrum Sf(k) , temporal spectrum Ωf(ω) , POD spectrum σℓ , and temporal autocorrelation Cf(τ) . Black solid: HW ground truth. Blue dashed: PreVAE-Turb prediction t∈[8000,13000] . Red dotted: PreVAE-Turb prediction t∈[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: ϕ - n cross-correlation Cϕn(τ) showing agreement between HW model (black) and PreVAE-Turb predictions at different time windows (blue dashed: t∈[8000,13000] , red dotted: t∈[18000,23000]) . Middle: joint probability distribution P(ϕ,n) for HW ground truth. Right: joint probability distribution P(ϕ,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 ϕ , density n , parallel temperature T∥ , and perpendicular temperature T⊥ . 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 ( ϕ , n , T∥ , T⊥ ). Columns: spatial spectrum Sf(k) , temporal spectrum Ωf(ω) , POD spectrum σℓ , and autocorrelation Cf(τ) . Black: GENE ground truth, t∈[400,900] . Blue: PreVAE-Turb, t∈[400,900] . Red: PreVAE-Turb, t∈[2400,2900] .
Figure 8: Ablation study: validation metrics for PreVAE-Turb trained without spectral loss ( wspec=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 Emani (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: ϕ - n (particle transport), ϕ - T⊥ (perpendicular heat transport), n - T⊥ (density-temperature coupling), and T∥ - T⊥ (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 Γ=⟨n~vr⟩ . Middle: parallel heat flux Q∥=⟨T~∥vr⟩ . Right: perpendicular heat flux Q⊥=⟨T~⊥vr⟩ . 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 Γ (left), parallel heat flux Q∥ (middle), and perpendicular heat flux Q⊥ (right). Black: GENE ground truth ( t∈[400,1400] ). Blue: PreVAE-Turb prediction ( t∈[400,1400] ). Red: PreVAE-Turb long-horizon prediction ( t∈[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 × 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×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.
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.
Gianluca Galletti, Gerald Gutenbrunner, William Hornsby +5
Institute for Machine Learning, JKU Linz · United Kingdom Atomic Energy Authority, Culham campus · Mistral AI
The use of machine learning for modeling, understanding, and controlling large-scale physics systems is quickly gaining in popularity, with examples ranging from electromagnetism over nuclear fusion reactors and magneto-hydrodynamics to fluid mechanics and climate modeling. These systems - governed by partial differential equations - present unique challenges regarding the large number of degrees of freedom and the complex dynamics over many scales both in space and time, and additional measures to improve accuracy and sample efficiency are highly desirable. We present an end-to-end equivariant surrogate model consisting of an equivariant convolutional autoencoder and an equivariant convolutional LSTM using G-steerable kernels. As a case study, we consider the three-dimensional Rayleigh-B'enard convection, which describes the buoyancy-driven fluid flow between a heated bottom and a cooled top plate. While the system is E(2)-equivariant in the horizontal plane, the boundary conditions break the translational equivariance in the vertical direction. Our architecture leverages vertically stacked layers of D4-steerable kernels, with additional partial kernel sharing in the vertical direction for further efficiency improvement. We demonstrate significant gains in sample and parameter efficiency, as well as a better scaling to more complex dynamics. The accompanying code is available under https://github.com/FynnFromme/equivariant-rb-forecasting.
Fynn Fromme, Hans Harder, Christine Allen-Blanchette +1
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
High-resolution turbulence modeling is essential for scientific computing, but remains constrained by the cost of direct numerical simulation and the scarcity of full-resolution data. Existing scientific compressors reduce storage but typically operate on per-frame representations, whereas learned compressors yield compact latents that are often resolution-dependent and weakly aligned with the physics of turbulence. This raises the need for a compression framework that reduces data size, preserves physical diagnostics, and transfers from low-resolution training fields to high-resolution test fields without retraining. In this paper, we propose Physics-Preserving Latent Compression (PPLC), a patch-local latent compressor for three-dimensional turbulence. Motivated by inertial-range scale similarity, PPLC treats fixed-size patches as transferable units and applies a shared variational autoencoder independently of the global grid size. It combines exact mean preservation, zero-mean fluctuation encoding, an invertible Haar wavelet front-end, shift-consistency regularization, and overlap-aware reconstruction. Instantiated on forced isotropic turbulence, PPLC is trained only on stride-downsampled 256^3 fields and transfers zero-shot to 1024^3 fields. Experiments show that PPLC improves the balance between reconstruction accuracy and physical fidelity over classical and learned baselines, keeping diagnostics such as dissipation, enstrophy, energy spectra, and incompressibility closer to the ground truth. Beyond turbulence compression, PPLC offers a general strategy for physics-preserving latent representations that support data-efficient scientific surrogate modeling.
Yilong Dai, Yiming Sun, Yiheng Chen +4
The University of Alabama · University of Pittsburgh · University of Maryland, College Park