Hardware-Aware Design of a GNN-Based Hit Filtering Algorithm for the Belle II Level-1 Trigger

Cet article présente un flux de travail de compression de modèle adapté au matériel pour un algorithme de filtrage de coups basé sur un réseau de neurones graphiques, permettant son déploiement efficace sur FPGA dans le déclencheur de niveau 1 de Belle II avec une réduction de la complexité computationnelle de plus de deux ordres de grandeur tout en préservant des performances élevées.

L'essentiel

Imaginez que le laboratoire Belle II est une immense usine de particules où des collisions se produisent des milliards de fois par seconde. C'est un peu comme une tempête de confettis géante, mais au lieu de papier, ce sont des particules subatomiques.

Le problème ? La plupart de ces "confettis" sont du bruit (des débris inutiles), et seuls quelques-uns contiennent le message secret que les physiciens cherchent (de nouvelles lois de la physique).

Le défi : Le tri à la volée
Pour ne pas être submergés par la quantité de données, l'usine possède un gardien de sécurité ultra-rapide appelé le "déclencheur de niveau 1" (Level-1 Trigger). Ce gardien a une tâche impossible : il doit examiner chaque collision en 5 millionièmes de seconde (5 microsecondes) et décider quoi garder et quoi jeter. S'il est trop lent, l'usine s'embouteille. S'il est trop bête, il jette les trésors avec les ordures.

La solution : Un détective intelligent mais économe
Les chercheurs de ce papier ont créé un détective numérique basé sur un type d'intelligence artificielle appelé "Réseau de Neurones à Graphes" (GNN).

• L'analogie du réseau : Imaginez que chaque collision est une toile d'araignée. Les points de la toile sont les particules, et les fils sont leurs liens. Le détective regarde cette toile pour repérer les motifs qui ressemblent à un "trésor" (une particule intéressante) et ignore les nœuds inutiles.

Le problème du matériel (Le FPGA)
Ce détective est très intelligent, mais il est aussi très gourmand. Pour le faire fonctionner dans le gardien de sécurité (qui est une puce électronique spéciale appelée FPGA), il doit être :

1. Très petit pour tenir dans la puce.
2. Très rapide pour respecter le délai de 5 microsecondes.
3. Économe en énergie (les puces ne peuvent pas chauffer trop).

Le modèle original était comme un éléphant dans un magasin de porcelaine : trop gros et trop lent pour la puce.

La recette de compression : Transformer l'éléphant en souris
Les auteurs ont appliqué une recette en trois étapes pour rendre ce détective compatible avec la puce, sans le rendre stupide :

