Patch Hierarchical Attention Transformer for Efficient Particle Jet Tagging

Abstract: 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 (\textsc{hls4ml}, JetClass, Top Tagging, and Quark--Gluon). Our code is available at https://github.com/aaronw5/PHAT-JeT.

• The paper introduces PHAT-JeT, which combines geometric message passing with hierarchical patch attention to achieve near-linear complexity for particle jet tagging.
• It demonstrates superior performance on multiple benchmarks, attaining high accuracy and background rejection at only about 1M FLOPs.
• Empirical evaluations show that the design maintains explicit intra-jet interactions and robust handling of constituent ordering under strict resource constraints.

Patch Hierarchical Attention Transformer for Efficient Particle Jet Tagging

Context and Motivation

Particle jet tagging is a central problem for real-time data reduction in high-energy physics detectors, notably at the LHC, where trigger systems must select signal events within stringent computational and latency constraints. While transformers have produced leading jet tagging performance, their quadratic attention cost limits on-device deployment at the trigger level. Existing efficient transformers—e.g., linear attention variants or cluster-based methods—sacrifice modeling expressivity, especially the capacity for explicit intra-jet constituent interactions, to achieve speed.

This work introduces the Patch Hierarchical Attention Transformer for Jet Tagging (PHAT-JeT) (2605.21789), which combines local geometric message passing with hierarchical patch-based attention, specifically targeting efficient jet tagging under hardware-constrained, low-latency requirements.

Architectural Components

PHAT-JeT architecture is composed of three synergistic modules: (1) a Geometric Message Passing (GMP) layer encoding detector-plane local geometry; (2) Local Patch-Based Self-Attention, computing full pairwise attention within patches; and (3) Hierarchical Patch-Level Global Attention, enabling efficient cross-patch communication. This design results in near-linear complexity with respect to jet constituent count, yet maintains explicit modeling of fine-grained particle interactions.

Figure 1: PHAT-JeT architecture schematic, illustrating GMP, patch partitioning, intra-patch attention, and global patch-token communication.

Geometric Message Passing

The GMP module projects jet constituents onto a coarse $(\eta,\phi)$ detector grid, aggregates features within cells, applies a depthwise 2D convolution, and samples features back per constituent, enforcing a lightweight inductive bias for local angular correlations. Empirically, GMP enhances class-specific discrimination and is robust to standard grid spacings. This spatial filtering produces geometry-aware embeddings before attention.

Local Patch-Based Self-Attention

Particles are partitioned, by fixed-ordering or batching, into patches of size $P \ll N$; exact multihead self-attention is applied within each patch. This bounds the main cost to $\mathcal{O}(NP)$, contrasting with the $\mathcal{O}(N^2)$ complexity of conventional transformers. The patch size hyperparameter directly tunes the accuracy-compute tradeoff. Crucially, explicit intra-patch pairwise structure is retained, unlike cluster/low-rank baselines.

Hierarchical Global Patch Attention

Patch tokens (pooling each patch, typically via mean or learned pooling) are constructed and global, multihead self-attention is applied across these. The globally-updated patch states are subsequently broadcast back to all particles in their patch, implementing efficient non-local communication. As the number of patch tokens is $O(N/P)$, this global stage introduces negligible overhead compared to the main local-attention block.

Empirical Evaluation

PHAT-JeT is evaluated across four standard benchmarks: hls4ml, JetClass, Top Tagging, and Quark–Gluon discrimination. The experimental regime enforces strict resource constraints, matching attention dimension and layer count across all baselines.

Patch Size–Efficiency Tradeoff

PHAT-JeT demonstrates substantial efficiency, achieving high tagging accuracy at $O(1\mathrm{M})$ FLOPs—about $0.4\%$ of a full Particle Transformer (ParT)—while outperforming alternative efficient baselines such as JEDI-Linear, Linformer, and SAL-T.

Figure 3: PHAT-JeT accuracy and computational cost as patch size increases; compared with JEDI-Linear baseline.

The empirical sweep in (Figure 2) shows strong accuracy retention even at modest patch sizes, saturating well before full global attention is required. Patch-token pooling via either mean or learned pooling performs similarly, indicating sufficient information transport with simple aggregation.

Attention Visualization and Robustness

