Particle Transformer for Jet Tagging

• Particle Transformer (ParT) is a transformer-based architecture that processes unordered particle sets to achieve precise jet tagging in collider experiments.
• It embeds physics-motivated pairwise interactions into a sparse attention mechanism, delivering state-of-the-art classification and interpretable substructure identification.
• ParT offers computational efficiency and scalability on large-scale collider data, enabling real-time applications and improved physical analyses.

The Particle Transformer (ParT) is a transformer-based architecture specifically optimized for tasks involving sets of particles, most notably jet tagging in high-energy physics. ParT advances beyond standard graph and vision transformer approaches by embedding physics-motivated pairwise interactions directly into the attention mechanism and by exhibiting a salient sparsity—an emergent “nearly binary” pattern—within its particle–particle attentions. ParT achieves state-of-the-art classification and tagging performance, provides interpretable internal representations aligned with physical substructure, and offers computational efficiencies relevant for large-scale collider data and real-time applications.

1. Architectural Principles and Model Formulation

ParT processes unordered sets of particle features (“clouds”) representing jets. Each input consists of $N$ particles with $d$-dimensional features such as four-momentum components, particle identification flags, and detector-specific variables. These are embedded via a small per-particle MLP (with or without convolution), producing representations $X \in \mathbb{R}^{N \times d}$.

The core computational block is the Particle Multi-Head Attention (P-MHA), which generalizes standard transformer self-attention by incorporating a learnable pairwise interaction matrix. For each block and head, ParT computes:

$$S_{jk}^{(i)} = \frac{(x W^Q_i)_j \cdot (x W^K_i)_k}{\sqrt{d_k}} + U_{jk}^{(i)}$$

where $W^Q_i, W^K_i$ are learnable projections and $U_{jk}^{(i)}$ is the output of a pairwise MLP acting on physics-derived features:

• $$\Delta_{jk} = \sqrt{(\eta_j - \eta_k)^2 + (\phi_j - \phi_k)^2}$$
• $k_{T,jk} = \min(p_{T,j}, p_{T,k}) \Delta_{jk}$
• $$z_{jk} = \frac{\min(p_{T,j},p_{T,k})}{p_{T,j} + p_{T,k}}$$
• $$m_{jk}^2 = (E_j + E_k)^2 - \|\mathbf{p}_j + \mathbf{p}_k\|^2$$

The attention weights are then

$d$0

yielding permutation-invariant representations, with residual connections, layer normalization, and pointwise feed-forward networks maintaining architectural depth and stability (Legge et al., 28 Nov 2025, Usman et al., 2024, Wang et al., 2024).

Most ParT variants use $d$1 blocks, $d$2 attention heads, and embedding dimensions $d$3, with adaptations to problem scale and dataset.

2. Emergence of Sparse (“Binary”) Attention

Comprehensive analyses show that ParT’s self-attention displays a pronounced sparse, nearly binary pattern: For complex jet tagging tasks (JetClass, Quark–Gluon), each query particle allocates nearly all its attention to a single particle, with $d$4 for one $d$5 and $d$6 elsewhere. This holds for $d$7 of queries, as measured by the criterion $d$8 (Legge et al., 28 Nov 2025, Wang et al., 2024).

A detailed attribution study disentangles the contributions of the physics-inspired pairwise bias ($d$9) and the typical $X \in \mathbb{R}^{N \times d}$0 term. In JetClass and Quark–Gluon, the ratio $X \in \mathbb{R}^{N \times d}$1 is $X \in \mathbb{R}^{N \times d}$2 almost everywhere, indicating that sparsity is driven predominantly by learned attention weights. The same “edge-like” one-to-one patterns appear in pre-softmax $X \in \mathbb{R}^{N \times d}$3 visualizations, and the addition of $X \in \mathbb{R}^{N \times d}$4 only perturbs the dominant connections for a minority of queries. In contrast, for simpler tasks with limited kinematic diversity (Top Landscape), ParT’s attention becomes smoother and the pairwise bias plays a greater role (Legge et al., 28 Nov 2025).

3. Physical Interpretability and Substructure Discovery

ParT’s sparse attention pattern provides a direct mapping between attention heads and physical correlations. In leptonic-top jets, trained ParT identifies the central lepton—even without explicit PID—by attending preferentially to its track (top fraction $X \in \mathbb{R}^{N \times d}$5 vs. $X \in \mathbb{R}^{N \times d}$6 for untrained models). In hadronic jets, different heads capture kinematic substructure, linking prongs corresponding to QCD-inspired jet splitting and organizing attention according to familiar observables such as $X \in \mathbb{R}^{N \times d}$7 scales, subjet angles, and energy asymmetries (Legge et al., 28 Nov 2025, Wang et al., 2024).

This emergent interpretability contrasts sharply with vision transformers, wherein attention is diffuse and lacks direct correspondence to known physics observables. In ParT, learned linkages mirror classic QCD subjet algorithms and enable post-hoc explanations for individual classification decisions (Wang et al., 2024).

4. Performance Benchmarks and Limitations

ParT achieves state-of-the-art results on large jet classification datasets:

