[论文解读] Solving Key Challenges in Collider Physics with Foundation Models

V. M. Mikuni, Benjamin Nachman|arXiv (Cornell University)|Apr 24, 2024

Computational Physics and Python Applications被引用 5

一句话总结

OmniLearn 训练一个大型有监督的基础模型，适用于喷气系统，一旦学成即可在跨数据集、探测器和碰撞系统的多项下游喷气物理任务中实现改进、加速和多任务能力，而无需针对具体任务的预训练。

ABSTRACT

Foundation Models are neural networks that are capable of simultaneously solving many problems. Large Language Foundation Models like ChatGPT have revolutionized many aspects of daily life, but their impact for science is not yet clear. In this paper, we use a new Foundation Model for hadronic jets to solve three key challenges in collider physics. In particular, we show how experiments can (1) save significant computing power when developing reconstruction algorithms, (2) perform a complete uncertainty quantification for high-dimensional measurements, and (3) search for new physics with model agnostic methods using low-level inputs. In each case, there are significant computational or methodological challenges with current methods that limit the science potential of deep learning algorithms. By solving each problem, we take jet Foundation Models beyond proof-of-principle studies and into the toolkit of practitioners.

研究动机与目标

• 通过利用对喷气的大规模有监督训练来学习通用、可迁移的喷气表示，提出喷气物理的基础模型方法。
• 证明单一的 OmniLearn 主干在其原始训练任务之外的多个下游任务上提高或改进。
• 展示在喷气类型、探测器仿真和碰撞系统上的泛化能力，并实现条件生成与弱监督能力。
• 强调训练效率提升及在真实对撞机分析中的实用性。

提出的方法

• 将喷气表示为点云并构建点-边 Transformer（PET）骨干网络。
• 在多分类喷气分类任务（10 类）上训练共享表示，输入采用时间条件扩散。
• 附加任务特定头：分类头和生成头，以及用于下游任务的扩散条件机制。
• 使用特征丢失和受 classifier-free 指引启发的层丢失，提升跨数据集与条件标签的鲁棒性。
• 结合多项损失项：分类的交叉熵、用于扩散风格生成的基于速度的损失，以及用于正则化学习的涂抹项。
• 在来自 JetClass 的 1 亿喷气上进行训练，采用 128 GPU 设置，使用扩散时间条件嵌入和 LayerScale 以实现稳定性。

Figure 1: Neural network architecture used to train OmniLearn . The main neural network blocks of the architecture are shown in the further left with detailed architecture design shown for each block in the right. See the text for more details.

实验结果

研究问题

• RQ1 一个有监督的大规模喷气基础模型是否能够在不同喷气类型、探测器仿真和碰撞系统之间实现泛化？
• RQ2 在一个任务上对单一 OmniLearn 主干进行预训练，是否会加速或提升对其他喷气物理任务（分类、生成、似然比估计、异常检测）的性能？
• RQ3 将 OmniLearn 转移到具有不同条件需求的新下游数据集时，训练效率提升和性能变化是多少？
• RQ4 在数据集之间的条件生成与重加权/展开（unfolding）情境中，OmniLearn 的表现如何？
• RQ5 OmniLearn 是否能够在喷气物理中实现稳健的弱监督与谐振异常检测？

主要发现

• OmniLearn 在顶夸克标记和类/夸克数据集的分类器从零开始训练的对比中表现更佳，达到更高的 AUC 和准确率。
• 在 CMS Open Data 的类/夸克标记任务中，OmniLearn 与 PET 分类器相匹配或略有超越，同时收敛速度更快。
• 对于 ATLAS 的顶夸标记，OmniLearn 实现了最先进的性能，且在用较少样本微调时也能达到或超过在完整数据集上训练的模型。
• OmniLearn 在多数据集上收敛更快，训练时间通常减少 2–3.5 倍。
• 在 DIS（电子-质子喷气）情境下，OmniLearn 的性能与非 OmniLearn 基线相当但收敛更快，体现跨系统泛化能力。
• OmniLearn 支持条件生成研究（JetNet），且在共享骨干上实现了有竞争力的粒子级与喷气级生成。

Figure 2: Validation loss curves obtained in the top quark tagging (left) and quark/ gluon (rights) datasets. The OmniLearn validation loss is compared with the PET classifier trained from scratch.