[论文解读] Calo-VQ: Vector-Quantized Two-Stage Generative Model in Calorimeter Simulation

Qibin Liu, C. O. Shimmin+4 more|arXiv (Cornell University)|May 10, 2024

Gaussian Processes and Bayesian Inference被引用 4

一句话总结

该论文提出 Calo-VQ，一种基于 VQ-VAE 和基于 GPT 的潜在序列模型的两阶段向量量化生成模型，用于实现快速量能器模拟。其相较于 GEANT4 的速度提升超过 2000 倍，同时保持了高物理保真度，在多种探测器几何结构和分辨率下，每事件生成时间低于 1.2 毫秒。

ABSTRACT

We introduce a novel machine learning method developed for the fast simulation of calorimeter detector response, adapting vector-quantized variational autoencoder (VQ-VAE). Our model adopts a two-stage generation strategy: initially compressing geometry-aware calorimeter data into a discrete latent space, followed by the application of a sequence model to learn and generate the latent tokens. Extensive experimentation on the Calo-challenge dataset underscores the efficiency of our approach, showcasing a remarkable improvement in the generation speed compared with conventional method by a factor of 2000. Remarkably, our model achieves the generation of calorimeter showers within milliseconds. Furthermore, comprehensive quantitative evaluations across various metrics are performed to validate physics performance of generation.

研究动机与目标

• 解决使用 GEANT4 进行传统量能器模拟所面临的计算瓶颈，其速度过慢，难以满足大规模高能物理工作流的需求。
• 开发一种快速、可微分且可扩展的模拟方法，同时保留关键物理指标，如总能量沉积和簇射形状。
• 通过将模拟分解为离散潜在空间生成与序列建模，实现在高分辨率量能器上的实时推理。
• 在多种探测器几何结构中实现泛化，包括不规则结构和圆柱形结构，最高可达 40,500个体素。
• 集成物理启发的设计选择，如 FFT 重采样、圆柱卷积和逐层归一化，以提升生成质量。

提出的方法

• 采用两阶段框架：首先，使用带有向量量化的卷积编码器-解码器将 3D 量能器簇射数据压缩到离散潜在空间。
• 应用圆柱卷积神经网络，以更好地建模具有空间结构的高分辨率圆柱形量能器（如 ds3，含 40,500 个体素），提升归纳偏差。
• 使用基于 GPT 的 Transformer 模型自回归地生成离散潜在标记序列，实现快速且连贯的采样。
• 引入物理启发的预处理方法：能量沉积的对数变换和基于 FFT 的重采样，以提升潜在表征质量。
• 在 VQ-VAE 中应用逐层归一化（LN），以稳定训练并提升重建保真度，尤其在高维空间中表现更优。
• 将模型分解为两个阶段：第一阶段（VQ-VAE）负责压缩，第二阶段（GPT）负责潜在生成，从而实现独立优化与可扩展性。

Figure 1: Demonstration of the vector-quantization based two-stage generative model. The upper and lower parts show the two stages of the model, respectively.

实验结果

研究问题

• RQ1 两阶段 VQ-VAE 与自回归模型是否能在保持关键物理可观测量的前提下，实现亚毫秒级推理速度？
• RQ2 该模型在不同探测器几何结构下的表现如何，包括不规则结构、低分辨率和高分辨率圆柱形设计？
• RQ3 物理启发的架构选择（如圆柱卷积和 FFT 重采样）在多大程度上提升了生成质量与稳定性？
• RQ4 在不同数据集规模（如 368 至 40,500 个体素）下，模型复杂度、推理速度与重建精度之间的权衡如何？
• RQ5 两阶段设计是否能实现可扩展且模块化的优化，特别是未来可通过在第二阶段探索潜在扩散模型实现进一步改进？

主要发现

• 模型相较基于 GEANT4 的模拟实现了 2000 倍的速度提升，低分辨率数据集的生成时间最短可达 0.25 毫秒，最高分辨率的 40,500 个体素数据集生成时间仅为 1.14 毫秒。
• 物理性能表现优异：能量沉积分离度指标为 0.01（最差/最佳），簇射形状分离度为 0.1，表明与参考 GEANT4 模拟具有高度保真度。
• 第一阶段（VQ-VAE）在所有数据集上表现出近乎恒定的推理时间，表明卷积压缩模块在输入尺寸变化下仍具备良好可扩展性。
• 第二阶段（GPT）的生成时间仅与潜在空间大小相关，而不依赖于输入维度，因此即使在高分辨率数据上也能实现高效推理。
• 在高分辨率 ds3 数据集上，采用逐层归一化（LN）的模型生成时间为 36.10 毫秒，但分离度指标仍保持在 0.095 的低水平，表明对架构变化具有强鲁棒性。
• 模型在不同数据集上泛化能力出色：在不规则结构（ds1）和圆柱形结构（ds2, ds3）上均保持一致的性能表现，体现出架构的灵活性。

Figure 2: Architecture of Stage-1 Model.