JEDI-linear: Fast and Efficient Graph Neural Networks for Jet Tagging on FPGAs

This paper introduces JEDI-linear, a novel linear-complexity Graph Neural Network architecture that utilizes fine-grained quantization and multiplier-free operations to achieve record-breaking latency and resource efficiency on FPGAs, successfully meeting the real-time jet tagging requirements for the HL-LHC CMS Level-1 trigger system.

The Gist

Imagine you are running a high-speed security checkpoint at a massive airport (the Large Hadron Collider). Every 25 nanoseconds, a new "flight" of particles crashes into the ground, creating a chaotic spray of debris. Your job is to instantly look at this spray and decide: "Is this a boring pile of trash, or is it a rare, valuable treasure?"

If you try to save every single piece of debris, you'll run out of storage space in a split second. So, you need a trigger system—a super-fast filter that makes split-second decisions to keep only the interesting events.

This is where the paper comes in. The authors built a new, super-fast "brain" (called JEDI-linear) to help these security guards make better decisions, and they managed to fit this brain onto a tiny, specialized computer chip (an FPGA) that has to work incredibly fast.

Here is the breakdown of their invention using simple analogies:

1. The Problem: The "Handshake" Bottleneck

Previous methods for sorting these particle sprays (called "jets") used a technique similar to a massive round-robin handshake.

• The Old Way: Imagine a room with 64 people. To understand the group, the old method required every single person to turn around and shake hands with every other person individually.
• The Result: If you have 64 people, that's over 4,000 handshakes. It takes too long, and the room gets too crowded with people trying to talk at once. In the world of particle physics, this "handshake" process is too slow and uses too much hardware space to be useful for real-time security checks.

2. The Solution: The "Group Huddle" (JEDI-linear)

The authors realized they didn't need everyone to shake hands individually. Instead, they invented a linear complexity approach.

• The New Way: Instead of individual handshakes, imagine everyone in the room simply raises their hand to share their current mood, and a single "captain" collects all those moods into one big summary. Then, the captain tells everyone, "Here is the vibe of the whole group."
• The Magic: Now, instead of 4,000 handshakes, you only need 64 people to speak once. The work scales up linearly (if you double the people, you double the work, not quadruple it). This is the "JEDI-linear" part: it keeps the group context without the messy, slow pairwise interactions.

3. The Hardware Hacks: Making it Fit on a Tiny Chip

Even with the new "huddle" method, the brain still needed to be small and fast enough to fit on a specific type of chip used in the security system. The authors used two clever tricks:

• The "Customized Uniform" Trick (Quantization):
  Usually, computers treat all numbers the same way (like giving every soldier the same heavy coat). The authors realized that some parts of the math are very sensitive and need high precision (a heavy coat), while others don't care much (a light t-shirt). They trained the system to wear a "customized uniform," giving tiny, efficient bit-widths to numbers that don't need much precision. This shrank the memory footprint significantly.

• The "No-Multiplier" Trick (Distributed Arithmetic):
  Standard chips use special, expensive "multiplier" blocks to do math, which are like heavy, power-hungry engines. The authors replaced these engines with a clever system of adders and shifters (like using a slide rule or a stack of blocks).

  • The Result: They completely eliminated the need for the heavy "multiplier engines" (DSP blocks). This saved massive amounts of space and power, allowing the system to run on a chip that previously couldn't handle the load.

4. The Results: Speed and Efficiency

When they tested this new system against the best existing methods:

• Speed: It is 3.7 to 11.5 times faster. It can make a decision in less than 60 nanoseconds (which is faster than a blink of an eye).
• Efficiency: It uses up to 150 times less "start-up time" between decisions and uses 6.2 times less space on the chip.
• Accuracy: Despite being smaller and faster, it is actually more accurate at identifying the rare particle jets than the previous, heavier models.

Why This Matters

The authors claim this is the first time an interaction-based AI model has been fast and small enough to be used in the Level-1 Trigger system at CERN's High-Luminosity Large Hadron Collider.

Think of it as upgrading the airport security from a slow, manual search to a super-fast, automated scanner that never misses a rare item but never slows down the line. This allows scientists to catch rare physics events that were previously too fast to see, all while using less hardware than a standard calculator.

In short: They took a complex, slow AI, simplified its math so it doesn't need to "talk to itself" constantly, dressed it in custom-fitted clothes to save space, and replaced its heavy engines with lightweight gears. The result is a super-fast, tiny brain that fits on a chip and can spot rare particles in real-time.

Technical Summary

1. Problem Statement

High-energy physics experiments at the CERN Large Hadron Collider (LHC) generate massive data volumes (hundreds of terabytes per second). To manage this, the Level-1 Trigger (L1T) system must filter events in real-time (within a few microseconds) using FPGAs.

