Patch Hierarchical Attention Transformer for Efficient Particle Jet Tagging

L'article présente PHAT-JeT, une architecture de transformateur novatrice qui combine un passage de messages géométrique inspiré par la physique avec un mécanisme d'attention hiérarchique basé sur des patches pour atteindre une précision de pointe dans l'étiquetage en temps réel des jets de particules tout en surmontant les contraintes computationnelles des transformateurs standards.

L'essentiel

La vue d'ensemble : le problème de l'aiguille dans une botte de foin

Imaginez le Grand collisionneur de hadrons (LHC) comme une usine massive et ultra-rapide qui fait s'écraser des particules 40 millions de fois par seconde. C'est comme un tuyau d'arrosage projetant un trillion de données chaque seconde.

Le problème ? L'usine ne peut pas sauvegarder toutes ces données. C'est trop. Ainsi, l'usine dispose d'un gardien de sécurité (appelé « système de déclenchement ») posté à la sortie. Ce gardien doit décider en microsecondes (plus vite qu'un clignement d'œil) quelles collisions sont suffisamment intéressantes pour être conservées et lesquelles ne sont que du bruit de fond ennuyeux à jeter.

Les collisions « intéressantes » impliquent souvent des particules de courte durée de vie qui se désintègrent en jets d'autres particules appelés jets. Le travail du gardien est d'examiner un jet et de dire : « Est-ce une particule rare et lourde (comme un quark Top) ou simplement un jet commun (comme un gluon) ? »

Le défi : Vitesse contre Intelligence

Pour ce faire, les scientifiques utilisent des modèles d'intelligence artificielle.

• Les modèles « Super-Cerveau » : Ils sont incroyablement intelligents et précis, mais ils sont énormes et lents. Ils mettent trop de temps à réfléchir, de sorte que le gardien de sécurité ne peut pas les utiliser avant que les données ne s'envolent.
• Les modèles « Rapides » : Ils sont minuscules et rapides, mais pas assez intelligents pour repérer les particules rares et complexes. Ils manquent les « aiguilles » dans la botte de foin.

L'objectif de ce document est de construire un modèle qui soit à la fois assez rapide pour le gardien de sécurité et assez intelligent pour trouver les aiguilles.

