PHAT-JeT: Hierarchical Attention Transformer

• The paper introduces PHAT-JeT, a novel transformer model that achieves high jet tagging accuracy under constrained compute budgets through a hierarchical attention design.
• It integrates a physics-inspired geometric message passing module with patch-based self-attention, reducing FLOPs while preserving local and global context.
• PHAT-JeT demonstrates superior performance on benchmarks like HLS4ML and JetClass, balancing fine-grained interactions with resource-efficient inference.

Searching arXiv for the cited PHAT-JeT and related HSA papers to ground the article in current records. Patch Hierarchical Attention Transformer (PHAT-JeT) is a transformer-based architecture for efficient particle jet tagging under the latency and resource constraints of real-time trigger systems at the Large Hadron Collider. It is designed for the setting in which transformer models can offer the highest jet-tagging accuracy when compute is unconstrained, but standard quadratic self-attention is restrictive on trigger budget. PHAT-JeT addresses this by combining a physics-inspired geometric message-passing module on the detector plane with a hierarchical patch-based attention scheme that computes exact attention within small particle groups while preserving global context through patch-token communication. On four benchmarks—\textsc{hls4ml}, JetClass, Top Tagging, and Quark--Gluon—it is reported to achieve the highest accuracy and background rejection among resource-constrained jet-tagging models within a strict $\le 1.5$ M-FLOP budget and $\lesssim 10$ K-parameter regime suitable for FPGA triggers (Wang et al., 20 May 2026).

1. Problem formulation and design objective

PHAT-JeT is motivated by a specific systems problem: real-time jet tagging must identify short-lived particle decays in high-throughput detectors, while the trigger system deciding which collision events to store imposes strict latency and accuracy constraints (Wang et al., 20 May 2026). In this regime, the principal architectural tension is between classification performance and inference cost.

The PHAT-JeT paper states that transformer architectures achieve the highest jet-tagging accuracy when compute is unconstrained, but that their quadratic self-attention cost makes inference restrictive on trigger budget. It further states that existing efficient variants reduce the computational cost, but hinder the classification performance. PHAT-JeT is therefore positioned as a resource-constrained alternative that retains transformer-style contextual modeling while reducing the dominant attention cost.

A central feature of this objective is that efficiency is not treated as an auxiliary consideration. The reported comparisons are explicitly made under a restricted budget, and the model is evaluated against deployable baselines such as JEDI-Linear, SAL-T, Linformer, HEPT, PointTransformer V3, and PointNet, while also being compared with unconstrained references such as ParT and LorentzNet. This framing is important for interpreting the reported results: PHAT-JeT is not presented as an overall absolute best-performing jet tagger in an unconstrained sense, but as the strongest model within a specified deployment envelope.

2. Architectural composition

PHAT-JeT consists of three main computational modules: Physics-Inspired Geometric Message Passing (GMP), local patch self-attention, and hierarchical global attention (Wang et al., 20 May 2026). The full block also includes residual connections, LayerNorm, a position-wise feed-forward network, global mean pooling, and an MLP classifier head.

The GMP module operates on particle embeddings $X \in \mathbb{R}^{N \times c}$ together with per-particle coordinates $(\eta_i,\phi_i)$. Each particle is quantized onto a coarse $2$D grid of spacing $\delta$,

$$u_i = \left\lfloor \frac{\eta_i-\eta_{\min}}{\delta} \right\rfloor,\qquad v_i = \left\lfloor \frac{\phi_i-\phi_{\min}}{\delta} \right\rfloor,$$

where $\eta_{\min} = \min_k \eta_k$ and $\phi_{\min} = \min_k \phi_k$. A grid feature map $G \in \mathbb{R}^{H \times W \times c}$ is formed by summing particle features per cell,

$\lesssim 10$0

A depth-wise $\lesssim 10$1D convolution is then applied,

$\lesssim 10$2

and the convolved grid is sampled back at the particle locations,

$\lesssim 10$3

The output is mixed pointwise and added residually,

$\lesssim 10$4

with $\lesssim 10$5. The paper characterizes this as a learned spatial filter on the $\lesssim 10$6 plane that propagates local angular context to each particle embedding with negligible extra FLOPs.

The second component is hierarchical patch-based attention. The particle sequence is partitioned into patches of fixed size $\lesssim 10$7, with padding or truncation so that $\lesssim 10$8 patches are formed, each of shape $\lesssim 10$9. Within each patch, standard multi-head self-attention is computed exactly:

$X \in \mathbb{R}^{N \times c}$0

$X \in \mathbb{R}^{N \times c}$1

$X \in \mathbb{R}^{N \times c}$2

This exact intra-patch attention is the local part of the hierarchy.

