Joint inference of weak lensing convergence map and cosmology with diffusion models

Résumé technique : Inférence conjointe des cartes de convergence de lentillage gravitationnel faible et de la cosmologie avec les modèles de diffusion

Énoncé du problème
Les analyses actuelles de lentillage gravitationnel faible reposent généralement sur des statistiques de résumé (par exemple, les fonctions de corrélation à deux points) qui délaissent l'information non gaussienne présente dans la distribution de matière, particulièrement sur les petites échelles non linéaires. Bien que les approches de champ complet offrent un moyen rigoureux de capturer toute l'information statistique, les méthodes existantes font face à une dichotomie :

1. Inférence marginale : Les méthodes basées sur les réseaux de neurones infèrent directement les paramètres cosmologiques ($\theta$) à partir des données, mais marginalisent sur le champ de convergence ($\kappa$) sous-jacent, échouant ainsi à reconstruire la carte de masse.
2. Reconstruction de champ : Les méthodes de cartographie de masse reconstruisent $\kappa$ à partir du cisaillement bruité ($\gamma_{obs}$), mais traitent les paramètres cosmologiques comme fixes, plutôt que de les inférer conjointement.
3. Inférence explicite (MCMC) : Les approches comme BORG permettent une inférence conjointe ($p(\theta, \kappa | \gamma_{obs})$) en utilisant l'échantillonnage de Monte Carlo par chaînes de Markov (MCMC). Cependant, celles-ci nécessitent un modèle direct différentiable reliant les conditions initiales aux observables, ce qui limite la complexité de la physique employée, et souffrent de coûts de calcul élevés dans des espaces de grande dimension.

Les auteurs visent à développer une méthode pour l'inférence conjointe qui reconstruit simultanément le champ de convergence et infère les paramètres cosmologiques sans nécessiter de modèle direct différentiable ou d'échantillonnage MCMC itératif lors de l'inférence.

Méthodologie : JADE
Les auteurs introduisent JADE (Joint Architecture for fielD and cosmological parameter Estimation), un cadre d'inférence implicite basé sur des modèles de diffusion conditionnels.

• Inférence implicite via les modèles de diffusion : Au lieu d'apprendre une fonction de vraisemblance, JADE apprend la distribution postérieure conjointe $p(\theta, \kappa | \gamma_{obs})$ directement à partir de simulations. La méthode utilise le Flow Matching, un cadre génératif où un réseau de neurones apprend un champ de vitesse pour transporter une distribution de base (bruit gaussien) vers la distribution cible. Le modèle est entraîné pour prédire l'échantillon « propre » ($x_1 = [\theta, \kappa]$) à partir d'un interpolant bruité ($x_t$) étant donné le champ de cisaillement observé.
• Architecture Transformer : Le cœur de JADE est un Vision Transformer en espace de pixels adapté aux données multimodales.
  • Séquence de jetons unifiée : Le champ de convergence $\kappa$ est décomposé en patchs, et les paramètres cosmologiques $\theta$ sont projetés en plongements (embeddings) de même dimension. Ces derniers sont concaténés dans une séquence unique de jetons (tokens).
  • Conditionnement : Le champ de cisaillement bruité observé $\gamma_{obs}$ est également patchifié et projeté, puis concaténé à la séquence. Le mécanisme d'attention du transformer permet aux jetons cibles (champ et paramètres) de prêter attention à l'entrée de conditionnement (cisaillement) à toutes les positions spatiales.
  • Sortie : Le réseau produit des estimations débruitées pour la carte de convergence et pour les paramètres cosmologiques simultanément.
• Inférence amortie : Une fois entraîné, le modèle peut générer des échantillons de la distribution postérieure pour toute nouvelle observation en intégrant une équation différentielle ordinaire (ODE) conditionnée par la $\gamma_{obs}$ spécifique. Ce processus est dit « amorti », ce qui signifie qu'aucun réentraînement ou optimisation par observation n'est requis (contrairement au MCMC).

Configuration de la simulation
La méthode a été validée sur des données de lentillage faible simulées :

• Cosmologie : Un modèle $w$CDM avec six paramètres ($\Omega_c, \Omega_b, \sigma_8, h_0, n_s, w_0$).
• Génération de données : Les cartes de convergence ont été générées à l'aide d'un modèle log-normal 2D (modèle phénoménologique de croissance non linéaire) plutôt que par des simulations N-corps complètes, afin de permettre une génération rapide du vaste ensemble d'entraînement (100 000 échantillons).
• Observables : Champs de cisaillement simulés de la profondeur LSST Année 10 (Y10) avec bruit de forme, couvrant une zone de $5 \times 5$ deg$^2$ avec $128 \times 128$ pixels.

