Title: GyroSwin: 5D Surrogates for Gyrokinetic Plasma Turbulence Simulations
Venue: NeurIPS 2025
Accepted: yes
Keywords: 5D Swin Transformer

Gyrokinetics

Nuclear Fusion

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


Peer reviews:
{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis 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.\n\nSTRENGTHS AND WEAKNESSES:\n*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.\n\nQUESTIONS:\n1. 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?\n\nLIMITATIONS:\nyes", "TITLE": "Official Review", "RECOMMENDATION": 4, "REVIEWER_CONFIDENCE": 3, "CLARITY": 3, "IMPACT": 2, "ORIGINALITY": 2, "DATE": "2025-07-02"}

comments: SUMMARY:
The paper presents a novel deep learning approach, GyroSwin, aimed at efficiently simulating nonlinear gyrokinetic plasma turbulence. 

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

STRENGTHS AND WEAKNESSES:
**Strengths**

 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.

2. Although the modules of the proposed surrogate model are existing, the design of each module is principled and strongly related to the task.

**Weaknesses**

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.

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.
TITLE: Official Review
CLARITY: 3
IMPACT: 3
ORIGINALITY: 3
DATE: 2025-07-03

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

STRENGTHS AND WEAKNESSES:
Strengths:
The approach appears sound. It makes use of vision transformer in 5D distribution function over time.
The paper explains the context well, and has demonstrated the significance of the problem in application contexts.

Weaknesses:
The paper can benefit from discussion of related works in neural approaches to solving Gyrokinetics problems, and their weaknesses compared to the proposed model. 

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.

The paper lacks ablation studies.

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

QUESTIONS:
what is the model computational complexity and how does that compare to the baselines? It seems less clear in the manuscript.

LIMITATIONS:
yes
TITLE: Official Review
CLARITY: 2
IMPACT: 2
ORIGINALITY: 3
DATE: 2025-07-05

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nIn 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.\n\nSTRENGTHS AND WEAKNESSES:\n**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.\n\nQUESTIONS:\n(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?\n\nLIMITATIONS:\nyes", "TITLE": "Official Review", "RECOMMENDATION": 4, "REVIEWER_CONFIDENCE": 4, "CLARITY": 4, "IMPACT": 3, "ORIGINALITY": 2, "DATE": "2025-07-14"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis 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.\n\nSTRENGTHS AND WEAKNESSES:\nStrength:\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.\n\nQUESTIONS:\n- 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?\n\nLIMITATIONS:\nYes.", "TITLE": "Official Review", "RECOMMENDATION": 3, "REVIEWER_CONFIDENCE": 4, "CLARITY": 3, "IMPACT": 3, "ORIGINALITY": 3, "DATE": "2025-07-15"}