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

Dieser Beitrag stellt JEDI-linear vor, eine neuartige Graph-Neural-Network-Architektur mit linearer Komplexität, die feingranulare Quantisierung und multipliziererfreie Operationen nutzt, um auf FPGAs eine latenz- und ressourceneffiziente Leistung von Rekordniveau zu erzielen und dabei die Echtzeit-Anforderungen für das Jet-Tagging im CMS-Level-1-Trigger-System des HL-LHC erfolgreich zu erfüllen.

Das Wesentliche

Stellen Sie sich vor, Sie leiten einen Hochgeschwindigkeits-Sicherheitskontrollpunkt an einem riesigen Flughafen (dem Large Hadron Collider). Alle 25 Nanosekunden stürzt ein neuer „Flug" von Teilchen auf den Boden und erzeugt ein chaotisches Sprühfeuer aus Trümmern. Ihre Aufgabe ist es, dieses Sprühfeuer sofort zu betrachten und zu entscheiden: „Ist das ein langweiliger Haufen Müll oder ein seltenes, wertvolles Schatzstück?"

Wenn Sie versuchen, jedes einzelne Trümmerteil zu speichern, werden Sie im Handumdrehen den Speicherplatz erschöpfen. Daher benötigen Sie ein Trigger-System – einen ultraschnellen Filter, der in Bruchteilen von Sekunden Entscheidungen trifft, um nur die interessanten Ereignisse zu behalten.

Hier kommt das Papier ins Spiel. Die Autoren haben ein neues, ultraschnelles „Gehirn" (genannt JEDI-linear) entwickelt, um diesen Sicherheitsbeamten bei besseren Entscheidungen zu helfen, und es geschafft, dieses Gehirn auf einen winzigen, spezialisierten Computerchip (einen FPGA) zu packen, der unglaublich schnell arbeiten muss.

Hier ist die Aufschlüsselung ihrer Erfindung mit einfachen Analogien:

1. Das Problem: Der „Handschlag"-Engpass

