Economical Jet Taggers -- Equivariant, Slim, and Quantized

本文提出了一种精简、量化且参数缩减版的 L-GATr 喷注分类器，该版本在仅有适度性能下降的情况下实现了数量级的能量成本降低，为大型强子对撞机（LHC）上高效的触发级喷注分类铺平了道路。

要点

想象一下大型强子对撞机（LHC）是一座庞大的、高速运转的粒子工厂。每一秒钟，它都会将质子撞击在一起，产生混乱的碎片喷流。物理学家需要从这些碎片中进行分类，以寻找隐藏在数十亿普通粒子中的特定稀有粒子（例如“顶夸克”）。这种分类过程被称为喷注标记（jet tagging）。

多年来，科学家一直使用复杂的计算机程序（机器学习）来进行这种分类。目前的佼佼者是“Transformer”——一种功能强大的 AI 模型，它们极其精确，但同时也庞大、缓慢且极其耗能。它们就像是一支由巨大的、耗油量极高的卡车组成的车队，为了递送一封信而奔波；它们确实能完成任务，但在数据采集的瞬间（即“触发器”层面），它们过于庞大且昂贵，无法投入使用。

这篇论文提出了一个简单的问题：我们能否将这些巨大的卡车缩减为小巧、省油的轻便摩托车，同时又不失去递送信件的能力？

以下是作者如何通过三种主要策略实现这一目标的：

1. “精简版”（L-GATr-slim）

原始的“L-GATr”模型就像一把瑞士军刀，携带了每一种可能的工具：标量、向量、张量等等。然而，作者意识到，对于大多数粒子物理任务，你其实只需要两种工具：标量（数值）和向量（带有方向的箭头）。

• 类比： 想象一位厨师，仅仅为了做一个简单的三明治，就坚持要使用一间配备了烤箱、搅拌机和混合器的工业级厨房。作者说：“让我们只用一把刀和一个案板吧。”
• 结果： 他们构建了一个“精简版（Slim）”AI，剥离了不必要的工具。它的表现与庞大的版本一样出色，但训练速度更快，占用的内存也更少。这就像是从一辆重型卡车换成了一辆灵巧的跑车，既能完成同样的工作，又更加轻便。

2. “微型版”（超小型标记器）

随后，作者问道：“我们能缩减到多小？”他们尝试将这些 AI 模型缩小到只有玩具车大小（大约 1,000 个参数，而原始模型拥有数百万个）。

• 类比： 想象试图将一整个图书馆的知识装进一张明信片里。通常情况下，你会丢失故事的内容。但作者发现，如果你能够正确地组织信息（使用遵循物理定律的特定“洛伦兹协变”规则），你就可以将核心知识装进极小的空间。
• 结果： 他们发现，对于极小的模型，“LLoCa”架构在缩减层数方面表现最好，而“L-GATr-slim”在缩减层级的宽度方面表现最好。即使在如此微小的规模下，它们仍然优于那些不具备物理感知能力的旧款 AI 模型。

3. “量化版”（低精度数学）

这是最显著的节能手段。标准的 AI 使用非常精确的数学计算（例如测量距离精确到十亿分之一毫米）。作者意识到，对于喷注标记任务，并不需要如此高的精度。通过大幅度舍入数字，依然可以获得不错的效果。

• 类比： 想象你正在仓库里清点苹果。
  • 标准 AI： 你会对每一个苹果进行称重，精确到微克。（准确，但耗时且消耗大量称重能量）。
  • 量化 AI： 你直接按整数进行计数。（快速，几乎不耗能，而且对于了解“有多少个苹果”这个目的来说，这完全足够了）。
• 方法： 他们使用了一种称为 PARQ（分段仿射正则化量化）的技术。可以将它理解为一种智能的舍入规则，在训练过程中轻轻地将数字引导至简单的数值（如 0、1 或 -1），而不是生硬地强制转换。
• 结果： 通过切换到这些“粗略”的数字，他们将运行 AI 的能量消耗降低了 10 倍（一个数量级）。AI 变得异常快速且高效，而准确度仅有微小的下降。

大局观

作者结合了这三种策略——精简架构、微型化尺寸以及量化数学——创造出了“经济型喷注标记器”。

• 为什么这很重要？ 目前，这些强大的 AI 模型体积太大，无法在决定哪些碰撞保留、哪些丢弃的硬件（即“触发器”）上实时运行。
• 目标： 通过将这些模型变得更小、更快、更节能，作者希望最终能将它们直接运行在触发器硬件上。这将使 LHC 能够利用 AI 在瞬间做出决定，筛选哪些粒子碰撞值得保存，从而发现那些此前因为数据被过快丢弃而错失的新物理现象。

简而言之：他们给一个庞大、耗能的 AI 进行了节食、缩小了体积，并教会了它用更少的有效数字进行数学运算，最终打造出一个微小、超高效的引擎，它依然能够识别出宇宙中最重要的粒子。

技术摘要

技术摘要：经济型喷注标记器——等变、精简与量化

