Locality-Sensitive Hashing-Based Efficient Point Transformer for Charged Particle Reconstruction¶

会议: NeurIPS 2025
arXiv: 2510.07594
代码: 有
领域: 3D 视觉 / 粒子物理
关键词: Point Transformer, Locality-Sensitive Hashing, 粒子追踪, 端到端学习

一句话总结¶

通过将 LSH 与 Point Transformer 结合，提出 HEPTv2 实现粒子轨迹重建的端到端学习，消除了 DBScan 聚类后处理瓶颈，在保持竞争性追踪效率的同时实现 28.9 倍加速。

领域现状：高能物理 LHC 实验中粒子轨迹重建是最计算密集的任务，传统 Kalman Filter 在高 pile-up 条件下性能下降。

现有痛点：GNN 虽然性能好，但存在图构造成本高 \(O(n^2)\)、不规则邻域聚合导致硬件低效、随机访存伤害缓存利用率等三大问题。HEPT 虽引入 LSH 实现线性复杂度，但需要额外的 DBScan 聚类占运行时间 90%。

核心矛盾：快速编码 vs 完整任务（需要追踪分配）；表达力 vs 硬件友好。

核心 idea：扩展 HEPT 为 HEPTv2，加入轻量级查询基础 Transformer 解码器，直接预测轨迹分配。

三阶段管线：(1) 度量学习（LSH 编码）——将检测器命中散列到 1D 序列；(2) 实例解码——查询基础解码器精化轨迹假设；(3) 分配与后处理——将命中分配至最可能的轨迹。

LSH 编码器
- 功能：E2LSH 方案将相近命中映射到相同 1D 桶，实现块对角 Attention
- 核心思路：OR 构造用 \(m_1\) 个独立哈希表，AND 构造每表级联 \(m_2\) 个哈希函数，\(h_j(x) = \lfloor(a_j \cdot x + b_j)/r\rfloor\)
- 设计动机：规则访存模式，GPU 友好，同桶内自注意力成本 \(O(1)\)
端到端轨迹分配解码器
- 功能：固定 3000 个可学习轨迹查询，通过自注意力和交叉注意力预测轨迹分配
- 核心思路：二元命中分类器判断是否属于轨迹；查询基础解码器（自注意力→交叉注意力→前馈层）；输出每查询置信度和密集掩码 logits
- 设计动机：消除 DBScan 后处理，仅增加 17% 计算开销（4ms），而 DBScan 需额外 1401ms
联合损失函数
- 功能：5 项损失联合优化
- 核心思路：\(\mathcal{L} = \lambda_{nce}\mathcal{L}_{NCE} + \lambda_{clf}\mathcal{L}_{CLF} + \lambda_{ce}\mathcal{L}_{CE} + \lambda_{mask}\mathcal{L}_{BCE} + \lambda_{dice}\mathcal{L}_{Dice}\)
- 设计动机：InfoNCE 对比损失聚集同粒子命中 + 分类损失 + 掩码损失，覆盖从嵌入到分配的完整流程

课程学习：早期优先训练干净可轨迹的命中，逐步引入难样品和低动量命中。

方法	追踪效率	假率	推理时间(ms)	相对加速
Exa.TrkX (GNN SOTA)	0.994	0.002	~800	基准
HEPT + DBScan	0.923	0.070	1425	0.56x
HEPTv2	0.993	0.113	27.7	28.9x