PaperlayerView exact revision on arXiv

Source-linked AI summary

GyroSwin: 5D Surrogates for Gyrokinetic Plasma Turbulence Simulations

Fabian Paischer, Gianluca Galletti, William Hornsby, Paul Setinek, Lorenzo Zanisi, Naomi Carey, Stanislas Pamela, Johannes Brandstetter

arXiv:2510.07314v4physics.plasm-phcs.AIstat.ML

TL;DR

Reduced-order models are used because nonlinear gyrokinetic simulations are costly, but they omit nonlinear physics important to turbulent transport. GyroSwin learns the evolving 5D distribution function and related physical quantities; the paper reports improved flux predictions over reduced numerical approaches and favorable scaling, while noting rollout error accumulation and restricted simulation scope.

  • Problem

    Nonlinear gyrokinetic simulations are costly, while reduced models used to approximate turbulent transport omit nonlinear effects such as zonal flows.

    Source: §1 · p. 2

  • Method

    GyroSwin is a multitask 5D neural surrogate that evolves the distribution function and predicts derived potential and heat-flux quantities using latent cross-dimensional modules and channelwise mode separation.

    Source: §1 · p. 2, §6 · p. 10

  • Results

    GyroSwin improves time-averaged flux predictions over state-of-the-art ROMs and scales favorably against other neural surrogates.

    Source: §1 · p. 2

  • Takeaways & Limitations

    Predicted physical quantities align with ground truth on unseen simulations, and GyroSwin is three orders of magnitude faster than GKW for the adiabatic electron approximation.

    Source: §1 · p. 2

    GyroSwin accumulates error during autoregressive rollouts, neglects the linear phase, and is evaluated only with the adiabatic electron approximation.

    Source: §6 · pp. 10–11

Abstract

from arXiv · show

Nuclear fusion plays a pivotal role in the quest for reliable and sustainable energy production. A major roadblock to viable fusion power is understanding plasma turbulence, which significantly impairs plasma confinement, and is vital for next-generation reactor design. Plasma turbulence is governed by the nonlinear gyrokinetic equation, which evolves a 5D distribution function over time. Due to its high computational cost, reduced-order models are often employed in practice to approximate turbulent transport of energy. However, they omit nonlinear effects unique to the full 5D dynamics. To tackle this, we introduce GyroSwin, the first scalable 5D neural surrogate that can model 5D nonlinear gyrokinetic simulations, thereby capturing the physical phenomena neglected by reduced models, while providing accurate estimates of turbulent heat transport. GyroSwin (i) extends hierarchical Vision Transformers to 5D, (ii) introduces cross-attention and integration modules for latent 3D$\leftrightarrow$5D interactions between electrostatic potential fields and the distribution function, and (iii) performs channelwise mode separation inspired by nonlinear physics. We demonstrate that GyroSwin outperforms widely used reduced numerics on heat flux prediction, captures the turbulent energy cascade, and reduces the cost of fully resolved nonlinear gyrokinetics by three orders of magnitude while remaining physically verifiable. GyroSwin shows promising scaling laws, tested up to one billion parameters, paving the way for scalable neural surrogates for gyrokinetic simulations of plasma turbulence.

1 Introduction

Gyrokinetic simulations capture nonlinear plasma turbulence but are computationally costly, while reduced models omit important nonlinear physics. GyroSwin learns the full 5D dynamics and predicts derived physical quantities, with reported gains in accuracy and computational efficiency.

1 Introduction

  • Nonlinear gyrokinetic simulations are prohibitively costly, so quasi-linear reduced models approximate turbulent transport but neglect zonal flows that substantially affect it.A reliable estimate of turbulent fluxes is described as attainable only through expensive nonlinear simulations.Source: §1 · pp. 1–2
  • GyroSwin combines a 5D Swin Transformer, latent 3D↔5D cross-attention and integration, and channelwise mode separation, trained on the distribution function and derived quantities.Its multitask targets include 3D potential fields and scalar fluxes.Source: §1 · p. 2
  • GyroSwin aligns predicted downstream quantities with ground truth on unseen simulations and runs three orders of magnitude faster than GKW under the adiabatic electron approximation.Reported downstream quantities include heat flux, turbulence-related spectra, and zonal flows.Source: §1 · p. 2
  • GyroSwin improves time-averaged flux predictions over state-of-the-art ROMs, captures nonlinear 5D dynamics, and scales favorably against other neural surrogates.The model was tested up to 1B parameters.Source: §1 · p. 2

