Machine Learning for neutron source distributions

Cet article propose et évalue une approche novatrice utilisant des modèles génératifs probabilistes pour estimer les distributions de sources de neutrons à partir de listes de particules Monte Carlo, permettant un échantillonnage efficace et indépendant de la mémoire une fois les modèles entraînés.

L'essentiel

Imaginez que vous essayez de recréer la recette parfaite d'un gâteau complexe et à plusieurs étages. Dans le monde de la science des neutrons, ce « gâteau » est un flux de neutrons (de minuscules particules) éjectés depuis une source, chacun ayant sa propre vitesse, sa propre direction, son propre énergie et son propre timing.

Traditionnellement, les scientifiques ont tenté de recréer ce flux de deux manières :

1. La méthode « Copier-Coller » : Ils exécutent une simulation informatique massive et lente pour générer une liste géante de chaque neutron individuel. Ils enregistrent cette liste (appelée fichier MCPL) et tentent de l'utiliser encore et encore. Le problème ? Si vous avez besoin de plus de neutrons que ce que la liste contient, vous copiez et collez les mêmes encore et encore. Cela crée des « bugs » ou des « points chauds » dans la simulation, comme si l'on voyait le même motif de miettes se répéter indéfiniment.
2. La méthode « Règle empirique » : Ils tentent de deviner la recette en examinant les ingrédients séparément (par exemple : « combien sont rapides ? », « combien sont lents ? »). Le problème ? Cela ignore la façon dont les ingrédients se mélangent. En réalité, un neutron rapide pourrait toujours se déplacer dans une direction spécifique, mais cette méthode les traite comme s'ils étaient sans rapport, perdant ainsi la « saveur » des données réelles.

La nouvelle approche : Le « Chef IA »
Ce papier présente une nouvelle façon de résoudre ce problème en utilisant l'apprentissage automatique. Au lieu de copier la liste ou de deviner les règles, les auteurs ont entraîné quatre types différents de « Chefs IA » (modèles génératifs) pour apprendre l'essence de la recette des neutrons.

Voici comment le papier décompose cela :

1. La phase d'entraînement (Apprendre la recette)