问题陈述
现代机器学习（ML）已经改变了大型强子对撞机（LHC）中的喷注标记（jet tagging）任务，洛伦兹等变变换器（Lorentz-equivariant transformers）已成为最先进的架构。然而，领先的模型如 L-GATr 计算成本高昂，需要大量的内存和训练时间。虽然工业界的趋势是扩大网络和数据集的规模，但 LHC 物理学面临着特定的约束，特别是关于触发硬件的内存和延迟要求。目前的喷注分类尚未在触发器中发挥作用，但作者认为它应该发挥作用。核心挑战在于：如何在减少现代等变喷注标记器的大小和计算成本的同时，最大限度地减少性能下降，从而使其能够部署在触发层级。

方法论
本文提出了一个两管齐下的策略来优化资源效率：架构精简与数值量化。

1. L-GATr-slim 架构：
   作者引入了一个精简版的洛伦兹等变变换器（L-GATr）。标准的 L-GATr 使用涉及标量、伪标量、向量、轴矢量和反对称二阶张量的几何代数表示。作者观察到，对于大多数 LHC 应用而言，伪标量、轴矢量和张量是不必要的。因此，L-GATr-slim 将潜在表示限制为仅包含标量和向量。

   • 线性层： 扩展以处理耦合的标量和向量表示，确保向量分量共享一个可学习的标量系数，以维持洛伦兹等变性。
   • 非线性： 通过将两个向量的内积乘以向量输出，对门控线性单元（GLU）进行适配。
   • 归一化： 修改 RMSNorm，使其对向量通道使用闵可夫斯基内积的绝对值。
   • 注意力机制： 构建标量注意力矩阵，采用特定的公式以避免使用全量 L-GATr 中计算昂贵的外积。
   • 实现： 该架构旨在通过 torch.compile 进行编译以提高效率。

2. 量化策略：
   作者应用低精度数据类型和权重量化来进一步降低成本。

   • 数据类型量化： 线性层的输入被量化为 int8（使用零点量化），同时在对精度敏感的操作和反向传播过程中保持 bfloat16。这被应用于 Transformer、ParT、L-GATr-slim 和 LLoCa-Transformer 的隐藏层。
   • 权重量化： 线性权重使用**近端梯度量化（PARQ）**被量化为二进制或三元值。该方法将量化视为一种正则化约束，使用近端算子来更新权重。作者将 PARQ 与直通估计（STE）进行了比较，发现 PARQ 提供了更好的稳定性和性能。
   • 等变性保持： 特别注意确保量化不会破坏洛伦兹等变性。对于 LLoCa，正交化和框架投影保持为全精度（float32），将低精度操作限制在洛伦兹不变量内。对于 L-GATr-slim，全量向量与量化权重相乘，这不会引入额外的对称性破坏。

3. 超小型缩放（Ultra-Mini Scaling）：
   作者研究了通过减少网络的块数（blocks）或宽度（通道数）将这些架构缩减至 1,000 个参数时的性能表现。

关键结果
本研究在三个任务上对所提方法进行了基准测试：顶标记（top tagging）、振幅回归和事件生成。

• 性能 vs. 效率 (L-GATr-slim)：

  • 在 JetClass 数据集（多类喷注标记）上，L-GATr-slim 的性能与全量 L-GATr 和 LLoCa-Transformer 持平（AUC ~0.9885），但将训练时间减少了 6 倍（从 H100 GPU 上的 166 小时降至 27 小时），并将内存消耗减少了 2 倍。
  • 在振幅回归（$Z + 4g$）中，L-GATr-slim 达到了与全量 L-GATr 相同的均方误差（MSE），但所需的训练操作减少了 20 倍，训练时间缩短了一半。
  • 在事件生成（$t\bar{t} + nj$）中，精简架构匹配了全量模型的负对数似然性能。

• 超小型标记器：

  • 当减少块数（深度）时，LLoCa-Transformer 在极小规模（例如 1,000 个参数）下优于 L-GATr-slim。
  • 当保持通道数（宽度）固定（10 个块）并减少通道时，L-GATr-slim 仅使用 2 个向量和 4 个标量通道即可保持超过 1,000 的背景抑制率，优于其他 1,000 参数规模的架构。

• 量化收益：

  • 将输入量化为 int8 并将权重量化为三元值，可在仅有轻微性能损失的情况下，将能耗降低约一个数量级（因子为 10）。
  • LLoCa-Transformer 和 L-GATr-slim 对量化具有鲁棒性，在标准 Transformer 可能性能大幅下降的情况下仍能保持高性能。
  • 在资源最受限的场景下（1 个块，16 维潜在空间，int8），量化后的 LLoCa-Transformer（全局规范化）保留了优于预图（pre-graph）标记器的性能，尽管其背景抑制率较其全尺寸版本下降了 2 倍。

意义与主张
本文声称，这些“经济型”等变变换器代表了迈向高亮度 LHC（HL-LHC）触发级喷注标记的可行路径。通过结合架构精简（移除不必要的几何代数分量）和激进的量化（PARQ 和 int8），作者证明了可以创建拥有约 1,000 个参数且保留洛伦兹等变性这一物理动机优势的标记器。

作者强调，虽然扩大规模是工业界的标准，但 LHC 物理学需要一种“物理感知的降规模”（physics-aware downscaling）方法。结果表明，小型、量化且等变的网络可以在不牺牲使其有效的基本对称性的情况下，部署在资源受限的硬件（如 FPGA）上，这可能为喷注子结构（jet substructure）的实时分析开辟新途径。