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

AI summary

Overview

Research area: Machine learning for high-energy physics — specifically efficient transformer architectures for real-time particle jet tagging at the CERN Large Hadron Collider.

Technical level: Advanced. The paper assumes familiarity with transformer attention mechanisms, linear/low-rank attention approximations, and particle physics concepts such as jets, pseudorapidity, azimuthal angle, and transverse momentum.

Scope: The paper introduces and benchmarks SAL-T, a physics-informed linear-attention transformer that replaces the quadratic attention of standard transformers with spatially partitioned linear attention plus a light convolutional step, evaluated on jet-tagging datasets and the ModelNet10 point cloud benchmark.

What This Paper Is About

Transformers capture global and local correlations in particle collisions well, but their quadratic attention cost makes them too slow and resource-hungry for the LHC's real-time trigger, where the LHC collides protons 40 million times per second and only about 1 in 40,000 events can be stored. The authors build SAL-T, a modification of the linear-complexity linformer that injects spatial awareness — sorting particles by a physics-motivated quantity, partitioning keys and values into localized groups, and convolving over attention scores — so that a linear-attention model can approach full-transformer accuracy at much lower cost.

Key Contributions

1. SAL-T architecture: A physics-inspired enhancement of linformer that keeps linear attention while adding spatial structure, built around a new module called linear partitioned particle multi-head attention (LPP-MHA).

2. Three concrete modifications to linformer: (1) sorting input particles in descending order by the metric k_T = p_T ΔR (where ΔR is the pseudoangular distance to the jet axis), (2) partitioning key and value projections into p groups so each projection head attends only to its own subset of particles, and (3) applying a small depthwise 2D convolution over each head's raw attention scores to mix information between neighboring particles without reintroducing quadratic complexity.

3. Complexity reduction: The paper states that partitioning reduces attention complexity from O(n²) to O(n·p), which is near linear when p << n, and that the full LPP-MHA has complexity O(n·p·f·c) for sequence length n, projection dimension p, f filters, and kernel width c.

4. Cross-domain validation: Beyond jet tagging on the hls4ml, Top Tagging, and Quark Gluon datasets, the authors benchmark on the generic point cloud dataset ModelNet10, using Morton-code sorting in place of k_T sorting.

5. Compute-matched and SOTA comparisons: The authors compare against scaled-down Particle Transformer (ParT), Point Transformer v3 (PTv3), Transolver, and PointNet, and run a controlled experiment reducing a standard Transformer to match SAL-T's FLOPs.

Main Findings

• On the hls4ml dataset (single layer, ≤150 constituent particles): SAL-T with k_T sorting reaches 81.18 ± 0.03% accuracy, 0.9593 ± 0.0002 ROC AUC, and 40.78 ± 0.57 average rejection (1/FPR at 0.8 TPR). The full-attention Transformer with p_T sorting reaches 81.27 ± 0.08% accuracy, 0.9589 ± 0.0004 AUC, and 42.02 ± 0.71 rejection, essentially parity within uncertainty.

• SAL-T beats linformer: Linformer with k_T sorting scores 81.00 ± 0.08% accuracy, 0.9585 ± 0.0003 AUC, and 38.41 ± 0.54 rejection; linformer with p_T sorting scores 79.90 ± 0.00%, 0.9545 ± 0.0004, and 28.06 ± 0.58. SAL-T with p_T sorting scores 78.82 ± 0.01%, 0.950 ± 0.0028, and 23.45 ± 2.12.

• k_T sorting helps beyond SAL-T: k_T sorting also improved PointTransformer v3 (81.37 ± 0.06% accuracy, 0.9599 ± 0.0002 AUC, 66.75 ± 0.18 rejection) over p_T sorting (81.32 ± 0.00%, 0.9597 ± 0.0001, 65.82 ± 1.42) and Morton code sorting (80.99 ± 0.15%, 0.9584 ± 0.0004, 60.19 ± 1.41).

• SAL-T beats earlier trigger-oriented models: On hls4ml, SAL-T outperformed Deep set (79.65 ± 0.23% accuracy, 0.8336 ± 0.0376 AUC, 7.41 ± 0.99 rejection), Interaction network (80.05 ± 0.46%, 0.9128 ± 0.0139, 22.01 ± 3.69), and MLP (53.09 ± 1.32%, 0.8054 ± 0.0120, 3.66 ± 0.12).

• Hardware efficiency at 150 particles: SAL-T uses 3,264 parameters, 739,918 FLOPs, 303.4 MB peak GPU memory, and 27.69 ± 0.32 µs inference time. The Transformer uses 2,009 parameters, 2,479,918 FLOPs, 4,357.1 MB, and 30.86 ± 0.14 µs. Linformer uses 6,809 parameters, 552,718 FLOPs, 245.8 MB, and 22.38 ± 0.33 µs. The Interaction network uses 10,064,702 parameters, 61,529,931 FLOPs, 4,283.9 MB, and 307.38 ± 0.44 µs.

• SAL-T outperforms scaled-down SOTA point cloud models: 1-layer ParT has 8,217,868 FLOPs and 211.20 ± 0.84 µs latency (an order of magnitude slower than SAL-T at 27 µs, with out-of-memory errors noted); PTv3 with k_T sorting uses 767,573 FLOPs and 94.10 ± 0.90 µs; Transolver uses 6,430,935 FLOPs and 88.44 ± 0.33 µs with 80.13 ± 0.09% accuracy; PointNet uses 78,809 FLOPs at 21.03 ± 0.10 µs but only 74.22 ± 0.07% accuracy. SAL-T is described as the most performant option under the inference latency budget.

