TrackFormers Part 2: Enhanced Transformer-Based Models for High-Energy Physics Track Reconstruction

Ce papier présente une version améliorée des modèles TrackFormers pour la reconstruction de traces en physique des hautes énergies, intégrant de nouveaux mécanismes d'attention, une conception combinant projection géométrique et clustering léger, ainsi qu'un conditionnement conjoint de la classification sur les prédictions de régression, afin d'augmenter la précision et l'efficacité face à l'augmentation du volume de données du HL-LHC.

L'essentiel

Imaginez que vous êtes dans une gigantesque discothèque (le Grand collisionneur de hadrons, ou LHC). À l'intérieur, des milliards de personnes (les particules) se bousculent, dansent et se croisent à une vitesse folle.

Le problème ? Il y a tellement de monde que c'est le chaos total. Les chercheurs veulent savoir : « Qui a dansé avec qui ? » (c'est-à-dire, quelles particules proviennent de la même collision). C'est ce qu'on appelle la reconstruction de trajectoires.

Jusqu'à présent, les méthodes traditionnelles pour trier cette foule étaient comme essayer de compter les gens un par un avec un stylo et un papier : c'est précis, mais trop lent pour la nouvelle version de la discothèque (le HL-LHC), où la foule sera encore plus dense.

Voici comment les auteurs de ce papier, avec leur nouveau modèle appelé TrackFormers, ont trouvé une solution plus intelligente et plus rapide.

1. Le nouveau terrain de jeu (Les Données)

Avant de pouvoir entraîner leur intelligence artificielle, ils ont dû créer un simulateur de foule ultra-réaliste.

• L'analogie : Imaginez qu'ils ont créé un jeu vidéo où ils simulent des collisions de particules (comme des voitures qui se percutent) avec différents niveaux de foule (de 0 à 200 personnes par mètre carré).
• Pourquoi ? Pour que leur IA apprenne à trier la foule dans des conditions extrêmes, pas seulement quand il y a peu de monde. Ils ont rendu ce jeu vidéo entièrement reproductible pour que d'autres chercheurs puissent l'utiliser.

2. La méthode du "Plan 2D" (Projection et Regroupement)

Le plus gros problème des modèles précédents était qu'ils essayaient de comparer chaque personne de la foule avec toutes les autres personnes. Si vous avez 100 000 personnes, faire toutes ces comparaisons prendrait une éternité (c'est ce qu'on appelle la "complexité quadratique").

La solution des auteurs :

