[论文解读] CaloScore v2: Single-shot Calorimeter Shower Simulation with Diffusion Models

V. M. Mikuni, Benjamin Nachman|ArXiv.org|Aug 7, 2023

Scientific Computing and Data Management参考文献 39被引用 4

一句话总结

CaloScore v2 提出了一种单次推理扩散模型，用于高能物理中高保真、快速的电磁量能器簇射模拟，通过渐进蒸馏和架构改进，将采样时间减少至单次前向传播，同时保持最先进的保真度。其通过将能量沉积与体素响应生成解耦，并优化扩散过程以降低方差和加快推理速度，实现实时模拟，与完整步长基线相比性能下降极小。

ABSTRACT

Diffusion generative models are promising alternatives for fast surrogate models, producing high-fidelity physics simulations. However, the generation time often requires an expensive denoising process with hundreds of function evaluations, restricting the current applicability of these models in a realistic setting. In this work, we report updates on the CaloScore architecture, detailing the changes in the diffusion process, which produces higher quality samples, and the use of progressive distillation, resulting in a diffusion model capable of generating new samples with a single function evaluation. We demonstrate these improvements using the Calorimeter Simulation Challenge 2022 dataset.

研究动机与目标

• 开发一种用于高能物理中量能器簇射模拟的快速、高保真代理模型。
• 解决扩散模型通常需要数百次去噪步骤导致的采样速度慢的问题。
• 通过架构改进和损失方差更低的修改版扩散过程，提升样本质量。
• 通过渐进蒸馏实现单步（单次）生成，极大缩短推理时间。
• 在2022年量能器模拟挑战赛数据集上展示最先进的性能。

提出的方法

• 该模型采用扩散生成过程，通过逐步去噪潜在表征来生成量能器簇射图像。
• 将生成过程分为两个阶段：首先预测每层的总能量沉积，然后基于该能量条件生成归一化的体素响应。
• 修改扩散过程以降低训练期间的损失方差，提升训练稳定性和样本质量。
• 应用渐进蒸馏，将完整的512步扩散过程压缩为单步推理，实现单次生成。
• 架构中引入注意力层和更小的卷积核，以提升特征学习能力，同时不增加推理时间。
• 使用独立的生成模型预测总能量沉积，提升整体能量分布的准确性。

实验结果

研究问题

• RQ1 能否在保持量能器簇射高保真模拟的前提下，将扩散模型压缩为单步推理？
• RQ2 将能量预测与体素响应生成解耦，是否能提升样本质量与训练稳定性？
• RQ3 渐进蒸馏能否有效应用于高能物理模拟中的扩散模型，以实现实时推理？
• RQ4 与完整步长基线相比，单步扩散模型在保真度和分类器可区分性方面的表现如何？
• RQ5 通过引入注意力层和更小的卷积核等架构改进，能否在不增加推理时间的前提下提升样本质量？

主要发现

• 在数据集3上，CaloScore v2 的单次生成时间仅为每簇射0.002秒，相比原始CaloScore实现2000倍加速。
• 单次生成模型在数据集1上的分类器测试AUC为0.714，数据集3为0.880，表明与Geant4模拟结果高度相似。
• 渐进蒸馏将采样步数从512步减少至1步，同时保持高样本质量，仅造成保真度的轻微下降。
• CaloScore v2的8步版本在数据集1上的JSD为0.066，在数据集3上为0.310，相比原始CaloScore表现出更高的保真度。
• 该模型在CaloChallenge 2022数据集上的表现表明，其属于当前最快的高保真量能器模拟方法之一。
• 注意力层和更小卷积核的使用提升了特征表征能力，且未增加推理时间，使采样速度超过原始CaloScore。