• Compute-matched comparison: When a standard Transformer is shrunk to match SAL-T's FLOP count, SAL-T reaches 81.18 ± 0.03% accuracy and 40.78 ± 0.57 rejection while SAL-T (No Conv) reaches 81.09 ± 0.10% and 39.32 ± 1.06, versus the reduced Transformer at 78.93 ± 0.51% and 26.19 ± 1.98.

• FlashAttention comparison: On an NVIDIA A10 GPU, SAL-T runs in 3.99 µs versus 4.42 µs for a FlashAttention-enabled Transformer at these sequence lengths.

• Top Tagging dataset: Two-layer SAL-T (6,939 parameters, 2,794,161 FLOPs) reaches 92.79 ± 0.03% accuracy, 0.9796 ± 0.0001 AUC, and 35.15 ± 0.17 rejection, beating two-layer linformer (16,329 parameters, 1,450,161 FLOPs, 92.61 ± 0.02%, 0.9786 ± 0.0001, 33.11 ± 0.22). The Transformer reaches 93.11 ± 0.03%, 0.9813 ± 0.0001, and 40.11 ± 0.23.

• Quark Gluon dataset: Two-layer SAL-T (6,171 parameters, 2,095,761 FLOPs) reaches 81.77 ± 0.02% accuracy, 0.8924 ± 0.0004 AUC, and 6.07 ± 0.01 rejection — surpassing the single-layer Transformer (1,941 parameters, 2,479,761 FLOPs, 81.64 ± 0.03%, 0.8913 ± 0.0001, 5.98 ± 0.03). Two-layer linformer reaches 81.60 ± 0.08%, 0.8906 ± 0.0006, 5.94 ± 0.05.

• Ablation results: SAL-T without partition reaches 81.02 ± 0.05% accuracy (with 6,848 parameters), SAL-T without convolution reaches 81.09 ± 0.10% (3,225 parameters), and full SAL-T reaches 81.18 ± 0.03% (3,264 parameters). Both components individually improve over linformer, and combining them improves further.

• ModelNet10 results: SAL-T scores 80.10 ± 1.11% accuracy at 5,047,180 FLOPs and 454.6 ± 3.1 MB peak GPU memory; linformer scores 77.86 ± 1.16% at 3,769,228 FLOPs and 397.8 ± 27.1 MB; the Transformer scores 82.56 ± 3.36% at 95,683,468 FLOPs and 6,221.1 ± 8.3 MB. Without sorting, SAL-T drops to 68.79 ± 1.16% and linformer to 60.72 ± 1.97%.

• Data context: The hls4ml dataset contains 504,000 training, 126,000 validation, and 240,000 test jets across five classes (light quark, gluon, W boson, Z boson, top quark). Top Tagging contains 1.2 million training, 0.4 million validation, and 0.4 million test jets. Quark Gluon contains 1.8 million training jets and 0.2 million test jets, with 20% of training randomly sampled for validation. ModelNet10 contains 4,899 training and 908 testing 3D objects across 10 categories, with 1,024 points sampled per object.

• Statistical caveat: The highest particle-count bin (115–150 particles) contains only 41 jets, less than 0.016% of the dataset, which the authors cite as the cause of performance variance there.

Methodology in Plain English

The authors start from linformer, an existing linear-attention transformer that compresses keys and values with learned projections, but which encodes no spatial information and is therefore blind to jet substructure. They add three changes.

First, instead of ordering particles by transverse momentum (the LHC default), they sort by k_T = p_T ΔR, a metric that is larger for particles with greater transverse momentum or closer alignment with the jet axis, and which underpins infrared- and collinear-safe jet clustering. The authors note explicitly that k_T is used only as a scalar sorting key, not as a full sequential recombination clustering algorithm, which would be infeasible inside trigger latency constraints.

Second, they split the sorted sequence into p partitions and give each partition its own learnable projection matrix for keys and values, so each projection row aggregates only over a localized group of particles. Complexity drops from O(n²) to O(n·p).

Third, after splitting into H heads, they apply a depthwise 2D convolution over the attention logits within each head. Because the sequence is k_T-sorted, a convolutional filter sees particles that are near each other in physical space. The filter spans the projection dimension and uses "same" padding, and the responses are averaged and passed through softmax. This adds only linear cost in sequence length.

All models embed constituents into a 16-dimensional space with H = 4 heads, use dynamic tanh instead of layer norm for faster inference, aggregate with max pooling, and end in a dense classification head. SAL-T uses rank-4 partitioned projections and three 2D convolutional filters with heights 1, 3, and 5. Training ran for 1400 epochs with a phased batch size schedule (200 epochs each at 128, 256, 512, 1024, and 2048, then 400 epochs at 4096), Adam optimizer at learning rate 0.001, and early stopping after 40 epochs without validation-loss improvement, on the National Research Platform Nautilus cluster using NVIDIA GTX 1080-Ti or GTX 3090 GPUs.

Why This Matters

Research impact: The work demonstrates that injecting known physics structure — sorting by a theoretically motivated metric and partitioning spatially — into a low-rank attention mechanism can recover accuracy lost by linear approximation. It also shows the sorting insight transfers beyond physics, improving PointTransformer v3 on jet tagging and being necessary for SAL-T and linformer to perform well on ModelNet10 (where

Authors’ 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.