Résultats clés

1. Reconstruction conjointe : Le modèle a généré avec succès des échantillons de la distribution postérieure de $\kappa$ et de $\theta$ conditionnés par un cisaillement bruité. L'inspection visuelle et les tests statistiques ont confirmé que les cartes de convergence reconstruites récupèrent les structures de vérité terrain tout en reflétant correctement l'incertitude résiduelle sur les petites échelles.
2. Cohérence statistique :
   • Spectres de puissance : Les spectres de puissance auto et croisés des échantillons de la distribution postérieure correspondent à ceux de simulations indépendantes tirées aux cosmologies échantillonnées, avec une erreur relative d'environ 6,7 % en moyenne sur les bacs et les échelles.
   • PDF à un point : Les distributions récupérées présentent la forme log-normale attendue, confirmant la capture des statistiques non gaussiennes.
   • Résidus : La statistique du $\chi^2$ réduit ($\chi^2_{red} \approx 1,002$) a confirmé que les résidus entre les observations et les échantillons de la distribution postérieure étaient dominés par le bruit de forme injecté, comme prévu.
3. Contraintes cosmologiques : Les distributions postérieures cosmologiques marginales ($p(\theta | \gamma_{obs})$) dérivées de JADE montrent un excellent accord avec les chaînes MCMC traditionnelles NUTS (No-U-Turn Sampler) en termes de régions de crédibilité et de directions de dégénérescence.
4. Calibration :
   • Score MIRA : La distribution postérieure conjointe a obtenu un score de calibration MIRA (Mass In Random Areas) de $0,635 \pm 0,017$ (où 0,667 est l'optimum), indiquant une bonne calibration avec une légère surconfiance. La distribution postérieure cosmologique marginale a obtenu un score de $0,659 \pm 0,0139$.
   • TARP : Les tests d'exactitude avec des points aléatoires (Tests of Accuracy with Random Points) ont confirmé que les régions de crédibilité atteignent une couverture fréquentiste correcte à travers l'espace des paramètres.
5. Efficacité : La génération d'un seul échantillon de la distribution postérieure conjointe prend environ 0,2 seconde (256 étapes d'ODE). Tirer 1 000 échantillons prend environ 3 minutes, contre le coût de calcul nettement plus élevé de l'exécution de chaînes MCMC (qui nécessitaient environ 2,18 millions d'évaluations de gradient pour le benchmark).

Signification et revendications
L'article affirme établir l'inférence implicite basée sur la diffusion comme une voie viable vers des analyses cosmologiques de plein champ. La principale signification réside dans l'unification de la reconstruction de la carte de masse et de l'estimation des paramètres cosmologiques au sein d'un cadre unique non différentiable.

Les principales revendications concernant les avantages de la méthode incluent :

• Pas de modèle direct différentiable requis : Contrairement à l'inférence de niveau de champ par MCMC, JADE ne nécessite pas de gradients analytiques du modèle direct, ce qui permet théoriquement d'utiliser des simulateurs arbitrairement complexes (par exemple, des codes N-corps complets ou hydrodynamiques) une fois les données d'entraînement disponibles.
• Amortissement : Le modèle est entraîné une seule fois et peut être appliqué instantanément à toute nouvelle observation, évitant ainsi l'optimisation par observation.
• Inférence conjointe : C'est la première méthode présentée qui infère conjointement la distribution postérieure complète $p(\theta, \kappa | \gamma_{obs})$ à l'aide d'un seul modèle de diffusion conditionnel, plutôt que de marginaliser sur le champ ou d'inférer les paramètres séparément.

Limites et travaux futurs
Les auteurs notent que la démonstration actuelle repose sur des simulations log-normales, qui sont plus rapides que les codes N-corps mais moins complètes physiquement. Le passage à des simulateurs réalistes et coûteux présente des défis concernant la densité de l'ensemble d'entraînement dans l'espace des paramètres cosmologiques. De plus, en tant que méthode amortie, JADE doit être réentraîné si le modèle de données (par exemple, le masque de l'étude, le modèle de bruit) change, contrairement aux méthodes d'inférence explicite qui peuvent s'adapter au moment de l'inférence. Les auteurs proposent comme travail futur d'étendre cette architecture à l'inférence conjointe des conditions initiales et des paramètres cosmologiques, et d'explorer des approches de diffusion latente pour gérer plus efficacement les champs à haute résolution.