Diff-PIC: Revolutionizing Particle-In-Cell Nuclear Fusion Simulation with Diffusion Models

## Abstract

Abstract The rapid development of AI highlights the pressing need for sustainable energy, a critical global challenge for decades. Nuclear fusion, generally seen as an ultimate solution, has been the focus of intensive research for nearly a century, with investments reaching hundreds of billions of dollars. Recent advancements in Inertial Confinement Fusion have drawn significant attention to fusion research, in which Laser-Plasma Interaction (LPI) is critical for ensuring fusion stability and efficiency. However, the complexity of LPI upon fusion ignition makes analytical approaches impractical, leaving researchers depending on extremely computation-demanding Particle-in-Cell (PIC) simulations to generate data, presenting a significant bottleneck to advancing fusion research. In response, this work introduces Diff-PIC, a novel framework that leverages conditional diffusion models as a computationally efficient alternative to PIC simulations for generating high-fidelity scientific LPI data. In this work, physical patterns captured by PIC simulations are distilled into diffusion models associated with two tailored enhancements: (1) To effectively capture the complex relationships between physical parameters and corresponding outcomes, the parameters are encoded in a physically-informed manner. (2) To further enhance efficiency while maintaining high fidelity and physical validity, the rectified flow technique is employed to transform our model into a one-step conditional diffusion model. Experimental results show that Diff-PIC achieves 16,200 × \times × speedup compared to traditional PIC on a 100 picosecond simulation, with an average reduction in MAE / RMSE / FID of 59.21% / 57.15% / 39.46% with respect to two other SOTA data generation approaches.

## 1 Introduction

Sustainable energy stands as one of the paramount challenges of our era, particularly with the rapid advancement of AI. The recent successful demonstration of fusion ignition [] underscores the transformative potential of fusion as a sustainable energy source. In 2023 and 2024, the National Ignition Facility (NIF) achieved groundbreaking milestones, generating 3.4 MJ and 5.2 MJ of fusion energy, respectively, from 2.2 MJ input energy. Given that the estimated output could reach $\sim$120 MJ [], there is a growing demand for a deeper understanding of the fundamental science behind ignition efficiency, especially the physical mechanisms governing the interaction between the laser and the plasma emitted when the laser bombards the fuel pellet. However, the Laser-Plasma Interaction (LPI) is a complex multi-body problem, which traditionally relies on time-stepping method, particularly, Particle-in-Cell (PIC) simulations [, , , , ]. Despite being the preeminent standard for modeling the physics of LPI, PIC simulations are exceedingly intensive in computation, often requiring tens of millions of CPU hours, consuming millions of dollars in order to obtain meaningful outputs [, , ]. The computational overhead of PIC simulations has become a daunting bottleneck in fusion research, raising the pressing need for innovative methodologies capable of generating high-quality scientific data with substantially reduced computational burden.

Over the years, numerous CPU-GPU implementations have been developed for PIC simulations [, , ]. While invaluable, these efforts remain within the scope of the time-stepping approach that iteratively executes over infinitesimal time intervals, falling short in addressing the inherent computational overhead of long-term simulations. Recent advancements in generative AI – diffusion models, however, present a novel approach to bypass the constraint. Diffusion models [, , ] have demonstrated exceptional capabilities in Computer Vision (CV) [, , ], synthesizing highly complex data distributions that match real data with high fidelity. From the perspective of energy-based models [, ], a diffusion model effectively constructs a highly complex energy field that governs the evolution of variables, analogous to the motion of particles in cells. In fact, diffusion models are rooted in the diffusion concept in physics, where particles move according to the energy field. This has sparked significant interest in their potential for generating scientific data, as recent applications of diffusion models in molecular dynamics simulations have demonstrated their promise in this domain [, ].

Although diffusion models exhibit outstanding compatibility for generating PIC simulation data, two critical research gaps must be addressed. ❶ Physical soundness must be ensured.
In contrast to traditional PIC simulations that directly take continuous physical parameters as constraints,
it remains unclear how diffusion models can effectively capture and distill complex physical patterns. ❷ Substantial efficiency improvement must be achieved. The step-by-step denoising process in diffusion models is computationally demanding. Although the requirement for infinitesimal time intervals has been relaxed, the process still presents significant challenges similar to those faced by time-stepping methods, limiting the practicality of diffusion models as advanced alternatives for PIC simulations.

Figure: Figure 1: Overview of the proposed method.
Refer to caption: extracted/5904548/figures/intro.png

In light of these challenges, we propose a distillation framework for physical patterns, titled Diff-PIC, that leverages diffusion models to efficiently generate a snapshot of arbitrary time, under arbitrary simulation parameters within certain ranges. To specifically address the two challenges above: ❶ We develop a conditional diffusion model with a Physically-Informed Parameter Encoder. This encoder allows the model to capture the relationship between continuous physical parameters and PIC simulation data, distilling physical phenomena into Diff-PIC. ❷ We employ the rectified flow technique to eliminate the requirement for the multi-step denoising, further optimizing the runtime efficiency of Diff-PIC. As highlighted in Fig. 1, orders-of-magnitude speedup is achieved compared to PIC simulations and other generative approaches including Generative Adversarial Networks (GAN) [] and Normalizing Flow (NF) [], dubbed "GAN-PIC" and "NF-PIC".

In summary, our work represents the first known effort to tackle the imperative challenges associated with generating high-quality PIC simulation data for LPI using diffusion models renowned in the field of generative AI. The core contributions of our work include:

- •
We propose Diff-PIC, a pioneering study that utilizes diffusion models as a computationally efficient alternative to PIC simulations (see §3.1). By making all resources publicly available, we aim to establish Diff-PIC as a robust baseline and a valuable benchmark for the research community, thereby accelerating the advancement in scientific data generation for nuclear fusion research.
- •
We develop a physically-informed conditional diffusion model (see §3.2) that seamlessly integrates physical simulation parameters into the diffusion model. The designed condition encoder facilitates the generalization of effective simulation data within and beyond existing simulation parameters, endowing the model with robust generalization capabilities and adaptable transferability.
- •
We implement the rectified flow technique to transform our model into a one-step diffusion model (see §3.3), thereby enhancing its efficiency in generating high-fidelity fusion data.
- •
Experimental results in §4.2 demonstrate that our method achieves a remarkable speedup of 16,200 times compared to traditional PIC simulations while preserving high fidelity and physical validity of the generated data – 0.014, 0.019, 1.03 in terms of MAE, RMSE, and FID, respectively, with an average reduction of 59.21% / 57.15% / 39.46% compared to two other SOTA generative models.

