[论文解读] Fast Point Cloud Generation with Diffusion Models in High Energy Physics

V. M. Mikuni, Benjamin Nachman+1 more|arXiv (Cornell University)|Apr 3, 2023

Generative Adversarial Networks and Image Synthesis被引用 4

一句话总结

该论文提出了一种面向高能物理喷射模拟的快速点云扩散模型（FPCD），可生成高保真度的可变长度点云。通过渐进蒸馏，FPCD相较基线模型将采样时间减少了450倍，实现了单步推理下的最先进性能，同时在顶夸克、W玻色子和胶子等多种喷射类型中保持了高精度。

ABSTRACT

Many particle physics datasets like those generated at colliders are described by continuous coordinates (in contrast to grid points like in an image), respect a number of symmetries (like permutation invariance), and have a stochastic dimensionality. For this reason, standard deep generative models that produce images or at least a fixed set of features are limiting. We introduce a new neural network simulation based on a diffusion model that addresses these limitations named Fast Point Cloud Diffusion (FPCD). We show that our approach can reproduce the complex properties of hadronic jets from proton-proton collisions with competitive precision to other recently proposed models. Additionally, we use a procedure called progressive distillation to accelerate the generation time of our method, which is typically a significant challenge for diffusion models despite their state-of-the-art precision.

研究动机与目标

• 为解决基于网格和固定长度的生成模型在模拟高能物理数据（特别是可变长度、排列不变的点云，如强子喷射）方面的局限性。
• 开发一种扩散模型，以保留喷射数据的连续性、随机性和对称性，避免像素化或固定尺寸的限制。
• 通过渐进蒸馏加速扩散模型推理（通常较慢），实现实时或近实时的大规模数据集模拟。
• 通过单一条件模型统一生成多种粒子类型的喷射（如顶夸克、胶子、W玻色子），避免为每类粒子单独训练模型。

提出的方法

• FPCD采用两阶段扩散模型：第一阶段生成喷射的运动学特征（如能量、横动量），第二阶段在给定喷射类型和运动学信息的条件下生成粒子级运动学特征（如四维动量）。
• 该模型采用基于得分的扩散框架，学习对数概率密度的梯度，支持凸优化并实现稳定训练，无需可逆流。
• 通过掩码策略，模型在训练期间基于掩码输入实现可变粒子多重性的喷射生成（粒子数范围30至150）。
• 应用渐进蒸馏，将完整的512步扩散过程蒸馏为仅需8步甚至1步即可生成的模型，显著加速推理过程。
• 模型通过粒子类型（如顶夸克、胶子）进行条件控制，使单一统一架构可生成多种喷射类型，无需独立模型。
• 该架构利用图神经网络处理排列不变性，并捕捉点云数据中的复杂拓扑关系。

实验结果

研究问题

• RQ1 扩散模型能否有效生成代表高能物理中强子喷射的可变长度、排列不变点云？
• RQ2 如何在不牺牲喷射模拟物理保真度的前提下，降低标准扩散模型的高计算成本？
• RQ3 单一统一模型能否通过粒子类型条件控制，高精度地生成多种喷射类型（如顶夸克、W玻色子、胶子）？
• RQ4 当将扩散采样步数从512步减少至1步时，渐进蒸馏在多大程度上保持了生成质量？
• RQ5 在生成保真度和速度方面，FPCD相较于现有生成模型（如EPiC-GAN和MP-GAN）在喷射生成任务中的表现如何？

主要发现

• FPCD在物理启发的度量上达到最先进性能，150粒子数据集上胶子的粒子多重性生成误差为0.157 ± 0.036，Z玻色子为0.241 ± 0.090。
• 经蒸馏的FPCD单步采样模型（FPCD 1）保持了高保真度，胶子的粒子多重性误差为0.157 ± 0.035，Z玻色子为0.239 ± 0.090。
• FPCD在150粒子下的生成时间从基线的31 × 10³ μs降至66 μs，相较完整模型提速450倍。
• FPCD 1的生成速度优于MP-GAN（150粒子下分别为11 μs vs. 12 μs），表明单步扩散模型在速度上可超越多步GAN，同时保持或超越保真度。
• 该模型成功使用单一架构生成五类喷射（顶夸克、W玻色子、Z玻色子、胶子、轻夸克），避免了为每类粒子单独训练模型。
• 附录B中的直方图对比表明，蒸馏模型在顶夸克引发的喷射所有运动学分布上均与完整模拟高度一致。