Increasing the Precision of Surrogate Models for Weak Lensing Mass Maps with Flow Matching

Résumé technique : Augmentation de la précision des modèles de substitution pour les cartes de masse du lentillage gravitationnel faible par appariement de flux

Énoncé du problème
Le lentillage gravitationnel faible fournit une sonde directe de la structure à grande échelle de l'univers, encodant des informations non gaussiennes dans les cartes de convergence (masse) qui sont inaccessibles aux statistiques à deux points traditionnelles. Une modélisation directe précise de ces cartes est essentielle pour l'inférence cosmologique, reposant généralement sur des simulations de corps N ($N$-body) coûteuses en calcul. Bien que des émulateurs d'apprentissage profond, en particulier les réseaux antagonistes génératifs (GAN), aient été proposés pour accélérer ce processus, les méthodes existantes souffrent de limitations significatives. Les modèles de substitution basés sur les GAN précédents échouent souvent à correspondre aux statistiques résumées avec une précision suffisante (erreurs typiques de 5 à 10 %), produisent des cartes trop lissées, peinent à capturer la distribution correcte des ensembles de cartes générées (échouant à reproduire la distribution probabiliste complète plutôt que simplement le « mode ») et présentent une instabilité lors de l'entraînement. De plus, la génération conditionnelle sur des paramètres cosmologiques continus (par exemple, la densité de matière $\Omega_m$ et l'amplitude du regroupement $\sigma_8$) reste difficile en raison de difficultés d'apprentissage hétérogènes dans l'espace des étiquettes.

Méthodologie
Les auteurs proposent un cadre d'Appariement de Flux Conditionnel aux Étiquettes Résiduelles (RLCFM) pour générer des cartes de convergence haute fidélité conditionnées par des paramètres cosmologiques. L'approche s'appuie sur les Flux Normalisants Continus (CNF) et l'Appariement de Flux (FM), qui apprennent un champ vectoriel dépendant du temps pour transporter une distribution a priori vers une distribution de données cible.

Les composants méthodologiques clés incluent :

1. Appariement de Flux Conditionnel aux Étiquettes Résiduelles : Au lieu d'entraîner un flux conditionnel aux étiquettes standard, les auteurs introduisent une reparamétrisation dans un espace résiduel. Ils soustraient la moyenne spécifique à l'étiquette ($\mu_0(y)$) des cartes cibles avant l'entraînement. Cela élimine les décalages de moyenne triviaux dépendants de l'étiquette, permettant au réseau de se concentrer sur l'apprentissage de la structure résiduelle transférable. Cette conception vise à stabiliser l'interpolation à travers les paramètres cosmologiques.
2. Prior Gaussien Spécifique à l'Étiquette : Le modèle utilise un prior gaussien spécifique à l'étiquette, piloté par les données $p(x|y)$, plutôt qu'un prior unique partagé. La moyenne et la variance du prior sont dérivées des statistiques de la distribution conditionnelle $q(x|y)$ à $t=0$, transportant ce prior spécifique vers la distribution cible à $t=1$.
3. Objectif d'Entraînement : Le réseau est entraîné pour minimiser la perte d'appariement de flux, en régressant le champ vectoriel dépendant du temps qui transporte le prior spécifique à l'étiquette vers les cartes de convergence cibles. L'entraînement est effectué sur un ensemble de données de 57 modèles cosmologiques ($\Omega_m, \sigma_8$) dans le cadre d'un $\Lambda$CDM plat, comprenant 684 000 cartes de $128 \times 128$ pixels.
4. Prétraitement des Données : Les cartes de convergence sont transformées par un logarithme (approximant un champ lognormal) et standardisées. Le modèle apprend à générer le champ résiduel $r = \tilde{x} - \mu_0(y)$, qui est ensuite reconstruit lors de l'inférence en réajoutant la moyenne spécifique à l'étiquette.

Contributions clés

• Architecture Nouvelle : Introduction de l'architecture RLCFM, qui combine des priors spécifiques à l'étiquette avec un entraînement dans l'espace résiduel pour répondre aux défis d'instabilité et d'interpolation de la modélisation générative conditionnelle en cosmologie.
• Fidélité Améliorée : La méthode atteint une précision nettement supérieure aux références GAN précédentes. Elle réduit les erreurs typiques dans les statistiques de base (distributions de pixels/sommets, spectre de puissance) de $<10\%$ à $<1\%$, et dans les statistiques d'ordre supérieur (bispectre, fonctionnels de Minkowski) de $<20\%$ à $<5\%$.
• Précision au Niveau de la Distribution : Contrairement aux GAN qui ajustent souvent le « mode » des données, le modèle RLCFM reproduit avec succès la distribution complète des statistiques résumées générées à partir de conditions initiales aléatoires, correspondant à la variabilité des simulations de corps N.
• Interpolation Robuste : Le modèle démontre des performances stables à travers l'espace des paramètres cosmologiques, y compris pour des cosmologies de test non vues et aux bords de la grille, avec des erreurs n'augmentant que légèrement aux frontières par rapport au centre.

Résultats
Les auteurs ont évalué le modèle RLCFM par rapport aux simulations de corps N en utilisant une suite complète de métriques :

• Qualité Visuelle : Les cartes générées sont visuellement indiscernables des références de simulations de corps N, avec une évolution morphologique cohérente lorsque $\Omega_m$ et $\sigma_8$ varient.
• Statistiques Résumées :
  • Distributions de Pixels et de Sommets : Les distances de Wasserstein-1 normalisées sont typiquement $<0,01$.
  • Spectre de Puissance ($C_\ell$) : Les différences fractionnaires sont typiquement $<1\%$.
  • Bispectre ($B_\ell$) : Les différences fractionnaires sont typiquement $<5\%$.
  • Fonctionnels de Minkowski ($V_0, V_1, V_2$) : Les différences fractionnaires sont généralement inférieures à 3 %.
  • Matrices de Corrélation : La norme de contraste de Frobenius entre les matrices de corrélation RLCFM et celles des simulations de corps N est typiquement $<3\%$.
  • Similarité Structurelle (MS-SSIM) : La signification de la différence dans les scores de similarité structurelle est faible, indiquant que le modèle ne souffre pas d'effondrement de mode et préserve le réalisme structurel multi-échelle.
• Efficacité : Le modèle a été entraîné en environ 7 heures (65 000 étapes) sur un seul GPU NVIDIA A100, nécessitant nettement moins d'époques pour atteindre des performances stables par rapport aux bases de référence GAN.

Signification et Revendications
L'article revendique que l'appariement de flux conditionnel, lorsqu'il est augmenté d'une formulation dans l'espace résiduel et d'un prior piloté par les données, fournit un cadre puissant et efficace pour générer des cartes de masse cosmologiques. La signification principale réside dans l'utilisabilité pratique de l'émulateur : il capture le signal cosmologique avec une haute fidélité à la fois pour les statistiques d'ordre faible et d'ordre élevé, tout en supportant plusieurs formes d'analyse de données. Les auteurs affirment que ce travail rapproche le domaine d'un émulateur de cartes de masse pratique capable de soutenir la modélisation directe pour les futures enquêtes à grand champ (par exemple, Euclid, LSST) sans le coût computationnel d'exécuter des simulations de corps N complètes pour chaque point de paramètre cosmologique. La méthode est présentée comme une étape vers des pipelines d'inférence cosmologique entièrement génératifs et basés sur des simulations.