Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

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

Ruoqing Zheng, Chang Sun, Qibin Liu, Lauri Laatu, Arianna Cox, Benedikt Maier, Alexander Tapper, Jose G. F. Coutinho, Wayne Luk, Zhiqiang Que

TL;DR

JetFormer tackles the need for scalable, hardware-friendly jet tagging by presenting an encoder-only Transformer capable of operating across offline-analytic accuracy and online FPGA-trigger latency regimes. It emphasizes a permutation-invariant input design, class-token aggregation, and hardware-aware modifications (batch normalization, ReLU) to ease FPGA deployment. A multi-objective hyperparameter optimization, structured pruning, and 1-bit quantization enable compact JetFormer variants (e.g., JetFormer-tiny) suitable for sub-microsecond latency, while maintaining competitive accuracy (within ~0.7% of ParT with 37.4% fewer FLOPs on JetClass). The work demonstrates practical deployment viability through Allo-based FPGA synthesis, validated across datasets (HLS4ML 150P and JetClass) and hardware measurements, highlighting a concrete path from high-performance modeling to realistic trigger-system integration at the LHC.

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.

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

TL;DR

JetFormer tackles the need for scalable, hardware-friendly jet tagging by presenting an encoder-only Transformer capable of operating across offline-analytic accuracy and online FPGA-trigger latency regimes. It emphasizes a permutation-invariant input design, class-token aggregation, and hardware-aware modifications (batch normalization, ReLU) to ease FPGA deployment. A multi-objective hyperparameter optimization, structured pruning, and 1-bit quantization enable compact JetFormer variants (e.g., JetFormer-tiny) suitable for sub-microsecond latency, while maintaining competitive accuracy (within ~0.7% of ParT with 37.4% fewer FLOPs on JetClass). The work demonstrates practical deployment viability through Allo-based FPGA synthesis, validated across datasets (HLS4ML 150P and JetClass) and hardware measurements, highlighting a concrete path from high-performance modeling to realistic trigger-system integration at the LHC.

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.
Paper Structure (31 sections, 4 equations, 5 figures, 9 tables)

This paper contains 31 sections, 4 equations, 5 figures, 9 tables.

Table of Contents

  1. Introduction
  2. Background and Related Work
  3. Jet Tagging
  4. Low Latency Jet Tagging on FPGAs
  5. Allo
  6. JetFormer
  7. JetFormer Architecture
  8. Training
  9. Hyperparameter Optimization (HPO)
  10. Optimization settings.
  11. Constraints.
  12. Sampler comparison.
  13. Pruning
  14. Quantization
  15. FPGA Implementation
  16. ...and 16 more sections

Figures (5)

  • Figure 1: The model architecture and transformer block design.
  • Figure 2: The architecture of BitLinear layer wang2023bitnet.
  • Figure 3: Illustration of replacing standard Linear layers with BitLinear layers in the JetFormer architecture.
  • Figure 4: Hypervolume of TPESampler, NSGAIISampler, and BoTorchSampler over 80 valid trials.
  • Figure 5: The Pareto front plot using NSGAIISampler on 8-particle 3-feature dataset.