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

Dit artikel introduceert JEDI-linear, een nieuwe Graph Neural Network-architectuur met lineaire complexiteit die fijnkorrelige kwantisatie en vermenigvuldigersvrije bewerkingen gebruikt om recordbrekende latentie en resource-efficiëntie op FPGAs te bereiken, waarmee succesvol aan de real-time jet-tagging-eisen voor het CMS Level-1-triggersysteem van de HL-LHC wordt voldaan.

De kern

Stel je voor dat je een snelheidscontrole voor beveiliging runt op een enorm vliegveld (de Large Hadron Collider). Elke 25 nanoseconden crasht een nieuwe "vlucht" deeltjes tegen de grond, waardoor een chaotische spervuur van puin ontstaat. Jouw taak is om direct naar dit puin te kijken en te beslissen: "Is dit een saaie hoop afval, of is het een zeldzame, waardevolle schat?"

Als je elk enkel stukje puin probeert op te slaan, zul je in een splitseconde je opslagruimte opraken. Je hebt dus een trigger-systeem nodig—een supersnelle filter die in een splitseconde beslissingen neemt om alleen de interessante gebeurtenissen te bewaren.

Hier komt het artikel om de hoek kijken. De auteurs bouwden een nieuwe, supersnelle "hersenen" (genaamd JEDI-linear) om deze beveiligingsagenten te helpen betere beslissingen te nemen, en ze slaagden erin om deze hersenen op een tiny, gespecialiseerde computerchip (een FPGA) te plaatsen die ongelooflijk snel moet werken.

Hier is de uitleg van hun uitvinding met eenvoudige analogieën:

1. Het Probleem: De "Handdruk"-Bottleneck

Vorige methoden om deze deeltjesstralen (genaamd "jets") te sorteren, gebruikten een techniek die lijkt op een massale round-robin handdruk.

• De Oude Manier: Stel je een kamer voor met 64 mensen. Om de groep te begrijpen, vereiste de oude methode dat elk persoon zich omdraaide en elke andere persoon individueel de hand schudde.
• Het Resultaat: Als je 64 mensen hebt, zijn dat meer dan 4.000 handdrukken. Het duurt te lang, en de kamer wordt te vol met mensen die tegelijkertijd proberen te praten. In de wereld van de deeltjesfysica is dit "handdruk"-proces te traag en gebruikt het te veel hardware-ruimte om bruikbaar te zijn voor real-time beveiligingscontroles.

2. De Oplossing: De "Groepsbespreking" (JEDI-linear)

De auteurs beseften dat ze niet nodig hadden dat iedereen individueel de hand schudde. In plaats daarvan bedachten ze een lineaire complexiteit-benadering.

• De Nieuwe Manier: In plaats van individuele handdrukken, stel je voor dat iedereen in de kamer gewoon zijn hand opsteekt om zijn huidige stemming te delen, en een enkele "kapitein" al die stemmingen verzamelt in één grote samenvatting. Vervolgens vertelt de kapitein iedereen: "Hier is de sfeer van de hele groep."
• De Magie: Nu heb je in plaats van 4.000 handdrukken slechts nodig dat 64 mensen één keer spreken. Het werk schaalt lineair op (als je het aantal mensen verdubbelt, verdubbel je het werk, niet verkwart). Dit is het "JEDI-linear"-deel: het houdt de groepscontext vast zonder de rommelige, trage interacties tussen paren.

3. De Hardware-Hacks: Het Op een Tiny Chip Krijgen

Zelfs met de nieuwe "bespreking"-methode moest de hersenen nog steeds klein en snel genoeg zijn om op een specifiek type chip te passen dat in het beveiligingssysteem wordt gebruikt. De auteurs gebruikten twee slimme trucs:

• De "Aangepaste Uniform"-Truc (Quantisatie):
  Normaal gesproken behandelen computers alle getallen op dezelfde manier (alsof elke soldaat dezelfde zware jas krijgt). De auteurs beseften dat sommige delen van de wiskunde zeer gevoelig zijn en hoge precisie nodig hebben (een zware jas), terwijl anderen er niet om geven (een licht T-shirt). Ze trainden het systeem om een "aangepast uniform" te dragen, waarbij ze kleine, efficiënte bit-breedtes toewezen aan getallen die niet veel precisie nodig hebben. Dit verkleinde de geheugenvoetafdruk aanzienlijk.

• De "Geen-Multiplicator"-Truc (Gedistribueerde Aritmetiek):
  Standaard chips gebruiken speciale, dure "multiplicator"-blokken om wiskunde te doen, die zware, stroomverslindende motoren zijn. De auteurs vervangen deze motoren door een slim systeem van optellers en verschuivers (zoals het gebruik van een rekenliniaal of een stapel blokken).

  • Het Resultaat: Ze elimineerden de behoefte aan de zware "multiplicatormotoren" (DSP-blokken) volledig. Dit bespaarde enorme hoeveelheden ruimte en stroom, waardoor het systeem kon draaien op een chip die eerder de last niet aankon.

