GyroSwin: 5D Surrogates for Gyrokinetic Plasma Turbulence Simulations

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.

We 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:

1. 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.
2. 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.
3. We added additional baslines, like a simple Transformer model and provide results for scaling GyroSwin to larger parameter counts which results in significant gains.
4. 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.
5. 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.
6. We added a separate table for ablation studies for different prediction heads on a reduced data setup.

Weaknesses

(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)?

We 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.

(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?

It 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.

It 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.

Importantly, 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.

(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.

Below 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. We also added a vanilla Transformer baseline to model the sequential evolution of the flux trace.

We 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.

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.

The 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.

(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.)?

Thank 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.

(2b) How would such an error distributions in the rollouts differ for baseline methods?

Baseline 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.

References

[1] Kumar N. et al. Turbulent transport driven by kinetic ballooning modes in the inner core of jet hybrid h-modes. Nuclear Fusion.

[2] Citrin J. et al. Tractable flux-driven temperature, density, and rotation profile evolution with the quasilinear gyrokinetic transport model QuaLiKiz.

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.

Weaknesses

(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.

We 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).

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.

Questions

(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.

It 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.

If 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.

We'd first like to thank DBSK for their review and their constructive feedback.

We 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:

1. 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.
2. 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.
3. We added additional baslines, like a simple Transformer model and provide results for scaling GyroSwin to larger parameter counts which results in significant gains.
4. 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.
5. 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.
6. We added a separate table for ablation studies for different prediction heads on a reduced data setup.

Weaknesses

(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.

For additional information and proper referencing, we would like to point the reviewer to lines 123 to 137 in the Related Works section. Literature 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.

(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.

Below 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. We also added a vanilla Transformer baseline to model the sequential evolution of the flux trace. We 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.

(3) The paper lacks ablation studies.

We 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.

Questions

(1) what is the model computational complexity and how does that compare to the baselines?

Theoretical 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$

making it linear in the input resolution.

In 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.

When 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). Nonetheless, we report numbers for the simple baselines below.

References

[1] Kumar N. et al. Turbulent transport driven by kinetic ballooning modes in the inner core of jet hybrid h-modes. Nuclear Fusion.

[2] Citrin J. et al. Tractable flux-driven temperature, density, and rotation profile evolution with the quasilinear gyrokinetic transport model QuaLiKiz.

We sincerely thank the reviewer for their comments. We address them in the following subsections.

We 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:

1. Conditioning the flux prediction head of GyroSwin, which resulted in substantial gains in flux prediciton.
2. 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.
3. Additional baselines, (flux sequence Transformer); scaling GyroSwin to larger sizes.
4. Information on the input shapes for each method; diagnostic fields to clarify that certain evaluations are not applicable to the most baselines.
5. 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.
6. A separate table ablating different heads on a reduced data setup.

Weaknesses

(1) Extending model capabilities to support higher dimensions is not entirely new.

It 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.

This highlights the need for more scalability. Hierarchical 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.

Furthermore, we incorporate inductive-biases, namely channelwise mode separation and latent integration modules. These are tailored to our application domain, but are simple ways of including inductive biases into model design, which may be relevant to the community.

(2) The concept of multitask prediction is a well-established technique in AI.

This 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.

(3) As an application-specific approach, the results in Table 2 are outperformed by generic and relatively simple approaches in many cases.

As mentioned above, the main cause of the result inconsistency at submission was that the flux head was not conditioned on operational parameters.

Below 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.

(4) Some acronyms are not properly defined. We address this point in question (1).

Questions

(1) What is W-MSA and SW-MSA in Figure 2? Thanks 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.

• 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.`shape` = `(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.

• 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)=`roll`(\text{W-MSA}(`roll`(x, −W/2)), +W/2)$, where $W$ is the window size and $`roll`$ is a matrix roll operation, which circularly shifts the input along the spatial dimensions by a given displacement (equivalent to np.roll).

Please refer to the 2021 Shifted Window Transformer paper [5] for simpler 2D visualizations.

(2) More insights into how cross-attention is related to 5D-3D integral?

The 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).

First, 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.

Then, 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.

Finally, 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.

(3) In Tab2, traditional GPR and MLP approaches seem to excel at ID and OOD scenarios. Is the problem (5D->0D) not sophisticated enough?

This 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.

(4) In Figure 3, why QL, GPR, and MLP results are flat lines?

Thank you for pointing out that this isn't clear. Generally, 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.

(5) Is it possible to use traditional approaches (e.g., MLP) to obtain 5D->1D evaluation results and show them together in Figure 4?

Yes, 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.

References

[1] Kumar N. et al. Turbulent transport driven by kinetic ballooning modes in the inner core of jet hybrid h-modes. Nuclear Fusion.

[2] Citrin J. et al. Tractable flux-driven temperature, density, and rotation profile evolution with the quasilinear gyrokinetic transport model QuaLiKiz.

[3] Wang M. et al. (2017). Factorized Convolutional Neural Networks. ICCV.

[4] Tran A. et al. (2023). Factorized Fourier Neural Operators. ICLR.

[5] Liu Z. et al. (2021). Swin Transformer: Hierarchical Vision Transformer using Shifted Windows. ICCV.

[6] Liu Z. et al. (2021). Video Swin Transformer.

[7] Kim P. et al., Moon T. (2023). SwiFT: Swin 4D fMRI Transformer. NeurIPS.

[8] Alkin B. et al. (2024). Universal Physics Transformers: A Framework For Efficiently Scaling Neural Operators. NeurIPS.

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.

We 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:

1. 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.
2. 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.
3. Additional baslines, like a simple Transformer model and provide results for scaling GyroSwin to larger parameter counts which results in significant gains.
4. 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.
5. 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.
6. A separate table for ablation studies for different prediction heads on a reduced data setup.

Weaknesses

(1) GyroSwin does not consistently outperform other methods (Table 2).

As 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. Below 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.

(2) Lacks theoretical analysis why the proposed approach should work well.

We 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.

(3) Efficiency comparison is not entirely fair.

We 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).

In 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.

(4) Training cost comparisons with other models are missing.

We 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.

(5) Lacks sufficient comparison with other methods.

Unfortunately, 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).

Scaling 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.

As 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].

Finally, GyroSwin is based on ViTs, in particular shifted window attention and special cross-attention layers, which are arguably among the latest deep learning techniques.

Questions

(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?

This is discussed in the answer to weakness (1).

(1) Ablation studies to better understand the contributions of each module.

We 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.

(2) Training cost, number of parameters, and GFLOPs for each model.

See response to weakness (4).

(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?

See response to weakness (5).

References

[1] Kumar N. et al. Turbulent transport driven by kinetic ballooning modes in the inner core of jet hybrid h-modes. Nuclear Fusion.

[2] Citrin J. et al. Tractable flux-driven temperature, density, and rotation profile evolution with the quasilinear gyrokinetic transport model QuaLiKiz.

[3] Najlaoui A. et al. Verifying turbulence model reduction in high β tokamak plasmas (2025) Plasma Phys. Control