IAFormer: Interaction-Aware Transformer network for collider data analysis

TL;DR

IAFormer tackles the challenge of efficiently leveraging pairwise particle interactions in collider data analysis. It introduces an interaction-aware Transformer with a dynamic sparse attention mechanism (differential attention) that replaces $Q$/$K$ with a trainable interaction matrix and selectively attends to the most informative constituents. The architecture achieves state-of-the-art or competitive performance on top tagging and quark–gluon tagging while using an order of magnitude fewer parameters and reduced FLOPs, with robust behavior across seeds. Interpretability analyses using attention maps and CKA demonstrate that IAFormer captures physically meaningful, layer-dependent representations, and the approach scales to larger JetClass benchmarks.

Abstract

In this paper, we introduce \texttt{IAFormer}, a novel Transformer-based architecture that efficiently integrates pairwise particle interactions through a dynamic sparse attention mechanism. \texttt{IAFormer} has two new mechanisms within the model. First, the attention matrix depends on predefined boost invariant pairwise quantities, reducing the network parameters significantly from the original particle transformer models. Second, \texttt{IAFormer} incorporates the sparse attention mechanism by utilizing the "differential attention", so that it can dynamically prioritize relevant particle tokens while reducing computational overhead associated with less informative ones. This approach significantly lowers the model complexity without compromising performance. Despite being computationally efficient by more than an order of magnitude than the Particle Transformer network, \texttt{IAFormer} achieves state-of-the-art performance in classification tasks on the top and quark-gluon datasets. Furthermore, we employ AI interpretability techniques, verifying that the model effectively captures physically meaningful information layer by layer through its sparse attention mechanism, building an efficient network output that is resistant to statistical fluctuations. \texttt{IAFormer} highlights the need for sparse attention in Transformer analysis to reduce the network size while improving its performance.

Figures (5)

• Figure 1: Schematic architecture ofIAFormer network.
• Figure 2: The left plot shows the $\beta$ distribution across the twelve IAFormer layers for different seed numbers on the test dataset. The right plot illustrates the uncertainty in the network output when trained with different seed numbers, 42, 0, 7, 123, and 12345. The results of the architectures of the Transformer $+\mathcal{I}_{ij}$(similar to ParT) and the plain Transformer are shown. The uncertainty is reported in terms of background rejection at 0.3 signal efficiency.
• Figure 3: $\beta$ distribution for quark-gluon test data set for three different seed numbers.
• Figure 4: Attention maps of the final self-attention layer of IAFormer (left) and a plain Transformer (right), generated using 10,000 test events from the top jet dataset. Each network comprises 16 attention heads, shown individually. The axes represent particle tokens, with each event containing 100 particles, while the color bar denotes the attention scores assigned to each particle pair, ranging from 0 to 1.
• Figure 5: Linear CKA similarity for IAFormer (left), Transformer $+\mathcal{I}_{i,j}$ (middle), and Plain Transformer (right) using 1000 test events from the top jet dataset. The axes represent the attention layers in each network, while the colour bar indicates the CKA values.