Generalizable Foundation Models for Calorimetry via Mixtures-of-Experts and Parameter Efficient Fine Tuning

Cet article présente un modèle fondamental généralisable pour la calorimétrie, basé sur des transformateurs de type « next-token » et combinant des mélanges d'experts avec un ajustement fin efficace des paramètres, permettant une adaptation modulaire et évolutive aux différents matériaux, types de particules et configurations de détecteurs sans oublier les connaissances précédentes.

L'essentiel

Imaginez que vous essayez de concevoir un nouveau type de voiture de course. Pour savoir si elle est rapide et sûre, vous devez simuler des millions de crashes virtuels. Dans le monde de la physique des particules, c'est la même chose : les scientifiques construisent des détecteurs géants pour capturer des particules invisibles, mais avant de les construire, ils doivent simuler comment ces particules vont réagir à l'intérieur.

Le problème ? Ces simulations sont comme essayer de prédire chaque goutte d'eau d'un tsunami avec une calculatrice. C'est lent, ça coûte une fortune en énergie informatique, et les ordinateurs actuels vont bientôt ne plus pouvoir suivre le rythme.

Voici comment les auteurs de cette recherche proposent de résoudre ce problème, en utilisant une idée tirée de l'intelligence artificielle moderne (comme les chatbots) :

1. Le "Chef Cuisinier" et ses "Spécialistes" (Le Modèle de Base)

Imaginez un grand restaurant. Au lieu d'avoir un seul chef qui essaie de tout cuisiner (des sushis, des pizzas, des plats indiens), vous avez un Chef de Cuisine Principal (le modèle de base) qui connaît les bases de la cuisine, et une équipe de Chefs Spécialistes (les "Experts").

• Le Chef Principal a déjà appris à cuisiner pour deux types de matériaux (le Tungstène et le Tantalum, qui sont comme des ingrédients lourds dans le détecteur). Il sait exactement comment une particule de lumière (un photon) réagit dans ces matériaux.
• L'astuce : Si vous voulez maintenant simuler un troisième ingrédient (le Plomb), vous n'avez pas besoin de réapprendre toute la cuisine à votre Chef Principal. Vous engagez simplement un nouveau Chef Spécialiste qui sait cuisiner le Plomb. Le Chef Principal reste le même, il ne perd pas ses compétences sur les autres plats. C'est ce qu'on appelle un "Mélange d'Experts".

