[论文解读] EPiC-ly Fast Particle Cloud Generation with Flow-Matching and Diffusion

Erik Buhmann, Cedric Ewen+8 more|arXiv (Cornell University)|Sep 29, 2023

Generative Adversarial Networks and Image Synthesis被引用 14

一句话总结

本文提出两种置换等变的喷流云生成器，EPiC-JeDi (score-based diffusion) 和 EPiC-FM (flow-matching CNF)，在 JetNet top-jet 数据集上实现了最先进的保真度，且生成速度很快。它还分析了无条件与条件生成，并在一个统一的连续时间框架中将扩散与基于流的观点联系起来。

ABSTRACT

Jets at the LHC, typically consisting of a large number of highly correlated particles, are a fascinating laboratory for deep generative modeling. In this paper, we present two novel methods that generate LHC jets as point clouds efficiently and accurately. We introduce \epcjedi, which combines score-matching diffusion models with the Equivariant Point Cloud (EPiC) architecture based on the deep sets framework. This model offers a much faster alternative to previous transformer-based diffusion models without reducing the quality of the generated jets. In addition, we introduce \epcfm, the first permutation equivariant continuous normalizing flow (CNF) for particle cloud generation. This model is trained with {\it flow-matching}, a scalable and easy-to-train objective based on optimal transport that directly regresses the vector fields connecting the Gaussian noise prior to the data distribution. Our experiments demonstrate that \epcjedi and \epcfm both achieve state-of-the-art performance on the top-quark JetNet datasets whilst maintaining fast generation speed. Most notably, we find that the \epcfm model consistently outperforms all the other generative models considered here across every metric. Finally, we also introduce two new particle cloud performance metrics: the first based on the Kullback-Leibler divergence between feature distributions, the second is the negative log-posterior of a multi-model ParticleNet classifier.

研究动机与目标

• 通过利用置换等变架构实现快速、可扩展采样来提升喷流云生成。
• 证明基于扩散的和基于流的连续时间模型可以在一个统一的框架下用于喷流生成。
• 在 JetNet top-jet 数据集上与现有模型进行基准对比，以确立保真度和速度提升。
• 研究条件与无条件生成的差异，以及喷流级条件对性能的影响。

提出的方法

• 将 EPiC 层与 score-matching diffusion 结合用于 EPiC-JeDi，以获得快速、准确的喷流云。
• 引入 EPiC-FM，这是首个使用 flow-matching 训练的置换等变 CNF，用于实现笔直、最优传输状轨迹。
• 在两种模型中使用将条件信息拼接到每个 MLP 输入的 EPiC Network 架构。
• 将喷流表示为点云，带有相对的Hadronic 坐标，并在需要时对喷流级观测量如 pT 和质量进行条件生成。
• 在条件设置中，使用条件正则化流（NormFlows）来生成诸如质量、pT 和组成粒子数等联合喷流属性。
• 讨论采样选项（基于 ODE 的确定性采样和基于 SDE 的随机采样），并将 flow-matching 与基于得分的目标联系起来。

Figure 2 : Model schema of the EPiC Network generator architecture used in both EPiC-JeDi and EPiC-FM . Each multi-layer perceptron (MLP) is a two-layer neural network with LeakyReLU activation. The pooling operation is a concatenation of both average and summation pooling.

实验结果

研究问题

• RQ1 EPiC-JeDi（diffusion-based）和 EPiC-FM（flow-matching CNF）在保持快速采样速度的同时，是否能够在 top-jet 生成方面达到最先进的保真度？
• RQ2 在条件生成与无条件生成方面，置换等变喷流云模型的保真度与效率有何差异？
• RQ3 在这一喷射物理场景中，使用 flow-matching 与 score-matching 目标对训练稳定性和采样效率有何影响？
• RQ4 所提出的模型是否在不同的 JetNet top-jet 配置（例如不同的成分粒数）和条件设定下具有泛化能力？

主要发现

• 相较于其他模型，EPiC-FM 在评估指标上始终提供更高的保真度。
• 条件化的 EPiC-FM 在测试配置中实现了最高的生成保真度。
• EPiC-JeDi 提供有竞争力的性能，并且采样速度快于基于 Transformer 的扩散模型。
• 两者在 top-jet 生成方面都表现强劲，EPiC-FM 在多项指标上通常优于其他基线。
• 本研究引入了两项新的粒子云指标：基于 KL 发散的特征分布以及一个多模型的 NLP 基准分类器。

Figure 6 : Generation time per jet as a function of generated particle multiplicity. The generation time is measured on the same hardware using an NVIDIA ® A100 40 GB graphics card, with per-jet times calculated as the average for all test events generated. The batch size was optimised for optimal g