---
title: 'JEDI-linear: FPGA GNN for Jet Tagging'
url: https://www.emergentmind.com/topics/jedi-linear
type: topic
---

# JEDI-linear: FPGA GNN for Jet Tagging

JEDI-linear is a graph neural network architecture for jet tagging on FPGAs that was designed for the latency, throughput, and resource constraints of the High-Luminosity Large Hadron Collider’s CMS Level-1 trigger. It is defined as a linear-complexity, interaction-based GNN that removes explicit pairwise interactions and replaces them with shared transformations and global aggregation, then co-designs the model with fine-grained quantization-aware training and multiplier-free distributed-arithmetic inference. In the reported FPGA implementation, JEDI-linear achieves 3.7 to 11.5 times lower latency, up to 150 times lower initiation interval, and up to 6.2 times lower LUT usage than prior designs, while eliminating DSP blocks entirely and, in some configurations, achieving less than 60 ns latency [2508.15468]. The term should be distinguished from other arXiv papers using the acronym “JEDI” for unrelated methods, several of which explicitly state that they do not define a “JEDI-linear” variant [2505.19698], [2603.10489].

## 1. Definition and problem setting

JEDI-linear targets jet tagging in the HL-LHC CMS Level-1 trigger, where classification must be performed at 40 MHz under strict latency and throughput constraints. The Level-1 system receives collisions every 25 ns and must make decisions within a few hundred nanoseconds, typically with initiation intervals at or below one clock cycle and within a constrained fraction of a single FPGA, often one SLR. In that regime, conventional Interaction Networks are difficult to deploy because they require explicit edge-wise processing over all node pairs, which induces quadratic computational complexity and irregular memory access patterns [2508.15468].

The task is multiclass jet tagging over five categories: gluons, light quarks, \(W\) bosons, \(Z\) bosons, and top quarks. Prior interaction-based GNNs had already shown strong physics performance, but the FPGA cost of edge traversal, dense multiply-accumulate pipelines, and deep irregular scheduling made them poorly matched to real-time trigger hardware. The stated motivation for JEDI-linear is therefore not merely model compression, but a structural redesign of the interaction mechanism so that the architecture remains compatible with static dataflow, deterministic initiation interval, and low on-chip resource usage [2508.15468].

This leads to a specific design objective: preserve the interaction-network viewpoint while replacing the explicit \(O(N^2)\) message-passing pattern by an \(O(N)\) computation over particles. The resulting model is described as the first interaction-based GNN to achieve less than 60 ns latency and as meeting the current requirements for use in the HL-LHC CMS Level-1 trigger system [2508.15468].

## 2. Linearization of interaction networks

JEDI-linear is derived from the standard Interaction Network formulation on a densely connected directed graph without self-loops. For \(N\) particles with \(P\) features each, the original edge construction is

\[
B_{ij} = I_i \| I_j \in \mathbb{R}^{2P}, \quad
E_{ij} = f_R(B_{ij}) \in \mathbb{R}^{D_E}, \quad
\bar{E}_i = \sum_{j \ne i} E_{ij} \in \mathbb{R}^{D_E}.
\]

In that form, complexity scales as \(\mathcal{O}(N^2 \cdot C_{f_R})\), because every ordered pair \((i,j)\) is processed by the edge network \(f_R\) [2508.15468].

The central architectural step in JEDI-linear is to constrain the edge function to be affine:

\[
f_R(B_{ij}) = W_1 I_i + W_2 I_j + C,
\]

with \(W_1, W_2 \in \mathbb{R}^{D_E \times P}\) and \(C \in \mathbb{R}^{D_E}\). Summing over \(j \ne i\) yields

\[
\bar{E}_i = \sum_{j \ne i} \left(W_1 I_i + W_2 I_j + C\right)
          = W_2 \sum_{j \ne i} I_j + (N-1)(W_1 I_i + C).
\]

After scaling by \(1/N\) and dropping \(O(1/N)\) terms for \(N \gg 1\), the edge aggregation becomes

\[
\bar{E}_i' = \frac{1}{N} \bar{E}_i \approx W_2 \left( \frac{1}{N} \sum_{j} I_j \right) + W_1 I_i + C.
\]

This converts the pairwise interaction pattern into a shared global summary plus a per-node linear term, thereby yielding linear complexity, \(\mathcal{O}(N)\) [2508.15468].