1. Réduire la taille (Le déménagement) : Ils ont simplifié la structure du détective. Au lieu d'avoir plusieurs étages de réflexion (couches cachées), ils ont réduit le nombre de pièces. C'est comme passer d'une grande villa à un studio meublé : on garde l'essentiel, on enlève le superflu.
2. Parler plus simplement (La quantification) : Le détective original parlait avec une précision infinie (des nombres à virgule flottante très complexes). Les chercheurs l'ont forcé à parler avec des nombres entiers simples, comme si on passait d'une conversation en haute définition à une conversation en "Morse" (4 bits). C'est beaucoup plus rapide à traiter pour la machine.
3. La taille du filet (L'élagage) : Ils ont coupé les liens inutiles dans la toile d'araignée. Au lieu de connecter chaque point à tout le monde, ils ont simplifié les connexions et supprimé les poids (les décisions) qui n'apportaient rien. C'est comme élaguer un arbre : on enlève les branches mortes pour qu'il pousse plus vite.

Le résultat : Un champion de la rapidité
Grâce à ces astuces, ils ont réussi à réduire la quantité de travail mathématique nécessaire par un facteur de 100 (deux ordres de grandeur !).

• Avant : Le détective prenait trop de temps et de place.
• Après : Il est si léger qu'il rentre parfaitement dans la puce électronique.

La performance : Pas de compromis sur la qualité
Le plus incroyable, c'est que malgré cette compression massive, le détective n'est pas devenu bête.

• Sur des données réelles de collisions, le modèle original avait un score de réussite de 97,4 %.
• Le modèle compressé a un score de 96,8 %.

C'est une différence infime ! Ils ont réussi à transformer un éléphant en souris sans perdre sa capacité à chasser les souris.

En résumé
Ce papier raconte l'histoire de comment les physiciens ont pris une intelligence artificielle très puissante mais trop lourde, et l'ont "compressée" comme un fichier ZIP pour qu'elle puisse courir à toute vitesse sur une puce électronique, tout en continuant à faire son travail de tri avec une précision quasi parfaite. C'est un exemple parfait de "co-conception" : on ne construit pas le logiciel, puis le matériel, mais on les façonne ensemble pour qu'ils soient parfaits l'un pour l'autre.

Résumé technique

1. Problématique

L'expérience Belle II, située au collisionneur SuperKEKB, opère à une luminosité élevée, générant un bruit de fond induit par le faisceau de plus en plus important. Pour gérer ce flux de données, le système d'acquisition de données (DAQ) repose sur un déclencheur matériel de niveau 1 (L1) qui doit sélectionner les événements physiques pertinents en temps réel.

• Contraintes sévères : Le déclencheur L1 doit fonctionner avec une latence stricte de 5 µs et respecter des limites de bande passante.
• Défi spécifique : Le sous-système du déclencheur pour la Chambre à Fils de Dérive Centrale (CDC) doit filtrer les "hits" (impulsions) avant la reconstruction des traces. Avec l'augmentation du bruit de fond, l'efficacité de reconstruction et la pureté des traces sont compromises.
• Objectif : Développer un algorithme de filtrage de hits basé sur des Réseaux de Neurones à Graphes (GNN) capable de s'exécuter sur des FPGA (Field-Programmable Gate Arrays) tout en respectant les contraintes de latence (sub-microseconde) et de ressources matérielles, tout en préservant la performance physique.

2. Méthodologie : Conception Logicielle-Hardware Co-Design

Les auteurs proposent un flux de travail itératif de compression de modèle, allant d'un modèle de référence en précision flottante à une version optimisée pour le matériel. Ce processus repose sur trois piliers principaux :

A. Architecture du Modèle de Base

Le modèle initial est une Interaction Network (un type de GNN léger) implémentée dans PyTorch Geometric.

• Structure : Trois blocs MLP séquentiels : un bloc d'arêtes ($R_1$) pour mettre à jour les caractéristiques des arêtes, un bloc de nœuds ($O$) pour mettre à jour les caractéristiques des nœuds, et un bloc final d'arêtes ($R_2$) pour la classification.
• Entrées : Les nœuds représentent les fils de détection (coordonnées $x, y$ et somme des comptes ADC). Les arêtes capturent les différences spatiales ($\Delta r, \Delta \phi$) et temporelles ($\Delta TDC$) entre les fils connectés.
• Paramètres initiaux : 495 paramètres entraînables.

B. Pipeline de Compression

Pour adapter le modèle aux contraintes du FPGA (AMD Ultrascale XCVU190), les étapes suivantes ont été appliquées :

1. Réduction de la taille du modèle et du graphe :
   • Réduction du nombre de couches cachées dans les MLP (de 2 à 1) et de la taille des couches (de 8 à 6 neurones), ramenant les paramètres de 495 à 211.
   • Passage d'arêtes bidirectionnelles à unidirectionnelles, réduisant de moitié le nombre d'arêtes et la charge de calcul.
2. Quantification (Précision Mixte) :
   • Utilisation de l'entraînement conscient de la quantification (QAT) via la bibliothèque Brevitas.
   • Conversion en arithmétique fixe (fixed-point) : entrées et poids en 4 bits, activations en 6 bits, biais en 16 bits, sorties en 8 bits.
   • Suppression de l'activation sigmoïde finale (remplacée par un seuillage simple) pour l'inférence.
3. Élagage (Pruning) :
   • Application d'un élagage non structuré basé sur la magnitude, augmentant linéairement la parcimonie des poids jusqu'à 65 % de poids nuls.

C. Métrique de Complexité : Bit Operations (BOPs)

Au lieu de se fier uniquement au nombre d'opérations (MAC), les auteurs utilisent le nombre d'Opérations de Bits (BOPs) comme métrique de complexité adaptée au matériel. Cette métrique prend en compte la largeur des bits des poids, des entrées et des accumulateurs, offrant une estimation plus précise de la consommation de ressources FPGA.

3. Résultats Clés

Performance Physique (Filtrage de Hits)

L'évaluation a été réalisée sur des données de collision réelles de Belle II (2024) et des simulations.

• Métrique : Courbes ROC et Score AUC (Area Under the Curve).
• Résultat : Le modèle complet (référence) obtient un AUC de 0,974. Le modèle final compressé et élagué atteint un AUC de 0,968.
• Impact : La dégradation de performance est minime. À une efficacité de hits de 95 %, le taux de rejet du bruit de fond passe de 94,2 % (référence) à 90,9 % (modèle compressé).

Complexité et Efficacité Matérielle

• Réduction des BOPs : Le nombre d'opérations de bits est réduit de plus de deux ordres de grandeur, passant de 116,6 MBOPs (modèle complet) à 1,8 MBOPs (modèle final) pour le plus grand secteur du CDC.
• Conformité aux contraintes : Cette valeur de 1,8 MBOPs se situe dans la plage cible définie (1,0 - 2,5 MBOPs) compatible avec les ressources du FPGA.
• Validation FPGA : Une implémentation "out-of-context" sur un FPGA AMD Ultrascale XCVU190 a confirmé :
  • Une latence totale de pipeline de 632,4 ns (à 128 MHz), respectant l'exigence de sub-microseconde.
  • Une utilisation modérée des ressources : 35,65 % des LUTs et 29,75 % des Flip-Flops.
  • Aucune utilisation de DSP (Digital Signal Processors), ce qui est un avantage majeur pour la flexibilité du design.

4. Contributions Principales

1. Flux de travail "Hardware-Aware" : Démonstration d'une approche où les limitations matérielles (précision fixe, ressources FPGA) guident directement la conception de l'architecture du réseau de neurones, plutôt que de simplement compresser un modèle existant a posteriori.
2. Optimisation pour le L1 Trigger : Première implémentation réussie d'un GNN pour le filtrage de hits dans le déclencheur L1 de Belle II, répondant aux contraintes de latence extrême.
3. Métrique BOPs : Utilisation et validation du nombre d'opérations de bits comme indicateur fiable pour prédire la faisabilité d'un modèle sur FPGA dans le contexte de la physique des hautes énergies.
4. Équilibre Performance/Ressources : Démonstration qu'une compression agressive (quantification 4 bits, élagage 65 %) permet de réduire drastiquement la complexité computationnelle sans sacrifier significativement la capacité de discrimination physique.

5. Signification

Ce travail est crucial pour l'avenir de l'expérience Belle II. En permettant un filtrage intelligent et efficace du bruit de fond directement au niveau du déclencheur matériel, il garantit que le système peut continuer à fonctionner à des luminosités élevées sans saturer les capacités de lecture des données. La méthodologie développée ici sert de modèle pour l'intégration future de l'apprentissage automatique (Machine Learning) dans les systèmes de déclencheurs matériels de haute performance, prouvant que les GNN peuvent être déployés avec une latence ultra-faible sur des FPGA commerciaux.