B-jet Tagging Using a Hybrid Edge Convolution and Transformer Architecture

This paper introduces the Edge Convolution Transformer (ECT), a hybrid deep learning architecture that combines edge convolutions and self-attention mechanisms to achieve state-of-the-art b-jet tagging performance (0.9333 AUC) with low inference latency on ATLAS simulation data, outperforming both ParticleNet and pure transformer baselines.

The Gist

Imagine you are a detective at a massive, high-speed train station (the Large Hadron Collider). Every second, thousands of trains (particles) crash into each other, creating a chaotic explosion of debris. Your job is to look at the piles of wreckage (called jets) and figure out exactly what kind of train caused the crash.

Some trains are made of heavy, slow-moving freight cars (bottom quarks), some are medium-sized delivery trucks (charm quarks), and some are just lightweight bicycles or empty carts (light quarks).

The problem? The heavy freight cars and the delivery trucks leave behind very similar-looking wreckage. They both have "secondary" pieces that fell off a bit later than the main crash. Distinguishing between a heavy freight car and a delivery truck is incredibly hard, but it's crucial for solving the biggest mysteries of the universe.

This paper introduces a new detective tool called ECT (Edge Convolution Transformer). Think of it as a super-smart AI detective that combines two different ways of thinking to solve the case better than anyone else.

The Two Detective Styles

Before this new tool, detectives used two main strategies:

1. The "Neighborhood Watch" (ParticleNet): This detective looks at a specific piece of debris and asks, "Who are my immediate neighbors?" It builds a map of who is standing next to whom. This is great for spotting local patterns, like a cluster of broken glass that fell from a specific spot.
2. The "Big Picture" Observer (Transformer): This detective steps back and looks at the whole crime scene at once. It asks, "How does the energy flow across the entire pile of wreckage?" It connects dots that are far apart, noticing the overall shape and structure of the crash.

The New Hybrid Detective: ECT

The authors realized that to solve the hardest cases (telling the difference between the heavy freight car and the delivery truck), you need both skills. You need to see the local details and the big picture simultaneously.

So, they built ECT, a hybrid detective that does both at the same time:

• Step 1: The Local Scan (Edge Convolution): First, the AI zooms in on small groups of particles. It looks at how they are arranged in space, just like checking if a group of people are huddled together in a tight circle. This helps it spot the tiny, specific "displaced" tracks left by heavy particles.
• Step 2: The Global Scan (Transformer): Next, the AI zooms out. It uses a "self-attention" mechanism (like a spotlight that can focus on any part of the room instantly) to see how all the particles relate to each other across the entire jet.
• Step 3: The Verdict: The AI combines these two views. It takes the local clues and the global context, mixes them together, and makes a final decision: "This is definitely a bottom-quark jet!"

Why Is This a Big Deal?

In the past, the "Neighborhood Watch" detectives were good at spotting the heavy freight cars but missed the subtle differences between them and the delivery trucks. The "Big Picture" detectives were great at spotting the light bicycles but sometimes missed the fine details needed to tell the heavy trucks apart.

ECT is the first to master both.

• The Result: In their tests, ECT became the best detective in the room. It correctly identified the heavy freight cars (bottom jets) 88.5% of the time when trying to distinguish them from delivery trucks (charm jets). The old "Neighborhood Watch" only got about 80%, and the "Big Picture" observer got about 86%.
• Speed: Even though it's doing twice the work (looking locally and globally), it's incredibly fast. It can analyze a jet in less than 0.06 milliseconds. That's faster than a human eye can blink, making it fast enough to be used in real-time at the LHC to decide which crashes to keep studying and which to ignore.

The Analogy of the "Displaced Vertex"

To understand why this is hard, imagine two people dropping a ball:

• The Light Quark (Bicycle): Drops the ball the instant they hit the ground. The ball lands right where they are.
• The Heavy Quark (Freight/Truck): They are heavy and slow. They hit the ground, stumble, and then drop the ball a few inches away.

The "displaced vertex" is that few inches of distance.

• Charm jets stumble a little bit (drop the ball 150 microns away).
• Bottom jets stumble a lot (drop the ball 460 microns away).

The difference is tiny—like trying to tell the difference between a 1-inch gap and a 3-inch gap in a dark room. The ECT model is like a detective with a high-powered microscope (EdgeConv) and a wide-angle lens (Transformer) that can measure that tiny gap perfectly while also understanding the context of the whole room.

Summary

This paper presents a new AI model that combines local detail-checking with global pattern-spotting. By doing so, it solves the "hard problem" of telling heavy particles apart from medium ones better than any previous method, all while running fast enough to help physicists discover new laws of physics in real-time.

Technical Summary

1. Problem Statement

In high-energy particle physics, specifically at the Large Hadron Collider (LHC), identifying the flavor of the parton that initiated a jet (bottom, charm, light quarks, or gluons) is critical for precision Standard Model measurements and searches for new physics.

• The Challenge: While distinguishing bottom jets ($b$-jets) from light jets is relatively straightforward due to the absence of secondary vertices in light jets, separating $b$-jets from charm jets ($c$-jets) is highly difficult. Both originate from heavy quarks with finite lifetimes, resulting in displaced secondary vertices. However, $b$-hadrons have a longer decay length ($c\tau_b \approx 460\,\mu\text{m}$) compared to $c$-hadrons ($c\tau_c \approx 150\,\mu\text{m}$).
• Current Limitations: Existing state-of-the-art models generally fall into two categories:
  1. Graph Neural Networks (e.g., ParticleNet): Excellent at modeling local geometric relationships (vertex topology) but may miss global correlations.
  2. Transformer Architectures (e.g., Particle Transformer/ParT): Excellent at capturing global jet-wide patterns via self-attention but may struggle with the fine-grained local vertex displacement differences required for $b$ vs. $c$ separation.
