Economical Jet Taggers -- Equivariant, Slim, and Quantized

Ce document présente une version svelte, quantifiée et à paramètres réduits du tagueur de jets L-GATr qui parvient à une réduction de l'énergie de coût d'un ordre de grandeur avec seulement une diminution modérée des performances, ouvrant la voie à un marquage de jets efficace au niveau du déclencheur au LHC.

L'essentiel

Imaginez le Grand Collisionneur de Hadrons (LHC) comme une immense usine à particules à haute vitesse. Chaque seconde, il fait s'entrechoquer des protons, créant un jet chaotique de débris. Les physiciens doivent trier ces débris pour trouver des particules spécifiques et rares (comme le « quark top ») cachées parmi des milliards d'autres ordinaires. Ce processus de tri est appelé étiquetage de jets (jet tagging).

Pendant des années, les scientifiques ont utilisé des programmes informatiques complexes (Apprentissage Automatique) pour effectuer ce tri. Les champions actuels sont les « Transformers » — des modèles d'IA puissants qui sont incroyablement précis mais aussi énormes, lents et gourmands en énergie. Ils sont comme une flotte de camions massifs et gourmands en carburant essayant de livrer une seule lettre ; ils font le travail, mais ils sont trop gros et trop coûteux pour être utilisés au moment même où les données sont collectées (au niveau du « déclencheur » ou trigger).

Cet article pose une question simple : Pouvons-nous transformer ces camions géants en petits scooters économes en carburant sans perdre la capacité de livrer la lettre ?

Voici comment les auteurs y sont parvenus, en utilisant trois stratégies principales :

1. La version « Élaguée » (L-GATr-slim)

Le modèle original « L-GATr » est comme un couteau suisse qui transporte tous les outils possibles : scalaires, vecteurs, tenseurs, et plus encore. Cependant, les auteurs ont réalisé que pour la plupart des tâches de physique des particules, vous n'avez réellement besoin que de deux outils : des scalaires (des nombres) et des vecteurs (des flèches avec une direction).

• L'analogie : Imaginez un chef qui insiste pour utiliser une cuisine industrielle complète avec des fours, des blenders et des mélangeurs juste pour préparer un simple sandwich. Les auteurs ont dit : « Utilisons simplement un couteau et une planche à découper. »
• Le résultat : Ils ont construit une version « Slim » (élaguée) de l'IA qui supprime les outils inutiles. Elle est aussi performante que la version géante, mais beaucoup plus rapide à entraîner et utilise moins de mémoire. C'est comme passer d'un camion lourd à une voiture de sport agile qui accomplit le même travail.

2. La version « Minuscule » (Ultra-mini Taggers)

Les auteurs se sont ensuite demandé : « Jusqu'à quelle petite taille pouvons-nous descendre ? » Ils ont tenté de réduire ces modèles d'IA à la taille d'une petite voiture miniature (environ 1 000 paramètres, contre des millions dans l'original).

• L'analogie : Pensez à essayer de faire tenir toute la connaissance d'une bibliothèque dans une simple carte postale. Habituellement, on perd l'histoire. Mais les auteurs ont découvert que si l'on organise l'information correctement (en utilisant des règles spécifiques « Lorentz-équivariantes » qui respectent les lois de la physique), on peut faire tenir l'essentiel de la connaissance dans un espace minuscule.
• Le résultat : Ils ont constaté que pour les modèles très petits, l'architecture « LLoCa » fonctionne mieux si l'on réduit le nombre de couches, tandis que le « L-GATr-slim » fonctionne mieux si l'on réduit la largeur des couches. Même à cette taille microscopique, ils surpassent toujours les anciens modèles d'IA qui ne sont pas conscients de la physique.

3. La version « Quantifiée » (Mathématiques à faible précision)

C'est l'économie d'énergie la plus spectaculaire. L'IA standard utilise des mathématiques très précises (comme mesurer une distance au milliardième de millimètre près). Les auteurs ont réalisé que pour l'étiquetage de jets, vous n'avez pas besoin d'une telle précision. Vous pouvez vous contenter d'arrondir les nombres de manière significative.

