Diff-PIC: Revolutionizing Particle-In-Cell Nuclear Fusion Simulation with Diffusion Models
Authors: Chuan Liu
Chunshu Wu
shihui cao
Mingkai Chen
James Chenhao Liang
Ang Li
Michael Huang
Chuang Ren
Ying Nian Wu
Dongfang Liu
Tong Geng
Venue: ICLR 2025 Poster (2025)
Keywords: Diffusion Models
PIC Simulations
Scientific Data Generation
Primary area: applications to physical sciences (physics, chemistry, biology, etc.)
TL;DR: We propose Diff-PIC, leveraging diffusion models to generate high-fidelity synthetic data for nuclear fusion research.

Abstract: The rapid development of AI highlights the pressing need for sustainable energy, a critical global challenge for decades. Nuclear fusion, generally seen as a promising 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 (ICF) 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 makes analytical approaches impractical, leaving researchers dependent on extremely computationally intensive Particle-in-Cell (PIC) simulations to generate data, posing a significant bottleneck to the advancement of 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 their corresponding outcomes, the parameters are encoded in a physically informed manner. (2) To further enhance efficiency while maintaining 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 a $\sim$16,200$\times$ speedup compared to traditional PIC on a 100 picosecond simulation, while delivering superior accuracy compared to other data generation approaches.

Published as a conference paper at ICLR 2025
DIFF-PIC: REVOLUTIONIZING PARTICLE-IN-CELI NU- CLEAR FUSION SIMULATION WITH DIFFUSION MODELS
Chuan Liu', Chunshu
Cao', Mingkai Chen', James Chenhao Liang'.
Ang Li², Michael Huang', Chuang Ren', Ying Nian Wu', Dongfang Liu', Tony (Tong) Geng! 'University of Rochester, 'Rochester Institute of Technology. 'Pacific Northwest National Laboratory, *University of California, Los Angeles
ABSTRACT
The rapid development of AI highlights the pressing need for sustainable energy, a critical global challenge for decades. Nuclear fusion, generally seen as promising solution, has been the focus of intensive research for nearly century, with invest- ments reaching hundreds of billions of dollars. Recent advancements in Inertial Confinement Fusion (ICF) 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 makes analytical approaches im- practical, leaving researchers dependent on extremely computationally intensive Particle-in-Cell (PIC) simulations to generate data, posing a significant bottleneck to the advancement of fusion research. In response, this work introduces Diff-PIC, novel framework that leverages conditional diffusion models as a computationally efficient alternative to PIC simulations for generating high-fiideltty scientific LPI data. In this work, physical patterns captured by PIC simulations are distilled into diffusion models associated with two tailored enhancements: tively capture the complex relationships between physical parameters and their corresponding outcomes, the parameters are encoded in physically informed manner. (2) To further enhance efficiency while maintaining physical validity, the rectified flow technique is employed to transform our model into one-step conditional diffusion model. Experimental results show that Diff-PIC achieves a ~16,200x speedup compared to traditional PIC on a 100 picosecond simulation, while delivering superior accuracy compared to other data generation approaches.
(1) To effec-
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 (Abu-Shawareb) ct al. 2024) underscores the 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 from 2.2 MJ input energy, respectively. Given that the estimated output could reach 120 MJ (Suter et al.) 2004), there a growing demand for deeper understanding of the fundamental science behind ignition, especially the physics governing the interaction between the laser and the plasma emitted when the laser bombards the fuel pellet. However, LPI is complex problem that is traditionally analyzed using time- sime-stepping method, particularly through PIC simulations (Tskhakaya et al., [2077] Langdon 2014 Arber et al., 2015). 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 to obtain meaningful outputs. The computational overhead of PIC simulations has become 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 based accelerations have been developed for PIC simula- tions (Bowers et al. 2008 Sgattoni et al. 2015). Although invaluable, these efforts remain within the scope of the iimestteppig approach that iteratively executes over infinitesimal time intervals, falling short in addressing the inherent computational overhead of long-term simulations. Recent

---

