PC-Droid: Faster diffusion and improved quality for particle cloud generation

TL;DR

PC-Droid presents an optimized diffusion-based generator for jet particle clouds, building on PC-JeDi with an EDM-style diffusion formulation, advanced solvers, and simultaneous training on all jet types to achieve state-of-the-art fidelity and speed. It introduces a cross-attention encoder to reduce computational cost, consistency distillation for ultra-fast generation, and normalizing flows to enable unconditional generation conditioned on jet kinematics. Across JetNet30 and JetNet150, PC-Droid demonstrates superior or competitive metrics, notably improving tail $p_T$ distributions and substructure observables such as $\tau_{32}$, while enabling generation times approaching real-time scales. The combination of diffusion reformulation, fast solvers, and CD-based one-shot generation substantially advances fast jet simulation, with potential applicability to other particle cloud tasks such as calorimeter shower modeling.

Abstract

Building on the success of PC-JeDi we introduce PC-Droid, a substantially improved diffusion model for the generation of jet particle clouds. By leveraging a new diffusion formulation, studying more recent integration solvers, and training on all jet types simultaneously, we are able to achieve state-of-the-art performance for all types of jets across all evaluation metrics. We study the trade-off between generation speed and quality by comparing two attention based architectures, as well as the potential of consistency distillation to reduce the number of diffusion steps. Both the faster architecture and consistency models demonstrate performance surpassing many competing models, with generation time up to two orders of magnitude faster than PC-JeDi and three orders of magnitude faster than Delphes.

Figures (13)

• Figure 1: PC-Droid network architecture and training setup.
• Figure 2: A single cross-attention encoder block updating both an input point cloud $x$ and global tokens $g$, using layer-norm (LN), multi-headed attention (MHA). Converging arrows represent vector addition and contextual information $c$ is injected into the network by concatenating to the inputs of the MLPs. The first attention operation effectively pools information from the point cloud into the global tokens. The second attention operation is inverted, and information from the updated globals token are distributed back to the point cloud.
• Figure 3: Performance as measured by FPND (left), W$_1^{\tau_{32}}$ (middle), and W$_1^\mathrm{M}$ (right) on the NFE for top jets with up to 30 constituents. We generate samples with PC-Droid using Heun and HeunSDE methods Karras2022, DPM2 DPMSolver, DPM2 with ancestral sampling (DPM2A), DPM++2M DPMSolverPP, and LMS ODEBook.
• Figure 4: Comparison of $p_\mathrm{T}$ distributions of the leading, fifth leading, and twentieth leading constituents of the generated top and gluon jets with up to 30 constituents.
• Figure 5: Mass and substructure distributions of the generated top jets with up to 30 constituents. The diagonal consists of the marginals of the distributions and the off-diagonal elements contain the joint distributions.