• L'analogie : Imaginez que vous comptez des pommes dans un entrepôt.
  • IA Standard : Vous pesez chaque pomme au microgramme près. (Précis, mais cela prend un temps infini et consomme beaucoup d'énergie de pesée).
  • IA Quantifiée : Vous les comptez simplement en nombres entiers. (Rapide, consomme presque aucune énergie, et pour l'objectif de savoir « combien il y a de pommes », c'est parfaitement suffisant).
• La méthode : Ils ont utilisé une technique appelée PARQ (Quantification Régularisée Pièce par Pièce Affine). Voyez cela comme une règle d'arrondi intelligente qui pousse doucement les nombres vers des valeurs simples (comme 0, 1 ou -1) pendant le processus d'entraînement, plutôt que de les forcer brutalement.
• Le résultat : En passant à ces nombres plus « grossiers », ils ont réduit le coût énergétique de l'exécution de l'IA d'un facteur 10 (un ordre de grandeur). L'IA est devenue incroyablement rapide et économe en énergie, avec seulement une légère baisse de précision.

La vue d'ensemble

Les auteurs ont combiné ces trois stratégies — l'élagage de l'architecture, la miniaturisation de la taille et la quantification des mathématiques — pour créer des « Étiqueteurs de Jets Économiques ».

• Pourquoi est-ce important ? Actuellement, ces modèles d'IA puissants sont trop volumineux pour être exécutés sur le matériel qui décide en temps réel quelles collisions conserver et lesquelles rejeter (le « déclencheur » ou trigger).
• L'objectif : En rendant ces modèles petits, rapides et économes en énergie, les auteurs espèrent pouvoir les faire fonctionner directement sur le matériel du déclencheur. Cela permettrait au LHC d'utiliser l'IA pour prendre des décisions en une fraction de seconde sur les collisions de particules à conserver, permettant potentiellement de découvrir une nouvelle physique qui aurait été manquée auparavant parce que les données ont été rejetées trop rapidement.

En résumé : Ils ont pris une IA géante et gourmande en énergie, lui ont imposé un régime, l'ont rétrécie et lui ont appris à faire des mathématiques avec moins de décimales, créant ainsi un moteur minuscule et super efficace capable de reconnaître les particules les plus importantes de l'univers.

Résumé technique

Résumé technique : Jet Taggers Économiques – Équivariants, Élagués et Quantifiés

Énoncé du Problème
L'apprentissage automatique (ML) moderne a transformé le marquage de jets (jet tagging) au Grand Collisionneur de Hadrons (LHC), les transformeurs Lorentz-équivariants émergeant comme des architectures de pointe. Cependant, les modèles de premier plan comme L-GATr sont coûteux en calcul, nécessitant une mémoire et un temps d'entraînement importants. Alors que les tendances de l'industrie favorisent l'augmentation de l'échelle des réseaux et des jeux de données, la physique du LHC fait face à des contraintes spécifiques, notamment concernant les exigences de mémoire et de latence du matériel de déclenchement (trigger). Le marquage de jets actuel ne joue pas encore de rôle dans le déclenchement, mais les auteurs soutiennent qu'il le devrait. Le défi central abordé est de savoir comment réduire la taille et le coût computationnel des modernes marqueurs de jets équivariants tout en minimisant la dégradation des performances, permettant potentiellement leur déploiement au niveau du déclencheur.

Méthodologie
Le papier propose une stratégie à deux volets pour optimiser l'efficacité des ressources : l'élagage architectural et la quantification numérique.

1. Architecture L-GATr-slim :
   Les auteurs introduisent une version rationalisée du transformeur Lorentz-équivariant (L-GATr). Le L-GATr standard utilise une représentation d'algèbre géométrique impliquant des scalaires, des pseudo-scalaires, des vecteurs, des axial-vecteurs et des tenseurs antisymétriques de rang deux. Les auteurs observent que les pseudo-scalaires, les axial-vecteurs et les tenseurs sont inutiles pour la plupart des applications du LHC. Par conséquent, L-GATr-slim restreint la représentation latente uniquement aux scalaires et aux vecteurs.

• Couches Linéaires : Étendues pour opérer sur des représentations couplées de scalaires et de vecteurs, garantissant que les composantes vectorielles partagent un coefficient scalaire apprenable unique pour maintenir l'équivariance de Lorentz.
• Non-linéarité : Adapte l'Unité Linéaire Portée (GLU) en appliquant la non-linéarité au produit scalaire de deux vecteurs, multiplié par la sortie vectorielle.
• Normalisation : Modifie RMSNorm pour utiliser la valeur absolue du produit scalaire de Minkowski pour les canaux vectoriels.
• Attention : Construit des matrices d'attention scalaires à l'aide d'une formulation spécifique qui évite le produit extérieur coûteux utilisé dans le L-GATr complet.
• Implémentation : L'architecture est conçée pour être compilée avec torch.compile pour plus d'efficacité.

2. Stratégies de Quantification :
   Les auteurs appliquent des types de données de faible précision et la quantification des poids pour réduire davantage les coûts.

• Quantification des Types de Données : Les entrées des couches linéaires sont quantifiées en int8 (en utilisant la quantification du point zéro) tout en conservant le bfloat16 pour les opérations sensibles à la précision et pour la passe arrière (backward pass). Cela est appliqué aux couches cachées du Transformer, de ParT, de L-GATr-slim et du LLoCa-Transformer.
• Quantification des Poids : Les poids linéaires sont quantifiés en valeurs binaires ou ternaires en utilisant la Quantification par Gradient Proximal (PARQ). Cette méthode traite la quantification comme une contrainte de régularisation, utilisant un opérateur proximal pour mettre à jour les poids. Les auteurs comparent PARQ à l'Estimation de Passage Direct (STE), trouvant que PARQ offre une meilleure stabilité et performance.
• Préservation de l'Équivariance : Un soin particulier est apporté pour s'assurer que la quantification ne viole pas l'équivariance de Lorentz. Pour LLoCa, l'orthonormalisation et les projections de cadre restent en pleine précision (float32), limitant les opérations de faible précision aux invariants de Lorentz. Pour L-GATr-slim, les vecteurs complets sont multipliés par des poids quantifiés, ce qui n'introduit pas de violations supplémentaires de symétrie.

3. Mise à l'Échelle Ultra-Mini :
   Les auteurs étudient les performances de ces architectures jusqu'à 1 000 paramètres en réduisant le nombre de blocs ou la largeur (canaux) du réseau.

Résultats Clés
L'étude évalue les méthodes proposées sur trois tâches : le top tagging, la régression d'amplitude et la génération d'événements.

• Performance vs Efficacité (L-GATr-slim) :

  • Sur le jeu de données JetClass (marquage de jets multi-classes), L-GATr-slim égale la performance du L-GATr complet et du LLoCa-Transformer (AUC ~0,9885) mais réduit le temps d'entraînement d'un facteur six (de 166h à 27h sur un GPU H100) et la consommation de mémoire d'un facteur deux.
  • En régression d'amplitude ($Z + 4g$), L-GATr-slim atteint la même erreur quadratique moyenne (MSE) que le L-GATr complet mais nécessite 20 fois moins d'opérations d'entraînement et la moitié du temps d'entraînement.
  • En génération d'événements ($t\bar{t} + nj$), l'architecture élaguée correspond aux performances de log-vraisemblance négative des modèles complets.

• Marqueurs Ultra-Mini :

  • Lorsqu'on réduit le nombre de blocs (profondeur), le LLoCa-Transformer surpasse le L-GATr-slim à de très petites tailles (par exemple, 1 000 paramètres).
  • En maintenant le nombre de blocs fixe (10) et en réduisant les canaux (largeur), L-GATr-slim maintient un taux de rejet de fond supérieur à 1 000 avec seulement 2 canaux vectoriels et 4 canaux scalaires, surpassant les autres architectures de 1 000 paramètres.

• Gains de Quantification :

  • La quantification des entrées en int8 et des poids en valeurs ternaires réduit la consommation d'énergie d'environ un ordre de grandeur (facteur 10) avec une perte de performance marginale.
  • Le LLoCa-Transformer et le L-GATr-slim sont robustes à la quantification, maintenant des performances élevées là où les transformeurs standards pourraient se dégrader davantage.
  • Pour le scénario le plus contraint en ressources (1 bloc, espace latent de 16 dimensions, int8), le LLoCa-Transformer quantifié (canonisation globale) conserve une performance supérieure aux marqueurs pré-graphes, malgré une réduction par deux du rejet de fond par rapport à son homologue de taille complète.

Signification et Revendications
Le papier affirme que ces versions "économiques" de transformeurs équivariants représentent une voie viable vers le marquage de jets au niveau du déclencheur (trigger-level) pour le Grand Collisionneur de Hadrons de Haute Luminosité (HL-LHC). En combinant l'élagage architectural (suppression des composants inutiles de l'algèbre géométrique) et une quantification agressive (PARQ et int8), les auteurs démontrent qu'il est possible de créer des marqueurs de ~1 000 paramètres qui conservent les avantages motivés par la physique de l'équivariance de Lorentz.

Les auteurs soulignent que si l'augmentation de l'échelle est la norme de l'industrie, la physique du LHC nécessite une approche de "réduction d'échelle consciente de la physique" (physics-aware downscaling). Les résultats suggèrent que des réseaux petits, quantifiés et équivariants peuvent être déployés sur du matériel aux ressources limitées (tels que les FPGA) sans sacrifier les symétries fondamentales qui rendent ces modèles efficaces, ouvrant potentiellement de nouvelles voies pour l'analyse en temps réel de la sous-structure des jets.