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

TL;DR

This paper introduces Double Metric Learning, a novel approach for constructing directed graphs in GNN-based tracking, effectively handling chain connections in particle tracks at the ATLAS ITk detector.

Contribution

The paper proposes a double metric learning method that learns two node representations to improve directed graph construction for particle tracking.

Findings

Better graph construction performance for high transverse momentum particles.

Accurate prediction of edge direction.

Improved handling of chain connections in particle tracks.

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.