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

This paper proposes a "Double Metric Learning" approach that learns two distinct node representations to resolve conflicts in constructing directed graphs with chain connections for the ATLAS ITk detector, demonstrating improved performance in graph construction and edge direction prediction for high transverse momentum particles compared to simple metric learning.

The Gist

Imagine you are trying to solve a massive, 3D jigsaw puzzle in a dark room. The pieces are tiny flashes of light (called "hits") left behind by subatomic particles zipping through a giant detector called the ATLAS ITk. Your goal is to figure out which flashes belong to the same particle and in what order they happened, so you can trace the particle's path.

To do this, scientists use a type of artificial intelligence called a Graph Neural Network (GNN). But before the AI can solve the puzzle, it needs to build a "map" (a graph) connecting the dots. The challenge is: How do you connect the dots without making a mess?

The Problem: The "Chain" Confusion

In the old way of doing things (called Simple Metric Learning), the AI tries to learn a special "address" for every flash of light. The rule is simple: if two flashes belong to the same particle, they should have similar addresses.

However, there's a catch. In particle physics, we only want to connect a flash to the very next flash in the line (like a chain: A connects to B, and B connects to C). We don't want to connect A directly to C, because that skips a step.

Here is where the old method gets confused, like a teacher giving contradictory instructions:

1. "Bring A and B together."
2. "Bring B and C together."
3. "But keep A and C far apart!"

Mathematically, if A is close to B, and B is close to C, then A must be close to C. The AI gets a headache trying to satisfy all three rules at once. It ends up building a messy map with too many connections, including "hopping" connections that skip steps, which slows everything down.

The Solution: The "Double Agent" Strategy

The authors of this paper propose a new method called Double Metric Learning.

Instead of giving every flash of light just one address, they give it two:

1. A "Source" address (where the light came from).
2. A "Target" address (where the light is going).

Think of it like a one-way street system.

• When the AI looks at the connection from A to B, it compares A's Source address with B's Target address.
• When it looks at B to C, it compares B's Source with C's Target.

This solves the confusion! The AI learns that A's Source is close to B's Target, and B's Source is close to C's Target. But there is no rule forcing A's Source to be close to C's Target. The "contradiction" disappears.

The Results: A Cleaner, Faster Map

The team tested this new method using simulations of the ATLAS detector (specifically looking at high-energy collisions). Here is what they found:

• Direction Matters: Because the method uses "Source" and "Target" addresses, the resulting map is directed. It knows exactly which way the particle is moving (like a one-way arrow), rather than just a fuzzy cloud of connections.
• Fewer Mistakes: The new method is much better at avoiding "hopping" errors (connecting A directly to C). It sticks strictly to the chain, keeping the map clean.
• High-Speed Performance: The method works especially well for particles moving very fast (high momentum). These are the hardest particles to track, and the new method builds a more accurate map for them than the old way.
• Efficiency: The final maps are smaller and less cluttered, which means the computer doesn't have to work as hard to solve the puzzle later.

The Bottom Line

The paper introduces a clever trick of giving particles two different "identities" (Source and Target) to teach the AI how to build a one-way map. This stops the AI from getting confused by the rules of the game, resulting in a cleaner, more accurate map of particle paths, especially for the fastest-moving particles.

Note: The paper focuses strictly on the construction of these maps for the ATLAS detector. It does not discuss medical applications or other future uses beyond improving the efficiency of particle tracking in this specific high-energy physics context.

Technical Summary

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

Problem Statement
Graph Neural Networks (GNNs) have become a standard for particle track reconstruction in high-energy physics, yet their efficacy relies heavily on the initial graph construction step. This step must balance two competing objectives: minimizing graph size to ensure computational tractability (as GNN inference costs scale with edge count) and maximizing the recovery of "true" edges to ensure robust downstream segmentation.

A specific challenge arises when applying the Metric Learning paradigm to the chain connection definition of true edges. In chain connections, only successive hits along a particle's trajectory are considered true pairs, excluding "hopping" connections between non-successive hits. Standard Metric Learning employs a symmetric distance metric in a latent embedding space, trained via a contrastive loss to pull true pairs together and push false pairs apart. However, under the chain definition, this creates a logical conflict: if hits $A$ and $B$ are true pairs, and $B$ and $C$ are true pairs, the symmetric metric forces $A$ and $C$ to be close (transitivity). Yet, $A$ and $C$ are not a true pair in the chain definition and should be pushed apart. This contradiction confuses the training objective, potentially degrading graph purity and efficiency, particularly for high-transverse momentum ($p_T$) particles where track curvature and hit density complicate reconstruction.

