Contrastive Metric Learning for Point Cloud Segmentation in Highly Granular Detectors

TL;DR

This paper introduces a contrastive metric learning approach for point cloud segmentation in highly granular detectors, improving stability, separation, and generalization over existing methods like object condensation.

Contribution

The paper presents a novel supervised contrastive metric learning method that enhances point cloud segmentation by learning a separable embedding space, outperforming object condensation in complex detector environments.

Findings

CML yields more stable and separable embeddings.

Improved separation and generalization in high-multiplicity scenarios.

Higher reconstruction efficiency and purity, better energy resolution.

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.