[论文解读] Dataset for flavour tagging R&D

I. Ochoa, S. B. Klein+2 more|arXiv (Cornell University)|Aug 20, 2024

Electron and X-Ray Spectroscopy Techniques被引用 5

一句话总结

本文通过去掉标记化、升级解码器，并探索多样的重建任务和集合对集合生成，来改进喷气物理中的掩码粒子建模（MPM），以在喷气数据上对基础模型风格的骨干网络进行预训练。它引入 MPMv2 和集合对集合流匹配，展示了包括 OOD 任务在内的强大下游性能。

ABSTRACT

In this work, we significantly enhance masked particle modeling (MPM), a self-supervised learning scheme for constructing highly expressive representations of unordered sets relevant to developing foundation models for high-energy physics. In MPM, a model is trained to recover the missing elements of a set, a learning objective that requires no labels and can be applied directly to experimental data. We achieve significant performance improvements over previous work on MPM by addressing inefficiencies in the implementation and incorporating a more powerful decoder. We compare several pre-training tasks and introduce new reconstruction methods that utilize conditional generative models without data tokenization or discretization. We show that these new methods outperform the tokenized learning objective from the original MPM on a new test bed for foundation models for jets, which includes using a wide variety of downstream tasks relevant to jet physics, such as classification, secondary vertex finding, and track identification.

研究动机与目标

• 以未标注的喷气数据为基础科学研究提供基础模型风格的预训练动机。
• 开发改进的掩码粒子建模（MPMv2），在不使用 VQVAE 标记化的情况下进行预训练。
• 评估多种重建任务，包括条件生成方法。
• 提出集合对集合流匹配作为喷气预训练的有竞争力范式。

提出的方法

• 修改 MPM，去除重复的掩码令牌并使用完整的 transformer 解码器。
• 仅在被掩码的元素之间提供位置编码，以避免使任务变得平庸。
• 研究五种连续特征重建任务和一种分类任务（粒子ID）。
• 探索替代目标：K-Means 标记化、CNF（条件正则化流）、流匹配（CFM），以及集合对集合流匹配（SSFM）。
• 使用 Delphes 模拟的 JetClass 和 BTag 数据集进行骨干表示的预训练和评估。
• 消融研究比较解码器类型、添加特征和训练设置（表 1）。

Figure 1 : A comparison of the original MPM encoder-decoder setup (left) and the new model configuration (right). The new model includes multiple reconstruction tasks, swaps the MLP decoder for a transformer, and only encodes the reduced set.

实验结果

研究问题

• RQ1 去除标记化并使用更强的解码器是否能使 MPM 相对于原始的 MPMv1 有所提升？
• RQ2 替代重建目标（CNF、K-Means、基于流的方法）是否与 VQVAE 标记化在 MPM 预训练中具有竞争力？
• RQ3 改进后的骨干在分布内、弱监督和分布外的喷气任务上表现如何？
• RQ4 集合对集合流匹配是否可以为无序的喷气粒子集合提供可行的预训练范式？
• RQ5 扩展训练、掩码率调整和附加特征对下游任务有哪些影响？

主要发现

• MPMv2 采用 transformer 解码器且输入基数减少，显著提升了相较于 MPMv1 的分类准确性。
• 切换到 MAE 风格的解码并限制位置编码可获得更好表现且 GPU 内存占用更低。
• 加入冲击参数特征和粒子 ID 进一步提升准确性（示例：回归从 62.2 提升至 80.4；k-means 从 70.2 提升至 83.0，增加后）。
• 完全基于 transformer 的解码器（MAE）在消融中实现了回归 79.2 和 k-means 81.4。
• 更长的训练、更深的解码器以及 40% 的掩码率可获得最佳结果：回归 83.3 和 k-means 84.0。
• 预训练骨干在分布内、弱监督和分布外任务中均优于随机初始化，显示出良好的泛化能力。

Figure 2 : A schematic overview of the SSFM model.