Source-linked AI summary
Particle Transformer for Jet Tagging
Huilin Qu, Congqiao Li, Sitian Qian
TL;DR
Jet tagging is a difficult particle-physics classification task, and progress is constrained by limited public datasets. The paper introduces the 100 M-jet JETCLASS dataset and Particle Transformer, whose pairwise-interaction attention substantially improves tagging over prior approaches and benchmark fine-tuning results.
Problem
Jet tagging is challenging because particle radiation creates cascades and smears the initiating particle’s characteristics, while existing public datasets are much smaller.
Method
The paper develops JETCLASS, a 100 M-jet dataset spanning 10 jet classes, and Particle Transformer, which augments attention with pairwise particle interactions.
Results
Particle Transformer substantially improves jet-tagging performance over previous state-of-the-art, and fine-tuned pre-trained models substantially improve results on two widely adopted benchmarks.
Takeaways & Limitations
JETCLASS provides a large open dataset, while Particle Transformer is proposed as a new jet-tagging baseline for future research.
Takeaways & Limitations
JETCLASS applies a decay-particle selection requiring all relevant quarks and charged leptons to satisfy ΔR(jet, q/ℓ) < 0.8.
Abstract
from arXiv · showhide
Jet tagging is a critical yet challenging classification task in particle physics. While deep learning has transformed jet tagging and significantly improved performance, the lack of a large-scale public dataset impedes further enhancement. In this work, we present JetClass, a new comprehensive dataset for jet tagging. The JetClass dataset consists of 100 M jets, about two orders of magnitude larger than existing public datasets. A total of 10 types of jets are simulated, including several types unexplored for tagging so far. Based on the large dataset, we propose a new Transformer-based architecture for jet tagging, called Particle Transformer (ParT). By incorporating pairwise particle interactions in the attention mechanism, ParT achieves higher tagging performance than a plain Transformer and surpasses the previous state-of-the-art, ParticleNet, by a large margin. The pre-trained ParT models, once fine-tuned, also substantially enhance the performance on two widely adopted jet tagging benchmarks. The dataset, code and models are publicly available at https://github.com/jet-universe/particle_transformer.
1. Introduction
Jet tagging classifies particle sprays to identify jets initiated by particles of interest, but radiation cascades obscure the initiating particle. The paper addresses dataset scale and modeling by introducing JETCLASS and Particle Transformer.
- Jet tagging classifies collimated particle sprays to distinguish jets arising from particles such as Higgs, W, Z, or top quarks.
- O(10) to O(100) particles can result from radiation cascades, smearing the initiating particle’s characteristics and making identification difficult.
- Particle Transformer is a Transformer-based jet-tagging architecture built on the JETCLASS dataset.
- Pairwise particle interactions augment attention, enabling ParT to outperform a plain Transformer and ParticleNet by large margins.
2. The JETCLASS Dataset
JETCLASS provides simulated, multiclass jet-tagging data with 10 jet categories and 100 M training jets, alongside features and evaluation metrics tailored to LHC analyses.
- JETCLASS contains 10 jet types, including background q/g jets and signal jets from top quarks or W, Z, and Higgs bosons.
- Jets are simulated with MAD-GRAPH5 aMC@NLO, PYTHIA, and DELPHES to approximate realistic ATLAS or CMS reconstructed jets.
- Each jet supplies all constituent particles, typically 10–100 particles, with kinematic, particle-identification, and trajectory-displacement features.
- The training set has 100 M jets equally distributed across 10 classes, with 5 M validation jets and 20 M test jets.
- Accuracy and AUC measure overall multiclass performance, while RejX is evaluated for each signal type against q/g background at specified TPR values.
- A factor-of-two increase in background rejection can yield about a 40% increase in discovery potential.
3. Related Work
Related jet-tagging work spans image, sequence, tree, graph, and particle-cloud representations, while public datasets remain smaller or narrower than JETCLASS. Transformers provide a broader architectural context.
- Deep-learning jet taggers use representations including images, sequences, trees, graphs, and permutation-invariant particle clouds.
- ParticleNet adapts Dynamic Graph CNN to represent jets as particle clouds.
- DeepAK8 uses a ResNet-based 1D CNN for tagging top-, Higgs-, W-, or Z-boson jets in CMS analyses.
- Existing public datasets range from approximately 500 k to 3.9 M signal or mixed jets and cover two to five classes in the listed examples.
- JETCLASS is substantially larger and more inclusive in jet types than existing datasets.
- Transformers have achieved strong results across natural-language-processing and computer-vision tasks, motivating their use as universal architectures.
4. Model Architecture
ParT processes particle features and pairwise interaction features with particle and class attention blocks. Its augmented attention incorporates pairwise interactions directly, while class attention aggregates particle information through a global class token.
- Particle and interaction inputs: ParT represents each jet with particle features and pairwise interaction features, embedding them into particle and interaction representations.Particle inputs have shape (N, C), while interaction inputs have shape (N, N, C′); both are projected by MLPs.
- Particle and interaction inputs: Particles are permutation invariant, so ParT omits ad-hoc positional encodings and includes spatial information directly in particle inputs.The model can also be viewed as a fully connected graph neural network whose particles are nodes and interactions are edge features.
- Particle interaction features: The interaction features are derived from particle energy-momentum 4-vectors and include four pairwise quantities transformed to (ln ∆, ln kT, ln z, ln m2).For particles a and b, the features are computed from their 4-vectors; logarithms address long-tail distributions.
- Particle attention block: The particle attention block uses normalized multi-head attention followed by a two-layer MLP, with residual connections after each stage.ParT replaces standard multi-head attention with P-MHA, while retaining the overall block structure based on NormFormer.
- Particle attention block: P-MHA adds the interaction matrix U to pre-softmax attention weights, allowing physics-motivated pairwise features to modify dot-product attention.Q, K, and V are linear projections of the particle embedding.
- Class attention block: The class attention block applies standard attention between a global class token and all particles after the particle-attention stack.The architecture uses a stack of particle attention blocks followed by class attention blocks.
5. Experiments
Experiments evaluate ParT on JETCLASS, compare it with established baselines, test dataset scale and model complexity, and assess transfer through fine-tuning on other benchmarks. ParT achieves the strongest reported performance while benefiting from pairwise interactions, large-scale training, and pre-training.
- JETCLASS results: ParT achieves the best performance on every JETCLASS metric and increases overall accuracy by 1.7% over ParticleNet.Background rejection improves by factors of 3 for t →bqq′, 2 for H →4q, and about 70% for H →c¯c.
- JETCLASS results: ParT’s background rejection exceeds ParticleNet’s across several signal types, while PFN and P-CNN lag substantially behind both models.The difference between the earlier baselines and ParT or ParticleNet reaches up to an order of magnitude in background rejection.
- Ablation and scaling: ParT’s pairwise-interaction attention improves accuracy by 1.2% and background rejection by 20–30% for most signals versus a plain Transformer.Replacing P-MHA with standard MHA discards the particle-interaction input.
- Ablation and scaling: Reducing the training set to 10% and 2% degrades ParT accuracy by 1.1% and 2.5%, respectively.The corresponding ParticleNet degradation is 0.7% when trained on a reduced dataset.
- Model complexity: ParT uses more than 5× as many trainable parameters as ParticleNet but 40% fewer FLOPs, while its FLOPs are 30% higher than ParT (plain).The additional computational cost relative to ParT (plain) mostly comes from pairwise-feature encoding, which scales quadratically with particle count.
- Fine-tuning for other datasets: Pre-trained ParT models achieve the best performance on both fine-tuning benchmark scenarios, whereas ParticleNet pre-training yields only a small improvement on top tagging.The quark-gluon benchmark evaluates both experimentally restricted and full particle-identification inputs.
6. Discussion and Conclusion
The work presents JETCLASS as a large, broad dataset for advancing deep learning in particle physics and introduces ParT as a substantially stronger jet-tagging architecture. ParT’s augmented self-attention incorporates physics-inspired pairwise interactions, while its full pairwise matrix increases computational time and memory consumption.
- JETCLASS dataset: 100 M simulated jets across 10 classes make JETCLASS about two orders of magnitude larger than existing public jet datasets.The dataset includes several jet types not previously studied with deep learning and is publicly available for broader particle-physics research.
- JETCLASS dataset: JETCLASS is positioned as a basis for unsupervised, self-supervised, generative, and regression research beyond jet tagging.The paper specifically mentions particle-physics applications including fast simulation and precision prediction of jet energy and momentum.
- Particle Transformer: ParT substantially improves jet-tagging performance over previous state-of-the-art models and is proposed as a baseline for future research.The architecture’s effectiveness is attributed mainly to augmented self-attention combining physics-inspired pairwise interactions with learned dot-product attention.
- Particle Transformer: Full pairwise interaction matrices increase computational time and memory consumption, motivating more efficient particle embeddings and self-attention methods.The paper identifies reducing this computational cost as an interesting direction for future research.