Visualization of per-head and aggregate attention distributions for PHAT-JeT reveals that the learned attention patterns remain structured and physics-relevant across different patchings and constituent orderings (e.g., $k_T$ and random). Both local (intra-patch) and global (patch-token) attention exhibit sparsity and class sensitivity, with no degradation under consistent training/testing order.

Figure 5: Per-head attention maps for a top jet with $k_T$ sorting: each column shows Transformer, PHAT-JeT local, PHAT-JeT global, and PHAT-JeT combined attention.

Figure 2: Per-head attention for a top jet under random sorting; matches the structured patterns seen under $k_T$ sorting.

Figure 4: Combined attention projected onto the detector $P \ll N$0 plane under $P \ll N$1 ordering; thick intra-subjet and dashed inter-subjet edges highlight selective correlations.

Figure 6: Detector-plane attention pattern for matched random ordering; attention structure and particle relationships are preserved as in the $P \ll N$2 case.

Figure 7: Head-averaged attention for sample jets across all hls4ml classes, for each attention variant; PHAT-JeT demonstrates consistent class-dependent structure and selectivity.

These observations support the empirical result that PHAT-JeT's performance is robust to a range of patching and ordering strategies, provided training and inference are consistent.

Numerical Results and Claims

PHAT-JeT achieves maximum reported accuracy and background rejection among all resource-constrained models at trigger scale across all benchmarks. On the hls4ml dataset:

• PHAT-JeT: $P \ll N$3 accuracy, $P \ll N$4 ROC AUC, $P \ll N$5 average background rejection, at $P \ll N$6M FLOPs (compared to ParT: $P \ll N$7 at $P \ll N$8M FLOPs)
• Outperforms JEDI-Linear (strongest explicit prior baseline): $P \ll N$9 accuracy, $\mathcal{O}(NP)$0 background rejection at similar FLOPs

On JetClass: $\mathcal{O}(NP)$1 accuracy and $\mathcal{O}(NP)$2 ROC AUC at $\mathcal{O}(NP)$3M FLOPs ($\mathcal{O}(NP)$4 absolute accuracy improvement over best alternative); similar superiority is observed for binary Top Tagging and Quark–Gluon tasks. These results corroborate the claim that explicit intra-patch attention and hierarchical pooling close the gap with unconstrained transformers, within a strict real-time hardware budget.

Ablation and Inductive Bias Analysis

Ablation studies conclusively establish the additive contributions of GMP and hierarchical patch-token attention. Removing GMP or hierarchical global pooling degrades both overall accuracy and background rejection, with cumulative loss exceeding $\mathcal{O}(NP)$5 relative. Furthermore, the architecture is robust to consistent constituent ordering: models trained and tested on any fixed ordering (including random) yield indistinguishable accuracy; mismatched train-test orderings, however, cause sharp deterioration, confirming that the model is not permutation invariant at deployment.

Theoretical and Practical Implications

PHAT-JeT demonstrates that local, physics-motivated inductive biases—here, detector-plane grid geometric processing and structured attention factorization—are superior in the low-capacity, resource-constrained regime, compared to generic low-rank, cluster, or linearized approximations. Practically, this closes the performance gap between resource-efficient models and the full transformer baseline, directly enabling more aggressive but reliable jet tagging at the trigger edge. Moreover, the architecture is suitable for deployment on FPGAs and similar devices without requiring custom pipelining for large attention maps. The geometric message passing module is transferable and may serve as a generic inductive bias for point-cloud-based HEP inference tasks.

Future Directions

This architecture is expected to generalize to other point cloud and irregular geometric domains. Directions for future research include:

• Tight FPGA/ASIC mapping and real firmware demonstrations
• Further scaling to multimodal fusion, e.g., integrating tracks, calorimeter, or timing layers
• Combining with self-supervised pretraining or foundation model objectives for cross-task transfer
• Extension to anomaly detection and regression tasks in HEP and beyond
• Exploring permutation-equivariant or order-agnostic global stages for scenarios where strict order invariance at inference is desired

Conclusion

PHAT-JeT introduces a new paradigm for efficient particle jet tagging with hierarchical attention transformers, combining geometric message passing, patch-based local attention, and global token pooling. It consistently achieves best-in-class tagging accuracy and background rejection within the strict FLOPs and parameter constraints of real-time HEP triggering, and exhibits robustness and interpretability in attention mapping. The architectural principles should inform the design of efficient, domain-adapted models in high-throughput scientific inference applications, especially as experimental datasets and trigger rates continue to scale.