[论文解读] Particle Transformer for Jet Tagging

H. Qu, C. Li+1 more|arXiv (Cornell University)|Feb 8, 2022

Astrophysics and Cosmic Phenomena被引用 45

一句话总结

本论文介绍 JetClass，一个大型喷射标记数据集，以及 ParT，一种基于 Transformer 的模型，具备成对交互感知的注意力，优于现有方法在 JetClass 上的表现，并在微调后提升基准。

ABSTRACT

Jet tagging is a critical yet challenging classification task in particle physics. While deep learning has transformed jet tagging and significantly improved performance, the lack of a large-scale public dataset impedes further enhancement. In this work, we present JetClass, a new comprehensive dataset for jet tagging. The JetClass dataset consists of 100 M jets, about two orders of magnitude larger than existing public datasets. A total of 10 types of jets are simulated, including several types unexplored for tagging so far. Based on the large dataset, we propose a new Transformer-based architecture for jet tagging, called Particle Transformer (ParT). By incorporating pairwise particle interactions in the attention mechanism, ParT achieves higher tagging performance than a plain Transformer and surpasses the previous state-of-the-art, ParticleNet, by a large margin. The pre-trained ParT models, once fine-tuned, also substantially enhance the performance on two widely adopted jet tagging benchmarks. The dataset, code and models are publicly available at https://github.com/jet-universe/particle_transformer.

研究动机与目标

• 动机：需要一个大规模公开喷射标记数据集，以推动对撞机物理中的深度学习方法的发展。
• 提出一种基于 Transformer 的架构（ParT），在注意力中纳入成对粒子相互作用，以提高标记性能。
• 提供对 JetClass 的全面评估，并通过微调在既有喷射标记基准上展示迁移收益。

提出的方法

• 以 1 亿喷射覆盖 10 种喷射类型定义 JetClass，其中包括此前未探索的类别。
• 引入 ParT，一种 Transformer 变体，使用两组输入：每粒子特征和每对相互作用特征。
• 用一个成对交互偏置 U 添加到 softmax 之前的注意力权重中，增强粒子自注意力（P-MHA）。
• 实现类注意力机制，以生成用于最终分类的全局喷射表示。
• 对来自粒子四维向量的成对交互特征进行编码，通过 8 维交互嵌入进行变换。
• 端到端训练 ParT，并与 PFN、P-CNN、ParticleNet 进行比较；同时评估在其他数据集上进行预训练再微调的效果。

实验结果

研究问题

• RQ1 基于 Transformer 的架构结合成对粒子相互作用信息，是否在大规模喷射数据集上超过以往的喷射标记模型？
• RQ2 注意力中的成对交互偏置（P-MHA）对总体性能的贡献是什么？
• RQ3 在 JetClass 上对 ParT 进行预训练，是否通过微调提升在已有喷射标记基准上的表现？

主要发现

• ParT 在 JetClass 上的准确率和 AUC 均高于 ParticleNet 等基线。
• 通过 P-MHA 将成对粒子相互作用引入，较普通 Transformer 有显著提升。
• 在 JetClass 上预训练并在其他喷射标记基准上微调时，ParT 提供了显著改进。
• 在多种信号类型下，ParT 相对于 ParticleNet 显著提升背景拒绝，某些喷射类型甚至有巨大增益。
• 消融研究表明用标准 MHA 替换 P-MHA 会降低准确率和背景拒绝，证实成对交互建模的价值。