Patch-token extraction provides the interface between local and global levels. Each patch is mean-pooled to a single token,

$X \in \mathbb{R}^{N \times c}$3

and these patch tokens are collected into $X \in \mathbb{R}^{N \times c}$4. Global inter-patch self-attention is then applied:

$X \in \mathbb{R}^{N \times c}$5

$X \in \mathbb{R}^{N \times c}$6

The resulting patch-level messages are broadcast back to particles in the corresponding patch,

$X \in \mathbb{R}^{N \times c}$7

This decomposition yields a two-scale communication pattern: exact token-token interactions are retained inside each patch, while global context is compressed into patch tokens. A plausible implication is that the model preserves fine-grained relational structure where the budget allows it, and uses a lower-resolution communication channel for long-range interactions.

3. Forward pass and computational structure

The single-block PHAT-JeT forward pass is specified explicitly. Given input particle embeddings $X \in \mathbb{R}^{N \times c}$8, the model applies, for each block $X \in \mathbb{R}^{N \times c}$9:

1. Geometric Message Passing:

$(\eta_i,\phi_i)$0

1. Local Patch Self-Attention: $(\eta_i,\phi_i)$1 is partitioned into patches $(\eta_i,\phi_i)$2 of size $(\eta_i,\phi_i)$3, and $(\eta_i,\phi_i)$4 is applied in parallel to each patch. The outputs are reassembled into a particle-level tensor $(\eta_i,\phi_i)$5.
2. Hierarchical Global Attention: patch tokens are extracted by mean pooling, $(\eta_i,\phi_i)$6 is computed, and each patch token is projected and added back to all particles in the patch to obtain $(\eta_i,\phi_i)$7.
3. Position-wise feed-forward network:

$(\eta_i,\phi_i)$8

After the final block, the jet-level classifier head applies

$(\eta_i,\phi_i)$9

The trigger-oriented configuration uses a single PHAT block, with $2$0 in the reported training settings. This is not merely an implementation detail: it indicates that the architecture is tuned for low-depth inference under strict budget constraints rather than for deep scaling.

The computational structure separates three distinct operations with different inductive biases. GMP injects detector-plane locality through quantization and depth-wise convolution. Local patch attention computes exact self-attention inside bounded particle groups. Global patch attention communicates only through pooled patch tokens. This organization is the defining sense in which PHAT-JeT is “hierarchical”: interaction granularity changes across stages instead of remaining uniformly tokenwise.

4. Complexity, resource profile, and deployment regime

The complexity analysis in the PHAT-JeT paper contrasts standard full self-attention with the model’s local-global decomposition (Wang et al., 20 May 2026). Standard full self-attention costs $2$1 per layer. For $2$2 and $2$3, the paper gives

$2$4

multiplies.

For PHAT-JeT, local attention costs

$2$5

and with $2$6, $2$7, $2$8, this becomes

$2$9

multiplies.

Global patch attention costs

$\delta$0

which, for $\delta$1 and $\delta$2, gives

$\delta$3

multiplies.

The total per block is therefore reported as approximately $\delta$4 versus $\delta$5 for full attention, corresponding to a $\delta$6 reduction in FLOPs. The paper further states that removing the quadratic term reduces on-chip memory needs and routing complexity, improving achievable clock frequency and end-to-end latency on FPGA.

The reported deployment regime is stringent: all four benchmark comparisons are framed under a strict $\delta$7 M-FLOP budget and $\delta$8 K-parameter regime suitable for FPGA triggers. This context is essential when interpreting both the performance gains over other deployable models and the remaining gap to large unconstrained models such as full ParT. A common misconception is to read efficient-attention results as claiming dominance over unconstrained transformers; the reported numbers do not support that reading. Instead, they support a narrower claim about performance under explicit hardware-oriented constraints.

5. Hyperparameters, training protocol, and benchmark performance