• The Challenge: Jet tagging (identifying the origin of particle sprays) is critical for this filtering. While Graph Neural Networks (GNNs), specifically Interaction Networks (INs) like JEDI-net, offer superior accuracy by modeling particle interactions, they are difficult to deploy on FPGAs for L1T due to:
  • Computational Complexity: Standard GNNs require explicit pairwise edge computations ($O(N^2)$), creating a bottleneck for jets with many particles.
  • Hardware Constraints: Strict latency requirements (<100 ns), limited FPGA resources (typically <1 Super Logic Region), and the need for low initiation intervals (II).
  • Resource Consumption: Existing FPGA-based GNNs often require thousands of Digital Signal Processing (DSP) blocks and large Look-Up Table (LUT) counts, making them impractical for real-world deployment alongside other algorithms.

2. Methodology

The authors propose JEDI-linear, a novel GNN architecture designed specifically for hardware efficiency, combined with advanced optimization techniques.

A. Algorithmic Innovation: Linear Complexity

• Core Concept: The original JEDI-net computes interactions between every pair of particles ($O(N^2)$). JEDI-linear reformulates the edge interaction function $f_R$ as a simple affine transformation (a single dense layer).
• Mathematical Derivation: By assuming $f_R(I_i \| I_j) = W_1 I_i + W_2 I_j + C$, the explicit pairwise summation can be rewritten as a global aggregation. The interaction embedding for particle $i$ becomes a function of the global average of all particle features plus a transformation of the individual particle's features.
• Result: This reduces computational complexity from quadratic $O(N^2)$ to linear $O(N)$, eliminating the need for explicit edge-level computations while preserving global context.

B. Hardware Optimization Strategies

1. Fine-Grained Quantization-Aware Training (QAT):

   • Unlike uniform quantization, the authors use a per-parameter bitwidth optimization approach.
   • Using a differentiable surrogate gradient, the training process automatically assigns specific bitwidths to each weight based on its impact on accuracy and hardware cost (measured by Effective Bit Operations, EBOPs).
   • This allows for mixed-precision models where many weights are pruned (bitwidth driven to zero) or reduced to 1-2 bits, significantly shrinking model size without sacrificing accuracy.

2. Distributed Arithmetic (DA) for Multiplier-Free MACs:

   • To further reduce resource usage, the implementation replaces conventional multipliers with Distributed Arithmetic.
   • DA decomposes matrix-vector multiplications into shift-add operations implemented via LUTs.
   • Outcome: The design completely eliminates the need for DSP blocks, relying solely on LUTs and registers, which are more abundant and flexible on FPGAs.

3. Fully Unrolled Dataflow Architecture:

   • The design employs a fully static, unrolled dataflow where every operation is mapped to dedicated hardware.
   • This avoids resource sharing and control overhead, enabling a 1-cycle initiation interval and deterministic, ultra-low latency.

3. Key Contributions

• JEDI-linear Architecture: The first interaction-based GNN for jet tagging that achieves linear complexity by removing explicit pairwise interactions, making it scalable to large particle counts.
• Hardware-Aware Co-Design: Integration of fine-grained mixed-precision quantization and Distributed Arithmetic to create a multiplier-free, DSP-free implementation.
• Automation Framework: An extended da4ml framework that automatically traces symbolic computation graphs and generates synthesizable Verilog for these complex, unrolled architectures.
• Open Source: Release of JEDI-linear templates and code to support reproducibility.

4. Experimental Results

The models were evaluated on AMD VU13P FPGAs targeting the CMS Level-1 Trigger system (Correlator Layer 2).

• Latency & Throughput:
  • Achieved <60 ns latency (e.g., 52 ns for 16 particles with 16 features) and 1-clock cycle initiation interval.
  • This is 3.7x to 11.5x lower latency and up to 150x lower initiation interval compared to state-of-the-art (SOTA) GNN designs (e.g., LL-GNN, JEDI-net variants).
• Resource Efficiency:
  • 0 DSP blocks used across all configurations (SOTA designs often use 5,000–9,000+ DSPs).
  • Up to 6.2x lower LUT usage compared to SOTA models.
  • Example: A 32-particle JEDI-linear model uses 6.2x fewer LUTs and achieves 11.5x lower latency than the GNN J5 model while offering higher accuracy (81.4% vs 79.9%).
• Accuracy:
  • Achieved up to 82.4% classification accuracy (on 16-feature inputs with 64 particles).
  • Outperforms DeepSets (DS) and previous GNN implementations across various particle counts (8 to 128).
  • Demonstrated superior scalability, maintaining high accuracy as particle count increases, whereas other models degrade or become infeasible.

5. Significance

• First Real-World Feasibility: This is the first GNN for jet tagging to meet the strict latency (<60 ns) and resource constraints of the HL-LHC CMS Level-1 Trigger. It enables the use of powerful GNNs in real-time hardware triggers, a task previously dominated by simpler, less accurate models.
• Scalability: The linear complexity allows the system to handle jets with large numbers of particles (up to 128) without exponential resource growth, future-proofing the system for higher luminosity runs.
• Broader Impact: The techniques (linearized processing, mixed-precision quantization, distributed arithmetic, and fully unrolled designs) are applicable beyond particle physics to other low-latency domains like trustworthy DNNs, VAEs, and Transformers.

In conclusion, JEDI-linear successfully bridges the gap between high-accuracy deep learning algorithms and the extreme constraints of real-time hardware triggers, paving the way for next-generation intelligent trigger systems at the LHC.