Vision Transformers and Graph Neural Networks for Charged Particle Tracking in the ATLAS Muon Spectrometer

Cet article présente deux approches basées sur l'apprentissage automatique pour améliorer le suivi des particules chargées dans le spectromètre à muons d'ATLAS : l'intégration de réseaux de neurones graphiques pour rejeter les bruits de fond et accélérer la reconstruction de 15 %, et une démonstration de concept utilisant des Vision Transformers pour un suivi de muons quasi instantané (2,3 ms) avec une efficacité de 98 %, afin de répondre aux défis posés par la haute luminosité du LHC.

L'essentiel

🚀 Le Grand Défi : Trouver une aiguille dans une botte de foin... qui explose !

Imaginez que vous êtes dans un immense stade (le Grand Collisionneur de Hadrons ou LHC) où des milliards de voitures (des protons) foncent l'une vers l'autre à une vitesse folle. Quand elles se percutent, cela crée une explosion de débris. La plupart de ces débris sont inutiles (du bruit), mais quelques-uns sont des trésors : des muons, des particules qui pourraient nous révéler les secrets de l'univers.

Le problème ? Avec la prochaine version de l'accélérateur (le HL-LHC), il y aura tellement de collisions en même temps que le stade sera rempli de bruit. C'est comme essayer d'entendre un chuchotement dans une discothèque où tout le monde crie en même temps.

Les physiciens de l'expérience ATLAS ont un détecteur géant (le Spectromètre à Muons) pour attraper ces muons. Mais avec autant de bruit, les ordinateurs actuels mettent trop de temps à trier les bons signaux des faux. C'est là que cette recherche intervient.

L'auteur, Jonathan Renusch, propose deux nouvelles méthodes basées sur l'intelligence artificielle (IA) pour résoudre ce casse-tête.

---

Méthode 1 : Le Tri-Intelligent (Les Réseaux de Neurones Graphiques)

L'analogie : Imaginez que vous avez une salle de classe remplie d'élèves (les points de données). Certains sont de vrais élèves (le signal), d'autres sont des fantômes ou des erreurs (le bruit). L'ordinateur actuel essaie de vérifier chaque élève un par un, ce qui prend du temps.

La solution : L'IA utilise un Réseau de Neurones Graphiques (GNN).

• Au lieu de regarder les élèves isolément, l'IA les regroupe par "bancs" (ce qu'on appelle des "bacs à muons").
• Elle observe les connexions entre les bancs. Si un banc est entouré de fantômes, l'IA comprend immédiatement : "Ah, ce banc est du bruit, je ne vais même pas le regarder de plus près."
• Le résultat : L'IA jette 97 % du bruit avant même que le vrai travail de reconstruction ne commence.
• Le gain : Le temps de traitement passe de 255 millisecondes à 217 ms. C'est comme si un coureur de 100 mètres gagnait 38 millisecondes : ça semble peu, mais dans un monde où tout va très vite, c'est énorme ! Cela permet de traiter plus de données sans ralentir le système.

---

Méthode 2 : Le Détective Visionnaire (Les Transformers)

L'analogie : Imaginez que vous devez reconstituer le trajet de deux cyclistes dans une ville bondée, mais vous avez une photo floue prise de nuit avec des milliers de points lumineux (des voitures, des réverbères, des néons). L'ancien système essaie de relier les points un par un avec des règles mathématiques strictes. C'est lent et compliqué.

La solution : L'auteur utilise une technologie appelée Vision Transformer (ViT), la même famille d'IA qui a révolutionné la reconnaissance d'images (comme dans les voitures autonomes ou les filtres Instagram).