## 2 Background

Inertial Confinement Fusion (ICF) is a method of achieving controlled nuclear fusion by using intense energy pulses to compress and heat small fuel pellets [, ], typically containing isotopes of hydrogen such as deuterium and tritium. This process unfolds the nuclear fusion reaction as delineated below:

$$ ${}^{2}\mathrm{H}+{}^{3}\mathrm{H}\rightarrow{}^{4}\mathrm{He}+{}^{1}\mathrm{ neutron}+\textbf{Energy}.$ (1) $$

Given the ubiquity of these hydrogen isotopes in the ocean, nuclear fusion holds immense potential to provide "near-infinite" energy by achieving the necessary temperature and pressure conditions to initiate fusion reactions, attaining a positive net energy gain (i.e., output energy surpassing the input). To optimize the sophisticated initiation of fusion, an advanced understanding of the underlying LPI mechanism is essential. For this purpose, PIC is considered a crucial tool to provide theoretical insights into LPI, due to its capability of predicting and interpreting physical phenomena.

Particle-in-Cell Simulations. The PIC method is a computational technique widely used in the study of plasma physics and fusion energy research [, , ]. Developed in the mid-20$th$ century, the PIC method has become a cornerstone in the simulation of complex plasma behaviors, enabling researchers to delve into the intricate dynamics of particles and electromagnetic fields [, ]. To highlight, PIC is especially useful in LPI studies [, , ], which involves complex dynamics of electrons and ions. PIC simulations track the trajectories and interactions of these charged particles under the influence of electromagnetic fields, providing insights into both shock wave formation and heating mechanisms that are essential for ICF.

In essence, PIC is an iterative time-stepping method applied to atomic particles such as electrons and ions. Within each iteration, particles are systematically arranged into discrete cells according to their spatial distribution, with their positions and velocities being updated over infinitesimally small time steps, typically on the scale of femtoseconds ($10^{-15}$ seconds). Unlike molecular dynamics widely applied in biology [, ], PIC simulations in LPI are characterized by intensive electromagnetic fields, which exert a significant influence on particle trajectories as follows:

$$ $\frac{d\mathbf{v}_{i}}{dt}=\frac{q_{i}}{m_{i}}\left(\mathbf{E}(\mathbf{r}_{i}, t)+\mathbf{v}_{i}\times\mathbf{B}(\mathbf{r}_{i},t)\right).$ (2) $$

For a particle at position $\mathbf{r}_{i}$ with charge $q_{i}$ and mass $m_{i}$, the equation describes how the velocity $\mathbf{v}_{i}$ evolves according to an energy field that consists of electric field $\mathbf{E}$ and magnetic field $\mathbf{B}$. To ensure accuracy, simulating LPI at the scale of mere hundreds of picoseconds (i.e., $10^{-10}$ seconds) requires hundreds of thousands of sophisticated PIC iterations. This imposes significant demands on computational storage and processing capabilities. As a result, the PIC methodology has emerged as a stringent bottleneck in fusion research, significantly constraining progress in this domain.

Diffusion models have emerged as a prominent class of generative models within the realm of artificial intelligence, offering an innovative methodology for the synthesis of high-fidelity data [, , ]. Named after the physical concept, diffusion models use the idea of diffusion, which in physics refers to the random movement of particles from regions of high concentration to regions of lower concentration, often driven by thermal energy.

In machine learning, diffusion models are generative models that progressively add noise to data in a forward process, then gradually remove it in the reverse process to generate new data. From the perspective of energy-based models, the reverse process can be seen as moving variables through an energy landscape, where the model transitions variables from high-energy, noisy states (where data are unstructured) to low-energy, clean states (where data are structured). In this sense, diffusion models generate data by allowing variables to evolve in a manner akin to particles diffusing within an energy field, granting the opportunity to effectively apply physical methods to evaluate the diffusion, specifically, Langevin dynamics:

$$ $d\mathbf{x}=-\nabla_{\mathbf{x}}E(\mathbf{x})dt+\sqrt{2\beta}d\mathbf{W_{t}}$ (3) $$

The evolution of variables $\mathbf{x}$ is influenced by both the energy field $E(\mathbf{x})$ and the stochastic motion term $\mathbf{W_{t}}$ regulated by the diffusion coefficient $\beta$. Through the dynamics, variables move towards low energy regions, conceptually mirroring the behavior of particles in PIC simulations. This similarity encourages an exploration of diffusion models as potential alternatives to traditional PIC simulations.

## 3 Diff-PIC

In this section, we introduce Diff-PIC, a physically-informed conditional diffusion model tailored for generating high-fidelity synthetic data for LPI in nuclear fusion. As illustrated in Fig. 2, accepting physical parameters as inputs, the parameters are encoded and integrated with the model main body through the physically informed parameter encoder. Additionally, the rectified flow mechanism is adopted to break free from the time-stepping paradigm in both PIC simulation and traditional diffusion models, unlocking the model’s capability to generate scientific data in one single step.

Figure: Figure 2: Workflow of the proposed Diff-PIC compared to traditional step-by-step PIC simulation.
Refer to caption: extracted/5904548/figures/overview_ng.png

### 3.1 The Overall Distillation Framework

Since LPI is governed by the behavior of electromagnetic field that dominates plasma dynamics, the tasks specifically assigned to Diff-PIC are to generate high-fidelity electric fields as representative cases under various physical parameters. Due to the polarization of the input laser, the resulting electric field only oscillates in a 2D plane, making 2D field descriptions sufficient to capture primary features of LPI. This approach is also a standard practice in PIC simulations.

