Contrastive Metric Learning for Point Cloud Segmentation in Highly Granular Detectors

Abstract

We propose a novel clustering approach for point-cloud segmentation based on supervised contrastive metric learning (CML). Rather than predicting cluster assignments or object-centric variables, the method learns a latent representation in which points belonging to the same object are embedded nearby while unrelated points are separated. Clusters are then reconstructed using a density-based readout in the learned metric space, decoupling representation learning from cluster formation and enabling flexible inference. The approach is evaluated on simulated data from a highly granular calorimeter, where the task is to separate highly overlapping particle showers represented as sets of calorimeter hits. A direct comparison with object condensation (OC) is performed using identical graph neural network backbones and equal latent dimensionality, isolating the effect of the learning objective. The CML method produces a more stable and separable embedding geometry for both electromagnetic and hadronic particle showers, leading to improved local neighbourhood consistency, a more reliable separation of overlapping showers, and better generalization when extrapolating to unseen multiplicities and energies. This translates directly into higher reconstruction efficiency and purity, particularly in high-multiplicity regimes, as well as improved energy resolution. In mixed-particle environments, CML maintains strong performance, suggesting robust learning of the shower topology, while OC exhibits significant degradation. These results demonstrate that similarity-based representation learning combined with density-based aggregation is a promising alternative to object-centric approaches for point cloud segmentation in highly granular detectors.

Figures (4)

• Figure 1: Event-level embedding-geometry distributions for EM and HAD showers at increasing generated particle multiplicity ($N_{\mathrm{particles}}$). Top: median inter-shower separation $d_{\mathrm{inter}}^{\mathrm{median}}$. Middle: intra-shower distance tail $d_{\mathrm{intra}}^{\mathrm{Q99}}$. Bottom: separation margin $\Delta = d_{\mathrm{inter}}^{\mathrm{median}} - d_{\mathrm{intra}}^{\mathrm{Q99}}$. Narrow positive or near-zero margins indicate a well-defined clustering scale, while broad or negative margins indicate increasing ambiguity between showers.
• Figure 2: Reconstruction efficiency (top) and purity (bottom) as functions of primary particle energy for electromagnetic (EM) and hadronic (HAD) showers. Columns show models trained on EM, trained on HAD, and a mixed-trained model evaluated separately on EM and HAD showers. The vertical dashed line indicates the upper boundary of the training energy range. The dominant differences are largely independent of energy, indicating that performance is controlled primarily by the learned representation rather than by the absolute shower energy.
• Figure 3: Reconstruction efficiency (top), purity (middle), and number ratio (bottom) as functions of particle multiplicity $N$ for electromagnetic (EM) and hadronic (HAD) showers. Columns show models trained on EM, trained on HAD, and a mixed-trained model evaluated separately on EM and HAD showers. The vertical dashed line indicates the upper boundary of the training multiplicity range. The performance gap between CML and OC increases strongly with multiplicity, showing that clustering stability in dense environments is determined by the underlying embedding geometry.
• Figure 4: Energy resolution as a function of primary particle energy for electromagnetic (EM) and hadronic (HAD) showers. Panels show models trained on EM, a mixed-trained model evaluated on EM, trained on HAD, and a mixed-trained model evaluated on HAD. The vertical dashed line marks the upper boundary of the training energy range. The black dashed curve denotes the ideal pattern-recognition limit, defined as the resolution obtained by a perfect clustering algorithm with no merging or splitting, and represents a lower bound on the achievable resolution given the detector response. Improvements in CML resolution follow directly from its improved clustering purity and reduced merging.