• L'analogie : Au lieu de regarder la foule en 3D (dans tout l'espace), ils projettent tout le monde sur un tapis roulant plat ou sur un mur cylindrique.
  • Imaginez que vous prenez une photo de la foule vue de dessus, puis une vue de côté. Sur ces surfaces plates, les gens qui marchent ensemble (les particules d'une même trajectoire) se retrouvent très proches les uns des autres.
• Le regroupement (Clustering) : Une fois sur ce "tapis", ils utilisent un algorithme rapide pour dire : « Hé, ces 5 personnes sont très proches, elles doivent être un groupe ! ».
• Le résultat : Au lieu de comparer 100 000 personnes entre elles, l'IA ne compare que les gens qui sont dans le même petit groupe. Cela réduit le travail de calcul de 400 fois ! C'est comme passer de l'envoi de 100 000 courriers individuels à l'envoi de 250 colis groupés.

3. Le "Cerveau" qui fait deux métiers à la fois (Régression et Classification)

Dans leur nouveau modèle, l'IA ne se contente pas de dire « Ces gens sont ensemble ». Elle fait deux choses en même temps, comme un détective qui observe et déduit :

1. Le Prévisionniste (Régression) : Il regarde une particule et devine immédiatement sa vitesse et sa direction (comme deviner où va une voiture juste en la regardant).
2. Le Trieur (Classification) : Il prend cette information de vitesse et l'utilise pour mieux décider si deux particules appartiennent au même groupe.

L'analogie : Imaginez un chef d'orchestre. Au lieu de juste écouter les musiciens (classification), il leur dit aussi : « Toi, tu dois jouer un peu plus fort » (régression). En utilisant cette information, il peut mieux coordonner l'ensemble du groupe.

• Le gain : En combinant ces deux tâches, leur modèle est devenu plus précis (environ 2 % de mieux), ce qui est énorme dans le monde de la physique.

4. La technologie magique (FlexAttention)

Pour que tout cela fonctionne sur un seul ordinateur puissant sans exploser la mémoire, ils ont utilisé une nouvelle technologie appelée FlexAttention.

• L'analogie : C'est comme passer d'un système de file d'attente rigide où tout le monde doit attendre son tour, à un système où les gens peuvent s'asseoir dans des salles d'attente de tailles différentes selon leurs besoins, sans bloquer les autres. Cela permet d'entraîner des modèles beaucoup plus grands et complexes sur le même matériel.

En résumé : Pourquoi c'est important ?

Ce papier montre comment transformer un problème impossible (trier des milliards de particules en temps réel) en une tâche gérable grâce à trois astuces :

1. Simplifier la vue (projeter la foule sur des murs plats).
2. Grouper intelligemment (ne comparer que les voisins proches).
3. Faire travailler l'IA en duo (prédire la vitesse ET trier en même temps).

Le résultat final ? Leur système peut reconstruire les trajectoires en quelques millisecondes (environ 100 ms), ce qui est 5 à 10 fois plus rapide que les anciennes méthodes, tout en étant plus précis. C'est une étape cruciale pour que les futurs accélérateurs de particules puissent fonctionner sans être étouffés par la quantité de données qu'ils produisent.

C'est comme passer d'un tri manuel de lettres à un système de tri postal automatisé ultra-rapide : tout devient fluide, rapide et précis.

Résumé technique

Résumé Technique : TrackFormers Part 2

1. Problématique

Les expériences de physique des hautes énergies, en particulier avec la montée en puissance du Grand collisionneur de hadrons à haute luminosité (HL-LHC), génèrent des volumes de données sans précédent. Cela pose un défi majeur pour la reconstruction des traces de particules, où des centaines de milliers de « hits » (détections) doivent être associés avec précision à leurs particules d'origine.

• Limites des méthodes traditionnelles : Bien que précises, les méthodes classiques peinent à s'adapter à ces débits de données massifs.
• Limites des approches ML précédentes : Les modèles basés sur les Transformers (introduits dans le travail précédent « TrackFormers ») souffrent d'une complexité quadratique par rapport au nombre de hits, les rendant impraticables pour des événements complets du HL-LHC sans optimisation.

2. Méthodologie

Les auteurs proposent une série d'améliorations architecturales et de nouvelles données pour surmonter ces obstacles.

A. Nouvelles Données (Niveau Hit)

• Création d'un jeu de données entièrement reproductible basé sur ACTS (A Common Tracking Software).
• Processus simulés : $pp \to t\bar{t}H$ (avec $H \to b\bar{b}$) et $pp \to t\bar{t}$ inclusif, générés via Pythia8.
• Simulation : Utilisation de la simulation rapide « Fatras » et numérisation des hits pour imiter un détecteur TrackML réaliste.
• Conditions : Jeux de données générés pour différents niveaux de pile-up (0, 5, 20, 50 et 200), chacun contenant 40 000 événements.

B. Conception Améliorée du Modèle
L'approche repose sur trois piliers techniques :

1. Projection Géométrique et Masquage (Masking) :

   • Pour éviter la complexité quadratique de l'attention, les hits sont projetés sur des surfaces simplifiées : un cylindre pour le barrel (coordonnées $R-\phi, z$) et deux plans pour les endcaps (coordonnées $x, y$).
   • Clustering léger : Un algorithme de fenêtrage itératif ou DBSCAN est appliqué sur ces surfaces projetées pour former des voisinages locaux.
   • FlexAttention : Les clusters définissent des masques de blocs (block masks) pour FlexAttention. Cela permet de n'associer que les paires de hits physiquement plausibles, réduisant la matrice d'attention effective d'un facteur jusqu'à 400x.
   • Note : Pour les endcaps, une réprojection itérative sur les positions candidates du sommet (vertex-z) affine l'alignement.

2. Apprentissage Contrastif (InfoNCE) :

   • Le modèle encode chaque hit dans un vecteur de 32 dimensions.
   • Il est entraîné avec une perte InfoNCE multi-positifs : les hits d'une même trace sont des positifs, les autres des négatifs.
   • À l'inférence, une matrice de similarité cosinus $N \times N$ est générée pour assembler les traces sans étape de clustering séparée.

3. Modèle Joint : Régression + Classification (JM)

   • Étape 1 (Régression) : Un encodeur (EncReg) prédit les paramètres de la trace ($\theta, \sin\phi, \cos\phi, q$) et quatre variables latentes.
   • Étape 2 (Classification) : Un second encodeur (EncCla) prend en entrée les coordonnées brutes concaténées avec les sorties de la régression. Il classe les hits dans des bins quantiles de $(\phi, \theta, p, q)$.
   • Entraînement Joint : Une perte combinée $L = \alpha L_{reg} + \beta L_{cla}$ est utilisée. Cela permet une inférence en un seul passage (one-shot) produisant à la fois les paramètres de la trace et la classification.

C. Optimisation Logicielle (FlexAttention)

• Remplacement de FlashAttention-2 par FlexAttention.
• Avantage : FlexAttention gère nativement des longueurs de séquence hétérogènes via des masques de blocs pré-calculés, éliminant la nécessité de padding manuel et permettant un entraînement par lots (batch) standard, contrairement à FlashAttention-2 qui limitait le batch size à 1.

3. Résultats Clés

• Performance de Latence :

  • Le temps d'inférence global est de l'ordre de quelques dizaines de millisecondes par événement (ex: ~47 ms pour l'assignation trace-hit).
  • C'est nettement plus rapide que les pipelines GNN existants (0,5–1 s) et comparable à l'état de l'art (~100 ms).
  • La latence CPU est stable à 0,1 ms, tandis que la latence GPU augmente linéairement avec la profondeur du modèle (~2,4 ms par couche).

• Performance Physique (Efficacité et Score) :

  • Efficacité : Le modèle atteint ~90 % d'efficacité (double-majority) dans le barrel et 91 % dans les endcaps après raffinement du vertex-z.
  • Comparaison avec l'état précédent :
    • Les modèles EncCla (classification seule) voient leur score TrackML passer de 78 % à 89 % avec l'augmentation de la profondeur (jusqu'à 15 couches).
    • L'ajout du module de régression (Modèle Joint JM) apporte un gain supplémentaire d'environ 2,4 % en précision et 2 % en score TrackML.
  • Le modèle JM 9:15 (9 couches régresseur / 15 couches classifieur) atteint un score TrackML de 91,4 %.

• Efficacité Mémoire :

  • Grâce à FlexAttention, il est possible d'entraîner conjointement le régresseur et le classifieur sur une seule carte NVIDIA A100 (40 Go), ce qui était impossible avec FlashAttention-2.

4. Contributions Principales

1. Jeu de données reproductible : Publication d'un pipeline ACTS complet générant des données de niveau « hit » pour divers processus physiques et niveaux de pile-up (0 à 200).
2. Architecture Scalable : Démonstration qu'un clustering basé sur la projection géométrique couplé à FlexAttention permet de réduire la complexité de l'attention de 400x, rendant les Transformers applicables aux densités de hits du HL-LHC.
3. Modèle Unifié (JM) : Introduction d'une architecture end-to-end combinant régression et classification en un seul passage, améliorant la précision sans sacrifier la simplicité de déploiement.
4. Optimisation Logicielle : Adoption de FlexAttention pour lever les goulots d'étranglement de mémoire et de batch processing des modèles de vision par ordinateur appliqués à la physique.

5. Signification et Perspectives

Ce travail marque une étape cruciale vers le déploiement de l'IA dans les expériences de physique des hautes énergies de nouvelle génération. En prouvant que les modèles basés sur les Transformers peuvent gérer les densités de données extrêmes du HL-LHC tout en maintenant une latence faible et une haute précision, les auteurs offrent une solution robuste pour remplacer ou compléter les méthodes de reconstruction traditionnelles. La capacité à entraîner des modèles profonds et complexes sur du matériel standard (A100) ouvre la voie à des architectures encore plus sophistiquées pour l'analyse future des collisions.