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

AI summary

Overview

Research area: Machine learning for high-energy physics — specifically Transformer-based jet tagging at the CERN Large Hadron Collider, combined with hardware-aware optimization and FPGA deployment for real-time triggering.

Technical level: Advanced. The paper assumes familiarity with Transformer architectures, quantization/pruning techniques, high-level synthesis (HLS) for FPGAs, and LHC trigger terminology.

Scope: The paper proposes a single encoder-only Transformer family (JetFormer) that can be scaled from large offline analysis models down to tiny pruned variants synthesized onto FPGAs, and evaluates it on the HLS4ML 150-particle and JetClass jet datasets.

What This Paper Is About

Jet tagging — classifying a collimated spray of particles by which particle produced it — is essential for isolating Higgs boson and top quark signatures from huge light-quark and gluon backgrounds at the LHC. Existing deep learning models tend to be specialized either for high-accuracy offline analysis or for ultra-low-latency FPGA trigger deployment, not both. The authors build JetFormer, an encoder-only Transformer that uses a learnable class token (as in BERT) and deliberately omits explicit pairwise particle interactions and positional encoding, then push it through a hardware-aware pipeline of hyperparameter search, structured pruning and 1-bit quantization so the same architectural family covers both regimes.

Key Contributions

1. A scalable encoder-only Transformer for jet tagging. JetFormer processes variable-length, unordered sets of particle features with no explicit pairwise interaction inputs and no positional encoding, aiming to serve compact trigger-scale models and large offline models from one design.
2. An automated hardware-aware optimization and compression pipeline. Multi-objective hyperparameter search with Optuna (maximizing accuracy, minimizing FLOPs), followed by structured pruning via the torch-pruning library and 1-bit quantization in a BitNet-style PyTorch framework.
3. Extensions to the Allo high-level synthesis framework. Support was added for Transformer operations missing from the Allo PyTorch frontend — log_softmax, relu3d, batchnorm1d for 2D and 3D inputs, and class-token operations (expansion, concatenation, slicing) — enabling end-to-end compilation of JetFormer onto FPGA targets.
4. Comprehensive evaluation across datasets and hardware stages. Accuracy, efficiency and compressibility are measured on HLS4ML 150-particle benchmarks and the large-scale JetClass dataset, with deployment validated through LLVM CPU, sw_emu and hw_emu stages.

Main Findings

• Competitive with interaction-rich ParT on JetClass. On a 2M-sample subset of JetClass, a large JetFormer (10 transformer encoder layers, embedding dimension 128, 8 attention heads, dropout 0.0) reaches 0.829 accuracy and 0.9827 AUC with 1.66M parameters and 213M FLOPs. ParT reports 0.836 accuracy, 0.9834 AUC, 2.14M parameters and 340M FLOPs — a 0.7% accuracy and 0.07% AUC decrease for JetFormer, at 37.4% fewer FLOPs. ParticleNet reports 0.828 accuracy, 0.9820 AUC, 370k parameters and 540M FLOPs.
• Outperforms MLP, Deep Sets and Interaction Networks on small benchmarks. On the HLS4ML 150P dataset truncated to 3 features, JetFormer reaches 67.1% (8 particles), 74.4% (16 particles) and 79.9% (32 particles) accuracy, versus 64.9%, 70.8% and 75.8% for the Interaction Network and lower values for MLP and Deep Sets — an improvement of roughly 3–4 percentage points, with higher per-class AUC values. JetFormer uses substantially more parameters (101,291 to 107,339) and FLOPs (933k to 4M) than the baselines.
• On the full 16-feature dataset with varying constituent counts, JetFormer reports accuracies of 81.77 (30 constituents), 82.88 (50), 83.05 (100) and 82.97 (150), with 4M–24M FLOPs and 107,403–238,443 parameters. JEDI-net comparison values are quoted from prior work.
• Optuna sampler comparison. Over 80 valid trials, TPESampler, NSGAIISampler and BoTorchSampler all converge to a comparable hypervolume of around 0.83; NSGAIISampler converges fastest and reaches the highest hypervolume, so it is selected. TPESampler and BoTorchSampler each used 20 random startup trials plus 60 sampler trials.
• JetFormer-tiny selection. Among 80 runs, the lowest-FLOP model with validation accuracy above the 0.65 feasibility threshold is Model 0: 4 transformer blocks, embedding dimension 8, 2 attention heads, dropout 0.00, 26,168 FLOPs and validation accuracy 0.6525; its final test accuracy is 0.656. The best model found (4 transformer blocks, embedding dimension 128, 8 heads, dropout 0.05) reaches 0.673 test accuracy at 4,842,384 FLOPs.
• Structured pruning is cheap in accuracy. At a global pruning ratio of 0.5, JetFormer-tiny drops from 26,168 to 13,784 FLOPs and from 3,080 to 1,997 parameters, with roughly 0.5% accuracy loss; GPU inference time at batch size 10,240 falls by 17.46%. Larger models lose less than 0.3% accuracy.
• 1-bit quantization compresses heavily at modest cost. Quantization shrinks model size by 82–92% at the cost of 1.5–3.5% absolute accuracy drop. Only the attention projections (query, key, value, output) and feedforward layers become BitLinear layers; the input embedding and the classification head dense output layer stay in full precision, and activations remain 8-bit.
• Hardware flow validated up to emulation. Both the original and 50%-pruned JetFormer-tiny are converted through Allo and validated on the LLVM CPU backend, sw_emu and hw_emu. Quantized models were not deployed because Allo lacks Vitis HLS support for quantized deployment. Backend hardware resource and latency results are not included in the provided text.

