Double Metric Learning for Building Directed Graphs with Chain Connections for the ATLAS ITk Detector

Abstract

Graph construction is an essential step in the Graph Neural Network (GNN) based tracking pipelines. The goal of the graph construction is to construct a graph that contains only the defined true edge connections between nodes (detector hits). A promising approach for the graph construction is through the Metric Learning approach, where a node representation in an embedding space is learned, and nodes are connected according to their distance in the embedding space. The loss function for the metric learning in this case is a contrastive loss encouraging the true pairs of nodes to be close to each other, and pulling away the false pairs of nodes. This approach presents a conflict of the learning objective for the hopping connections when a true edge is defined as a chain connection in a particle track. To address the conflict for this case, we propose a ``Double Metric Learning'' approach, where two node representations are learned. A directed graph can then be constructed based on the distance between the two representations from two nodes respectively. We test this idea with the ATLAS ITk detector at the HL-LHC using the ATLAS ITk simulation and show better graph construction performance particularly for particles with high transverse momentum compared to the Simple Metric Learning approach. We also show that Double Metric Learning is able to accurately predict edge direction.