Reconstruction of overlapping electromagnetic showers in calorimeters using Transformers

Abstract

Accurate clustering of electromagnetic energy deposits is essential for reconstructing photons and electrons in modern hadron collider experiments, where boosted topologies and pileup cause overlapping showers and ambiguous energy assignment. We present deep learning-based clustering approaches that reconstruct particle energy and position directly from calorimeter readout. The study includes a two-step strategy in which candidate seed windows are identified and then jointly processed via distance-weighted message passing or attention mechanism and a single-step graph transformer, ClusTEX, which performs candidate selection and reconstruction in one inference stage. ClusTEX uses a novel positional encoding scheme that separates local coordinates within the graph from global detector coordinates, enabling efficient, geometry-aware inference. Models are trained on GEANT4 simulations of a simplified (toy) and an ECAL-inspired topology with an explicit $η-φ$ dependence. Performance is evaluated using efficiency, energy and position resolutions and splitting rate - reconstruction of two objects for a single photon. In the toy calorimeter, attention-based interactions improve the reconstruction of overlapping showers relative to both the standard algorithm and distance-driven message passing, while maintaining performance on isolated photons and reducing splitting without multi-pass inference. For boosted $π^0\toγγ$, the attention-based model retains di-photon mass reconstruction capability, where the standard algorithm becomes inefficient. In the ECAL-inspired topology, ClusTEX provides the best overall performance, yielding improved energy resolution and reduced splitting compared to two-step approaches and the standard algorithm. It also remains robust under localized detector failures, showing improved stability and partial recovery of energy in non-responsive channels.

Figures (20)

• Figure 1: Left: realistic CMS ECAL simulation geometry, where each PbWO$_4$ crystal is quasi-projective and aligned toward the interaction point with a 3$^\circ$ tilt in both the $z$ and $\phi$ directions. Right (top): example calorimeter window for a single-photon event showing the energy deposits in the crystals. Right (bottom): the same calorimeter window after applying the simulated read-out noise and an energy cut at 50 MeV to mitigate contributions dominated by electronic noise
• Figure 2: The optimized CNN based SeedFinder (SF) model architecture. The seed window candidates passing the threshold on deposited energy are taken as inputs to the SF model. The model assigns a probability whether the window is centred around the true seed crystal of a generated photon. The network consists of three consecutive convolutional blocks with a kernel size of three with decreasing filter sizes and three fully connected (dense) layers with residual connection in between
• Figure 3: Architecture overview of the attention-based PoEN. The two flavours differ only by the window encoder (flattened pixels vs convolutional)
• Figure 4: The ratio of the total reconstructed energy to the generated energy for the 1-photon sample as obtained by GAT-$\text{PoEN}_\mathrm{flat}$, GNN-based PoEN and PFClustering (PF)
• Figure 5: Di-photon invariant mass reconstructed from the $\pi^0$ sample and divided into four regions according to pion momentum. Comparison between GAT-$\text{PoEN}_\mathrm{flat}$, GNN-based PoEN and PFClustering