GyroSwin: 5D Surrogates for Gyrokinetic Plasma Turbulence Simulations
Authors: Fabian Paischer
Gianluca Galletti
William Hornsby
Paul Setinek
Lorenzo Zanisi
Naomi Carey
Stanislas Pamela
Johannes Brandstetter
Venue: NeurIPS 2025 poster ()
Keywords: 5D Swin Transformer
Gyrokinetics
Nuclear Fusion
PDE
Primary area: machine_learning_for_sciences
TL;DR: We propose GyroSwin, a 5D Swin Transformer to learn a 5D PDE commonly encountered in Gyrokinetics to simulate turbulence in a nuclear fusion reactor.

Abstract: 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.

GyroSwin: 5D Surrogates for Gyrokinetic Plasma Turbulence Simulations
Fabian Paischer 1,3
Gianluca Galletti
William Hornsby²
Paul Setinek
Lorenzo Zanisi?
Naomi Carey?
Stanislas Pamela 2
Johannes Brandstetter
ELLIS Unit, Institute for Machine Learning, JKU Linz
United Kingdom Atomic Energy Authority, Culham campus
EMMI AI, Linz
gaaaacherrgllettti,, brandstetter))m pp j.u. at
ml Ll-Juu/heural-ggrookeeteee all-jku/gyroswit
Abstract
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 gy- rokinetic equation, which evolves a 5D distribution function over time. Due to its high computational cost, reduced-order models are often employed in practice approximate turbulent transport 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 hi- erarchical Vision Transformers 5D, (ii) introduces cross-altention and integra- tion modules for latent 3D++5 interactions between electrostatic potential fields and the distribution function, and (iii) performs channelwise mode separation in- spired by nonlinear physics. We demonstrate that GyroSwin outperforms widely used reduced numerics on heat flux prediction, captures the turbulent energy cas- cade, 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 neu- ral surrogates for gyrokinetic simulations of plasma turbulence.
Introduction
Nuclear fusion promises sustainable energy by fusing hydrogen isotopes within an ionized plasma. Because the plasma reaches temperatures of hundreds of millions of degrees, it must be confined by a magnetic field. During confinement, turbulence can arise from micro-instabilities, leading to energy and particle transport toward the reactor walls. This causes plasma to leak from its magnetic cage, resulting in heat and density loss, and therefore impaired energy production. The design and control of plasma scenarios strictly require knowledge of turbulent transport, which be obtained via nonlinear gyrokinetic simulations that evolve a 5D Partial Differential Equation (PDE) over time.
The computational cost of nonlinear gyrokinetic simulations is prohibitive. Therefore, QuasiLLinear approaches (QL), such as QuaLiKiz (Bourdelle et al., 2015; Citrin al., 2017) and TGLF (Staebler
Equal contribution
39th Conference on Neural Information Processing Systems (NcurIPS 2025).

---

f(t)
GyroSwin
quasilineat
Linear OK
lnb Yol Ty
LBB ta
Saturation rule
Multitask Loss witdh.f m444/0.01
naturated
5D
2Dx3D
(X X
((1)
33
27
Q
50
Data
@D
3D
3D
3D
Nonlinear Q Data
Figure I: Left: GyroSwin models the 5D distribution function of nonlinear gyrokinetics and incor- porates integration blocks predict 3D electrostatic potential fields and scalar heat flux. Right: ROMs (quasilinear) solve a cartesian product of 2D modes in spectral space and 3D fields. Further- more, they rely on saturation rules approximate the nonlinear flux spectrum.
et al., 2007; Staebler & Kinsey, 2010), are commonly used to approximate turbulent transport. They are based Reduced-Order Model (ROM) that operates in 3D and adopt so-called saturation rule to estimate nonlinear fluxes. These saturation rules usually employ free parameters that are fitted to nonlinear flux data. However, QL models neglect essential parts of the nonlinear physics that contribute substantially to evolving turbulent transport, namely zonal flows. Therefore, a reliable estimate of turbulent fluxes to date is only attainable via expensive nonlinear gyrokinetic simulations.
To tackle the prohibitive cost of nonlinear gyrokinetic simualtions, we introduce GyroSwin, scal- able neural surrogate model efficiently approximate turbulent transport. GyroSwin based on three essential ingredients: (i) extension of a Swin transformer (Liu et al., 2021) 5D data, (ii) ila- tent cross-attention and integration modules for interaction between 5D and 3D fields, as well as 5D) + 3D integration, and (iii) channelwise mode separation informed by nonlinear physics. To ground GyroSwin on physically meaningful quantities, we train it multitask manner on the distribution function and derived quantities thereof, such 3D potential fields, and scalar fluxes.
GyroSwin accurately captures nonlinear physics of 5D gyrokinetics. To verify this, infer phys- ical downstream quantities of the predicted 5D PDE, such as heat flux, turbulence related spectra, and zonal flows. We find that the predicted quantities align well with the ground truth for unseen simulations and GyroSwin scales favorably compared to other neural surrogate approaches. Further- more, GyroSwin three orders of magnitude faster than the numerical code GKW for the adiabatic electron approximation. In summary, our contributions are as follows.
We introduce GyroSwin, multitask hierarchical Vision Transformer that handles data of arbitrary dimensionality. It is trained on the adiabatic electron approximation using channelwise mode separation predict the 5D PDE, 3D electrostatic potentials, and scalar fluxes simultaneously.
We propose latent antss-attention between fields of varying dimensionality, well as integration layers that enable integrating along various dimensions of the 5D field.
We demonstrate that GyroSwin captures nonlinear dynamics in the 5D field, improves time- averaged flux predictions over state-of-the-art ROMs, and significantly outperforms deep learning surrogates plasma turbulence modelling.
demonstrate promising scaling laws of GyroSwin, tested up to IB parameters.
2 Background and Related Work
SD Gyrokinetics. To determine turbulent transport, one must simulate the evolution of electrons and ions in the plasma over time. In principle, this could be achieved by modelling each species as particle in the plasma (Birdsall & Langdon, 2005; Tskhakaya, 2008). However, because the number

---

of particles in plasma can exceed 10a0, this approach is computationally expensive. An alternative approach is statistically model ensembles of particles in the plasma using a distribution function. This approach is termed plasma kinetics and is desirable to analyse stability in a bulk of plasma.
Plasma kinetics models the time evolution of electrons and ions via the distribution function based on 3D coordinates, their parallel and perpendicular velocities, and the angle around the field lines. Gyrokinetics (Krommes, 2012) is reduced form of plasma kinetics that more efficient by averaging out the angle around the field lines and only considering a particle's guiding center. Local gyrokinetics focuses on perpendicular spatial scales comparable to the gyroradius and on frequencies much lower than the particle cyclotron frequency, the circular motion of charged particles around magnetic field lines due to the Lorentz force. Hence, the 5D gyrokinetic distribution function can be written as f = fkkz, k.y. 8, 2||, ")), where K2 and ky are spectral coordinates (for the spatial æ and y). 8 is the toroidal coordinate along a field line, and 2I and i represent parallel and perpendicular velocity components, respectively. The perturbed time-evolution of f. for each species (ions and electrons), is governed by
as Øt + (eij b UD) vf JBB B VB aS
x V = S
(1)
B²
du
Linear
Nonlincar
where og. is the motion along magnetic field lines, b B/B B/Bii the unit vector along the magnetic field B with magnitude B², UD the magnetic drift due to gradients and curvature in B, and Vx describes drifts arising from the E x B force, key driver of plasma turbulence. The nonlinear term models the interaction between the distribution function f and the electrostatic potential, - = E, which comes from the velocity space integral of itself, and describes turbulent advection. The resulting nonlinear coupling constitutes the computationally most expensive term. Finally, $ is the source term that represents collisions between particles. For more detailed derivation of the gyrokinetic equation, see Appendix A.
Quantities of interest. In practice, the main quantity of interest is the radial transport of energy (heat) towards the reactor wall, which is crucial for reactor design. It can be obtained by integrals on the 5D distribution function as
Q /= vos du dp drdyds,
f Jos dondp
f
Φ $=
A
where A, C RExXxy comprise geometric and operating parameters, v R denotes parti- cle energy, and Jo denotes the zeroth-order Bessel function and Φ E R=*xxxy is the electrostatic potential. In tokamak, B largely points in the toroidal direction, ka encodes the radial direction, and ky is perpendicular to k (binormal ~ poloidal) in Fourier space. Electrostatic fluctuations that drive turbulence occur mainly the ky direction, while hg encodes the radial structural properties of turbulence, i.c. the size of emerging eddies.
Radial turbulent transport occurs only if the amplitude of the hy modes is non-zero, as this results in x B drift, and consequently in fully developed nonlinear turbulence. A turbulent simulation usually follows certain pattern. In the linear phase, different modes in ky start growing due to underlying micro-instability, resulting in an initial increase in heat flux. Afterwards, the simulations enter the nonlinear (saturated) regime where modes start interacting and the system converges into a statistically steady state. This state is controlled by zonal flows that shear and break up arising eddies, effectively dampening turbulence (Itoh al., 2006).
To analyse turbulence, practitioners usually investigate various spectra along the by direction, as those provide insights into turbulent transport. Most importantly, ((kyy) provides insight into which modes dominate the turbulence energy:
Q(ky)
Σ Q( key
(3)
where E X2 is the flux field, which aggregates to the scalar flux Q in Equation (2). In addition (ky)) E Rtv describes the intensity of turbulence along y, and zonal flows dampen
We adopt uppercase notation for vector fields E and B to adhere with literature.
3

---

turbulence resulting in the statistically steady state which can be isolated as the first mode of Φ in ky direction:
(ky)) = E
(z, = Σ Sk. k, e ethiee
(4)
s.ke
Ri
Since the system converges to a statistically steady state, practitioners usually investigate time- averaged quantities of (kk,), W (ky)), ΦZF. and scalar heat flux trace Q
State-of-the--at numerical approximations. Solving Equation (1) numerically is computationally expensive, especially for high-fidelity simulations across ion and electron scales. Quasilinear (QL) models mitigate the main source of computational cost and are commonly adopted in integrated modelling pipelines (Mulders al., 2021; Citrin et al., 2024). They are based on linear simulations that neglect the nonlinear term in Equation (1) and solve for multiple f for each mode in hi x ky
independently. This results in a speedup as the order of f is reduced to f= ["]]] 14, s). linear simulations do not account for interaction between modes, but only compute linear growth rates and potential transport contributions for each of them. This mode-wise transport contribution combined via so-called saturation rules by
As result,
Q = (kk,))
(5)
where Qoi (ky) is the flux contribution per mode, and W(ky) modelled weighting function which approximates W(y)) and whose free parameters are fit to nonlinear simulations. Since QL models are based on linear simulations, they neglect nonlinear physics that substantially impact turbulence. Another limitation of QL models are operating parameter regions leading to strong turbulence (Staebler et al., 2024; Kiefer al., 2021; Dimits et al., 2000b; Bourdelle et al., 2008).
Machine learning for Gyrokinetics. Machine Learning offers a fruitful alternative ROMs. Most literature to date has focused on multilayer perceptrons (van de Plassche et al., 2020; Citrin et al., 2023; Zanisi et al., 2024, MLP) or gaussian processes (Hornsby et al., 2024). They usually map from operating parameters, such as temperature or density gradients to Q predicted via ROMs. Hence, they are inherently limited by the capabilities of ROMs used to produce the training data. Few works have attempted to train machine learning surrogates for nonlinear gyrokinetics. Narita et al. (2022) leverage convolutional neural networks on 2D wavenumber (kiz x ky) slices to predict time to saturation and heat flux. Building on this idea, Honda et al. (2023) include snapshots of electro- static potentials as additional channels. More recently. Wan et al. (2025) investigate transfer from low-fidelity to higher-fidede simulations in reduced ID space. Our proposed GyroSwin is funda- mentally different, as is trained directly on the 5D distribution function nonlinear gyrokinetics.
Neural surrogates. Over recent years, deep neural network-based surrogates have emerged as a computationally efficient alternative in science and engineering (Thuerey et al., 2021; Zhang et al., 2023; Brunton et al., 2020), impacting weather forecasting (Kurth et al., 2023; Bi et al., 2023; Lam et al., 2023; Nguyen et al., 2023; Bodnar et al., 2024), protein folding (Jumper et al., 2021; Abramson et al., 2024), material design (Merchant et al., 2023; Zeni et al., 2025; Yang et al., 2024), and multi-physies modelling Alkin et al. (2024b). These success stories share the common thread of deep learning surrogates overcoming seemingly insurmountable challenges (Brandstetter, 2024). Especially for weather modelling, Vision Transformer (Dosovitskiy et al., 2021, ViT) and especially Swin Transformer (Liu et al., 2021) have shown exceptional performance, which manifested in the first medium-range weather modelling surpassing numerical accuracy (Bi et al., 2023), and the current state-of-the-art foundation model of the atmosphere(Bodnart et al., 2024). In gyrokinetics face similar challenges in terms of fidelity, complexity and local couplings. However, these challenges are even exacerbated for gyrokinetics due to its 5D nature.
3 GyroSwin
State-of-the-arr ROMs neglect 5D physical phenomena that are crucial for reliable turbulent trans- port estimates. To remedy this, we develop a neural surrogate, GyroSwin, that directly learns to evolve the 5D distribution function of nonlinear gyrokinetics over time. The most important aspect

---

50Embed
50MA
5DWA
5 3DMix
Unnerge
SDWA
Ner ge
3DMix
Encoder
VVapacee
+ Head!
Q
((0130))
V(20-50)
X(30)
Vapuca
Uniterge
30MA
Head
(a) High-level GyroSwin architecture.
(b) 5D attention, latent mixing and integratior layers.
Figure 2: Left:: GyroSwin receives the 5D distribution function as input and predicts the evolved 5D distribution function, as well as the respective 3D potential and heat flux at the next timestep. Right: Essential building blocks and integrator layers that enable multitask training. The latent 5D space is integrated over velocity space to obtain latent 3D field for potential prediction via cross-attention.
in designing such a surrogate model is scalability as nonlinear gyrokinetic equations can span mul- tiple turbulence modalities for different species at varying resolutions. Furthermore, downstream physical quantities, such as electrostatic potential fields and scalar fluxes should be consistent with the predicted 5D field. Finally, understanding the structure of the 5D field enables us bake induc- tive biases into the model architecture that improve capturing nonlinear dynamics, i.e., zonal flows. Based on these observations we pose the following desiderata for designing GyroSwin.
1. Scalability. We identify two main candidates: (1) Convolution baased (Fukushima, 1980; Le- Cun et al., 1989, CNN) FNOs (Li et al., 2021). and (2) Vision Transformers (Dosovitskiy et al., 2021, ViT). To preserve locality of the 5D field, CNNs or FNOs require factorized im- plementations (Wang al., 2017; Tran et al., 2023) that do not scale well (see Table 3). ViTs are in principle applicable to 5D. however flatteming 5D field results in extremely long patch sequences. Due their quadratic complexity (Vaswani et al., 2017), ViTs do not scale well to high resolution input. Hierarchical processing with linear attention, as in Swin (Liu et al., 2021), provides an efficient way to preserve locality via local window attention and was successfully applied to 3D (Liu et al., 2022) or 4D (Kim et al., 2023) input. Finally, patch embeddings can be compressed lower resolution to facilitate scalability.
2. Modelling latent integrals. Electrostatic potentials and heat fluxes computed integrals (Equation (2)). We replicate this operation in the latent space of GyroSwin. To obtain the 3D latent fields, we introduce latent integrator module that aggregates over the velocity space. Similarly, taking inspiration from Equation (2), the flux prediction is based on cross-attention pooling of 3D and 5D latents.
3. Inductive biases. The main benefit of GyroSwin over ROMs is that nonlinear dynamics can be captured that severely affect emerging turbulence. To bias GyroSwin towards nonlinear zonal flows, we perform channelwise mode separation. Specifically, we separate the zonal flow mode from Equation (4) from the other modes and transfer the spectral space real space to retain the same dimensions. Then add real and imaginary parts of the zonal flow additional channel.
Taking into account these requirements, we design GyroSwin as Swin-based UNet (Ronneberger et al., 2015) with multiple branches to accommodate multitask predictions, and physical priors. The following paragraphs explain the main architectural components of GyroSwin. We describe them in the domain-specific case of gyrokinetics with 5D shifted Window Attention (5DWA) and 5D) up/downsampling layers, initially introduced by Galletti et al. (2025). All of them are generalized to n-dimensions through adaptive window partitioning and local convolution.
5D shifted window attention. The main characteristic of Swin Transformer (Liu et al., 2021) local attention in fixed window sizes. This has significant advantage over ViTs, reduc- ing the quadratic complexity in token count to near-linear. More information on complexity for Swin and VIT can be found in Appendix H.I. The core component of GyroSwin is 5DWA (see Figure 2b). In 5DWA, attention is performed across tokens within 5D windows of size
5

---

