1. 项目概述当Transformer遇上粒子流在大型强子对撞机LHC的每一次质子对撞中探测器会记录下海量的数据点——来自硅像素探测器的带电粒子径迹、来自电磁量能器的光子簇射、来自强子量能器的强子簇射。粒子流算法的核心任务就是将这些零散的“点”重新拼凑成一幅完整的“粒子级”事件图像告诉你这次对撞究竟产生了哪些粒子它们的能量、动量、种类分别是什么。这听起来像是一个复杂的拼图游戏而传统的玩法也就是CMS实验过去十几年一直使用的基于规则的粒子流算法就像一位经验丰富的老师傅凭借一套固定的“如果…那么…”规则比如如果一个径迹的末端指向一个簇射那么它们很可能来自同一个粒子来拼接。这套方法很稳健但面对LHC运行三期及未来高亮度升级带来的更高对撞频率和更复杂的“堆积”背景即一次质子束团对撞中夹杂的数十次额外软对撞规则系统的局限性开始显现它难以捕捉探测器响应中那些非线性的、微妙的关联模式。这正是机器学习特别是Transformer模型大显身手的地方。我这次要分享的就是我们在CMS实验中如何用基于Transformer和FlashAttention的机器学习粒子流算法来尝试“教”计算机自己学会玩这个拼图游戏。简单来说我们不再手动编写规则而是给模型输入所有径迹和簇射的低层特征如动量、位置、能量沉积形状让它直接输出重建出的粒子列表包括种类、是否来自堆积、四动量。这背后的直觉是模型能从海量的模拟数据中发现那些人类难以用简单规则描述的复杂模式。更关键的是通过集成FlashAttention这一针对Transformer的优化内核我们在保持甚至提升物理性能的前提下将单事件处理时间缩短了近一半。这对于未来每秒需要处理数百万次事件的实时数据处理系统而言意义重大。2. 核心思路从规则驱动到数据驱动2.1 传统粒子流算法的瓶颈传统的CMS粒子流算法是一个精心设计的、多步骤的规则引擎。它的工作流程可以概括为首先分别重建出径迹带电粒子和簇射中性粒子或未完全被追踪的带电粒子然后通过一系列基于几何接近性、能量匹配等启发式规则将径迹和簇射关联起来形成“粒子流候选者”最后进行堆积抑制如使用PUPPI算法剔除那些来自非主对撞顶点即堆积的粒子贡献。这套系统的优势在于可解释性强每一步都有清晰的物理图像对应。但其瓶颈也显而易见规则复杂性为了处理各种边界情况如重叠的簇射、转换电子等规则集变得异常复杂难以维护和优化。信息利用不充分规则通常基于少数几个高阶特征如能量、方向而探测器原始数据中大量低层特征如簇射形状的细节、径迹拟合的协方差矩阵被丢弃或仅被简单使用。全局关联能力弱基于局部最近邻的关联方式难以处理全局性的事件拓扑结构例如在密集喷注核心区多个粒子靠得非常近时容易产生误关联。2.2 机器学习粒子流的基本范式机器学习粒子流的核心思想是端到端学习。我们不再分步进行关联、分类和能量回归而是构建一个模型其输入是所有重建好的径迹和簇射的低层特征向量输出则是直接预测的粒子列表。每个预测粒子包含其粒子ID、四动量以及一个关键的“是否来自堆积”的标签。这实际上将一个复杂的、多步骤的推理问题转变为一个集合到集合的预测问题。输入是一个大小可变的径迹和簇射集合输出也是一个大小可变的粒子集合。Transformer架构因其强大的序列建模和全局注意力机制非常适合处理这种不定长、元素间存在复杂关系的集合数据。注意这里有一个关键概念叫“目标定义”。模型在训练时需要知道什么是“正确”的答案。我们使用Geant4全探测器模拟来生成“目标粒子”。但需要注意的是模拟本身有阈值和有限粒度因此“目标粒子”与理论上的“真实粒子”之间存在不可避免的“模糊”。此外在包含堆积的样本中目标粒子既包含硬散射过程的粒子也包含堆积粒子这要求模型必须学会区分它们。2.3 为何选择Transformer与FlashAttention最初版本的MLPF使用的是图神经网络。GNN通过边来传递信息天然适合处理粒子探测器这种图结构的数据径迹和簇射作为节点它们之间的空间关系作为边。然而GNN的信息传递通常局限于局部邻域要获取全局信息需要多层堆叠计算开销大。Transformer则通过自注意力机制允许序列中的任何一个元素直接与所有其他元素交互一步到位地捕获全局上下文。这对于粒子重建至关重要因为一个远离的μ子径迹的信息可能有助于判断某个孤立的量能器沉积是否来自中性强子。Transformer的这种全局建模能力理论上能更有效地处理密集环境下的关联问题。但是标准Transformer的自注意力计算复杂度是输入序列长度的平方级。在高能物理事件中径迹和簇射的数量动辄成千上万这会导致巨大的计算和内存开销。FlashAttention的引入正是为了解决这个问题。它是一种IO感知的精确注意力算法通过巧妙地重组计算顺序将中间结果保留在SRAM高速缓存中大幅减少了对慢速HBM内存的访问次数。这使得我们能够用Transformer处理以前无法承受的长序列将理论优势转化为实际可用的模型。3. 模型架构与训练策略深度解析3.1 输入特征工程给模型一双“慧眼”模型的性能上限很大程度上由输入特征决定。我们为径迹和簇射精心设计了一套特征集旨在尽可能全面地表征探测器的响应。对于径迹我们提供的不仅仅是动量pT, η, φ和电荷。为了帮助模型理解测量的不确定性我们加入了动量、方向角的不确定度。考虑到粒子从产生点到量能器的飞行过程我们提供了径迹在参考点、以及外推到电磁量能器和强子量能器入口处的坐标和方向角。此外径迹的“质量”信息也很重要比如径迹上重建出的击中点数越多通常意味着径迹质量越高、径迹的顶点坐标用于区分主顶点和次级顶点。一个特别的技巧是使用λ π/2 - θ来代替极角θ因为在高能物理中粒子在η赝快度方向分布更均匀而λ与η在相对论极限下是线性相关的这样的变换有时能使优化更稳定。对于簇射特征则更加丰富。除了基本的横向能量ET、位置η, φ, x, y, z和时间我们区分了电磁量能器和强子量能器的能量沉积E_ECAL, E_HCAL。能量校正值及其不确定度是关键特征反映了重建算法的置信度。簇射的“形状”特征尤为重要我们计算了簇射中所有击中点在x, y, z方向上的加权位置标准差作为簇射大小的度量。更进一步我们计算了簇射在η和φ方向上、随量能器深度层数变化的形状参数。这相当于给了模型一个“三维视觉”让它能区分例如一个窄而深的电磁簇射可能来自电子/光子和一个宽而浅的强子簇射。实操心得特征标准化至关重要。所有输入特征在送入模型前都经过减去均值、除以标准差的处理。这个均值和标准差是从训练数据集中计算得到的。对于像pT、E这类跨度很大的特征通常先取对数再进行标准化可以改善训练稳定性。确保训练集和验证集使用完全相同的标准化参数否则会导致模型性能的意外下降。3.2 损失函数设计多任务学习的艺术MLPF模型需要同时完成四个任务判断一个输入元素是否对应一个真实粒子存在性分类、识别该粒子的类型PID分类、判断该粒子是否来自堆积PU分类、以及精确回归该粒子的四动量。因此我们的损失函数是这四个子损失的加权和。存在性分类损失这是一个二分类问题是粒子/不是粒子使用标准的二元交叉熵损失。一个常见的陷阱是正负样本极度不平衡非粒子背景远多于真实粒子。我们通过在每个训练批次内对负样本进行随机下采样来解决确保正负样本比例大致平衡。PID分类损失这是一个多类分类问题如电子、光子、带电强子、中性强子等。我们使用了Focal Loss。Focal Loss是交叉熵损失的一个变体它通过降低易分类样本的权重让模型更专注于难分类的样本。在高能物理中某些粒子类型如μ子可能很容易识别而另一些如中子和K_L介子则很难区分。Focal Loss能自动调整这种关注度显著提升对稀有或难例粒子的分类性能。堆积识别损失同样是一个二分类问题使用二元交叉熵。这里的关键是训练数据的标签。我们利用模拟中的顶点信息为每个目标粒子标记其是否来源于主相互作用顶点。模型需要从粒子特征中学习堆积粒子的共性例如通常来自更远的顶点动量较低方向更各向同性。四动量回归损失使用均方误差损失来回归粒子的pT, η, sinφ, cosφ, E。这里有一个重要的技巧pT加权。我们在pT和E的回归损失项前乘以一个√pT的因子。这是因为高pT的粒子在物理分析中通常更重要例如产生希格斯玻色子的衰变产物它们的能量测量误差对整体物理性能如喷注质量、不变质量分辨率影响更大。通过这个加权我们引导模型在训练时更关注高能区域的回归精度。3.3 训练流程与技巧模型使用PyTorch框架实现在单块NVIDIA A100 80GB GPU上训练了8个epoch耗时约260小时。我们使用了AdamW优化器并采用了带热重启的余弦退火学习率调度。初始学习率设置为3e-4每个epoch结束时验证损失并保存性能最佳的模型。注意事项监控所有四个损失分量随训练epoch的变化曲线至关重要。理想情况下它们应该同步平稳下降并在验证集上逐渐收敛。如果某个损失特别是回归损失出现剧烈震荡或很早就停止下降可能需要调整该损失的权重或检查特征输入。在我们的训练中所有损失都平滑收敛没有出现过拟合迹象训练损失持续下降而验证损失上升。大约在6个epoch后性能提升变得微乎其微因此选择8个epoch作为最终模型。最终的模型包含约400万个可学习参数。这个规模对于Transformer模型来说相对适中确保了在推理时的效率同时又提供了足够的表达能力来捕捉粒子重建中的复杂模式。4. 物理性能验证与经典算法的正面较量模型训练好后一切都要用物理性能来说话。我们在模拟样本上进行了全面的测试并与传统的PF算法进行了对比。4.1 单粒子性能效率与假阳性我们首先在无堆积的ttbar样本中检查模型对单粒子的重建能力。我们将重建出的粒子与模拟中的“真实”粒子进行匹配在ΔR 0.1的范围内然后计算重建效率和假阳性率。重建效率定义为被正确重建的“真实”粒子数 / 总的“真实”粒子数。假阳性率定义为重建出的、无法与任何“真实”粒子匹配的粒子数 / 总的重建粒子数。从性能图中对应原文图2我们可以看到几个关键结论整体一致性MLPF重建出的粒子在pT和η谱分布上与真实粒子高度一致说明模型没有引入明显的整体性偏差。中性粒子的优势对于光子和中性强子MLPF显示出比传统PF更高的重建效率尤其是在中等pT区域10-100 GeV。这是一个重要的突破因为中性粒子仅通过量能器测量传统规则方法在区分重叠簇射或估计非补偿量能器响应方面存在困难而MLPF能够从簇射形状的复杂特征中学习到更有效的模式。假阳性率可控在提升效率的同时MLPF的假阳性率与传统PF基本持平甚至在部分区域略有降低。这表明模型没有以大量“虚报”为代价来换取效率提升其学习到的关联是可靠的。4.2 喷注与丢失横动量性能集成效应的考验单个粒子重建得好最终要服务于物理对象如喷注和丢失横动量。我们使用anti-kT算法R0.4将PF或MLPF重建出的粒子聚类成喷注。喷注能量响应与分辨率我们将重建喷注与生成器级别的“真实”喷注匹配计算响应重建pT / 真实pT的平均值和分辨率响应的标准差。在包含55-75个堆积事件的ttbar和QCD样本中MLPF的表现与传统PFPUPPI算法不相上下。这里有一个关键设计MLPF如何做堆积抑制传统PF使用独立的PUPPI算法在粒子聚类前对每个粒子赋予一个权重来抑制堆积贡献。MLPF则将堆积识别集成到了模型中对于带电粒子如果它有匹配的径迹我们沿用传统方法利用径迹的顶点信息这是最可靠的堆积判别器。对于中性粒子和没有径迹匹配的带电粒子MLPF使用其内部输出的“PU分数”作为权重。也就是说模型在重建粒子的同时已经判断了它来自堆积的可能性并可以直接输出一个用于抑制的权重。这种“一体化”设计更加简洁优雅。丢失横动量这是许多新物理寻找的关键变量。我们比较了使用MLPF粒子与使用传统PF粒子计算的丢失横动量分布。在ttbar和QCD样本中两者在分布形状和分辨率上都非常接近。这说明MLPF在重建可见粒子的能量和方向上是准确的从而保证了由其衍生的丢失横动量也是可靠的。实操心得喷注性能的对比必须考虑能量校正。无论是PF还是MLPF重建出的喷注能量都需要经过一系列基于模拟数据的校正如绝对能量标度、堆积能量密度偏移等才能与数据比较。在比较MLPF和PF时我们通常比较“原始”重建性能即校正前或者对两者应用相同的校正流程后的性能。确保比较的公平性至关重要。5. 在真实数据上的试运行模拟世界再完美也需要在真实数据上接受检验。我们在2024年采集的一小部分CMS Run 3数据上对MLPF进行了试运行。我们选择了双喷注触发的事例并应用了离线的双喷注选择。我们比较了MLPF和传统PF算法在几个关键观测量上的分布丢失横动量两种算法出的未修正的pTmiss分布高度吻合。双喷注不对称性定义为 (pT1 - pT2) / (pT1 pT2)这个变量对喷注能量标度非常敏感。MLPF和PF计算出的不对称性分布也显示出良好的一致性。这些在真实数据上的初步一致性是MLPF迈向实际应用的关键一步。它表明模型在模拟中学到的知识能够有效地迁移到真实探测器的数据上没有出现严重的系统性偏差。当然这只是一个开始后续需要在更多数据、更多物理过程上进行验证。6. 工程集成与性能优化从模型到生产系统一个再好的模型如果不能高效地集成到现有的、庞大的实验软件框架中也只是纸上谈兵。CMS的软件框架CMSSW是一个复杂的、模块化的C系统。6.1 集成路径ONNX与onnxruntime我们的集成策略是模型导出将训练好的PyTorch模型导出为ONNX格式。ONNX是一种开放的模型交换格式得到了多种推理引擎的支持。运行时集成在CMSSW中我们通过onnxruntimeC API来加载和运行这个ONNX模型。onnxruntime针对不同硬件提供了优化的执行提供程序。无缝替换MLPF模块被设计为直接替换原有PF算法中的核心重建步骤。它接收标准的径迹和簇射集合作为输入输出标准的粒子流候选者集合。这样下游的所有模块喷注聚类、堆积抑制后处理、b标签等都无需任何修改可以透明地使用MLPF的输出。6.2 FlashAttention的威力GPU推理加速为了发挥Transformer模型特别是集成FlashAttention的优势我们将推理部署在GPU上。我们使用了onnxruntime的CUDA执行提供程序并确保FlashAttention内核被正确调用。性能基准测试结果令人振奋对应原文图5。我们在一个专用的测试机上运行6个并行的CMSSW作业每个作业使用8个CPU线程和1/7的A100 GPU资源。对比传统PF算法完全在CPU上运行传统PF平均每事件处理时间约为90毫秒。MLPF (GPU)平均每事件处理时间约为40毫秒。这意味着MLPF带来了超过2倍的加速这不仅仅是算法效率的提升更是计算范式的转变——将计算密集型任务从CPU卸载到GPU。考虑到未来高亮度LHC条件下事件率的飙升这种加速对于满足实时数据处理触发系统的严苛时间预算具有战略意义。注意事项GPU推理的延迟不仅包括模型计算本身还包括数据从CPU内存到GPU显存的传输时间。在我们的实现中需要将径迹和簇射的特征向量从CMSSW的C数据结构复制到连续的GPU内存缓冲区。这部分开销必须被最小化。我们通过批量处理事件尽管在在线触发中批量大小通常为1和优化内存布局来降低传输开销。此外需要确保onnxruntime的会话配置如线程数、内存分配策略针对我们的特定工作负载进行了优化。6.3 稳定性与可重复性在生产环境中算法的稳定性与可重复性至关重要。我们测试了模型在不同硬件不同型号的GPU上、以及多次重复运行同一批数据时的输出。结果表明MLPF的输出是确定性的在不同硬件上的差异完全在浮点数精度的预期范围内。这确保了物理分析结果不会因为推理环境的不同而产生不可控的变化。7. 挑战、应对与未来展望7.1 实际部署中的挑战输入特征的实时计算MLPF依赖大量低层特征其中一些如簇射的形状参数在标准重建链中可能不是默认计算的。我们需要确保这些特征能在在线触发系统和离线重建中以可接受的资源开销被实时计算出来。模型校准与系统误差机器学习模型可能存在校准问题即其预测的置信度不能真实反映误差。我们需要研究如何对模型输出的四动量进行校准并评估其引入的系统性不确定性这需要大量的模拟和数据-模拟比较研究。对模拟的依赖性MLPF完全在模拟数据上训练其性能依赖于模拟的保真度。任何模拟与数据之间的差异“模拟失配”都可能被模型学习并放大。持续的、基于数据的模型验证和可能的微调是必要的。7.2 常见问题排查思路问题模型在真实数据上输出的pTmiss整体偏高。排查首先检查输入特征特别是量能器簇射的能量校正值是否被正确计算和应用。然后在模拟中注入已知的“模拟失配”如量能器能量响应偏移几个百分点观察模型输出如何变化。这有助于定位是某个探测器子系统的问题还是模型整体的响应偏差。问题GPU推理时出现内存不足错误。排查检查单个事件输入特征的大小。Transformer的注意力内存开销与序列长度的平方成正比。如果事件非常密集如高堆积的重离子碰撞可能需要动态裁剪输入序列例如只保留pT最高的前N个径迹和簇射或者研究更高效的内存管理策略。同时确认onnxruntime配置没有内存泄漏。问题对某类稀有粒子如K_L重建效率极低。排查回顾训练数据中该类粒子的样本数量是否充足。检查Focal Loss中该类别的权重因子。可以尝试使用加权的采样策略在训练时过采样稀有粒子类别或者为该类别的分类损失赋予更高的权重。7.3 未来发展方向这项工作只是一个起点。未来的方向非常广阔端到端优化目前的MLPF仍以传统重建的径迹和簇射作为输入。更激进的思路是开发“端到端”的重建直接从原始探测器信号如像素击中、量能器塔能量开始用单个模型同时完成击中聚类、径迹拟合、粒子重建等所有步骤。这将最大化利用信息但模型复杂度和数据需求也将呈指数级增长。面向HL-LHC的优化高亮度LHC将带来前所未有的堆积水平平均200个堆积事件。需要专门针对这种极端环境训练和优化模型可能需要在架构上创新例如引入更显式的顶点信息处理模块。推理极致优化探索模型量化将FP32精度降低为INT8或FP16、剪枝、知识蒸馏等技术在尽可能保持精度的前提下进一步压缩模型大小、降低延迟以满足一级触发系统的纳秒级时间要求。可解释性与不确定性估计开发工具来理解模型为何做出某个特定的关联或分类决策。同时让模型不仅输出预测值还能输出其预测的不确定性这对于后续的物理分析至关重要。从基于规则到基于学习的粒子流算法我们正在经历高能物理实验数据处理范式的一次深刻变革。Transformer与FlashAttention的结合让我们在精度和速度的平衡木上找到了一个新的支点。虽然前路仍有诸多工程与物理上的挑战但初步的结果已经证明这条道路充满希望它或许将重塑我们“看见”对撞机中那些稍纵即逝的基本粒子的方式。
