PaperlayerView exact revision on arXiv

Source-linked AI summary

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

arXiv:2601.17215v1cs.LGhep-ex

TL;DR

Jet tagging requires models that retain high accuracy while meeting the severe latency and resource constraints of LHC online triggers. JetFormer addresses this with a scalable encoder-only Transformer and hardware-aware optimization, matching ParT accuracy within 0.7% with 37.4% fewer FLOPs and enabling compressed FPGA variants. The paper also identifies frontend and hardware-optimization limitations that motivate further co-design.

  • Problem

    LHC Level-1 triggers must select interesting events under stringent timing and resource constraints, while existing Transformer jet taggers are often too large for FPGA deployment.

  • Method

    JetFormer is a scalable encoder-only Transformer paired with multi-objective optimization, structured pruning, 1-bit quantization, and Allo-based FPGA synthesis.

  • Results

    JetFormer matches ParT accuracy within 0.7% on JETCLASS while using 37.4% fewer FLOPs, and compressed variants are synthesized for FPGA deployment.

  • Takeaways & Limitations

    A single modular architecture can support both high-accuracy offline jet tagging and low-latency online triggering through scaling and hardware-aware compression.

  • Takeaways & Limitations

    Allo conversion has unsupported Transformer-related operations, and the FPGA implementation used no dedicated hardware optimizations, leaving further latency reduction for future work.

Abstract

from arXiv · show

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.

1 Introduction

JetFormer is introduced as a scalable encoder-only Transformer for jet tagging across offline analysis and FPGA-based online triggering. It combines competitive performance with hardware-aware compression and synthesis to support this range of deployment requirements.

  • FPGA deployment matters because the LHC Level-1 trigger must rapidly reduce collision data under stringent timing and buffer constraints.The Level-1 trigger uses FPGAs for low latency, throughput, and custom design.
  • JetFormer matches ParT accuracy within 0.7% while using 37.4% fewer FLOPs on the large-scale JETCLASS dataset.The comparison targets the interaction-rich ParT model.
  • JetFormer processes unordered particle-feature sets without explicit interaction terms and is designed for both offline analysis and ultra-low-latency online triggering.Its modular architecture supports different model sizes for differing accuracy and latency requirements.
  • The hardware-aware pipeline combines multi-objective hyperparameter tuning, structured pruning, and 1-bit quantization for FPGA-based trigger deployment.The target setting requires sub-microsecond latency.
  • Pruning cuts computational cost by 50% with less than 0.5% accuracy loss, while quantization reduces model size by 82–92% with a 1.5–3.5% absolute accuracy drop.These compressed variants, including JetFormer-tiny, are synthesized onto FPGA hardware.
  • The work extends Allo to support Transformer operations and evaluates JetFormer’s accuracy, efficiency, and compressibility across multiple datasets.The stated contributions include end-to-end FPGA synthesis of compressed models and comprehensive evaluations.

2 Background and Related Work

Jet tagging research spans graph, particle-cloud, and Transformer architectures, while deployment needs motivate hardware-aware designs. Allo provides the compiler framework used to transform models into FPGA-oriented implementations.

  • JEDI-net models jets as graphs with pairwise particle interactions, whereas ParticleNet treats them as unordered particle clouds using dynamic graph convolutions.Both are established approaches for jet-tagging classification.
  • ParT achieves state-of-the-art jet-tagging performance using particle multi-head attention with pairwise features as an attention bias.Other Transformer variants target reduced computational cost or incorporate spatial information.
  • Transformer jet-tagging models are often too large for FPGA-based Level-1 triggers and are therefore better suited to offline or higher-latency trigger systems.This deployment gap motivates adapting and optimizing Transformers for FPGA use.
  • Earlier FPGA jet taggers used MLPs with limited accuracy, while later GNN work applied quantization, pruning, and efficient graph representations under Level-1 trigger constraints.JEDI-net was the first cited FPGA implementation of a GNN-based jet tagger.
  • Prior real-time jet-tagging work has reported transformer latencies of O(1µs) and O(100ns) for particle-level inputs.The cited approaches include Multi-Head Attention and Linformer.
  • Allo lowers PyTorch or Hugging Face models through MLIR to LLVM IR or HLS C/C++ for CPU and FPGA targets.Its frontend defines reusable neural-network kernels, while its backend schedules and optimizes them.

3 JetFormer