M == Mo x Ms, x M, x M2 x My. Let X R x xaxxxyxd be 5D field of batch size and hidden dimension d. After partitioning X into non-overlapping windows, we obtain X E ER. Nu May Map Mr My) xd with Nwin = l/l / " -,n/Ma,,... ... denoting the num- ber of windows. Self-attention is performed within each window in parallel across all windows. To enable interaction across neighboring windows we stack Window Multi-Head Attention (W-MSA) and Shifted Window MSA (SW-MSA) blocks. This results in approximate global attention with increasing depth. We define W-MSA for any n-dimensional
(X (Xw
XX WK)T Vd
M-MSA(X) = W SoftMax
(X w
(6)
hEheads
where : W W are the head-wise query, key and value projection matrices for head h, and Wo is the output matrix. The shifted version SW-MSA is defined as spatially shifting the window partition stencil prior to Equation (6), and reversing the shift after
M ) .t)+/ 2 .+2)
M
SW (W-MSA(x) = roll W-MSA roll x,
(xx=2
(7)
2
where roll (X, + is matrix roll operation, which circularly shifts each dimension of æ by + M the respective dimension. For details on batched implementation the cyclie shift we refer to Liu et al. (2021) and Kim et al. (2023).
Up/Downsampling, Following (Dosovitskiy et al., 2021), we employ patch embedding to split the input into non-overlapping patches via 5D local convolution and embed them into tokens. The inputs are then spatially downsampled by the patch size and channels are projected to dimension C. The downsampling path of the UNet interleaves 5DWA blocks with patch merging layers to compress the 5D field at each stage. We extend these patch merging layers from Liu et al. (2021) to 5D: they concatenate features 12x2x2x2x neighboring patches and apply an MLP, reducing resolution by 2 times. The output dimension is set to 2C, which doubles the channels after each merge. We store intermediate feature maps for skip connections to produce hierarchical representations at different resolutions, resulting in growing receptive field and lower computational cost lower stages. For upsampling we employ patch expansion layers, reversing the patch embedding (and merging) via a linear projection and rearrangement to the original field size. Figure 2a sketches the architecture of GyroSwin, incorporating patch anbbedding/merging/epnnssss
blocks along with SDWA.
Multitask training. To ensure that GyroSwin adheres to downstream integrated quantities, such as electrostatic potential fields & and scalar flux Q, we add prediction tasks for each of them in addition to predicting f. The ø-head is constructed of L blocks, each of which employs an integral layer to reduce the 5D latent to 3D latent followed by 3DWA layers and expansion layers until the original spatial resolution is recovered. For the Q-head, we perform max pooling over the integrated 3D space followed by MLP head. Importantly, we also employ latent 5D < 3D mixing layers each block to facilitate latent communication among the different heads. For multitask training of GyroSwin, we introduce loss terms for each head with weighting factors:
L= == # (gL + ue (Q)
5D++3D mixing and integrator modules. Multitask training on 3D potential fields requires re- ducing the 5D latent space to a latent 3D space. To this end, we incorporate latent integrator mod- ules (see fyspace block in Figure 2b). They perform wross-attention between learnable ID query Q R. Ixd, which is broadcasted over the velocity space of K, V (B(bba-ryy) x xaxd This pro- cedure is reminiscent of perceiver-style pooling (Jaegle et al., 2021) with single fixed query across velocity dimensions. Moreover, we allow interaction between 3D and SD latents via cross-attention in each block L, where Q e xxxxxxxxyxd and K, V example of latent 3D + 5D interaction. by swapping Q with K and V. we also obtain 5D - 3D cross-attention. We stack both cros-attention layers to obtain our 51.DMix bloks.
Figure 2b shows an
4 Experiments
In this section we elaborate on our experimental setup, ranging from data generation to baselines and our evaluation setup. For implementation details see Appendix D.

---

Data Generation. We run nonlinear simulations using the numerical code GKW (Peeters et al., 2009), varying noise amplitude of initial conditions and four operating parameters, namely the safety factor q, the magnetic shear 6, the ion temperature gradient R/ I, and the density gradient R/ T... To reduce the computational burden of data generation, we consider the adiabatic electron approxima- tion at a resolution of (32 x 8 x 16 x 85 x 32), i.e. we only consider turbulence originating from ion temperature gradients. We use latin hypercube sampling to uniformly populate the parameter space (McKay et al., 1979) and chose parameter regions to ensure that the resulting simulations are highly turbulent. The distribution of operating parameters and corresponding heat flux can be observed in Figure 5. We run each simulation for total of 31,920 steps, which averaged every 40 steps and subsampled every third step, resulting in a total of 266 snapshots. Each snapshot comes with two channels, representing the real and imaginary parts of the ballooning transform, commonly used for plasma coordinates. We neglect the first 80 snapshots of each simulations as those correspond to the linear phase where turbulence is not fully established yet. The entire dataset comprises 255 simula- tions, based on which we assemble two training subsets, one comprising 48 simulations and another one comprising 241 simulations. We use the small subset for comparison to baselines and the latter for scaling experiments. In total, this results in 44,585 training samples, of which 8880 are used for the small subset. We provide visualizations of sample snapshots of the 5D fields in Appendix C.
Baselines. We implement three types of baselines: (i) As state-of-the-an art ROM, we implement the QuaLiKiz saturation rule (Bourdelle et al., 2015) and fit it to nonlinear fluxes of our training set, similarly Kumar et al. (2021). For details, see Appendix B. (fi) Tabular regressors, such as Gaussian Process Regression (GPR) for nonlinear fluxes as proposed by (Hornsby et al., 2024), and a MLP trained in the same manner, akin to van de Plassche et al. (2020). (iii) Neural surrogates, trained to predict the 5D density function: Fourier Neural Operator (Li et al., 2021, FNO), PointNet (Qi al., 2016). Transolver (Wu et al., 2024), and vanilla ViT (Dosovitskiy et al., 2021). Additional information on baselines can be found in Appendix D.1.
Evaluation. The promise of GyroSwin is that it can replace QL approaches as it is trained on the full 5D distribution function, but maintains efficiency and scalability. To properly evaluate whether our GyroSwin yields improvements over QL models, we compile a set of nonlinear simulations in a high ion temperature gradient regime, ensuring strong turbulence regime. We set aside 14 of the 255 simulations that we generate in total to evaluate for in-distribution (ID) and (OOD) generalization. For the ID set, we identify region in the 4D parameter set that lives within the convex hull of the training set, but is unseen to the model. Conversely, we ensure that the 0OD set is not governed by the convex hull of the training simulations. In total, we compile select six simulations for the ID set and five simulations for the OOD set. The remaining three simulations are used as a validation set during training. All of them are excluded from the training set.
5 Results
As shown in Table different methods are restricted to certain evaluations, i.e. tabular regressors can only be evaluated on nonlinear flux prediction (scalars), while QL models can additionally be used to evaluate for diagnosties. Neural surrogates that predict the full SD field, such as GyroSwin, are the only class that can also be evaluated for zonal flows. In line with this observation, we provide results for each of the three categories based on their capabilities.
Table 1: Comparison of different surrogate approaches by capabilities.
Method
Average Flux
Diagnostics
Zonal Flows
Turbulence
Tabular Regressors, e.g. GPR, MLP
1D ID--00
X
X
X
SOTA Reduced Numerical modelling, e.g., QL
3D 3D-000
3D>+10
X
X
Neural Surrogates, e.g. GyroSwin (Ours)
5D--0D
5D-+ID
5D-11D
5D-5D
5D 5D-55 Turbulence modelling, To evaluate for 5D turbulence modelling capabilities, we perform an autoregressive rollouts with the neural surrogates and measure correlation time. Following Alkin et al. (2024a), we define correlation time as the number of snapshots that can be predicted while maintaining certain pearson correlation т. This metric, demonstrates what methods are capable of performing stable autoregressive rollouts without drifting too far from the ground truth. We report
7

---

n/a n/a Table 2: Evaluation for 5D turbulence modelling and nonlinear heat flux prediction. We evaluate all methods across six in-distribution (ID) and five out-of-distributon (OOD) simulations. For Q we report RMSE of tvme-averaged predictions after an autoregressive rollout. For I we report correlation time for autoregressive rollouts with threshold = 0.1. Higher correlation time is better.
Method
Input
(f)
OOD (t)
ID (
OOD (L)
SOTA Reduced Numerical modelling
QL (Bourdelle et al., 2007)
3D
n/a
n/a
89.48
11.77
95.18 21.58
Classical Regression Techniques
GPR (Hornsby al., 2024)
0D
n/a
43.82 = 10.84
59.28 + 17.55
MLP
0D
n/a
50.90 = 10.87
62.88 +: 18.58
Neural Surrogate Models (48 simalations)
FNO (Li et al., 2021)
3D
9.33 0.56
9.20 1: 0.58
119.88 + 13.15
124.96
23.27
PointNet (Qi et al., 2016)
5D
8.5 1.26
9.0 0.89
119.90 + 13.16
125.01
= 23.28
Transolver (Wu al., 2024)
5D
9.66 0.94
10.2
119.92 + 13.15
125.02
23.28
VIT (Desovitskiy et al., 2021)
5D
16.83 + 1.49
19,20 1.36
119.63 + 13.13
125.13 = 23.29
GyroSwin (Ours)
5D
26.50 3.55
28.60 8.82
67.68 10.28
70.48
17.21
Scaling GyroSwin to 241 simulations
GyroSwinssall (Ours)
5D
98.00 + 27.53
76.40 + 17.60
23.72 + 4.05
53.54 + 18.10
GyroSwinAtdum (Ours)
5D
94.17 + 21.96
91.20 + 18.61
37.24 1 9.60
44.17 + 17.68
GyroSwin (Ours)
5D
110.33 + 19.74
111.80 = 23.86
18.35 + 1.56
26.43 = 9.49
correlation times for the neural surrogates in Table 2. We observe a clear trend that GyroSwin is by far the most stable autoregressive method compared to other neural surrogates. Remarkably, when scaling GyroSwin in terms of data and model size, we attain stable rollouts for over 100 timesteps, even for OOD simulations.
5D OD average flux. The task for this evaluation is predict the average heat flux Q over the last 80 timesteps of simulation from the full 5D field after an autoregressive rollout. We present results for both ID and OOD evaluation sets in Table 2. On the reduced training set GyroSwin yields the best performance compared to the QL model and alternative 5D neural surrogates. Interestingly, most neural surrogates converge to similar heat flux Q of around 120. This value is obtained if the model diverges and the flux predictions for each simulation are consistently zero, which would be the case if no turbulence occurs all. This indicates that all competitors fail to capture turbulent transport as they do not predict any turbulent flux transport after longer rollouts, leading to similar errors for Q. Furthermore, when scaling GyroSwin to more data and larger model sizes, we observe a drastic improvement in nonlinear flux prediction, achieving significantly lower error than currently existing surrogates. In Appendix G we also show that the composition of all components of GyroSwin results in the best performance.
Scalability. In integrated plasma simulations, turbulent heat fluxes must be repeatedly computed across thousands of runs, radial locations, and time intervals, underscoring the need for scalable sur- rogate architectures. To assess scalability of neural surrogates, we report inference speed, memory consumption, and number of parameters on a single NVIDIA H100 80GB HBM3 in Table We exclude the 3D FNO as it is based on collapsing velocity dimensions into channels (see Appendix D) which is inherently unscalable to higher resolutions. Factorized variants of FNO and CNN are also slow and memory heavy. Field-based surrogates (PointNet, Transolver) require subsampling during training and chunked inference, limiting scalability. ViTs suffer from quadratic complexity and sim- ilar memory consumption as GyroSwin despite using half of its parameters. Scaling VIT to the same parameter count as GyroSwin results ~ 18ms inference speed which is 52.5% slower than GyroSwin PeroSwinguann Small- GyroSwin, small is roughly three orders of magnitude faster than the numerical code GKW (4200 vs. 756 GFLOPs). When scaled to ~BB parameters (Figure 3), GyroSwin continues to improve training and validation loss for both f and o prediction, demonstrating superior scalability and strong potential for scaling to higher-fidelity simulations.
5D>11 flux spectrum. A major advantage of quasilinear approaches over tabular regressors physical verifiability. As they are based on 3D linear simulations, flux contributions for each mode
8

---

GyroSwinsmal
Gyr
GyroSwinuarge
Train
Validation
Relative Norm MSE
Distribution function
Potential Field
10"
102
10
10
10-1
10-1
10-2
100 200
300
400
500
0
50
100
150
200
250
Number of epochs
Number of epochs
Figure 3: Scaling GyroSwin to ~ IB parameters trained on 241 simulations amounting to approxi- mately 6TB of data. We show train/validation error for predicting the 5D distribution function (Ieft) and the 3D electrostatic potential field (right).
CT
GyroSwinuarse
Quasilinear
Transolver
PointNet
VIT
104
1.0
102
Amplitude
0.5
w(kk)
0.0
10
-055
10-2
-1.0
10
10
0
10
20
30
Figure 4: Left: W(ky) averaged over time and OOD simulations for different 5D neural surrogates. Competitors tend to underestimate while GyroSwin matches the spectrum well with slight dis- crepancy on higher frequencies. Right: Tvme-averaged zonal flow profile for slice along s across radial coordinates x for a selected OOD simulation. GyroSwin captures the zonal flow profile.
in by (Q(ke))) from Equation (3) can be inspected. This provides insights as to whether modes are captured that contribute most to heat transport. Similarly, we assess whether the general structure of Q(ky) is captured in predicted 5D snapshots of neural surrogates. We report pearson correlation to the ground-truth Q(k,,) for all neural surrogates in Table We observe that all neural sur- rogates exhibit higher correlation than the QL approach highlighting their potential. Furthermore, we demonstrate that GyroSwin exhibits the best correlation on the OOD test cases. However, performs slightly worse the ID cases than competitors. When scaling GyroSwin the large training set, we find boost in performance, resulting in almost perfect reproduction of the shape of Q(ky). This means that GyroSwin, Large captures the energy transport por-mode almost perfectly. In Figure 10 visualize Q(ky) for all ID well OOD cases.
As mentioned in Section 2 another useful diagnostic is the turbulence intensity spectrum (k.,)) which is not necessarily aligned with (.,y). i.c the mode transporting the most energy might not be the same where turbulence is most intense. Therefore we visually showcase W(k,) Figure (left) averaged across time and all OOD cases. In addition we provide visualizations for each test case of ID and OOD separately in Figure We observe that while all neural surrogates decently reproduce the shape of the spectrum, they heavily underestimate the magnitude. Furthermore, GyroSwin, aagge matches shape and magnitude almost perfectly, up until few higher frequency components. We attribute this finding to the spectral bias of neural networks towards lower frequency (higher energy) modes (Rahaman et al., 2019).

---

Table 3: Evaluation for scalability, diagnostics, and nonlinear physics. We evaluate all methods across six in in-distribution (ID), and five out-of-distributio (OOD) simulations. For diagnostics and nonlinear physies we report time-averaged pearson correlation for autoregressive rollout. GyroSwin is the only method that can be scaled to ~1B parameters while maintaining reasonable inference speed and improving correlation to diagnostics.
(kky)
Method
Fwd[ms]
Mem[GB]
Params[M]
ID (1)
OOD (1)
SOTA Reduced Nunerical modelling
QL (Bourdelle al., 2007)
N/A
NIA
N/A
0.51 - 0.4
0.58 0.33
Neural Surrogate Models (48 simulations)
F-CNN
569.6
11.1
2.0
N/A
N/A
F-FNO
963.3
36.9
1.3
N/A
N/A
PointNet (Qi ul., 2016)
19.7
39.7
43.9
0.49 + 0.02
0.54 0.09
Transolver (Wu et al., 2024)
3103
35.8
40.6
0.50 + 0.02
0.53 + 0.09
VIT (Dosovitskiy al., 2021)
6.3
2.4
46.1
0.52 + 0.05
0.63 0.11
GyroSwin (Ours)
11.8
90.2
0.59 + 0.06
0.69 0.06
Scaling GyraSwin 241 simulations
GyraSwinsmall (Ours)
11.8
2.8
90.2
0.82 + 0.10
0.79 + 0.09
GyroSwin Medium (Ours)
12.1
5.3
250.9
0.72 + 0.08
0.76 + 0.12
GyroSwin Lare (Ours)
15.4
998.3
0.87 + 0.11
0.91 0.08
New capability: 5D zonal flow modelling. Zonal flows have a significant impact on the turbulence dynamics. To evaluate whether GyroSwin is capable capturing zonal flows, we visualize the zonal flow profile of test case of the OOD set. We again time-average the profiles and compare to the time-averaged predicted zonal flow profile of GyroSwin Large" We find that GyroSwin Large accurately captures the zonal flow profile. This capability has so far been unreachable by any other surrogate modelling technique. We provide additional visualizations for each test case of the ID and OOD sets in Figure 11.
6 Conclusions and Limitations
We present GyroSwin, scalable neural surrogate model for nonlinear gyrokinetic equations mod- elling turbulent transport in Plasmas. Unlike existing surrogate models, GyroSwin operates directly 5D space and evolves the 5D distribution function of gyrokinetics. It is based on latent cross- attention and integration modules for 3D+5D latent interaction and trained in multitask fashion predict the distribution function, electrostatic potentials and heat fluxes. Furthermore, we perform channelwise mode separation to incorporate an inductive bias towards essential nonlinear phenom- ena (zonal flows) that are essential for modelling turbulence. We show that GyroSwin outperforms reduced numerical appraoches and scales favorably compared to other neural surrogates. Further- more, GyroSwin accurately captures nonlinear phenomena self-connistently. capability not present in prior surrogate reduced numerical models. As result, GyroSwin offers fruitful alternative to efficient approximation of turbulent transport.
Currently, the main limitation GyroSwin that it does not take into account the chaotic and therefore distributional nature turbulence. Although we find that GyroSwin produces stable au- toregressive rollouts over 100 timesteps, it suffers from eтror accumulation. Generative modelling offers remedy to this problem by directly predicting snapshots in the saturated phase. We aim to incorporate such distributional approaches in future work. Furthermore, we neglect the linear phase of the simulation. The reason for this is that our focus primarily lies in modelling turbulence. In the future we aim to extend GyroSwin to the linear phase as well. Finally, we only consider the adiabatic electron approximation, as allows generation of a relatively large training set at rather low cost. Still, each simulation produces considerable data volumes, making full coverage of the 4D parameter space difficult. Beyond that, we envision a high-fidelity surrogate model by transfer learning from low-fidelity approximate simulations.
10

---

Acknowledgments and Disclosure of Funding
We extend special thanks our colleagues Wei Lin for the discussion on computer vision and multi- task learning, and Anna Zimmel and Gerald Gutenbrunner for conversations on the spectral behavior of neural networks.
This work has been funded by the Fusion Futures Programme. As announced by the UK Gov- ermment in October 2023, Fusion Futures aims provide holistic support for the development of the fusion sector. The ELLIS Unit Linz, the LIT AT Lab, the Institute for Machine Learn- ing, are supported by the Federal State Upper Austria. We thank the projects FWF AIRI FG 9- N (10.55776/FG9), A[4GreenHcatingGrio (FFG- 899943), Stars4Waters (HORIZON-CL6-200) CLIMATE-01-O11 FWF Bilateral Artificial Intelligence (10.5577//CEE12) We acknowledge Eu- roHPC Joint Undertaking for awarding us access to Leonardo at CINECA, Italy, and Deucalion at MACC, Portugal.
References
Josh Abramson, Jonas Adler, Jack Dunger, Richard Evans, Tim Green, Alexander Pritzel, Olaf Ronneberger, Lindsay Willmore, Andrew J Ballard, Joshua Bambrick, et al. Accurate structure prediction of biomolecular interactions with alphafold 3. Nature, pp. 1-3, 2024.
Benedikt Alkin, Andreas Fürst, Simon Schmid, Lukas Gruber, Markus Holzleitner, and Johannes Brandstetter. Universal physics transformers. CoRR, abs/2402.12365, 2024a. doi: 10.48550/
1RX1V.2402.1265.
Benedikt Alkin, Tobias Kronlachner, Samuele Papa, Stefan Pirker, Thomas Lichtenegger, and Johannes Brandstetter. Neuraldem-real-ting simulation of industrial particulate flows. arXiv preprint Ttt::2411.09672. 2024b.
Kaifeng Bi, Lingxi Xie, Hengheng Zhang, Xin Chen, Xiaotao Gu, and Qi Tian. Accurate medium- range global weather forecasting with 3d neural networks. Nat., 619(7970):533-5388 2023. doi: 10. 10.1038/S41586-02. 023-666----------3.
C. K. Birdsall and A. B. Langdon. Plasma physics via computer simulation. New York: Taylor and Francis, first edition, 2005.
Cristian Bodnar, Wessel P. Bruinsma, Ana Lucie, Megan Stanley, Johannes Brandstetter, Patrick Garvan, Maik Riechert, Jonathan A. Weyn, Haiyu Dong, Anna Vaughan, Jayesh K. Gupta, Kit Thambiratnam, Alex Archibald, Elizabeth Heider, Max Welling, Richard E. Turner, and Paris Perdikaris. Aurora: A foundation model of the atmosphere. CoRR, abs/2405. 13063, 2024. doi: 0048550/ARXIV.2405..00
C. Bourdelle, X. Garbet, F. Imbeaux, Casati, N. Dubuit, Guirlet, and Parisot. new gyrokinetic quasilinear transport model applied to particle transport in tokamak plasmas. Physics of Plasmas, 14(11):11250)
2007. ISSN 1070-664X. doi: 10. (063/1.280066)
C. Bourdelle, A. Casati, X. Garbet, F. Imbeaux, J. Candy, F. Clairet, G. Dif-Pradalier, G. Falchetto, T. Gerbaud, V. Grandgirard, P. Hennequin, R. Sahot, Y. Sarazin, L. Vermare, and R. E. Waltz. Validity of quasi-linear transport model. Proceedings of the 22nd IAEA Fusion Energy Conference, pp. 227, Vienna, Austria, 2008. International Atomic Energy Agency, Paper TH/P8- 7.
C Bourdelle, J Citrin, B Baiocchi, A Casati, P Cottier, X Garbet, and F Imbeaux and. Core turbulent transport in tokamak plasmas: bridging theory and experiment with QuaLiKiz. Plasma Physics and Controlled Fusion. 58(1):014036, December 2015. doi: 10.108807741-33355511100000
Johannes Brandstetter. Envisioning better benchmarks for machine learning pde solvers. Nature Machine Intelligence, 1-2, 2024.
Steven L.. Brunton, Bernd R. Noack, and Petros Koumoutsakos. Machine learning for fluid mechan- ics. Annual Review of Fluid Mechanies, 52(Volume 52, 2020):477-508, 2020. ISSN 1545-4479. doi: https://doi.org/10.I

---

J Citrin, C Bourdelle, F J Casson, C Angioni, N Bonanomi, Y Camenen, X Garbet, L Garzotti, Т Görler, Gürean, Koechl, F Imbeaux, Linder, van de Plassche, Strand, and G Szepesi and. Tractable flux-driven temperature, density, and rotation profile evolution with the quasilinear gyrokinetic transport model Qual Plasma Physics and Controlled Fusion, 59(11):124005 November 2017. doi:
1188//3366-65//aa8aeeh
J. Citrin, P. Trochim, T. Goerler, D. Pfau, K. L. van de Plassche, and F. Jenko. Fast transport simulations with higher-fidelity surrogate models for ITER. Physics of Plasmas, 30(6), jun 2023. doi: N0.1063/5.0136755
Jonathan Citrin, Ian Goodfellow, Akhil Raju, Jeremy Chen, Jonas Degrave, Craig Donner, Federico Felici, Philippe Hamel, Andrea Huber, Dmitry Nikulin, David Pfau, Brendan Tracey, Martin Riedmiller, and Pushmeet Kohli. Torax: A fast and differentiable tokamak transport simulator in jax, 2024.
A. M. Dimits, G. Bateman, M. A. Beer, B. I. Cohen, W. Dorland, G. W. Hammett, C. Kim, E. Kinsey, M. Kotschenreuther, A. H. Kritz, L.. L. Lao, J. Mandrekas, W. M. Nevins, S. H. Parker, Sydora, and Weiland. Comparisons and physics basis of tokamak transport models and turbulence simulations. Physics of Plasmas, 733)969-988,, 03 2000a. ISSN 1070-664X. doi: 10.1063/1.87330
A. J. Redd, D. E. Shumaker,
A. M. Dimits, G. Bateman, M. A. Beer, B. I. Cohen, W. Dorland, G. W. Hammett, C. Kim, J. E. Kinsey, M. Kotschenreuther, A. H. Kritz, L. Lao, J. Mandrekas, W. M. Nevins, S. E. Parker, A. J. Redd, D. E. Shumaker, R. Sydora, and J. Weiland. Comparisons and physics basis of tokamak transport models and turbulence simulations. Physics of Plasmas, 7(3):669-995 983, 03 2000b. ISSN 1070-664X. doi: 00.1063/1.87380
Alexey Dosovitskiy, Lucas Beyer, Alexander Kolesnikov, Dirk Weissenborn, Xiaohua Zhai, Thomas Unterthiner, Mostafa Dehghani, Matthias Minderer, Georg Heigold, Sylvain Gelly, Jakob Uszko- reit, and Neil Houlsby. An image worth 16x words: Transformers for image recognition scale. In 9th International Conference on Learning Representations, ICLR 2021, Virtual Event, Austria, May 3-7, 2021. OpenReview.net, 2021.
Kunihiko Fukushima. Neocognitron: A self-organizing neural network model for mechanism of pattern recognition unaffected by shift in position. Biological cybernetics, 36(4):193-200 1980.
Gianluca Galletti, Fabian Paischer, Paul Setinek, William Hornsby, Lorenzo Zanisi, Naomi Carey, Stanislas Pamela, and Johannes Brandstetter. 5d neural surrogates for nonlinear gyrokinetic sim- ulations of plasma turbulence, 2025.
Mitsuru Honda, Emi Narita, Shinya Maeyama, and Tomo-Hiko Watanabe. Multimodal convo- lutional neural networks for predicting evolution of gyrokinetic simulations. Contributions to Plasma Physics, 63(5-6)xe00220013) 2023. doi: http::/doi..rr// 10. 0002/ctpp.202200133
W. A Hornsby, A. Gray, J. Buchanan, B. S. Patel, D. Kennedy, F. J. Casson, C. M. Roach, M. B. Lykkegaard, H. Nguyen, N. Papadimas, B. Fourcin, and J. Hart. Gaussian process regression models for the properties of micro micro-tearing modes in spherical tokamaks. Physies of Plasmas, 31 (1), jan 2024. ISSN 1089-7674. doi: 10. 10.1063/5.01744/)
K. Itoh, S.-I. Itoh, P. H. Diamond, T. Hahm, A. Fujisawa, G. R. Tynan, M. Yagi, and Y. Nagashima. Physies of zonal flowsa). Physics of Plasmas, 13(5):055502, 05 2006. ISSN 1070-664X. doi: 00.1063/1.2178779
Andrew Jaegle, Felix Gimeno, Andy Brock, Oriol Vinyals, Andrew Zisserman, and João Carreira. Perceiver: General perception with iterative attention. In Marina Meila and Tong Zhang (eds.). Proceedings of the 38th International Conference on Machine Learning, ICML 2021, 18-24 July 2021, Virtual Event, volume 139 of Proceedings of Machine Learning Research, pp. 4651-4664 PMLR, 2021.
John Jumper, Richard Evans, Alexander Pritzel, Tim Green, Michael Figurnov, Olaf Ronneberger, Kathryn Tunyasuvunakool, Russ Bates, Augustin Žídek, Anna Potapenko, al. Highly accurate protein structure prediction with alphafold. nature, 596(7873):583-58)) 2021.
12