The key hyperparameters are fixed and compact. The number of PHAT blocks is $\delta$9. The hidden dimension is $$u_i = \left\lfloor \frac{\eta_i-\eta_{\min}}{\delta} \right\rfloor,\qquad v_i = \left\lfloor \frac{\phi_i-\phi_{\min}}{\delta} \right\rfloor,$$0 for \textsc{hls4ml} and Top/QG, and $$u_i = \left\lfloor \frac{\eta_i-\eta_{\min}}{\delta} \right\rfloor,\qquad v_i = \left\lfloor \frac{\phi_i-\phi_{\min}}{\delta} \right\rfloor,$$1 for JetClass. The number of attention heads is $$u_i = \left\lfloor \frac{\eta_i-\eta_{\min}}{\delta} \right\rfloor,\qquad v_i = \left\lfloor \frac{\phi_i-\phi_{\min}}{\delta} \right\rfloor,$$2. The patch size is $$u_i = \left\lfloor \frac{\eta_i-\eta_{\min}}{\delta} \right\rfloor,\qquad v_i = \left\lfloor \frac{\phi_i-\phi_{\min}}{\delta} \right\rfloor,$$3. The GMP grid spacing is $$u_i = \left\lfloor \frac{\eta_i-\eta_{\min}}{\delta} \right\rfloor,\qquad v_i = \left\lfloor \frac{\phi_i-\phi_{\min}}{\delta} \right\rfloor,$$4. Optimization uses Adam with $$u_i = \left\lfloor \frac{\eta_i-\eta_{\min}}{\delta} \right\rfloor,\qquad v_i = \left\lfloor \frac{\phi_i-\phi_{\min}}{\delta} \right\rfloor,$$5. The learning-rate schedule is a constant learning rate with batch-size ramp from $$u_i = \left\lfloor \frac{\eta_i-\eta_{\min}}{\delta} \right\rfloor,\qquad v_i = \left\lfloor \frac{\phi_i-\phi_{\min}}{\delta} \right\rfloor,$$6, early stopping on validation loss with patience $$u_i = \left\lfloor \frac{\eta_i-\eta_{\min}}{\delta} \right\rfloor,\qquad v_i = \left\lfloor \frac{\phi_i-\phi_{\min}}{\delta} \right\rfloor,$$7, and a maximum of $$u_i = \left\lfloor \frac{\eta_i-\eta_{\min}}{\delta} \right\rfloor,\qquad v_i = \left\lfloor \frac{\phi_i-\phi_{\min}}{\delta} \right\rfloor,$$8 epochs. Each configuration is run with $$u_i = \left\lfloor \frac{\eta_i-\eta_{\min}}{\delta} \right\rfloor,\qquad v_i = \left\lfloor \frac{\phi_i-\phi_{\min}}{\delta} \right\rfloor,$$9 random-seed replicates. Code is reported at https://github.com/aaronw5/PHAT-JeT (Wang et al., 20 May 2026).

On the \textsc{hls4ml} benchmark with five classes $\eta_{\min} = \min_k \eta_k$0, PHAT-JeT reports $\eta_{\min} = \min_k \eta_k$1, $\eta_{\min} = \min_k \eta_k$2, $\eta_{\min} = \min_k \eta_k$3, $\eta_{\min} = \min_k \eta_k$4, and $\eta_{\min} = \min_k \eta_k$5. The corresponding results listed for JEDI-Linear are $\eta_{\min} = \min_k \eta_k$6, $\eta_{\min} = \min_k \eta_k$7, $\eta_{\min} = \min_k \eta_k$8, $\eta_{\min} = \min_k \eta_k$9, and $\phi_{\min} = \min_k \phi_k$0; for the one-layer Transformer, $\phi_{\min} = \min_k \phi_k$1, $\phi_{\min} = \min_k \phi_k$2, $\phi_{\min} = \min_k \phi_k$3, $\phi_{\min} = \min_k \phi_k$4, and $\phi_{\min} = \min_k \phi_k$5; for SAL-T, $\phi_{\min} = \min_k \phi_k$6, $\phi_{\min} = \min_k \phi_k$7, $\phi_{\min} = \min_k \phi_k$8, $\phi_{\min} = \min_k \phi_k$9, and $G \in \mathbb{R}^{H \times W \times c}$0; for Linformer, $G \in \mathbb{R}^{H \times W \times c}$1, $G \in \mathbb{R}^{H \times W \times c}$2, $G \in \mathbb{R}^{H \times W \times c}$3, $G \in \mathbb{R}^{H \times W \times c}$4, and $G \in \mathbb{R}^{H \times W \times c}$5. The unconstrained reference ParT reports $G \in \mathbb{R}^{H \times W \times c}$6, $G \in \mathbb{R}^{H \times W \times c}$7, $G \in \mathbb{R}^{H \times W \times c}$8, $G \in \mathbb{R}^{H \times W \times c}$9 parameters, and $\lesssim 10$00 FLOPs.

On JetClass with ten classes, PHAT-JeT reports $\lesssim 10$01, $\lesssim 10$02, $\lesssim 10$03, $\lesssim 10$04 parameters, and $\lesssim 10$05 FLOPs. The listed baselines are Linformer at $\lesssim 10$06, $\lesssim 10$07, $\lesssim 10$08, $\lesssim 10$09, and $\lesssim 10$10; SAL-T at $\lesssim 10$11, $\lesssim 10$12, $\lesssim 10$13, $\lesssim 10$14, and $\lesssim 10$15; JEDI-Linear at $\lesssim 10$16, $\lesssim 10$17, $\lesssim 10$18, $\lesssim 10$19, and $\lesssim 10$20; Small ParT at $\lesssim 10$21, $\lesssim 10$22, $\lesssim 10$23, $\lesssim 10$24, and $\lesssim 10$25; and the full ParT reference at $\lesssim 10$26, $\lesssim 10$27, $\lesssim 10$28, $\lesssim 10$29, and $\lesssim 10$30.

