AllShowers：基于深度学习的多粒子探测器模拟框架

1. 粒子探测器模拟的挑战与机遇

高能物理实验的核心任务之一，就是精确模拟粒子与探测器材料的相互作用过程。传统方法依赖于Geant4等蒙特卡洛模拟工具包，它们通过逐步跟踪每个粒子的运动轨迹和相互作用过程来构建完整的物理图像。这种"第一性原理"方法虽然精度高，但计算代价极其昂贵——单个大型强子对撞机（LHC）事件的完整模拟可能需要数小时CPU时间。

随着未来对撞机计划（如国际线性对撞机ILC）的推进，探测器将具备更高的粒度和更复杂的结构。以ILD探测器为例，其电磁量能器（ECAL）采用30层钨吸收体和硅传感器，单元尺寸仅5×5 mm²；强子量能器（HCAL）则有48层不锈钢吸收体和塑料闪烁体，单元尺寸3×3 cm²。这种高粒度设计虽然提升了测量精度，却使得传统模拟方法的计算成本呈指数级增长。

关键痛点：在ILC等未来实验中，Geant4模拟单个事例的计算成本预计将比LHC时代高出100-1000倍，完全依赖传统方法将无法满足实验所需的统计量。

深度学习生成模型为解决这一困境提供了新思路。近年来，生成对抗网络（GAN）、变分自编码器（VAE）和标准化流（NF）等模型已成功应用于量能器模拟。但这些方案存在三个根本局限：

1. 多数模型仅针对单一粒子类型（如光子或π介子）设计
2. 通常只模拟ECAL或HCAL中的一种
3. 难以处理高粒度探测器产生的超大点云（>6000个能量沉积点）

2. AllShowers架构设计精要

2.1 整体框架创新

AllShowers采用分而治之的双模块架构（图1），巧妙化解了多粒子联合模拟的复杂度：

PointCountFM模块
负责预测每层的能量沉积点数，输入为入射粒子信息（类型、能量、角度），输出78个整数（对应ECAL+HCAL的78层）。其核心创新在于：

• 取消去量化噪声：传统流模型需要对离散数据添加噪声以便训练，而采用条件流匹配（Conditional Flow Matching）方法后，直接建模离散分布效果更优
• 动态角度处理：将入射角度编码为单位球面上的向量，与固定角度方案相比，可适应更广的物理场景

CNF-Transformer模块
根据PointCountFM指定的点数，生成每个点的(x,y,能量)分布。其突破性设计包括：

• 层感知的位置编码：为每层学习独立的64维嵌入向量，自动捕捉材料预算、距探测器表面距离等物理特征
• 受限注意力机制：仅允许相邻±2层内的点相互关注，将计算复杂度从O(N²)降至O(N)
• 层间最优传输映射：通过解耦各层的OT问题，显著降低ODE求解器的积分步数

# 简化版CNF-Transformer的PyTorch伪代码 class CNFTransformer(nn.Module): def __init__(self): self.layer_embed = nn.Embedding(78, 64) # 层嵌入 self.particle_embed = nn.Embedding(12, 64) # 粒子类型嵌入 self.encoder_blocks = nn.ModuleList([ TransformerBlock(attention_mask=create_layer_mask()) for _ in range(4)]) def forward(self, x, particle_type, layer_idx): # x: [batch, points, 3] (x,y,logE) h = self.layer_embed(layer_idx) + self.particle_embed(particle_type) for block in self.encoder_blocks: x = block(x, h) return x # 输出向量场

2.2 关键技术创新解析

动态稀疏注意力机制
传统Transformer处理6000点云需要计算3600万对注意力权重。AllShowers的创新方案（图4）包含：

1. 按层索引对点云排序
2. 定义层邻域半径r=2的注意力窗口
3. 采用PyTorch的FlexAttention模块实现稀疏计算

实测表明，该方案不仅提速20倍，还因引入合理的归纳偏置（相邻层物理相关性更强）提升了生成质量。例如在π介子模拟中，该机制能准确保持穿越多个层的最小电离粒子（MIP）径迹连续性。

层间最优传输映射
传统CNF训练时随机匹配潜空间与数据点，导致ODE轨迹弯曲（高曲率κ），需要更多积分步数。AllShowers的创新在于：

1. 对每层的点云独立求解OT问题
2. 使用POT库计算3D预处理空间中的欧氏距离成本
3. 构建点对点的直线映射路径

这种方法使平均轨迹长度缩短37%，训练收敛速度提升2.3倍。如图5所示，即便是复杂的中性K介子（K⁰_L） showers，也能用更少的NFE（Number of Function Evaluations）获得精细结构。

3. 实现细节与工程优化

3.1 数据预处理管道

