Ultra fast, event-by-event heavy-ion simulations for next generation experiments

TL;DR

The paper tackles the computational bottleneck of event-by-event heavy-ion simulations by introducing HEIDi, a conditional diffusion-based generator operating on fixed-size, multi-feature point clouds that represent complete collision events. It combines a PointNet-based encoder, a normalizing-flow decoder, and a diffusion model to produce UrQMD-like events across 26 hadron species, conditioned on global event properties. Empirical results show HEIDi reproduces key momentum, rapidity, and multiplicity distributions and event-level correlations, while delivering roughly two orders of magnitude speedups on GPUs, enabling rapid Bayesian inference and real-time model-data comparisons. The approach is adaptable to varying collision conditions and broader data-processing tasks, laying groundwork for a foundation model for heavy-ion collisions and beyond.

Abstract

We present a novel deep generative framework that uses probabilistic diffusion models for ultra fast, event-by-event simulations of heavy-ion collision output. This new framework is trained on UrQMD cascade data to generate a full collision event output containing 26 distinct hadron species. The output is represented as a point cloud, where each point is defined by a particle's momentum vector and its corresponding species information (ID). Our architecture integrates a normalizing flow-based condition generator that encodes global event features into a latent vector, and a diffusion model that synthesizes a point cloud of particles based on this condition. A detailed description of the model and an in-depth analysis of its performance is provided. The conditional point cloud diffusion model learns to generate realistic output particles of collision events which successfully reproduce the UrQMD distributions for multiplicity, momentum and rapidity of each hadron type. The flexible point cloud representation of the event output preserves full event-level granularity, enabling direct application to inverse problems and parameter estimation tasks while also making it easily adaptable for accelerating any event-by-event model calculation or detector simulation.

Figures (9)

• Figure 1: (Color online) HEIDi: network structure. The red and green arrows show the flow of information during training and generation, respectively. Only the parts shaded in gray are used during the generation process. A collision event output is represented by a point cloud containing 1084 points where each point is a final state particle. $\sigma$ and $\mu$ represent the mean and standard deviation, respectively, of a 128 dimensional Gaussian distribution used to construct the latent condition vector $z$. The generation process starts with random samples from the standard normal distribution $\mathcal{N}(0,I)$ (mean= zero vector, covariance = identity matrix $I$), which is processed by the normalizing flow-based decoder and the diffusion model to generate the final state event output point cloud. The labels '$1 \times 128$' and '$1084 \times 32$' refer to the dimensionality of the standard normal distribution from which the random samples are drawn.
• Figure 2: (Color online) Momentum distributions of various hadron types generated by HEIDi for Au-Au collisions with $b=1$ fm at 10 $A$GeV. Each row corresponds to one hadron type. The columns 1, 2 and 3 represent the distributions for $p_x$, $p_y$, and $p_z$ respectively. UrQMD and HEIDi results are shown in blue and red curves respectively.
• Figure 3: (Color online) Momentum distributions of various hadron types generated by HEIDi for $b=1$ fm, Au-Au collisions at 10 $A$GeV. The color scheme and plot layouts are similar to figure \ref{['mom1']}.
• Figure 4: (Color online) Transverse momentum distribution of various hadrons. The results are for 10 $A$GeV Au-Au collisions with impact parameter $b=1$ fm. The HEIDi distributions are shown as red curves while the UrQMD distributions are shown in blue. The dashed curves in the upper left two plots are the $p_T$ distributions for the spectator nucleons where the solid curves in the plots show the $p_T$ distributions for participant nucleons.
• Figure 5: (Color online) $\langle p_T \rangle$ distribution of various hadrons for Au-Au collisions with $b=1$ fm at 10 $A$GeV. The HEIDi distributions are shown as red curves, while the UrQMD distributions are shown in blue. The dashed curves in the upper left two plots are the distributions for spectator nucleons and the solid curves in the plots represent the distributions for participant nucleons.