2. Apprendre sans tout oublier (Pas d'Amnésie)

Dans les anciennes méthodes d'IA, si on apprenait à un modèle à cuisiner le Plomb, il oubliait souvent comment cuisiner le Tungstène. C'est comme si un étudiant apprenait l'espagnol et oubliait tout son français.

Ici, grâce à leur méthode, le modèle est modulaire. Quand on ajoute un nouveau spécialiste (pour le Plomb), on ne touche pas au cerveau du Chef Principal. On ajoute juste un nouveau module. Si demain on a besoin de revenir au Tungstène, on désactive simplement le module Plomb. Rien n'est perdu. C'est comme ajouter une nouvelle recette dans un livre de cuisine sans effacer les pages précédentes.

3. Changer de type de particule (Le "LoRA")

Et si on veut simuler non pas de la lumière (photons), mais des électrons ? C'est un peu différent, car les électrons se comportent différemment.

Pour cela, les chercheurs utilisent une technique appelée LoRA (Low-Rank Adaptation). Imaginez que le Chef Principal a un chapeau magique. Au lieu de réécrire tout son livre de cuisine, on lui donne juste un petit carnet de notes (le LoRA) qui dit : "Attention, pour les électrons, il faut ajouter un peu plus de sel ici et cuire plus vite là".

• Le livre de base reste intact.
• Le petit carnet ajuste le comportement pour les électrons.
• C'est rapide, léger et efficace.

4. La Vitesse : De la Tortue à la Formule 1

Avant, simuler ces particules prenait des heures sur des supercalculateurs lents (des CPU). C'était comme essayer de traverser l'océan à la rame.

Grâce à cette nouvelle architecture et à des astuces de programmation empruntées aux grands modèles de langage (comme ceux qui font fonctionner les chatbots), leur modèle est devenu une Formule 1.

• Ils utilisent des techniques comme le "cache" (garder les informations en mémoire pour ne pas avoir à les recalculer à chaque fois).
• Résultat : Ils génèrent des simulations des centaines de fois plus vite que les méthodes traditionnelles, tout en restant aussi précises.

En résumé

Cette recherche propose un modèle universel de cuisine pour la physique.

1. Il apprend une fois les bases.
2. Il peut ajouter de nouveaux ingrédients (matériaux) ou de nouveaux types de plats (particules) sans tout réapprendre.
3. Il ne perd jamais ses compétences précédentes.
4. Il est incroyablement rapide.

Cela permet aux scientifiques de tester des milliers de designs de détecteurs en quelques heures au lieu de quelques mois, accélérant ainsi la découverte de nouveaux secrets de l'univers.

Résumé technique

1. Problématique

Les expériences de physique des particules modernes font face à une demande croissante de simulations de détecteurs à haute fidélité. À mesure que les luminosités augmentent, les besoins en calcul pour les simulations Monte Carlo traditionnelles (comme Geant4) risquent de dépasser les ressources informatiques disponibles. Bien que des modèles génératifs profonds (GANs, VAEs, modèles de diffusion) aient été proposés comme substituts, ils souffrent souvent d'un manque de généralisabilité : un modèle entraîné sur un matériau ou une particule spécifique ne peut pas facilement s'adapter à de nouveaux scénarios sans un réentraînement coûteux, risquant l'oubli catastrophique (catastrophic forgetting) des connaissances précédentes.

L'objectif est de développer un modèle fondamental (Foundation Model) pour la calorimétrie capable de :

• Générer des gerbes électromagnétiques pour différents matériaux (Tungstène, Tantalum, Plomb) et différentes particules (photons, électrons) au sein d'une seule architecture.
• S'adapter de manière modulaire et additive à de nouvelles configurations sans réentraîner l'ensemble du modèle.
• Être compétitif en termes de temps d'inférence par rapport aux méthodes génératives standard.

2. Méthodologie et Architecture

Les auteurs proposent une architecture basée sur des Transformers de prédiction du prochain token (inspirés des LLM), combinant deux stratégies clés : les Mélanges d'Experts (MoE) et le Raffinement Fin Efficace en Paramètres (PEFT).

A. Architecture de Base et Tokenisation

• Représentation : Les données du détecteur (grille 3D de voxels) sont tokenisées. La dimension spatiale est discrétisée en indices de pixels (vocabulaire de 27k tokens), et l'énergie est binnée linéairement (vocabulaire de 25k tokens).
• Séquence : Les gerbes sont générées séquentiellement, triées par énergie décroissante.
• Conditionnement : L'énergie initiale de la particule est injectée via des tokens de contexte préfixés. Des embeddings spatiaux, cinématiques et des identifiants de particule guident la génération.
• Positionnement : Utilisation d'encodages de position rotatifs (RoPE) pour gérer les séquences de longueur variable, avec une précision 32 bits pour éviter les erreurs d'accumulation.

B. Adaptation Modulaire via MoE et PEFT

L'architecture repose sur un backbone Transformer pré-entraîné et figé (frozen), auquel sont ajoutés des modules adaptatifs :

1. Extension aux nouveaux matériaux (MoE) :

   • Une couche Mixture-of-Experts (MoE) est utilisée. Un routeur fixe dirige les entrées vers des experts spécialisés par matériau.
   • Pour ajouter un nouveau matériau (ex: Plomb), un seul nouvel expert est ajouté et affiné, tandis que le reste du modèle (backbone + autres experts) reste figé. Cela évite l'oubli des matériaux précédents.

2. Extension aux nouvelles particules (PEFT) :

   • Pour changer de type de particule (ex: Photons $\to$ Électrons), la structure d'interaction change fondamentalement.
   • LoRA (Low-Rank Adaptation) : Des mises à jour de faible rang sont appliquées aux projections d'attention (Q, K, V, Output) pour modifier la dynamique de génération sans toucher aux poids de base.
   • Têtes de vocabulaire modulaires : Des têtes de sortie spécifiques à chaque particule sont introduites pour gérer les changements de probabilités dans l'espace des tokens, évitant des ajustements de rang élevé coûteux sur une projection partagée.

C. Optimisation de l'Inférence

Pour rendre l'inférence autoregressive compétitive, les auteurs appliquent des techniques d'optimisation issues des LLM :

• Cache KV (Key-Value) : Stockage des projections des tokens précédents pour réduire la complexité de $O(n^2)$ à $O(n)$.
• Pré-allocation mémoire et CUDA Graphs : Élimination de la gestion dynamique de la mémoire à chaque étape pour maximiser l'utilisation du GPU.

3. Résultats Clés

• Performance de Génération : Le modèle pré-entraîné sur le Tungstène et le Tantalum reproduit avec une haute fidélité les distributions de Geant4 (énergie visible, multiplicité des impacts, centre de gravité longitudinal, profil radial). Les résultats sont comparables ou supérieurs aux modèles autoregressifs existants (comme Omnijet-$\alpha_c$) qui sont limités à un seul matériau.
• Efficacité du Raffinement Fin (Fine-Tuning) :
  • Nouveaux matériaux : L'ajout d'un expert pour le Plomb avec seulement 1 000 à 10 000 échantillons permet d'atteindre une fidélité équivalente à un entraînement complet, sans dégrader les performances sur les matériaux originaux.
  • Nouvelles particules : L'adaptation aux électrons nécessite environ 50 000 échantillons pour une haute fidélité. Une légère correction post-hoc (décalage des dépôts d'énergie) est parfois nécessaire pour corriger les biais systématiques liés à la profondeur de la gerbe.