---

C.K. Kiefer, C. Angioni, G. Tardini, N. Bonanomi, B. Geiger, P. Mantica, T. Pütterich, E. Fable, P.A. Schneider, ASDEX Upgrade Team EUROfusion MSTI Team and JET Contributors Validation of quasi-linear turbulent transport models against plasmas with dominant electron heating for the prediction of iter pfpo-1 plasmas. Nuclear Fusion, (6):0066,,, may 2021. doi: 10. 0108//1741-336//bbcce
Peter Kim, Junbeom Kwon, Sunghwan Joo, Sangyoon Bae, Donggyu Lee, Yoonho Jung, Shin- jae Yoo, Jiook Cha, and Taesup Moon. Swift: Swin 4d fmni transformer. Advances in Neural Information Processing Systems, 06:42011-420337 2023.
Diederik P. Kingma and Jimmy Ba. Adam: A method for stochastic optimization. In Yoshua Bengio and Yann LeCun (eds.), 3rd International Conference Learning Representations, ICLR 2015, San Diego, CA, USA, May 7-9, 2015, Conference Track Proceedings. 2015.
Nikola B. Kovachki, Zongyi Li, Burigede Liu, Kamyar Azizzadenesheli, Kaushik Bhattacharya, Andrew M. Stuart, and Anima Anandkumar. Neural operator: Learning maps between function spaces with applications to pdes. J. Mach. Learn. Res., 24:89: 188997,, 2023.
John A. Krommes. The gyrokinetic description of microturbulence in magnetized plasmas. Annual Review of Fluid Mechanics, 44(Volume 44, 2012):175-201 2012. ISSN 1545-4479. doi: https: org//. 146/annurev-fuud- 1207 207100101223.
N. Kumar, Y. Camenen, S. Benkadda, C. Bourdelle, A. Loarte, A.R. Polevoi, F. Widmer, and JET contributors. Turbulent transport driven by kinetic ballooning modes in the inner core of jet hybrid h-modes. Nuclear Fusion, 61(3):036005, jan 2021. doi: 10. 101100///411432//aadtt
Thorsten Kurth, Shashank Subramanian, Peter Harrington, Jaideep Pathak, Morteza Mardani, David Hall, Andrea Miele, Karthik Kashinath, and Anima Anandkumar. Fourcastnet: Accelerating global high-resolution weather forecasting using adaptive fourier neural operators. In Proceedings of the Platform for Advanced Scientific Computing Conference, PASC '23, New York, NY, USA, 2023. Association for Computing Machinery. ISBN 9798400701900. doi: 10.1145/3592979. 3593412.
Remi Lam, Alvaro Sanchez-Gonzalez, Matthew Willson, Peter Wirnsberger, Meire Fortunato, Fer- Alet, Suman Ravuri, Timo Ewalds, Zach Eaton-Rosen, Weihua Hu, Alexander Merose, Stephan Hoyer, George Holland, Oriol Vinyals, Jacklynn Stott, Alexander Pritzel, Shakir Mo- hamed, and Peter Battaglia. Learning skillful medium mrangee---angeeeeee global weather forecasting. Science, (82(6677):1416-1444 2023. doi: 10. 122//ciiene.aai23333..
Yann LeCun, Bernhard E. Boser, John S. Denker, Donnie Henderson, Richard E. Howard, Wayne E. Hubbard, and Lawrence D. Jackel. Backpropagation Applied to Handwritten Zip Code Recogni- tion. Neural Comput., 144)5541-55, 1989. doi: 10. 166/necoo..o. 1989.1.4.444
Zongyi Li, Nikola B. Kovachki, Kamyar Azizzadenesheli, Burigede Liu, Kaushik Bhattacharya, Andrew M. Stuart, and Anima Anandkumar. Neural operator: Graph kernel network for partial differential equations. CoRR, abs/2003.03488 2020.
Zongyi Li, Nikola Kovachki, Kamyar Azizzadenesheli, Burigede Liu, Kaushik Bhattacharya, An- drew Stuart, and Anima Anandkumar. Fourier neural operator for parametric partial differential equations, 2021.
Ze Liu, Yutong Lin, Yue Cao, Han Hu, Yixuan Wei, Zheng Zhang. Stephen Lin, and Baining Guo. Swin transformer: Hierarchical vision transformer using shifted windows. In 2021 IEEE/CVF International Conference on Computer Vision, ICCV 2021, Montreal, OC, Canada, October 10-17, 2021, pp. 9992-10002. IEEE, 2021. doi: 10. 10.100/CCC4499222.000 202.00099
Ze Liu, Jia Ning, Yue Cao, Yixuan Wei, Zheng Zhang, Stephen Lin, and Han Hu. Video swin transformer. In IEEE/CVF Conference on Computer Vision and Pattern Recognition, CVPR 2022, New Orleans, LA, USA, June 18-24, 2022, pp. 3192-3201. IEEE, 2022. doi: 10. 2VPPR52688.0022.0022
M. D. McKay, R. J. Beckman, and W. J. Conover. A comparison of three methods for selecting values of input variables the analysis of output from a computer code. Technometrics, 21(2): 239-245, May 1979, ISSN 0040- 1706. doi: 10.2307/1268522 OSTI 5236110.
13

---

Amil Merchant, Simon L. Batzner, Samuel S. Schoenholz, Muratahan Aykol, Gowoon Cheon, and Ekin Dogus Cubuk. Scaling deep learning for materials discovery. Nat., 624(79900::0--3)) 2023. doi: 10. (8888//41586-023-667777
S. Van Mulders, F. Felici, O. Sauter, J. Citrin, A. Ho, M. Marin, and K.L. van de Plassche. Rapid optimization of stationary tokamak plasmas in RAPTOR: demonstration for the ITER hybrid scenario with neural network surrogate transport model QLKNN. Nuclear Fusion, 61 (1))0886019, July 2021. doi: 10. (088/1741-4426/aall)
E. Narita, M. Honda, S. Maeyama, and T.-H. Watanabe. Toward efficient runs of nonlinear gyroki- netic simulations assisted by convolutional neural network model recognizing wavenumber- space images. Nuclear Fusion, 62(8):086037, jun 2022. doi: 10.1088/1741 431/////022/////
Tung Nguyen, Johannes Brandstetter, Ashish Kapoor, Jayesh K. Gupta, and Aditya Grover. Climax: A foundation model for weather and climate. In Andreas Krause, Emma Brunskill, Kyunghyun Cho, Barbara Engelhardt, Sivan Sabato, and Jonathan Scarlett (eds.), International Conference Machine Learning. ICML 2023, 23-29 July 2023, Honolulu, Hawaii, USA, volume 202 of Proceedings of Machine Learning Research, pp. 25904-25938, PMLR, 2023.
William Peebles and Saining Xie. Scalable diffusion models with transformers. In JEEE/CVE International Conference on Computer Vision, ICCV 2023, Paris, France, October 1-6, 2023, pp. 4172-4182, IEEE, 2023, doi: 10. (O09//CCV51000.22333..... 00387.
A.G. Peeters, Y. Camenen, FJ Casson, W.A. Hornsby, A.P. Snodin, D. Strintzi, and G. Szepesi. The nonlinear gyro-kinetic flux tube code gkw. Computer Physics Communications, 180(12): 2650-2672, 2009. ISSN 0010-4655. doi: https://doi. orgiiirr///// 10, 10.00.111111 1016//.cpc. OF CPC: A celebratory issue focused on quality software for high performance, grid and novel computing architectures.
40 YEARS
Ethan Perez, Florian Strub, Harm de Vries, Vincent Dumoulin, and Aaron C. Courville. Film: Vi- sual reasoning with a general conditioning layer. In Sheila A. Mcllraith and Kilian Q. Wein- berger (eds.), Proceedings of the Thirty-Second AAAI Conference on Artificial Intelligence, (AAAI-18), the 30th innovative Applications of Artificial Intelligence (IAAI-18), and the Sth AAAI Symposium Educational Advances in Artificial Intelligence (EAAI- 11)),,)),..0 New Orleans, Louisiana, USA, February 2-7, 2018, pp. 3942-3951. AAAI Press, 2018. doi: 10. A609/AAA V3211.11671.
Charles R Qi, Hao Su, Kaichun Mo, and Leonidas J Guibas. Pointnet: Deep learning point sets for 3d classification and segmentation. arXiv preprint arXiv:1612.00593,, 2016.
Nasim Rahaman, Aristide Baratin, Devansh Arpit, Felix Draxler, Min Lin, Fred A. Hamprecht, Yoshua Bengio, and Aaron C. Courville. On the spectral bias of neural networks. In Kamalika Chaudhuri and Ruslan Salakhutdinov (eds.), Proceedings of the 36th International Conference on Machine Learning. ICML 2019, 9-15 June 2019, Long Beach, California, USA, volume 97 of Proceedings of Machine Learning Research, pp. 5301-5310. PMLR, 2019.
Carl Edward Rasmussen and Christopher K. I. Williams. Gaussian processes for machine learning. Adaptive computation and machine learning. MIT Press, 2006. ISBN 026218253X.
Olaf Ronneberger, Philipp Fischer, and Thomas Brox. U-net: Convolutional networks for biomedi- cal image segmentation. In Nassir Navab, Joachim Hornegger, William M. Wells III, and Alejan- dro F. Frangi (eds.), Medical Image Computing and Computer-Assisted Intervention MICCAI 2015 18th International Conference Munich, Germany, October 5- 9, 2015, Proceedings, Part III, volume 9351 of Lecture Notes in Computer Science, pp. 234 241. Springer, 2015. doi: 10. 10.10077978-3-311 4
Paul Setinek, Gianluca Galletti, Thomas Gross, Dominik Schnürer, Johannes Brandstetter, and Werner Zellinger. Simshift: A benchmark for adapting neural surrogates distribution shifts, 2025.
G. Staebler, C. Bourdelle, J. Citrin, and R. Waltz. Quasilinear theory and modelling of gy- rokinetic turbulet transport in tokamaks. Nuclear Fusion, 64(10):103001, sep 2024. doi:
14

