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

Ce papier présente JEDI-linear, une nouvelle architecture de réseau de neurones à graphes à complexité linéaire qui utilise une quantification fine et des opérations sans multiplicateur pour atteindre une latence et une efficacité des ressources record sur les FPGA, répondant ainsi avec succès aux exigences de marquage des jets en temps réel pour le système de déclenchement de niveau 1 CMS du HL-LHC.

L'essentiel

Imaginez que vous gérez un poste de contrôle de sécurité à haute vitesse dans un aéroport massif (le Grand collisionneur de hadrons). Toutes les 25 nanosecondes, un nouveau « vol » de particules s'écrase au sol, créant un éparpillement chaotique de débris. Votre travail consiste à examiner instantanément cet éparpillement et à décider : « Est-ce un tas de déchets ennuyeux, ou est-ce un trésor rare et précieux ? »

Si vous essayez de sauvegarder chaque morceau de débris, vous manquerez d'espace de stockage en une fraction de seconde. Vous avez donc besoin d'un système de déclenchement — un filtre ultra-rapide qui prend des décisions en un éclair pour ne conserver que les événements intéressants.

C'est ici que l'article intervient. Les auteurs ont construit un nouveau « cerveau » ultra-rapide (appelé JEDI-linear) pour aider ces gardes de sécurité à prendre de meilleures décisions, et ils ont réussi à intégrer ce cerveau sur une puce informatique spécialisée et minuscule (un FPGA) qui doit fonctionner à une vitesse incroyable.

Voici la décomposition de leur invention à l'aide d'analogies simples :

1. Le Problème : Le Goulot d'étranglement de la « Poignée de main »

Les méthodes précédentes pour trier ces éparpillements de particules (appelés « jets ») utilisaient une technique similaire à une poignée de main massive en roue libre.

• L'Ancienne Méthode : Imaginez une pièce avec 64 personnes. Pour comprendre le groupe, l'ancienne méthode exigeait que chaque personne se retourne et serre la main de chaque autre personne individuellement.
• Le Résultat : Si vous avez 64 personnes, cela représente plus de 4 000 poignées de main. Cela prend trop de temps, et la pièce devient trop encombrée de personnes essayant de parler en même temps. Dans le monde de la physique des particules, ce processus de « poignée de main » est trop lent et utilise trop d'espace matériel pour être utile aux contrôles de sécurité en temps réel.

2. La Solution : Le « Rassemblement de Groupe » (JEDI-linear)

Les auteurs ont réalisé qu'ils n'avaient pas besoin que tout le monde se serre la main individuellement. Au lieu de cela, ils ont inventé une approche de complexité linéaire.

• La Nouvelle Méthode : Au lieu de poignées de main individuelles, imaginez que tout le monde dans la pièce lève simplement la main pour partager son humeur actuelle, et qu'un seul « capitaine » rassemble toutes ces humeurs en un seul grand résumé. Ensuite, le capitaine dit à tout le monde : « Voici l'ambiance du groupe entier. »
• La Magie : Maintenant, au lieu de 4 000 poignées de main, vous n'avez besoin que de 64 personnes pour parler une fois. Le travail évolue de manière linéaire (si vous doublez le nombre de personnes, vous doublez le travail, pas vous le quadruplez). C'est la partie « JEDI-linear » : elle maintient le contexte du groupe sans les interactions paires désordonnées et lentes.

3. Les Astuces Matérielles : Faire tenir le tout sur une minuscule puce

Même avec la nouvelle méthode de « rassemblement », le cerveau devait toujours être assez petit et rapide pour tenir sur un type spécifique de puce utilisé dans le système de sécurité. Les auteurs ont utilisé deux astuces ingénieuses :

• L'Astuce de l'« Uniforme Personnalisé » (Quantification) :
  Habituellement, les ordinateurs traitent tous les nombres de la même manière (comme donner à chaque soldat le même manteau lourd). Les auteurs ont réalisé que certaines parties des mathématiques sont très sensibles et nécessitent une haute précision (un manteau lourd), tandis que d'autres s'en fichent (un t-shirt léger). Ils ont entraîné le système à porter un « uniforme personnalisé », attribuant des largeurs de bits minuscules et efficaces aux nombres qui n'ont pas besoin de beaucoup de précision. Cela a considérablement réduit l'empreinte mémoire.

• L'Astuce du « Sans-Multiplicateur » (Arithmétique Distribuée) :
  Les puces standard utilisent des blocs « multiplicateurs » spéciaux et coûteux pour faire des mathématiques, qui sont comme des moteurs lourds et gourmands en énergie. Les auteurs ont remplacé ces moteurs par un système astucieux d'additionneurs et de décalages (comme utiliser une règle à calcul ou une pile de blocs).

  • Le Résultat : Ils ont éliminé complètement le besoin de lourds « moteurs multiplicateurs » (blocs DSP). Cela a économisé d'énormes quantités d'espace et d'énergie, permettant au système de fonctionner sur une puce qui ne pouvait auparavant pas supporter la charge.

