Differentiable Surrogate for Detector Simulation and Design with Diffusion Models

This paper introduces a differentiable conditional denoising-diffusion surrogate for electromagnetic calorimeter simulations that achieves high-fidelity energy-deposition maps with minimal retraining via Low-Rank Adaptation and provides accurate gradients for gradient-based detector design optimization.

The Gist

Imagine you are an architect trying to design the ultimate building to catch falling rain. In the world of high-energy physics, this "building" is a particle detector (specifically an electromagnetic calorimeter), and the "rain" is a storm of subatomic particles.

To make sure the building works, scientists usually run massive, incredibly slow computer simulations (called GEANT4) to see how the "rain" hits the walls. It's like trying to predict exactly how a single raindrop will splash by simulating every single molecule of water in the atmosphere. It's accurate, but it takes so long that you can't test thousands of different building designs in a reasonable time.

This paper introduces a super-smart, fast "AI apprentice" that learns to mimic these slow simulations so you can design detectors much faster. Here is how it works, broken down into simple concepts:

1. The "AI Apprentice" (The Diffusion Model)

Think of the AI as a student who has watched millions of hours of "rain hitting buildings" videos.

• How it learns: Instead of memorizing the videos, the AI learns to "unscramble" them. Imagine taking a clear photo of a rain splash and slowly adding static noise until it's just gray fuzz. The AI learns to reverse this process: starting with gray fuzz, it learns to remove the noise step-by-step until a perfect, realistic rain splash image appears.
• The Magic Trick: This AI is differentiable. In plain English, this means it doesn't just give you an answer; it can also explain why it gave that answer. If you ask, "What happens if I make the wall 1mm thicker?" the AI can instantly calculate how the splash changes, not just by guessing, but by understanding the mathematical relationship. This allows architects to use "gradient-based optimization"—basically, sliding down a hill to find the lowest point (the best design) automatically, rather than guessing and checking.

2. The Two-Stage Training (Pre-training + LoRA)

Training an AI to understand every possible building design from scratch would take forever and require a supercomputer the size of a city. The authors used a clever two-step strategy:

• Stage 1: The Generalist (Pre-training): First, they teach the AI on a huge dataset of many different building sizes and materials. The AI becomes a "Generalist" who understands the basic physics of how rain splashes on almost any wall. It learns the general rules of the universe.
• Stage 2: The Specialist (LoRA Adaptation): Now, imagine you need to design a very specific, weird-shaped building that the AI hasn't seen before. Instead of retraining the whole AI (which is expensive), they use a technique called LoRA (Low-Rank Adaptation).
  • The Analogy: Think of the Generalist AI as a master chef who knows how to cook Italian, Chinese, and Mexican food. You want them to cook a very specific regional dish they've never made. Instead of sending them back to culinary school for 4 years, you just give them a specialized recipe card (the LoRA adapter) that tweaks their existing skills slightly. Now they can cook that specific dish perfectly, using only a tiny bit of new data.

3. The Results: Fast, Accurate, and Helpful

The team tested this AI apprentice:

• Accuracy: When they compared the AI's "fake" rain splashes to the real, slow computer simulations, they matched almost perfectly (within 2% error). The total energy, the spread of the splash, and the shape were all spot on.
• Speed: The AI generates these results in a fraction of a second, whereas the real simulation takes minutes or hours.
• Design Help: Because the AI is differentiable, they could ask it to optimize the detector. The AI successfully told them which way to tweak the design to get better results, matching the "direction" of the truth, even if the exact numbers were slightly smoothed out.

Why Does This Matter?

In the future, we will build massive particle colliders (like the High-Luminosity LHC or a Muon Collider) to discover new physics. These machines are incredibly complex and expensive.

Before, designing them was like trying to find the best route through a maze by walking every single path one by one. It took too long.
With this new Diffusion Surrogate, scientists can now use a "GPS" that instantly calculates the best route. They can test thousands of design variations in the time it used to take to test one. This means we can build better, more sensitive detectors faster, potentially leading to bigger discoveries in the universe.

In a nutshell: They built a fast, smart AI that learns the rules of particle physics, can quickly adapt to new designs with a tiny "cheat sheet," and helps engineers optimize particle detectors by calculating the perfect design changes instantly.

Technical Summary

1. Problem Statement

In High-Energy Physics (HEP), particularly for future collider experiments like the High-Luminosity LHC (HL-LHC) and muon colliders, the design and optimization of particle detectors (specifically electromagnetic calorimeters) rely heavily on accurate simulations of particle showers.

• The Bottleneck: The standard simulation tool, GEANT4, provides high-fidelity physics-based results but is computationally expensive and inherently non-differentiable.
• The Challenge: Modern design optimization requires gradient-based methods to navigate high-dimensional parameter spaces (geometry, materials, granularity). Traditional "black-box" optimization (e.g., Bayesian optimization, evolutionary algorithms) becomes inefficient as dimensionality increases. Furthermore, existing machine learning surrogates (GANs, Normalizing Flows) often lack differentiability or require retraining for every new detector configuration, making them unsuitable for rapid, end-to-end co-design workflows.

2. Methodology

The authors propose a conditional denoising-diffusion surrogate framework that combines high-fidelity generation with differentiability. The approach consists of three core components:

A. Diffusion-Based Surrogate Architecture

• Model Type: A conditional Denoising Diffusion Probabilistic Model (DDPM) trained to learn the conditional distribution $p_\theta(x|y)$, where $x$ is the calorimeter shower (energy deposition map) and $y$ represents conditioning variables (incident energy, cell size, material type).
• Architecture: A U-Net backbone with skip connections.
  • Conditioning: Integrates diffusion time embeddings and calorimeter-specific embeddings (energy, cell dimensions, material) into every residual block.
  • Sampling: Uses Denoising Diffusion Implicit Models (DDIM) for inference. This allows for deterministic, differentiable sampling in fewer steps compared to standard DDPM, which is crucial for backpropagation.
• Differentiability: The entire pipeline (from conditioning parameters to generated shower) is differentiable, enabling the computation of gradients $\nabla_y L(y)$ for a utility function $L$ via automatic differentiation.

B. Two-Stage Training Strategy (Pre-training + LoRA)

To address the computational cost of training a single model on the entire design space:

1. Global Pre-training: The model is trained on a diverse dataset of GEANT4 simulations covering a wide range of cell sizes and energies (1–100 GeV) for a nominal geometry. This establishes a global representation of the simulation space.
2. Local Adaptation (LoRA): For a specific, unseen detector geometry, the model is adapted using Low-Rank Adaptation (LoRA).
   • Only low-rank matrices within the convolutional layers are trained; the original weights are frozen.
   • This requires a very small post-training dataset (10,000 events) and minimal compute, allowing rapid specialization to new configurations without catastrophic forgetting.

C. Gradient-Based Utility Analysis

The authors define a differentiable utility function to test the surrogate's ability to guide design:

• Reconstruction: A soft mask discriminates signal from background (simulated beam-induced background).
• Utility: Defined as the inverse of a stabilized Mean Squared Error (MSE) between the reconstructed energy and the true energy.
• Gradient Flow: Gradients of this utility with respect to design parameters (e.g., cell dimensions) are computed via backpropagation through the DDIM sampling process and compared against Finite Difference (FD) references from GEANT4.

3. Key Contributions

1. Differentiable Diffusion Surrogate: First application of a conditional diffusion model for calorimeter simulation that is explicitly designed for gradient-based optimization, enabling end-to-end detector co-design.
2. Efficient Adaptation Framework: A novel two-stage strategy combining broad pre-training with LoRA-based fine-tuning. This allows the model to generalize across diverse configurations while achieving high local accuracy with minimal data and compute.
3. Validation of Gradients: Demonstrated that the surrogate can reproduce the qualitative structure and directional trends of the true utility landscape (gradient signs and geometric dependencies), making it viable for sensitivity analysis.
4. Deterministic Sampling for Optimization: Utilization of DDIM to ensure the mapping from design parameters to outputs is smooth and differentiable, avoiding the stochastic noise that typically hinders gradient estimation in generative models.

4. Results

• Generative Fidelity:
  • The model generates high-fidelity energy deposition maps that closely match GEANT4 ground truth.
  • Metrics: Relative Root Mean Square Error (RRMSE) for total energy, energy-weighted radius, and shower dispersion is below 2% for high-energy cases (70–100 GeV).
  • Visuals: Generated longitudinal and transverse profiles match the ground truth across all energy levels, capturing both spatial structure and stochastic variability.
• Adaptation Performance:
  • When applied to an unseen geometry ($2.5 \times 2.5 \times 6$ cm$^3$), the pre-trained model showed systematic underestimation in longitudinal profiles.
  • After LoRA post-training, the model achieved significant improvement, with RRMSE values dropping closer to the ground truth (e.g., total energy RRMSE improved from ~0.73 to ~0.63 at 50 GeV).
• Gradient Analysis:
  • Sign and Trend: The surrogate correctly captures the sign and geometric dependence of gradients compared to Finite Difference references.
  • Magnitude: Gradients are slightly smoother and underestimated compared to GEANT4 (due to the deterministic nature of DDIM averaging out fluctuations), but post-training improves magnitude consistency.
  • Cosine Similarity: The directional alignment between surrogate gradients and true gradients is generally good, though some sign flips occur at extreme energies, indicating the method is promising but requires further refinement for full optimization loops.

5. Significance and Future Outlook

• Accelerated Design: This work provides a pathway to replace computationally expensive GEANT4 simulations with fast, differentiable surrogates, enabling gradient-based optimization of detector geometries that was previously infeasible.
• Scalability: The LoRA approach solves the "data hunger" problem of training models for every new detector configuration, making the framework scalable for future collider experiments.
• Limitations & Future Work:
  • Current limitations include a restricted set of materials, 2D down-sampling of 3D showers, and idealized homogeneous calorimeters.
  • Future work aims to integrate stochastic backgrounds, model cell boundary effects, and perform full end-to-end optimization loops to identify optimal calorimeter configurations for muon colliders.

In conclusion, the paper establishes that diffusion-based surrogates, when combined with parameter-efficient adaptation (LoRA) and deterministic sampling, are a viable and powerful tool for the next generation of differentiable detector design in high-energy physics.