---

G. M. Staebler and J. E. Kinsey. Electron collisions in the trapped gyro-landau fluid transport model. Physies of Plasmas, 17(12). dec 2010. ISSN 1089-7674. doi: (0.1063/1.3505388.
G. M. Staebler, J. E. Kinsey, and R. E. Waltz. A theory-based transport model with comprehensive physics. Physics of Plasmas, 14(5), may 2007. ISSN 1089-7674. doi: 10. 0666/..243682.
Nils Thuerey, Philipp Holl, Maximilian Müller, Patrick Schnell, Felix Trost, and Kiwon Um. Physics-based deep learning. CoRR, abs/2109.05237, 2021.
Alasdair Tran, Alexander Mathews, Lexing Xie, and Cheng Soon Ong. Factorized fourier neural operators, 2023.
David Tskhakaya. The Particle-in Method, pp. 161-189. Springer Berlin Heidelberg, Berlin, Heidelberg, 2008. ISBN 978-3-540-74686-77 doi: 10.
K. L. van de Plassche, J. Citrin, C. Bourdelle, Y. Camenen, F. J. Casson, V. L. Dagnelie, F. Felici, A. Ho, S. Van Mulders, and JET Contributors. Fast modeling of turbulent transport in fusion plasmas using neural networks. Physics of Plasmas, 27(2):022300, 02 2020. ISSN 1070-664X. doi: 10.1063/1.5134122
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N. Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention is all you need. In Isabelle Guyon, Ulrike von Luxburg, Samy Bengio, Hanna M. Wallach, Rob Fergus, S. V. N. Vishwanathan, and Roman Garnett (eds.), Advances in Neural Information Processing Systems 30: Annual Conference on Neural Information Processing Systems 2017, December 4-9, 2017, Long Beach, CA, USA, pp- 5998-6008, 2017.
Chenguang Wan, Youngwoo Cho, Zhisong Qu, Yann Camenen, Robin Varemnes, Kyungtak Lim, Kunpeng Li, Jiangang Li, Yanlong Li, and Xavier Garbet. A high-fidelity surrogate model for the ion temperature gradient (itg) instability using a small expensive simulation dataset. Nuclear Fusion, 65(5):054001, apr 2025. doi: 10.1088/174114422//ad7777
Min Wang, Baoyuan Liu, and Hassan Foroosh. Factorized convolutional neural networks. In 2017 IEEE International Conference on Computer Vision Workshops (ICCVW), pp. 345-553, 2017. doi: 10.1 10.110//CCV..2211..777
Haixu Wu, Huakun Luo, Haowen Wang, Jianmin Wang, and Mingsheng Long. Transolver: A fast transformer solver for pdes on general geometries. In International Conference on Machine Learning, 2024.
Han Yang, Chenxi Hu, Yichi Zhou, Xixian Liu, Yu Shi, Jielan Li, Guanzhi Li, Zekun Chen, Shuizhou Chen, Claudio Zeni, Matthew Horton, Robert Pinsler, Andrew Fowler, Daniel Zügner, Tian Xie, Jake Smith, Lixin Sun, Qian Wang, Lingyu Kong, Chang Liu, Hongxia Hao, and Ziheng Lu. Mattersim: A deep learning atomistic model across elements, temperatures and pressures. arXiv preprint 24ix:2405.00967, 2024.
L. Zanisi, A. Ho, J. Barr, T. Madula, J. Citrin, S. Pamela, J. Buchanan, FJ. Casson, and V. Gopakumar. Efficient training sets for surrogate models of tokamak turbulence with ac- tive deep ensembles. Nuclear Fusion, 64(3):036022, February 2024. ISSN 1741-4326. doi: 10.1088/1741 111188//7711-4422//22400000
Claudio Zeni, Robert Pinsler, Daniel Zügner, Andrew Fowler, Matthew Horton, Xiang Fu, Zilong Wang, Aliaksandra Shysheya, Jonathan Crabbé, Shoko Ueda, et al. A generative model for inor- ganic materials design. Nature, pp. 1-3, 2025.
Xuan Zhang, Limei Wang, Jacob Helwig, Youzhi Luo, Cong Fu, Yaochen Xie, Meng Liu, Yuchao Lin, Zhao Xu, Keqiang Yan, Keir Adams, Maurice Weiler, Xiner Li, Tianfan Fu, Yucheng Wang, Haiyang Yu, Yuqing Xie, Xiang Fu, Alex Strasser, Shenglong Xu, Yi Liu, Yuanqi Du, Alexan- dra Saxton, Hongyi Ling, Hannah Lawrence, Hannes Stärk, Shurui Gui, Carl Edwards, Nicholas Gao, Adriana Ladera, Tailin Wu, Elyssa F. Hofgard, Aria Mansouri Tehrani, Rui Wang, Ameya Daigavane, Montgomery Bohde, Jerry Kurtin, Qian Huang, Tuong Phung. Minkai Xu, Chai- tanya K. Joshi, Simon V. Mathis, Kamyar Arizzadenesheli, Ada Fang, Alán Aspuru-Guzik, Erik J. Bekkers, Michael M. Bronstein, Marinka Zitnik, Anima Anandkumar, Stefano Ermon,
15

---

Pietro Liö, Rose Yu. Stephan Günnemann. Jure Leskovec, Heng Ji, Jimeng Sun, Regina Barzilay, Tommi S. Jaakkola, Connor W. Coley, Xiaoning Qian, Xiaofeng Qian, Tess E. Smidt, and Shui- wang Ji. Artificial intelligence for science in quantum, atomistic, and continuum systems. CoRR, hax/2307.08422. 2023. doi: 10.45550/ARXXV223370022222
16

---

A Derivation of the Gyrokinetic equation
We begin with the Vlasov equation for the distribution function f(r, v,t);
of aL + V EE B) =
m
The Vlasov equation describes the conservation of particles in phase space in the absence of col- lisions. Here, r (xr, y, s) and - (Uz, Vy, Us) correspond coordinates in the spatial and the velocity domain, respectively. Hence the Vlasov equation is a 7D (including time) PDE representing the density of particles in phase space position r, velocity v, and time. The term V. describes the response of the distribution function to accelerations of particles and 2 (E + x B) denotes the Lorentz force, which depends on particle charge q and mass m. well electric field and magnetic field B. Finally, the advection (or convection) term v vV f describes transport of the distri- bution functon through space due velocities.
To derive the gyrokinetic equation, we transform from particle coordinates to guiding center coordi- mu? nates (R, lll p. 0), where f 2 is the magnetic moment, the gyrophase, which describes the position of particle around its guiding center as it gyrates along field line, and R is the coordinate of the guiding center.
Assuming the time scale L at which the background field changes is much longer than the gyroperiod with small Larmor radius p < L, we can gyroaverage remove the dependency on the gyrophase 0, yielding:
0s Ot
0s 0 Oun
R R.
(10)
Linear Terms
The unperturbed (background) motion of the guiding center is governed by:
R = id b + UD
(1I)
in = H b VB
(12)
m
Here, b = B/B is the unit vector along the magnetic field, and up represents magnetic drifts. Substituting into the kinetic equation yields
0f 0t
" "
V 0f = Oep
+ (ia + UD) Vf
(13)
m
We can express the magnetic gradient term using:
B VB
b VB
(14)
so that:
F b V m
#B B VB
(15)
m
B²
Nonlinear Term
Fluctuating electromagnetic potentials 8ø, 8.A induce ExB and magnetic flutter drifts. We define
the gyroaveraged generalized potential
x= (0- 1 A
(16)
where A is the aprallel component of the vector potential. (-) denotes the gyroaverage, and c the speed of light, which is added to ensure correct units. ó is the electrostatic potential, the computation of which involves an integral of f over the velocity space (see eq. 1.41 in the GKW manual for complete description).
17

---

This gives rise to the drift
Vx VX,
(17)
and yields the nonlinear advection term Vx V..
Final Equation
We arrive at the gyrokinetic equation in split form:
o aL
#B B
VB O
- (1lb vD) vs
Vx VS
= S
(18)
Y7t
B
Nonlincar
Linear
Here, SS represents external sources, collisions, or other drive terms. To enhance the tractability of Equation (I), the distribution function f is usually split into equilibrium and perturbation
7 fo
ah
= fo + of fo
+
(19)
T
Öt
Adiabatic
Kinetic
where is background or equilibrium distribution function, T the particle temperature, Z the par- ticle charge, á the electrostatic potential, and 6f the total perturbation to the distribution function, which comprises of adiabatic and kinetic response. The adiabatic term describes rapid and passive responses to the electrostatic potential that do not contribute to turbulent transport, while the ki- netic term governs irreversible dynamics that facilitate turbulence. Numerical codes, such GKW (Peeters et al., 2009), rely on solving for öf instead of f. A common simplification is to assume that electrons are adiabatic, which allows to neglect the kinetic term in the respective of. Hence, the respective f for electrons (fe) does not need to be modelled, effectively halving the computational cost.
B Quasilinear models
We used the QuaLiKiZ saturation rule (Bourdelle et al., 2007) applied to lienar GKW simulations, following Kumar et al. (2021). The QuaLiKiz saturation rule estimates turbulent transport based on linear gyrokinetic stability analysis, using quasilinear theory and empirical saturation rules. The quasilinear ion heat flux Q is modelled as:
Q Σ A&Wa
(20)
where A is the linear weight spectrum that quantifies the phase relationship between electrostatic potential fluctuations and that is retrieved by solving the linear gyrokinetic equation, and We is the turbulence intensity spectrum which in QuaLiKiz is parametrised by shape function S₁ and normalisation factor calibrated on nonlinear gyrokinetic simulations.
To fit the normalisation factor C, we assume access to a vector nonlinear fluxes E F and set the right-hand side of Equation (20) to z E Rd, then the optimal solution for C can be obtained via a least-squares
E, Ig C= Σ Σ
(21)
In our setup, we used flux estimates from the small training set (48 simulations) to compute the fit that resulted (2021). We attribute the difference to the different sampling strategy of simulations. The QuaLiKiz saturation rule retains key physics from gyrokinetics while allowing efficient prediction of turbulent transport in integrated modelling frameworks. Furthermore, we found it beneficial to center the nonlinear flux vector q which alters the normalisation constant C, but leads to reduced error on flux predictions on separate validation set.
7.93. This value differs slightly from the one reported Kumar et al.
18

---

C Data Generation and Visualisation
Data Generation.
In this work, we mainly consider turbulence driven by ion temperature gradients, also typically referred to the cyclone base case (Dimits et al., 2000a). We leverage the numerical code GKW (Peeters et al., 2009) for generating nonlinear and linear gyrokinetic simulations, by varying four parameters: R/ Lt. R/LL, 3, and q, which significantly affect emerging turbulence in the Plasma. R is the ion temperature gradient, which is the main driver of turbulence. R/LL, is the density gradient, which has more stabilizing effect but can sometimes also lead to enhanced turbulence. The parameter q denotes the so-called safety factor, which is the inverse of the rotational transform and describes how often particle takes poloidal turn before taking toroidal turn. & denotes magnetic shearing. It has a stabilizing effect as more magnetic shearing results in improved isolation of the Plasma.
In total we perform two data generation passes. For the first pass, we specify the ranges to sample the four parameters as RT [3, 12], R/ [1, 7], [1.9], and á E [0.5, 5]. In addition, we also always vary the noise amplitude of the initial condition within [le 5, le - 3] and the initial condition itself as one of [sin, cos, random]. We generate 100 simulations in total with this setup, but found that only approximately half of them actually exhibit turbulence. The turbulent configurations depicted as the sparse point cloud in Figure 5. Therefore, we adapt the parameter ranges for the second data generation pass to R/ LT E [6, 12). Ln E [0, 2]. q E [5, 9). and å E [0.5, 2]. Effectively, this change in parameters reduces stabilizing factors such as magnetic shearing and density gradients to force simulations to exhibit turbulence. We sample another 200 simulations from these parameter ranges and found that all of these develop strong turbulence (see distribution of Q in Figure 5). Eventually, we end up with around 250 simulations of which we randomly select three as validation set during training and six as in-distribution evaluation set, to end up with final training set of 241 simulations. Importantly, we still include the initially generated turbulent simulations as data generation is computationally expensive. Finally, we use the same parameter combination of each nonlinear simulation to generate their linear counterparts for the reduced-ord order model.
5D Data Visualization. We show a visual illustration of the 5D distribution function in Figure 6. For visualization purposes, we always show combinations of the different axes while averaging or slicing across the remaining ones. This way, we end up with 2D planes that are easy to visualize.
Furthermore, we also highlight the difference between nonlinear and linear simulations in Figure 7. We plot 2D planes of the 5D field of nonlinear and linear simulations side-by-side for three different timesteps. Since colorbars are shared across each pair of 2D planes, the difference between nonlinear and linear simulations can be clearly observed. Specifically, the structure in the linear simulation is symmetric, while the nonlinear simulation exhibits asymmetries. Furthermore, the difference is particularly obvious in wavenumber space, where (1) magnitudes of nonlinear simulations much larger, and (2), there is a lot more interaction between modes. Hence, there is structure in the nonlinear simulations that is not captured in linear ones, which a drawback of approximations based on linear simulations, i.c. for quasilinear models.
D Implementation details
Data preprocessing A first measure we take is standardizing all fields and fluxes zero mean and unit variance. We observed that this positively affects training. To enable autoregressive rollouts, we simply accumulate all statistics across all training simulations, and use them to normalize and denormalize during inference. Furthermore, another measure that had significant impact on training dynamics transferring from spectral kzz ky coordinates to real æ, y. This positively affects normal- ization as the difference in magnitudes is less pronounced in real space. While electrostatic potential fields coming from GKW W are already in real space, they are padded for the real FFT and result in dif- ferent spatial dimensions to the 5D distribution function. This is inconvenient for up/downsampling so we unpad the Fourier-space electrostatic potentials to match spatial dimensions of the 5D field, before transforming them back to real space.
19

---

150
2.5
5.0
7.5
2
5.0
7.5
10.0
100
200
Figure 5: Distribution input parameters 5, g, R/Ln, and R/L, along with average heat flux Q. The sampled parameter space is evenly distributed.
GyroSwin. Vision Transformers usually employ convolutions for patch operations (Dosovitskiy et al., 2021; Liu al., 2021, 2022), but they are not scalable to SD. Whereas convolutions the ker- nel is shared across every location, this is harder to parallelize in general nD setting. Therefore, we employ less parameter-efficiem local convolutions with independent kernels each location, imple- mented through fully connected layers. The n-dimensional input grids are first tiled with adaptive reshapes, then the flattened patch dimension embedded with shallow MIP. Patch embedding. merging, and expansions are implemented as linear layers or MI.s.. Furthermore, we add relative positional biases and condition all Swin layers on the 4 parameters as well as the current timestep via FiLM (Perez al., 2018). We experimented with DiT-style conditioning (Peebles & Xie, 2023), but found no improvements over FiLM.
We train GyroSwin for next-step prediction of the 5D distribution function of nonlinear gyrokinetic simulations. We use the Adam optimizer (Kingma & Ba, 2015) with weight decay of le-5 and cosine learning rate scheduler with linear warmup and peak at 3e-4, decayed to 0. During training we employ automatic mixed precision and cast to b with gradient clipping to a magnitude of 1. Due to the bulk of training data we perform lazy dataloading which results in substantial overhead, however is not possible to all data in RAM. In our special case of adiabatic electrons, each direction in the magnetic moment is independent of each other. Therefore, this dimension
20

---

0.0
2.5
5.0
7.5
0.0000
0.0001
-0.002
0.000
0.002
1e-5
1
110
t50
2200
H
kx
ky
Figure 6: Visualization of the 5D distribution function of nonlinear gyrokinetics (ground truth, Fourier space along k2 and kyy) We show different combinations of axes of the 5D field while averaging over the remaining ones for different timesteps. Colorbars are shared columnwise.
Nonlinear
Linear
Nonlinear
Linear
Nonlinear
Linear
Nonlinear
Linear
t10
V
H
S
Kx
t20
too
Kx
ky
Figure 7: Comparison of 5D distribution functions for linear and nonlinear simulations (ground truth, Fourier space along k2 and kyy) We show different combinations of axes of the 5D field while averaging over the remaining ones for different timesteps. Colorbars are shared across each 2D nonlinear and linear plane.
21

---

can be decoupled from the remaining ones and viewed as additional channels. This results in a substantial speedup without any loss of information.
For the smaller training set (48 simulations), we train our model for 200 epochs and evaluate every 20 epochs the three holdout trajectories based on which we perform model selection. During inference, we roll out the model autoregressively for the entire duration of each simulation (185 timesteps). GynoSwin is trained on four H100 GPUs with 80GB VRAM using PyTorch's Distributed Data Parallel (DDP) for approximately 120 hours.
For the larger training set (241 simulations), train GyroSwin for 500 epochs on 16 H100 GPUs with 80GB VRAM. We employ a scheduler for the different loss terms mentioned in Equation (8). Specifically, we start training only y for the first half of training, then switch on Ly and linearly increase We from 0 to from epoch 250 to 350 and keep it to afterwards. From epoch 350 onwards we also activate "q and again increase it linearly from 0 to between epoch 350 and 450 and continue with w = We = uq = for the last 50 epochs. Depending on the model size, we trained up until 336 hours (GyroSwin
D.1 Baselines
Gaussian Process Regression (GPR). We use four-parameter GPR trained on the operating pa- rameters (R/Lz, R/L,, Ln. q. and 8). It uses a Matern 3/2 kernel (Rasmussen & Williams, 2006) to predict the nonlinear flux. This is the model proposed by Hornsby al. (2024).
Multi-Layer Perceptron (MLP). We use 3 layer MLP with 128 hidden dimension and GELU activations. The inputs are the four operating parameters LT, R/Ln, 9. and 8), which are em- bedded in continuous sin-cos space, and the output is the auneraeergged scalar heat flux. This is similar to QLKNN proposed by van de Plassche et al. (2020), but trained on nonlinear fluxes instead of quasiliner ones.
FNO. The neuraloperator library* (Kovachki et al., 2023) contains tensorized nD implementa- tions of Fourier convolution layers. However, as shown in Table 3 these are too expensive to be trained on our data. To enable training FNO (Li et al., 2020), we utilize much faster 3D FNO which operates only on the spatial dimensions, while flattening the velocity space in the channels. This baseline is shown in Table 2 and has 256 latent dimension, 4 layers and considers of the total modes on each dimension (5,004 flat spatial modes).
PointNet. We use PointNet's (Qi et al., 2016) implementation from SIMSHIFT (Setinek et al., 2025), adapted to work on regular grid data. The total number of points on the grid is approxi- mately 22.3 million considering both channel dimensions. Therefore, to train this coordinate-bact baseline, randomly subsample the grid to 65,536 points for each sample. We also experimented with larger number of coordinates per-batch but observed diminishing returns beyond 65,536 points. Since the task is to predict the evolution of the distribution function, we provide its values at the cor- responding coordinates at time t to predict those for t + 1. For embedding grid coordinates we employ per-dimension embedding layers of the corresponding resolution and concatenate them to obtain the final coordinate embedding. Each of these embeddings is of dimension 64 leading to a hidden dimension of 320 for 5D input. We use GELU activations and condition the concatenated values and coordinates on the four operating parameters. The encoder projects the conditioned co- ordinates/values to dimensionality 640. Finally, we project to 12.800 global pooling features which are concatenated with the local coordinate features and serve as input to the decoder that predicts the field values at the local coordinates for the next timestep.
Transolver: We use Transover's (Wu et al., 2024) implementation from SIMSHIFT (Setinek al., 2025), adapted to work on regular grid data. We follow the same sampling and coordinate strategy as for PointNet. The Transolver embedding dimension set to 512 with 8 transformer blocks using 8 attention heads and 256 slices per head. Again, we provide the values of the distribution function time as input to predict the values at + 1.
22