In this work, we focus on the following critical physical parameters in LPI: “Electron Temperature ($T_{e}$) in $keV$,” “Ion Temperature ($T_{i}$) in $keV$,” and “Laser Intensity ($I$) in $W/m^{2}$.” Additionally, to enable direct and versatile generation of the electric field at a specific simulation time (a snapshot), we include a user-defined parameter $t_{as}$ that represents the simulation time for an arbitrary snapshot. Through a standard learning process in diffusion models [, ], the relationships between the four parameters and the resulting electric field are captured, as the model learns to progressively transform Gaussian noise into realistic electric fields. Once trained, the model takes the four parameters as inputs and produce electric field snapshots $\mathbf{E}(t_{as},\theta)$ corresponding to the specified parameters $\theta=\{T_{e},T_{i},I\}$. These generated snapshots can be used for multiple purposes, including data augmentation, parameter exploration, and studying the effects of different physical parameters. Furthermore, the diffusion model aids in the interpretation and visualization of complex plasma phenomena, providing valuable insights for researchers in the field of nuclear fusion. Essentially, the proposed techniques offer two advantages:

- •
Data Dependency Relaxation. We treat snapshots from various parameter combinations as distinct distributions. This approach decouples data dependency in model training, enabling the model to efficiently learn from individual snapshots while generalizing across a wide array of scenarios.
- •
Systemic efficiency. Unlike traditional PIC simulations that generate data sequentially over time, the proposed diffusion model can directly produce data for any target snapshot (see Fig. 4). This non-sequential behavior allows for substantially more efficient data generation and analysis (see Table 3), enabling researchers to focus on specific times of interest without needing to simulate the entire LPI process from the beginning.

### 3.2 Physically-Informed Parameter Encoder

PIC simulations employ continuous physical parameters as inputs, which necessitate a seamless and continuous transition in the resulting synthesized data as input parameters are adjusted. Consequently, an encoder is considered exceptionally beneficial in this scenario, responsible for transforming domain-specific inputs into embeddings comprehensible by the model, meanwhile preserving the physical continuity of output snapshots.
In particular, these inputs comprise the simulation parameters $\theta$ and the target simulation snapshot $t_{as}$ that the conditional diffusion model aims to generate.

Given the extensive range of physically feasible parameters and the limited data available during training, an optimal encoder must excel in both interpolation and extrapolation — critical measures of the model’s ability of generalization. Interpolation capability refers to the encoder’s proficiency in generating suitable embeddings for new parameters that, although not encountered during training, lie between observed parameters. Extrapolation capability, conversely, pertains to generating embeddings for parameters that fall outside the range of those observed during training. Notably, both capabilities are indispensable for addressing the LPI problem, in order to cover a large enough, and fine-grained enough parameter space to provide sufficient insight into further LPI evaluations.

Figure: Figure 3: The proposed PIPE.
Refer to caption: x1.png

To meet the two essentials, we introduce a Physically-Informed Parameter Encoder (PIPE) as shown in Fig. 3. To encode the simulation parameters $\theta$, we incorporate two distinct types of encoders. For interpolation, we employ Positional Encoding [] (denoted “$\sim$” in a circle) to learn a continuous representation of input parameters, enabling smooth transitions between observed parameters. To augment extrapolation capability, we enhance the encoding process with a polynomial encoder (“P” in a circle), utilizing transformation functions constructed as a linear combination of polynomial basis functions $f_{i}(\theta)$ of varying degrees:

$$ $\mathcal{P}(\theta)=\sum_{i=0}^{n}f_{i}(\theta),$ (4) $$

where $n$ denotes the maximum order of polynomial terms, and the polynomial $\mathcal{P}(\theta)$ can be chosen as Chebyshev polynomials and Legendre polynomials, based on the characteristics of the parameter space. This polynomial enhancement allows the encoder to generate plausible embeddings for parameters well beyond those encountered during training, ensuring robust performance across a broader spectrum of simulation scenarios.
Subsequently, we concatenate (“C” in a circle) the embeddings from these two encoders and apply a Multi-Layer Perceptron (MLP) to further refine the embeddings. The MLP, with its trainable parameters, learns to combine and transform the concatenated embeddings, resulting in a more informative representation of the input parameters.
For encoding the simulation time step $t_{as}$, we utilize Positional Encoding  [] followed by an MLP layer. This approach is specifically chosen to learn continuous representations that facilitate smooth transitions between consecutive snapshots, thereby enhancing the model’s temporal coherence.
In summary, this design offers the following advantages:

- •
Algorithmic generalization. PIPE improves the generalizability of the conditional diffusion model (see Table 2). The dual-encoding strategy captures non-linear relationships by incorporating both positional and polynomial encodings, empowering the model to adeptly manage a diverse array of simulation parameters, ranging from scenarios encountered during the training phase to parameters that lie beyond the spectrum of the training data.
- •
Adaptive transferability. By fine-tuning the pre-trained model with a new dataset, this methodology facilitates adaptation to other fields where precise and efficient simulations are imperative for deciphering intricate physical phenomena.

### 3.3 Rectified Flow-Based Acceleration

To enable rapid generation of high-fidelity synthetic data, we employ the Rectified Flow Acceleration (RFA) technique in model optimization. Based on the principles of rectified flow [, , ], RFA converts the original complex denoising trajectory from initial noise $\epsilon$ to the target electric field snapshot $\mathbf{E}(t_{as})$ into a direct and straight path – the shortest route between two distributions.
During training, RFA minimizes the following objective function with straight ordinary differential equations:

$$ $\operatorname*{arg\,min}_{\zeta}\quad\mathbb{E}\left[\int_{0}^{1}\left\|( \mathbf{E}(t_{as},\theta)-\epsilon)-\zeta(\mathbf{E}_{t},t~{}|~{}t_{as},\theta )\right\|^{2}dt\right],$ (5) $$