JetFormer is an encoder-only Transformer for particle-jet tagging, adapted for both accuracy and hardware deployment through architectural simplifications and multi-objective optimization.

  • Architecture: JetFormer models particle sequences with input embedding, stacked Transformer blocks, and a class-token classification head while omitting positional encoding.The design preserves permutation invariance and uses the class-token representation for jet classification.
  • Architecture: Transformer blocks combine multi-head self-attention, feedforward layers, layer normalization, SiLU activation, and residual connections.The block structure is illustrated in Figure 1b.
  • Hardware adaptations: Batch normalization replaces layer normalization and ReLU replaces SiLU to simplify inference-time computation for hardware deployment.Batch-normalization parameters are fixed during inference, while ReLU avoids exponential and division operations.
  • Training: OneCycleLR is selected after scheduler comparison because it provides the best validation performance and fastest convergence.Training uses AdamW, early stopping, and dataset-specific batch sizes.
  • Hyperparameter optimization: Multi-objective Optuna optimization searches for models that maximize accuracy while minimizing FLOPs under a minimum validation-accuracy constraint.The workflow compares NSGAIISampler, TPESampler, and BoTorchSampler using Pareto-front hypervolume.

1-Bit Quantization

JetFormer uses quantization-aware training to replace selected linear layers with BitLinear layers, binarizing weights while retaining higher-precision activations for stability.

  • Quantization: 1-bit quantization restricts weight matrices to +1 or −1, converting matrix multiplication into addition-only computation.The method is contrasted with FP16 arithmetic, which requires multiplication and addition.
  • Weight quantization: Weights are centralized, binarized with a sign function, and rescaled with β to approximate their original values.The cited equations define the centering, binarization, and dequantization stages.
  • Activation quantization: Absmax quantization maps activation vectors to −127 through 127 using γ based on the infinity norm, then dequantizes them by inverse scaling.The procedure rounds scaled activations to the nearest integer.
  • Quantization-aware training: The straight-through estimator treats quantization as an identity during backpropagation, enabling training despite non-differentiable sign and round operations.Forward passes simulate low-precision computation through quantization and dequantization.
  • JetFormer integration: BitLinear layers replace attention projections and feedforward layers, while input embedding and classification output layers remain unquantized.Weights use 1-bit precision and activations use 8-bit precision to preserve training stability and probability outputs.

4 FPGA Implementation

The FPGA workflow uses Allo to validate, extend, and synthesize JetFormer models, with compiler support added for transformer operations missing from the original frontend and kernels.

  • Deployment workflow: The Allo workflow first validates a particle MLP, then extends the framework and translates JetFormer-tiny for hardware deployment.Validation proceeds through progressively more hardware-oriented stages.
  • Frontend limitations: Allo initially fails on class-token operations and several operators, including log softmax, ReLU, and 3D batch normalization.These unsupported or partially implemented operations cause compilation failures during model translation.
  • Allo extension: Frontend parsing and kernel support are extended for log_softmax, relu3d, batchnorm1d, and class-token expansion, concatenation, and slicing.The additions provide compatibility with Allo’s compilation flow.
  • Validation model: The reference particle MLP contains hidden layers of 64, 32, and 32 neurons and is trained on the 8-particle, 3-feature dataset.It serves as the initial hardware-feasibility model.
  • JetFormer-tiny deployment: The selected JetFormer-tiny has 4 transformer blocks, embedding dimension 8, 2 attention heads, 26,168 FLOPs, and 3,080 parameters before compression.The 50% compressed version uses 13,784 FLOPs and 1,997 parameters.

5 Evaluation

JetFormer is evaluated across HLS4ML 150-particle and JETCLASS datasets using accuracy, efficiency, and compressibility benchmarks. On HLS4ML datasets, it consistently outperforms prior compact models, while the evaluation motivates compression for deployment.

  • The HLS4ML evaluation compares JetFormer with MLP, Deep Sets, Interaction Network, and JEDI-net baselines under matched particle-count and feature settings.The 3-feature comparison uses truncated 8-, 16-, and 32-particle datasets, while the JEDI-net comparison uses the full 16 features and up to 150 particles.
  • JetFormer improves accuracy by approximately 3% to 4% across the evaluated HLS4ML datasets, with higher AUC scores for each jet class.Comparisons cover 3-feature datasets with 8, 16, and 32 particles and 16-feature datasets with 30, 50, 100, and 150 particles.
  • JetFormer’s higher accuracy is accompanied by a significantly larger model size than the other evaluated models, motivating subsequent compression for hardware deployment.The evaluation frames model compression as necessary to reduce computational cost and inference latency.

Results on JETCLASS Dataset