---

6f pred
Gf
pred
OfGT
5f
GfGT
pf
ōfGT
ts
220
Vi lI
H
250
100
S
Figure 8: Side-by Side-by-side comparison of autoregressive rollout predictions with GyroSwin compared to ground truth (ær and in real space), in the saturated phase with shared colorbars. GyroSwin preserves the high-level structure within the first rollout steps. After larger amount of rollout steps emor accumulates, but predictions remain stable and do not diverge.
1916
(a) ID
(b) 00D
Figure 9: Comparison of W W(ky) In- In-Distribution and
E
Additional Results
Finally, we include qualitative results for autoregressive rollouts of GyroSwin. Specifically, we visualize model predictions for for the test case ID, and compare it to the ground truth in Figure for timesteps {10, 100, 200). We observe that the overall structure is more or less preserved within the first 10 rollout steps, however, afterwards the model suffers from error accumulation. This is especially pronounced for the velocity space, where the model seems to overpredict. Interestingly, though, the predictions in wavenumber space remain somewhat stable and do not diverge, evan after 100 autoregressive prediction steps. Generally we can say that GyroSwin yields stable predictions for long rollouts well beyond 100 timesteps.
F Diagnostics
We provide visualizations for the time-averaged turbulence-intense spectrum W for the six different ID test cases Figure and the five different OOD test cases in Figure 9b. Furthermore, we show flux spectra Qkk,) for the six different ID test cases in Figure 10a and the five different OOD test cases Figure 10b. Generally, we observe that GyroSwin sushtantially improves over the state-of-the-an numericcl quasilinear model on W (ky), which is reflected in the pearson
23

---

014/
1.10
(a) ID
(b) 00D
Figure 10: Comparison of flux spectra In-Distribution and Out-Of-Distribuiin
Qyazels
Qyto2eit:
Ampltudde
Amplluude
Amplitude
Amplltude 7
Ampltruse
Amplttude R 2
(a) ID
(b) OOD
Figure 11: Comparison of zonal flow profiles In-Distribution and Out-Of-Distriuttion.
correlation coefficients shown in Table 3. Interestingly, the fit for Q(ky) first glance appears be improved as well, however the correlation is approximately the same as for the quasilinear model. The reason for this is that high-frequencies are not well captured by GyroSwin. We expect this be a result of the inherent spectral bias of neural network architectures (Rahaman et al., 2019).
In addition to the time-averaged spectra, we provide visualizations for the time-averaged zonal flow profile for all ID simulations as well as OOD simulations for GyroSwin, We observe that GyroSwin Large usually tends to overestimate the amplitude therefore we normalize it for visualization purposes. The results for the ID and OOD test cases be observed in Figure 1 lla, and Figure ll,, respectively. Generally, we observe that the predicted zonal flow profile of GyroSwin, Large appears apperrs as a smoothed out version of the ground-truth. Again, our intuition is that this is due to the inherent low- frequency bias. Interestingly, some cases the zonal flow profile entirely off, which explains the rather high variance observed in the main table. Overall, we can conclude that GyroSwin, Large can capture the zonal flow of unseen simulations in an autoregressive manner even though there is plenty of room for improvement.
G Ablation Studies
To justify the design choices GyroSwin, we conduct set of ablation studies as follows. We pro- gressively add components of GyroSwin to a 5D Swin transformer and evaluate based on correlation

---

time to the 5D distribution function f and RMSE for the time-averaged flux trace Q. We report our results in Appendix G. Interestingly. the 5D-Swin already appears to be capable of performing sta- ble rollouts, however it completely fails in capturing the nonlinear flux trace. When moving to UNet like structure, we observe slight drop in correlation time, but improved flux predictions. Furthermore, adding channelwise mode separation gives a boost in corrllation time, particularly on OOD simulations. Moreover, there slight improvement in correlation time when adding our latent cross-attention and integrator modules. Finally, the most significant boost is observed when adding the flux prediction head, which results in our final design of GyroSwin, which achieves the best combination of stable autoregressive rollouts and nonlinear flux prediction.
Table Ablation study on different components GyroSwin, i.c. channel- mode separation of zonal flows and latent cross- puss-attentioonintegaaat modules. We report correlation time with r = 0.1 for the 5D distribution function and RMSE for Q on ID and OOD test cases. All methods are trained
on reduced dataset of 48 simulations.
f
Method
Params [M]
ID (t)
0OD (1)
ID (4)
OOD (L)
5D-Swin
46.2
24.83 + 3.19
25.00 2.39
121.34 = 13.36
125.74 + 23.23
GyroSwin
36.3
19.50 + 1.86
18.60 + 0.51
119.85 + 13.15
124.91 + 23.2
GyroSwin, + ZF channel
44.2
19.33 + 1.84
22.00
+
119.19 = 13.65
122.76 +: 23.63
GyroSwin f+
77.8
20.17 - 0.75
21.80
1:800::0 1.20
106.26 14.09
112.95 +: 25.22
GyroSwin LatentCA
87.3
21.00
2.62
2.89
116.49 12.11
121.41 23.45
GyroSwin 7464Q
91.1
22.83
+ 2.63
25.00
0.71
67.68 + 10.28
70.48 17.21
H Learning saturation rule
We perform an additional ablation study where we discard the nonlinear term in Equation (1), e.g. train GyroSwin linear simulations predict nonlinear fluxes. This effectively minrrr saturation in the Quasilinear setting, except that it does not rely on growth rates and reconstruction of the ky spectrum. Since linear simulations essentially converge fixed structure in the statistically steady state, we only use the last snapshot of a simulation and provide it as input to the model to predict the average nonlinear flux Q. We call this variant GyroSwin and compare it the Quasilinear saturation rule as well as the tabular regressors. Since this baseline is relatively cheap to train, we train it for three seeds and for both training sets, containing 48 and 241 simulations and report our results in Equation (5).
Interestingly, observe consistent trend that machine learning based methods outperform the reduced numerical model. An interesting aspect of this result is that the quasilinear saturation rule leverages additonal data such as linear growth rates of modes and flux contributions per mode which are hidden to tabular regressors and the learned saturation rule. Furthermore, with increasing train- ing data GyroSwin outperforms tabular regressors GPR and MLP that map from operating parameters to nonlinear heat flux. Since GyroSwin is only trained on the linear part of Equa- tion (1), therefore can conclude that it capable of generalizing over linear mode structures to infer nonlinear fluxes. Finally, all methods experience improvements for larger training sets. For the QL model the improvement can be traced back larger scaling factor that is fitted on more simulations, whereas all other methods are machin-learnin based, hence naturally improve with more available training data.
H.1 Vision Transformer Complexity
Computational complexity of 5D swin layers, with an input of resolution & x x channels and window size M (assumed square for simplicity) is
(MSSA) (zyzhw)² d² - 2xyzhww) d
OWW-MSA) zyzhw d2 2M ryzhw
removing the squared dependency on the sequence length zyzhw, and replacing it with much smaller window complexity of M/5 This consideration is the motivation behind the original swin paper by (Liu et al., 2021), as it makes attention images close to linear in resolution.
25

---

Table 5: Learning a saturation rule with GyroSwin. We report RMSE for the average flux Q on ID and OOD test cases. GyroSwin significantly outperforms the reduced numerical model (QL).
48 sims
241 sims
Method
(L)
OOD (L)
(L)
OOD (L)
Reduced numerical methods
QL
89.48 + 11.77
95.18 + 21.58 56.68 + 14.09
56.06 + 14.35
Tabular regressors
GPR
43.82 + 10.84
59.28
17.55
36.95 9.24 49.54 + 13.36
8.78 34.99 + 12.80
25.89
MLP
50.90 + 10.87
62.88 + 18.58
Learned saturation rule
GyroSwin Linear inner
76.74 + 11.79
82.95 + 19.58
42.85 + 12.65
23.40 + 5.89
NeurIPS Paper Checklist
Claims
Question: Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?
Answer: [Yes]
Justification: see Section and Section 5
Guidelines:
The answer NA means that the abstract and introduction do not include the claims made in the paper.
The abstract and/or introduction should clearly state the claims made, including the contributions made in the paper and important assumptions and limitations. A No or
NA answer to this question will not perceived well by the reviewers.
The claims made should match theoretical and experimental results, and reflect how much the results
be expected to generalize other settings.
fine to include aspirational goals motivation as long it is clear that these goals are not attained by the paper.
2. Limitations
Question: Does the paper discuss the limitations the work performed by the authors?
Answer: [Yes]
Justification: We discuss limitations in the conclusion section.
3. Theory assumptions and proof's
Question: For each theoretical result, does the paper provide the full set of assumptions and complete (and correct) proof?
Answer: [NA]
Justification: We do not provide theoretical results
4. Experimental result reproducibility
Question: Does the paper fully disclose all the information needed to reproduce the main experimental results of the paper to the extent that it affects the main claims and/or conclu- sions of the paper (regardless of whether the code and data are provided or not)?
Answer: [Yes]
Justification: Code is added as supplementary material, instructions for data generation are provided in Appendix C
5. Open access to data and code
26

---

Question: Does the paper provide open access to the data and code, with sufficient instruc- tions faithfully reproduce the main experimental results, as described in supplemental material?
Answer: [Yes]
Justification: Data generation pipelines is reproducible, numerical code is open source, trained models and snapshots for evaluation are publicly available on the huggingface hub.
6. Experimental setting/details
Question: Does the paper specify all the training and test details (e.g., data splits, hyper- parameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?
Answer: [Yes]
Justification: Section 4 provides details on experimental design, code added as supple- mentary material
7. Experiment statistical significance
Question: Does the paper report error bars suitably and correctly defined or other appropri- ate information about the statistical significance of the experiments?
Answer: [Ys]]
Justification: We provide error bars across evaluation runs for different simulations, how- ever we usually train only one seed due to computational constraints.
8. Experiments compute resources
For each experiment, does the paper provide sufficient information on the com- puter resources (type of compute workers, memory, time of execution) needed reproduce the experiments?
Question:
Answer: [Yes]
Justification: Section 4 and Appendix provide details on data generation and storage, Section 4 and Appendix D provide requirements for training GyroSwin
9. Code of ethics
Question: Does the research conducted in the paper conform, in every respect, with the NeurlPS Code of Ethics https: //tp:r://////pp////eeeee httppppllc//cccccciieleeeeeeeeee
Answer: [Yes]
Justification: Our work conforms with the code of ethics.
10. Broader impacts
Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?
Answer: [Yes]
Justification:
Discussed in Section 6
11. Safeguards
Does the paper describe safeguards that have been put in place for responsible release of data models that have high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?
Question:
Answer: [Yes]
Justification: There is no high risk for misuse.
12. Licenses for existing assets
Are the creators original owners of assets (e.g., code, data, models), used in the paper, properly credited and are the license and terms of use explicitly mentioned and properly respected?
Question:
Answer: [Yes]
Justification:
All properly credited.
27

---

13. New assets
Question: Are new assets introduced in the paper well documented and is the documenta- tion provided alongside the assets?
Answer: [Yes]
Justification: Data generation pipeline is documented in Section 4 and Appendix C, model is documented in Section 3 and Appendix D.
14. Crowdsourcing and research with human subjects
Question: For crowdsourcing experiments and research with human subjects, does the pa- per include the full text instructions given to participants and screenshots, applicable, well as details about compensation (if any)?
Answer: [NA]
Justification: No human subjects.
15. Institutional review board (IRB) approvals or equivalent for research with human subjects
Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed the subjects, and whether Institutional Review Board (IRB) approvals (or equivalent approval/review based on the requirements of your country or institution) were obtained?
Answer: [NA]
Justification: No human subjects.
16. Declaration of LLM usage
Question: Does the paper describe the usage of LLMs if it is an important, original, or non-standard component of the methods in this research? Note that if the LLM is used only for writing, editing, or formatting purposes and does not impact the core methodology, scientific rigorousness, originality of the research, declaration is not required.
Answer: [NA]
Justification: LLMs were majorly used for writing and editing.
28


