Influencer Loss: End-to-end Geometric Representation Learning for Track Reconstruction

Abstract

Significant progress has been made in applying graph neural networks (GNNs) and other geometric ML ideas to the track reconstruction problem. State-of-the-art results are obtained using approaches such as the Exatrkx pipeline, which currently applies separate edge construction, classification and segmentation stages. One can also treat the problem as an object condensation task, and cluster hits into tracks in a single stage, such as in the GravNet architecture. However, condensation with such an architecture may still require non-differentiable operations, and arbitrary post-processing. In this work, I extend the ideas of geometric attention to the task of fully geometric (and therefore fully differentiable) end-to-end track reconstruction in a single step. To realize this goal, I introduce a novel condensation loss function called Influencer Loss, which allows an embedded representation of tracks to be learned in tandem with the most representative hit(s) in each track. This loss has global optima that formally match the task of track reconstruction, namely smooth condensation of tracks to a single point, and I demonstrate this empirically on the TrackML dataset. The model not only significantly outperforms the physics performance of the baseline model, it is up to an order of magnitude faster in inference.