Published as a conference paper at ICLR 2025
advancements in generative AI diffusion models, however, present novel approach to bypass the constraint. Diffusion models have demonstrated exceptional capabilities Computer Vision (CV), synthesizing highly complex data distributions that match real data with high fidelity (Ho et al. 2020, [Song & Ermon, 2019). Rooted in the diffusion concept in physics, diffusion models be understood as constructing highly complex field that governs the evolution of variables, analogous to the motion of particles in cells. This has sparked significant interest in their potential for generating scientific data, as recent applications of diffusion models in molecular dynamics simulations demonstrated (Wu & Li 2023 Petersen et al. 2023
Although diffusion models exhibit promising potential for generating PIC simulation data, two critical research gaps remain. O Physical soundness: It remains unclear how to effectively integrate continuous physical parameters into diffusion models to generate data that faithfully represent the underlying physics. @ Substantial efficiency improvement: Although the requirement for infinitesimal time intervals PIC simulations has been relaxed, the step-by-step denoising process diffusion models is still computationally demanding, limiting the practicality of diffusion models as advanced alternatives to PIC simulations.
To address these challenges, propose Diff-PIC, which leverages diffusion models to efficiently generate snapshot of arbitrary time, under arbi- trary physical parameters within certain ranges. Specifically, ! develop conditional diffu- sion model with Physically-Informee Parameter Encoder. This encoder enables the model cap- ture the relationship between continuous physical parameters and PIC simulation data, effectively distilling the underlying physics into Diff-PIC. We employ the rectified flow technique (Liu et al. 2022) to eliminate the requirement for multi- step denoising, further optimizing the runtime effi- ciency of PICI-PIC. As highlighted in Fig. Diff- PIC achieves onders-of-magnitee pared to PIC simulations, and delivering superior fidelity compared baselines "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. The core contributions of our work include:
Particle-In- In-Clll
simulation
Cost Time & Energy
10000x
NF-PIC
GAN-PIC
Diff-PIC
Fidelity
Figure 1: Overview of this work.
speedup com-
We propose Diff-PIC, a pioneering study that utilizes diffusion models as a computationally efficient alternative to PIC simulations. accelerating the scientific data generation for nuclear fusion research.
We develop poysically-informer conditional diffusion model that seamlessly integrates physical parameters into the diffusion model. The designed condition encoder effectively facilitates the generalization of simulation data within and beyond seen parameters, endowing the model with robust generalization capabilities.
We implement the rectified flow technique to make our model one-step diffusion model, further enhancing its efficiency generating high-fidelity fusion data.
Experimental results demonstrate that our method achieves a remarkable speedup of ~ 16,200 times compared to traditional PIC simulations while preserving physical validity of the generated data.
BACKGROUND
Inertial Continement Fusion (ICF) method of achieving controlled nuclear fusion by using intense energy pulses to heat small fuel pellets (Keefe, 1982. Betti & Hurricane, 2016). typically containing isotopes of hydrogen such as deuterium and tritium. This process unfolds the nuclear fusion reaction as delineated below:
2 H 3H + 'He + neeuroo + Energy.
Given the ubiquity of these hydrogen isotopes, nuclear fusion holds immense potential to provide "near-infinite" energy by achieving the necessary temperature and pressure conditions to initiate fusion
2

---

Published as a conference paper at ICLR 2025
reactions, attaining positive energy gain (i.e., output energy surpassing the input). To optimize the sophisticated initiation of fusion, an advanced understanding of the underlying LPI is essential. For this purpose, PIC simulation is considered crucial tool for providing theoretical insights into LPI, due to its capability to predict and interpret physical phenomena.
widely used in the study of plasma physics and fusion energy re- search (Tikhakaya et al. 2007 Derouillat et al., 2018). Developed in the mid-20th century, the PIC method has become cornerstone in the simulation of complex plasma behaviors, enabling researchers to delve into the underlying physics. To highlight, PIC is especially useful in LPI stud- ies (Arber al. 2013: Klimo et al., 2010), which involve 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 for 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. To ensure accuracy, simulating LPI over only few hundred picoseconds (i.e., 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 stringent bottleneck in fusion research, significantly restricting progress in this domain.
PIC Simulations
Diffusion models have emerged powerful generative models, capable of generating high-fidelity data (Ho et al. 2020; Song & Ermon 2019). Diffusion models involve two processes: the forward process and the reverse process. In the forward process, Gaussian noise is gradually added clean data, transforming the original data distribution into a tractable Gaussian distribution. Mathematically, this can be represented as:
((-))=)1) = (xx;; va (1 -aa)))
(2)
where ay regulates the noise level at each time step f. The reverse process reconstructs the original data from the noisy distribution by iteratively denoising the data. This procedure is typically parameterized by a neural network, which trained to approximate the inverse of the forward diffusion process:
po x (x.) -> uo l), Ex(x,, ()),
where 140 represent the mean predicted by the neural network and Eg denotes time-dependent parameters that control the variance or uncertainty at each denoising step.
3 DIFF-PIC
This section introduces Diff-PIC, physically-nnfomen conditional diffusion model tailored for generating high-fidelity synthetic data for LPI nuclear fusion. As illustrated in Fig. [ accepting physical parameters as inputs, the parameters are encoded and integrated with the model through the physically informed parameter encoder. Additionally, the rectified flow method is adopted to break free from the multi-step denoising process in diffusion models, unlocking the model's capability to generate scientific data in one single step.
3.1 THE OVERALL DISTILLATION FRAMEWORK
Since LPI is governed by the behavior of electromagnetic fields that dominate plasma dynamics, as a representative case, the task assigned to Diff-PIC is to generate high-fidelity 2D electric fields under various physical parameters. which is also a standard practice in PIC simulations. Specifically, we focus on the following critical physical parameters in LPI: "Electron Temperature (T.) in keV," "Ion Temperature (T₃) in keV," and "Laser Intensity (T) in W/m?." Additionally, to enable direct generation of the electric field at specific simulation time (a snapshot), we include an additional parameter tas that represents the simulation time for an arbitrary snapshot. Through training, the model captures the relationships between the four parameters and the resulting electric fields. Once trained, the model takes the four parameters as inputs and produces electric field snapshots E(tae. corresponding to the specified physical parameters 0 = {Te, T,, 1}. These generated snapshots can be used for multiple purposes, including data augmentation and parameter exploration. Essentially, the proposed distillation framework offers two advantages:
3

---

Published as a conference paper at ICLR 2025
Diff-PIC Simulation
Hf Simulation Parameters
Physically Informed Parameter Encoder
Rectifid
Acceleration
83.3
One-step Generation
Simulation Results
Distiliation Paradigm
Step-by- Step PIC Simulation
Hi Simulation Parameters
Stop
Step
Step
3
Step
Simulation Results
Figure 2: Workflow of the proposed Diff-PIC compared to traditional step-by-step PIC simulations.
Systemic efficiency. Unlike traditional PIC simulations that generate data sequentially over time, the proposed diffusion model can directly produce data for any target snapshots by varying tas- This non non-sequential behavior allows for substantially more efficient data generation and analysis (see Table 3), enabling researchers focus specific times of interest without needing simulate the entire LPI process from the beginning.
Data dependency relaxation. We treat snapshots from various parameter combinations as distinct distributions. This approach decouples data dependency model training, enabling the model to efficiently learn from individual snapshots while generalizing across wide array of parameters.
3.2 PHYSICALLY-INFORMER PARAMETER ENCODER
PIC simulations employ continuous physical parameters as inputs, which necessitates seamless and continuous transition in the resulting synthesized data as input parameters are adjusted. Consequently, an encoder is considered exceptionally important in this scenario, responsible for transforming domain-specific inputs into effective embeddings. In particular, these inputs comprise the simulation parameters and the target simulation snapshot tas- Furthermore, the encoder should also excel in both interpolation and extrapolation - critical measures of the model's ability of generalization. Interpolation capability refers to the encoder's proficiency generating suitable embeddings for new parameters that, although not encountered during training. lie within the range of parameters seen during training. Extrapolation capability, conversely, pertains to generating embeddings for parameters that fall outside the range of those seen during training. Notably, both capabilities are indispensable for the LPI problem, in order to explore a large enough parameter space to provide sufficient insight into further LPI evaluations.
To meet the essentials, we introduce Parameter Encoder (PIPE) as shown in Fig. To encode the simulation parameters 0, we employ two distinct types of en- coders tailored for interpolation and extrapolation tasks. For interpolation, we employ Positional Encoding (Vaswani et al. 2017) (denoted by "~~" in circle), which leverages sinusoidal functions to encode the input parameters, facilitating smooth transitions between observed parameters. To enhance extrap- olation capability, we incorporate polynomial encoder ("P" in circle). Polynomial encoders are widely used approxi- mate a wide range of functions effectively, capturing nonlinear correlations and unbounded growth patterns in the data, which is crucial for extrapolation. This is achieved through transfor- mation constructed as a linear combination polynomial basis functions f((0) of varying degrees:
Poysically-Informmn
Physieally-lafoemed Informed Parsmeter Encoder
Posinumall Enoowke Pelynomial Feccder Cuneaterate
Figure 3: The proposed PIPE
P(O) Σ f((0).
where denotes the maximum order of polynomial terms, and the polynomial P(0) 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 beyond those encountered during training, ensuring robust performance across a broader spectrum of simulation scenarios. Subsequently, concatenate ("C" in circle) the embeddings from
4

---

Published as a conference paper at ICLR 2025
these two encoders and apply 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 time step lass we utilize Positional Encoding (Vaswani et al. (017) followed by an MLP layer. This approach is specifically chosen 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:
informative representation of the input parameters. For encoding the simulation
Algorithmic generalization. PIPE improves the generalizability of the conditional diffusion model (see Table and 2). The dual-encoding strategy captures nonlinear relationships by incorpo- rating both positional and polynomial encoders, empowering the model to adeptly manage a diverse array of parameters, ranging from scenarios encountered during the training phase to parameters that lie beyond the spectrum of the training data.
Adaptive transferabiliity. By fine-tuning the pre-trained encoder with new dataset, it facilitates adaptation to other related tasks.
3.3 RECTIFEED FLOW ACCELERATION
To enable rapid generation of nigh-fidelity synthetic data, we employ the Rectified Flow Acceleration (RFA) technique. Based on the principles of rectified flow (Liu al. [202]] 2024). RFA converts the original complex denoising trajectory from initial noise to the target electric field snapshot E( 0) into a direct and straight path. During training, Diff-PIC minimizes the following objective function:
[/ IIEE(tas, 0) -) <(E,, fass 1)
arg min
E
where E, = /E(tas, 0) (1 f)e denotes the linear interpolation between E(tas: 0) and across the diffusion timeline, with ranging from 0 to 1. Using a modified U-Net backbone, ( learns the denoising trajectory by minimizing the expectation of the squared deviations between E(las. 0)
and ((E,, ta.. 0). Once < is trained, we further straighten the learned trajectories through an interactive reflow procedure (Liu et al. 2024). In summary, the RFA module provides additional benefits for our method:
Streamlined denoising process. RFA significantly accelerates the denoising process (see Table by converting the complex multi-step denoising process into one-step denoising process.
Robust optimization. By streamlining the denoising process, RFA reduces the possibility of error accumulation that occurs with the multi-step denoising trajectory.
EVALUATION
4.1 EXPERIMENTAL SETUP
Datasets. We construct a dataset comprising 6,615 simulations across varied physical parameters. each containing 80 snapshots of electric fields along two orthogonal directions denoted as E1 and E2. The data are generated using OSIRIS (Fonseca et al. 2002). well-established PIC simulation software suite. The dataset encompasses diverse physical parameters, including Te. Tú, and I, all of which are critical influencing the resultant electric fields.
Metrics. To validate physical soundness, Mean Absolute Error (MAF) and Root Mean Squared Error (RMSE) are used to evaluate the electric field difference and energy difference between the Diff-PIC generated data and the ground truth produced by PIC simulations. To further evaluate the difference in the generated and the ground truth data distributions, the Fréchet Inception Distance (FID) metric is also employed.
Baselines. We compare Diff-PIC with two other types of generative models (arras et al. 2020 Zhang & Chen 2021). 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. More details provided in Appendix A.2
5

---

Published as a conference paper at ICLR 2025
Visualization -
50
E1 Distribution
PPC DIs-II Sindatone
Simibrions
DIFPIC
Frequency
20CK
Frequency
60k
60k
00
0.5
-0.5 Value
0.5
Value
Visunlization t, 50
E2 Data Distribution
PIC DREPIC Sinskakkoo
PIC Simabiiions
Diff PIC
Feqqeency
40k
25k
25
20k
20k
0.0 -3.5
0.3
0.5 0.0 0.5
Value
Value
Figure 4: Visualization and Comparison of PIC simulations and Diff-PIC.
Table 1: Quantitative results for interpolation evaluation. MAE and RMSE are on the order of 10-2.
Training for E1
Testing Set for E1
Training Set for E2
Testing Set for
Method
MAE 4 RMSE,
FID
MAEІ
RMSEJ
FID
MAE,
RMSEJ
FID
MAEJ
RMSE,
FID
GAN-PIC
4.59
5.31
2.32
4.73
5.84
2.51
1.82
2.07
0.973
1.97
2.18
1.03
NF-PIC
4.46
5.12
2.06
4.61
5.35
2.42
1.70
2.03
0.914
1.86
2.45
0.986
Diff-PIC
1.56
2.67
1.21
1.68
2.29
1.62
0.795
0.932
0.328
0.826
1.03
0.341
Diff-PIC Configurations. The architectural foundation of our model is a modified U-Net frame- work (Ronneberger et al.] (2015), comprising three down-sampling blocks and three up-sampling blocks, strategically integrated with attention mechanisms to capture both local and global dependen- cies. For parameter encoding, the positional encoders generate 16-dimensional embeddings using sinusoidal functions, facilitating smooth interpolation of input parameters. Polynomial encoders incorporate polynomial terms up to the fourth degree capture nonlinear relationships, resulting in comprehensive 20-dimensional embedding for each parameter after concatenation and transfor- mation via single-layer MLP. These refined embeddings concatenated with the simulation data and fed into the U-Net backbone. More implementation details are provided in Appendix A.I
4.2 MAIN RESULTS
In this section, 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 energy efficiency compared traditional PIC simulations.
Interpolation and Extrapolation. To evaluate the interpolation capability of Diff-PIC, we sample a specified range for each simulation parameter (T., T, and 2), totaling 500 simulations and 500 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 | To better illustrate the relative error, the metrics are evaluated after normalizing the generated data the range of [0,1]. The reasonably low MAE, RMSE, and FID scores indicate that the proposed Diff-PIC is able to synthesize high- scientific data similar to what PIC generates, while significantly outperforming baselines on all three metrics. On average. Diff-PIC achieves 59.25% reduction in MAE, 57.77% reduction in RMSE, and 49.21% 2% reduction in FID across all testing sets compared to the baselines.
In addition, Fig. 4 compares the results of one randomly selected simulation produced by DICI-PIC and PIC simulations in the testing set, respectively. Throughout the snapshots (c.g., fas = 1 - 30 + 50 - 80), the visualizations of the synthetic data closely follow the ground truths, indicating that the physical evolution is preserved over time. Additionally, the distributions on the right demonstrate that proposed model successfully captures the data distributions of the ground truth.

---

Published as a conference paper at ICLR 2025
Table Quantitative results for extrapolation. MAE and RMSE are on the order of 10-2,
EI
10%
EI 20
E2 10 %
F.2 %
Method
MAE]
RMSEL
FIDI
MAEL
RMSEL
FID
MAEL
RMSEL
FID
MAEL
RMSEL
FID
GAN-PIC
5.24
6.32
1.97
5.73
6.98
2.15
2.18
3.25
1.15
2.93
3.89
1.29
NF-PIC
4.74
5.31
1.85
5.41
6.46
2.08
1.90
3.16
1.04
2.63
3.42
1.17
Diff-PIC
1.83
2.40
1.74
2.18
2.62
1.82
0.947
1.36
0.536
1.13
1.85
0.673
10
wwwwmm
0.5
0.5
0.5
0.5
10-40
25
50 Snapshots
25
50
25
50
25
50
75
Snapshots
Snapshots
Snapshots
Figure 5: Energy evaluation of electric fields for training and test sets with 80 snapshots. From left to right: E1 Training Set, E1 Testing Set, E2 Training Set, and E2 Testing Set.
In terms of extrapolation capability evaluation, we progressively extend the range of physical parameters selected for training. In particular, the ranges are extended by 10% and 20% respectively case studies. The results in Table 2 demonstrate that Diff-PIC achieves approximately 2% relative absolute error in extrapolation tasks, significantly outperforming other generative counterparts. On average, Diff-PIC achieves 59.16% reduction in MAE, 56.53% reduction in RMSE, and 29.70% reduction in FID across all test cases compared to the baselines.
Physical Validity. In addition to ensuring the high quality of the synthetic data, assess the physical validity of the generated electric fields by evaluating their energy evolution-a fundamental property that characterizes physical systems. In particular, training and testing sets for E1 and E2. The energy profiles of these electric fields are depicted in Fig. 5 where we compare the ground-truth PIC simulations with synthetic data generated by Diff-PIC. GAN-PIC, and NF-PIC. Notably, Diff-PIC closely aligns with the ground truth, demonstrating low errors across all datasets. Specifically, for the E2 electric field, Diff-PIC accurately preserves the energy oscillations in the ground truth, outperforming baselines. This evaluation highlights that Diff-PIC capable of generating sequentially continuous data, demonstrating the effectiveness proposed distillation paradigm and the Poysically-Informee Parameter Encoder.
randomly select simulations from both the
Speedup and Energy Consumption Reduction. In addition to traditional PIC simulations, baselines also included in the comparison shown in Table 3 The PIC simulations executed on the Perlmutter supercomputer at the National Energy Research Scientific Computing (NERSC) facility, utilizing AMD EPYC 7763 CPUs. Since essential fusion phenomena typically appear approxi- mately 100 ps, the cost of PIC simulation at 100 ps is selected as the reference for comparisons. The GPU results for Diff-PIC and the baselines are obtained an Nvidia RTX 4090 GPU to demonstrate the accessibility of this approach to general users. The CPU results for these approaches are acquired on Intel 13th Gen i9- 13900KF CPU. The results highlight that Diff-PIC-GPU achieves speedup of ~ 16, 200x compared to traditional PIC simulations, along with a similar orders-on diderr-off-magiuuee magnitude reduction energy consumption.
Ablation Studies on PIPE. To further demonstrate the effectiveness of PIPE, we conduct ablation studies by replacing the PIPE component in Diff-PIC with other commonly used encoders: MLP and Transformer decoder. Detailed implementations and comparison results are presented in Ap- pendix A.3 The results show that PIPE outperforms both MLP and Transformer decoder both
Table 3: Speedup comparison across different methods.
Method
1PIC-100pt
DIP-PIC-GPU DiY-PIC-CPU GAN GPU GIC-PIC NPU NF CPU
Speedup
1.00x
1.624xx 1.0144x
$19x
1.56elx 1.1364x
$23x 1.47e3x
9,21e2x 8.14e2x
24x
Energy Reduction
1.05e3x

---

Published as a conference paper at ICLR 2025
interpolation and extrapolation tasks. MLP struggles with both interpolation and extrapolation. The Transformer decoder, while powerful for sequence modeling and capturing relationships in discrete token sequences, is not inherently designed to effectively process continuous physical parameters, resulting in inferior performance.
In addition, to provide more insight into the effectiveness of PIPE, we conduct another ablation study to evaluate the contributions of the positional encoder and polynomial encoder separately. Detailed implementations and comparison results are shown in Appendix AA.3] highlighting the individual importance of these encoders. Specifically, combining both positional and polynomial encoders consistently outperforms using either encoder alone in both interpolation and extrapolation tasks. In interpolation tasks, the positional encoder demonstrates better performance than the polynomial encoder, suggesting its important role in interpolation scenarios. Conversely, extrapolation tasks, particularly at larger ranges (20%), the polynomial encoder shows relatively better performance, indicating its importance for extrapolation capabilities. These results validate our design choice of combining both encoders in PIPE to leverage their complementary strengths.
4.3 DISCUSSION
To highlight the value of this work, it is worth noting that our approach exhibits outstanding scalability compared traditional PIC simulations. In PIC, for N particles, the computing complexity can generally reach a formidable T'Nlog.V with I' time steps. In contrast, Diff-PIC is not sensitive to the number of particles in space or the number of time steps. For larger particle spaces, the speedup achieved by Diff-PIC: can be easily increased by extra orders of magnitude, further accelerating the research of fusion, or other research areas involving large-scale PIC simulations.
Furthermore, to provide insight into the quality of the synthetic data generated by Diff-PIC and its applicability for ICF research, we reference a SOTA ICF modeling method (Ejaz et al. 2024), which suggests prediction error of approximately 12% is considered effective. In experiments, Diff-PIC achieves MAE and RMSE values of approximately 1-2% relative to the PIC simulations. This high level of precision is sufficient for the generated data to be useful in ICF research, such as preliminary analysis, parameter exploration, and prediction modeling.
Additionally, Diff-PIC shows new insights for machine learning methods in science by using diffusion models to address common challenges in complex scientific simulations. Our work demonstrates the power and effectiveness of Diff-PIC a computationally efficient alternative expensive ICF simulations. This achievement not only validates the effectiveness of diffusion models in enhancing simulation efficiency but also paves the way for diffusion models augment or potentially replace a wide range of computationally expensive scientific simulations, enabling more efficient and scalable simulation methods.
RELATED WORK
Particle-in-Cal Cell Simulations have long been fundamental to modeling physical processes in fusion research Taccognaa & Minelli 2017 Garrigues et al., 2016). However, the computational intensity of
PIC simulations presents significant challenges ( /Verboncoeur 2005). To mitigate these computational constraints, various methods have explored GPU and hy brid CPU-GPU acceleration technologies.
Studies such as (Abreu et al.) 2010 Burau et al.) 2010 Decyk & Singh 201 Kong et al., 2011 Suzuki al. 2011) have utilized parallel computing. high memory bandwidth, and multiple processors to accelerate simulations. Architecturally, the simulator optimized for the Kepler GPU architecture, as discussed in (Shah al., 2017), underscores the potential of specific GPU architectures to enhance simulation efficiency. For more intricate simulations, research efforts like (Xu et al., 2012 Chen al. 2012) have developed hybrid CPU-GPU implementations, and (Wang & Zhu (021) introduced a hybrid approach for multi-core and multi-GPU systems, highlighting the continuous integration and evolution of these technologies
accelerating 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 time-stepping property 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 solutions (Kube al. 2021), and learning the probability of interactions
8

---

Published as a conference paper at ICLR 2025
between potential collision pairs (Bilbao et al. 2022) However, the approach by (Kube al. 2021 depends on a pre-computed vector space and may not generalize well to novel scenarios, while (Bilbao et al. 2022)* 's method focuses on binary interactions, overlooking the complex many- interactions in PIC simulations.
Contrastingly, our proposed method overcomes these limitations by employing a conditional diffusion model to distill the complex physics captured by PIC simulations. By leveraging diffusion models, our approach can efficiently generate high-fidelity synthetic data (see Fig. 4 and Table [ without the computational expense of traditional PIC algorithms (see Table B 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-tune tenettned with minimal additional data to suit various physical parameters. This flexibility not only renders our approach suitable for a wide range of applications in fusion research but also extends its potential use to other related domains where efficient and accurate simulations are crucial for understanding complex physics.
Diffusion Models in Scientific Research have emerged as promising tool across myriad of scientific domains. In materials science and chemistry, diffusion models have been applied to enhance molecular dynamics simulations. For example, (Wu & Li, (023) introduced diffusion model for molecular dynamics simulations, showing good results in generating molecular trajectories. (Arts) et al. 2023) presented an approach that integrates diffusion models with coarse-grained molecular dynamics to develop new force field to simulate protein dynamics. (Duan et al. 2023) introduced an object-aware diffusion model capable of rapidly generating accurate 3D transition state structures, significantly reducing the computational burden typically associated with quantum chemistry calcu- lations. In astrophysics, diffusion models have been employed to generate realistic galaxy images. (mith et al. 2022) demonstrated that diffusion models could produce highly detailed and accurate representations of galaxies, which are essential for the analysis of large-scale sky surveys and for enhancing the understanding of the universe. Diffusion models have also found applications in climate science and earth system modeling. For instance, (Oyama et al. 2023) employed a deep generative model to enhance the spatial resolution of global climate simulations, which is crucial for long-term climate projections and infrastructure development planning. (Li et al. 2024) explored the generative emulation of weather forecast ensembles with diffusion models, illustrating their effectiivneess 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, (Imani et al. 2024) introduced a diffusion model for generating Sigh-quaiity Liquid Argon Time Projection Chamber (LArTPC) images, showcasing the model's ability 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, the potential of diffusion models within fusion research, specifically as an alternative to PIC simulations, remains largely underexplored. Our work aims to bridge this gap by proposing Diff-PIC, a conditional diffusion model that generates high-fidelity synthetic data for fusion research. Diff-PIC offers a computationally efficient alternative to traditional PIC methods, paving the way for more scalable and resource-efficient scientific simulations.
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, thereby offering a computationally efficient alternative to conventional PIC simulations in nuclear fusion research. By integrating a Maysscally-IInoomee Parameter Encoder and applying the Rectified Flow Acceleration, Diff-PIC significantly augments the diffusion model's capacity to manage physical parameters, thereby accelerating the generation of high-fidelity synthetic data. These advancements facilitate rapid, resource-efficient exploration of the design space, markedly reducing the computational overhead associated with PIC simulations. Our research not only catalyzes accelerated scientific discoveries within fusion research but also sets a novel precedent for the application of generative AI models in scientific simulations. Future investigations may focus further optimizing the distillation paradigm, harmonizing the simulation time tas with the diffusion time t, and refining the condition encoder to incorporate a wider range of physical parameters.
9

---

Published as a conference paper at ICLR 2025
ACKNOWLEDGMENTS
This work is supported by NSF under Award SHF-2326494, DMS-2015577, and RMS-2415226.
REFERENCES
Paulo Abreu, Ricardo A Fonseca, João M Pereira, and Luís O Silva. Pic codes new processors: full relativistic pic code in cuda-enabled hardware with direct visualization. IEEE Transactions Plasma Science, 39(2):675-60) 685, 2010.
H Jbu-Shawareb, R Acree, P Adams, J Adams, B Addis, R Aden, P Adrian, BB Afeyan, M Aggleton, L Aghaian, et al. Achievement of target gain larger than unity in an inertial fusion experiment. Physical Review Letters, 132(6):065102, 2024.
TD Arber, Keith Bennett, CS Brady, A Hewrence-Douglee MG Ramsay, Nathan John Sircombe, Paddy Gillies, RG Evans, Holger Schmitz, AR Bell, et al. Contemporary particle-in-cell approach to laser-plasma modelling. Plasma Physics and Controlled Fusion, 57(1 (1:1):111 113001,, 2015.
Marloes Arts, Victor Garcia Satorras, Chin-Wei Huang, Daniel Zuegner, Marco Federici, Cecilia Clementi, Frank Noé, Robert Pinsler, and Rianne van den Berg. Two for one: Diffusion models and force fields for yoarse-grained molecular dynamics, 2023.
Betti and OA Hurricane. Inertial-confinement fusion with lasers. Nature Physics, 12(5):435-44 448, 2016.
PJ Bilbao, C Badiali, and F Cruz. Towards deep learning accelerated -nntille-in-cell simulations: application to compton scattering. OSIRIS, 10:3, 2022.
Kevin J Bowers, Brian James Albright, Lilan Yin, B Bergen, and Thomas JT Kwan. Ultrahigh performance dimee-dimnnii electromagnetic relativistic kinetic plasma simulation. Physics of Plasmas, 15(5), 2008
Heiko Burau, Renée Widera, Wolfgang Hönig, Guido Juckeland, Alexander Debus, Thomas Kluge, Ulrich Schramm, Tomas E Cowan, Roland Sauerbrey, and Michael Bussmann. Picongpu: a fully relativistic particle-in indlee-i--elll code for gpu cluster. IEEE Transactions on Plasma Science, 38(10): 2831-2839, 2010.
Guangye Chen, Luis Chacón, and Daniel C Barnes. An efficient implementation of a nonlinearly implicit one-dimensional particle-in-cell algorithm. Journal of Computational Physics, 23 (3(())63374-53388 2012.
hybrid cpu gpu
Viktor K Decyk and Tajendra V Singh. Adaptable particle-in-cell algorithms for graphical processing units. Computer Physics Communications, 182(3):641-66 201
Julien Derouillat, Amaud Beck, Frédéric Pérez, Tommaso Vinci, M Chiaramello, Anna Grassi, M Flé, Guillaume Bouchard, I Plotnikov, Nicolas Aunai, et al. Smilei: A collaborative, open-source, multi-purposee particle-in-cell code for plasma simulation. Computer Physics Communications, 222:351-373, 2018.
Chenru Duan, Yuanqi Du, Haojun Jia, and Heather J. Kulik. Accurate transition state generation with an object-aware equivariant elementary reaction diffusion model, 2023,
Rahman Ejaz, Varchas Gopalaswamy, A Lees, C Kanan, D Cao, and R Betti. Deep learning-basse predictive models for laser direct drive the omega laser facility. Physics of Plasmas, 31(5). 2024.
Ricardo A Fonseca, Luis O Silva, Frank S Tsung, Viktor K Decyk, Wei Lu, Chuang Ren, Warren B Mori, Shaogui Deng, Shiyoun Lee, Katsouleas, et al. Osiris: A dinensiiinsional, fully relativis- tic particle in cell code for modeling plasma based accelerators. In Computational Science- -ICCS 2002: International Conference Amsterdam, The Netherlands, April 21-24, 2002 Proceedings, Part III 2, pp. 242-351. Springer, 2002.
10

---

Published as a conference paper at ICLR 2025
Laurent Garrigues, Gwénaël Fubiani, and Jean-Pierre Boeuf. Negative ion extraction via particle simulation for fusion: critical assessment of recent contributions. Nuclear Fusion, 57(1)0014003, 2016.
Jonathan Ho, Ajay Jain, and Pieter Abbeel. Denoising diffusion probabilistic models. In NeurIPS, 2020.
Zeviel Imani, Shuchin Aeron, and Taritree Wongjirad. Score-based diffusion models for generating liquid argon time projection chamber images, 2024.
Tero Karras, Samuli Laine, Miika Aittala, Janne Hellsten, Jaakko Lehtinen, and Timo Aila. Analyzing and improving the image quality stylegan. In Proceedings of the IEEE/CVF conference on computer vision and pattern recognition, pp. 111088119 2020.
Denis Keefe. Inertial confinement fusion. Annual Review of Nuclear and Particle Science, 32(1): 391-441, 1982.
Ondrej Klimo, Shlomo Weber, VT Tikhonchuk, and Limpouch. Particle-in-cell simulations of laser- laser-plasma interaction for the shock ignition scenario. Plasma Physics and Controlled Fusion, 52(5):055013, 2010.
Xianglong Kong, Michael C Huang, Chuang Ren, and Viktor Decyk. Particle-in-cell simulations with charge-conserving current deposition on graphic processing units. Journal of Computational Physics, 230(4):667611655 2011.
Kube, RM Churchill, and B Sturdevant. Machine learning accelerated particle-in-cell plasma simulations. arXiv preprint arXiv:21 (0.::2/00.12400 124444 2021.
A Bruce Langdon. Evolution of particle-in-cell plasma simulation. IEEE Transactions on Plasma Science, 42(5):1317-1320, 2014.
Lizao Li, Robert Carver, Ignacio Lopez-Gomes Fei Sha, and John Anderson. Generative emulation of weather forecast ensembles with diffusion models. Science Advances, 10(13)eaa)44489 2024.
Xingchao Liu, Chengyue Gong, and Qiang Liu. Flow straight and fast: Learning to generate and transfer data with rectified flow, 2022.
Xingchao Liu, Xiwen Zhang, Jianzhu Ma, Jian Peng, and Qiang Liu. Instaflow: One step is enough for high-quality diffusion- text-to-image generation. 2024.
Norihiro Oyama, Noriko N Ishizaki, Satoshi Koide, and Hiroaki Yoshida. Deep generative model super-resolves spatially correlated multiregional climate data. Scientific Reports, 13(1):5992, 2023.
Magnus Petersen, Gemma Roig, and Roberto Covino. Dynamicsdiffusion: Generating and rare event sampling of molecular dynamic trajectories using diffusion models. In NeurlPS 2023 Al for Science Workshop, 2023.
Olaf Ronneberger, Philipp Fischer, and Thomas Brox. U-net: Convolutional networks for biomedical image segmentation. In MICCAI, 2015.
Andrea Sgattoni, Luca Fedeli, Stefano Sinigardi, Alberto Marocchino, Andrea Macchi, Volker Weinberg, and Anupam Karmakar. Optimising piccante-an open source particle-in-cell code for advanced simulations tier-O systems. arXiv preprint arxiv:1503.02465: 054.40443.0444 2015.
Harshil Shah. Siddharth Kamaria, Riddhesh Markandeya, Miral Shah, and Bhaskar Chaudhury. A novel implementation of 2d3v particle-in-cell (pic) algorithm for kepler gpu architecture, In 2017 IEEE 24th International Conference on High Performance Computing (HiPC), pp. 378-387. IEEE, 2017.
Michael J Smith, James Geach, Ryan A Jackson, Nikhil Arora, Connor Stone, and Stéphane Courteau. Realistic galaxy image simulation via basee-based generative models. Monthly Notices of the Royal Astronomical Society, 51 (2): 111118088881 1818, January 2022. ISSN 1365-2966.
Yang Song and Stefano Ermon. Generative modeling by estimating gradients of the data distribution. In NewrlPS, 2019.
11

---

Published as a conference paper at ICLR 2025
LJ Suter, S Glenzer, S Haan, B Hammel, K Manes, N Meezan, Moody, M Spaeth, L Divol, K Oades, et al. Prospects for high-gain, high yield national ignition facility targets driven by 2w (green) light. Physies of Plasmas, 11(5):273-2277 2745, 2004.
Junya Suzuki, Hironori Shimazu, Keicchir Fukazawa, and Mitsue Den. Acceleration of pic simulation with gpu. Plasma and Fusion Research, 6:2401075-2401075, 2011.
F Taccogna and P Minelli. Pic modeling of negative ion sources for fusion. New Journal of Physics, 19(1):015012, 2017.
David Tskhakaya, Konstantin Matyash, Ralf Schneider, and Francesco Taccogna. The particle-in-cell method. Contributions to Plasma Physics, 47(8-9):533-594 2007.
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention is all you need. In NeurIPS, 2017.
John P Verboncoeur. Particle simulation of plasmas: review and advances. Plasma Physics and Controlled Fusion, 47(5A):A231, 2005.
Patrick von Platen, Suraj Patil, Anton Lozhkov, Pedro Cuenca, Nathan Lambert, Kashif Rasul, Mishig Davaadorj. Dhruv Nair, Sayak Paul, William Berman, Yiyi Xu, Steven Liu, and Thomas Wolf. Diffusers: state-of-the-ae diffusion models.
diffusers 2022.
Pan Wang and Xiangqin Zhu. Hybrid cpu-and gpu-based implementation for particle-in-cell sim- ulation on multicore and multi-gpu systems. In 2021 Photonics & Electromagnetics Research Symposium (PIERS), pp. 155-161. IEEE, 2021.
Fang Wu and Stan Z Li. Diffmd: a geometric diffusson model for molecular dynamics simulations. In AAAI, 2023.
Ming Xu, Feiguo Chen, Xinhua Liu, Wei Ge, and Jinghai Li. Discrete particle simulation of gas-solid two-phase flows with multi-scale epu-gpu hybrid computation. Chemical engineering journal, 207:746-757, 2012.
Qinsheng Zhang and Yongxin Chen. Diffusion normalizing flow. Advances neural information processing systems, 34:16280-16291 2021.
12

---

Published as a conference paper at ICLR 2025
A APPENDIX
A.I IMPLEMENTATION DETAILS OF DIFF-PIC
Physically-informedd parameter encoder (PIPE). The PIPE module comprises two encoding strategies (i.e. positional encoders and polynomial encoders) followed by single-layer MLP. The positional encoder (Vaswani et al. 2017), widely used in Large Language Models, utilizes sinusoidal functions to encode the input parameters, facilitating smooth interpolation. For an input parameter æ, the positional encoder of PIPE generates 1 16-dimensional embedding using sine and cosine functions at varying frequencies:
sin
if is even,
(os;(x) =
10000 #
cos
10000 A
if is odd,
where 0 S 16 and d = 16. Polynomial encoders include polynomial terms up to the fourth degree: (x, 2², xx). These polynomial terms are widely used to approximate wide range of functions effectively, capturing nonlinear correlations and unbounded growth patterns in the data, which is crucial for extrapolation. For parameters T,, I, each of them is processed separately through parallel positional and polynomial encoders, concatenated as a 20-dimensional embedding, and transformed by a MIP, resulting in refined 20-dimensional embedding. The parameter las undergoes a positional encoder followed by MLP transformation, resulting in refined 16-dimensional embedding. Finally, embeddings from these four parameters are concatenated produce a final encoded representation, which is then combined with the simulation data as inputs to the Diff-PIC.
The architecture and configurations of the modified U-Net backbone. The visualization of the modified U-Net backbone is shown in Fig. O built using the HuggingFace Diffusers framework (von Platen et al., 2022). The U-Net backbone processes input images through an encoder-decoder architecture with skip connections. The encoding path comprises three downsampling blocks: standard Down module followed by two AttnDownBlock2D modules, progressively reducing spatial dimensions by 4x (from HxW to H/4x W/4) while increasing the channel dimension from 64 to 256. Attention mechanisms are strategically integrated into both down- wwwn-samping and up- sampling paths to capture long-range dependencies, with the attention operations performed on flattened spatial dimensions. The decoder path mirrors the encoder with three upsampling blocks (two Attn Attnokkoek2D modules and one UpBlock2D module), using skip connections to preserve fine spatial details. Finally, the model applies a convolution layer to map the representation back to the original input dimension and resolution.
A.2 DETAILS OF BASELINES
GAN-PIC is based on StyleGAN2 (Karras et al. 2020), SOTA variant of Generative Adversarial Networks (GANs). It comprises two main components: a Generator and a Discriminator. The Generator StyleGAN2 is enhanced with style-based architecture that uses a mapping network to convert latent vectors into style vectors. These style vectors are then applied at various layers of the generator. The Discriminator's role to differentiate between real and generated images, driving the adversarial training process that improves the fidelity of the generated outputs. NF-PIC implemented based on Zhang & Chen 202T). NF-PIC has drift network and a score network. both modeled using U-Net architectures. Apart from the integration of PIPE, both baselines are implemented following the experimental setups detailed in their respective original papers.
Input Image
Corv In
DownBlock2D
AttnDownBlock2D
AltnDownBBloc200
Skip Connection
Skip Connection
Skip Connection
UNeiMidBlock2D
Ouput Image
Conv Out
UpBlock2D
AttnUpBlock2D
AltnUpBlock20
Figure The U-Net backbone.
13

---

Published as a conference paper at ICLR 2025
Table Comparisons between MLP, Trans, and PIPE in interpolation tasks.
Training Set for E1
Testing Set for E1
Training Set for E2
Testing Set for F2
Method
MAEL
RMSEL
FIDI
MAE]
RMSEL
FID
MAEL
RMSEJ
FID
MAEJ
RMSE,
FID
MLP
3.72e-2
4.16e-2
2.06
3.90e-2
4.57e-2
2.15
1.71e-2
1.92e-2
0,926
185e-2
1.9e-2
0.932
Trans
1.6le-2
2.80e-2
1.35
1.72e-2
2.31e-2
1.70
8.61e-3
1.02e-2
0.375
9.41e-3
1.18e22
0.409
PIPE
1.56e-2
2.67e-2
1.21
1.68-22
2.29e-2
1.62
7.95e-3
9.32e-3 0.328
8.26e-3
1.0ee-2 0.341
Table 5: Comparisons between MLP, Trans, and PIPE in extrapolation tasks.
10%
E1 20%
E2 10%
E2 20%
Method
MAE
RMSEL
FID.
MAEI
RMSEL
FID
MAEL
RMSEL
FID
MAE
RMSEL
FID
MLP
5.12e-2
6.01e-2
1,93
5,42e-2
5.64e-2
1.95e-2
2.62e-2
1.09
2.32e-2
2.91e-2
1.16
Trans
1.95-2
2.61e-2
1.82
2.29e-2
2.73e-2
1.89
1.02e-2
1.55e-2
0.615
1.32e-2
2.0e--2
0.761
PIPE
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.85-2
0.673
A.3
ABLATION STUDIES ON THE PROPOSED PIPE
To further demonstrate the effectiveness of PIPE, we conduct ablation studies by replacing the PIPE component in Diff-PIC with other commonly used encoders: MLP (denoted as MLP) and Transformer decoder (denoted as Trans). Specifically, the two-layer MLP has 128 hidden units and ReLU activations. For the Transformer decoder, we first apply positional encoder to generate 16-dimensional embedding for each parameter (the same as the positional encoder in PIPE). This embedding is then processed by the Transformer decoder, which comprises multi-head self self-attention and a two-layer MLP with ReLU activations. The comparison results are presented in Table 4 for interpolation tasks and Table 5 for extrapolation tasks. The results show that PIPE significantly outperforms both MLP and Transformer decoder in both interpolation and extrapolation tasks. MLP struggles with both interpolation and extrapolation. The Transformer decoder, while powerful for sequence modeling and capturing relationships in discrete token sequences, is not inherently designed to effectively process continuous physical parameters, resulting in inferior performance.
To provide more insights into the effectiveness of PIPE, we conduct another ablation study to evaluate the contributions of the positional encoder (only pos), polynomial encoder (only poly) separately. The results are shown in Table Table| (interpolation tasks) and Table 7 (extrapolation tasks), highlighting the individual importance of these encoders. Specifically, combining both positional and polynomial encoders consistently outperforms using either encoder alone in both interpolation and extrapolation tasks. In interpolation tasks, the positional encoder shows better performance than the polynomial encoder, suggesting its important role in interpolation scenarios. Conversely, in extrapolation tasks, particularly at larger ranges (20%), the polynomial encoder shows relatively better performance, indicating its importance for extrapolation capabilities. These results validate our design choice of combining both encoders in PIPE to leverage their complementary strengths.
Table Ablation studies on different components of PIPE in interpolation tasks.
Training Set for E1
Testing Set for E1
Training Set for E2
Testing Set for E2
Method
MAE]
RMSEJ FIDJ
MAEІ
RMSEJ
FIDJ
MAEJ RMSEL FIDJ
MAEІ
RMSEJ
FID
only pos
1.64e-2
2.82e-2
1.37
1.76e-2
2.35e-2 1.72
8,63e-3
1.05e-2 0.379
9.42e-3
1.20e-2
0.412
only poly
1.76e-2
2.94e-2
1.52
1.92e-2
2.6le-2
1.14-22
1.27e-2
0.517
1.13e-2
1.35e-2
0.546
PIPE
1.56e-2
2,67e-2
1.21
1,68e-2
2.29c-2
1.62
7.95e-3
9.33e-3 0.328
8.26e-3
1.03--2 0.341
Table 7: Ablation studies on different components of PIPE in extrapolation tasks.
E1
E1 20%
10%
20%
Method
MAE!
RMSE,
FID
MAEL
RMSEL
FIDJ,
MAEJ
RMSEL
FIDI
MAEL
RMSEL
FIDL
only pos
1.97e-2
2.64e-2
1.83
2.31e-2
2.76e-2
1.93
1.03e-2
1.57e-2
0.624
1.35e-2
2.10e-2
0.812
only poly
1.91e-2
2.52e-2
1.78
2.25e-2
2.70e-2
1.89
9.75e-3
1.42-22
0.581
1.26c-2
1.92e-2
0.724
PIPE
1.83e-2
2.40e-2
1.74
2.18
2.62e-2
1.82
9.47e-3
1.36e-2
0.536
1.13e-2
1.85e-2
0.673
14


Reviews:
[{"review_id": "eOtsQs7eT1", "reviewer": "Reviewer_FMwc", "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.", "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.", "limitations": "", "rating": 8, "confidence": 3, "soundness": 4, "presentation": 3, "contribution": 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.\n\nThe 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.", "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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "HoubJyPp4R", "reviewer": "Reviewer_Mxd3", "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.", "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\n\nPerhaps authors could add uncertainty estimates to numbers in Tables? such as standard deviation over multiple runs or confidence intervals\n\nCan 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", "limitations": "", "rating": 6, "confidence": 2, "soundness": 2, "presentation": 3, "contribution": 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.", "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.\n\nthe main conclusion from paper's experiments seems to be that diffusion with u-net performs better than Gan or normalizing flow, for this task.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "nUKjT0B76l", "reviewer": "Reviewer_4k8H", "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.\n\n\n========== post rebuttal\n\nI 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.", "questions": "Please see the weakness part.\n\nCompared with previous work, \"diffusion models in molecular dynamics simulations \"\n\nI hope the author can clarify:\n1. The above cited paper is first paper apply diffusion to a similar domain.\n2. 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?\n3. \"Substantial efficiency improvement\" I got this achieved by apply normalizing flow.", "limitations": "", "rating": 6, "confidence": 2, "soundness": 2, "presentation": 2, "contribution": 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.", "weaknesses": "It's unclear to me what is the ML contribution for this paper. \nTo me it seems most interesting part is the newly designed encoder. Can the author explain:\n\"Algorithmic generalization. PIPE improves the generalizability of the conditional diffusion model\"\nI don't quite get why such design improved generalizability compared with normal mlp/transformer layers. I hope the author would clarify it.\n\nThe 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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "dNa9QmBjaC", "reviewer": "Reviewer_qGUg", "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). \n\nResults 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. \n\nPost rebuttal\n=========\nThe 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:\n- it works as well or better than competitive methods \n- it is significantly faster than mainstream computational approaches used in plasma physics (owing to the rectified flow acceleration technique)\nHowever, 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.\n\n[1] Xingchao Liu, Chengyue Gong, and Qiang Liu. Flow straight and fast: Learning to generate and\ntransfer data with rectified flow, 2022.\n\n[2] Patrick Esser, Sumith Kulal, Andreas Blattmann, Rahim Entezari, Jonas Muller, Harry Saini, Yam ¨\nLevi, Dominik Lorenz, Axel Sauer, Frederic Boesel, et al. Scaling rectified flow transformers for\nhigh-resolution image synthesis. arXiv preprint arXiv:2403.03206, 2024.\n\n[3] Xingchao Liu, Xiwen Zhang, Jianzhu Ma, Jian Peng, and Qiang Liu. Instaflow: One step is enough\nfor high-quality diffusion-based text-to-image generation, 2024.\n\n[4] GAN: Alankrita Aggarwal, Mamta Mittal, and Gopi Battineni. Generative adversarial network: An overview\nof theory and applications. International Journal of Information Management Data Insights, 1(1):\n100004, 2021.\n\n[5] NF: Qinsheng Zhang and Yongxin Chen. Diffusion normalizing flow. Advances in neural information\nprocessing systems, 34:16280–16291, 2021.", "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? \n\n- 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?\n\n- 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).", "limitations": "", "rating": 5, "confidence": 4, "soundness": 3, "presentation": 2, "contribution": 2, "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. \n\n- The methodology appears to be sound, and the evaluation follows the techniques used in generative modelling generally. \n\n- The evaluation fares quite well, compared to competitors (GAN and normalizing flow), achieving better performance and speed.", "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. \n\n- 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. \n\n- 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? \n\n- Network Architecture: Perhaps I am missing this, but I would like to see a block containing the network design - transformer blocks, CNNs, etc.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "RGTPMBXDuw", "reviewer": "Reviewer_gLQJ", "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.", "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.", "limitations": "", "rating": 8, "confidence": 4, "soundness": 3, "presentation": 4, "contribution": 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.", "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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (Poster)