Official implementation of PHAT-JeT (Patch Hierarchical Attention Transformer for Jet Tagging), an efficient attention-based architecture for real-time particle jet classification in high-energy physics experiments.

PHAT-JeT addresses the challenge of performing accurate jet tagging under the extreme latency and resource constraints imposed by real-time trigger systems at the Large Hadron Collider (LHC). While transformer models achieve state-of-the-art accuracy in particle-level jet classification, their quadratic computational cost hinders deployment in real-time applications where decisions must be made within ~10 microseconds.

Particle jets in the detector coordinate system (η, φ). Each jet consists of collimated particle showers that must be efficiently classified in real-time.

• Geometric Message Passing (GMP): Physics-inspired module that encodes local detector-plane structure using lightweight 2D convolutions on a coarse grid to capture energy flow patterns
• Hierarchical Patch-Based Attention: Replaces global self-attention with efficient patch-wise attention that computes exact attention within small particle groups while preserving global context through lightweight patch-token communication
• Near-Linear Scalability: Tunable patch size enables near-linear scaling with the number of particles while maintaining expressive capacity
• State-of-the-Art Performance: Achieves superior accuracy and background rejection compared to existing efficient models at comparable computational budgets

The PHAT block partitions input particles into patches before computing attention, enabling efficient processing while maintaining global context through patch-token interactions.

Performance comparison on 5-class jet classification (q, g, W, Z, t):

Avg Bkg Rej = Average background rejection at 80% signal efficiency (higher is better)

PHAT-JeT achieves the highest accuracy, ROC AUC, and background rejection while maintaining a competitive parameter count and computational cost suitable for real-time deployment.

• Python 3.8+
• TensorFlow 2.11+
• NumPy
• scikit-learn
• matplotlib

Train PHAT-JeT on the HLS4ML dataset:

Key training arguments:

• --dataset: Dataset to use (hls4ml, top, QG, jetclass)
• --model_size: Model configuration (small, medium, large)
• --sort_by: Particle ordering method (kt, pt, eta, morton)
• --use_flash_attention: Enable Flash Attention for faster training
• --ffn_activation: Activation function (gelu, relu)

Evaluate a trained model:

The model architecture is defined in models/PHAT_JeT.py:

We evaluate PHAT-JeT on three jet tagging benchmarks:

1. HLS4ML LHC Jet Dataset: 5-class classification (q, g, W, Z, t) designed for CMS Level-1 trigger emulation

   • 504K training, 126K validation, 240K test jets
   • Each jet has up to 150 particles with (pT, Δη, Δφ) features

2. Top Tagging Dataset: Binary classification (top quarks vs. light quarks/gluons)

   • 1.2M training, 400K validation, 400K test jets

3. Quark-Gluon Dataset: Binary classification (quark-initiated vs. gluon-initiated jets)

   • 1.8M training, 200K test jets

Detailed view of PHAT-JeT's core components: Geometric Message Passing (GMP) and Hierarchical Patch Attention mechanism.

The GMP module injects local angular context into particle embeddings:

1. Quantize particles to a 2D detector grid (η, φ)
2. Apply depthwise 2D convolution to capture local energy flow
3. Gather features back to particles

This provides the model with physics-informed positional information about local detector structure.

Attention is computed at two levels:

1. Local Attention: Exact self-attention within small patches of particles (e.g., 10-25 particles)
2. Global Attention: Lightweight attention between patch tokens for global context

This hierarchical design reduces computational cost from O(N²) to approximately O(N) while maintaining representational capacity.

This project is licensed under the MIT License - see the LICENSE file for details.