Economical Jet Taggers -- Equivariant, Slim, and Quantized

This paper presents a slim, quantized, and parameter-reduced version of the L-GATr jet tagger that achieves an order-of-magnitude reduction in energy cost with only a moderate performance decrease, paving the way for efficient trigger-level jet tagging at the LHC.

The Gist

Imagine the Large Hadron Collider (LHC) as a massive, high-speed particle factory. Every second, it smashes protons together, creating a chaotic spray of debris. Physicists need to sort through this debris to find specific, rare particles (like the "top quark") hidden among billions of ordinary ones. This sorting process is called jet tagging.

For years, scientists have used complex computer programs (Machine Learning) to do this sorting. The current champions are "Transformers"—powerful AI models that are incredibly accurate but also huge, slow, and hungry for energy. They are like a fleet of massive, fuel-guzzling trucks trying to deliver a single letter; they get the job done, but they are too big and expensive to use at the very moment the data is being collected (the "trigger" level).

This paper asks a simple question: Can we shrink these giant trucks into tiny, fuel-efficient scooters without losing the ability to deliver the letter?

Here is how the authors did it, using three main strategies:

1. The "Slim" Version (L-GATr-slim)

The original "L-GATr" model is like a Swiss Army knife that carries every possible tool: scalars, vectors, tensors, and more. However, the authors realized that for most particle physics jobs, you only really need two tools: scalars (numbers) and vectors (arrows with direction).

• The Analogy: Imagine a chef who insists on using a full industrial kitchen with ovens, blenders, and mixers just to make a simple sandwich. The authors said, "Let's just use a knife and a cutting board."
• The Result: They built a "Slim" version of the AI that strips away the unnecessary tools. It performs just as well as the giant version but is much faster to train and uses less memory. It's like switching from a heavy-duty truck to a nimble sports car that gets the same job done.

2. The "Tiny" Version (Ultra-mini Taggers)

The authors then asked, "How small can we go?" They tried to shrink these AI models down to the size of a tiny toy car (around 1,000 parameters, compared to the millions in the original).

• The Analogy: Think of trying to fit a whole library's worth of knowledge into a single postcard. Usually, you lose the story. But the authors found that if you organize the information correctly (using specific "Lorentz-equivariant" rules that respect the laws of physics), you can fit the essential knowledge into a tiny space.
• The Result: They found that for very small models, the "LLoCa" architecture works best if you shrink the number of layers, while the "L-GATr-slim" works best if you shrink the width of the layers. Even at this microscopic size, they still outperformed older, non-physics-aware AI models.

3. The "Quantized" Version (Low-Precision Math)

This is the most dramatic energy saver. Standard AI uses very precise math (like measuring a distance to the billionth of a millimeter). The authors realized that for jet tagging, you don't need that much precision. You can get away with rounding numbers off significantly.

