JetFormer: A Scalable and Efficient Transformer for Jet Tagging from Offline Analysis to FPGA Triggers

JetFormer is a scalable, encoder-only Transformer architecture designed for particle jet tagging that achieves high accuracy and computational efficiency across both high-precision offline analysis and ultra-low-latency FPGA-based online triggers.

The Gist

Imagine you are a security guard at a massive, high-speed stadium entrance. Every second, thousands of people rush through the gates. Most are just regular fans, but occasionally, a "troublemaker" (a specific type of particle) tries to sneak in.

To keep the stadium safe, you need to identify these troublemakers instantly. If you take too long to check IDs, the crowd piles up and the whole system crashes. If you are too careless, the troublemakers get in.

This paper introduces JetFormer, a new "smart scanner" designed for the world’s largest particle accelerator (the LHC). Here is the breakdown of how it works using everyday concepts.

1. The Problem: The "Needle in a Haystack" at Warp Speed

In particle physics, when protons collide, they explode into sprays of particles called "jets." Scientists want to know what caused these jets (was it a Higgs boson? A top quark?).

There are two ways to do this:

• The "Slow & Careful" Way (Offline Analysis): Like a detective examining high-resolution photos of a crime scene. It’s incredibly accurate but takes a long time.
• The "Fast & Furious" Way (Online Triggering): Like a security camera that has to decide in a fraction of a microsecond whether to sound an alarm. It has to be lightning-fast because the data is coming in too quickly to save everything.

Until now, it was hard to have a single "brain" that could do both. High-accuracy models were too "heavy" and slow for the fast gates, and fast models were too "simple" and inaccurate.

2. The Solution: JetFormer (The Versatile Brain)

The researchers created JetFormer, a specialized AI architecture based on a "Transformer" (the same technology behind ChatGPT, but redesigned for physics).

Think of JetFormer like a Swiss Army Knife.

• If you need a heavy-duty saw for a big project, you can use the "Large JetFormer" for deep, slow scientific research.
• If you just need a tiny pocket knife to cut a string quickly, you can use "JetFormer-tiny" for the high-speed security gates.

It’s "scalable," meaning the same basic design works whether you need massive intelligence or extreme speed.

3. The "Diet" Plan: Pruning and Quantization

To make the AI small enough to fit into the "security gate" hardware (called an FPGA), the researchers put the model on a strict diet using two methods:

• Pruning (The "Marie Kondo" Method): Imagine a massive textbook. Pruning is like going through and ripping out every page that doesn't provide essential information. You end up with a much thinner book that still tells the same story. The researchers cut the "computational weight" by 50% while barely losing any accuracy.
• Quantization (The "Sketch Artist" Method): Instead of a high-definition, 4K photograph that requires massive memory, quantization turns the image into a simplified sketch using only black and white. It’s much "lighter" to carry, and while you lose some fine detail, you can still clearly tell if the person in the photo is a troublemaker. They used "1-bit quantization," which is the ultimate extreme version of this.

4. The Result: Fast, Smart, and Efficient

The researchers proved that JetFormer is a champion in both worlds:

1. In the Lab: It is just as smart as the current "gold standard" models but uses much less "brain power" (FLOPs) to get the job done.
2. At the Gate: They successfully shrunk it down so it can live on specialized hardware, making it ready for the next generation of particle accelerators.

In short: JetFormer is a master of disguise—it can be a heavyweight scientist or a lightweight security guard, all while using the same brilliant architectural blueprint.

Technical Summary

Technical Summary: JetFormer

JetFormer: A Scalable and Efficient Transformer for Jet Tagging from Offline Analysis to FPGA Triggers

1. Problem Statement

In High-Energy Physics (HEP) experiments like those at the Large Hadron Collider (LHC), "jet tagging"—the classification of particle jets to identify heavy particles (e.g., Higgs bosons or top quarks)—is critical. While Transformer-based architectures (like ParT) have set new state-of-the-art (SOTA) benchmarks for accuracy in offline analysis, they are typically too computationally expensive and structurally complex for online triggering.

The Level-1 Trigger (L1T) system must make decisions in real-time using FPGAs under sub-microsecond latency constraints. Existing models for FPGAs (like MLPs or GNNs) often struggle to match the accuracy of Transformers, while existing Transformer implementations are often too large for resource-constrained hardware. There is a need for a unified architecture that scales from high-precision offline use to ultra-low-latency hardware deployment.

2. Methodology

The authors propose JetFormer, a versatile, encoder-only Transformer architecture. The methodology is divided into three main pillars:

• Architectural Design:
  • Encoder-only Structure: Inspired by BERT, it uses a learnable [CLS] token to aggregate global jet features.
  • Permutation Invariance: It processes unordered sets of particle features without positional encodings.
  • Hardware-Friendly Modifications: To facilitate FPGA synthesis, the authors replaced Layer Normalization with Batch Normalization (precomputable parameters) and swapped the SiLU activation function for ReLU (avoiding expensive exponential/division operations).
• Optimization & Compression Pipeline:
  • Multi-Objective Hyperparameter Optimization (HPO): Using the Optuna framework, the authors performed a search to maximize accuracy while minimizing FLOPs. They utilized the NSGAIISampler to identify a Pareto front of optimal models, leading to the creation of JetFormer-tiny.
  • Structured Pruning: A pipeline using the torch-pruning library and Taylor importance was used to remove entire channels/filters, reducing computational costs by ~50% with minimal accuracy loss.
  • 1-Bit Quantization: Using the BitNet framework, the authors implemented Quantization-Aware Training (QAT) to compress weights to 1-bit (+1 or -1) while keeping activations at 8-bit, significantly shrinking the model size.
• Hardware Implementation:
  • The authors extended the Allo framework (built on MLIR) to support Transformer-specific operations (e.g., class token concatenation, log-softmax, and 3D batch normalization) to enable end-to-end FPGA synthesis via High-Level Synthesis (HLS).

3. Key Contributions

• Unified Framework: A single architectural foundation that scales from tiny FPGA-ready models to large-scale offline models.
• Hardware-Aware Pipeline: An automated end-to-end workflow encompassing HPO, structured pruning, and 1-bit quantization specifically for HEP applications.
• Allo Extension: Contributions to the Allo compiler to bridge the gap between complex Transformer layers and FPGA-compatible HLS code.

4. Results

• Accuracy & Efficiency (Offline): On the large-scale JETCLASS dataset, JetFormer achieves an accuracy of 82.9% (comparable to ParT's 83.6%) but uses 37.4% fewer FLOPs. On the HLS4ML 150P dataset, it outperforms MLPs, Deep Sets, and Interaction Networks by 3–4% in accuracy.
• Compression Effectiveness:
  • Pruning: Reduced FLOPs by ~50% with an accuracy loss of $<0.5\%$.
  • Quantization: Achieved a model size reduction of 82–92% with a manageable accuracy drop of only 1.5–3.5%.
• Hardware Feasibility: The JetFormer-tiny model was successfully synthesized onto FPGAs. While the initial implementation (without manual pipelining) showed higher latency, the low resource utilization (e.g., $<10\%$ LUT/DSP usage) demonstrates significant headroom for further optimization through parallelism and pipelining.

5. Significance

JetFormer provides a practical pathway for bringing the power of Transformers into the real-time decision-making processes of particle physics experiments. By proving that Transformers can be aggressively compressed without losing their competitive edge, the work bridges the gap between cutting-edge deep learning research and the stringent, real-world constraints of high-speed hardware triggers at the LHC.