• Le concept : Au lieu de suivre des règles rigides, l'IA "regarde" l'ensemble de la scène d'un coup d'œil global. Elle utilise une attention intelligente (comme un détective qui sait où regarder) pour ignorer instantanément les néons inutiles et se concentrer uniquement sur les pistes des cyclistes.
• La magie : Elle transforme chaque point de données en un "mot" et lit toute la phrase pour comprendre le trajet.
• Le résultat :
  • Elle nettoie le bruit avec une précision incroyable (elle garde 99,7 % des vrais muons et jette presque tout le reste).
  • Elle reconstruit le trajet complet en 2,3 millisecondes sur une carte graphique de bureau (une machine qu'on peut acheter pour 1 500 CHF, pas un super-ordinateur de 1 million !).
  • C'est 100 fois plus rapide que la méthode actuelle pour un seul événement.

---

Pourquoi c'est important pour nous ?

1. Vitesse et Économie : Ces méthodes permettent de traiter des données beaucoup plus vite. Cela signifie que l'expérience ATLAS pourra capturer plus de phénomènes rares sans avoir besoin de construire des ordinateurs plus gros et plus chers.
2. Le Futur : Avec l'arrivée du HL-LHC après 2030, le nombre de collisions va exploser. Sans ces nouvelles techniques d'IA, les physiciens seraient submergés et manqueraient des découvertes potentielles.
3. L'Innovation : Ce papier montre que l'on peut utiliser des outils développés pour les images (comme les Transformers) pour résoudre des problèmes de physique des particules, prouvant que l'IA peut transformer la science fondamentale.

En résumé :
Les physiciens ont créé deux nouveaux "filtres" intelligents. Le premier est un tri rapide qui jette le gros du bruit avant le travail principal. Le second est un détective ultra-rapide capable de reconstruire des trajectoires complexes en une fraction de seconde, même avec un matériel informatique grand public. C'est une victoire majeure pour préparer l'avenir de la physique des particules !

Résumé technique

1. Contexte et Problématique

L'expérience ATLAS au Grand collisionneur de hadrons (LHC) fait face à des défis majeurs avec l'avènement de la phase à haute luminosité (HL-LHC) après 2030. Le nombre de collisions proton-proton par croisement de paquets (pile-up, $\langle\mu\rangle$) passera de 60 à 200. Cette densité accrue entraîne une occupation élevée dans le spectromètre à muons d'ATLAS, rendant la reconstruction des particules chargées (notamment les muons) extrêmement difficile pour le système de déclenchement en temps réel (Event Filter).

Les défis spécifiques au spectromètre à muons incluent :

• Un rapport signal/bruit très faible : Dans certaines régions (comme la nouvelle roue petite, NSW), ce rapport peut chuter à 0,6 %.
• Hétérogénéité des détecteurs : L'intégration de cinq technologies différentes (MDT, RPC, TGC, MM, sTGC) avec des précisions et des vitesses de lecture variées.
• Complexité combinatoire : La nécessité d'associer des milliers de "hits" (détections) aux trajectoires correctes dans un champ magnétique inhomogène.

L'algorithme de reconstruction actuel (baseline), exécuté sur CPU, prend environ 255 ms par événement à $\langle\mu\rangle=200$, ce qui est trop lent pour les contraintes de latence futures.

2. Méthodologie

L'article propose deux approches distinctes basées sur l'apprentissage automatique pour surmonter ces obstacles :

A. Réjection des bruits de fond par Réseaux de Neurones à Graphes (GNN)

Cette méthode vise à améliorer l'algorithme de reconstruction existant en pré-filtrant les données.

• Approche : Au lieu de traiter chaque hit individuellement (trop coûteux), les hits sont regroupés en "buckets" (clusters) selon leur proximité spatiale. Un graphe est construit dynamiquement où les nœuds sont ces buckets et les arêtes représentent les connexions spatiales entre secteurs adjacents.
• Architecture : Utilisation d'une architecture EdgeConv (inspirée de TrackML) intégrée dans le framework ATLAS Athena. Le modèle apprend à distinguer les buckets contenant des segments de muons réels de ceux contenant uniquement du bruit de fond.
• Intégration : Le modèle est déployé via ONNX (avec une implémentation personnalisée pour EdgeConv) avant l'étape de reconnaissance de motifs (Hough transform).

B. Reconstruction de bout en bout par Vision Transformers (ViT)

Cette approche est une preuve de concept pour remplacer la reconstruction traditionnelle par un modèle entièrement basé sur l'apprentissage profond.

• Architecture : Adaptation de l'architecture Mask2Former (issue de la vision par ordinateur) pour la physique des hautes énergies.
• Traitement des données : Les hits individuels sont traités comme des "tokens". Pour gérer la nature dispersée et bruyante des données, une étape de filtrage des hits (classification binaire) précède l'étape de suivi.
• Mécanisme d'attention : Utilisation de l'Attention Flash à fenêtre (Windowed Flash Attention) avec un tri préalable des hits selon l'angle azimutal ($\phi$). Cela permet de capturer les corrélations globales tout en réduisant la complexité computationnelle de $O(N^2)$ à $O(W \times N)$.
• Fonctionnement : Le modèle effectue une détection d'objets (identification des trajectoires candidates) et une régression des paramètres de la trajectoire ($p_T$, $\eta$, $\phi$, charge) via un mécanisme d'attention croisée conditionnée par des masques.

3. Résultats Clés

Résultats des GNN (Filtrage des buckets)

• Performance : À $\langle\mu\rangle=60$, le taux de rejet des buckets de bruit de fond est d'environ 97 %.
• Gain de vitesse : L'intégration du filtre GNN réduit le temps de reconstruction par événement de 255 ms à 217 ms (soit une amélioration de 15 %) sur un CPU AMD EPYC 9654 couplé à un GPU NVIDIA H100.
• Impact physique : Aucune dégradation de l'efficacité de reconstruction ou de la précision des paramètres cinématiques n'est observée.

Résultats des ViT (Reconstruction de bout en bout)

• Filtrage des hits : Le modèle atteint une aire sous la courbe ROC (AUC) de 0,9997. À un point de fonctionnement de 99 % d'efficacité, il rejette 99,7 % du bruit de fond, réduisant l'occupation moyenne de 6 900 hits à 55 hits par événement.
• Efficacité de suivi :
  • Efficacité de détection du signal : 98,0 %.
  • Efficacité de double correspondance (efficacité et pureté > 50 %) : 94,59 %.
  • Précision de classification de la charge : 96,35 %.
• Vitesse d'inférence : Sur une carte graphique grand public (NVIDIA RTX 3090), le temps de traitement est de 2,3 ms par événement (pour un lot de 200 événements). Sur un GPU H100, ce temps descend à 0,9 ms.
• Limites actuelles : La régression précise des paramètres de la trajectoire n'est pas encore aussi performante que l'algorithme de base, et la vitesse est actuellement limitée par les surcoûts de lancement des noyaux GPU (kernel launch overheads) plutôt que par le calcul pur.

4. Contributions et Signification

• Validation de nouvelles architectures : L'article démontre pour la première fois l'efficacité des GNN et des Transformers (ViT) dans le contexte spécifique et complexe du spectromètre à muons d'ATLAS, un domaine où ces méthodes étaient peu explorées.
• Optimisation pour le HL-LHC : Les deux approches répondent directement au besoin de réduire la latence et la charge computationnelle pour gérer le taux de collision élevé du HL-LHC.
• Approche hybride et modulaire : La méthode GNN offre une amélioration immédiate et compatible avec l'infrastructure actuelle, tandis que la méthode ViT propose une voie radicalement nouvelle pour une reconstruction entièrement basée sur l'IA.
• Faisabilité matérielle : La démonstration de résultats performants sur du matériel grand public (RTX 3090) suggère que ces solutions pourraient être déployées économiquement dans les systèmes de déclenchement futurs.

Conclusion

Ce travail établit que les méthodes d'apprentissage automatique, en particulier les mécanismes d'attention (Transformers) et les réseaux de graphes, sont des candidats prometteurs pour remplacer ou optimiser les algorithmes traditionnels de reconstruction de muons. Bien que des optimisations supplémentaires soient nécessaires pour la régression des paramètres et l'intégration complète dans le pipeline de déclenchement, ces résultats ouvrent la voie à des systèmes de traitement de données plus rapides, robustes et évolutifs pour la physique des hautes énergies de la prochaine décennie.