4. De Resultaten: Snelheid en Efficiëntie

Toen ze dit nieuwe systeem testten tegen de beste bestaande methoden:

• Snelheid: Het is 3,7 tot 11,5 keer sneller. Het kan in minder dan 60 nanoseconden een beslissing nemen (wat sneller is dan een knipoog).
• Efficiëntie: Het gebruikt tot 150 keer minder "starttijd" tussen beslissingen en gebruikt 6,2 keer minder ruimte op de chip.
• Nauwkeurigheid: Ondanks dat het kleiner en sneller is, is het eigenlijk nauwkeuriger in het identificeren van de zeldzame deeltjesjets dan de eerdere, zwaardere modellen.

Waarom Dit Belangrijk Is

De auteurs beweren dat dit de eerste keer is dat een interactie-gebaseerd AI-model snel en klein genoeg is om te worden gebruikt in het Level-1 Trigger-systeem bij de High-Luminosity Large Hadron Collider van CERN.

Denk hierbij aan het upgraden van de luchthavenbeveiliging van een trage, handmatige zoektocht naar een supersnelle, geautomatiseerde scanner die nooit een zeldzaam item mist maar de lijn nooit vertraagt. Dit stelt wetenschappers in staat om zeldzame natuurkundige gebeurtenissen te vangen die eerder te snel waren om te zien, allemaal terwijl ze minder hardware gebruiken dan een standaard rekenmachine.

Kortom: Ze namen een complexe, trage AI, vereenvoudigden de wiskunde zodat deze niet constant met zichzelf hoeft te "praten", kleedden het in op maat gemaakte kleding om ruimte te besparen, en vervingen de zware motoren door lichtgewicht tandwielen. Het resultaat is een supersnelle, tiny hersenen die op een chip past en zeldzame deeltjes in real-time kan opsporen.

Technische samenvatting

1. Probleemstelling

Experimenten in de hoge-energiefysica bij de Large Hadron Collider (LHC) van CERN genereren enorme datavolumes (honderden terabytes per seconde). Om dit te beheersen, moet het Level-1 Trigger (L1T)-systeem gebeurtenissen in real-time filteren (binnen enkele microseconden) met behulp van FPGA's.

• De Uitdaging: Jet-tagging (het identificeren van de oorsprong van deeltjesstralen) is cruciaal voor deze filtering. Hoewel Graph Neural Networks (GNN's), met name Interaction Networks (IN's) zoals JEDI-net, superieure nauwkeurigheid bieden door de interacties tussen deeltjes te modelleren, zijn ze moeilijk te implementeren op FPGA's voor L1T vanwege:
  • Rekencomplexiteit: Standaard GNN's vereisen expliciete berekeningen van paren randen ($O(N^2)$), wat een knelpunt creëert voor jets met veel deeltjes.
  • Hardware-beperkingen: Strikte latentie-eisen (<100 ns), beperkte FPGA-bronnen (typisch <1 Super Logic Region) en de noodzaak van lage startintervallen (II).
  • Ressourcenverbruik: Bestaande op FPGA gebaseerde GNN's vereisen vaak duizenden Digital Signal Processing (DSP)-blokken en grote Look-Up Table (LUT)-aantallen, waardoor ze onpraktisch zijn voor real-world implementatie naast andere algoritmen.

2. Methodologie

De auteurs stellen JEDI-linear voor, een nieuw GNN-architectuur specifiek ontworpen voor hardware-efficiëntie, gecombineerd met geavanceerde optimalisatietechnieken.

A. Algorithmische Innovatie: Lineaire Complexiteit

• Kernconcept: De originele JEDI-net berekent interacties tussen elk paar deeltjes ($O(N^2)$). JEDI-linear herschrijft de rand-interactiefunctie $f_R$ als een eenvoudige affiene transformatie (een enkele dichte laag).
• Wiskundige Afleiding: Door aan te nemen dat $f_R(I_i \| I_j) = W_1 I_i + W_2 I_j + C$, kan de expliciete sommatie over paren worden herschreven als een globale aggregatie. De interactie-inbedding voor deeltje $i$ wordt een functie van het globale gemiddelde van alle deeltjeskenmerken plus een transformatie van de individuele kenmerken van het deeltje.
• Resultaat: Dit reduceert de rekencomplexiteit van kwadratisch $O(N^2)$ naar lineair $O(N)$, waardoor de noodzaak voor expliciete berekeningen op randniveau wordt geëlimineerd terwijl de globale context behouden blijft.

B. Hardware-optimalisatiestrategieën

1. Fin-granulaire Quantisatie-bewuste Training (QAT):

   • In tegenstelling tot uniforme quantisatie gebruiken de auteurs een per-parameter bitbreedte-optimalisatie aanpak.
   • Met behulp van een differentieerbare surrogate-gradiënt wijst het trainingsproces automatisch specifieke bitbreedtes toe aan elke gewicht op basis van de impact op nauwkeurigheid en hardwarekosten (gemeten in Effective Bit Operations, EBOPs).
   • Dit maakt mixed-precision-modellen mogelijk waarbij veel gewichten worden gesneden (bitbreedte gedreven naar nul) of gereduceerd tot 1-2 bits, wat de modelgrootte aanzienlijk verkleint zonder in te leveren op nauwkeurigheid.