Les chefs IA sont nourris avec un échantillon de la simulation informatique originale et lente (les « données d'entraînement »). Ils ne se contentent pas de mémoriser la liste ; ils apprennent les relations complexes entre toutes les variables.

• L'analogie : Imaginez montrer à un chef mille photos d'un type spécifique de nuage. Il ne se contente pas de mémoriser les photos ; il apprend ce qui fait qu'un nuage ressemble à ce nuage – la façon dont les bords s'enroulent, la densité, et comment la lumière le frappe. Une fois qu'il a appris cela, il peut peindre un nouveau nuage qui n'a jamais existé auparavant mais qui semble parfaitement juste.

2. Les quatre Chefs IA

Les auteurs ont testé quatre types différents de modèles IA pour voir lequel apprenait le mieux la recette :

• Flux de normalisation (NF) : Imaginez cela comme un chef capable d'étirer et de comprimer parfaitement un morceau de pâte. Ils commencent avec une boule de pâte simple et uniforme (bruit aléatoire) et l'étirent pour former la forme complexe exacte du nuage de neutrons. Le papier a constaté que c'était le meilleur chef, créant les « nouveaux » neutrons les plus précis qui correspondaient parfaitement aux données originales.
• Autoencodeurs variationnels (VAE) : Ce chef tente de compresser la recette en un résumé puis de la reconstruire. C'est rapide et efficace pour les formes complexes, mais parfois le gâteau reconstruit ressort un peu « flou » ou moins net que l'original.
• Réseaux antagonistes génératifs (GAN) : C'est une « lutte de traction » entre deux chefs. L'un tente de cuire un gâteau factice, et l'autre tente de repérer le faux. Ils continuent de rivaliser jusqu'à ce que le gâteau factice soit indiscernable du vrai. Ce papier les a trouvés un peu difficiles à entraîner et sujets à la « triche » (répétant les mêmes quelques motifs).
• Modèles de diffusion (DM) : Ce chef commence avec un gâteau bruyant et désordonné et le nettoie lentement, étape par étape, jusqu'à ce qu'il soit parfait. Cela fonctionne bien mais est très lent et coûteux en calcul, comme essayer de nettoyer une pièce en ramassant un grain de poussière à la fois.

3. Les résultats : Pourquoi cela compte

Le papier a testé ces chefs IA sur deux scénarios réels :

• Scénario A (Le jeu de données TDR) : Une source de neutrons complexe et à haute énergie. Les chefs IA ont appris la recette si bien qu'ils pouvaient générer des millions de nouveaux neutrons qui semblaient statistiquement identiques à la simulation originale, mais sans les bugs de « copier-coller ».
• Scénario B (Le jeu de données de référence) : Une expérience réelle où ils ont comparé les neutrons générés par l'IA aux mesures réelles prises dans un laboratoire. L'IA (spécifiquement le Flux de normalisation) correspondait aux données réelles presque parfaitement.

L'avantage clé :
Une fois que le chef IA a appris la recette, la liste géante et lourde des neutrons originaux n'est plus nécessaire. Le modèle IA est minuscule (comme quelques kilo-octets) et peut générer instantanément un nombre illimité de nouveaux neutrons statistiquement parfaits. Cela économise d'énormes quantités de temps et de mémoire informatiques.

Ce que le papier ne dit pas

Les auteurs prennent soin de préciser que ces modèles sont basés sur les données. Ils apprennent strictement à partir des données qui leur sont fournies.

• Si la simulation originale manquait d'un certain type de neutron, l'IA ne l'inventera pas (sauf si le modèle est spécifiquement ajusté pour deviner au-delà des données, ce que le papier note comme une fonctionnalité spécifique d'autres méthodes, et non l'objectif principal ici).
• Le papier ne prétend pas que ces modèles peuvent prédire une nouvelle physique ou corriger de mauvaises données ; ce sont des outils pour recréer efficacement des modèles de données existants afin de les utiliser dans la conception d'instruments à neutrons.

En résumé :
Le papier démontre que nous pouvons remplacer les listes lourdes et sujettes aux bugs de données de neutrons par de minuscules modèles IA intelligents. Ces modèles apprennent l'« ADN » du flux de neutrons et peuvent générer des neutrons frais et réalistes à la demande, rendant la conception des futures expériences de neutrons plus rapide, moins chère et plus précise. Parmi les quatre modèles testés, le Flux de normalisation était le grand gagnant.

Résumé technique

Résumé technique : Apprentissage automatique pour les distributions de sources de neutrons

Énoncé du problème
En optique neutronique et en conception d'instruments, la caractérisation précise des distributions d'espace de phase des neutrons (position, direction, temps, énergie, poids et polarisation) à des surfaces spécifiques est cruciale. Traditionnellement, ces distributions sont dérivées de simulations de Monte Carlo (MC) (par exemple, MCNP, PHITS, OpenMC) et stockées sous forme de fichiers Monte Carlo Particle List (MCPL). Bien que les fichiers MCPL servent d'entrées standard pour les logiciels de traçage de rayons (par exemple, Vitess, McStas), leur réutilisation directe présente des limitations significatives :

1. Artéfacts statistiques : Le sous-échantillonnage et le bruit statistique dans le fichier MCPL d'origine se propagent dans les simulations en aval, créant des structures artificielles (par exemple, des « points chauds ») et une convergence médiocre.
2. Perte de corrélation : Les approches alternatives qui ajustent des modèles analytiques ou semi-analytiques à des variables individuelles échouent souvent à préserver les corrélations complexes et de haute dimension entre les variables de l'espace de phase, essentielles pour un transport réaliste, en particulier dans les sources de haute brillance.
3. Coût computationnel : Le réexamen de simulations MC à haute statistique pour générer de nouvelles listes de sources pour différentes configurations d'instruments est chronophage et coûteux en ressources informatiques.

Les solutions d'apprentissage automatique existantes, telles que l'estimation de densité par noyau (KDSource), répondent à certains problèmes mais nécessitent souvent un réglage manuel des poids d'importance et peinent à estimer l'espace de phase multivarié complet, en particulier lorsque les poids des particules varient.

Méthodologie
Les auteurs proposent un cadre utilisant des modèles génératifs probabilistes (PGM) pour apprendre la distribution de probabilité conjointe des variables de l'espace de phase des neutrons directement à partir de fichiers MCPL. Une fois entraînés, ces modèles peuvent générer de nouveaux échantillons de neutrons statistiquement indépendants qui préservent les corrélations de la distribution originale sans nécessiter la liste de particules d'origine.

• Ensembles de données : Deux ensembles de données distincts ont été utilisés pour évaluer les modèles :

  1. Ensemble de données TDR : Une simulation PHITS du modérateur d'hydrogène para de la source de haute brillance (HBS) (676 558 neutrons). Cet ensemble de données inclut les poids des particules et couvre un large spectre d'énergies (froid, thermique, rapide), présentant un défi de modélisation complexe et de haute dimension (7 variables).
  2. Ensemble de données de référence : Une simulation PHITS de la plateforme neutronique JULIC (118 219 neutrons), mise en regard avec des mesures expérimentales de 2023. Cet ensemble de données a des poids de particules constants (la dimensionnalité est réduite à 6 variables) et permet une validation directe par rapport aux spectres expérimentaux.
  3. Prétraitement : Toutes les variables ont été normalisées à l'aide d'une mise à l'échelle MinMax. Des transformations logarithmiques ont été appliquées aux variables de temps et d'énergie pour gérer leur large plage dynamique.

• Modèles génératifs : Quatre architectures PGM distinctes ont été comparées :

  1. Flux de normalisation (NF) : Plus précisément, les flux de couplage et les flux autorégressifs masqués (MAF). Ces modèles apprennent des transformations inversibles entre une distribution de référence simple (Gaussienne) et la distribution de données complexe, permettant une évaluation exacte de la vraisemblance et l'échantillonnage.
  2. Autoencodeurs variationnels (VAE) : Une architecture encodeur-décodeur apprenant une représentation latente des données, optimisée via une perte de reconstruction et un régularisateur de divergence de Kullback–Leibler (KL).
  3. Réseaux antagonistes génératifs (GAN) : Un réseau générateur entraîné contre un réseau discriminateur pour produire des échantillons indiscernables des données réelles.
  4. Modèles de diffusion (DM) : Modèles apprenant à inverser un processus de bruitage avant pour reconstruire des données à partir d'un bruit aléatoire.

• Implémentation : Les modèles ont été entraînés à l'aide de PyTorch sur le cluster booster JUWELS. Un module C++ personnalisé, source AI, a été développé pour le logiciel MC Vitess afin de charger les modèles PyTorch entraînés (via TorchScript) et d'effectuer un échantillonnage en temps réel, remplaçant efficacement l'entrée de fichier MCPL.

• Métriques d'évaluation :

  • Discrépance moyenne maximale (MMD) : Utilisée pour quantifier la distance entre la distribution apprise et la distribution MCPL originale dans un espace de Hilbert à noyau reproduisant (RKHS).
  • Vitesse d'échantillonnage : Temps nécessaire pour générer 20 000 neutrons.
  • Validation physique : Comparaison des spectres de longueur d'onde simulés avec les données expérimentales de la plateforme JULIC.

Résultats clés

• Fidélité de la distribution : Tous les modèles ont appris avec succès les distributions sous-jacentes, mais les flux de normalisation (NF) ont démontré des performances supérieures. Les NF ont obtenu les scores MMD les plus bas sur les deux ensembles de données (par exemple, $1,19 \times 10^{-4}$ pour TDR et $1,20 \times 10^{-4}$ pour la référence), surpassant les VAE, les GAN, les modèles de diffusion et la référence KDSource.
• Préservation des corrélations : L'analyse visuelle des histogrammes 2D a confirmé que les NF ont préservé avec succès les corrélations complexes entre les variables (par exemple, temps-énergie, angle-position) qui étaient perdues dans les méthodes d'ajustement plus simples.
• Efficacité d'échantillonnage : Bien que KDSource offre les temps d'échantillonnage les plus rapides, les NF offrent un équilibre favorable entre vitesse et précision. Les VAE et les GAN étaient également rapides, tandis que les modèles de diffusion étaient significativement plus lents à échantillonner en raison de leur nature itérative de débruitage.
• Validation expérimentale : Dans l'ensemble de données de référence, les sources générées par les NF ont reproduit le spectre de longueur d'onde expérimental avec une haute fidélité, correspondant aux performances de KDSource. Cependant, contrairement à KDSource, le modèle NF adhère strictement aux limites des données d'entraînement (en raison de la couche d'activation Sigmoid), empêchant la génération d'extrapolations non physiques (par exemple, des neutrons en dehors du guide ou avec des longueurs d'onde non observées) sauf s'il est explicitement conçu pour le faire.
• Robustesse : Les mêmes architectures de modèles ont nécessité un réglage minimal lorsqu'elles ont été appliquées aux deux ensembles de données différents, démontrant une robustesse face aux variations de dimensionnalité et de complexité de distribution.

