Scalable Multi-Task Learning for Particle Collision Event Reconstruction with Heterogeneous Graph Neural Networks

This paper proposes a scalable Heterogeneous Graph Neural Network (HGNN) that employs a multi-task learning paradigm to simultaneously perform particle vertex association and graph pruning, thereby significantly improving beauty hadron reconstruction performance and inference efficiency for complex particle collision events at the Large Hadron Collider.

The Gist

Imagine the Large Hadron Collider (LHC) as the world's most powerful, high-speed camera. It takes pictures of particles smashing into each other at nearly the speed of light. But here's the problem: the camera is getting faster and faster. In the past, it took a "snapshot" of a few particles. Now, and in the future, it's taking snapshots of thousands of particles all at once, creating a chaotic, crowded mess.

The scientists in this paper are trying to build a smarter way to sort through this chaos. They call their new tool a Heterogeneous Graph Neural Network (HGNN). Let's break down what that means using some everyday analogies.

1. The Problem: The "Crowded Party"

Imagine a massive, chaotic party where thousands of people (particles) are dancing, talking, and bumping into each other.

• The Goal: You want to find specific groups of friends (Beauty Hadrons) who arrived together, danced together, and left together.
• The Challenge:
  • Too many people: There are so many guests that it's hard to see who is with whom.
  • Wrong groups: Sometimes, a person from Group A accidentally bumps into a person from Group B, making it look like they are friends when they aren't.
  • Multiple entrances: The party has many different doors (Primary Vertices). People enter from different doors, and you need to know which door each person came through to figure out their story.
  • Speed limit: The party is moving so fast that you have to sort everyone out in a fraction of a second, or the data gets lost.

2. The Old Way: The "One-Size-Fits-All" Filter

Previously, scientists used a method called a "Homogeneous Graph Neural Network." Think of this like a security guard who treats everyone at the party exactly the same.

• The guard looks at every person and every interaction with the same set of rules.
• They try to guess who is with whom by looking at everyone's position.
• The Flaw: Because the guard treats a VIP (a particle from a specific decay) the same as a random guest, they often get confused. They also have to look at everyone before making a decision, which is too slow for the crowded future parties.

3. The New Solution: The "Specialized Detective" (HGNN)

The authors propose a new system, the Heterogeneous Graph Neural Network. Instead of treating everyone the same, this system acts like a team of specialized detectives who understand the different types of people and relationships at the party.

Here is how it works, step-by-step:

A. Understanding Different Roles (Heterogeneity)