2 Background and Related Work

The background describes gyrokinetics as a reduced kinetic model for studying turbulent transport, whose nonlinear 5D dynamics make simulations costly. Common quasi-linear approaches reduce computation but omit nonlinear mode interactions, while prior machine-learning approaches largely rely on reduced data or lower-dimensional inputs.

Background

  • Gyrokinetics evolves a 5D distribution function over spectral coordinates, position along a field line, and parallel and perpendicular velocities.It reduces plasma kinetics by averaging over the angle around field lines and following particle guiding centers.Source: §2 · p. 3
  • The nonlinear term couples the distribution function to electrostatic potential and describes turbulent advection, making it the most computationally expensive term.The potential is obtained from a velocity-space integral of the distribution function.Source: §2 · p. 3

Quantities of interest

  • Heat transport toward the reactor wall is a central quantity of interest, and turbulence analysis uses ky spectra to examine mode contributions and energy.Zonal flows shear and break up eddies, damping turbulence in the statistically steady state.Source: §2 · pp. 3–4

State-of-the-art numerical approximations

  • Quasi-linear models solve modes independently and combine their contributions with fitted saturation rules, but omit nonlinear interactions and can be limited in strong-turbulence regions.These models reduce computation by solving for a reduced distribution function for each mode.Source: §2 · p. 4

Machine learning for Gyrokinetics

  • Most prior machine-learning gyrokinetics work predicts fluxes from operating parameters or uses 2D wavenumber slices, potential snapshots, or a reduced 1D space.GyroSwin instead trains directly on the 5D distribution function of nonlinear gyrokinetics.Source: §2 · p. 4

3 GyroSwin

GyroSwin is a Swin-based, hierarchical 5D surrogate designed to scale while preserving local structure and incorporating physical priors. It predicts the distribution function alongside integrated 3D potential and scalar flux, using cross-dimensional latent communication and channelwise mode separation.

3 GyroSwin

  • GyroSwin directly evolves the 5D distribution function, addressing 5D phenomena omitted by state-of-the-art ROMs.Its downstream physical outputs are designed to remain consistent with the predicted 5D field.Source: §3 · pp. 4–5

Inductive biases

  • Channelwise mode separation isolates the zonal-flow mode and adds its real and imaginary parts as channels to bias the model toward nonlinear dynamics.GyroSwin is designed as a Swin-based UNet with multiple branches for multitask predictions and physical priors.Source: §3 · p. 5

5D shifted window attention

  • Local window attention reduces ViT's quadratic token complexity toward near-linear scaling, while shifted windows enable interaction across neighboring windows.Stacked W-MSA and SW-MSA blocks provide approximate global attention with increasing depth.Source: §3 · pp. 5–6

Up/Downsampling

  • The hierarchical 5D UNet uses patch embedding, merging, and expansion with skip connections to form multiresolution representations.Merging concatenates neighboring 2×2×2×2×2 patches, reduces resolution by 25 times, and doubles channels.Source: §3 · p. 6

5D←→3D mixing and integrator modules

  • Latent integrators aggregate velocity dimensions into a 3D latent, and cross-attention mixing lets 3D and 5D prediction branches communicate.The 3D potential head uses integral layers, 3D window-attention layers, and expansion layers; the flux head pools integrated features before an MLP.Source: §3 · p. 6

4 Experiments

The experiments generate nonlinear simulations under the adiabatic electron approximation and compare GyroSwin with reduced, tabular, and neural baselines. Evaluation uses held-out simulations selected for in-distribution and out-of-distribution generalization.

Data Generation

  • At resolution 32×8×16×85×32, the dataset comprises 255 simulations and 44,400 training samples assembled from subsets of 48 and 241 simulations for baseline comparisons and scaling experiments.Each simulation ran for 31,920 steps, averaged every 40 steps and subsampled every third step to yield 266 snapshots total; the first 80 snapshots were omitted as the linear phase.Source: §4 · p. 7
  • Simulations vary initial noise and four operating parameters, and use the adiabatic electron approximation to focus on ion-temperature-gradient turbulence.The parameter space was populated using Latin hypercube sampling, with parameter regions chosen to produce highly turbulent simulations.Source: §4 · p. 7