Signification et revendications
L'article prétend présenter la première étude comparative intégrant des modèles génératifs modernes d'apprentissage automatique dans les pipelines de Monte Carlo pour neutrons et les validant par rapport à des données expérimentales. La signification principale réside dans la démonstration que :

1. Les PGM sont des substituts viables : Les modèles génératifs probabilistes peuvent remplacer les fichiers MCPL bruts, offrant une méthode efficace en mémoire (modèles de la taille de kilooctets contre fichiers de gigaoctets) et statistiquement robuste pour la représentation des sources.
2. Les flux de normalisation sont optimaux : Parmi les architectures testées, les NF offrent le meilleur compromis pour cette application spécifique, fournissant un apprentissage de distribution haute fidélité avec la capacité d'intégrer des contraintes physiques strictes (par exemple, un espace de phase borné) directement dans l'architecture.
3. Intégration du flux de travail : Le module source AI développé dans Vitess prouve que ces modèles peuvent être intégrés de manière transparente dans les flux de travail de simulation neutronique existants, permettant un échantillonnage rapide et indépendant sans nécessiter les données d'entraînement originales.

Les auteurs concluent que, bien que les modèles génératifs ne remplacent pas le besoin de simulations MC initiales haute fidélité, ils fournissent un outil efficace, flexible et physiquement significatif pour la conception et l'optimisation d'instruments en aval, en particulier dans les scénarios nécessitant des simulations répétées ou un échantillonnage à haute statistique où la réutilisation directe de MCPL est impraticable.