In the old system, a "track" (a person's path) and a "vertex" (the door they entered) were just generic dots.
In the new system, the AI knows the difference:

• Tracks are the people walking around.
• Vertices are the doors they entered.
• Edges are the connections between them.
  The AI has special "glasses" for tracks and different "glasses" for doors. It understands that a connection between a person and a door is different from a connection between two people. This helps it make smarter guesses about who belongs to which group.

B. The "Smart Pruning" (Cutting the Noise)

This is the most creative part. Imagine the party is so crowded that the room is full of noise.

• Old Way: The detective tries to listen to every single conversation to find the right group. This takes forever.
• New Way (Pruning): The new AI has a "volume knob." As it processes the party, it instantly realizes, "Hey, these 90% of the people are just background noise; they aren't part of the story I'm looking for."
• It silences (prunes) the irrelevant people and conversations while it is thinking. It only keeps the important clues. This makes the process incredibly fast, even when the room is packed.

C. The "Multi-Task" Superpower

The old AI tried to do one thing at a time: first find the groups, then figure out who entered where.
The new AI is a Multi-Task Learner. It does everything at once, like a chef who is chopping vegetables, stirring the soup, and plating the dish simultaneously.

1. Reconstruction: It finds the specific groups of friends (Beauty Hadrons).
2. Association: It figures out which door (Vertex) each person came through.
3. Pruning: It ignores the noise.
   By doing these three things together, the AI learns from each task to help the others. For example, knowing which door a person came through helps the AI figure out which group they belong to.

4. The Results: Why This Matters

The paper shows that this new "Specialized Detective" is a huge upgrade:

• Faster: Because it "prunes" (ignores) the noise early on, it can process huge crowds much faster than the old methods. It's like finding a needle in a haystack by first throwing away 90% of the hay.
• Smarter: It is much better at figuring out which door people entered, which solves a major problem where particles get mixed up in crowded collisions.
• More Accurate: It reconstructs the "stories" of the particles (how they decayed) with much higher accuracy, finding the perfect groups that the old methods missed.

The Bottom Line

As the LHC gets more powerful and creates even more crowded particle collisions, we need a smarter way to sort the data. This paper introduces an AI that doesn't just look at the data; it understands the types of data, ignores the noise instantly, and solves multiple puzzles at the same time. It's the difference between trying to find a friend in a crowded stadium by shouting at everyone, versus having a smart guide who knows exactly who to talk to and who to ignore.

Technical Summary

1. Problem Statement

The Large Hadron Collider (LHC), specifically the LHCb experiment, is facing a "luminosity frontier" challenge. Upgrades to the experiment (Upgrade I and II) will increase collision rates by a factor of 10 to 50, leading to:

• High Particle Multiplicity: Events will contain up to 1,000 charged particle tracks, compared to ~150 in previous runs.
• Data Acquisition Strain: Storing full events is becoming infeasible due to storage limits (O(10 PB/year)) and latency constraints (O(100 ms/event)).
• Primary Vertex (PV) Misassociation: With multiple proton-proton interactions occurring simultaneously, tracks are frequently misattributed to the wrong primary vertex. This degrades the resolution of key observables, such as beauty-hadron decay flight distances, crucial for CP violation studies.
• Limitations of Current Methods: Existing approaches like the Deep Full Event Interpretation (DFEI) rely on multi-stage, homogeneous Graph Neural Networks (GNNs). These struggle with scalability, high latency, and do not natively solve the PV misassociation problem.

2. Methodology

The authors propose a Heterogeneous Graph Neural Network (HGNN) architecture designed for Multi-Task Learning (MTL) to simultaneously perform three tasks:

1. Beauty Hadron Reconstruction: Identifying decay chains of beauty hadrons.
2. Graph Pruning: Removing background nodes and edges to reduce computational load.
3. Primary Vertex (PV) Association: Correctly linking tracks to their originating interaction point.

Key Architectural Components

• Heterogeneous Graph Representation:
  • Unlike homogeneous GNNs that treat all nodes identically, this model uses distinct node types: Tracks and Primary Vertices (PVs).
  • It utilizes distinct edge types: Track-Track (indicating shared decay ancestry) and PV-Track (indicating origin).
  • This allows the network to learn relationships dynamically rather than appending static PV coordinates to track features (a limitation of previous methods).
• HGNN Layer Updates:
  • The architecture extends the Message Passing Neural Network (MPNN) framework by Battaglia et al.
  • It employs type-specific update functions ($\phi_{tr}, \phi_{pv-tr}$) and aggregation functions ($\rho$) for different node and edge types, enabling richer context-specific representations.
• Integrated Graph Pruning:
  • To ensure scalability, the model includes pruning layers within the HGNN structure.
  • It predicts scores ($\hat{y}$) for nodes and edges using MLPs. During training, these scores act as continuous weights in message passing (weighted message passing). During inference, a threshold ($y_{cut}$) is applied to hard-prune irrelevant connections.
  • This reduces the graph size early in the inference pipeline, significantly lowering latency.
• Multi-Task Loss Function:
  • The model is trained with a combined loss function:
    $$L = L_{CE}^{LCA} + \beta_{etr} L_{BCE}^{prune} + \beta_{vtr} L_{BCE}^{prune} + \beta_{pv-tr} L_{BCE}^{PV}$$
  • $L_{CE}^{LCA}$: Class-weighted cross-entropy for reconstructing the Lowest Common Ancestor (LCA) hierarchy of decay chains.
  • $L_{BCE}$ terms: Binary Cross-Entropy losses for pruning tracks/edges and associating PVs.
  • Weighted Message Passing: The model uses predicted pruning scores as weights during message passing, which stabilizes performance even when aggressive pruning is applied.

3. Key Contributions

1. Unified Framework: First application of HGNNs in particle physics to jointly solve reconstruction, pruning, and PV association in a single end-to-end framework.
2. Scalability via Pruning: Integration of pruning directly into the network layers allows the model to handle high-multiplicity events (up to 1,000 tracks) with sub-linear inference time scaling, meeting strict latency requirements.
3. Superior PV Association: By explicitly modeling PV-Track edges, the model solves the PV misassociation problem, which is critical for high-luminosity environments where simple "minimum impact parameter" methods fail.
4. Performance Gains: Demonstrates that multi-task learning improves the primary reconstruction task (LCA) by providing auxiliary signals (pruning and PV association) that help the model learn better representations.

4. Results

The model was evaluated on simulated LHCb Run 3 data (mimicking high-luminosity conditions).

• Reconstruction Performance:
  • The proposed HGNN achieved a 22.4% perfect reconstruction rate for inclusive beauty hadrons, compared to 4.7% for the previous DFEI method (a 4.8x improvement).
  • It significantly reduced "not isolated" reconstructions (background contamination) from 76.1% (DFEI) to 44.1%.
• Pruning Efficiency:
  • Inference Time: Early pruning (layers 1-3) reduced CPU inference time by 5x and GPU time by 2-3x for events with >400 tracks.
  • Signal vs. Background: Tight pruning ($\hat{y} > 0.2$) suppressed background tracks while maintaining high signal efficiency, even at high multiplicities.
• PV Association:
  • The HGNN achieved 99.88% accuracy in associating tracks to the correct PV, compared to 95.56% for the standard "minimum impact parameter" method.
  • For beauty hadrons, accuracy reached 99.85%, a massive improvement over the 96.14% baseline.
  • Performance remained robust even as the number of PVs per event increased (up to 15), whereas baseline methods degraded significantly.
• Exclusive Decays:
  • When trained with a small subset of exclusive decay modes, the model achieved >90% perfect reconstruction for rare topologies (e.g., $B^0 \to K^*\mu\mu$), demonstrating strong generalization.

5. Significance

• Enabling Future Physics: This architecture is essential for the LHCb Upgrade II, where data rates will be too high for traditional storage and reconstruction methods. It enables the retention of valuable physics events while discarding background in real-time.
• Precision Measurements: Improved PV association directly enhances the precision of time-dependent CP violation measurements and searches for decays with missing energy (e.g., neutrinos) by correctly determining the flight direction of beauty hadrons.
• Generalizability: The approach of using HGNNs for heterogeneous data with integrated pruning is applicable to other particle physics experiments and domains involving complex, multi-relational data structures.
• Efficiency: The combination of multi-task learning and weighted message passing allows for aggressive graph sparsification without sacrificing reconstruction accuracy, solving the trade-off between speed and precision.

In conclusion, this work presents a scalable, holistic solution for the next generation of particle physics data analysis, moving beyond sequential, homogeneous approaches to a unified, heterogeneous, and pruning-aware neural network architecture.