Reviews:
[{"review_id": "IdKK9914WH", "reviewer": "Reviewer_n9nG", "summary": "This paper presents GyroSwin, a deep learning surrogate model for 5D gyrokinetic plasma turbulence simulations; a critical and computationally demanding problem in fusion research. Leveraging a hierarchical Swin Transformer architecture adapted for 5D data, GyroSwin captures key nonlinear effects, including zonal flows, that are often missed by traditional reduced models. GyroSwin predicts the 5D distribution function, 3D electrostatic potential, and scalar fluxes, achieving accuracy close to high-fidelity simulations while operating orders of magnitude faster. The approach enables new capabilities for data-driven plasma modeling.", "questions": "- In Table 2, GyroSwin does not consistently outperform classical baselines such as GPR, MLP, and LSTM, both for in-distribution and out-of-distribution data. Can the authors comment on the scenarios where GyroSwin underperforms, and explain why it does not achieve better accuracy in these cases?\n\n- The paper lacks theoretical analysis and does not provide ablation studies. Could the authors include more theoretical justification and perform ablation studies to better understand the contributions of each module?\n\n- The efficiency comparison is not entirely fair, as the numerical solver is evaluated on CPU while GyroSwin uses an A100 GPU. Could the authors provide a more balanced speed comparison, and also report the training cost, number of parameters, and GFLOPs for each model in Table 2?\n\n- The paper does not compare GyroSwin with recently developed numerical methods or state-of-the-art neural network models. Could the authors expand their experimental comparison to include such baselines for a fairer assessment?", "limitations": "Yes.", "rating": 3, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 3, "significance": 3, "originality": 3, "strengths_and_weaknesses": "Strength:\n- The paper is clearly written and well organized, making complex concepts accessible to the reader. Also, the work addresses a significant and challenging problem in fusion plasma physics, which is accurate and efficient modeling of 5D gyrokinetic turbulence. \n\n- GyroSwin introduces a multitask hierarchical Vision Transformer framework capable of processing high-dimensional scientific data. The approach leverages channel-wise mode separation, latent cross-attention across fields of varying dimensionality, and novel integration layers.\n\n- The method simultaneously predicts the 5D distribution function, 3D electrostatic potential, and scalar fluxes, and demonstrates improved accuracy in flux prediction and nonlinear dynamics compared to reduced numerical models.\n\nWeaknesses:\n- In the main results table (Table 2), GyroSwin does not consistently outperform other methods. In fact, it often fails to achieve better results than classical baselines such as GPR, MLP, and LSTM, both in in-distribution (ID) and out-of-distribution (OOD) settings.\n\n- The paper lacks theoretical analysis or evidence to explain why the proposed approach should work well. There is also no ablation study or detailed investigation into the factors driving the model’s performance.\n\n- The efficiency comparison is not entirely fair: the numerical solver runs on CPU while GyroSwin is evaluated on an A100 GPU. A direct speed comparison under these conditions is questionable. Additionally, training cost comparisons with other models are missing; it would be helpful to report the number of parameters and GFLOPs for each method. Given the scale of the architecture, the training cost for GyroSwin also appears to be quite high.\n\n- The paper lacks sufficient comparison with other methods, especially recently developed numerical approaches and neural network-based models. For a fair evaluation, it would be important to benchmark GyroSwin against both state-of-the-art numerical solvers and the latest deep learning models. In general, it is difficult for surrogate models to surpass the accuracy of modern numerical methods and advanced neural architectures.", "was_revised": true, "final_justification": "The authors have made an effort to provide a more detailed explanation of the experimental results in the rebuttal. However, the baseline models appear to be weak baselines, and despite the large difference in the number of parameters, the results do not show a substantial improvement. As these main concerns remain unresolved, I will maintain my original borderline reject score.", "rebuttal": "We appreciate the constructive feedback, recognizing both our technical contributions (accurate and efficient modeling of 5D gyrokinetic turbulence, multitask hierarchical ViT framework) and opportunities to strengthen the presentation of our work. \n\nWe noticed that the biggest critique is wrt Tab 2, in line with concerns raised by other reviewers (WQGA, DBSK, vuY7). From this we distilled that we have failed to bring across the right message in Table 2. Our intention was to showcase that GyroSwin can indeed model turbulence in 5D fields `(2x32x8x16x85x32)` as opposed to widely used reduced numerical models or surrogates that predict only quantities resulting from turbulence. The baselines are well established regression tools which predict the flux directly from operation parameters ([R/Lt, R/Ln, s, q] = 4-valued vector). Furthermore, the quasilinear model is state-of-the-art and the most widely used reduced numerical model in practice [1,2]. The task of GyroSwin is multiple orders of magnitude harder, and thus a direct comparison to baselines might have been misleading. Most importantly, 3D and 5D phenomena are not captured by most of the baselines. We have now made this much clearer by: \n1) Added conditioning to the flux prediction head of GyroSwin as all baselines are also conditioned on the operation parameters, which resulted in substantial gains in flux prediciton.\n2) Pearson correlation coefficient for diagnostics and zonal flow evaluation in Table 2, including error bars for all evaluations to demonstrate the significance of our results. \n3) Additional baslines, like a simple Transformer model and provide results for scaling GyroSwin to larger parameter counts which results in significant gains.\n4) Adding information on the input shapes each method achieves and n/a fields to make it clear that certain evaluations are not applicable to the most baselines.\n5) Additional tables for runtime/memory/GFLOPs for different methodologies that could be applied to 5D input, but simply do not scale well, hence result in prohibitive cost.\n6) A separate table for ablation studies for different prediction heads on a reduced data setup.\n\n### Weaknesses\n__(1) GyroSwin does not consistently outperform other methods (Table 2).__\n\nAs mentioned above, the main reason for the inconsistency in the results at the time of submission was that the flux prediction head was not conditioned on the operational parameters.\nBelow we report the restructured version of Table 2 with the improved flux results, diagnostics and zonal flow predictions. _n/a_ means that the model is not capable of producing the result. We hope that the updated Table + updated experiments help to bring across our main point, i.e., **only GyroSwin can model turbulence**. Other methods can only predict quantities which arise from turbulence.\n\nMethod|Params|Input Shape|$Q$ ID↓|$Q$ OOD↓|$W(k_y)$↑|$Q_i(k_y)$↑|Zonal Flow↑\n---|---|---|---|---|---|---|---\nGPR|6|(4,)|36.95±9.24|25.89±8.78|n/a|n/a|n/a\nMLP|37.3K|(4,)|59.52±15.01|29.59±11.19|n/a|n/a|n/a\nLSTM|5.3M|(10,)+(5,)|57.99±15.62|73.14±23.63|n/a|n/a|n/a\nTransformer|14.1|(10,)+(5,)|74.86±20.26|77.67±28.41|n/a|n/a|n/a|\nQL|1|(16,85,32)|40.74±7.47|53.3±14.77|0.04±0.15|0.54±0.46|n/a\nGyroSwin-Lin|51.3M|(32\\,8,16,85,32)|38.03±10.73|31.19±8.73|n/a|n/a|n/a\nGyroSwin-S|90.2M|(32,8,16,85,32)|36.13±7.55|59.76±17.62|0.98±0.03|0.4±0.37|0.11±0.19\nGyroSwin-M|250.9M|(32,8,16,85,32)|36.15±7.45|76.24±16.99|0.98±0.03|0.3±0.37|0.15±0.22\nGyroSwin-L|998.3M|(32,8,16,85,32)|24.38±6.79|18.49±6.15|0.98±0.03|0.35±0.4|0.32±0.18\n\n__(2) Lacks theoretical analysis why the proposed approach should work well.__\n\nWe appreciate the reviewer’s request for theoretical justification. However, a general theoretical framework for selecting optimal surrogate models in high-dimensional nonlinear gyrokinetics does not currently exist. Our results show that GyroSwin successfully generalizes over ID and OOD evaluation scenarios and accurately captures some of the underlying physics. This is the strongest form of validation that is currently feasible. Our work can be viewed as an initial step towards establishing practical surrogate modeling techniques in this domain, and we hope it will motivate future theoretical study.\n\n__(3) Efficiency comparison is not entirely fair.__\n\nWe acknowledge that a comparison purely based on clocktime between CPU and GPU is not fair. A better comparison might be with floating-point operation count betwen the two simulations: to produce a single 5D dump, the numerical code GKW takes between 20-80 seconds on 64 CPU cores, with the theoretical peak GFLOP is up to 4200 GFLOPs (double precision, Ice Lake 4.1GHz 76 cores CPU, 2 FMA x 8 DP = 16 FLOPs/cycle). For GyroSwin, the runtime (NVIDIA H100 80GB HBM3) is ~30ms 756 GFLOPs. More details on the machine learning side is in the table in point (5).\n\nIn a practical scenario of full integrated plasma simulations, it is typically necessary to compute turbulent heat fluxes multiple times, with up to thousands of simulation runs (at multiple radial locations and at regular intervals). Additionally, when adding more complex phenomena to the system, the numerical simulation will increase computation time substantially while surrogate models scale more favorably.\n\n__(4) Training cost comparisons with other models are missing.__\n\nWe report training cost in terms of GFLOPS, number of parameters and peak memory usage during training for each model in the table below. GFLOPs are computed as the number of updates times the GFLOPs for a single forward pass times two (considering forward and backward pass). Importantly, though, as outlined above, the increase in GFLOPs comes from the fact that GyroSwin processes the whole 5D field, whereas the baselines merely process vectors of scalars.\n\nModel|Memory[GB]|GFLOPs|Params[M]\n---|---|---|---\nMLP|0.002|6.1|<0.0\nLSTM|2.7|600|5.3\nTransformer|6.4|15K|9.0\nGyroSwin|26.6|263.8G|90.2\nGyroSwin(Large)|60.9|3.3B|998.3\n\n__(5) Lacks sufficient comparison with other methods.__\n\nUnfortunately, as far as we know, there are no other neural networks that can reliably be trained on the same task, scale and setup as GyroSwin. On one hand there is very limited research in general to extending neural networks beyond 4D- inputs, and on the other even when common layers such as convolution are implemented in 5D, they struggle to scale to higher resolutions due to the shear data size. To drive the second point, the following table compares very basic (factorized) CNN and FNO models against GyroSwin when applied to a single data sample of shape `(2, 32, 8, 16, 85, 32)`.\n\nModel|Inference[ms]|Memory[GB]|GFLOPs|Params[M]\n---|---|---|---|---\nF-CNN|569.6|11.1|19408.7|2.0\nF-FNO|963.3|36.9|N/A|1.3\nViT|87.7|1.2|1356.3|90.2\nViT-L|1141.9|5.5|14316.3|936.8\nGyroSwin-S|30.4|5.0|756.0|90.2\nGyroSwin-L|262.4|11.4|9513.5|998.3\n\nScaling other popular methods such as modern conv-UNet and FNO to our task is not trivial, and probably requires a great deal of engineering, something that would arguably deserve it's own research work. This should both motivate the choice of our architecture and the reason why comparison with alternative neural networks is absent. \n\nAs for other numerical approaches, the aim of GyroSwin is not outperform sota numerical methods on accuracy. Our goal is to speed up the computation of quantities of interest and enable practictioners to afford more runs for the same compute time. The numerical code we employ to generate data, GKW, is considered state of the art in computational plasma physics [3].\n\nFinally, GyroSwin is based on ViTs, in particular shifted window attention and special cross-attention layers, which are arguably among the latest deep learning techniques.\n\n### Questions\n__(1) In Table 2, GyroSwin does not consistently outperform classical baselines such as GPR, MLP, and LSTM, both for ID and OOD data. Can the authors comment on the scenarios where GyroSwin underperforms, and explain why it does not achieve better accuracy in these cases?__\n\nThis is discussed in the answer to weakness (1).\n\n__(1) Ablation studies to better understand the contributions of each module.__\n\nWe agree with the reviewer that ablation studies are vital to determine crucial components of our model design. Therefore we report results on ablating the different heads added to the model for the different fields (5D distribution function, 3D potential fields, scalar flux). Importantly, ablating window attention is not feasible within the rebuttal period, as replacing it with full attention results in a drastic increase of computational burden (see Table in weakness (5)).  We conduct the ablations on a limited evaluation setup as extensive ablation studies on the whole dataset are infeasible due to the massive bulk of data that each variant is trained on (~2TB). We observe that the best OOD generalization performance is attained when combining all prediction heads.\n\nModel|Params[M]|ID|OOD\n---|---|---|---\nGyroSwin-OnlyFlux|51.3|37.8±9.18|54.38±16.48\nGyroSwin-Onlydf|48.4|132.45±12.12|84.8±13.69\nGyroSwin-df+phi|94.7|41.16±7.16|53.48±15.91\nGyroSwin-All|90.2|58.52±14.68|38.52±9.77\n\n__(2) Training cost, number of parameters, and GFLOPs for each model.__\n\nSee response to weakness (4).\n\n__(3) The paper does not compare GyroSwin with recently developed numerical methods or state-of-the-art neural network models. Could the authors expand their experimental comparison to include such baselines for a fairer assessment?__\n\nSee response to weakness (5).\n\n### References\n[1] Kumar N. et al. Turbulent transport driven by kinetic ballooning modes in the inner core of jet hybrid h-modes. Nuclear Fusion.\n\n[2] Citrin J. et al. Tractable flux-driven temperature, density, and rotation profile evolution with the quasilinear gyrokinetic transport model QuaLiKiz. \n\n[3] Najlaoui A. et al. Verifying turbulence model reduction in high β tokamak plasmas (2025) *Plasma Phys. Control*"}, {"review_id": "BkUzROK7SH", "reviewer": "Reviewer_WQGA", "summary": "In magnetic confinement fusion, plasma turbulence is a critical phenomenon that can result in heat and particle transport towards the reactor walls. However, this phenomenon is based on a 5D PDE that is computationally intensive. While quasi-linear approximation approaches neglect some nonlinear physics that are crucial to modeling the turbulence. Therefore, the authors propose GyroSwin, a multitask hierarchical ViT to process the 5D data to capture nonlinear dynamics.", "questions": "(1) What is W-MSA and SW-MSA in Figure 2? \\\n(2) Could you provide more insights into how cross-attention is related to 5D-3D integral? \\\n(3) In Table 2, traditional GPR and MLP approaches seem to excel at in-distribution and out-of-distribution scenarios. Is it because the problem  (5D->0D) is not sophisticated enough? \\\n(4) In Figure 3, why QL, GPR, and MLP results are flat lines? Is it possible to use these models to predict a trend like what GyroSwin does? \\\n(5) Is it possible to use traditional approaches (e.g., MLP) to obtain 5D->1D evaluation results and show them together in Figure 4? For example, taking 5D data as input and getting 1D data as output?", "limitations": "yes", "rating": 4, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 4, "significance": 3, "originality": 2, "strengths_and_weaknesses": "**Strengths:**\\\n(1) The motivation of this work is clear. \\\n(2) The model design goals are reasonable and clearly listed in Section 3. \\\n(3) The inclusion of physics could be inspiring to researchers on AI for science, especially in terms of high-dimensional PDEs. \\\n\\\n**Weaknesses:**\\\n(1) Extending model capabilities to support higher dimensions is not entirely new. \\\n(2) The concept of multitask prediction is a well-established technique in AI. \\\n(3) As an application-specific approach, the results in Table 2 are outperformed by generic and relatively simple approaches in many cases. \\\n(4) Some acronyms are not properly defined.", "was_revised": false, "final_justification": "", "rebuttal": "We sincerely thank the reviewer for their comments. We address them in the following subsections.\n\nWe noticed that one of the biggest critiques is wrt Tab 2, in line with concerns raised by reviewers (n9nG, DBSK, vuY7). From this we distilled that we have failed to bring across the intended message, i.e., to show that GyroSwin can indeed model turbulence in 5D fields `(2x32x8x16x85x32)` as opposed to widely used reduced numerical models or surrogates that predict only quantities resulting from turbulence. The baselines are well established regression tools which predict flux directly from operation parameters ([R/Lt, R/Ln, s, q] = 4 scalars). Furthermore, the quasilinear model is SOTA and the most widely used reduced numerical model in practice [1,2]. The task of GyroSwin is considerably harder, and a direct comparison to baselines was misleading. __3D and 5D phenomena are not captured by most of the baselines__. We make this clearer by adding: \n1) Conditioning the flux prediction head of GyroSwin, which resulted in substantial gains in flux prediciton.\n2) Pearson correlation coefficient for diagnostics and zonal flow evaluation in Table 2, including error bars for all evaluations to demonstrate the significance of our results. \n3) Additional baselines, (flux sequence Transformer); scaling GyroSwin to larger sizes.\n4) Information on the input shapes for each method; diagnostic fields to clarify that certain evaluations are not applicable to the most baselines.\n5) Additional tables for runtime/memory/GFLOPs for different methodologies that could be applied to 5D input but do not scale well and result in prohibitive cost.\n6) A separate table ablating different heads on a reduced data setup.\n\n### Weaknesses\n__(1) Extending model capabilities to support higher dimensions is not entirely new.__\n\nIt is true that other architectures have been extended to higher dimensions. E.g., Convolutions and FNOs have factorized implementations ([3, 4]), making them in theory generalizable to n-dimensions. Also, vanilla ViTs can in principle be applied to nD data, similarly to how we apply Swin. However, these methods do not natively scale well to higher dim inputs as we show in the table below.\n\nModel|Inference [ms]|Memory [GB]|GFLOPs|Params [M]\n---|---|---|---|---\nF-CNN|569.6|11.1|19408.7|2.0\nF-FNO|963.3|36.9|N/A|1.3\nViT|87.7|1.2|1356.3|90.2\nGyroSwin|30.4|5.0|756.0|90.2\n\nThis highlights the need for more scalability.\nHierarchical ViTs such as Swin [5] have been extended to Videos (3D [6]) and medical images (4D [7]), but not beyond that. We generalize Swin by implementing n-dimenional _local convolutions_, used in patching/merging/expansion layers and in window attention. To the best of our knowledge, this is novel.\n\nFurthermore, we incorporate inductive-biases, namely channelwise mode separation and latent integration modules.\nThese are tailored to our application domain, but are simple ways of including inductive biases into model design, which may be relevant to the community.\n\n__(2) The concept of multitask prediction is a well-established technique in AI.__\n\nThis is absolutely true, we did not claim that the multitask setup is one of our contributions. It is simply a tool to enhance our performance.\n\n__(3) As an application-specific approach, the results in Table 2 are outperformed by generic and relatively simple approaches in many cases.__\n\nAs mentioned above, the main cause of the result inconsistency at submission was that the flux head was not conditioned on operational parameters.\n\nBelow we report the restructured Table 2 with improved flux results, diagnostics and zonal flow predictions. _n/a_ means that the model is incapable of producing the result. We hope that the updated Table + updated experiments help to bring across our main point, i.e., **only GyroSwin can model turbulence**. Others can only predict quantities which arise from turbulence.\n\nMethod|Params|Input Shape|$Q$ ID↓|$Q$ OOD↓|$W(k_y)$↑|$Q_i(k_y)$↑|Zonal Flow↑\n---|---|---|---|---|---|---|---\nGPR|6|(4,)|36.95±9.24|25.89±8.78|n/a|n/a|n/a\nMLP|37.3K|(4,)|59.52±15.01|29.59±11.19|n/a|n/a|n/a\nLSTM|5.3M|(10,)+(5,)|57.99±15.62|73.14±23.63|n/a|n/a|n/a\nTransformer|14.1|(10,)+(5,)|74.86±20.26|77.67±28.41|n/a|n/a|n/a|\nQL|1|(16,85,32)|40.74±7.47|53.3±14.77|0.04±0.15|0.54±0.46|n/a\nGyroSwin-Lin|51.3M|(32\\,8,16,85,32)|38.03±10.73|31.19±8.73|n/a|n/a|n/a\nGyroSwin-S|90.2M|(32,8,16,85,32)|36.13±7.55|59.76±17.62|0.98±0.03|0.4±0.37|0.11±0.19\nGyroSwin-M|250.9M|(32,8,16,85,32)|36.15±7.45|76.24±16.99|0.98±0.03|0.3±0.37|0.15±0.22\nGyroSwin-L|998.3M|(32,8,16,85,32)|24.38±6.79|18.49±6.15|0.98±0.03|0.35±0.4|0.32±0.18\n\n__(4) Some acronyms are not properly defined.__\nWe address this point in question (1).\n\n### Questions\n\n__(1) What is W-MSA and SW-MSA in Figure 2?__\nThanks for noticing this. These are indeed not defined in the paper, and they were borrowed from the original swin paper. We now formalize those as they are used in our setup.\n\n- W-MSA is Window Multihead Self Attention (MSA). It uses the even-layer window partition scheme, and performs self-attention only within the windows (in parallel). As a slightly lax formulation, $\\text{W-MSA}(x)=\\left[\\sigma\\left( \\frac{(x_w W_h^q)(x_w W_h^k)^\\top}{\\sqrt{d}} + B \\right) (x_w W_h^v)\\right]_{h \\in \\text{heads}}W^o$, where $x_w$ is the window partitioned input, with $x_w.\\texttt{shape} = \\texttt{(batch*nwindows, w0, w1, ..., channels)}$, and $W_h^q, W_h^k, W_h^v, W_h^o$ are the q, k, v and out projection per-head $h$ projections of self-attention, and $B$ is a bias which encodes relative positional information.\n\n- SW-MSA is Shifted Window MSA, it uses the shifted, odd-layer partition scheme. This introduces some information exchange with the neighbors. $\\text{SW-MSA}(x)=\\texttt{roll}(\\text{W-MSA}(\\texttt{roll}(x, −W/2)), +W/2)$, where $W$ is the window size and $\\texttt{roll}$ is a matrix roll operation, which circularly shifts the input along the spatial dimensions by a given displacement (equivalent to `np.roll`).\n\nPlease refer to the 2021 Shifted Window Transformer paper [5] for simpler 2D visualizations. \n\n__(2) More insights into how cross-attention is related to 5D-3D integral?__\n\nThe inspiration for the 5D-3D mix and the $\\int_{V_{space}}$ modules comes from the integral expressed in Eq. 2. It consists of a velocity space integral (to obtain potential fields) and a spatial integral (to obtain fluxes).\n\nFirst, 3D latents of the electrostatic potentials are computed (starting from 5D latents) with the $\\int_{V_{space}}$ block. This uses cross-attention, with a learned scalar query which reduces the velocity space of the 5D key-value, (akin to perceiver-pooling [8]). The output of this operation is a latent 3D field.\n\nThen, the 5D-3D mix module is used to allow exchange of information at different scales. Here, cross-attention is performed between the latents of the 3D potential and 5D distribution function (connection expressed by the $v_\\chi$ term in Eq. 2). The output is another 3D latent field.\n\nFinally, this 3D latent is used both for aiding potential reconstruction (3D potential fields) and for flux estimation (scalar). Fluxes are obtained by pooling the latents obtained with 5D-3D mixing spatially, and finally applying an MLP that predicts a scalar value.\n\n__(3) In Tab2, traditional GPR and MLP approaches seem to excel at ID and OOD scenarios. Is the problem (5D->0D) not sophisticated enough?__\n\nThis is a very legitimate question, because the terminology of “5D->0D” evaluation for Table 2 is partially inaccurate. In fact, only GyroSwin treats it as an actual “5D->0D” problem, while all other baselines do not and cannot process 5D data at all. __Refer to weakness (3) for full answer.__\n\n__(4) In Figure 3, why QL, GPR, and MLP results are flat lines?__\n\nThank you for pointing out that this isn't clear.\nGenerally, it is __not__ possible to extend GPR and QL to predict trend over time. The target for __QL, GPR and MLP is the time-averaged heat flux__ $\\bar{Q}$ in the statistically stable (saturated) phase. These models do not have the concept of time, and are trained to regress directly from physical parameters (GPR, MLP) or simplified linear equations (QL) to the saturated average flux $\\bar{Q}$. While theoretically it would be possible to extend MLP to predict flux over time by adding timestep information, we expect this to not perform better than LSTM/Transformer due to its lack of expressivenes. Due to its autoregressive nature, GyroSwin is able to reproduce the evolution of turbulent flux into the saturated phase after a certain number of steps. \n\n__(5) Is it possible to use traditional approaches (e.g., MLP) to obtain 5D->1D evaluation results and show them together in Figure 4?__\n\nYes, theoretically it is possible, but it would require flattening out the 5D field into a vector. Assuming a hidden size of the MLP of 512, a single weight matrix would carry over 10B parameters. Since 5D->1D evaluation requires a predicted 5D field, two of those are needed. The flattening operation destroys all locality in the field, which is of utmost importance.\n\n### References\n[1] Kumar N. et al. Turbulent transport driven by kinetic ballooning modes in the inner core of jet hybrid h-modes. Nuclear Fusion.\n\n[2] Citrin J. et al. Tractable flux-driven temperature, density, and rotation profile evolution with the quasilinear gyrokinetic transport model QuaLiKiz. \n\n[3] Wang M. et al. (2017). Factorized Convolutional Neural Networks. *ICCV*.\n\n[4] Tran A. et al. (2023). Factorized Fourier Neural Operators. *ICLR*.\n\n[5] Liu Z. et al. (2021). Swin Transformer: Hierarchical Vision Transformer using Shifted Windows. *ICCV*.\n\n[6] Liu Z. et al. (2021). Video Swin Transformer.\n\n[7] Kim P. et al., Moon T. (2023). SwiFT: Swin 4D fMRI Transformer. *NeurIPS*.\n\n[8] Alkin B. et al. (2024). Universal Physics Transformers: A Framework For Efficiently Scaling Neural Operators. *NeurIPS*."}, {"review_id": "oHINvRkoYh", "reviewer": "Reviewer_DBSK", "summary": "The paper proposes an application-specific model to learn a type of PDE used for plasma simulation. GyroSwin builds a hierarchical Vision Transformer in 5D space and is trained on the 5D distribution function for the adiabatic electron approximation", "questions": "what is the model computational complexity and how does that compare to the baselines? It seems less clear in the manuscript.", "limitations": "yes", "rating": 3, "confidence": 2, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 2, "significance": 2, "originality": 3, "strengths_and_weaknesses": "Strengths:\nThe approach appears sound. It makes use of vision transformer in 5D distribution function over time.\nThe paper explains the context well, and has demonstrated the significance of the problem in application contexts.\n\nWeaknesses:\nThe paper can benefit from discussion of related works in neural approaches to solving Gyrokinetics problems, and their weaknesses compared to the proposed model. \n\nIn table 2, standard deviation needs to be reported, to assess significance of the results. Furthermore, a transformer-based methods could potentially help readers judge the advantages of proposed methods. LSTM methods might be less effective in large-scale training. Other advanced methods such as state-space methods could be considered as well.\n\nThe paper lacks ablation studies.\n\nThe paper includes a lot of domain-specific details for the application, not commonly encountered in the deep learning community. Unfortunately I lack the physics background to assess the relevance and importance of the proposed simulation model.", "was_revised": false, "final_justification": "", "rebuttal": "We'd first like to thank DBSK for their review and their constructive feedback.\n\nWe noticed that one of the biggest critiques is w.r.t. Table 2. This is in line with concerns raised by other reviewers (n9nG, WQGA, vuY7). From this feedback we distilled that we have failed to bring across the right message in Table 2. Our intention was to showcase that GyroSwin can indeed model turbulence in 5D fields `(2x32x8x16x85x32)` as opposed to widely used reduced numerical models or surrogates that predict only quantities resulting from turbulence. The baselines are well established regression tools which predict the flux directly from operation parameters ([R/Lt, R/Ln, s, q] = 4-valued vector). Furthermore, the quasilinear model is state-of-the-art and the most widely used reduced numerical model in practice [1,2]. The task of GyroSwin is multiple orders of magnitude harder, and thus a direct comparison to baselines might have been misleading. Most importantly, 3D and 5D phenomena are not captured by most of the baselines. We have now made this much clearer by: \n1) We added conditioning to the flux prediction head of GyroSwin as all baselines are also conditioned on the operation parameters, which resulted in substantial gains in flux prediciton.\n2) We added Pearson correlation coefficient for diagnostics and zonal flow evaluation in Table 2, including error bars for all evaluations to demonstrate the significance of our results. \n3) We added additional baslines, like a simple Transformer model and provide results for scaling GyroSwin to larger parameter counts which results in significant gains.\n4) We added information on the input shapes each method achieves and n/a fields to make it clear that certain evaluations are not applicable to the most baselines.\n5) We added additional tables for runtime/memory/GFLOPs for different methodologies that could be applied to 5D input, but simply do not scale well, hence result in prohibitive cost.\n6) We added a separate table for ablation studies for different prediction heads on a reduced data setup.\n\n\n### Weaknesses\n__(1) The paper can benefit from discussion of related works in neural approaches to solving Gyrokinetics problems, and their weaknesses compared to the proposed model.__\n\nFor additional information and proper referencing, we would like to point the reviewer to lines 123 to 137 in the Related Works section.\nLiterature on machine learning approaches for gyrokinetics is rather sparse. Most works focus on predicting average fluxes computed via reduced quasilinear models from physical parameters, employing traditional machine learning methods like gaussian processes or basic neural techniques such as MLPs. More recently there have been a few attempts to learn nonlinear gyrokinetics through reduced and integrated quantities, such as wavenumbers and slices of the electrostatic potential field. Note that these methods are limited to scalars, 1D and 2D information, so they cannot properly model the gyrokinetics nonlinear behaviour. This distinguishes GyroSwin from any previous work on this topic.\n\n__(2) In table 2, standard deviation needs to be reported, to assess significance of the results. Furthermore, a transformer-based methods could potentially help readers judge the advantages of proposed methods. LSTM methods might be less effective in large-scale training. Other advanced methods such as state-space methods could be considered as well.__\n\nBelow we report the restructured version of Table 2 with the improved flux results, diagnostics and zonal flow predictions. _n/a_ means that the model is not capable of producing the result. We hope that the updated Table + updated experiments help to bring across our main point, i.e., **only GyroSwin can model turbulence**. Other methods can only predict quantities which arise from turbulence.\nWe also added a vanilla Transformer baseline to model the sequential evolution of the flux trace.\nWe agree that state-space models could be added as well, but our main point is that we require specialized architectures like GyroSwin that scale well to process the 5D data which allows to actively model turbulence instead of just predicting quantities that result from it.\n\nMethod|Params|Input Shape|$Q$ ID↓|$Q$ OOD↓|$W(k_y)$↑|$Q_i(k_y)$↑|Zonal Flow↑\n---|---|---|---|---|---|---|---\nGPR|6|(4,)|36.95±9.24|25.89±8.78|n/a|n/a|n/a\nMLP|37.3K|(4,)|59.52±15.01|29.59±11.19|n/a|n/a|n/a\nLSTM|5.3M|(10,)+(5,)|57.99±15.62|73.14±23.63|n/a|n/a|n/a\nTransformer|14.1|(10,)+(5,)|74.86±20.26|77.67±28.41|n/a|n/a|n/a|\nQL|1|(16,85,32)|40.74±7.47|53.3±14.77|0.04±0.15|0.54±0.46|n/a\nGyroSwin-Lin|51.3M|(32\\,8,16,85,32)|38.03±10.73|31.19±8.73|n/a|n/a|n/a\nGyroSwin-S|90.2M|(32,8,16,85,32)|36.13±7.55|59.76±17.62|0.98±0.03|0.4±0.37|0.11±0.19\nGyroSwin-M|250.9M|(32,8,16,85,32)|36.15±7.45|76.24±16.99|0.98±0.03|0.3±0.37|0.15±0.22\nGyroSwin-L|998.3M|(32,8,16,85,32)|24.38±6.79|18.49±6.15|0.98±0.03|0.35±0.4|0.32±0.18\n\n__(3) The paper lacks ablation studies.__\n\nWe agree with the reviewer that ablation studies are vital to determine crucial components of our model design. Therefore we report results on ablating the different heads added to the model for the different fields (5D distribution function, 3D potential fields, scalar flux). Importantly, ablating window attention is not feasible within the rebuttal period, as replacing it with full attention results in a drastic increase of computational burden (see Table in Question (1)).  We conduct the ablations on a limited evaluation setup as extensive ablation studies on the whole dataset are infeasible due to the massive bulk of data that each variant is trained on (~2TB). We observe that the best OOD generalization performance is attained when combining all prediction heads.\n\n| Model | Params[M] | ID | OOD |\n| --- | --- | --- | --- |\nGyroSwin-OnlyFlux | 51.3 | 37.8 ± 9.18 | 54.38 ± 16.48 | \nGyroSwin-Onlydf | 48.4 | 132.45 ± 12.12 | 84.8 ± 13.69 |\nGyroSwin-df+phi | 94.7 | 41.16 ± 7.16 | 53.48 ± 15.91 |\nGyroSwin-All | 90.2 | 58.52 ± 14.68 | 38.52 ± 9.77 | \n\n### Questions\n__(1) what is the model computational complexity and how does that compare to the baselines?__\n\nTheoretical computational complexity of 5D swin layers, with an input of resolution $v \\times m \\times s \\times x \\times y$ with $C$ channels and window size $W$ (assumed square for simplicity) is $O(\\text{W-MSA}) = 4 \\: vmsxy \\: C^2 + 2W^5 \\: vmsxy \\: C$\n\nmaking it linear in the input resolution.\n\nIn practice, we measure inference time, memory and GFLOPs for our models and compare with commonly used architectures (factorized CNN and factorized FNO). Numbers are reported on an NVIDIA H100 80GB HBM3, run on a single sample with resolution `(2, 32, 8, 16, 85, 32)`. FLOP counts are computed with the _fvcore_ library.\n\nModel | Inference [ms] | Memory [GB] | GFLOPs | Params [M]\n--- | --- | --- | --- | ---\nF-CNN | 569.6 | 11.1 | 19408.7 | 2.0\nF-FNO | 963.3 | 36.9 | N/A | 1.3\nGyroViT  | 87.7 | 1.2 | 1356.3 | 90.2\nGyroViT (Large) | 1141.9 | 5.5 | 14316.3 | 936.8\nGyroSwin | 30.4 | 5.0 | 756.0 | 90.2\nGyroSwin (Large) | 262.4 | 11.4 | 9513.5 | 998.3\n\nWhen processing the entire 5D fields, GyroSwin is considerably faster and less memory demanding than even small and simple convolution-based models. However, if we compare to the baselines that we originally reported in Table 2, GyroSwin is obviously much more expensive since __the input of GyroSwin is ~5 orders of magnitudes larger__: `(2, 32, 8, 16, 85, 32)` vs `(4,)` (or `(10,) + (4,)` for the sequence models, LSTM and Transformer).\nNonetheless, we report numbers for the simple baselines below.\n\n| Model | Inference [ms] | Memory [GB] | GFLOPs | Params [M] |\n| --- | --- | --- | --- | --- |\n| MLP | 0.1 | <0.0 | <0.0 | <0.0 |\n| LSTM | 0.4 | 0.1 | 0.1 | 5.3 |\n| Transformer | 0.5 | 0.1 | 2.5 | 9.0 |\n\n\n### References\n[1] Kumar N. et al. Turbulent transport driven by kinetic ballooning modes in the inner core of jet hybrid h-modes. Nuclear Fusion.\n\n[2] Citrin J. et al. Tractable flux-driven temperature, density, and rotation profile evolution with the quasilinear gyrokinetic transport model QuaLiKiz."}, {"review_id": "GZdQf8uAPs", "reviewer": "Reviewer_386x", "summary": "The paper presents a novel deep learning approach, GyroSwin, aimed at efficiently simulating nonlinear gyrokinetic plasma turbulence. \n\nGyroSwin incorporates advanced techniques like latent cross-attention for interactions between 3D and 5D fields and channel-wise mode separation to isolate important nonlinear phenomena like zonal flows. It uses shifted window attention (Swin Transformer) to preserve locality and improve computational efficiency, allowing for accurate predictions in high-dimensional spaces. The model is trained using a multitask approach, predicting the 5D distribution function, 3D electrostatic potential, and heat flux.", "questions": "The proposed model structure seems to be much more complex than the baseline, such as MLP and LSTM. Can the author show them, as well as the parameter count and prediction speed of GyroSwin? This would be meaningful.", "limitations": "The authors discuss the relevant limitations.", "rating": 5, "confidence": 2, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 3, "significance": 3, "originality": 3, "strengths_and_weaknesses": "**Strengths**\n\n 1. This article is well written, with rich background descriptions and an overview of related work. Although I am not an expert in the field of controlled nuclear fusion, this introduction can help me quickly familiarize myself with the context.\n\n2. Although the modules of the proposed surrogate model are existing, the design of each module is principled and strongly related to the task.\n\n**Weaknesses**\n\nIt lacks, the ablation of each module, such as window attention, integral module, Multi-task training, etc., ablation experiments can better show the contribution of each module and inspire future work.", "was_revised": true, "final_justification": "The author's ablation resolved my concerns. I think this manuscript does make a contribution because each component design seems to have a task related motivation, even though I am not an expert in plasma turbulence simulation. Thus, I keep my score.", "rebuttal": "Thank you for the positive evaluation of our submission, and the constructive feedback on how to improve it. We address the raised concerns as follows.\n\n### Weaknesses\n__(1) It lacks, the ablation of each module, such as window attention, integral module, Multi-task training, etc., ablation experiments can better show the contribution of each module and inspire future work.__\n\nWe agree with the reviewer that ablation studies are vital to determine crucial components of our model design. Therefore we report results on ablating the different heads added to the model for the different fields (5D distribution function, 3D potential fields, scalar flux). Importantly, ablating window attention is not feasible within the rebuttal period, as replacing it with full attention results in a drastic increase of computational burden (see Table below).\n\nModel|Inference [ms]|Memory [GB]|GFLOPs|Params [M]\n---|---|---|---|---\nF-CNN|569.6|11.1|19408.7|2.0\nF-FNO|963.3|36.9|N/A|1.3\nViT|87.7|1.2|1356.3|90.2\nViT (Large)|1141.9|5.5|14316.3|936.8\nGyroSwin|30.4|5.0|756.0|90.2\nGyroSwin (Large)|262.4|11.4|9513.5|998.3\n\nWe conduct the ablations on a limited evaluation setup as extensive ablation studies on the whole dataset are infeasible due to the massive bulk of data that each variant is trained on (~2TB). We observe that the best OOD generalization performance is attained when combining all prediction heads.\n\n| Model | Params[M] | ID | OOD |\n| --- | --- | --- | --- |\nGyroSwin-OnlyFlux | 51.3 | 37.8 ± 9.18 | 54.38 ± 16.48 | \nGyroSwin-Onlydf | 48.4 | 132.45 ± 12.12 | 84.8 ± 13.69 |\nGyroSwin-df+phi | 94.7 | 41.16 ± 7.16 | 53.48 ± 15.91 |\nGyroSwin-All | 90.2 | 58.52 ± 14.68 | 38.52 ± 9.77 | \n\n### Questions\n__(1) The proposed model structure seems to be much more complex than the baseline, such as MLP and LSTM. Can the author show them, as well as the parameter count and prediction speed of GyroSwin? This would be meaningful.__\n\nIt is true that GyroSwin is more complex than baselines such as MLP or LSTM, it is important to note though, that the task for MLP, GPR, and LSTM/Transformer is much simpler as they do not process 5D fields: __the input of GyroSwin is ~5 orders of magnitudes larger__: `(2, 32, 8, 16, 85, 32)` vs `(4,)` (or `(10,) + (4,)` for the sequence models, LSTM and Transformer). We realize that this has not been conveyed as desired in our submission.\n\nIf we compare it with other architectures processing the same amount of data, the advantage of GyroSwin becomes more pronounced. Numbers are reported on an NVIDIA H100 80GB HBM3, run on a single sample with resolution `(2, 32, 8, 16, 85, 32)`. FLOP counts are computed with the _fvcore_ library.\n\n| Model | Input shape | Inference [ms] | Memory [GB] | GFLOPs | Params [M] |\n| --- | --- | --- | --- | --- | --- |\n| MLP | (4,) | 0.1 | <0.0 | <0.0 | <0.0 |\n| LSTM | (10,) + (5,) | 0.4 | 0.1 | 0.1 | 5.3 |\n| Transformer | (10,) + (5,) | 0.5 | 0.1 | 2.5 | 9.0 |\n| F-CNN | (2, 32, 8, 16, 85, 32) | 569.6 | 11.1 | 19408.7 | 2.0 |\n| F-FNO | (2, 32, 8, 16, 85, 32) | 963.3 | 36.9 | N/A | 1.3 |\n| ViT | (2, 32, 8, 16, 85, 32) | 87.7 | 1.2 | 1356.3 | 90.2 |\n| ViT (Large) | (2, 32, 8, 16, 85, 32) | 1141.9 | 5.5 | 14316.3 | 936.8 |\n| GyroSwin | (2, 32, 8, 16, 85, 32) | 30.4 | 5.0 | 756.0 | 90.2 |\n| GyroSwin (Large) | (2, 32, 8, 16, 85, 32) | 262.4 | 11.4 | 9513.5 | 998.3 |"}, {"review_id": "5Y0Glsqn4d", "reviewer": "Reviewer_vuY7", "summary": "This paper presents GyroSwin, a deep learning model for simulation of 5D gyrokinetic plasma turbulence, a challenging problem crucial for fusion reactor design but is computationally prohibitive for traditional direct numerical simulation methods. Current reduced models, like Quasilinear (QL) approaches, are faster but neglect important nonlinear phenomena. GyroSwin tries to instead directly learn the evolution of the 5D distribution function to improve prediction accuracy. Modeling 5D fields is challenging using existing deep learning approaches such as CNNs or Vision Transformers (ViTs), so the modeling is based on Swin Transformers extended to 5D. It also uses physics-informed features such as integration of latent variables (reminiscent of Latent ODEs and traditional reduced order models). The model is trained in a multitask setup to predict the future 5D distribution, 3D potential fields and scalar heat flux. Experiments compare various in-distribution and out-of-distribution datasets and show that GyroSwin is comparable or the most accurate compared to standard DL, and non DL approaches.", "questions": "1. OOD Qi Performance: Why could explain why GyroSwin struggles in some OOD cases compared to GPR/MLP (Table 2), even though the modeling of GyroSwin has significantly more consideration for the underlying physics? I would have expected GyroSwin would overfit less.\n\n1. While Fig 3 shows a successful rollout, would it be possible to visualize the distribution of errors over multiple rollouts as a function of time (different dataset, different random seed etc.)? How would such an error distributions in the rollouts differ for baseline methods?", "limitations": "yes", "rating": 4, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 2, "clarity": 3, "significance": 2, "originality": 2, "strengths_and_weaknesses": "*Strengths*\n- This work addresses a significant and challenging problem in computational physics (nuclear fusion) where breakthroughs would have a large impact.\n- The deep learning model demonstrates significant speedups compared to numerical solvers.\n\n*Weaknesses*\n\n- Ablations of the various model components are missing. What happens if the cross-attention modules, or latent integration modules, or the physics-informed loss terms are removed (and use simple squared error loss instead)?\n- While UNets/CNN, ViTs and FNOs may not be scalable, they would still serve as baselines for small scale datasets. Particularly, it would be informative to see the quality ceiling that may be achieved using a non-physics informed deep learning approach. A simple MLP architecture is used, but that may be a weak baseline (also what are the parameter counts?)\n- Even a simple baseline such as Gaussian process and MLPs achieve comparable quality on some datasets, even though they are missing several features such as physics informed learning dynamics. To me, the Table 2 appears too noisy to decide that GyroSwin is indeed the best performing model.", "was_revised": true, "final_justification": "Authors have addressed my concerns (lack of error bars, additional concerns), however a stronger evaluation would be even better.", "rebuttal": "We thank vuY7 for the time invested into the review and the constructive insights. We proceed to clarify the misconceptions and address the raised concerns.\n\nWe noticed that the biggest critique is w.r.t. Table 2. This is in line with concerns raised by other reviewers (n9nG, WQGA, DBSK). From this feedback we distilled that we have failed to bring across the right message in Table 2. Our intention was to showcase that GyroSwin can indeed model turbulence in 5D fields `(2x32x8x16x85x32)` as opposed to widely used reduced numerical models or surrogates that predict only quantities resulting from turbulence. The baselines are well established regression tools which predict the flux directly from operation parameters ([R/Lt, R/Ln, s, q] = 4-valued vector). Furthermore, the quasilinear model is state-of-the-art and the most widely used reduced numerical model in practice [1,2]. The task of GyroSwin is multiple orders of magnitude harder, and thus a direct comparison to baselines might have been misleading. Most importantly, 3D and 5D phenomena are not captured by most of the baselines. We have now made this much clearer by: \n1) We added conditioning to the flux prediction head of GyroSwin as all baselines are also conditioned on the operation parameters, which resulted in substantial gains in flux prediciton.\n2) We added Pearson correlation coefficient for diagnostics and zonal flow evaluation in Table 2, including error bars for all evaluations to demonstrate the significance of our results. \n3) We added additional baslines, like a simple Transformer model and provide results for scaling GyroSwin to larger parameter counts which results in significant gains.\n4) We added information on the input shapes each method achieves and n/a fields to make it clear that certain evaluations are not applicable to the most baselines.\n5) We added additional tables for runtime/memory/GFLOPs for different methodologies that could be applied to 5D input, but simply do not scale well, hence result in prohibitive cost.\n6) We added a separate table for ablation studies for different prediction heads on a reduced data setup.\n\n### Weaknesses\n__(1) Ablations of the various model components are missing. What happens if the cross-attention modules, or latent integration modules, or the physics-informed loss terms are removed (and use simple squared error loss instead)?__\n\nWe agree with the reviewer that ablation studies are vital to determine crucial components of our model design. Therefore we report results on ablating the different heads added to the model for the different fields (5D distribution function, 3D potential fields, scalar flux). Importantly, ablating window attention is not feasible within the rebuttal period, as replacing it with full attention results in a drastic increase of computational burden (see Table in weakness (2)).  We conduct the ablations on a limited evaluation setup as extensive ablation studies on the whole dataset are infeasible due to the massive bulk of data that each variant is trained on (~2TB). We observe that the best OOD generalization performance is attained when combining all prediction heads.\n\n| Model | Params[M] | ID | OOD |\n| --- | --- | --- | --- |\nGyroSwin-OnlyFlux | 51.3 | 37.8 ± 9.18 | 54.38 ± 16.48 | \nGyroSwin-Onlydf | 48.4 | 132.45 ± 12.12 | 84.8 ± 13.69 |\nGyroSwin-df+phi | 94.7 | 41.16 ± 7.16 | 53.48 ± 15.91 |\nGyroSwin-All | 90.2 | 58.52 ± 14.68 | 38.52 ± 9.77 | \n\n__(2) While UNets/CNN, ViTs and FNOs may not be scalable, they would still serve as baselines for small scale datasets. Particularly, it would be informative to see the quality ceiling that may be achieved using a non-physics informed deep learning approach. What are the parameter counts?__\n\nIt is true that other architectures such as CNNs/FNOs/ViTs can be extended to higher dimensions as well. However, these methods do not natively scale well to such high dimensional input. To demonstrate this, we report inference speed, peak memory consumption, GFLOPs, and parameter counts for those baselines in the table below. Numbers are reported on an NVIDIA H100 80GB HBM3, run on a single sample with resolution `(2, 32, 8, 16, 85, 32)`. FLOP counts are computed with the _fvcore_ library. F-CNN and F-FNO are rather small, with considerably less parameters than GyroSwin and ViT. This is because of memory constraints.\n\nIt is important to note here that enabling these methods to efficiently operate at scale is not trivial, and probably requires a great deal of engineering, something that would arguably deserve it's own research work. This should both motivate the choice of our architecture and the reason why comparison with alternative neural networks is absent.\n\nModel|Inference [ms]|Memory [GB]|GFLOPs|Params [M]\n---|---|---|---|---\nF-CNN|569.6|11.1|19408.7|2.0\nF-FNO|963.3|36.9|N/A|1.3\nViT|87.7|1.2|1356.3|90.2\nViT (Large)|1141.9|5.5|14316.3|936.8\nGyroSwin|30.4|5.0|756.0|90.2\nGyroSwin (Large)|262.4|11.4|9513.5|998.3\n\n\nImportantly, the dataset we compiled IS already _\"small scale\"_, as it is based on simplifying assumptions (single turbulence mode, adiabatic electrons, etc.). The spatial resolution used is among the smallest possible seen in gyrokinetics simulations, and that still results in around 100MB per time slice. For reference, higher fidelity simulations require resolutions that would amount to 10s of GBs per time slice. Therefore, running the suggested baselines even at our small scale is prohibitive. This is the reason for introducing GyroSwin, which effectively balances performance and efficiency. We will clarify this in our manuscript.\n\n\n__(3) Even a simple baseline such as Gaussian process and MLPs achieve comparable quality on some datasets, even though they are missing several features such as physics informed learning dynamics. To me, the Table 2 appears too noisy to decide that GyroSwin is indeed the best performing model.__\n\nBelow we report the restructured version of Table 2 with the improved flux results, diagnostics and zonal flow predictions. _n/a_ means that the model is not capable of producing the result. We hope that the updated Table + updated experiments help to bring across our main point, i.e., **only GyroSwin can model turbulence**. Other methods can only predict quantities which arise from turbulence.\nWe also added a vanilla Transformer baseline to model the sequential evolution of the flux trace.\n\nMethod|Params|Input Shape|$Q$ ID|$Q$ OOD|$W(k_y)$|$Q_i(k_y)$|Zonal Flow\n---|---|---|---|---|---|---|---\nGPR|6|(4,)|36.95±9.24|25.89±8.78|n/a|n/a|n/a\nMLP|37.3K|(4,)|59.52±15.01|29.59±11.19|n/a|n/a|n/a\nLSTM|5.3M|(10,)+(5,)|57.99±15.62|73.14±23.63|n/a|n/a|n/a\nTransformer|14.1|(10,)+(5,)|74.86±20.26|77.67±28.41|n/a|n/a|n/a|\nQL|1|(16,85,32)|40.74±7.47|53.3±14.77|0.04±0.15|0.54±0.46|n/a\nGyroSwin-Lin|51.3M|(32\\,8,16,85,32)|38.03±10.73|31.19±8.73|n/a|n/a|n/a\nGyroSwin-S|90.2M|(32,8,16,85,32)|36.13±7.55|59.76±17.62|0.98±0.03|0.4±0.37|0.11±0.19\nGyroSwin-M|250.9M|(32,8,16,85,32)|36.15±7.45|76.24±16.99|0.98±0.03|0.3±0.37|0.15±0.22\nGyroSwin-L|998.3M|(32,8,16,85,32)|24.38±6.79|18.49±6.15|0.98±0.03|0.35±0.4|0.32±0.18\n\nWe agree that Table 2 was indeed too noisy and not too informative to ultimately decide the best performing model. By adding the conditioning and reporting error bars we demonstrate that the results obtained with GyroSwin are indeed significant.\n\n### Questions\n__(1) OOD Qi Performance: Why could explain why GyroSwin struggles in some OOD cases compared to GPR/MLP (Table 2), even though the modeling of GyroSwin has significantly more consideration for the underlying physics? I would have expected GyroSwin would overfit less.__\n\nThe main reason for this was that the flux head was missing the conditioning on the 4 physical parameters, which is provided to all the baselines. Hence the comparison was inherently unfair. After conditioning the flux head via FiLM on those parameters, we observe vast improvements as reported in the table above.\n\n__(2a) While Fig 3 shows a successful rollout, would it be possible to visualize the distribution of errors over multiple rollouts as a function of time (different dataset, different random seed etc.)?__\n\nThank you for that interesting suggestion. It is indeed possible to visualize the evolution of the error for GyroSwin across timesteps. Unfortunately, though, since GyroSwin models are expensive to train we do not yet have multiple models trained on multiple seeds at our disposal. However, we are planning to add such a visualization in the future.\n\n__(2b) How would such an error distributions in the rollouts differ for baseline methods?__\n\nBaseline methods such as MLP, GPR, or QL do not evolve over time, they predict average flux in the saturated state of the simulation. For the LSTM we observed that it tends to predict fluctuations close to the average of the distribution of fluxes. This is because the flux traces are very noisy and predicting the average is a strong local optima. We also observed that GyroSwin always predicts turbulent transport, even in the very beginning of the simulation where actually no turbulent transport occurs. We attribute this to the imbalanced nature of the dataset, as training samples of the saturated (turbulent) phase of the simulation are overrepresented.\n\n### References\n[1] Kumar N. et al. Turbulent transport driven by kinetic ballooning modes in the inner core of jet hybrid h-modes. Nuclear Fusion.\n\n[2] Citrin J. et al. Tractable flux-driven temperature, density, and rotation profile evolution with the quasilinear gyrokinetic transport model QuaLiKiz."}]

Author rebuttal:

Meta-review: This paper presents GyroSwin, a deep learning surrogate model for 5D gyrokinetic plasma turbulence simulations, which is a critical and computationally demanding problem in fusion research. GyroSwin extends a hierarchical Vision Transformer to 5D data and makes necessary adaptation. GyroSwin predicts the 5D distribution function, 3D electrostatic potential, and scalar fluxes, showing competitive performance against baselines. 

The reviewers recognize that the paper has good motivation and clear writing, and also raise important concerns, mainly including that the results in the main Table 2 are outperformed by several baselines, the paper lacks ablation study, and lacks sufficient comparison with other methods. After the rebuttal, the authors provided an updated Table 2, where their method is now also conditioned on operational parameters (previously did not) similar to baselines. The updated Table 2 shows that the proposed GyroSwin outperforms baselines, and is the only method that can model turbulence. The authors also provide ablation study on hyperparameters. Overall the main concerns raised by the reviewers are addressed. As an application-driven paper and addresses the very important problem of fusion simulation, this paper is above the acceptance bar. The authors are required to incorporate the rebuttal improvements into the camera-ready version of the paper.
Decision: Accept (poster)