where $\mathbf{E}_{t}=t\mathbf{E}(t_{as},\theta)+(1-t)\epsilon$ denotes the linear interpolation between $\mathbf{E}(t_{as},\theta)$ and $\epsilon$ across the diffusion timeline, with $t$ ranging from 0 to 1. The score-based model $\zeta$, approximated using the U-Net, defines the learned trajectory. By minimizing the expectation of the squared deviations between the straight path $\mathbf{E}(t_{as},\theta)-\epsilon$ and the learned trajectory $\zeta(\mathbf{E}_{t},t~{}|~{}t_{as},\theta)$, RFA promotes the adoption of the shortest and most direct path in the denoising process, thus significantly reducing the denoising time.
Once this time-dependent score-based model $\zeta$ is trained, we further straighten the learned trajectories through an interactive reflow procedure  [].

In summary, the RFA module provides additional benefits for our paradigm:

- •
Streamlined Denoising Process. RFA significantly accelerates the denoising process (see Table 3) by converting the complex data trajectory from initial noise to the target snapshot into a direct denoising step. By distilling the typically winding diffusion path into the shortest route between two distributions, RFA greatly reduces the time required for generating high-fidelity synthetic data.
- •
Robust Optimization.
RFA leverages the principles of rectified flow to minimize deviations between the winding data trajectory and the optimal, shortest path. This direct path approach reduces the possibility of error accumulation that can occur with more winding, iterative methods.

## 4 Evaluation

### 4.1 Experimental Setup

Datasets.
We provide a new dataset comprising 6,615 simulations across varied physical simulation parameters, each containing 80 snapshots of electric fields along two orthogonal directions denoted E1 and E2. The data were generated by OSIRIS [], a well-established PIC simulation software suite. The dataset covers diverse parameters, including $T_{e}$, $T_{i}$, and $I$, all of which are critical parameters influencing the resultant electric fields. To foster further advancements in fusion and scientific data generation research, we will release the dataset publicly.

Metrics. To validate physical soundness, Mean Absolute Error (MAE) and Root Mean Squared Error (RMSE) are used to evaluate the electric field difference and the energy difference between the Diff-PIC generated and the ground truth produced by PIC simulations. To better demonstrate the relative error, the dataset is normalized to [0,1]. To further evaluate the difference in the generated and the ground truth data distributions, the Fréchet Inception Distance (FID) metric is also employed, reflecting the fidelity of the electric fields produced by DiffPIC.

Baselines. We compare Diff-PIC with two other SOTA generative models, Generative Adversarial Networks [] and Normalizing Flow []. The baseline models are implemented based on the setups provided in the original papers. Since neither of them originally supports learning meaningful embeddings for the physical parameters, for fair comparison, we equip them with the proposed PIPE to establish two baselines: GAN-PIC and NF-PIC.

Diff-PIC configurations. The architectural foundation of our model is predicated on the U-Net framework [], incorporating a series of three down-sampling blocks followed by three up-sampling blocks. The training regimen of our model encompasses an extensive 600 epochs, employing a batch size of 64, a configuration empirically validated to secure model convergence. Furthermore, the training protocol adheres to a fixed learning rate of $5\times 10^{-4}$, optimized via the Adam optimizer []. The optimization objective is delineated in Eq. 5. In addition, in terms of the encoders, the positional encoding mechanism yields a 16-dimensional embedding from a single input, while the polynomial encoder operates with a highest polynomial order of 4.

Figure: Figure 4: Visualization and Comparison of PIC simulations and Diff-PIC.
Refer to caption: x2.png

### 4.2 Main Results

In this section, we evaluate Diff-PIC on three key aspects for comprehensive comparisons. (1) The interpolation ability and extrapolation ability. (2) The physical validity of the Diff-PIC generated data. (3) The speedup and power efficiency compared to traditional PIC simulations.

**Table 1: Quantitative results for interpolation evaluation.**
| Method | Training Set for E1 | Testing Set for E1 | Training Set for E2 | Testing Set for E2 |  |  |  |  |  |  |  |  |
| --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- |
| MAE$\downarrow$ | RMSE$\downarrow$ | FID$\downarrow$ | MAE$\downarrow$ | RMSE$\downarrow$ | FID$\downarrow$ | MAE$\downarrow$ | RMSE$\downarrow$ | FID$\downarrow$ | MAE$\downarrow$ | RMSE$\downarrow$ | FID$\downarrow$ |  |
| GAN-PIC | 4.59e-2 | 5.31e-2 | 2.32 | 4.73e-2 | 5.84e-2 | 2.51 | 1.82e-2 | 2.07e-2 | 0.973 | 1.97e-2 | 2.18e-2 | 1.03 |
| NF-PIC | 4.46e-2 | 5.12e-2 | 2.06 | 4.61e-2 | 5.35e-2 | 2.42 | 1.70e-2 | 2.03e-2 | 0.914 | 1.86e-2 | 2.45e-2 | 0.986 |
| Diff-PIC | 1.56e-2 | 2.67e-2 | 1.21 | 1.68e-2 | 2.29e-2 | 1.62 | 7.95e-3 | 9.32e-3 | 0.328 | 8.26e-3 | 1.03e-2 | 0.341 |

Interpolation and Extrapolation. To evaluate the interpolation capability of Diff-PIC, we sample a specified range for each simulation parameters ($T_{e}$, $T_{i}$, and $I$), totally 500 simulations and $500\times 80=40,000$ snapshots. Then, we randomly split these 500 simulations into training and testing set with the ratio of 80% and 20%. We train Diff-PIC and baselines on the training set, and report the performance of Diff-PIC on the training and testing set in Table 1. The reasonably low MAE, RMSE, and FID scores indicate that the proposed Diff-PIC is able to synthesize high-quality scientific data similar to what PIC generates, meanwhile significantly outperforming baselines on all three metrics. On average, Diff-PIC achieves 59.25% reduced MAE, 57.77% reduced RMSE, and 49.21% reduced FID for both testing sets, compared to the baselines.

In addition, Fig. 4 compares the results for one randomly selected simulation produced by Diff-PIC and PIC respectively in the testing set. Throughout the snapshots ($t_{as}=1\rightarrow 30\rightarrow 50\rightarrow 80$ for example), the distributions of the synthetic data closely follow the ground truths, indicating that physical continuity is maintained over time. Additionaly, the distributions on the right demonstrate that our proposed model successfully captures the data distributions in the ground truths.

