Interpreting Transformers for Jet Tagging

TL;DR

The paper tackles interpretability of the Particle Transformer (ParT) for jet tagging in LHC data. It studies how ParT learns from jet substructure by inspecting attention heat maps and eta-phi correlations, revealing a nearly binary attention pattern where each particle attends to at most one other. The analysis shows that ParT's focus on important substructures varies with decay mode, suggesting the model learns traditional jet-substructure observables, and demonstrates inter-subjet and intra-subjet interactions through attention graphs. These findings point to potential efficiency gains by pruning attention and motivate further exploration of the interaction matrix and intermediate layers to enhance both interpretability and computational efficiency in high-energy physics transformers.

Abstract

Machine learning (ML) algorithms, particularly attention-based transformer models, have become indispensable for analyzing the vast data generated by particle physics experiments like ATLAS and CMS at the CERN LHC. Particle Transformer (ParT), a state-of-the-art model, leverages particle-level attention to improve jet-tagging tasks, which are critical for identifying particles resulting from proton collisions. This study focuses on interpreting ParT by analyzing attention heat maps and particle-pair correlations on the $η$-$φ$ plane, revealing a binary attention pattern where each particle attends to at most one other particle. At the same time, we observe that ParT shows varying focus on important particles and subjets depending on decay, indicating that the model learns traditional jet substructure observables. These insights enhance our understanding of the model's internal workings and learning process, offering potential avenues for improving the efficiency of transformer architectures in future high-energy physics applications.

Figures (3)

• Figure 1: Distribution of attention scores (left) illustrating the binary nature. Heat map of attention scores for example jets in the $t\to b \ell \nu$ (center) and $t\to b q q'$ classes (right) between the constituent particles.
• Figure 2: Visualization of attention values in the $\eta$-$\phi$ plane for pairs of particles (✖: muon, ▲: charged hadron, ▼: neutral hadron, ●: photon, and ✚: electron) within jets. The transparency of each point is proportional to the transverse momentum of the particle ($p_{\text{T}}\xspace$). The intensity of the connecting lines represents the magnitude of the attention scores. Solid lines depict intra-subjet connections, while dotted lines represent inter-subjet connections. The plots in the first column correspond to $t \to b\ell\nu$, the second column to $t \to bqq'$, and the third column to $H \to 4q$ processes.
• Figure 3: Distribution of the proportion of attention values that attend to leptons in $t\to b \ell \nu$ (left). Distribution of the proportion of attention in between subjets in the $t\to b q q'$ (center) and $H\to 4 q$ classes (right)