On Top Tagging for binary $\lesssim 10$31 versus QCD classification, PHAT-JeT reports $\lesssim 10$32, $\lesssim 10$33, background rejection $\lesssim 10$34, $\lesssim 10$35 parameters, and $\lesssim 10$36 FLOPs. SAL-T is listed at $\lesssim 10$37, $\lesssim 10$38, $\lesssim 10$39; Linformer at $\lesssim 10$40, $\lesssim 10$41, $\lesssim 10$42; JEDI-Linear at $\lesssim 10$43, $\lesssim 10$44, $\lesssim 10$45; while ParT and LorentzNet report $\lesssim 10$46, $\lesssim 10$47, $\lesssim 10$48 and $\lesssim 10$49, $\lesssim 10$50, $\lesssim 10$51, respectively.

On Quark--Gluon, PHAT-JeT reports $\lesssim 10$52, $\lesssim 10$53, and $\lesssim 10$54. SAL-T and Linformer both report $\lesssim 10$55-level accuracy with $\lesssim 10$56 and $\lesssim 10$57 AUC, and background rejection $\lesssim 10$58; JEDI-Linear reports $\lesssim 10$59, $\lesssim 10$60, and $\lesssim 10$61; ParT and LorentzNet report $\lesssim 10$62, $\lesssim 10$63, $\lesssim 10$64 and $\lesssim 10$65, $\lesssim 10$66, $\lesssim 10$67.

These results support a precise interpretation. Within the stated low-FLOP, low-parameter trigger regime, PHAT-JeT is reported to achieve the highest accuracy and background rejection among deployable models. At the same time, unconstrained models remain stronger in absolute terms on several benchmarks. This distinction is central to the reported contribution.

6. Relation to hierarchical attention and conceptual context

A broader theoretical context is provided by “Hierarchical Self-Attention: Generalizing Neural Attention Mechanics to Multi-Scale Problems” (Amizadeh et al., 18 Sep 2025). That work formalizes multi-scale data as nested signals, defines a tree-structured signal hierarchy, and derives a hierarchical attention mechanism from the first principle of entropy minimization. In its formulation, the conditional entropy upper bound is optimized using a Boltzmann distribution, and the resulting update recovers standard Softmax attention in the flat case. For nested signals, the interaction energy between unrelated subtree roots $\lesssim 10$68 and $\lesssim 10$69 is defined by

$\lesssim 10$70

with a recursive subtree energy $\lesssim 10$71 and a resulting block-structured attention matrix that is stated to be exactly the KL-projection of the flattened Softmax matrix onto the space of hierarchical stochastic matrices. The HSA paper also provides a dynamic-programming algorithm with overall time $\lesssim 10$72 and space $\lesssim 10$73.

PHAT-JeT is not presented as an implementation of that dynamic-programming HSA mechanism. Its local and global attention modules are specified directly through standard attention equations over patches and patch tokens, not through the HSA recursion. Nevertheless, the comparison is conceptually informative. HSA argues that attention can be generalized to user-supplied multi-scale hierarchy without flattening; PHAT-JeT instantiates a concrete hierarchical decomposition for jet constituents by restricting exact attention to patches and routing global communication through patch tokens. This suggests a shared methodological direction: replacing a single uniform quadratic interaction pattern with structured multi-scale communication.

The HSA summary also includes recommendations for hierarchy configuration under the label “PHAT-JeT” in the context of vision-style patch hierarchies: a bottom level of non-overlapping $\lesssim 10$74 patches, higher levels formed by grouping adjacent windows, typical depth $\lesssim 10$75–$\lesssim 10$76, and decreasing branching factors such as $\lesssim 10$77, with the guidance to choose $\lesssim 10$78 so that $\lesssim 10$79 (Amizadeh et al., 18 Sep 2025). A plausible implication is that the PHAT-JeT jet-tagging architecture belongs to a larger design family in which efficiency is obtained by matching the attention pattern to a known hierarchy.

The limitations stated for HSA are also useful as interpretive context. HSA requires a fixed, known tree hierarchy a priori, does not learn it, and reports that early layers are sensitive in zero-shot replacement settings. PHAT-JeT likewise fixes its patch size $\lesssim 10$80 and detector-plane quantization spacing $\lesssim 10$81 rather than learning a hierarchy. This does not establish an identical limitation, but it indicates that hierarchy specification remains an important design choice rather than an automatically inferred structure.