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

Diego F. Vasquez Plaza; Vidya Manian

hep-phMar 22, 2026

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

Authors: Diego F. Vasquez Plaza, Vidya Manian

Organizations: Department of Electrical and Computer Engineering, University of Puerto Rico Mayaguez, PR 00681-9000, USA

Abstract

Jet flavor tagging plays an important role in precise Standard Model measurement enabling the extraction of mass dependence in jet-quark interaction and quark-gluon plasma (QGP) interactions. They also enable inferring the nature of particles produced in high-energy particle collisions that contain heavy quarks. The classification of bottom jets is vital for exploring new Physics scenarios in proton-proton collisions. In this research, we present a hybrid deep learning architecture that integrates edge convolutions with transformer self-attention mechanisms, into one single architecture called the Edge Convolution Transformer (ECT) model for bottom-quark jet tagging. ECT processes track-level features (impact parameters, momentum, and their significances) alongside jet-level observables (vertex information and kinematics) to achieve state-of-the-art performance. The study utilizes the ATLAS simulation dataset. We demonstrate that ECT achieves 0.9333 AUC for b-jet versus combined charm and light jet discrimination, surpassing ParticleNet (0.8904 AUC) and the pure transformer baseline (0.9216 AUC). The model maintains inference latency below 0.060 ms per jet on modern GPUs, meeting the stringent requirements for real-time event selection at the LHC. Our results demonstrate that hybrid architectures combining local and global features offer superior performance for challenging jet classification tasks. The proposed architecture achieves good results in b-jet tagging, particularly excelling in charm jet rejection (the most challenging task), while maintaining competitive light-jet discrimination comparable to pure transformer models.

Figures & tables

Data	Jet feature	Total quantity
Train 62.23%	$p_{T}^{\text{jet}}$ , $\eta^{\text{jet}}$ , $\phi^{\text{jet}}$ , $M^{\text{jet}}$ , $N_{\text{trk}}$ , $N_{\text{vtx}}$ ,	1,071,940
Validation 18.88%	$L_{3D}^{\text{max}}$ , $N_{\text{trk}}^{\text{vtx,max}}$	325,289
Test 18.88%		325,290

Table 1: Total jet feature count

Data	Track feature	Total quantity
Train		7,624,594
Validation	$p_{T}^{\text{trk}}$ , $d0$ , $z0$ , $d_{0}/\sigma_{d_{0}}$ , $z_{0}/\sigma_{z_{0}}$ , $ip3D$ , $\text{IP}_{3D}/\sigma_{\text{IP}_{3D}}$	2,312,445
Test		2,311,995

Table 2: Total track feature count

Panel of Figure 1 — Figure 1: Distributions of jet-level features for $b$ -jets (red), $c$ -jets (green), and light jets (blue) in the ATLAS simulation dataset. Top left: Jet pseudorapidity ( $\eta^{\mathrm{jet}}$ ). Top right: Jet invariant mass ( $M^{\mathrm{jet}}$ ). Middle left: Jet azimuthal angle ( $\phi^{\mathrm{jet}}$ ). Middle right: Jet transverse momentum ( $p_{\mathrm{T}}^{\mathrm{jet}}$ ). Bottom: Jet flavor distribution. Solid lines represent the training set (62.2% of data), while dashed lines show the validation set (18.9%). All distributions are normalized to unit area for comparison. The broader $p_{\mathrm{T}}$ distribution of $b$ -jets reflects the higher mass of bottom quarks ( $m_{b}\approx 4.2$ GeV/ $c^{2}$ ), while vertex multiplicity differences are evident in the flavor distribution.

Panel of Figure 2 — Figure 2: Distributions of track-level features for charged particles associated with $b$ -jets (red), $c$ -jets (green), and light jets (blue). Top row: Track transverse momentum ( $p_{\mathrm{T}}^{\mathrm{trk}}$ ) showing $p_{\mathrm{T}}>1$ GeV/ $c$ selection, and track electric charge distribution. Second row: Transverse ( $d_{0}$ ) and longitudinal ( $z_{0}$ ) impact parameters, demonstrating displaced vertices characteristic of heavy-flavor decays. The pronounced tails for $b$ - and $c$ -jets arise from secondary vertex displacements ( $c\tau_{b}\approx 460~\mu$ m, $c\tau_{c}\approx 150~\mu$ m), while light jets peak sharply near zero, consistent with prompt tracks from the primary vertex. Bottom row: Track pseudorapidity ( $\eta^{\mathrm{trk}}$ ) and azimuthal angle ( $\phi^{\mathrm{trk}}$ ) showing angular distributions. where smaller uncertainties enable higher-significance discrimination of displaced vertices. Solid lines represent training data (7,624,594 tracks), dashed lines show validation data (2,312,445 tracks). All distributions are normalized to probability density. Heavy-flavor jets exhibit significantly broader impact parameter distributions, providing the primary discriminative power for $b$ -jet identification

Architecture
Parameter	ECT	ParticleNet	ParT
Embedding dimension	128	128	128
Attention heads	8	–	8
Self-attention blocks	4	–	8
Class-attention blocks	2	–	2
EdgeConv blocks	3	3	–