• The Analogy: Imagine you are counting apples in a warehouse.
  • Standard AI: You weigh every single apple to the microgram. (Accurate, but takes forever and uses a lot of scale energy).
  • Quantized AI: You just count them in whole numbers. (Fast, uses almost no energy, and for the purpose of knowing "how many apples," it's perfectly fine).
• The Method: They used a technique called PARQ (Piecewise-Affine Regularized Quantization). Think of this as a smart rounding rule that gently nudges the numbers to be simple (like 0, 1, or -1) during the training process, rather than forcing them abruptly.
• The Result: By switching to these "rougher" numbers, they reduced the energy cost of running the AI by 10 times (an order of magnitude). The AI became incredibly fast and energy-efficient, with only a tiny drop in accuracy.

The Big Picture

The authors combined these three strategies—Slimming the architecture, Miniaturizing the size, and Quantizing the math—to create "Economical Jet Taggers."

• Why does this matter? Currently, these powerful AI models are too big to run on the hardware that decides in real-time which collisions to keep and which to discard (the "trigger").
• The Goal: By making these models small, fast, and energy-efficient, the authors hope to eventually run them directly on the trigger hardware. This would allow the LHC to use AI to make split-second decisions about which particle collisions to save, potentially discovering new physics that was previously missed because the data was discarded too quickly.

In short: They took a giant, energy-hungry AI, gave it a diet, shrank it down, and taught it to do math with fewer decimals, resulting in a tiny, super-efficient engine that can still recognize the most important particles in the universe.

Technical Summary

Technical Summary: Economical Jet Taggers – Equivariant, Slim, and Quantized

Problem Statement
Modern machine learning (ML) has transformed jet tagging at the Large Hadron Collider (LHC), with Lorentz-equivariant transformers emerging as state-of-the-art architectures. However, leading models like L-GATr are computationally expensive, requiring significant memory and training time. While industry trends favor upscaling networks and datasets, LHC physics faces specific constraints, particularly regarding the memory and latency requirements of event triggering hardware. Current jet classification does not yet play a role in triggering, but the authors argue it should. The central challenge addressed is how to reduce the size and computational cost of modern equivariant jet taggers while minimizing performance degradation, potentially enabling their deployment at the trigger level.

Methodology
The paper proposes a two-pronged strategy to optimize resource efficiency: architectural slimming and numerical quantization.

1. L-GATr-slim Architecture:
   The authors introduce a streamlined version of the Lorentz-equivariant transformer (L-GATr). Standard L-GATr utilizes a geometric algebra representation involving scalars, pseudo-scalars, vectors, axial-vectors, and antisymmetric rank-two tensors. The authors observe that pseudo-scalars, axial-vectors, and tensors are unnecessary for most LHC applications. Consequently, L-GATr-slim restricts the latent representation to only scalars and vectors.

   • Linear Layers: Extended to operate on coupled scalar and vector representations, ensuring vector components share a single learnable scalar coefficient to maintain Lorentz equivariance.
   • Nonlinearity: Adapts the Gated Linear Unit (GLU) by applying nonlinearity to the inner product of two vectors, multiplied by the vector output.
   • Normalization: Modifies RMSNorm to use the absolute value of the Minkowski inner product for vector channels.
   • Attention: Constructs scalar attention matrices using a specific formulation that avoids the computationally expensive outer product used in the full L-GATr.
   • Implementation: The architecture is designed to be compiled with torch.compile for efficiency.

2. Quantization Strategies:
   The authors apply low-precision data types and weight quantization to further reduce costs.

   • Data Type Quantization: Inputs to linear layers are quantized to int8 (using zero-point quantization) while maintaining bfloat16 for precision-sensitive operations and the backward pass. This is applied to the hidden layers of Transformer, ParT, L-GATr-slim, and LLoCa-Transformer.
   • Weight Quantization: Linear weights are quantized to binary or ternary values using Proximal Gradient Quantization (PARQ). This method treats quantization as a regularization constraint, using a proximal operator to update weights. The authors compare PARQ against Straight-Through Estimation (STE), finding PARQ offers better stability and performance.
   • Equivariance Preservation: Special care is taken to ensure quantization does not violate Lorentz equivariance. For LLoCa, orthonormalization and frame projections remain in full precision (float32), limiting low-precision operations to Lorentz invariants. For L-GATr-slim, full vectors are multiplied by quantized weights, which does not introduce additional symmetry violations.

3. Ultra-Mini Scaling:
   The authors investigate the performance of these architectures down to 1,000 parameters by reducing the number of blocks or the width (channels) of the network.

Key Results
The study benchmarks the proposed methods on three tasks: top tagging, amplitude regression, and event generation.

• Performance vs. Efficiency (L-GATr-slim):

  • On the JetClass dataset (multi-class jet tagging), L-GATr-slim matches the performance of the full L-GATr and LLoCa-Transformer (AUC ~0.9885) but reduces training time by a factor of six (from 166h to 27h on an H100 GPU) and memory consumption by a factor of two.
  • In amplitude regression ($Z + 4g$), L-GATr-slim achieves the same Mean Squared Error (MSE) as full L-GATr but requires 20 times fewer training operations and half the training time.
  • In event generation ($t\bar{t} + nj$), the slim architecture matches the negative log-likelihood performance of the full models.

• Ultra-Mini Taggers:

  • When reducing the number of blocks (depth), the LLoCa-Transformer outperforms L-GATr-slim at very small sizes (e.g., 1,000 parameters).
  • When keeping the number of blocks fixed (10) and reducing channels (width), L-GATr-slim maintains a background rejection rate above 1,000 with only 2 vector and 4 scalar channels, outperforming other 1,000-parameter architectures.

• Quantization Gains:

  • Quantizing inputs to int8 and weights to ternary values reduces energy consumption by approximately an order of magnitude (factor of 10) with only marginal performance loss.
  • The LLoCa-Transformer and L-GATr-slim are robust to quantization, maintaining high performance where standard transformers might degrade more significantly.
  • For the most resource-constrained scenario (1 block, 16-dimensional latent space, int8), the quantized LLoCa-Transformer (global canonicalization) retains performance superior to pre-graph taggers, despite a factor-of-two reduction in background rejection compared to its full-size counterpart.

Significance and Claims
The paper claims that these "economical" versions of equivariant transformers represent a viable path toward trigger-level jet tagging at the High-Luminosity LHC (HL-LHC). By combining architectural slimming (removing unnecessary geometric algebra components) and aggressive quantization (PARQ and int8), the authors demonstrate that it is possible to create taggers with ~1,000 parameters that retain the physics-motivated benefits of Lorentz equivariance.

The authors emphasize that while upscaling is the industry standard, LHC physics requires a "physics-aware downscaling" approach. The results suggest that small, quantized, and equivariant networks can be deployed on resource-constrained hardware (such as FPGAs) without sacrificing the fundamental symmetries that make these models effective, potentially opening new avenues for real-time analysis of jet substructure.