[Paper Review] Particle Transformer for Jet Tagging

H. Qu, C. Li+1 more|arXiv (Cornell University)|Feb 8, 2022

Astrophysics and Cosmic PhenomenaPhysics and Astronomy45 citations

TL;DR

This paper introduces JetClass, a large jet tagging dataset, and ParT, a Transformer-based model with pairwise interaction-aware attention that outperforms prior methods on JetClass and improves benchmarks after fine-tuning.

ABSTRACT

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.

Motivation & Objective

• Motivate the need for a large-scale public jet tagging dataset to advance deep learning methods in collider physics.
• Propose a Transformer-based architecture (ParT) that incorporates pairwise particle interactions in attention for improved tagging performance.
• Provide a comprehensive evaluation on JetClass and demonstrate transfer gains on established jet tagging benchmarks through fine-tuning.

Proposed method

• Define JetClass with 100 million jets across 10 jet types, including previously unexplored categories.
• Introduce ParT, a Transformer variant that uses two inputs: per-particle features and per-pair interaction features.
• Augment the particle self-attention (P-MHA) with a pairwise interaction bias U added to pre-softmax attention weights.
• Implement a class attention mechanism to produce a global jet representation for final classification.
• Encode pairwise interaction features derived from particle 4-vectors, transformed via an 8-dimensional interaction embedding.
• Train ParT end-to-end and compare against PFN, P-CNN, and ParticleNet; also assess pre-training + fine-tuning on other datasets.

Experimental results

Research questions

• RQ1 Does a Transformer-based architecture with pairwise particle interaction information outperform previous jet tagging models on a large-scale jet dataset?
• RQ2 What is the contribution of the pairwise interaction bias in attention (P-MHA) to overall performance?
• RQ3 Can pre-training ParT on JetClass improve performance on existing jet tagging benchmarks via fine-tuning?

Key findings

• ParT achieves higher accuracy and AUC on JetClass than ParticleNet and other baselines.
• Incorporating pairwise particle interactions via P-MHA yields noticeable gains over a plain Transformer.
• ParT provides substantial improvements when pre-trained on JetClass and fine-tuned on other jet tagging benchmarks.
• Across several signal types, ParT significantly improves background rejection compared to ParticleNet, with large gains for some jet types.
• Ablation shows replacing P-MHA with standard MHA reduces accuracy and background rejection, confirming the value of pairwise interaction modeling.