The same idea can be written in message-passing form. Let \(x_i = I_i\) denote node features. JEDI-linear computes a shared transform and a permutation-invariant global aggregation,

\[
z_j = \psi(x_j), \quad G = \frac{1}{N} \sum_{j=1}^N \psi(x_j),
\]

and then updates each node using the broadcast global context,

\[
h_i' = \phi(h_i, G) = W_{\text{node}} h_i + W_{\text{global}} G + b.
\]

This suggests that JEDI-linear is not an arbitrary approximation to an Interaction Network, but a specific linearized interaction mechanism in which the collective effect of all other particles is summarized by mean pooling over a shared transformed representation [2508.15468].

Architecturally, the reported stack is: input projection, shared transform, global average pooling, broadcast, second shared dense transform for node update, final global average pooling, and an MLP classifier head producing five logits for \(g/q/W/Z/t\). The paper does not report residual or normalization variants, and it uses mean pooling rather than alternative aggregators such as max or sum [2508.15468].

## 3. Hardware-oriented model design

The hardware significance of the linearization lies in regularizing both computation and memory access. Conventional Interaction Networks require explicit traversal of pairwise edges, along with repeated concatenation of source and target features and irregular access to edge-associated data. JEDI-linear instead applies a shared dense transform over a stream of node features, accumulates a single global context vector through mean pooling, and broadcasts that compact summary back to all nodes. No adjacency indexing or per-edge memory traversal is required [2508.15468].

This regularity supports static dataflow and a deterministic initiation interval of one clock cycle. The implementation uses fully unrolled, pipelined RTL generated through the da4ml toolflow, with pipeline balancing such as inserting registers every two adders. The resulting designs operate at reported frequencies from 203 MHz to 381.7 MHz, with all data kept on-chip and no off-chip memory access during inference [2508.15468].

A second hardware component is the replacement of multiplier-heavy MAC units by distributed arithmetic. For \(B\)-bit inputs, the multiply-accumulate operation is written as

\[
y = \sum_{n} w_n x_n
  = \sum_{k=0}^{B-1} 2^k \, \left( \sum_{n} w_n \, b_{n,k} \right),
\]

where the inner sums are constant-coefficient partial sums implemented by LUT-based adder trees, and the accumulation over bit planes is realized by shift-add logic. Common subexpression elimination is used to reduce repeated partial sums across outputs. In the reported FPGA realizations, this removes the need for DSP blocks entirely [2508.15468].

The target hardware platform is the AMD Xilinx Virtex UltraScale+ VU13P used in CMS Correlator Trigger Layer 2. CTL2 is described as 30 VU13P FPGAs organized into 5 slices in a round-robin arrangement, with five FPGAs available concurrently for algorithms that must coexist with jet clustering and other trigger components. JEDI-linear is placed downstream of jet preprocessing and processes multiple jets in parallel while synchronizing its output tags with parallel feature streams [2508.15468].

## 4. Quantization-aware training and numerical representation

JEDI-linear is trained with fine-grained quantization-aware training based on High Granularity Quantization. Rather than assigning a uniform precision per layer, the method learns a per-parameter bitwidth jointly with the weights. The training objective is

\[
\mathcal{L} = \mathcal{L}_{\text{task}} + \lambda \cdot \text{EBOPs},
\]

where EBOPs is a differentiable proxy for hardware area and latency, and \(\lambda\) is increased along a single training trajectory to trace an accuracy–cost Pareto frontier. Parameters driven to zero bitwidth are pruned automatically, so quantization and unstructured pruning are integrated in the same optimization [2508.15468].

The paper states that surrogate gradients and straight-through estimators are used to propagate gradients through discretized bitwidths and quantizers. A reported empirical outcome is that the majority of non-zero weights are quantized to at most 3 bits in the 64-particle and 128-particle, 16-feature models, while preserving or improving accuracy relative to uniform 8-bit baselines. This indicates that the hardware savings are not achieved solely through architectural simplification, but also through highly nonuniform mixed-precision allocation [2508.15468].

A plausible implication is that JEDI-linear’s linear-complexity message passing and HGQ are mutually reinforcing. Because the architecture reduces the arithmetic graph to shared transforms and global pooling, the remaining dense operations can be mapped especially efficiently to distributed arithmetic, and the learned low-bit representation further shrinks the required adder-graph logic. The paper frames this as co-design rather than sequential optimization [2508.15468].

## 5. Dataset, evaluation protocol, and empirical results