4. Les Résultats : Vitesse et Efficacité

Lorsqu'ils ont testé ce nouveau système contre les meilleures méthodes existantes :

• Vitesse : Il est 3,7 à 11,5 fois plus rapide. Il peut prendre une décision en moins de 60 nanosecondes (ce qui est plus rapide qu'un clignement d'œil).
• Efficacité : Il utilise jusqu'à 150 fois moins de « temps de démarrage » entre les décisions et occupe 6,2 fois moins d'espace sur la puce.
• Précision : Malgré sa taille réduite et sa rapidité, il est en réalité plus précis pour identifier les jets de particules rares que les modèles précédents, plus lourds.

Pourquoi cela compte

Les auteurs affirment qu'il s'agit de la première fois qu'un modèle d'IA basé sur les interactions est assez rapide et assez petit pour être utilisé dans le système de déclenchement de niveau 1 du Grand collisionneur de hadrons à haute luminosité du CERN.

Pensez-y comme à la mise à niveau de la sécurité aéroportuaire d'une recherche manuelle et lente vers un scanner automatisé ultra-rapide qui ne manque jamais un objet rare mais ne ralentit jamais la file. Cela permet aux scientifiques de capturer des événements physiques rares qui étaient auparavant trop rapides pour être vus, le tout en utilisant moins de matériel qu'une calculatrice standard.

En résumé : Ils ont pris une IA complexe et lente, simplifié ses mathématiques pour qu'elle n'ait pas besoin de « se parler » constamment, l'ont habillée dans des vêtements sur mesure pour économiser de l'espace, et remplacé ses moteurs lourds par des engrenages légers. Le résultat est un cerveau minuscule et ultra-rapide qui tient sur une puce et peut repérer des particules rares en temps réel.

Résumé technique

1. Énoncé du problème

Les expériences de physique des hautes énergies au Grand collisionneur de hadrons (LHC) du CERN génèrent des volumes de données massifs (centaines de téraoctets par seconde). Pour gérer cela, le système de déclenchement de niveau 1 (L1T) doit filtrer les événements en temps réel (en quelques microsecondes) en utilisant des FPGA.