Table 3: Model hyperparameters and architectural configuration. ECT combines EdgeConv blocks (from ParticleNet) with transformer self-attention and class-attention mechanisms. All models trained on NVIDIA RTX A5000 GPU.

Figure 3: Architecture of the Edge Convolution Transformer (ECT) for $b$ -jet tagging, showing complete information flow and tensor dimensions at each stage.

Figure 4: Training (Train) and validation (Val) metrics for the ECT model on the $b$ vs $c$ + $light$ classification task over 100 epochs. Top left: Cross-entropy loss converges smoothly from 0.37 to 0.31 over the first 50 epochs, then stabilizes, indicating effective optimization without oscillations. Top right: Training (blue) and validation (orange) accuracy (acc) reach $\sim$ 87% with minimal gap (< 0.2%), demonstrating that the model generalizes well without overfitting. Bottom left: Area Under Curve (AUC) improves steadily from 0.92 to 0.93, plateauing around epoch 89 where early stopping was triggered (patience = 25). The small train-validation AUC gap confirms robust generalization. Bottom right: F1-score exhibits higher variance during early training due to threshold sensitivity, then stabilizes at 0.81, confirming balanced precision-recall trade-off. Dashed horizontal lines mark the maximum values achieved: training accuracy 87.2%, validation accuracy 87.0%, training AUC 0.930, validation AUC 0.928. These curves demonstrate stable convergence and effective regularization.

Model	Task	loss	Acc	AUC	F1-score	Inference	Total
						time (ms)	time (h:mm:ss)
ECT	$b$ vs. $c$	0.4303	81.72	0.8853	0.8183	0.060	1:17:53
	$b$ vs. $light$	0.1195	96.05	0.9883	0.9599	0.057	1:43:01
	$b$ vs. $c$ + $light$	0.3080	87.75	0.9333	0.8146	0.057	2:09:04
ParticleNet	$b$ vs. $c$	0.5402	73.47	0.8023	0.7545	12.231	4:27:37
	$b$ vs. $light$	0.2728	89.08	0.9451	0.8973	17.882	4:29:31

Table 4: Performance comparison on ATLAS simulation dataset. Accuracy, AUC, and F1-score evaluated on held-out test set (325,290 jets). Latency measured on NVIDIA RTX A5000 GPU with batch size 1024, averaged over 100 iterations after 10 warmup runs using CUDA events.

Figure 5: Misidentification rate versus signal efficiency for the ECT model across three

b

-jet classification tasks:

b

vs.

c

(orange),

b

vs.

light

(green), and

b

vs.

c

light

(blue). Horizontal dashed lines indicate standard working points used in ATLAS analyses: Loose (10% mistag rate), Medium (1%), and Tight (0.1%). The

b

vs.

light

discrimination achieves the best performance, reaching signal efficiencies above 65% even at the Tight working point, reflecting the distinct signatures of

light

jets compared to heavy-flavor jets. The more challenging

b

vs.

c

task, where both jet types contain displaced vertices from heavy-hadron decays, shows reduced but still competitive performance. ECT achieves inference throughput of approximately 17,380 jets/s on a single NVIDIA RTX A5000 GPU, corresponding to less than 60

\mu

s per jet, well within LHC High-Level Trigger latency requirements.

Figure 6: Comparison of misidentification rate versus signal efficiency between ECT (solid lines) and ParticleNet (dotted lines) for

b

-jet tagging. Three classification scenarios are evaluated:

b

vs.

c

(orange),

b

vs.

light

(green), and

b

vs.

c

light

(blue). ECT consistently outperforms ParticleNet across all tasks and working points. The most significant improvement is observed in

b

vs.

c

discrimination, where ECT achieves an AUC of 0.885 compared to 0.802 for ParticleNet, representing a 10% relative improvement. At the Medium working point (1% mistag rate), ECT provides approximately 8 percentage points higher

b

-tagging efficiency than ParticleNet for the combined

b

vs.

c

+light classification. This improvement stems from ECT’s hybrid architecture, which combines ParticleNet’s edge convolution operations for local geometric feature extraction with transformer self-attention for capturing global track correlations.

Figure 7: Comparison of misidentification rate versus signal efficiency between ECT (solid lines) and Particle Transformer (ParT, dash-dotted lines) for

b

-jet tagging. Both architectures employ transformer-based attention mechanisms; however, ECT additionally incorporates EdgeConv blocks for explicit local feature extraction in the

(\eta,\phi)

coordinate space. ECT demonstrates improved charm-jet rejection (AUC: 0.885 vs. 0.863 for

b

vs.

c

) while maintaining comparable

light

-jet rejection performance (AUC: 0.988 for both models in

b

vs.

light

). The EdgeConv component enables ECT to better exploit local geometric relationships among tracks originating from displaced vertices, which is particularly beneficial for distinguishing the decay topologies of

b

-hadrons from those of

c

