--- title: 'CaloDiT-2: Transformer Diffusion for Fast Calorimeter Simulation' url: https://www.emergentmind.com/topics/calodit-2 type: topic --- # CaloDiT-2: Transformer Diffusion for Fast Calorimeter Simulation CaloDiT-2 is a transformer-based diffusion model for fast simulation of electromagnetic calorimeter showers in collider experiments, introduced in "A Generalisable Generative Model for Multi-Detector Calorimeter Simulation" [2509.07700]. It is designed as a generative surrogate for Geant4 calorimeter simulation, with an emphasis not only on speed and fidelity but also on cross-detector generalisation. The model is formulated around detector-agnostic data representations, transformer diffusion blocks, and a pre-training-plus-adaptation regime on multiple detector geometries. In the reported experiments, it produces showers in approximately \(100\) ms on CPU and approximately \(3\) ms on GPU, matches Geant4 on critical shower observables, and adapts to novel detectors with substantially reduced data and training requirements [2509.07700]. ## 1. Research setting and stated objective High-energy physics experiments, including those at the Large Hadron Collider, rely on the Geant4 toolkit for high-fidelity calorimeter shower simulation, but the CPU cost is described as seconds per shower and becomes prohibitive at high luminosity [2509.07700]. CaloDiT-2 is positioned as a response to this computational bottleneck: a diffusion model using transformer blocks that directly emulates detector responses while preserving the observables that matter for calorimeter simulation. The reported objective is broader than accelerating a single fixed detector geometry. CaloDiT-2 can be applied to specific geometries, as is the case for other models explored for the task, but its stated strength lies in generalisation across detectors through pre-training on multiple geometries and rapid adaptation to new ones. The work frames this as an ambition toward a "foundation model" for FastSim: train once on varied geometries, then adapt cheaply to novel detectors under design [2509.07700]. The scope of the presented model is specifically electromagnetic calorimeter shower generation. The conditioning variables are defined for incident photons with energy \(E \in [1,1000]\) GeV and angular variables \(\phi \in [0,2\pi]\) and \(\theta \in [0.87,2.27]\). A plausible implication is that the reported generalisation claims are anchored to this simulation regime rather than to calorimeter simulation in full generality. ## 2. Detector-agnostic representation and conditioning interface A central component of CaloDiT-2 is a detector-agnostic data representation built from a virtual cylindrical scoring mesh aligned with the incident particle, exploiting approximate azimuthal symmetry [2509.07700]. In the Par04/LEMURS setup, voxels are defined in \((r,\phi,z)\) with fixed resolution \(9\times16\times45\). This mesh is independent of any real read-out geometry: all energy deposits from Geant4 are scored into the mesh, yielding a tensor \(x \in \mathbb{R}^{9\times16\times45}\). This representation separates the learning problem from any one detector segmentation scheme. Instead of binding the network to a specific read-out layout, the model operates on a common volumetric coordinate system into which multiple detectors can be projected. That design is the basis for the multi-detector pre-training strategy. The conditioning variables are the incident photon energy \(E\), the particle angles \(\phi\) and \(\theta\), and a detector identifier \(G \in \mathrm{one\text{-}hot}(K+1)\), where \(K\) is the number of pre-training detectors and the extra bit flags an "unknown" geometry. Voxel preprocessing is given by $$ \hat x_i = \frac{\log(x_i + \epsilon) - \mu}{2\sigma}, $$ where \(\mu\) and \(\sigma\) are the mean and standard deviation of log-energies over the entire pre-training set, and \(\epsilon = 10^{-6}\). The scalar conditions are normalised as \(E/E_{\max}\) and \(\theta/\pi\), while \(\phi\) is encoded as \((\sin\phi,\cos\phi)\) to preserve cyclicity [2509.07700]. The use of an "unknown" geometry flag is notable because it makes the conditioning scheme explicitly compatible with adaptation scenarios. This suggests that the representation is intended not merely for interpolation among known detectors but also for transfer to unseen geometries. ## 3. Diffusion formulation and consistency distillation CaloDiT-2 follows Karras et al.'s EDM framework in continuous time [2509.07700]. The forward stochastic differential equation is $$ dx = \mu(x,t)\,dt + \sigma(t)\,dW, $$ with \(\mu(x,t)=0\), \(\sigma(t)=\sqrt{2}\,t\), and \(W\) a Wiener process. At \(t=0\), the process begins at the data distribution \(p_0(x)\); at \(t=T\), the distribution \(p_T(x)\approx \mathcal N(0,I)\). The reverse-time SDE is $$ dx = \bigl[\mu(x,t) - \sigma(t)^2 \nabla_x \log p_t(x)\bigr]\,dt + \sigma(t)\,d\overline W, $$ and the corresponding probability-flow ODE is $$ dx = \bigl[\mu(x,t) - \tfrac12 \sigma(t)^2 \nabla_x \log p_t(x)\bigr]\,dt. $$ The score \(\nabla_x \log p_t(x)\) is approximated by a neural network \(s_\phi(x,t)\), reparameterised through a network \(v_\theta\) as $$ s_\phi(x,t) = c_{\rm skip}(t)\,x \;+\; c_{\rm out}(t)\;v_\theta\bigl(c_{\rm in}(t)\,x,\;c_{\rm noise}(\sigma(t))\bigr), $$ with $$ c_{\rm skip}(t)=\frac{\sigma_{\rm data}^2}{\sigma_{\rm data}^2 + t^2}, \qquad c_{\rm in}(t)=\frac{t}{\sqrt{\sigma_{\rm data}^2 + t^2}}, \qquad c_{\rm out}(t)=\frac{1}{\sqrt{\sigma_{\rm data}^2 + t^2}}, \qquad c_{\rm noise}(t)=\tfrac{\kappa}{4}\ln t, $$ where \(\sigma_{\rm data}=0.5\) and \(\kappa=10^4\). The EDM training loss is the weighted mean-squared error $$ \mathcal L_{\rm EDM} = \mathbb{E}_{x_0\sim p_0,\;t\sim\mathrm{noiseSchedule}} \bigl[\lambda(t)\,\|s_\phi(x_t,t)-x_0\|_2^2\bigr], $$ with $$ \lambda(t)=\frac{t^2+\sigma_{\rm data}^2}{t\,\sigma_{\rm data}}. $$ To reduce sampling cost, the work applies consistency distillation in the sense of Song et al. to collapse \(32\) EDM sampling steps into \(1\). With discretised times \(\epsilon=t_0<\cdots