[Paper Review] CaloScore v2: Single-shot Calorimeter Shower Simulation with Diffusion Models

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

Scientific Computing and Data Management39 references4 citations

TL;DR

CaloScore v2 introduces a single-shot diffusion model for high-fidelity, fast calorimeter shower simulation in high-energy physics, leveraging progressive distillation and architectural improvements to reduce sampling time to one inference step while maintaining state-of-the-art fidelity. It achieves this by decoupling energy deposition from voxel response generation and optimizing the diffusion process for lower variance and faster inference, enabling real-time simulation with minimal performance degradation compared to full-step baselines.

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.

Motivation & Objective

• To develop a fast, high-fidelity surrogate model for calorimeter shower simulation in high-energy physics.
• To overcome the slow sampling speed of diffusion models, which typically require hundreds of denoising steps.
• To improve sample quality through architectural refinements and a modified diffusion process with reduced loss variance.
• To enable single-step (single-shot) generation via progressive distillation, drastically reducing inference time.
• To demonstrate state-of-the-art performance on the Calorimeter Simulation Challenge 2022 dataset.

Proposed method

• The model uses a diffusion generative process that gradually denoises a latent representation to generate calorimeter shower images.
• It decouples the generation into two stages: first predicting total energy deposition per layer, then generating normalized voxel responses conditioned on that energy.
• The diffusion process is modified to reduce loss variance during training, improving training stability and sample quality.
• Progressive distillation is applied to compress the full 512-step diffusion process into a single-step inference, enabling single-shot generation.
• The architecture incorporates attention layers and smaller convolutional kernels to improve feature learning without increasing inference time.
• A separate generative model is used to predict total energy deposition, improving the accuracy of the overall energy distribution.

Experimental results

Research questions

• RQ1 Can a diffusion model be compressed to single-step inference while preserving high-fidelity simulation of calorimeter showers?
• RQ2 Does decoupling energy prediction from voxel response generation improve sample quality and training stability?
• RQ3 Can progressive distillation be effectively applied to diffusion models in high-energy physics simulation to achieve real-time inference?
• RQ4 How does the performance of a single-step diffusion model compare to full-step baselines in terms of fidelity and classifier distinguishability?
• RQ5 Can architectural improvements such as attention layers and smaller kernels enhance sample quality without increasing inference time?

Key findings

• CaloScore v2 achieves a single-step generation time of 0.002 seconds per shower on dataset 3, representing a 2000× speedup over the original CaloScore.
• The single-shot model achieves an AUC of 0.714 on dataset 1 and 0.880 on dataset 3 in a classifier test, indicating strong similarity to Geant4 simulations.
• Progressive distillation reduces the number of sampling steps from 512 to 1 while maintaining high sample quality, with only a modest drop in fidelity.
• The 8-step version of CaloScore v2 achieves a JSD of 0.066 on dataset 1 and 0.310 on dataset 3, showing improved fidelity over the original CaloScore.
• The model’s performance on the CaloChallenge 2022 dataset demonstrates that it is among the state-of-the-art in fast, high-fidelity calorimeter simulation.
• The use of attention layers and smaller kernels improves feature representation without increasing inference time, enabling faster sampling than the original CaloScore.