Baselines

  • Baselines include the QuaLiKiz saturation rule, Gaussian process and MLP regressors, and FNO, PointNet, Transolver, and vanilla ViT neural surrogates.The ROM and neural surrogates are trained or fit using nonlinear fluxes or 5D distribution-function targets, respectively.Source: §4 · p. 7

Evaluation

  • Evaluation reserves 14 simulations: six ID cases within the training-set convex hull, five OOD cases outside it, and three for validation.The evaluation cases are selected in a high ion-temperature-gradient, strong-turbulence regime.Source: §4 · p. 7

5 Results

GyroSwin is evaluated on autoregressive 5D turbulence modelling, heat-flux prediction, scalability, and physical diagnostics, with evaluation scope varying by model capability. Results show improved flux prediction when scaling the model and data, stable long rollouts, and reproduction of turbulence diagnostics.

Evaluation scope

  • Evaluation scope depends on model capability: only full-5D neural surrogates can also be assessed for zonal flows.Tabular regressors are evaluated only on scalar nonlinear-flux prediction, while QL models also support diagnostics.Source: §5 · p. 7

5D→5D Turbulence modelling

  • GyroSwin is the most stable autoregressive neural surrogate, sustaining rollouts for over 100 timesteps even on OOD simulations as data and model size scale.Correlation time counts snapshots predicted while maintaining the specified Pearson-correlation threshold.Source: §5 · pp. 7–8

5D→0D average flux

  • Scaling GyroSwin to more data and larger models improves nonlinear heat-flux prediction, while competing surrogates’ errors cluster near 120 as their rollout flux decays toward zero.On the reduced training set, GyroSwin performs best against the QL model and alternative 5D surrogates; Table 2 covers six ID and five OOD simulations.Source: §5 · p. 8, Table 2 · p. 8

Scalability

  • GyroSwinSmall is roughly three orders of magnitude faster than GKW, while GyroSwin is the only evaluated method scaled to about 1B parameters with reasonable inference speed and improved diagnostic correlation.The scalability evaluation reports inference speed, memory consumption, and parameter count; other approaches face scaling constraints.Source: §5 · p. 8, Table 3 · p. 10

Scaling laws

  • Increasing GyroSwin from 239M to 893M parameters at full data reduces loss by 12%, and the model-size and data axes must be scaled together.Across the tested configurations, model scaling to 893M enters the tail of the scaling law, while data scaling does not saturate.Source: §5 · pp. 8–9, Figure 3 · p. 9

Diagnostics and nonlinear physics

  • GyroSwinLarge reproduces the OOD turbulence-intensity spectrum’s shape and magnitude nearly perfectly except at some higher frequencies, and captures the zonal-flow profile.Other neural surrogates tend to underestimate the intensity spectrum; GyroSwinLarge’s zonal-flow profile is shown for a selected OOD case.Source: §5 · p. 9, Figure 4 · p. 10

6 Conclusions and Limitations

GyroSwin captures nonlinear phenomena self-consistently and offers an alternative to reduced numerical approaches. Its scope is limited by omitted distributional turbulence, rollout error accumulation, exclusion of the linear phase, and the adiabatic-electron approximation.

6 Conclusions and Limitations

  • GyroSwin captures nonlinear phenomena self-consistently and outperforms reduced numerical approaches while scaling favorably against other neural surrogates.It is presented as an alternative for efficient approximation of turbulent transport.Source: §6 · p. 10
  • GyroSwin does not account for the chaotic, distributional nature of turbulence.The authors identify this as its main limitation.Source: §6 · p. 10
  • Autoregressive rollouts are stable over 100 timesteps, but suffer from error accumulation.The authors propose generative modelling of saturated-phase snapshots as a future remedy.Source: §6 · p. 11
  • The model omits the linear simulation phase, with the stated focus primarily on modelling turbulence.The authors aim to extend GyroSwin to the linear phase in future work.Source: §6 · p. 11
  • GyroSwin uses only the adiabatic-electron approximation, and data volumes make full coverage of the 4D parameter space difficult.The approximation allows a relatively large training set at lower cost.Source: §6 · p. 11