La solution : PHAT-JeT (L'organisateur intelligent)

Les auteurs ont créé une nouvelle architecture d'IA appelée PHAT-JeT. Imaginez-la comme une équipe intelligente d'organisateurs tentant de trier un tas chaotique de jouets mélangés (les particules dans un jet).

Au lieu d'essayer de comparer chaque jouet individuel à chaque autre jouet individuel (ce qui prend une éternité), PHAT-JeT utilise trois astuces ingénieuses :

1. La surveillance de quartier (Passage de messages géométrique)

Imaginez que les jouets sont dispersés sur un sol. Avant même que les organisateurs ne commencent à trier, ils observent le sol et remarquent que les jouets proches les uns des autres appartiennent souvent au même groupe.

• L'analogie : PHAT-JeT trace une grille sur le sol. Si un bloc rouge et un bloc bleu se trouvent dans le même carré, ils « parlent » immédiatement. Cela aide le système à comprendre la forme locale du jet (comme une étoile à plusieurs branches) sans avoir besoin de regarder toute la pièce d'un coup. C'est comme se dire : « Hé, ces trois jouets sont regroupés ensemble ; ils viennent probablement de la même boîte à jouets. »

2. Les réunions de petits groupes (Attention par patchs locaux)

Maintenant, les organisateurs divisent les jouets en petits groupes (patchs).

• L'analogie : Au lieu d'une seule réunion géante où 150 personnes essaient de parler à tout le monde (ce qui crée le chaos et prend une éternité), ils se séparent en petits cercles de 10 personnes. À l'intérieur de chaque cercle, tout le monde peut parler à tout le monde parfaitement. Cela capture les détails fins du groupe sans le coût computationnel d'une réunion massive.

3. Les capitaines d'équipe (Attention globale hiérarchique)

Les petits groupes ont un problème : ils ne savent pas ce que font les autres groupes.

• L'analogie : Chaque petit groupe choisit un « Capitaine d'équipe » (un jeton de résumé). Ces capitaines se réunissent dans une salle séparée et plus petite pour partager la vue d'ensemble. Une fois que les capitaines ont compris l'histoire globale, ils retournent dans leurs groupes et disent à tout le monde : « D'accord, en fonction de ce que font les autres groupes, voici le contexte dont vous avez besoin. »
• Le résultat : Le système obtient le meilleur des deux mondes : les détails fins des petits cercles et la vue d'ensemble de la réunion des capitaines.

Pourquoi cela compte

Le document a testé ce nouveau système sur quatre ensembles de données d'« examen » différents (HLS4ML, JetClass, Top Tagging et Quark–Gluon).

• Le résultat : PHAT-JeT a battu tous les autres modèles « rapides ». Il était presque aussi précis que les modèles géants et lents de « Super-Cerveau », mais il fonctionnait assez vite pour s'exécuter sur le matériel spécialisé (FPGA) utilisé par les gardiens de sécurité du LHC.
• L'idée clé : En combinant des petits « cercles » locaux avec une « réunion de capitaines » et en ajoutant une « surveillance de quartier » pour les formes locales, ils ont réussi à comprimer un maximum d'intelligence dans un package minuscule et rapide.

Résumé

PHAT-JeT est une nouvelle façon d'organiser les données qui permet aux expériences de physique des particules de repérer des événements rares et excitants en temps réel. Il le fait en décomposant un problème massif et chaotique en petits groupes locaux gérables, en permettant à ces groupes de communiquer entre eux, puis en faisant en sorte que quelques représentants partagent la vue d'ensemble. C'est la différence entre essayer d'organiser un stade rempli de gens en criant à tout le monde en même temps et les organiser en petites équipes avec des capitaines d'équipe.

Note : Le document se concentre entièrement sur l'amélioration des algorithmes logiciels pour le filtrage des données en physique des particules. Il ne prétend pas modifier la façon dont le matériel est construit, ni ne discute d'applications médicales ou d'autres applications réelles en dehors de la physique des hautes énergies.

Résumé technique

Résumé technique : Patch Hierarchical Attention Transformer pour l'étiquetage efficace des jets de particules (PHAT-JeT)

Énoncé du problème
L'étiquetage des jets en temps réel au Grand collisionneur de hadrons (LHC) constitue un goulot d'étranglement critique pour l'identification des désintégrations de particules à courte durée de vie. Le LHC génère des flux de données dépassant 1 pétaoctet par seconde, mais les systèmes de déclenchement doivent décider en environ 10 microsecondes s'il faut enregistrer un événement. Cela impose des contraintes strictes de latence et de ressources (spécifiquement sur les circuits intégrés à matrice de portes programmables, ou FPGA) qui empêchent le déploiement de modèles hautement expressifs et de pointe comme le Particle Transformer (ParT), qui souffrent d'une complexité computationnelle quadratique ($O(N^2)$) par rapport au nombre de particules $N$. À l'inverse, les modèles efficaces existants qui rentrent dans le budget de déclenchement manquent souvent de capacité de représentation pour distinguer des sous-structures de jets complexes, créant un fossé entre la précision réalisable et la vitesse d'inférence déployable.

Méthodologie : Architecture PHAT-JeT
Les auteurs proposent le Patch Hierarchical Attention Transformer (PHAT-JeT), une architecture conçue pour équilibrer l'efficacité computationnelle avec la préservation des interactions fines entre particules. Le modèle se compose de trois composants principaux :

1. Passage de messages géométrique (GMP) :
   Pour encoder la structure locale du plan du détecteur inhérente à la physique des jets, le modèle introduit un module GMP inspiré par la physique. Les jets sont représentés comme des nuages de points dans le plan $(\eta, \phi)$. Le module GMP quantifie les particules en une grille 2D grossière du détecteur, agrège les caractéristiques au sein des cellules de la grille et applique une convolution 2D légère de type depthwise. Cela propage l'information entre les régions angulaires voisines, permettant aux particules d'intégrer un contexte géométrique local avant d'entrer dans le mécanisme d'attention. Cette étape injecte des a priori structurels concernant les dépôts d'énergie à plusieurs branches sans nécessiter de construction de graphe coûteuse.

2. Auto-attention locale basée sur des patches :
   Pour réduire le coût quadratique de l'auto-attention, PHAT-JeT partitionne les $N$ particules en $N/P$ patches non chevauchants de taille $P$. Au sein de chaque patch, une auto-attention multi-têtes standard est calculée exactement. Cela restreint les interactions par paires à des groupes locaux, réduisant la complexité de $O(N^2)$ à $O(N \cdot P)$. Contrairement à d'autres méthodes de patching qui reposent sur une sérialisation spatiale ou des grilles fixes, PHAT-JeT traite les patches comme une abstraction computationnelle ; les résultats empiriques montrent que la performance est robuste à l'ordre spécifique des particules (par exemple, $p_T$, $k_T$ ou aléatoire) tant que les ordonnancements d'entraînement et de test sont cohérents.

3. Attention hiérarchique au niveau des patches :
   Pour restaurer le contexte global perdu en restreignant l'attention aux patches locaux, le modèle emploie une étape de communication hiérarchique. Chaque patch est regroupé (via un pooling moyen) en un seul « token de patch » représentatif. Un mécanisme d'auto-attention global léger est ensuite appliqué à la séquence de ces tokens de patch. Le contexte global résultant est diffusé vers les particules individuelles au sein de chaque patch. Puisque le nombre de patches ($N/P$) est beaucoup plus petit que $N$, cette étape globale opère avec un coût négligeable par rapport à l'étape locale, préservant une mise à l'échelle globale quasi linéaire.

Contributions clés
L'article fait quatre contributions principales :

1. Conception d'architecture : L'introduction de PHAT-JeT, qui conserve les interactions par paires exactes au sein des patches locaux sous des contraintes de ressources strictes, contrairement aux transformers efficaces qui approximent l'attention via des projections de faible rang ou du clustering.
2. Biais inductif géométrique : Le développement du module GMP, qui améliore les performances sur plusieurs architectures basées sur l'attention en encodant explicitement la structure locale du plan du détecteur.
3. Compromis efficacité-expressivité : Démonstration que l'attention hiérarchique basée sur des patches préserve les interactions fines entre particules à un coût quasi linéaire tout en restant robuste aux ordonnancements de tri des particules (à condition que l'entraînement et le test soient cohérents).
4. Validation complète : Évaluation extensive sur quatre benchmarks (HLS4ML, JetClass, Top Tagging et Quark–Gluon) et études d'ablation confirmant la nécessité à la fois de l'étape globale de token de patch et du module GMP.