Methodology
To resolve this conflict, the authors propose Double Metric Learning, an evolution of the standard approach (termed "Simple Metric Learning").

• Dual Representations: Unlike Simple Metric Learning, which maps each hit to a single latent vector, Double Metric Learning maps each hit into two distinct representations: a source representation ($\vec{S}$) and a target representation ($\vec{T}$).
• Asymmetric Distance Metric: The model learns a directed relationship. A hit acts as a source when it is the origin of a connection and as a target when it is the destination. The distance metric is defined as $|\vec{S}_i - \vec{T}_j|$, which is generally asymmetric ($|\vec{S}_i - \vec{T}_j| \neq |\vec{S}_j - \vec{T}_i|$).
• Loss Function: The training utilizes a contrastive hinge loss over directed pairs ($h_i \to h_j$). For true pairs (where $h_j$ is the immediate successor of $h_i$), the loss minimizes $|\vec{S}_i - \vec{T}_j|$. For non-successive pairs, it maximizes this distance.
• Resolution of Conflict: By decoupling the representations, the model can minimize the distance between $\vec{S}_A$ and $\vec{T}_B$, and between $\vec{S}_B$ and $\vec{T}_C$, without transitively forcing $\vec{S}_A$ and $\vec{T}_C$ to be close. This eliminates the mathematical contradiction inherent in the symmetric approach when applied to chain connections.
• Graph Construction: Edges are formed from hit $i$ to hit $j$ if the distance $|\vec{S}_i - \vec{T}_j|$ falls within a fixed radius (determined via Fixed Radius Nearest Neighbor search). This naturally yields a directed graph that aligns with the physical directionality of particle tracks.

Experimental Setup
The methods were evaluated using ATLAS ITk simulation data for $t\bar{t}$ events at the High-Luminosity LHC (HL-LHC) with an average pileup of $\langle\mu\rangle = 200$.

• Models: Both Simple and Double Metric Learning models utilized nearly identical architectures (4 hidden layers, 1024 neurons, tanh activation) to ensure a controlled comparison. The primary difference was the output dimensionality: Simple ML output a single 12D vector, while Double ML output two 12D vectors (source and target). Both models had approximately 3.2 million trainable parameters.
• Evaluation: Graphs were constructed using Fixed Radius Nearest Neighbor (FRNN) algorithms. Performance was measured by graph construction efficiency ($\epsilon$), defined as the ratio of reconstructed true edges to total ground-truth edges, and mean graph size (number of edges).

Key Results

• Efficiency and Graph Size: At comparable levels of total construction efficiency ($\epsilon \approx 0.997$), Double Metric Learning produced undirected graphs with a smaller mean size ($6.59 \times 10^6$ edges) compared to Simple Metric Learning ($6.97 \times 10^6$ edges).
• High-$p_T$ Performance: Double Metric Learning demonstrated significantly higher graph construction efficiency for particles with high transverse momentum ($p_T$) compared to the Simple Metric Learning approach.
• Edge Purity: The ratio of "hopping edges" (connections between non-successive hits) to true reconstructed edges was lower for Double Metric Learning. The asymmetric metric successfully penalized non-physical connections that skip track layers.
• Directionality: The directed graphs produced by Double Metric Learning showed high alignment with physical track directionality. The number of edges in the directed version was similar to the undirected version, indicating that the model rarely produced redundant reversed edges.

Significance and Claims
The paper claims that Double Metric Learning successfully resolves the training conflicts inherent in applying standard Metric Learning to chain-connected graphs. By learning dual latent representations, the method enables the construction of directed graphs that inherently respect the physical directionality of particle tracks.

The authors assert that this approach offers two primary advantages:

1. Improved Performance: It enhances graph construction efficiency, particularly for high-$p_T$ particles, and reduces the inclusion of non-physical hopping edges.
2. Computational Efficiency: By reducing the number of redundant edges and enabling smaller graph sizes while maintaining high purity, the approach lowers the computational overhead of message-passing operations. This is presented as a critical improvement for graph neural networks operating in high-luminosity environments.

The paper notes that while the current implementation produces directed graphs, future work is required to develop arbitration logic to handle bidirectional edges or closed loops, as some downstream algorithms (e.g., walkthrough) require strictly directed acyclic graphs. The authors also suggest that this method could benefit object-condensation frameworks by transitioning from cluster-based to chain-based connectivity to reduce graph size.