In terms of extrapolation capability evaluation, we progressively extend the range of physical parameters from the range selected for training. In particular, the ranges are extended by 10% and 20% respectively as case studies. The results on Table 2 demonstrate that Diff-PIC achieves approximately only 2% relative absolute error in extrapolation tasks, significantly outperforming the other generative counterparts. On average, Diff-PIC achieves 59.16% reduced MAE, 56.53% reduced RMSE, and 29.70% reduced FID for all test cases, compared to the baselines.

**Table 2: Quantitative results for extrapolation evaluation.**
| Method |  | E1 10% |  |  | E1 20% |  |  | E2 10% |  |  | E2 20% |  |
| --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- |
| MAE$\downarrow$ | RMSE$\downarrow$ | FID$\downarrow$ | MAE$\downarrow$ | RMSE$\downarrow$ | FID$\downarrow$ | MAE$\downarrow$ | RMSE$\downarrow$ | FID$\downarrow$ | MAE$\downarrow$ | RMSE$\downarrow$ | FID$\downarrow$ |  |
| GAN-PIC | 5.24e-2 | 6.32e-2 | 1.97 | 5.73e-2 | 6.98e-2 | 2.15 | 2.18e-2 | 3.25e-2 | 1.15 | 2.93e-2 | 3.89e-2 | 1.29 |
| NF-PIC | 4.74e-2 | 5.31e-2 | 1.85 | 5.41e-2 | 6.46e-2 | 2.08 | 1.90e-2 | 3.16e-2 | 1.04 | 2.63e-2 | 3.42e-2 | 1.17 |
| Diff-PIC | 1.83e-2 | 2.40e-2 | 1.74 | 2.18e-2 | 2.62e-2 | 1.82 | 9.47e-3 | 1.36e-2 | 0.536 | 1.13e-2 | 1.85e-2 | 0.673 |

Physical Validity. In addition to the high quality of the synthetic data, we test the physical validity of the generated electric fields at each snapshot with energy, a fundamental concept used to describe the general property of a physical system. In particular, we randomly pick a simulation in testing set, the energy results of its electric fields are illustrated in Fig. 5. We observe that the pattern of the ground truth is captured by Diff-PIC, with fairly low MAE and RMSE incurred. Especially, for electric field E2, the oscillation in energy is properly preserved in the synthetic data. Furthermore, this evaluation highlights that Diff-PIC is capable of generating sequentially continuous data, demonstrating the effectiveness of the proposed distillation paradigm and the Physically-Informed Parameter Encoder.

Figure: Figure 5: Energy evaluation of electric fields for training and test sets with 80 snapshots.
Refer to caption: x3.png

**Table 3: Speedup and Energy Consumption Reduction.**
| Method | PIC-100 ps | Diff-PIC-GPU | Diff-PIC-CPU | GAN-PIC-GPU | GAN-PIC-CPU | NF-PIC-GPU | NF-PIC-CPU |
| --- | --- | --- | --- | --- | --- | --- | --- |
| Speedup | 1.00$\times$ | 1.62e4$\times$ | 519$\times$ | 1.56e4$\times$ | 523$\times$ | 9.21e2$\times$ | 24$\times$ |
| Energy Reduction | 1.00$\times$ | 1.01e4$\times$ | 1.05e3$\times$ | 1.13e4$\times$ | 1.47e3$\times$ | 8.14e2$\times$ | 53$\times$ |

Speedup and Power Efficiency. In addition to the traditional PIC approach, baselines are also included in the comparison shown in Table 3. The PIC simulations are run on the Perlmutter supercomputer in the National Energy Research Scientific Computing (NERSC) facility, with AMD EPYC 7763 CPUs. As essential fusion phenomena typically appear at approximately 100 ps, the costs of PIC simulation at 100 ps are selected as the reference across the comparisons. The GPU results for Diff-PIC and baselines are obtained on an Nvidia RTX 4090 GPU to demonstrate the availability of this approach to general users. The CPU results for these approaches are acquired on an Intel 13th Gen i9-13900KF CPU. The results highlight that Diff-PIC-GPU achieves over $10^{4}\times$ speedup versus traditional PIC simulation, as well as $10^{4}\times$ in terms of reduced energy cost.

### 4.3 Discussion

To provide further insights and highlight the value of this work, it is worth noting that our approach exhibits outstanding scalability compared to traditional PIC simulations. In PIC, for $N$ particles, the computing complexity can generally reach a formidable $TN\text{log}N$ with $T$ time steps. In contrast, Diff-PIC is not sensitive to the number of particles in space nor the number of time steps, since it focuses on generating the macroscopic data (e.g., electric field, which is usually considered more informative than individual particle status) directly for a specific time. For larger particle spaces, the speedup achieved by Diff-PIC can be easily improved by extra orders of magnitude, further accelerating the research of fusion, or other research areas involving large-scale PIC simulations.

## 5 Related Work

Particle-in-Cell Simulations have long been fundamental to modeling physical processes in fusion research [, ]. However, the computational intensity of PIC simulations presents significant challenges [].
To mitigate these computational constraints, various methods have explored GPU and hybrid CPU-GPU acceleration technologies. Studies such as [, , , , ] have utilized parallel computing, high memory bandwidth, and multiple processors to expedite simulations. Architecturally, the simulator optimized for the Kepler GPU architecture, as discussed in [], underscores the potential of specific GPU architectures to enhance simulation efficiency. For more intricate simulations, research efforts like [, ] have developed hybrid CPU-GPU implementations, [] introduced a hybrid approach for multi-core and multi-GPU systems, highlighting the continuous integration and evolution of these technologies in advancing PIC simulations. Despite these advancements, these approaches remain reliant on the fundamental PIC framework, which may not completely address the computational burden due to the inherent algorithmic complexity of the PIC method.
In recent years, rapid advancements in deep learning have opened new pathways for accelerating scientific simulations. Machine learning-based approaches have emerged, such as predicting a vector space that approximates the PIC system solution [], and learning the probability of interactions between potential collision pairs []. However, the approach by [] depends on a pre-computed vector space and may not generalize well to novel scenarios, while []’s method focuses on binary interactions, overlooking the complex many-body interactions in PIC simulations.