Résultats
PHAT-JeT a été évalué sur quatre benchmarks standards d'étiquetage de jets par rapport à des références contraintes par les ressources (incluant JEDI-Linear, Linformer, SAL-T et Point Transformer V3) et à des références sans contrainte (ParT, LorentzNet).

• Benchmark HLS4ML : PHAT-JeT a atteint la précision la plus élevée (81,80 %), la meilleure aire sous la courbe ROC (0,962) et le rejet de fond moyen le plus élevé (71,6) parmi tous les modèles contraints par les ressources avec des FLOPs similaires (~1,3 M). Il surpasse nettement la référence déployable la plus performante précédente, JEDI-Linear.
• Benchmark JetClass : Sur un problème plus difficile à 10 classes, PHAT-JeT a atteint 65,38 % de précision et 43,94 de rejet de fond, surpassant substantiellement les autres modèles dans la même plage de calcul.
• Top Tagging et Quark–Gluon : PHAT-JeT est resté le modèle le plus performant dans le régime contraint par les ressources, atteignant 92,69 % de précision sur Top Tagging et 81,80 % sur Quark–Gluon.
• Études d'ablation : La suppression de l'étape globale de token de patch a réduit le rejet de fond de 1 à 3 points, et la suppression du GMP l'a réduit d'environ 5 points, confirmant la valeur complémentaire des deux composants. Le modèle a montré une robustesse à l'ordre des particules (par exemple, $k_T$ contre aléatoire) tant que l'ordre était cohérent entre l'entraînement et le test.

Importance et affirmations
L'article affirme que PHAT-JeT réduit l'écart entre les modèles réalisables pour le déclenchement et les étiqueteurs de haute précision sans contrainte. En combinant une attention locale exacte, une communication globale légère et un passage de messages géométrique, l'architecture atteint des performances de pointe parmi les modèles contraints par les ressources sans dépendre du simple nombre de paramètres ou de l'échelle des réseaux à usage général. Les auteurs soulignent que les a priori architecturaux explicites (comme le GMP) sont particulièrement précieux dans les régimes de faible capacité où les modèles ne peuvent pas compter sur l'échelle pour compenser les limitations architecturales. Le travail est présenté comme une étape vers la synthèse matérielle, notant que bien que le modèle soit compatible avec les FPGA et conçu pour le budget de déclenchement, le déploiement réel de bout en bout sur FPGA est laissé pour un travail futur. Les résultats suggèrent que l'attention basée sur des patches sert de factorisation efficace du mécanisme d'attention qui ne dépend pas d'un ordre spécifique motivé par la physique, à condition que la cohérence soit maintenue.