This is SParTA: A modular, efficient and scalable transformer architecture for jet-tagging

Rajneil Baruah; Subhadeep Mondal; Sunando Kumar Patra; Satyajit Roy

hep-phOct 19, 2026

This is SParTA: A modular, efficient and scalable transformer architecture for jet-tagging

Authors: Rajneil Baruah, Subhadeep Mondal, Sunando Kumar Patra, Satyajit Roy

Organizations: Department of Physics, SEAS, Bennett University, Greater Noida, Uttar Pradesh, 201310, India · Department of Physics, Bangabasi Evening College, Kolkata, 700009, West Bengal, India

Abstract

We present the Sparse Particle-interaction Transformer Architecture (SParTA), a modular, scalable, and efficient transformer for jet tagging at high-energy colliders. Exploiting the steeply falling transverse-momentum ( $p_T$ ) spectrum of jet constituents, SParTA factorizes constituent processing into two topological regimes. Ordered by descending $p_T$ , only the $N$ hardest constituents (the hard core) participate in quadratic $O(N^2)$ differential interaction attention driven by pairwise physical and learned relational features. The remaining soft constituents are integrated via linear $O(lN)$ cross-attention to the hard core, preserving soft-radiation and color-flow information while avoiding the $O(l^2)$ memory bottleneck of full self-attention over all $l$ constituents. Pairwise interaction features are computed dynamically in-network from constituent four-momenta and conditioned on global jet-level substructure observables. On the quark-gluon tagging benchmark, a 300K-parameter SParTA model achieves an AUC of 92.02% and accuracy of 85.06%, matching the 2.14M-parameter Particle Transformer (ParT) at one-seventh the parameter count and outperforming IAFormer by 21% in background rejection at 30% signal efficiency. On the boosted top-tagging benchmark, SParTA outperforms ParT and ParticleNet with 511K parameters, while we analyze the architectural trade-offs relative to fully dense multi-layer interaction networks.

Figures & tables

Panel of Figure 1 — Figure 1: Schematics of the SParT Architecture.

	Quark–gluon	Top tagging
Signal / background	$q$ / $g$	$t$ / QCD
Generator	Pythia 8	Pythia 8
Detector sim.	none	Delphes (ATLAS)
Jet radius $R$	0.4	0.8
Jet $p_{\mathrm{T}}$ (GeV)	500–550	550–650
Rapidity cut	$\|y\|<1.7$	$\|\eta\|<2$

Table 1: Properties of the two datasets and of the inputs built from them.

	Feature	Description
1	$\Delta\eta=\eta-\eta_{\mathrm{jet}}$	pseudorapidity offset from the jet axis
2	$\Delta\phi=\phi-\phi_{\mathrm{jet}}$	azimuthal offset from the jet axis
3	$m$	constituent mass
4	$\log p_{\mathrm{T}}$	transverse momentum
5	$\log E$	energy
6	$\log(\pt/ \pt^{\mathrm{jet}})$	momentum fraction