原始Geant4模拟输出包含数千万个能量沉积点（step），AllShowers通过以下步骤实现高效压缩：

1. 空间离散化：将ECAL和HCAL分别划分为5×5 mm²和30×30 mm²的网格单元
2. 能量聚合：每个非空单元内，保留最高能量沉积的位置坐标，累加该单元总能量
3. 噪声过滤：剔除<10 keV的能量沉积和>200 ns的晚到信号
4. 坐标变换：平移所有点使入射轨迹看似垂直进入量能器中心
5. 对数标准化：对能量取对数后进行Z-score标准化

经过处理，平均每事例从原始约50000个Geant4 steps压缩到2306个有效点，同时保留99.7%的总能量信息。这种压缩比对于保持模型可训练性至关重要。

3.2 训练策略与调参

模型训练采用Ranger优化器（Lookahead+RAdam）组合，关键配置包括：

• 基础学习率：3e-4（余弦退火调度）
• 批量大小：256（跨16块A100 GPU）
• 训练周期：200 epoch（约24小时）
• 权重衰减：0.01（解耦形式）

特别值得注意的是能量校准技巧：训练后发现模型系统性地低估总能量3.3%。不同于简单缩放输出能量（会扭曲分布），AllShowers选择在推理时调整输入条件能量（×1.033），这种"逆向校准"既修正了系统偏差，又保持了内部能量分布的相对关系。

4. 性能评估与对比

4.1 单事例级可视化对比

图5展示了6种典型粒子在Geant4与AllShowers中的模拟结果对比：

• 电磁簇射（e⁺、γ）：在ECAL中呈现典型的云状分布，AllShowers准确再现了横向扩展和纵向发展
• 带电强子（π⁺、π⁻）：清晰再现了MIP径迹和后续簇射，包括在3.5T磁场中的正确偏转方向
• 中性强子（n、K⁰_L）：成功模拟了强子簇射的特征漏斗状结构，中子事例显示出立即发生的核相互作用

特别值得注意的是K⁰_L的模拟——作为长寿命中性粒子，它需要穿越相当距离后才发生相互作用。AllShowers无需特殊处理就能自动学习这种复杂行为模式，证明其嵌入表示（图3）确实捕捉到了粒子类型的本质特征。

4.2 统计分布一致性

图6对比了四种粒子（e⁻、γ、π⁻、n）在三个关键维度上的分布：

单元能量谱

• 准确复现MIP峰（~0.1 GeV）
• 高能尾部（>10 GeV）偏差稍大，主要源于训练数据稀疏

纵向剖面

• 电磁簇射（e⁻、γ）在ECAL前20层快速衰减
• 强子簇射（π⁻、n）穿透更深，HCAL贡献显著
• 交替层效应（因机械结构导致的调制）也被部分保留

径向分布

• 核心区域（<50 mm）精度达95%以上
• 尾部（>200 mm）偏差控制在10%内

4.3 与专用模型的对比

与专为光子设计的CaloClouds3模型相比（图7），AllShowers展现出显著优势：

虽然专用模型在推理速度上略快15%，但AllShowers在保持多粒子通用性的同时，反而在多数指标上领先。特别是在径向分布建模中，得益于Transformer的点间相关性建模能力，其尾部精度比iid假设的扩散模型高出近3倍。

5. 应用前景与优化方向

5.1 实际部署考量

AllShowers当前版本在NVIDIA A100上达到38事例/秒的推理速度，已比Geant4快3-4个数量级。进一步优化建议包括：

• 量化压缩：将FP32模型转为FP16，预计可提升1.8倍速度
• 缓存机制：预生成常见粒子/能量组合的 showers
• 分布式流水线：将PointCountFM与CNF-Transformer部署在不同计算节点

5.2 物理分析验证

初步测试表明，该模型在π⁰→γγ重建等关键物理分析中表现良好：

• 不变质量分辨率偏差<2%
• 双光子开口角重建误差<0.5 mrad
• 能量线性度在10-100 GeV范围内偏差<1%

但需注意，在极高能区（>100 GeV）和极浅入射角（θ<10°）等边缘情况下，可能需要额外的针对性训练。

5.3 未来扩展方向

1. 动态粒度支持：当前固定网格处理可扩展为自适应空间划分
2. 噪声建模：加入电子学噪声和pile-up效应
3. 条件扩展：支持磁场强度、探测器错位等额外条件
4. 混合模拟：与Geant4组成级联模拟系统，关键区域用MC，其余用ML

经验提示：在实际部署时，建议对每个实验设置进行特定的能量校准。我们发现即使3%的系统偏差，在精确测量如希格斯玻色子质量时也会引入显著偏差。简单的线性缩放因子通常足够，但最好用真实数据验证。