Methodology in Plain English

Architecture. Each particle in a jet is treated as one token, and each jet is a sequence. Particles are projected into an embedding vector; no positional encoding is added, so the model stays invariant to particle ordering. A randomly initialized, learnable class token is prepended to the sequence and, after the final Transformer block, its hidden state summarizes the whole jet and feeds a classification head with log-softmax. The blocks use multi-head self-attention plus a two-layer feedforward network, each preceded by normalization, with residual connections. For hardware friendliness, every layer normalization is replaced by batch normalization (whose parameters are fixed at inference), SiLU is replaced by ReLU, and no explicit pairwise interaction features are used.

Training. AdamW with weight decay 0.01 and initial learning rate 0.001; three schedulers are compared and OneCycleLR gives the best validation performance and fastest convergence. Batch size is 256 for the 150-particle dataset and 128 for JetClass, with early stopping. Training the JetClass 2M subset takes about three days on a single NVIDIA GeForce RTX 3090.

Hyperparameter search. Optuna explores 1–6 transformer blocks, eight predefined embedding-dimension/head combinations from (8,2) to (128,8), and dropout of 0.0 or 0.05, on the 8-particle, 3-feature dataset with a 9:1 train/validation split. Each trial runs up to 25 epochs with early stopping at patience 4. Solutions with validation accuracy below 0.65 are treated as infeasible and excluded from the Pareto front. Multi-objective search maximizes accuracy while minimizing FLOPs, and hypervolume is evaluated at 40, 60, 80 and 100 trials.

Compression. Structured pruning uses the torch-pruning library with Taylor importance scores, removing entire channels/filters/groups across input embedding, attention key/query/value projections and feedforward linear layers in five steps of 5 fine-tuning epochs each. Quantization uses a BitNet PyTorch framework: weights are centralized to zero mean, binarized with the sign function and rescaled, while activations use absmax quantization into the range −127 to 127, with a straight-through estimator for gradients. The quantized model is trained from scratch (quantization-aware training) for 80 epochs with an initial learning rate of 8e-4 and ReduceLROnPlateau (factor 0.8, patience 5, min_lr 1e-4).

Deployment. A small particle MLP (three hidden layers of 64, 32, 32 neurons and a 5-neuron output, all ReLU) is converted first as a feasibility check. Then original and pruned JetFormer-tiny are converted with batch size 16, after adding the missing Allo operators and frontend parsing logic.

Datasets. The HLS4ML LHC Jet dataset (150 particles) has 620,000 training and 260,000 test samples with balanced classes (q, g, W, Z, t) and up to 150 particles with 16 kinematic features; normalization uses Welford's online algorithm computed on the training set only. JetClass has 100M jets (100M/5M/20M for train/validation/test), 10 jet categories and 17 per-particle features plus interaction features; JetFormer uses the 17 standard features but excludes the pairwise interaction features, and only a subset is used for training and evaluation.

Why This Matters