• Vitesse d'Inférence : Grâce aux optimisations (KV-cache, CUDA graphs), le modèle atteint un temps d'inférence de ~10 ms sur un GPU A100. Cela représente un facteur d'accélération de 392x par rapport à Geant4 (CPU) et est compétitif avec des méthodes comme CaloClouds II ou L2LFlows, tout en conservant la haute fidélité des Transformers.

4. Contributions Principales

1. Modèle Fondamental Généralisable : Création d'un premier modèle capable de gérer simultanément plusieurs matériaux et particules via une architecture unifiée.
2. Stratégie d'Adaptation Additive : Démonstration qu'il est possible d'étendre les capacités d'un modèle (nouveaux matériaux ou particules) uniquement en ajoutant et en affinant des modules légers (experts, LoRA, têtes de vocabulaire), préservant ainsi l'intégrité du modèle de base.
3. Efficacité des Données : Preuve que des modèles fondés sur des Transformers peuvent être adaptés à des régimes de faible données (quelques milliers d'échantillons), crucial pour les phases de conception de détecteurs où les simulations complètes sont rares.
4. Optimisation Système : Application réussie des techniques d'optimisation d'inférence des LLM aux modèles de physique, rendant la génération autoregressive viable pour des workflows de simulation en temps réel.

5. Signification et Impact

Ce travail établit une voie prometteuse pour l'avenir de la simulation en physique des hautes énergies. En passant d'une approche de simulation Monte Carlo lourde et répétitive à un paradigme de modèles fondés sur des Transformers adaptatifs, les chercheurs peuvent :

• Réduire drastiquement les coûts de calcul et de stockage pour l'optimisation des détecteurs.
• Explorer rapidement de nombreuses configurations de matériaux et de géométries sans réentraîner des modèles entiers.
• Intégrer des connaissances physiques de manière incrémentale au fur et à mesure que de nouvelles données de simulation deviennent disponibles.

Cette approche ouvre la voie à des workflows de conception de détecteurs plus agiles, durables et évolutifs, en particulier pour les futurs collisionneurs comme le ILC (International Linear Collider).