Table 2: Per-constituent input features. The PID feature is available for the QG dataset only (

v=9

); top tagging uses the first eight (

v=8

Group	Features	No.
Jet kinematics	$p_{\mathrm{T}}^{\mathrm{jet}}$ , $\eta_{\mathrm{jet}}$ , $\phi_{\mathrm{jet}}$ , $m_{\mathrm{jet}}$	4
Multiplicity	$N_{\mathrm{cons}}$	1
N-subjettiness ( Thaler and Van Tilburg, 2011 ; Thaler and Van Tilburg, 2012 )	$\tau_{21}^{\beta}$ , $\tau_{32}^{\beta}$ , $\tau_{43}^{\beta}$ , $\tau_{54}^{\beta}$ , $\tau_{65}^{\beta}$ for $\beta=0.5$ , $1.0$ , $2.0$	15
$C$ -series ( Larkoski et al., 2013 )	$C_{1},C_{2},C_{3}$ ; $\beta=1,2$	6
$D$ -series ( Larkoski et al., 2014 ; Larkoski et al., 2016 ; Larkoski et al., 2015 ; Bhattacherjee et al., 2024 )	$D_{2}^{(\alpha,\beta)}$ , $\alpha,\beta\in\{1,2\}$ ; four $D_{3}$ variants	8
$U$ -series ( Moult et al., 2016 )	$U_{1},U_{2},U_{3}$ ; $\beta=0.5,1,2$	9

Table 3: The 53 jet-level (global) features used for both datasets.

\beta

is the angular exponent;

(\alpha,\beta)

label the two exponents of

D_{2}^{(\alpha,\beta)}

Figure 2: Schematics of the SParTA Block.

Hyperparameter	Symbol	Values scanned
Embedding dimension	$E$	40, 60, 80
Hard-core size	$N$	60, 80
Feed-forward width	$F$	$3E$ , $4E$
Number of SParTA blocks	$B$	2, 4, 6

Table 4: Hyperparameters scanned. The number of attention heads is fixed at 20.

Rank	$B$	$E$	$N$	Dp	FFN	Params (k)	FLOPs (M)	AUC (%)	$R_{50}$	$R_{30}$
1	6	40	80	0.2	120	185.3	416.1	91.86	46.2	126.5
2	6	80	80	0.3	320	757.3	695.8	91.80	44.8	124.1
3	6	40	80	0.2	160	214.9	424.4	91.79	45.1	126.3
4	4	40	80	0.2	160	146.9	293.1	91.78	45.4	121.4
5*	4	60	80	0.3	240	300.4	376.5	91.77	45.2	123.1
6	4	80	80	0.3	320	511.5	474.6	91.77	44.9	126.5

Table 5: All 30 scanned SParTA configurations on the QG test set, ranked by AUC, before fine-tuning.

B

: blocks;

E

: embedding dimension;

N

: hard-core size; Dp: dropout; FFN: feed-forward width

F

R_{50}

and

R_{30}

are the background rejections

1/\epsilon_{B}

\epsilon_{S}=0.5

and

0.3

. ⋆ Configuration used, after fine-tuning, for the results of Table 7 .

Step	$n$	$\langle\Delta\mathrm{AUC}\rangle$	Range	Up	$\langle\Delta R_{30}\rangle$
$B\!:2\to 4$	10	$+0.202$	$[+0.08,+0.43]$	10	$+5.5$
$N\!:60\to 80$	12	$+0.123$	$[+0.01,+0.26]$	12	$+1.1$
$F\!:3E\to 4E$	15	$+0.059$	$[-0.08,+0.33]$	10	$+2.7$
$B\!:4\to 6$	10	$-0.011$	$[-0.28,+0.14]$	6	$-0.5$
$E\!:60\to 80$	6	$-0.023$	$[-0.21,+0.09]$	3	$-0.5$
$E\!:40\to 60$	12	$-0.105$	$[-0.28,-0.01]$	0	$-3.5$

Table 6: Change in test AUC (percentage points) and in

R_{30}

for one hyperparameter step, over all

n

pairs of configurations that differ only in that step; “Up” counts the pairs whose AUC improved.

F

is compared as a multiple of

E

; the

E\!:60\to 80

pairs all have

N=80

Figure 3: Background rejection against cost for the scanned SParTA configurations on the QG test set. Point colour indicates the number of blocks (blue: 2, green: 4, orange: 6); red points joined by the dashed line form the Pareto front.

Figure 4: Main-effect variance fraction $\eta_{i}^{2}$ , Eq. ( 6 ), of each scanned hyperparameter for the test AUC (yellow), $R_{50}$ (blue) and the number of trainable parameters (green), over the 30 configurations of Table 5 .

Figure 5: Linear CKA, Eq. ( 7 ), between the outputs of the six blocks (B1–B6) of the fine-tuned $E_{60}H_{20}N_{80}F_{4}D_{2}B_{4}$ model on the QG test set.

		Accuracy (%)		AUC (%)		$1/\epsilon_{B}$		# of Params
		Best	Dist.	Best	Dist.	$\epsilon_{S}=0.5$	$\epsilon_{S}=0.3$
	ParticleNet Qu and Gouskos (2020)	$84.0$		$91.16$		$39.8(2)$	$98.6(1.3)$	370K
	PFN Komiske et al. (2019b)	–		$90.52$		$37.4(7)$	–	86.1K
	rPCN Shimmin (2021)	–		$90.81$		$38.6(5)$	–	–
Lorentz Invariance-	PELICAN Bogatskiy et al. (2024)	$\mathbf{85.51}$		$\mathbf{92.52}$		$\mathbf{52.3(3)}$	$\mathbf{149.8(2.4)}$	208K
based	LorentzNet Gong et al. (2022)	$84.40$		$91.56$		$42.4(4)$	$110.2(1.3)$	224K

Table 7: Performance of jet taggers on the QG test set. Numbers for SParTA are from this work: “Best” is the best of [

10

] trainings and “Dist.” the mean and standard deviation over them. All other numbers are taken from the cited references. The subscript full marks models that use additional weights for charged and neutral hadrons ( Qu et al., 2022b ) , and f.t. marks models pre-trained on JetClass and fine-tuned. Bold marks the best value in each column and italics the best within each group.

		Accuracy (%)		AUC (%)		$1/\epsilon_{B}$		# of Params
		Best	Dist.	Best	Dist.	$\epsilon_{S}=0.5$	$\epsilon_{S}=0.3$
	ParticleNet Qu and Gouskos (2020)	$94.0$		$98.58$		$397(7)$	$1615(93)$	370K
	PFN Komiske et al. (2019b)	–		$98.19$		$247(3)$	$888(17)$	86.1K
	rPCN Shimmin (2021)	–		$98.40$		$364(9)$	$1642(93)$	–
Lorentz Invariance-	PELICAN Bogatskiy et al. (2024)	94.26		98.70		–	2250(75)	208K
based	LorentzNet Gong et al. (2022)	$94.2$		$98.68$		$498(18)$	$2195(173)$	224K

Table 8: Performance of jet taggers on the top-tagging test set. Conventions as in Table 7 .

Appendix figures & tables1 asset

Supplementary material from the paper’s appendix.

Appendix

Feature	Mixer-Net	ABCNet	ECT	L-GATr	ParT	MIParT	IAFormer	SParTA
Symmetry and representation
Exact Lorentz equivariance	–	–	–	✓	–	–	–	–
Longitudinal-boost-invariant pairwise interaction inputs	–	–	–	–	✓	✓	✓	✓
Geometric algebra (multivector) representation	–	–	–	✓	–	–	–	–
Tunable symmetry breaking (reference vectors)	–	–	–	✓	–	–	–	–
Permutation invariance	✓	✓	✓	✓	✓	✓	✓	–

Appendix

Table 10: AB: Additive Bias; AS: Attention Score; GC: Graph Convolution ✓ indicates the feature is present as described in the corresponding publication (and, where available, its reference implementation); – indicates it is absent or not applicable.

\dagger

Feature not implemented in the current architecture, but is an option.

*

A convolution operation takes place when the interaction tensor is used in the Particle Attention Blocks.

\sim

denotes architectures that learn pairwise interaction features from embedded features but do not combine them with the physical interaction features in a single tensor. Marks refer to the architecture as published; several reference implementations expose additional options that are disabled in the published configuration. The “ Activation in FFN ” row reports the nonlinearity of the position-wise feed-forward sub-layer only; embedding-MLP and attention-coefficient activations may differ.

Transformer

Transformers are neural architectures that use attention to model relationships among sequence elements in language, vision, and scientific data.

See all 540 papers

Patch Hierarchical Attention Transformer for Efficient Particle Jet Tagging
May 20, 2026Aaron Wang, Zihan Zhao, Alan Xia +5Jet Tagging Efficient Attention
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
Particle-Lund Multimodality in Jet Taggers
May 26, 2026Loukas Gouskos, Benedikt MaierUnified Multimodal Models Jet Tagging
Brown University, Providence, USA · Imperial College of Science, Technology and Medicine, London, United Kingdom
B-jet Tagging Using a Hybrid Edge Convolution and Transformer Architecture
Mar 22, 2026Diego F. Vasquez Plaza, Vidya ManianJet Tagging High-Energy Physics
Department of Electrical and Computer Engineering, University of Puerto Rico Mayaguez, PR 00681-9000, USA

May 20, 2026J/K move · Enter open · S save

Patch Hierarchical Attention Transformer for Efficient Particle Jet Tagging

Aaron Wang, Zihan Zhao, Alan Xia, Chang Sun+4

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

Open paper AskPDF arXiv

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.

Abstract

Transformer

Explore similar work

Patch Hierarchical Attention Transformer for Efficient Particle Jet Tagging

Particle-Lund Multimodality in Jet Taggers

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