A Derivation of the Gyrokinetic equation

The derivation moves from particle phase-space dynamics to a gyroaveraged guiding-center description, then presents the gyrokinetic equation and a quasilinear transport model.

A Derivation of the Gyrokinetic equation

  • The Vlasov equation describes particle conservation in phase space without collisions; its terms include acceleration response and spatial advection.The distribution function is defined over spatial and velocity coordinates, with time as an additional dimension.Source: A Derivation of the Gyrokinetic equation · p. 17
  • The derivation transforms particle coordinates to guiding-center coordinates and gyroaverages away gyrophase under a long background timescale and small-Larmor-radius assumption.The guiding-center coordinates include position, parallel velocity, magnetic moment, and gyrophase.Source: A Derivation of the Gyrokinetic equation · p. 17

Linear Terms

  • The linear terms include guiding-center motion along the magnetic field and magnetic drifts; the magnetic-gradient term is expressed separately.Here b is the unit vector along the magnetic field and vD represents magnetic drifts.Source: Linear Terms · p. 17

Nonlinear Term

  • Fluctuating electromagnetic potentials induce E×B and magnetic-flutter drifts, with the gyroaveraged generalized potential used to express the resulting drift.The electrostatic potential is computed from an integral of f over velocity space.Source: Nonlinear Term · pp. 17–18

Final Equation

  • The split gyrokinetic equation includes source, collision, or drive terms, and represents f through equilibrium and perturbation components.The perturbation contains adiabatic and kinetic responses; the kinetic term governs irreversible dynamics that facilitate turbulence.Source: Final Equation · p. 18

B Quasilinear models

  • The QuaLiKiz saturation rule estimates turbulent transport from linear stability information and a turbulence-intensity spectrum calibrated using nonlinear fluxes.In this setup, the normalization factor was fitted using flux estimates from 48 training simulations, yielding C = 7.93.Source: B Quasilinear models · p. 18

C Data Generation and Visualisation

The simulations focus on ion-temperature-gradient turbulence and sample four operating parameters, adjusting ranges after many initial configurations failed to become turbulent. Visualizations show that nonlinear simulations contain asymmetric, stronger, more-interacting wavenumber structures than linear simulations.

C Data Generation and Visualisation

  • The study focuses on ion-temperature-gradient-driven turbulence and varies R/Lt, R/Ln, q, and magnetic shear.R/Lt is described as the main turbulence driver; R/Ln and magnetic shear can have stabilizing effects.Source: C Data Generation and Visualisation · p. 19
  • After only about half of the first 100 simulations became turbulent, the authors adjusted parameter ranges and generated 200 more, all strongly turbulent.The final training set contains 241 simulations, with three used for validation and six for in-distribution evaluation.Source: C Data Generation and Visualisation · p. 19
  • The 5D visualizations show selected pairs of field axes as 2D planes, averaging or slicing over the remaining axes.Nonlinear and linear planes are displayed side by side at three timesteps with shared colorbars for comparison.Source: C Data Generation and Visualisation · p. 19
  • Nonlinear simulations show asymmetric structures, larger wavenumber-space magnitudes, and more mode interaction than linear simulations.The contrast indicates structure in nonlinear simulations that linear-simulation-based approximations do not capture.Source: C Data Generation and Visualisation · p. 19

D Implementation Details

GyroSwin uses interacting branches to model the 5D distribution function and 3D electrostatic potential, with separate procedures for flux prediction and experiments scaling model and data size. The implementation also compares several baseline model families and reports stable, though increasingly error-prone, autoregressive rollouts.