Impact on research. The work argues that a single Transformer family can serve offline analysis and online triggering, rather than requiring separate bespoke architectures per deployment regime. It also tests whether explicit pairwise interaction features — central to ParT and JEDI-net-style models — are necessary for state-of-the-art jet tagging accuracy, reporting near-parity with ParT at 37.4% fewer FLOPs without them. On the tooling side, the Allo extensions push open-source HLS flows one step closer to supporting Transformer inference on FPGAs.

Real-world applications:

• L1 trigger event selection at the CMS Phase-2 upgrade, where track information enables jet reconstruction and jet tagging could run directly in hardware under extreme pileup.
• Offline and high-level-trigger physics analyses that need high jet-tagging accuracy with lower compute budgets.
• General deployment of compact Transformer inference on resource-constrained FPGA accelerators in latency-critical systems.
• Methods transfer: the pruning, 1-bit quantization and multi-objective hyperparameter search pipeline is applicable to other scientific and industrial streaming-classification tasks.

Industry relevance. The paper targets the general problem of fitting Transformer models into sub-microsecond, resource-limited hardware — relevant to edge inference, real-time sensor and detector pipelines, and any setting where accuracy and latency must be co-optimized. The use of Optuna, torch-pruning, BitNet-style quantization and the MLIR-based Allo compiler places the work within mainstream open-source ML systems tooling. Affiliations include Imperial College London, the California Institute of Technology and SLAC National Accelerator Laboratory, and code is released at https://github.com/walkieq/JetFormer.

Future Directions

1. Deploying quantized models on hardware. Allo currently lacks Vitis HLS support for quantized deployment, so the 1-bit models were not synthesized; adding this support is a natural next step.
2. Reporting full hardware metrics. The provided text covers the deployment workflow and validation stages but not the FPGA resource-utilization and latency tables, which the authors state are reported in Section 5.3.2 — obtaining these would confirm whether the small JetFormer-tiny variants meet L1T-style sub-microsecond timing budgets.
3. Scaling the JetClass study. Only a 2M-sample subset was used, so how JetFormer behaves at the full 100M-jet training scale — and whether the 0.7% accuracy gap to ParT closes or widens — remains open.
4. Broadening architectural exploration. The search space was limited to 1–6 blocks, eight predefined embedding/head pairings and two dropout values; larger design spaces, alternative quantization schemes and harder (post-training) quantization could be explored.

Target Audience

This paper suits researchers and engineers at the intersection of machine learning and high-energy physics: physicists developing jet taggers for ATLAS or CMS, ML practitioners working on FPGA/edge deployment of Transformers, and hardware-software co-design researchers interested in high-level synthesis flows. Readers with a background in deep learning but not in HEP will follow the architecture and compression sections easily, while the trigger and hardware sections require some familiarity with LHC data rates, level-1 triggers and FPGA design.

Authors’ abstract

We present JetFormer, a versatile and scalable encoder-only Transformer architecture for particle jet tagging at the Large Hadron Collider (LHC). Unlike prior approaches that are often tailored to specific deployment regimes, JetFormer is designed to operate effectively across the full spectrum of jet tagging scenarios, from high-accuracy offline analysis to ultra-low-latency online triggering. The model processes variable-length sets of particle features without relying on input of explicit pairwise interactions, yet achieves competitive or superior performance compared to state-of-the-art methods. On the large-scale JetClass dataset, a large-scale JetFormer matches the accuracy of the interaction-rich ParT model (within 0.7%) while using 37.4% fewer FLOPs, demonstrating its computational efficiency and strong generalization. On benchmark HLS4ML 150P datasets, JetFormer consistently outperforms existing models such as MLPs, Deep Sets, and Interaction Networks by 3-4% in accuracy. To bridge the gap to hardware deployment, we further introduce a hardware-aware optimization pipeline based on multi-objective hyperparameter search, yielding compact variants like JetFormer-tiny suitable for FPGA-based trigger systems with sub-microsecond latency requirements. Through structured pruning and quantization, we show that JetFormer can be aggressively compressed with minimal accuracy loss. By unifying high-performance modeling and deployability within a single architectural framework, JetFormer provides a practical pathway for deploying Transformer-based jet taggers in both offline and online environments at the LHC. Code is available at https://github.com/walkieq/JetFormer.