[论文解读] Point Cloud Generation using Transformer Encoders and Normalising Flows

Benno Käch, D. Krücker+1 more|arXiv (Cornell University)|Jan 1, 2022

Generative Adversarial Networks and Image Synthesis被引用 11

一句话总结

该论文提出了一种基于归一化流（NFs）并经变压器编码器优化的新型粒子喷射点云生成模型，采用对抗性训练方式，使用第二个变压器判别器进行训练。该方法在JetNet顶夸克数据集上实现了最先进性能，显著提升了相关性建模能力，并将喷射生成速度提升至先前模型的4.6倍，同时保持了训练的稳定性。

ABSTRACT

Data generation based on Machine Learning has become a major research topic in particle physics. This is due to the current Monte Carlo simulation approach being computationally challenging for future colliders, which will have a significantly higher luminosity. The generation of collider data is similar to point cloud generation, but arguably more difficult as there are complex correlations between the points which need to be modelled correctly. A refinement model consisting of normalising flows and transformer encoders is presented. The normalising flow output is corrected by a transformer encoder, which is adversarially trained against another transformer encoder discriminator/critic. The model reaches state-of-the-art performance while yielding a stable training.

研究动机与目标

• 通过开发基于深度学习的高效数据生成方法，缓解高亮度对撞机物理中蒙特卡洛模拟的计算负担。
• 克服归一化流在建模点云数据中复杂相关性（尤其是喷射子结构）方面的局限性。
• 通过使用归一化流作为强初始化的精炼机制，稳定对抗性变压器架构的训练。
• 在JetNet顶夸克喷射数据集上，实现分布相似性和不变质量建模方面的最先进性能。
• 为下游物理分析提供快速、可微且稳定的现实喷射事件生成能力。

提出的方法

• 使用归一化流（具体为15层有理二次样条耦合层）通过可逆、可微变换将训练数据映射到标准正态分布。
• 应用后处理变压器编码器，通过学习加法校正来改进相关性建模，以修正归一化流的输出。
• 采用对抗性训练设置，其中第二个变压器编码器作为判别器/批评家，用于区分真实与精炼后的生成喷射。
• 使用类似WGAN的训练目标，结合梯度惩罚和LSGAN损失以实现稳定优化，学习率调度采用RMSprop和两阶段余弦调度。
• 将喷射数据处理为(η, φ, pT)相对空间中的3D点云，通过零填充和掩码处理可变粒子数量。
• 将不变质量与能量流多项式（EFP）特征作为全局指标用于评估。

实验结果

研究问题

• RQ1 尽管归一化流在建模复杂点云相关性方面存在局限，它们是否仍可作为高能物理中生成建模的稳定且有效基础？
• RQ2 与标准归一化流相比，添加基于变压器的精炼网络是否能显著提升对喷射子结构相关性的建模能力？
• RQ3 采用两个变压器编码器的对抗性训练设置是否能在保持训练稳定性的同时，超越现有基于GAN的模型在喷射点云生成上的表现？
• RQ4 与最先进模型（如MP-GAN）相比，所提模型在分布相似性以及关键物理可观测量（如不变质量与EFP）方面表现如何？
• RQ5 该模型的推理速度如何？与现有生成模型相比，其计算效率如何？

主要发现

• 所提模型在JetNet顶夸克数据集上实现了最先进性能，所有指标均优于MP-GAN和基线NF模型。
• 精炼后的变压器模型（TF）的FPND得分为0.08，显著优于基线NF（VNF）的5.61，表明其在复杂相关性建模方面表现更优。
• 该模型每生成一个喷射约需7.6微秒，在A100 GPU上比MP-GAN模型快4.6倍。
• 所有粒子的η、φ和pT的边缘分布，以及归一化后的喷射质量分布，生成数据与真实数据高度一致。
• LSGAN训练目标在训练动力学方面最为稳定，优于Vanilla GAN和带梯度惩罚的WGAN。
• 该模型在不变质量与EFP等全局可观测量上保持高性能，证实其能有效保留关键物理信息。