[论文解读] PC-JeDi: Diffusion for Particle Cloud Generation in High Energy Physics

Matthew Leigh, D. Sengupta+4 more|arXiv (Cornell University)|Mar 9, 2023

Particle physics theoretical and experimental studies参考文献 86被引用 5

一句话总结

PC-JeDi 引入了一种基于扩散的生成模型，采用基于分数的学习方法与 Transformer 架构，用于在高能物理中生成高保真度的粒子云喷射流。该模型在喷射流子结构保真度和条件生成方面达到当前最先进水平，尽管推理速度慢于 GAN 模型，但仍提供了一种灵活且排列等变的框架，适用于复杂喷射流模拟，相比传统蒙特卡洛方法提速 10 倍。

ABSTRACT

In this paper, we present a new method to efficiently generate jets in High Energy Physics called PC-JeDi. This method utilises score-based diffusion models in conjunction with transformers which are well suited to the task of generating jets as particle clouds due to their permutation equivariance. PC-JeDi achieves competitive performance with current state-of-the-art methods across several metrics that evaluate the quality of the generated jets. Although slower than other models, due to the large number of forward passes required by diffusion models, it is still substantially faster than traditional detailed simulation. Furthermore, PC-JeDi uses conditional generation to produce jets with a desired mass and transverse momentum for two different particles, top quarks and gluons.

研究动机与目标

• 开发一种快速、准确的喷射流模拟代理模型，相比传统蒙特卡洛方法显著降低计算成本。
• 通过利用扩散模型与排列等变架构，解决以点云形式生成真实喷射流子结构的挑战。
• 实现对特定物理属性的喷射流条件生成，例如具有特定横动量与质量的顶夸克喷射流和胶子喷射流。
• 通过全面的指标评估模型性能，包括此前生成喷射流模型中未报告的子结构可观测量。
• 探索基于扩散的点云生成中生成速度与保真度之间的权衡，为未来优化提供可能。

提出的方法

• PC-JeDi 采用基于分数的扩散过程，通过逆转噪声扰动轨迹，逐步学习从标准正态先验中去噪粒子动量。
• 该模型采用基于 Transformer 的架构，具备排列等变性，天然适用于处理喷射事例中无序的粒子云数据。
• 通过在目标喷射流质量与横动量上条件化扩散过程，实现条件生成，从而实现对特定类型喷射流的可控采样。
• 去噪过程通过自适应步长的 ODE/SDE 求解器实现，利用数据分布上的分数匹配估计分数函数。
• 训练过程通过在一系列时间步上最小化去噪损失实现，网络在每一步预测分数，以引导反向过程。
• 推理过程涉及 200 步的迭代去噪，从随机噪声开始，逐步优化粒子动量，直至生成逼真的喷射流构型。

Figure 2 : Diagram of the PC-JeDi generation procedure. Random noise is sampled at the beginning of the loop from a standard normal distribution. Then any chosen integration sampler is iteratively applied in order to fully denoise the input towards actual data, using the conditional model as a noise

实验结果

研究问题

• RQ1 基于 Transformer 架构的扩散生成模型能否生成与真实喷射流统计特性匹配的高保真度粒子云喷射流？
• RQ2 在子结构可观测量与保真度指标方面，PC-JeDi 相较于当前最先进模型 MPGAN 的表现如何？
• RQ3 PC-JeDi 在生成具有特定物理属性的喷射流（如具有特定质量与横动量的顶夸克喷射流或胶子喷射流）方面的能力有多强？
• RQ4 PC-JeDi 相较于基于 GAN 的模型与传统蒙特卡洛模拟，其计算成本如何？
• RQ5 该模型能否捕捉喷射流子结构中的复杂相关性，例如非包含型顶夸克喷射流或二维边缘分布？

主要发现

• PC-JeDi 在多个指标上与当前最先进模型 MPGAN 表现相当，包括衡量喷射流子结构保真度与分布相似性的指标。
• 该模型在若干关键指标上优于 MPGAN，尤其在捕捉子结构可观测量（如 N-子喷射性与平面流）分布方面表现更优。
• PC-JeDi 能够成功生成具有指定横动量与质量的条件喷射流，证明其对输出物理属性具有精确控制能力。
• 尽管每一样本需要约 200 次前向传播，PC-JeDi 的生成速度仍约为传统蒙特卡洛模拟的 10 倍。
• 在 GPU 上生成 1000 个喷射流的推理时间约为 500ms，相较蒙特卡洛模拟的约 4721ms，显示出显著的速度优势。
• 基于 Transformer 的架构在批处理中展现出更优的可扩展性，推理时间随批量大小线性增长，而某些 GAN 模型则表现出非线性增长。

Figure 3 : The relative transverse momentum (left) and invariant mass (right) of gluon jets generated with MPGAN (orange) and PC-JeDi (DDIM solver, green; EM solver, red) compared to the MC simulation (shaded blue). Calculated from the leading 30 $p_{\mathrm{T}}$ constituents using the constituent $