Differentiable Surrogate for Detector Simulation and Design with Diffusion Models

TL;DR

This work develops a conditional denoising-diffusion surrogate for electromagnetic calorimeter showers that is differentiable with respect to detector-design parameters. The model combines a diffusion-based generator with a two-stage pre-training and LoRA-based post-training, enabling rapid adaptation to new geometries while preserving high-fidelity energy-deposition maps conditioned on energy, cell size, and material. Fidelity assessments show relative RMSE below 2% for key observables at high energies, and gradient analyses demonstrate qualitative agreement with finite-difference references, enabling gradient-based optimization in detector design pipelines. The approach offers a scalable path toward differentiable, physics-aware surrogate modeling for fast, gradient-informed detector optimization in future collider experiments.

Abstract

In this work, we present a conditional denoising-diffusion surrogate for electromagnetic calorimeter showers that is trained to generate high-fidelity energy-deposition maps conditioned on key detector and beam parameters. The model employs efficient inference using Denoising Diffusion Implicit Model sampling and is pre-trained on GEANT4 simulations before being adapted to a new calorimeter geometry through Low-Rank Adaptation, requiring only a small post-training dataset. We evaluate physically meaningful observables, including total deposited energy, energy-weighted radius, and shower dispersion, obtaining relative root mean square error values below 2% for representative high-energy cases. This is in line with state-of-the-art calorimeter surrogates which report comparable fidelity on high-level observables. Furthermore, we compare gradients of a reconstruction-based utility function with respect to design parameters between the surrogate and finite-difference references. The diffusion surrogate reproduces the qualitative structure and directional trends of the true utility landscape, providing usable sensitivities for gradient-based optimization. These results show that diffusion-based surrogates can accelerate simulation-driven detector design while enabling differentiable, gradient-informed analysis.

Figures (17)

• Figure 1: Architecture of the proposed conditional diffusion model for calorimeter shower simulation. The model employs a U-Net backbone that processes noisy shower images, with skip connections linking encoder and decoder stages. Calorimeter-specific conditions (energy, cell sizes, and material) together with diffusion time steps are embedded and injected into each residual block, enabling the network to learn energy deposition patterns across detector configurations. Numbers above each block indicate the number of feature maps, while numbers below indicate the spatial resolution of the feature maps.
• Figure 2: Schematic of a ResBlock in the conditional U-Net architecture with LoRA adapters. The input feature map passes through the first convolution block (GroupNorm → Swish → Conv), then time and conditioning embeddings are added. The second convolution block (GroupNorm → Swish → Dropout → Conv) follows, after which the residual shortcut and LoRA adapter are applied.
• Figure 3: Schematic of the differentiable reconstruction–utility pipeline. The forward path flows from the design parameters $\theta$ into either the GEANT4 simulation or the DDPM surrogate, producing calorimeter results that are passed to the reconstruction and utility function. The backward path illustrates gradient backpropagation of the utility $U(x)$ with respect to the design parameters $\theta$, which occurs only through the differentiable DDPM surrogate.
• Figure 4: GEANT4 visualization of a 4 $\times$ 4 $\times$ 10 $\mathrm{cm}^3$$PbF_2$ scintillator cell used in the pre-training dataset. The incident photon (red line) initiates an electromagnetic shower whose energy deposits are recorded in the voxelized calorimeter grid (white). Green tracks represent secondary particles generated during the cascade.
• Figure 5: 2×5 grid of X–Z energy-deposition maps for a 4 $\times$ 4 $\times$ 10 $cm^3$ scintillator cell of $PbF_2$. Top row shows ground-truth histograms from single photon interactions of our .h5 dataset (down-sampled to 32×32), bottom row shows single samples generated by our conditional DDPM under the same normalized ($xy$,$z$,material,energy) conditions. All panels use the inferno colormap from Matplotlib; columns are labeled by primary energy.