Contrastingly, our proposed method overcomes these limitations by employing a conditional diffusion model to distill the complex patterns captured by PIC simulations from a limited training dataset. Utilizing a time-dependent score-based model, our approach can efficiently generate high-fidelity synthetic data (see Fig. 4 and Table 1) without the computational expense of traditional PIC algorithms (see Table 3). This results in a significant reduction in computational cost while maintaining high simulation accuracy. Moreover, our method is highly adaptable, as it can be fine-tuned with minimal additional data to suit various physical parameters. This flexibility renders our approach suitable for a wide range of applications in fusion research and beyond, where efficient and accurate simulations are crucial for understanding complex physical phenomena.

Diffusion Models in Scientific Research have emerged as a formidable tool across a myriad of scientific domains. These models, which employ a stochastic process to incrementally convert a pristine data sample into a noise-distributed version and subsequently reverse this process. For instance, in materials science and chemistry,  [] introduced a diffusion model for molecular dynamics simulations, demonstrating the generalizability in generating molecular trajectories.  [] presented an approach that integrates diffusion models with coarse-grained molecular dynamics to develop a new force field for simulating protein dynamics. By leveraging score-based generative models, they trained a model on coarse-grained structures to produce a force field that enhances the performance and realism of protein simulations without requiring force inputs during training.  [] introduced an object-aware SE(3)-equivariant diffusion model for rapidly generating accurate 3D transition state structures, significantly reducing the computational burden typically associated with quantum chemistry calculations.
In astrophysics, diffusion models have been utilized to generate synthetic observations and simulate complex astrophysical phenomena.  [] proposed a diffusion model for generating realistic galaxy images, aiding in the analysis of large-scale sky surveys. Diffusion models have also found applications in climate science and Earth system modeling. For instance,  [] employed a deep generative model to super-resolve spatially correlated multiregional climate data, enhancing the spatial resolution of global climate simulations, which is crucial for long-term climate projections and infrastructure development planning.  [] explored the generative emulation of weather forecast ensembles with diffusion models, illustrating their effectiveness as scalable and cost-efficient alternatives to traditional ensemble forecasts, thus improving the reliability and accuracy of predictions for extreme weather events.
In particle physics,  [] introduced a diffusion model for generating high-quality Liquid Argon Time Projection Chamber (LArTPC) images, showcasing the model’s ability to handle the challenges of sparse but locally dense particles.

While these studies highlight the expanding interest and application of diffusion models across various scientific domains, their potential within fusion research, specifically as an alternative to PIC simulations, remains underexplored. Our work aims to bridge this gap by proposing Diff-PIC, a conditional diffusion model that integrates the capability of PIC for generating high-fidelity synthetic data in fusion research. Diff-PIC leverages the inherent advantages of diffusion models to provide a computationally efficient alternative to traditional PIC simulations (see Table 3).

## 6 Conclusions

This paper presents Diff-PIC, a pioneering approach that leverages the capabilities of diffusion models to generate high-fidelity synthetic data for LPI, offering a computationally efficient alternative to conventional PIC simulations in nuclear fusion research. By integrating a Physically-Informed Parameter Encoder and applying the Rectified Flow Acceleration, Diff-PIC significantly augments the diffusion model’s capacity to manage diverse experimental parameters, thereby expediting the generation of high-fidelity synthetic data.
These advancements facilitate rapid, resource-efficient exploration of the design space, markedly diminishing the computational demands associated with PIC simulations. Our research not only catalyzes accelerated scientific discoveries within the realm of fusion research but also sets a novel precedent for the application of generative AI models in scientific simulations. Future investigations may focus on optimizing the distillation paradigm, harmonizing the simulation time $t_{as}$ with the diffusion time $t$, and refining the condition encoder to encompass a broader spectrum of physical parameters.

## Peer Reviews

confidence: 3
contribution: 4
initial_rating: 8
presentation: 3
questions: The paper introduces innovative components like the PIPE with polynomial encodings, a U-Net for score-based modeling in the RFA. However, it lacks specific details on critical elements: it does not fully describe the polynomial encoder (such as the choice of polynomial degree or basis functions), the architecture and configurations of the U-Net, or the exact loss function guiding the training process. These missing details limit a clear understanding of how each component is optimized for effective high-fidelity data generation.
soundness: 4
strengths: The submission is clear in its objectives and introduces novel elements like the Physically-Informed Parameter Encoder (PIPE) and Rectified Flow Acceleration (RFA) Technique, which together can generate high-fidelity data by applying a diffusion model. The experimental setup is robust, providing a detailed comparison of Diff-PIC against state-of-the-art generative models (GANs and Normalizing Flow) on multiple performance metrics (MAE, RMSE, FID). The authors demonstrate significant speedups and error reductions, substantiating the model’s effectiveness with quantitative evidence. Diff-PIC makes a strong alternative to traditional Particle-in-Cell (PIC) simulations, since it achieves orders-of-magnitude speedup, while retaining high fidelity and reducing the computational cost of PIC.

