GIGA-Lens 2.0: Strong-Lens Modeling on Multiple GPU Nodes

Résumé Technique de GIGA-Lens 2.0 : Modélisation de Lentilles Fortes sur des Nœuds Multi-GPU

Énoncé du Problème
Le lentillage gravitationnel fort constitue une sonde puissante pour la cosmologie, incluant les contraintes sur la matière noire de sous-structure, la constante de Hubble ($H_0$) et l'énergie noire. Cependant, réaliser ce potentiel nécessite une inférence bayésienne complète avec un échantillonnage de la distribution postérieure rigoureux. À mesure que les modèles de lentilles deviennent plus flexibles et que les ensembles de données s'accroissent, garantir l'exploration robuste de la distribution postérieure et la convergence numérique impose des exigences strictes en termes de performance de calcul et de scalabilité. Le cadre original GIGA-Lens (GIGA-Lens 1.0) était limité à des configurations GPU sur un seul nœud, entravant son application à de grands échantillons de systèmes, des systèmes complexes (ex: lentilles à l'échelle des amas) et des analyses conjointes impliquant la dynamique stellaire. Ces limitations forcent souvent des hypothèses simplificatrices ou restreignent l'exploration des incertitudes systématiques.

Méthodologie
GIGA-Lens 2.0 introduit une mise à jour majeure du cadre, permettant une inférence de lentille forte bayésienne entièrement basée sur la modélisation directe (forward-modeling) à travers plusieurs nœuds GPU (jusqu'à 128 nœuds ou 512 GPU A100). Le pipeline utilise l'API shard_map de JAX pour distribuer le calcul tout en maintenant un modèle de programmation unifié. La méthodologie traite trois étapes distinctes du pipeline de modélisation :

1. Stratégies de Parallélisation :

   • MAP et HMC : La parallélisation est réalisée en distribuant des particules indépendantes (pour l'estimation MAP multi-départ) ou des chaînes indépendantes (pour l'échantillonnage de Monte Carlo Hamiltonien - HMC) à travers les dispositifs.
   • Inférence Variationnelle Stochastique (SVI) : Cette étape est plus complexe car elle nécessite des mises à jour de gradients coordonnées entre les dispositifs pour optimiser une distribution postérieure de substitution partagée. L'implémentation tire une fraction d'échantillons sur chaque dispositif, calcule les gradients locaux via la différenciation automatique, agrège ces gradients à travers tous les nœuds et applique une étape d'optimiseur. Cela nécessite une communication inter-nœuds, gérée via JAX.

2. Améliorations Algorithmiques :

   • Sélection de l'Optimiseur : Les auteurs ont identifié que l'optimiseur Adam par défaut causait souvent une instabilité dans la SVI, caractérisée par des augmentations brusques de la perte avant récupération. Ils l'ont remplacé par AdaBelief, qui adapte les tailles de pas en fonction de la déviation du gradient, ce qui permet une convergence plus stable et une distribution de substitution plus proche de la distribution réelle.
   • Adaptation de la Matrice de Masse : Pour les systèmes difficiles où la substitution de la SVI reste médiocre, GIGA-Lens 2.0 implémente une nouvelle étape d'adaptation de la matrice de masse. Cela implique une phase de burn-in plus longue (environ 10$\times$ la phase typique) suivie de l'ajustement d'une nouvelle gaussienne de substitution sur les derniers 90 % des échantillons de burn-in afin de mieux initialiser le HMC.

3. Diagnostics de Convergence :
   Le cadre emploie la métrique de Gelman-Rubin ($\hat{R}$) et la taille d'échantillon effectif (ESS) pour assurer la convergence. L'implémentation prend correctement en compte la variance inter-chaînes dans les calculs d'ESS.

Résultats Clés

• Performance de Scalabilité : Le pipeline démontre une mise à l'échelle quasi linéaire. Pour de lourdes charges de travail, la fraction parallélisable ($p$) atteint 0,983 pour MAP, 0,979 pour SVI et 0,968 pour HMC. Même avec des charges de travail modestes, l'efficacité reste supérieure à 85 %.
• Systèmes Simulés : Les auteurs ont modélisé 100 systèmes de lentilles fortes simulés sur 16 GPU. Avec les nouveaux hyperparamètres et les améliorations algorithmiques, tous les systèmes sauf un ont convergé vers $\hat{R} < 1,1$ ; le dernier système a convergé après l'application de la nouvelle adaptation de la matrice de masse.
• Application à un Système Réel : Le cadre a été appliqué à un système réel, DESI J238.5690+04.7276. En utilisant huit nœuds, le temps d'exécution a bénéficié d'un gain de vitesse de cinq fois par rapport à un seul nœud.
• Résolution et Précision : Contrairement aux applications précédentes de GIGA-Lens 1.0 sur des données réelles HST qui utilisaient un facteur de suréchantillonnage (supersampling) de 1, GIGA-Lens 2.0 a modélisé avec succès ce système avec un suréchantillonnage = 2. De plus, alors que les travaux précédents atteignaient $\hat{R} < 1,1$, ce travail a atteint $\hat{R} < 1,01$ pour tous les paramètres du système réel, indiquant une convergence supérieure.

Signification et Revendications
L'article affirme que GIGA-Lens 2.0 surmonte le principal goulot d'étranglement pour les prochains relevés d'imagerie à large champ : la capacité de modéliser les lentilles de manière robuste et à grande échelle. La signification des gains de vitesse n'est pas seulement l'efficacité, mais la capacité de poursuivre la rigueur et la précision. En permettant une inférence bayésienne de haute fidélité à des échelles sans précédent, le cadre permet aux chercheurs de :

• Modéliser des systèmes complexes et de haute dimension (ex: la "Cosmic Carousel Lens") qui étaient auparavant computationnellement prohibitifs.
• Explorer systématiquement les hypothèses de modélisation et les paramétrisations pour quantifier les incertitudes systématiques, plutôt que de s'appuyer sur des hypothèses simplificatrices.
• Atteindre des métriques de convergence plus serrées ($\hat{R} < 1,01$) et une modélisation à plus haute résolution (suréchantillonnage = 2) pour les données observationnelles réelles.

Les auteurs soulignent que le cadre est désormais accessible, avec une documentation améliorée et des instructions de configuration reproductibles, rendant la modélisation de lentilles fortes, rigoureuse et scalable, puissante et disponible pour les grands projets d'inférence cosmologique.