Lawrence Berkeley National Laboratory

Heavy-ion collisions produce final states with thousands to tens of thousands of particles, making their simulation among the most computationally intensive tasks in high-energy nuclear physics. We present a fast, high-fidelity generative model for heavy-ion events based on a score-driven diffusion process and the point-edge transformer architecture within the OmniLearn framework. A two-stage training strategy is performed: Stage 1 trains on lower-multiplicity O-O collisions, allowing the model to learn a stable event-and-particles representation; this is followed in Stage 2 by fine tuning on challenging high-multiplicity Pb-Pb collisions. We benchmark the generator with a broad set of closure checks, including agreement of event- and particle-level observables in one and two dimensions, flow consistency reconstructed from the generated particles, end-to-end jet finding with FastJet including key jet and substructure observables, and a classifier-based application to quantify the sample fidelity. The results are promising, showing that a compact generative model can produce realistic, high-multiplicity heavy-ion events, at a level that makes local-scale generation for heavy-ion collisions at high energies a practical goal.