[论文解读] PC-Droid: Faster diffusion and improved quality for particle cloud generation

Matthew Leigh, Debajyoti Sengupta+3 more|arXiv (Cornell University)|Jul 13, 2023

Advanced Neural Network Applications参考文献 59被引用 5

一句话总结

PC-Droid 提出了一种用于高能物理中粒子云生成的最先进扩散模型，通过改进的 EDM 噪声调度、先进的 ODE 求解器以及一致性蒸馏技术，实现了相比 PC-JeDi 快两个数量级、相比 Delphes 快三个数量级的生成速度，同时在所有喷射类型和指标下保持或提升了生成质量。

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.

研究动机与目标

• 开发一种更快、更高保真度的扩散模型，用于在高能物理模拟中生成作为点云的粒子喷射。
• 在保持或提升生成质量的前提下，减少推理时间，相比现有的代理模型如 PC-JeDi 和 Delphes。
• 探索基于注意力机制的架构与一致性模型在生成速度与质量之间的权衡。
• 通过在所有喷射类型上联合训练，提升泛化能力并实现无条件生成能力。
• 使用标准指标和真实世界硬件评估，建立粒子云生成的新基准。

提出的方法

• 采用 EDM（去噪得分基于）扩散公式，使用简化的 SDE 和 ODE，其中 $\gamma(t) = 1$ 且 $\sigma(t) = t$，从而实现更直的反向轨迹并减少截断误差。
• 采用噪声率 $\sigma \sim \mathcal{N}(-1.2, 1.2)$ 的对数正态采样，将训练重点集中在性能提升最显著的中间噪声水平。
• 引入 $\sigma$-依赖的缩放与跳跃连接（$c_{\text{in}}, c_{\text{out}}, c_{\text{skip}}$），以保持单位方差，并允许在低噪声水平下绕过网络。
• 将损失函数重新表述为 $\mathcal{L} = \|D_\theta(\boldsymbol{x}; \sigma, \boldsymbol{y}) - \boldsymbol{x}_0\|^2_2$，直接最小化重建误差，而非噪声预测。
• 使用 k-diffusion 库中的先进 ODE 积分求解器，其中 DPM2 在一致性模型训练中表现更优。
• 利用 PC-Droid 作为教师模型，训练一致性蒸馏模型（CD1、CD5），实现一步生成且保真度可接受，且在五步内性能趋于饱和。

Figure 1 : PC-Droid network architecture and training setup.

实验结果

研究问题

• RQ1 重新构想的扩散公式是否能显著减少推理时间，同时提升粒子云生成的质量？
• RQ2 现代 ODE 求解器与改进的网络归一化技术如何影响喷射粒子云生成的保真度与速度？
• RQ3 一致性模型在多大程度上能够实现一步生成，且质量与标准扩散模型相当？
• RQ4 在不同架构与蒸馏策略下，生成速度与质量之间的性能权衡如何？
• RQ5 在所有喷射类型上联合训练是否能提升泛化能力，并实现无性能下降的无条件生成？

主要发现

• PC-Droid 在所有喷射类型和评估指标上均达到最先进性能，优于 PC-JeDi 和 Delphes 的保真度与速度。
• 最快模型（CD1）的顶喷射生成时间约为 EPiC-GAN 的三倍，但其 FPND 得分比理想 Delphes 基准更近一个数量级。
• CD5（五步一致性模型）比八步 FPCD（PD8）模型快 2.5 倍，同时在所有指标上均有提升，包括 $\mathrm{W}_1^{\mathrm{M}}$ 和 $\mathrm{W}_1^{\tau_{32}}$。
• 完整版 PC-Droid 模型比 EPiC-GAN 慢 250 倍，但在除 $\mathrm{W}_1^{\tau_{32}}$ 外的所有指标上均接近理想性能，后者仍略有偏差。
• 一致性模型表明性能在五步后趋于饱和，CD1 虽存在质量展宽，但子结构再现可接受，尤其在 $\tau_{32}$ 上表现良好。
• 改进的噪声调度与网络归一化技术减少了截断误差，支持更少反向步数的稳定训练，从而在不损失质量的前提下实现更快推理。

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