Spatially Aware Linear Transformer (SAL-T) for Particle Jet Tagging

TL;DR

The paper introduces SAL-T, a linear transformer variant tailored for particle jet tagging that reduces computational complexity while maintaining high accuracy by incorporating spatial awareness and local correlations.

Contribution

SAL-T is a physics-inspired linear transformer that efficiently captures spatial and local features in particle physics data, outperforming standard models in jet classification tasks.

Findings

SAL-T outperforms standard linformer in jet classification.

SAL-T achieves accuracy comparable to full-attention transformers.

SAL-T requires fewer resources and has lower inference latency.

Abstract

Transformers are very effective in capturing both global and local correlations within high-energy particle collisions, but they present deployment challenges in high-data-throughput environments, such as the CERN LHC. The quadratic complexity of transformer models demands substantial resources and increases latency during inference. In order to address these issues, we introduce the Spatially Aware Linear Transformer (SAL-T), a physics-inspired enhancement of the linformer architecture that maintains linear attention. Our method incorporates spatially aware partitioning of particles based on kinematic features, thereby computing attention between regions of physical significance. Additionally, we employ convolutional layers to capture local correlations, informed by insights from jet physics. In addition to outperforming the standard linformer in jet classification tasks, SAL-T also achieves classification results comparable to full-attention transformers, while using considerably fewer resources with lower latency during inference. Experiments on a generic point cloud classification dataset (ModelNet10) further confirm this trend. Our code is available at https://github.com/aaronw5/SAL-T4HEP.