[论文解读] Zero-Permutation Jet-Parton Assignment using a Self-Attention Network

J. S. H. Lee, Inkyu Park+2 more|arXiv (Cornell University)|Dec 7, 2020

Particle physics theoretical and experimental studies参考文献 38被引用 5

一句话总结

本文提出 SaJa，一种用于高能物理中零排列喷射子-部分子关联的自注意力神经网络，可直接预测喷射子-部分子关联，无需组合优化。SaJa 在完全强子化的 top-antitop 事件中表现出优于传统基于似然的方法（如 KLFitter）的性能，通过批量处理实现高达 100 倍的推理速度提升，展现出更高的准确率与效率。

ABSTRACT

In high-energy particle physics events, it can be advantageous to find the jets associated with the decays of intermediate states, for example, the three jets produced by the hadronic decay of the top quark. Typically, a goodness-of-association measure, such as a $χ^2$ related to the mass of the associated jets, is constructed, and the best jet combination is found by optimizing this measure. As this process suffers from a combinatorial explosion with the number of jets, the number of permutations is limited by using only the $n$ highest $p_T$ jets. The self-attention block is a neural network unit used for the neural machine translation problem, which can highlight relationships between any number of inputs in a single iteration without permutations. In this paper, we introduce the Self-Attention for Jet Assignment (SaJa) network. SaJa can take any number of jets for input and outputs probabilities of jet-parton assignment for all jets in a single step. We apply SaJa to find jet-parton assignments of fully-hadronic $t\bar{t}$ events to evaluate the performance. We show that SaJa achieves better performance than a likelihood-based approach.

研究动机与目标

• 为解决如完全强子化 t\bar{t} 衰变等高多重性事例中喷射子-部分子关联的组合爆炸问题。
• 开发一种排列不变的方法，避免对喷射子进行穷举排列扫描。
• 在传统基于似然的方法（如 KLFitter）的基础上进一步提升喷射子-部分子关联的准确率。
• 通过 GPU 硬件上的批量处理实现实现高效、可扩展的推理。
• 通过注意力权重可视化提供模型可解释性，为物理洞察提供支持。

提出的方法

• SaJa 网络使用自注意力机制在单次前向传播中建模喷射子之间的依赖关系，消除了排列枚举的需求。
• 将喷射子-部分子关联问题建模为每个喷射子的多分类问题，共五种类别：b1、W1、b2、W2 和其他喷射子。
• 模型采用可学习的位置编码和多头自注意力模块，以捕捉顶夸克内部及顶夸克之间的喷射子关联。
• 目标函数使用交叉熵损失，对独热编码的真实标签进行优化，通过对称的注意力计算确保排列不变性。
• 网络在蒙特卡洛生成的 t\bar{t} 事例上进行训练，喷射子通过 anti-kT 算法重建，并引入喷射子形状特征以提升性能。
• 推理通过小批量进行，可在 GPU 硬件上实现高吞吐量，延迟按不同批量大小测量每个事例。

实验结果

研究问题

• RQ1 自注意力网络能否在喷射子-部分子关联性能上超越传统基于似然的方法（如 KLFitter）？
• RQ2 SaJa 网络在不同事件生成器（使用不同部分子喷射模型）下是否保持高性能？
• RQ3 引入喷射子形状信息在多大程度上提升了关联准确率与泛化能力？
• RQ4 SaJa 的推理延迟与 KLFitter 等顺序似然拟合方法相比如何？
• RQ5 SaJa 中的注意力权重能否提供对喷射子-部分子拓扑结构与关系的可解释性洞察？

主要发现

• SaJa 在复杂完全强子化 t\bar{t} 事例（最多 12 个喷射子）中实现了高于 KLFitter 的喷射子-部分子关联准确率。
• 引入喷射子形状信息可提升性能，但不同事件生成器之间的结果略有差异，这与部分子喷射模型的差异有关。
• 与不包含形状特征的基线相比，加入喷射子形状信息的 SaJa 在正确关联率上提升了 1.5%。
• 即使在批量大小为 1 时，SaJa 的推理时间也快于 KLFitter 的顺序处理，在 GPU 上批量大小为 1000 时实现 100 倍加速。
• 注意力权重可视化显示，早期块学习到顶夸克内部的喷射子关联，而后期块提取与关联相关的关键特征，证实了模型的可解释性。
• 当不包含喷射子形状特征时，模型对事件生成器差异具有鲁棒性，表明其在不同模拟模型间具备良好的泛化能力。