• Goal: Develop a model that achieves state-of-the-art performance in $b$-jet tagging (specifically $b$ vs. $c$ discrimination) while maintaining the low inference latency (< 1 ms) required for real-time High-Level Trigger (HLT) systems at the LHC.

2. Methodology: The Edge Convolution Transformer (ECT)

The authors propose ECT, a hybrid deep learning architecture that integrates the local feature extraction capabilities of Edge Convolution (EdgeConv) with the global context modeling of Transformers.

Data and Features

• Dataset: ATLAS simulation data ($pp \to t\bar{t}$ events at $\sqrt{s}=14$ TeV) generated via Pythia8 and Delphes.
• Input Features:
  • Track-Level (7 features): Transverse momentum ($p_T$), transverse ($d_0$) and longitudinal ($z_0$) impact parameters, their significances, and 3D impact parameters. These capture the displacement of tracks from the primary vertex.
  • Jet-Level (8 features): Global kinematics ($p_T$, $\eta$, $\phi$, mass) and vertex statistics (number of vertices, max displacement, max tracks per vertex).
• Preprocessing: Features are normalized (Z-score, log-transform, min-max), and tracks are zero-padded to a maximum of 40 per jet with a validity mask.

Architecture Overview

The ECT processes data through six sequential stages:

1. Feature Embedding: Track features are embedded into 128 dimensions via an MLP; jet-level features are processed via a separate MLP.
2. Local Feature Extraction (EdgeConv): Three EdgeConv blocks aggregate information from $K$-Nearest Neighbor (KNN) graphs constructed in $(\eta, \phi)$ space ($K=16$). This captures local geometric relationships crucial for identifying displaced vertices.
3. Global Interaction (Transformer): Four self-attention layers (8 heads) capture long-range dependencies between particles across the entire jet.
4. Jet-Level Aggregation: A learned "class token" attends to all particle representations via two class-attention layers, creating a permutation-invariant jet embedding.
5. Fusion: The class token embedding is fused with the jet-level feature embedding via element-wise addition.
6. Classification: A final linear classifier with Softmax outputs the probability for the target class.

Key Design Rationale: The EdgeConv blocks handle the local vertex topology (critical for $b$ vs. $c$), while the Transformer layers handle global jet patterns (critical for $b$ vs. light).

3. Key Contributions

1. Novel Hybrid Architecture: Introduction of the ECT model, which unifies EdgeConv and Transformer self-attention mechanisms into a single end-to-end trainable network.
2. Comprehensive Evaluation: Rigorous testing on ATLAS simulation data across three binary classification tasks:
   • $b$ vs. $c$ (most challenging)
   • $b$ vs. light
   • $b$ vs. $c$ + light (combined background)
3. Performance Benchmarking: Direct comparison against two leading baselines: ParticleNet (Graph-based) and Particle Transformer (ParT) (Attention-based).
4. Latency Analysis: Demonstration that the model meets strict LHC trigger latency requirements (< 0.060 ms per jet).
5. Insight into Architecture: Empirical proof that EdgeConv is essential for heavy-flavor separation ($b$ vs. $c$), while Transformers excel at light-jet rejection.

4. Results and Performance

The ECT model was evaluated on a test set of ~325,000 jets using an NVIDIA RTX A5000 GPU.

Metric	Task	ECT	ParticleNet	ParT
AUC	$b$ vs. $c$	0.8853	0.8023	0.8634
AUC	$b$ vs. light	0.9883	0.9451	0.9876
AUC	$b$ vs. $c$+light	0.9333	0.8904	0.9216
Inference Latency	Per Jet	0.060 ms	12.23 ms	0.146 ms

Key Findings:

• Superior $b$ vs. $c$ Discrimination: ECT outperformed ParticleNet by 8.3% and ParT by 2.2% in AUC for $b$ vs. $c$. At the standard "Medium" working point (1% misidentification rate), ECT achieved 65% signal efficiency, compared to 52% for ParticleNet and 60% for ParT.
• Light-Jet Rejection: ECT and ParT achieved nearly identical, excellent performance (>98.7% AUC), significantly outperforming ParticleNet.
• Efficiency: ECT achieved a throughput of ~17,380 jets/second. While ParT is slightly faster in inference, ECT is significantly faster than ParticleNet and offers a much better balance of accuracy and speed.
• Training Time: ECT trained in ~2.1 hours, comparable to ParT (1.5h) and much faster than ParticleNet (4.5h).

5. Significance and Conclusion

The paper demonstrates that hybrid architectures combining local and global feature extraction are superior for complex jet classification tasks.

• Physics Insight: The results confirm that distinguishing $b$-jets from $c$-jets relies heavily on local geometric relationships (captured by EdgeConv) to resolve subtle differences in decay lengths, whereas distinguishing heavy-flavor jets from light jets relies on global topological signatures (captured by Transformers).
• Practical Impact: The ECT model achieves state-of-the-art accuracy while maintaining inference latency well within the limits for real-time event selection at the LHC. This makes it a viable candidate for deployment in future High-Level Trigger systems, potentially improving the efficiency of physics analyses involving heavy-flavor jets (e.g., Higgs boson studies, top quark properties, and supersymmetry searches).

The authors conclude that the Edge Convolution Transformer represents a promising direction for heavy-flavor jet tagging, successfully bridging the gap between local vertex modeling and global jet context.