Frühere Methoden zum Sortieren dieser Teilchensprühregen (genannt „Jets") verwendeten eine Technik, die einem massiven Rundum-Handschlag ähnlich war.

• Der alte Weg: Stellen Sie sich einen Raum mit 64 Personen vor. Um die Gruppe zu verstehen, verlangte die alte Methode, dass jede einzelne Person sich umdreht und jede andere Person einzeln begrüßt.
• Das Ergebnis: Bei 64 Personen sind das über 4.000 Handschläge. Das dauert zu lange, und der Raum wird zu voll mit Leuten, die gleichzeitig versuchen zu sprechen. In der Welt der Teilchenphysik ist dieser „Handschlag"-Prozess zu langsam und verbraucht zu viel Hardware-Platz, um für Echtzeit-Sicherheitskontrollen nützlich zu sein.

2. Die Lösung: Die „Gruppenhuddle" (JEDI-linear)

Die Autoren erkannten, dass nicht jeder einzeln Handschläge austauschen musste. Stattdessen entwickelten sie einen Ansatz mit linearer Komplexität.

• Der neue Weg: Anstatt individueller Handschläge stellt man sich vor, dass jeder im Raum einfach die Hand hebt, um seine aktuelle Stimmung zu teilen, und ein einzelner „Kapitän" all diese Stimmungen in einer großen Zusammenfassung sammelt. Dann teilt der Kapitän allen mit: „Hier ist die Stimmung der ganzen Gruppe."
• Die Magie: Anstatt 4.000 Handschläge müssen nun nur noch 64 Personen einmal sprechen. Die Arbeit skaliert linear (verdoppelt man die Personen, verdoppelt sich die Arbeit, nicht vervierfacht sie). Dies ist der Teil „JEDI-linear": Er behält den Gruppenzusammenhang bei, ohne die unordentlichen, langsamen paarweisen Interaktionen.

3. Die Hardware-Tricks: Damit es auf einen winzigen Chip passt

Selbst mit der neuen „Huddle"-Methode musste das Gehirn klein und schnell genug sein, um auf einen bestimmten Chip-Typ zu passen, der im Sicherheitssystem verwendet wird. Die Autoren nutzten zwei clevere Tricks:

• Der „Maßgeschneiderte Uniform"-Trick (Quantisierung):
  Normalerweise behandeln Computer alle Zahlen gleich (wie jedem Soldaten denselben schweren Mantel zu geben). Die Autoren erkannten, dass einige Teile der Mathematik sehr empfindlich sind und hohe Präzision benötigen (ein schwerer Mantel), während andere sich wenig darum kümmern (ein leichtes T-Shirt). Sie trainierten das System, eine „maßgeschneiderte Uniform" zu tragen, und wiesen Zahlen, die keine hohe Präzision benötigen, winzige, effiziente Bit-Breiten zu. Dies verkleinerte den Speicherbedarf erheblich.

• Der „Kein-Multiplikator"-Trick (Verteilte Arithmetik):
  Standardchips verwenden spezielle, teure „Multiplikator"-Blöcke für Mathematik, die wie schwere, stromfressende Motoren sind. Die Autoren ersetzten diese Motoren durch ein cleversystem aus Addierern und Verschiebern (wie ein Rechenschieber oder ein Stapel Blöcke).

  • Das Ergebnis: Sie eliminierten vollständig den Bedarf an schweren „Multiplikator-Motoren" (DSP-Blöcke). Dies sparte enorme Mengen an Platz und Leistung, sodass das System auf einem Chip laufen konnte, der die Last zuvor nicht bewältigen konnte.

4. Die Ergebnisse: Geschwindigkeit und Effizienz

Als sie dieses neue System gegen die besten bestehenden Methoden testeten:

• Geschwindigkeit: Es ist 3,7 bis 11,5 Mal schneller. Es kann in weniger als 60 Nanosekunden eine Entscheidung treffen (was schneller ist als ein Augenblinzeln).
• Effizienz: Es benötigt bis zu 150 Mal weniger „Startzeit" zwischen den Entscheidungen und verbraucht 6,2 Mal weniger Platz auf dem Chip.
• Genauigkeit: Trotz seiner geringeren Größe und höheren Geschwindigkeit ist es tatsächlich genauer bei der Identifizierung seltener Teilchenjets als die früheren, schwereren Modelle.

Warum das wichtig ist

Die Autoren behaupten, dies sei das erste Mal, dass ein interaktionsbasiertes KI-Modell schnell und klein genug ist, um im Level-1-Trigger-System am High-Luminosity Large Hadron Collider des CERN eingesetzt zu werden.

Stellen Sie sich vor, Sie rüsten die Flughafensicherheit von einer langsamen, manuellen Durchsuchung auf einen ultraschnellen, automatisierten Scanner um, der niemals ein seltenes Objekt verpasst, aber die Schlange niemals verlangsamt. Dies ermöglicht es Wissenschaftlern, seltene physikalische Ereignisse zu fangen, die zuvor zu schnell waren, um sie zu sehen, und das alles mit weniger Hardware als ein normaler Taschenrechner.

Kurz gesagt: Sie nahmen ein komplexes, langsames KI-Modell, vereinfachten seine Mathematik, sodass es nicht ständig „mit sich selbst sprechen" muss, kleideten es in maßgeschneiderte Kleidung, um Platz zu sparen, und ersetzten seine schweren Motoren durch leichte Zahnräder. Das Ergebnis ist ein ultraschnelles, winziges Gehirn, das auf einem Chip Platz findet und in Echtzeit seltene Teilchen erkennen kann.

Technische Zusammenfassung

1. Problemstellung

Experimente der Hochenergiephysik am Large Hadron Collider (LHC) des CERN erzeugen massive Datenmengen (hunderte Terabyte pro Sekunde). Um diese zu bewältigen, muss das Level-1-Trigger-System (L1T) Ereignisse in Echtzeit (innerhalb weniger Mikrosekunden) unter Verwendung von FPGAs filtern.

• Die Herausforderung: Das Jet-Tagging (Identifizierung des Ursprungs von Teilchenschauern) ist für diese Filterung entscheidend. Während Graph Neural Networks (GNNs), insbesondere Interaktionsnetzwerke (INs) wie JEDI-net, durch die Modellierung von Teilchenwechselwirkungen eine überlegene Genauigkeit bieten, sind sie aufgrund folgender Faktoren schwierig für den Einsatz auf FPGAs im L1T zu implementieren:
  • Rechenkomplexität: Standard-GNNs erfordern explizite paarweise Kantensberechnungen ($O(N^2)$), was einen Engpass für Jets mit vielen Teilchen darstellt.
  • Hardware-Einschränkungen: Strenge Latenzanforderungen (<100 ns), begrenzte FPGA-Ressourcen (typischerweise <1 Super Logic Region) und die Notwendigkeit niedriger Startintervalle (II).
  • Ressourcenverbrauch: Bestehende FPGA-basierte GNNs erfordern oft Tausende von Digital Signal Processing (DSP)-Blöcken und große Look-Up Table (LUT)-Anzahlen, was sie für den praktischen Einsatz neben anderen Algorithmen unpraktisch macht.

2. Methodik

Die Autoren schlagen JEDI-linear vor, eine neuartige GNN-Architektur, die speziell für Hardware-Effizienz entwickelt wurde und mit fortschrittlichen Optimierungstechniken kombiniert ist.

A. Algorithmische Innovation: Lineare Komplexität

• Kernkonzept: Das ursprüngliche JEDI-net berechnet Wechselwirkungen zwischen jedem Teilchenpaar ($O(N^2)$). JEDI-linear reformuliert die Kantenswechselwirkungsfunktion $f_R$ als einfache affine Transformation (eine einzelne dichte Schicht).
• Mathematische Herleitung: Unter der Annahme, dass $f_R(I_i \| I_j) = W_1 I_i + W_2 I_j + C$ gilt, kann die explizite paarweise Summation als globale Aggregation umgeschrieben werden. Die Wechselwirkungseinbettung für Teilchen $i$ wird zu einer Funktion des globalen Durchschnitts aller Teilchenmerkmale plus einer Transformation der Merkmale des einzelnen Teilchens.
• Ergebnis: Dies reduziert die Rechenkomplexität von quadratisch $O(N^2)$ auf linear $O(N)$, wodurch die Notwendigkeit expliziter Kantenebenenberechnungen entfällt, während der globale Kontext erhalten bleibt.

B. Hardware-Optimierungsstrategien

1. Feingranulare quantisierungsbewusste Training (QAT):

   • Im Gegensatz zur uniformen Quantisierung verwenden die Autoren einen Ansatz zur Optimierung der Bitbreite pro Parameter.
   • Unter Verwendung eines differenzierbaren Surrogatgradienten weist der Trainingsprozess automatisch spezifische Bitbreiten für jedes Gewicht basierend auf dessen Einfluss auf Genauigkeit und Hardwarekosten (gemessen durch Effective Bit Operations, EBOPs) zu.
   • Dies ermöglicht gemischte Präzisionsmodelle, bei denen viele Gewichte beschnitten werden (Bitbreite auf Null reduziert) oder auf 1-2 Bits reduziert werden, was die Modellgröße erheblich verringert, ohne die Genauigkeit zu beeinträchtigen.

2. Distribuierte Arithmetik (DA) für Multiplikator-freie MACs:

   • Um den Ressourcenverbrauch weiter zu reduzieren, ersetzt die Implementierung herkömmliche Multiplizierer durch Distribuierte Arithmetik.
   • DA zerlegt Matrix-Vektor-Multiplikationen in Shift-Add-Operationen, die über LUTs implementiert werden.
   • Ergebnis: Das Design eliminiert vollständig die Notwendigkeit von DSP-Blöcken und verlässt sich ausschließlich auf LUTs und Register, die auf FPGAs häufiger und flexibler verfügbar sind.

3. Vollständig entrollte Datenflussarchitektur:

   • Das Design verwendet eine vollständig statische, entrollte Datenflussarchitektur, bei der jeder Vorgang dedizierter Hardware zugeordnet ist.
   • Dies vermeidet Ressourcenteilung und Steuerungs-Overhead und ermöglicht ein Startintervall von 1 Taktzyklus sowie deterministische, ultra-niedrige Latenz.

3. Hauptbeiträge

• JEDI-linear-Architektur: Das erste interaktionsbasierte GNN für Jet-Tagging, das durch die Entfernung expliziter paarweiser Wechselwirkungen eine lineare Komplexität erreicht und somit auf große Teilchenzahlen skalierbar ist.
• Hardware-bewusstes Co-Design: Integration feingranularer gemischter Präzisionsquantisierung und distributierter Arithmetik zur Schaffung einer multiplikatorfreien, DSP-freien Implementierung.
• Automatisierungsframework: Ein erweitertes da4ml-Framework, das symbolische Berechnungsgraphen automatisch verfolgt und synthetisierbaren Verilog-Code für diese komplexen, entrollten Architekturen generiert.
• Open Source: Veröffentlichung von JEDI-linear-Vorlagen und Code zur Unterstützung der Reproduzierbarkeit.

4. Experimentelle Ergebnisse

Die Modelle wurden auf AMD VU13P FPGAs bewertet, die auf das CMS Level-1-Trigger-System (Correlator Layer 2) abzielen.

• Latenz & Durchsatz:
  • Erreichte <60 ns Latenz (z. B. 52 ns für 16 Teilchen mit 16 Merkmalen) und ein Startintervall von 1 Taktzyklus.
  • Dies entspricht einer 3,7- bis 11,5-fach niedrigeren Latenz und einem bis zu 150-fach niedrigeren Startintervall im Vergleich zu State-of-the-Art (SOTA) GNN-Designs (z. B. LL-GNN, JEDI-net-Varianten).
• Ressourceneffizienz:
  • 0 DSP-Blöcke in allen Konfigurationen verwendet (SOTA-Designs verwenden oft 5.000–9.000+ DSPs).
  • Bis zu 6,2-fach geringerer LUT-Verbrauch im Vergleich zu SOTA-Modellen.
  • Beispiel: Ein 32-Teilchen-JEDI-linear-Modell verwendet 6,2-mal weniger LUTs und erreicht eine 11,5-fach niedrigere Latenz als das GNN J5-Modell, bei gleichzeitig höherer Genauigkeit (81,4 % vs. 79,9 %).
• Genauigkeit:
  • Erreichte bis zu 82,4 % Klassifizierungsgenauigkeit (bei 16-Merkmal-Eingaben mit 64 Teilchen).
  • Übertrifft DeepSets (DS) und vorherige GNN-Implementierungen bei verschiedenen Teilchenzahlen (8 bis 128).
  • Zeigte überlegene Skalierbarkeit und behielt hohe Genauigkeit bei steigender Teilchenzahl bei, während andere Modelle an Genauigkeit verlieren oder unpraktikabel werden.

5. Bedeutung

• Erste Realisierbarkeit in der Praxis: Dies ist das erste GNN für Jet-Tagging, das die strengen Latenz (<60 ns) und Ressourcenbeschränkungen des HL-LHC CMS Level-1-Trigger erfüllt. Es ermöglicht den Einsatz leistungsfähiger GNNs in Echtzeit-Hardware-Trigger-Systemen, eine Aufgabe, die bisher von einfacheren, weniger genauen Modellen dominiert wurde.
• Skalierbarkeit: Die lineare Komplexität ermöglicht es dem System, Jets mit großen Teilchenzahlen (bis zu 128) zu verarbeiten, ohne exponentielles Ressourcenwachstum, und macht das System für höhere Luminositätsläufe zukunftssicher.
• Breitere Auswirkungen: Die Techniken (linearisierte Verarbeitung, gemischte Präzisionsquantisierung, distribuierte Arithmetik und vollständig entrollte Designs) sind über die Teilchenphysik hinaus auf andere Low-Latenz-Bereiche wie vertrauenswürdige DNNs, VAEs und Transformer anwendbar.

Zusammenfassend überbrückt JEDI-linear erfolgreich die Lücke zwischen hochgenauen Deep-Learning-Algorithmen und den extremen Einschränkungen von Echtzeit-Hardware-Trigger-Systemen und ebnet den Weg für intelligente Trigger-Systeme der nächsten Generation am LHC.