D Implementation Details

  • The 5D δf field uses four input/output channels, while the 3D electrostatic potential uses one channel; fields and fluxes are standardized using training-set statistics.δf statistics are computed separately for each µ index, while ϕ statistics are global.Source: D Implementation Details · p. 19
  • The δf branch uses shifted-window attention over a 4D token grid, while the ϕ decoder obtains latents through velocity-space integration and exchanges information with δf through cross-attention.The branches exchange information at the bottleneck, decoder stages, and before unpatching.Source: D Implementation Details · p. 20
  • The training objective combines normalized relative squared errors for δf and ϕ, with the ϕ loss weight ramped from zero to one during training.The weight is zero before 50% progress, increases linearly between 50% and 70%, and stays at one thereafter.Source: D Implementation Details · p. 22
  • For heat-flux fine-tuning, a flux head is trained on frozen pretrained latents to predict heat flux averaged over the final 80 snapshots.Only the head is trained, using L1 loss for 200 epochs; the checkpoint with lowest validation flux MSE is retained.Source: D Implementation Details · p. 22
  • The baselines include parameter-to-flux GPR and MLP models, a spatial 3D FNO, and coordinate-based PointNet and Transolver models predicting the next 5D-field timestep.PointNet and Transolver use the same grid-sampling and coordinate strategy.Source: D Implementation Details · pp. 22–23

E Scaling Configurations

  • Scaling experiments vary model size and training-data subsets while holding Appendix D hyperparameters fixed.Model-axis configurations use two seeds; the Medium data-axis configuration uses two seeds and XS uses one.Source: E Scaling Configurations · p. 23
  • GyroSwin preserves high-level rollout structure initially, while errors accumulate later without predictions diverging.The rollout comparison is shown for the saturated phase.Source: Figure 8 · p. 24

F Additional Results

GyroSwin maintains stable long autoregressive rollouts, although error accumulates and velocity-space predictions tend to overpredict.

F Additional Results

  • GyroSwin’s wavenumber-space predictions remain somewhat stable without diverging after 100 autoregressive steps, despite accumulated errors and velocity-space overprediction.Overall structure is more or less preserved within the first 10 rollout steps.Source: F Additional Results · p. 24

G Diagnostics

GyroSwin improves turbulence-intensity spectra over the reduced numerical quasilinear model, while flux-spectrum correlation is approximately unchanged. Its zonal-flow profiles capture unseen simulations but can be smoothed or inaccurate in individual cases.

G Diagnostics

  • GyroSwin substantially improves W(ky) over the state-of-the-art reduced-numerical quasilinear model, but Q(ky) correlation is approximately unchanged.High frequencies are not well captured in Q(ky); the authors suggest neural-network spectral bias as a possible explanation.Source: G Diagnostics · p. 24, Figure 9 · p. 25, Figure 10 · p. 25
  • GyroSwinLarge can capture zonal-flow profiles for unseen simulations, but profiles are often smoothed and can be entirely off in some cases.The authors report overestimation of amplitude and normalize the profiles for visualization; some cases contribute to high variance in the main table.Source: G Diagnostics · pp. 24–25

H Ablation Studies

The ablations show that GyroSwin’s components contribute differently to rollout stability and nonlinear-flux prediction, with the flux head yielding the strongest combination. A separate saturation-rule comparison favors GyroSwinLinear over QL, while the complexity discussion motivates windowed 5D attention.

H Ablation Studies

  • Adding the flux-prediction head gives the largest ablation gain, yielding the best combination of stable autoregressive rollouts and nonlinear-flux prediction.The 5D-Swin baseline rolls out stably but misses the nonlinear flux trace; a UNet-like structure improves flux prediction with a slight drop in correlation time. Channelwise mode separation especially boosts correlation time on OOD simulations, while latent cross-attention and integration add a slight improvement.Source: H Ablation Studies · p. 25, Table 7 · p. 26

I Learning a saturation rule

  • GyroSwinLinear significantly outperforms the reduced numerical model (QL), and with more training data it outperforms tabular regressors GPR and MLP.The reported trend also finds machine-learning methods outperforming the reduced model; all methods improve with larger training sets.Source: I Learning a saturation rule · p. 26, Table 8 · p. 27

I.1 Vision Transformer Complexity

  • The 5D Swin-layer complexity replaces a squared dependency on sequence length xyzhw with a smaller window complexity M^5.The complexity discussion frames windowed attention as the reason attention can scale more favorably with resolution.Source: §I.1 · p. 27

Related papers

Loading 2510.07314v4…