On a 2M-sample JETCLASS subset, a larger JetFormer configuration achieves competitive performance with ParT without using pairwise interaction features. It nearly matches ParT while requiring substantially fewer FLOPs.

  • 0.7% lower accuracy and 0.07% lower AUC than ParT are reported for JetFormer without additional pairwise interaction features.Both models are evaluated on a 2M-sample training and validation subset.
  • JetFormer uses standard per-particle JETCLASS features while excluding the dataset’s pairwise interaction features.This preserves the deployment-oriented input design while enabling comparison with the interaction-rich ParT model.
  • 37.4% fewer FLOPs than ParT demonstrates JetFormer’s computational efficiency on the JETCLASS comparison.The larger JetFormer uses 10 transformer encoder layers, embedding dimension 128, and 8 attention heads.

Sampler Selection

The sampler study compares TPESampler, NSGAIISampler, and BoTorchSampler using hypervolume over 80 valid trials. NSGAIISampler converges faster and is selected for the final optimization process.

  • NSGAIISampler achieves the highest hypervolume during the study, with faster early convergence and steady improvement.All three samplers eventually reach a comparable hypervolume of around 0.83 over 80 valid trials.
  • TPESampler starts more slowly but eventually approximates NSGAIISampler, while BoTorchSampler lags slightly in early and later phases.The comparison uses hypervolume, which captures both convergence and diversity of Pareto solutions.

Best Trials

The NSGAIISampler search identifies Pareto-optimal accuracy–FLOPs trade-offs and selects a low-cost trial as JetFormer-tiny. Subsequent pruning and 1-bit quantization substantially reduce deployment costs with bounded accuracy losses.

  • Best Trials: The first trial is selected as JetFormer-tiny because it has the fewest FLOPs while maintaining validation accuracy above 0.65.The Pareto front marks optimal trials in red and infeasible trials below 0.65 accuracy in gray.
  • Best Trials: A four-block model with embedding dimension 128, eight attention heads, and dropout 0.05 achieves 0.673 test accuracy.This is the best model identified by the hyperparameter search.
  • Structured Pruning: 50% pruning reduces JetFormer-tiny to 13.78k FLOPs and 1,997 parameters with around 0.5% accuracy loss.Larger models lose less than 0.3% accuracy after pruning, and JetFormer-tiny GPU inference time decreases by 17.46%.
  • 1-Bit Quantization: 1-bit quantization reduces model size by 82–92% while causing a 1.5–3.5% accuracy drop.The relative accuracy loss reaches 3.5% for inputs with 32 constituents.

5.3 Hardware Evaluation

The hardware evaluation compares FPGA latency and resource utilization across MLP and JetFormer-tiny implementations, while identifying substantial optimization headroom from pipelining and parallelism.

  • MLP hardware evaluation: 0.585 ms latency is reported for the particle MLP at batch size 32, with each hardware resource using less than 1%.The reported latency is relatively high because Allo does not provide full pipelining.
  • JetFormer-tiny hardware evaluation: 0.404 ms latency is achieved by pruned JetFormer-tiny at batch size 2, although latency remains high without pipelining.At batch size 16, latency decreases from 4.767 ms before pruning to 2.705 ms after pruning.
  • JetFormer-tiny hardware evaluation: 50% lower computational cost from pruning is accompanied by reduced JetFormer-tiny latency and resource utilization, while remaining FPGA capacity leaves optimization room.The implementation reports low resource utilization and no dedicated hardware optimizations.
  • Future hardware optimization: Increasing parallelism through deeper pipelining or spatial replication is identified as a future route to reduce latency and use available on-chip resources more efficiently.These hardware optimizations were not applied in the reported implementation.
  • Pipelining comparison: 1.662 us latency after full pipelining improves on 7.439 us without pipelining, but resource utilization increases significantly.The comparison motivates additional pipelining methods for accelerating transformer models within hardware constraints.

6 Conclusion and Future Work

The conclusion presents JetFormer as a unified, scalable transformer for offline accuracy and online latency, combining competitive tagging performance with hardware-oriented compression. Future work targets lower FPGA latency, model–hardware co-design, and extension to other HEP tasks.

  • Conclusion: JetFormer unifies high-accuracy offline jet tagging and low-latency online triggering within one scalable encoder-only transformer architecture.The conclusion frames this unified architecture as the central achievement.
  • Conclusion: 37.4% fewer FLOPs are used while full-scale JetFormer matches ParT accuracy within 0.7% on JETCLASS.Compact variants use multi-objective search, pruning, and 1-bit quantization for FPGA deployment under stringent latency constraints.
  • Future work: Future work includes further FPGA hardware optimization, joint model–hardware co-design, and extending JetFormer to other HEP tasks.These directions are stated as ways to advance the current implementation and scope.
Loading 2601.17215v1…