2. Gedistribueerde Aritmetiek (DA) voor Vermenigvuldiger-vrije MAC's:

   • Om het bronverbruik verder te reduceren, vervangt de implementatie conventionele vermenigvuldigers door Gedistribueerde Aritmetiek.
   • DA decomposeert matrix-vectorvermenigvuldigingen in verschuif- en opteloperaties die worden geïmplementeerd via LUT's.
   • Resultaat: Het ontwerp elimineert volledig de noodzaak voor DSP-blokken, en vertrouwt uitsluitend op LUT's en registers, die op FPGA's ruimer voorhanden en flexibeler zijn.

3. Volledig Ontrollen Dataflow-architectuur:

   • Het ontwerp maakt gebruik van een volledig statische, ontrollen dataflow waarbij elke operatie wordt toegewezen aan dedicated hardware.
   • Dit vermijdt bronnen-deling en besturingsoverhead, en maakt een startinterval van 1 cyclus en deterministische, ultra-lage latentie mogelijk.

3. Belangrijkste Bijdragen

• JEDI-linear Architectuur: De eerste interactie-gebaseerde GNN voor jet-tagging die lineaire complexiteit bereikt door expliciete paren-interacties te verwijderen, waardoor het schaalbaar is voor grote aantallen deeltjes.
• Hardware-bewust Co-Design: Integratie van fin-granulaire mixed-precision quantisatie en Gedistribueerde Aritmetiek om een vermenigvuldiger-vrije, DSP-vrije implementatie te creëren.
• Automatiseringsframework: Een uitgebreid da4ml-framework dat automatisch symbolische berekeningsgrafieken traceert en synthetiseerbare Verilog genereert voor deze complexe, ontrollen architecturen.
• Open Source: Publicatie van JEDI-linear-sjablonen en code ter ondersteuning van reproduceerbaarheid.

4. Experimentele Resultaten

De modellen werden geëvalueerd op AMD VU13P FPGA's gericht op het CMS Level-1 Trigger-systeem (Correlator Layer 2).

• Latentie & Doorvoer:
  • Bereikte <60 ns latentie (bijv. 52 ns voor 16 deeltjes met 16 kenmerken) en een startinterval van 1 klokcycli.
  • Dit is 3,7x tot 11,5x lagere latentie en tot 150x lager startinterval vergeleken met state-of-the-art (SOTA) GNN-ontwerpen (bijv. LL-GNN, JEDI-net-varianten).
• Efficiëntie van Bronnen:
  • 0 DSP-blokken gebruikt in alle configuraties (SOTA-ontwerpen gebruiken vaak 5.000–9.000+ DSP's).
  • Tot 6,2x lager LUT-gebruik vergeleken met SOTA-modellen.
  • Voorbeeld: Een JEDI-linear-model met 32 deeltjes gebruikt 6,2x minder LUT's en bereikt 11,5x lagere latentie dan het GNN J5-model, terwijl het hogere nauwkeurigheid biedt (81,4% vs 79,9%).
• Nauwkeurigheid:
  • Bereikte tot 82,4% classificatienauwkeurigheid (op invoer met 16 kenmerken en 64 deeltjes).
  • Presteert beter dan DeepSets (DS) en eerdere GNN-implementaties over verschillende aantallen deeltjes (8 tot 128).
  • Toonde superieure schaalbaarheid, waarbij hoge nauwkeurigheid behouden blijft naarmate het aantal deeltjes toeneemt, terwijl andere modellen degraderen of onhaalbaar worden.

5. Betekenis

• Eerste Real-world Haalbaarheid: Dit is de eerste GNN voor jet-tagging die voldoet aan de strikte latentie- (<60 ns) en bronbeperkingen van de HL-LHC CMS Level-1 Trigger. Het maakt het gebruik van krachtige GNN's in real-time hardware-triggers mogelijk, een taak die voorheen werd gedomineerd door eenvoudigere, minder nauwkeurige modellen.
• Schaalbaarheid: De lineaire complexiteit stelt het systeem in staat om jets met grote aantallen deeltjes (tot 128) te verwerken zonder exponentiële groei van bronnen, waardoor het systeem toekomstbestendig is voor runs met hogere luminositeit.
• Breder Impact: De technieken (gelineariseerde verwerking, mixed-precision quantisatie, gedistribueerde aritmetiek en volledig ontrollen ontwerpen) zijn van toepassing buiten de deeltjesfysica op andere domeinen met lage latentie, zoals betrouwbare DNN's, VAE's en Transformers.

Concluderend slaagt JEDI-linear erin de kloof te overbruggen tussen nauwkeurige deep learning-algoritmen en de extreme beperkingen van real-time hardware-triggers, en eert de weg voor intelligente trigger-systemen van de volgende generatie bij de LHC.