The reported experiments use the hls4ml LHC Jet dataset, with 620k training jets and 260k test jets balanced across the five target classes. Each jet contains up to 150 particles with 16 kinematic features per particle, and a reduced 3-feature regime using \(p_T\), \(\eta\), and \(\phi\) with a \(p_T > 2\) GeV threshold is also evaluated. For the permutation-invariant variant, the top-\(N\) particles by \(p_T\) are selected; for non-invariant variants, particles are sorted by \(p_T\) [2508.15468].

Performance is reported using accuracy together with post-place-and-route timing and resource usage. The paper does not report AUC or mistag-efficiency curves, and it explicitly notes that alternative aggregators and normalization variants are not explored [2508.15468].

| Configuration | Accuracy | Latency / resources |
|---|---:|---|
| 32 particles, 16 features | 81.4% | 79 ns, 147k LUT, 71k FF, 0 BRAM, 0 DSP |
| 64 particles, 16 features | 82.4% | 93 ns, 192k LUT, 92k FF, 0 BRAM, 0 DSP |
| 128 particles, 16 features | 82.1% | 110 ns, 243k LUT, 111k FF, 0 BRAM, 0 DSP |
| 16 particles, 16 features, non-invariant | 77.6% | 52 ns, 38k LUT, 20k FF, 0 BRAM, 0 DSP |
| 16 particles, 3 features, non-invariant | 71.9% | 54 ns, 44k LUT, 22k FF, 0 BRAM, 0 DSP |

These results are compared against prior FPGA GNN implementations, including LL-GNN J5 and Ultrafast GNN. LL-GNN J5, for example, is reported at 79.9% accuracy, 905 ns latency, initiation interval 150, 9,833 DSPs, and 911k LUTs; Ultrafast GNN at 16 particles and 3 features is reported at 70.8% accuracy, about 180 ns latency, initiation interval 3, 5,362 DSPs, and 1,388k LUTs. Across the benchmark set, JEDI-linear is reported as achieving 3.7 to 11.5 times lower latency, up to 150 times lower initiation interval, up to 6.2 times lower LUT usage, and zero DSP usage, while also improving model accuracy [2508.15468].

The latency results are especially notable in the 16-particle non-invariant configurations, which reach 52 ns and 54 ns at 381.7 MHz and 354.4 MHz respectively. Those are the configurations underlying the claim that JEDI-linear is the first interaction-based GNN to achieve less than 60 ns latency [2508.15468].

## 6. Trigger relevance, limitations, and terminological scope

JEDI-linear matters because it converts interaction-based jet tagging from a largely offline or relaxed-latency concept into a trigger-compatible hardware block. The reported designs satisfy the main Level-1 requirements simultaneously: deterministic latency, initiation interval \(=1\), zero DSP usage, BRAM \(=0\), and resource footprints small enough to fit within CTL2 budgets while leaving room for other algorithms. In the permutation-invariant configurations, latency ranges from 67 ns to 110 ns depending on particle count and feature set; in selected non-invariant configurations, it is below 60 ns [2508.15468].

The principal trade-off is expressivity. By constraining the edge function to be affine and replacing explicit pairwise nonlinear interactions with global mean aggregation, JEDI-linear sacrifices the full relational flexibility of a standard Interaction Network. The paper states that there exist tasks where pairwise edges or attention may remain important, especially when fine-grained local structures dominate. It also notes that very large particle counts could increase accumulator width and make timing closure or buffering more challenging, despite the linear asymptotic scaling [2508.15468].

A second limitation is evaluative rather than architectural. The paper reports accuracy as the primary physics metric, but not AUC or mistag-rate efficiency curves, and it does not study alternative aggregators or normalization variants. This suggests that the current work is optimized around deployability under trigger constraints rather than exhaustive exploration of the model-design space [2508.15468].

Finally, the term “JEDI-linear” is specific to this jet-tagging FPGA architecture. Other contemporary arXiv papers use “JEDI” for unrelated ideas in model-based reinforcement learning, neuroscience, test-time diffusion adaptation, and encoder-only NLI, and several of those papers explicitly clarify that they do not define a “JEDI-linear” variant [2505.19698], [2603.10489], [2509.18901]. In current arXiv usage, therefore, JEDI-linear denotes not a generic linearized JEDI family, but a particular linear-complexity, FPGA-oriented graph architecture for real-time jet tagging [2508.15468].

Source: https://www.emergentmind.com/topics/jedi-linear