• Le défi : Le taggage de jets (identification de l'origine des gerbes de particules) est crucial pour ce filtrage. Bien que les réseaux de neurones graphiques (GNN), en particulier les réseaux d'interaction (IN) comme JEDI-net, offrent une précision supérieure en modélisant les interactions entre particules, ils sont difficiles à déployer sur FPGA pour le L1T en raison de :
  • Complexité computationnelle : Les GNN standards nécessitent des calculs explicites d'arêtes par paires ($O(N^2)$), créant un goulot d'étranglement pour les jets comportant de nombreuses particules.
  • Contraintes matérielles : Des exigences de latence strictes (<100 ns), des ressources FPGA limitées (généralement <1 Super Logic Region) et le besoin de faibles intervalles d'initiation (II).
  • Consommation de ressources : Les GNN existants basés sur FPGA nécessitent souvent des milliers de blocs de traitement de signal numérique (DSP) et de grands nombres de tables de recherche (LUT), les rendant peu pratiques pour un déploiement réel aux côtés d'autres algorithmes.

2. Méthodologie

Les auteurs proposent JEDI-linear, une architecture GNN novatrice conçue spécifiquement pour l'efficacité matérielle, combinée à des techniques d'optimisation avancées.

A. Innovation algorithmique : Complexité linéaire

• Concept de base : JEDI-net original calcule les interactions entre chaque paire de particules ($O(N^2)$). JEDI-linear reformule la fonction d'interaction d'arête $f_R$ comme une simple transformation affine (une seule couche dense).
• Dérivation mathématique : En supposant $f_R(I_i \| I_j) = W_1 I_i + W_2 I_j + C$, la sommation explicite par paires peut être réécrite comme une agrégation globale. L'embedding d'interaction pour la particule $i$ devient une fonction de la moyenne globale de toutes les caractéristiques des particules plus une transformation des caractéristiques de la particule individuelle.
• Résultat : Cela réduit la complexité computationnelle de quadratique $O(N^2)$ à linéaire $O(N)$, éliminant le besoin de calculs explicites au niveau des arêtes tout en préservant le contexte global.

B. Stratégies d'optimisation matérielle

1. Entraînement conscient de la quantification à granularité fine (QAT) :

   • Contrairement à la quantification uniforme, les auteurs utilisent une approche d'optimisation de la largeur de bits par paramètre.
   • En utilisant un gradient de substitution différentiable, le processus d'entraînement attribue automatiquement des largeurs de bits spécifiques à chaque poids en fonction de son impact sur la précision et le coût matériel (mesuré par les opérations effectives en bits, EBOPs).
   • Cela permet des modèles de précision mixte où de nombreux poids sont élagués (largeur de bits ramenée à zéro) ou réduits à 1-2 bits, réduisant considérablement la taille du modèle sans sacrifier la précision.

2. Arithmétique distribuée (DA) pour des MAC sans multiplicateur :

   • Pour réduire davantage l'utilisation des ressources, l'implémentation remplace les multiplicateurs conventionnels par de l'arithmétique distribuée.
   • La DA décompose les multiplications matrice-vecteur en opérations de décalage-addition implémentées via des LUT.
   • Résultat : La conception élimine complètement le besoin de blocs DSP, reposant uniquement sur des LUT et des registres, qui sont plus abondants et flexibles sur les FPGA.

3. Architecture de flux de données entièrement déroulée :

   • La conception emploie un flux de données statique et entièrement déroulé où chaque opération est mappée sur du matériel dédié.
   • Cela évite le partage de ressources et la surcharge de contrôle, permettant un intervalle d'initiation de 1 cycle et une latence déterministe ultra-faible.

3. Contributions clés

• Architecture JEDI-linear : Le premier GNN basé sur l'interaction pour le taggage de jets qui atteint une complexité linéaire en supprimant les interactions par paires explicites, le rendant évolutif pour de grands nombres de particules.
• Co-conception consciente du matériel : Intégration de la quantification de précision mixte à granularité fine et de l'arithmétique distribuée pour créer une implémentation sans multiplicateur et sans DSP.
• Cadre d'automatisation : Une extension du framework da4ml qui trace automatiquement les graphes de calcul symboliques et génère du Verilog synthétisable pour ces architectures complexes et déroulées.
• Open Source : Publication des modèles et du code JEDI-linear pour soutenir la reproductibilité.

4. Résultats expérimentaux

Les modèles ont été évalués sur des FPGA AMD VU13P ciblant le système de déclenchement de niveau 1 du CMS (Couche de corrélation 2).

• Latence et débit :
  • Une latence <60 ns a été atteinte (par exemple, 52 ns pour 16 particules avec 16 caractéristiques) et un intervalle d'initiation de 1 cycle d'horloge.
  • Cela représente une latence 3,7 à 11,5 fois plus faible et un intervalle d'initiation jusqu'à 150 fois plus faible par rapport aux conceptions GNN de l'état de l'art (SOTA) (par exemple, LL-GNN, variantes de JEDI-net).
• Efficacité des ressources :
  • 0 bloc DSP utilisé dans toutes les configurations (les conceptions SOTA utilisent souvent 5 000 à 9 000+ DSP).
  • Jusqu'à 6,2 fois moins d'utilisation de LUT par rapport aux modèles SOTA.
  • Exemple : Un modèle JEDI-linear à 32 particules utilise 6,2 fois moins de LUT et atteint une latence 11,5 fois plus faible que le modèle GNN J5 tout en offrant une précision supérieure (81,4 % contre 79,9 %).
• Précision :
  • Une précision de classification allant jusqu'à 82,4 % a été atteinte (sur des entrées à 16 caractéristiques avec 64 particules).
  • Il surpasse DeepSets (DS) et les implémentations GNN précédentes pour divers nombres de particules (de 8 à 128).
  • Il démontre une évolutivité supérieure, maintenant une haute précision à mesure que le nombre de particules augmente, tandis que d'autres modèles se dégradent ou deviennent irréalisables.

5. Importance

• Première faisabilité réelle : Il s'agit du premier GNN pour le taggage de jets à répondre aux contraintes strictes de latence (<60 ns) et de ressources du déclenchement de niveau 1 CMS HL-LHC. Il permet l'utilisation de GNN puissants dans des déclencheurs matériels en temps réel, une tâche précédemment dominée par des modèles plus simples et moins précis.
• Évolutivité : La complexité linéaire permet au système de gérer des jets avec un grand nombre de particules (jusqu'à 128) sans croissance exponentielle des ressources, rendant le système prêt pour les runs de plus grande luminosité.
• Impact plus large : Les techniques (traitement linéarisé, quantification de précision mixte, arithmétique distribuée et conceptions entièrement déroulées) sont applicables au-delà de la physique des particules à d'autres domaines à faible latence tels que les DNN fiables, les VAE et les Transformers.

En conclusion, JEDI-linear comble avec succès le fossé entre les algorithmes d'apprentissage profond à haute précision et les contraintes extrêmes des déclencheurs matériels en temps réel, ouvrant la voie à des systèmes de déclenchement intelligents de nouvelle génération au LHC.