• JetClass (10-way, 100M samples): accuracy $X \in \mathbb{R}^{N \times d}$8, macro ROC-AUC $X \in \mathbb{R}^{N \times d}$9 (Usman et al., 2024)
• Top Tagging (2M jets): accuracy $$S_{jk}^{(i)} = \frac{(x W^Q_i)_j \cdot (x W^K_i)_k}{\sqrt{d_k}} + U_{jk}^{(i)}$$0, ROC-AUC $$S_{jk}^{(i)} = \frac{(x W^Q_i)_j \cdot (x W^K_i)_k}{\sqrt{d_k}} + U_{jk}^{(i)}$$1, background rejection at 50% signal efficiency $$S_{jk}^{(i)} = \frac{(x W^Q_i)_j \cdot (x W^K_i)_k}{\sqrt{d_k}} + U_{jk}^{(i)}$$2 (Rai et al., 10 Aug 2025)
• Quark flavor tagging (6-way, ILC simulation): c-background acceptance at 80% b-efficiency: 0.48% (vs. 6.3% for conventional software), d-background at 0.14% (vs. 0.79%) (Tagami et al., 2024)

These gains are coupled to efficient training (O(10 GPU-hrs for flavor tagging) and tractable inference costs ($$S_{jk}^{(i)} = \frac{(x W^Q_i)_j \cdot (x W^K_i)_k}{\sqrt{d_k}} + U_{jk}^{(i)}$$31 ms/jet on GPU). For tasks with limited substructure or feature diversity, ParT’s binary-sparsity collapses, and physical biases ($$S_{jk}^{(i)} = \frac{(x W^Q_i)_j \cdot (x W^K_i)_k}{\sqrt{d_k}} + U_{jk}^{(i)}$$4) become more necessary (Legge et al., 28 Nov 2025).

5. Physics-Informed Biases and Extensions

The physics-inspired bias matrix $$S_{jk}^{(i)} = \frac{(x W^Q_i)_j \cdot (x W^K_i)_k}{\sqrt{d_k}} + U_{jk}^{(i)}$$5 (or $$S_{jk}^{(i)} = \frac{(x W^Q_i)_j \cdot (x W^K_i)_k}{\sqrt{d_k}} + U_{jk}^{(i)}$$6) augments ParT’s attention scores, encoding pairwise kinematic and Standard Model couplings. Experiments incorporating energy-dependent SM interaction strengths (“running matrix”) yield a further 10% absolute background rejection and $$S_{jk}^{(i)} = \frac{(x W^Q_i)_j \cdot (x W^K_i)_k}{\sqrt{d_k}} + U_{jk}^{(i)}$$716% increase in signal significance beyond purely kinematic bias (Builtjes et al., 2022). ParT and state-of-the-art graph architectures (e.g. ParticleNet) achieve comparable classification AUCs ($$S_{jk}^{(i)} = \frac{(x W^Q_i)_j \cdot (x W^K_i)_k}{\sqrt{d_k}} + U_{jk}^{(i)}$$8) but ParT retains strict permutation invariance and is typically computationally heavier at large $$S_{jk}^{(i)} = \frac{(x W^Q_i)_j \cdot (x W^K_i)_k}{\sqrt{d_k}} + U_{jk}^{(i)}$$9.

Variants like ParMAT introduce multi-axis and parallel attention mechanisms for improved scalability; quantized versions (BitParT) enable 1-bit weight and activation variants suitable for deployment on resource-constrained hardware without compromising tagging accuracy (Usman et al., 2024, Rai et al., 10 Aug 2025).

6. Computational Efficiency, Scaling, and Future Directions

ParT’s inherent attention sparsity points toward further computational optimizations. Constraining heads to top-$W^Q_i, W^K_i$0 attention (e.g., $W^Q_i, W^K_i$1), retraining achieves $W^Q_i, W^K_i$2 AUC with $W^Q_i, W^K_i$3, implying 4–10$W^Q_i, W^K_i$4 reduction in FLOPs without major performance loss (Wang et al., 2024). Removing or sparsifying the physics bias $W^Q_i, W^K_i$5 is viable except for a minority of interaction-dependent queries; this would further decrease cost. Approaches inspired by dynamical systems and ODE solvers (TransEvolve) “precompute” attention operators for up to 50% reduction in parameter count and 2–3$W^Q_i, W^K_i$6 training speedup on long sequences, often matching or exceeding the accuracy of regular Transformers (Dutta et al., 2021).

A plausible implication, given these findings, is that future ParT architectures could enforce sparsity or use learned pairwise biases exclusively, thereby increasing interpretability and efficiency in massive jet-tagging deployments.

7. Impact on Collider Physics and Downstream Applications

ParT has materially advanced precision measurement and event selection in collider experiments:

• In ILC studies, flavor tagging upgrades via ParT lead to orders-of-magnitude improvements in background suppression, directly translating into improved precision for Higgs couplings and self-coupling measurements (Tagami et al., 2024).
• At the LHC, ParT enhances signal significance in rare channels, offering efficiency gains equivalent to substantial increases in integrated luminosity (Builtjes et al., 2022).
• ParT’s interpretability facilitates robust post-hoc analyses crucial for experimental collaboration workflows.

Current research investigates ParT’s integration with online inference platforms, its role in extracting physics observables from attention maps, and its extension to other set-based physical systems. Attention pattern taxonomy and the origin of sparsity remain research frontiers.

---

References

• "[Why Is Attention Sparse In Particle Transformer?]" (Legge et al., 28 Nov 2025)
• "[Particle Multi-Axis Transformer for Jet Tagging]" (Usman et al., 2024)
• "[Application of Particle Transformer to quark flavor tagging in the ILC project]" (Tagami et al., 2024)
• "[Investigating 1-Bit Quantization in Transformer-Based Top Tagging]" (Rai et al., 10 Aug 2025)
• "[Attention to the strengths of physical interactions: Transformer and graph-based event classification for particle physics experiments]" (Builtjes et al., 2022)
• "[Interpreting Transformers for Jet Tagging]" (Wang et al., 2024)
• "[Redesigning the Transformer Architecture with Insights from Multi-particle Dynamical Systems]" (Dutta et al., 2021)