-hadrons. Despite the additional EdgeConv processing, ECT maintains comparable inference latency to ParT, making it suitable for real-time trigger applications.

Explore similar work

May 20, 2026

Patch Hierarchical Attention Transformer for Efficient Particle Jet Tagging

Real-time jet tagging is critical for identifying short-lived particle decays in the high-throughput detectors of the Large Hadron Collider, where real-time trigger systems responsible for deciding which collision events to store impose strict latency and accuracy constraints. While transformer architectures achieve the highest jet tagging accuracy when compute is unconstrained, their quadratic self-attention cost makes inference restrictive on trigger budget. Existing efficient variants reduce the computational cost, but hinder the classification performance. To address this limitation, we introduce the Patch Hierarchical Attention Transformer (PHAT-JeT), which combines two mechanisms: a physics-inspired geometric message-passing module that encodes local detector-plane structure, and a hierarchical patch-based attention scheme that computes exact attention within small particle groups while preserving global context through lightweight patch-token communication. Within a restricted budget, PHAT-JeT achieves state-of-the-art accuracy and background rejection among all resource-constrained jet tagging models on four benchmarks (hls4ml, JetClass, Top Tagging, and Quark--Gluon). Our code is available at https://github.com/aaronw5/PHAT-JeT.

Aaron Wang, Zihan Zhao, Alan Xia +5

University of Illinois Chicago Chicago, IL 60607, USA · University of California San Diego La Jolla, CA 92093, USA · California Institute of Technology Pasadena, CA 91125, USA +1

Jet Tagging Efficient Attention High-Energy Physics Attention Mechanisms

Ask PDF

May 26, 2026

Particle-Lund Multimodality in Jet Taggers

The Lund plane offers a physics-motivated, hierarchical representation of QCD radiation within jets, while transformer-based taggers have reached state-of-the-art performance by learning directly from raw particle constituents and their pairwise relations. We investigate whether transformers implicitly capture hierarchical QCD structure from constituent-level inputs, or whether explicit physics representations remain complementary. To test this, we introduce PLuM, a multimodal architecture that projects particle constituents and Lund plane splittings into a shared latent space, processing both jointly with a unified transformer. Cross-attention allows the model to probe whether structured QCD information provides discriminating power beyond what particles alone encode. We observe systematic gains for top-quark and

\mathrm{H}\to\mathrm{b}\bar{\mathrm{b}}

tagging, while finding no comparable improvement for

\mathrm{H}\to\mathrm{c}\bar{\mathrm{c}}

\mathrm{H}\to 4\mathrm{q}

topologies. This selective enhancement suggests that explicit hierarchical information about b-jet formation remains complementary to raw particle representations even in highly expressive architectures, while other topologies are already well-captured at constituent level. For high-impact LHC analyses such as Lorentz-boosted di-Higgs searches in the four

\mathrm{b}

quark final state (

\mathrm{H}\mathrm{H}(4\mathrm{b})

), the gains are substantial: at a

25\%

di-Higgs efficiency working point, PLuM achieves

25\%

higher background rejection than the baseline. Our results indicate that physically structured representations of QCD radiation retain discriminating value in the transformer era, motivating further study into how different aspects of jet dynamics are encoded by deep learning algorithms.

Loukas Gouskos, Benedikt Maier

Brown University, Providence, USA · Imperial College of Science, Technology and Medicine, London, United Kingdom

Unified Multimodal Models Jet Tagging High-Energy Physics Multimodal Fusion

Ask PDF

Jun 12, 2026

JetParticle-JEPA: An Efficient Self-Supervised Representation Learning method for Jet Tagging in High-Energy Physics

Jet tagging at the Large Hadron Collider increasingly relies on deep learning models trained on massive simulated datasets, leading to high computational costs and limited robustness to detector mismodeling. We introduce JetParticle-JEPA (JP-JEPA), a self-supervised Joint-Embedding Predictive Architecture that learns physically meaningful jet representations directly from continuous particle clouds without tokenization or reconstruction of raw inputs. Built on a Particle Transformer backbone, JP-JEPA predicts latent representations of masked particles while preserving fine-grained kinematic correlations. On the JetClass benchmark, JP-JEPA achieves performance comparable to fully supervised state-of-the-art methods on the full dataset, surpasses supervised baselines in low-label regimes, and significantly outperforms existing SSL approaches. On Top Quark and Quark-Gluon Tagging benchmarks, it remains on par with supervised methods. The learned representations also exhibit strong robustness to missing detector information and improved uncertainty behavior, highlighting JP-JEPA as a promising foundation-model framework for robust and data-efficient jet physics at the LHC.

Guillaume Letellier, Antonin Vacheret, Frédéric Jurie

LPCC · GREYC, Normandy University, Unicaen, ENSICAEN, UMR CNRS 6072, F-14000 Caen, France · LPC, Normandy University, Unicaen, ENSICAEN, IN2P3, UMR CNRS 6534, F-14000 Caen, France

Self-Supervised Learning Joint-Embedding Predictive Architecture Jet Tagging High-Energy Physics Data-Efficient Learning

Ask PDF