The PIPE provides a novel way to embed physical constraints, ensuring the generated data remains physically consistent without modifying the core diffusion process. This approach keeps the model flexible, allowing it to adapt to various scenarios, such as handling different particle types (e.g., electrons or ions), while avoiding the complexity of embedding physics-informed parameters directly within the diffusion model itself. While RFA is becoming more known in specific domains like high-resolution image generation, it remains novel in scientific applications, particularly for complex simulations like PIC. The inclusion of RFA in Diff-PIC enhances efficiency, offering significant speedup without sacrificing fidelity for generating realistic. Despite some missing details on specific model components or implementation—the paper’s contributions are substantial and have the potential to drive advancements in scientific data generation for PIC simulations and could become a tool for the broader community.
summary: The paper describes a computationally efficient alternative to Particle-in-Cell (PIC) simulations for nuclear fusion research, specifically for generating Laser-Plasma Interaction (LPI) data. This is achieved through the application of a diffusion model, designed to replace the resource-intensive and time-consuming nature of PIC simulations by introducing a Physically-Informed Parameter Encoder and the Rectified Flow Technique. The work primarily aims to advance nuclear fusion research by making high-fidelity data generation feasible at a fraction of the traditional computational cost.
weaknesses: The authors discuss how diffusion models align with the problem and highlight PIPE as a benefit for Diff-PIC, but GAN-PIC and NF-PIC include PIPE in their comparisons as well. Also, Diff-PIC’s performance heavily depends on PIPE’s ability to generalize from limited training data, which may require retraining or fine-tuning the encoder if the physical parameter ranges shift. While the paper evaluates Diff-PIC using error metrics and continuity checks, it would benefit from additional validation against physics-specific benchmarks and more interpretability studies. Also, architectural comparisons to GAN-PIC and NF-PIC in the LPI data generation context would provide insights into the superior performance of Diff-PIC.

---

confidence: 2
contribution: 2
initial_rating: 5
presentation: 3
questions: Perhaps authors could add an ablation study on various pieces of the physically informed parameter encoder? such as the positional encoding or polynomial transformation

Perhaps authors could add uncertainty estimates to numbers in Tables? such as standard deviation over multiple runs or confidence intervals

Can authors add energy evaluation on Gan-Pic and NF-Pic onto Fig 5? Would be good to have same energy evaluation metrics (MAE and RMSE) included for GAN-PIC and NF-PIC in Fig 5, allowing for a direct comparison across all methods
soundness: 2
strengths: The paper is original in using diffusion model for simulating particle-in-cell data. Paper tackles an important problem of design in nuclear fusion research. Paper is clearly written.
summary: Paper uses a diffusion model to generate data simulating a Particle-in-cell simulation for nuclear fusion research. The diffusion model takes parameters (electron temperature, ion temperature, and laser intensity) and timestep and generates an electric field snapshot. The paper designs a parameter encoder which uses positional encoding for timestep, and positional encoding concatenated with a polynomial transformation for electron temperature, ion temperature, and laser intensity. The model is trained with rectified flow-based acceleration using a u-net for learned trajectory. Paper includes experiments of quantitative interpolation, extrapolation, energy evaluation vs snapshot, qualitative comparison with ground truth data, and speedup vs particle-in-cell simulation.
weaknesses: A weakness is that it is unclear how close the generated data should be to the true data to be useful for nuclear fusion design/research. hard to evaluate the scale of quantitative error. also do the fine details matter in the E field? because they do not look close in Fig 4 visualization. could authors provide context on what level of accuracy is required for the generated data to be useful in nuclear fusion research, and discuss the importance of fine details in the electric field and how this impacts the utility of their approach.

the main conclusion from paper's experiments seems to be that diffusion with u-net performs better than Gan or normalizing flow, for this task.

---

confidence: 2
contribution: 2
initial_rating: 5
presentation: 2
questions: Please see the weakness part.

Compared with previous work, "diffusion models in molecular dynamics simulations "

I hope the author can clarify:
1. The above cited paper is first paper apply diffusion to a similar domain.
2. Among the two unsolved problem, can the author explain more how this paper address "Physical soundness must be ensured"? Does it more empirical or theoriotical?
3. "Substantial efficiency improvement" I got this achieved by apply normalizing flow.
soundness: 2
strengths: The paper try to apply diffusion model for PARTICLE-IN-CELL NUCLEAR FUSION SIMULATION. Although I'm not sure how important of this problem, i think using diffusion model for simulation in science is an important and exciting topic. Also the the speed-up looks impressive.
summary: This paper proposed an new approach to tackling the computational challenges in nuclear fusion simulations by applying diffusion model. The contribution including: (1) Proposing PIPE (Physically-Informed Parameter Encoder) that improves the model's ability to understand the relationship between physical parameters and simulation outcomes. (2) Applied rectified flow acceleration. The simulated results shows this diffusion based approach significantly speed-up compared with traditional method.


========== post rebuttal

I still feel the paper's contribution on ML side is not significant enough. However, I'm less familiar with the application side. I'm lean towards to accept, if other reviewer can champion it on ML for science side.
weaknesses: It's unclear to me what is the ML contribution for this paper. 
To me it seems most interesting part is the newly designed encoder. Can the author explain:
"Algorithmic generalization. PIPE improves the generalizability of the conditional diffusion model"
I don't quite get why such design improved generalizability compared with normal mlp/transformer layers. I hope the author would clarify it.

The application value of this paper beyond my expertise, I hope the author could also explain it and also explain how this approach generate new insights for other ML for science problem.

---

confidence: 4
contribution: 2
initial_rating: 3
presentation: 2
questions: - See 'weaknesses' above: I would like clarification on whether the generative model is learning 2D snapshots of the electric field. What is the ground truth representation? Is the equation (2) at all relevant? 

- The authors mention 'limited' samples. But shouldn't the PIC simulator be able to generate a large number of samples to train the model? How many samples do we think are necessary to qualify as a large enough dataset?

- Did the authors implement the alternative (GAN and NF) methods themselves? I do not see a corresponding reference for these methods for PIC (what is provided is the main method describing GAN and NF).
soundness: 3
strengths: - This appears to be a first of its kind (at least according to the authors) PIC simulation using generative model. I cannot verify this claim, but no other references of this type of work are compared with in the paper. 

- The methodology appears to be sound, and the evaluation follows the techniques used in generative modelling generally. 

- The evaluation fares quite well, compared to competitors (GAN and normalizing flow), achieving better performance and speed.
summary: This paper investigates laser plasma interactions in Particle In Cell simulations (PIC) through diffusion models. They study the evolution of charged particle trajectories, influenced by the electric and magnetic fields (equation 2). The neural network designed uses the diffusion model methodology, with some novel fittings - a physically designed parameter encoder, positional encodings and legendre polynomials, and a distillation framework. The network - from what I see - takes as input electron and ion temperatures $T_e$, $T_i$ and laser intensity $I$ (i.e. $\theta = (T_e, T_i, I)$ to generate the electric field $E$, evolving it over time $t_{as}$ to obtain $E(\theta, t_{as})$. The network optimization takes place through the rectified flow acceleration technique proposed in [1, 2, 3] in equation (5). 

Results are presented to show time evolution snapshots of the electric field, comparing them with two other generative models (GAN and normalizing flow). Metrics checked are MAE, FID and RMSE. They also examine the 'energy' evolution over time. They call this the interpolation and extrapolation capabilities of the generative model. Data distribution is also compared with the 'ground truth' which is obtained from a PIC simulator. The claims are that the model beats the competitors (GAN [4] and normalizing flow [5] ), and is much faster owing to the more efficient rectified flow acceleration technique used. 

Post rebuttal
=========
The authors have taken great pains to explain the method and provide additional ablations and experiments, all of which supports their case that the method works rather well for PIC. The paper paper adds to the body of research work in PIC simulations:
- it works as well or better than competitive methods 
- it is significantly faster than mainstream computational approaches used in plasma physics (owing to the rectified flow acceleration technique)
However, as stated before, the work does not add to novelty in terms of the method used in generative modelling and diffusion models. To this end, I am raising the score but am not (entirely) convinced about it being accepted. That being said, I think it is a marginal call.

[1] Xingchao Liu, Chengyue Gong, and Qiang Liu. Flow straight and fast: Learning to generate and
transfer data with rectified flow, 2022.

[2] Patrick Esser, Sumith Kulal, Andreas Blattmann, Rahim Entezari, Jonas Muller, Harry Saini, Yam ¨
Levi, Dominik Lorenz, Axel Sauer, Frederic Boesel, et al. Scaling rectified flow transformers for
high-resolution image synthesis. arXiv preprint arXiv:2403.03206, 2024.

[3] Xingchao Liu, Xiwen Zhang, Jianzhu Ma, Jian Peng, and Qiang Liu. Instaflow: One step is enough
for high-quality diffusion-based text-to-image generation, 2024.

[4] GAN: Alankrita Aggarwal, Mamta Mittal, and Gopi Battineni. Generative adversarial network: An overview
of theory and applications. International Journal of Information Management Data Insights, 1(1):
100004, 2021.

[5] NF: Qinsheng Zhang and Yongxin Chen. Diffusion normalizing flow. Advances in neural information
processing systems, 34:16280–16291, 2021.
weaknesses: - Novelty: The main novelty appears to be the application of an existing method in a new domain (plasma physics). While this is not a complaint in itself, it seems (at least to me) that it is only a replacement of the modalities from more well studied setups involving images. 

- Method and physical validity: It is hard for me to understand how the equation (2) (velocity field evolution) is being solved by the diffusion model. From what I see, we are learning the electric field from snapshots in time. 

- Ablations: The components used in the model should, I think be ablated to show effectiveness of each piece. For instance, how does the positional embedding help? What is the performance without the Chebychev polynomial addition? 

- Network Architecture: Perhaps I am missing this, but I would like to see a block containing the network design - transformer blocks, CNNs, etc.

---

confidence: 4
contribution: 3
initial_rating: 8
presentation: 4
questions: * Can you elaborate on the time encoding? At one point in the paper, it was written that the method works from snapshots. Thus, long-time series are not necessary. However, later in the paper, the authors mention that they use encoding to capture information about time. Is there some dynamic information necessary for the model function? If so, why not use a more recurrent kind of architecture? I can imagine, if dynamics are required for a complete description, that a single snapshot can occur many times in different trajectories. Therefore, a time-only encoding would not capture this information, and a state-based memory would be required. I am not an expert in these kinds of simulations, so I am likely misunderstanding something but more clarity there would be great.
soundness: 3
strengths: I find the paper to be very well-written and clear. Further, the topic being approached is very important and introduced in such a way. The experiments performed to validate their method and present prior work to highlight the difference in their approach.
summary: This paper introduces an approach for accelerating particle in cell simulations for laser in plasma simulations by applying a diffusion model to the computation of electric fields. These fields are produced based on key parameters in the simulation and show good accuracy when compared to literature and non-accelerated simulations.
weaknesses: The paper is undoubtedly more geared towards nuclear physicists than for people in the ML community. What was missing for me was the advance in the ML direction. What can be taken from this work and applied to improve knowledge in the ML community? By making this connection, I feel it will be better received at ICLR.

Meta-review: The paper introduces Diff-PIC, a novel framework that leverages diffusion models to generate high-fidelity simulation data for Laser-Plasma Interaction (LPI) in nuclear fusion. Diff-PIC addresses the computational bottleneck of traditional Particle-in-Cell (PIC) simulations by offering a 16,200× speedup while maintaining high accuracy. The paper contributes two key things: 1) a Physically-Informed Parameter Encoder (PIPE), which captures relationships between physical parameters and PIC data, enabling smooth interpolation and extrapolation of unseen parameter values; and 2) a Rectified Flow Acceleration (RFA) technique that allows for one-step data generation, further enhancing efficiency. The authors also provide an open-source dataset of 6,615 LPI simulations, positioning Diff-PIC as a new baseline for scientific data generation.

In terms of the strengths of the approach, as highlighted by the reviewers as well, include very good computational efficiency, strong generalisation to new parameter ranges, and a reduction in computational energy usage by 10,000×. Diff-PIC outperforms state-of-the-art methods like GAN-PIC and NF-PIC in terms of mean absolute error (MAE), root mean squared error (RMSE), and Fréchet Inception Distance (FID), showing some strong physical validity in its generated data.

However a minor weakness concerns the framework’s generalisation which is only tested on limited setups, therefore restricting its wider impact on machine learning methodologies. 

Overall, Diff-PIC represents a significant step forward for ML applications in nuclear fusion research, providing an efficient and accurate alternative to traditional PIC simulations, with the potential to accelerate research in fusion and other areas requiring computationally intensive simulations.

Decision: Accept (Poster)