Three-dimensional Conditional Diffusion Models for Cosmological 21 cm Lightcone Emulation

Résumé technique : Modèles de diffusion conditionnelle tridimensionnels pour l'émulation de cônes de lumière cosmologiques 21 cm

Énoncé du problème
Le signal 21 cm décalé vers le rouge provenant de l'hydrogène neutre est une sonde critique de l'Aube Cosmique et de l'Ère de la Réionisation (EoR). L'interprétation des futures données d'imagerie provenant d'instruments comme HERA et SKA nécessite des modèles directs rapides et physiquement fiables couvrant de vastes espaces de paramètres astrophysiques. Bien que les simulations semi-numériques comme 21cmFAST soient la norme, elles sont coûteuses en calcul pour des balayages de paramètres à grande échelle ou pour l'inférence basée sur des simulations. Les émulateurs d'apprentissage automatique, en particulier les modèles génératifs, offrent une solution. Cependant, l'extension de ces modèles de tranches d'images 2D à des cônes de lumière 3D complets introduit des défis significatifs. L'ajout d'une troisième dimension (profondeur le long de la ligne de visée) augmente considérablement les coûts mémoire, imposant des micro-lots très petits qui amplifient le bruit d'optimisation. De plus, la distribution sous-jacente des voxels des températures de brillance 21 cm est fortement asymétrique et à queue longue, rendant les distorsions dans l'espace des transformations et le contrôle de la plage dynamique critiques pour un entraînement stable et une fidélité physique.

Méthodologie
Les auteurs étudient la modélisation par diffusion conditionnelle pour l'émulation de cônes de lumière 21 cm 3D en utilisant un ensemble de données de 25 600 cônes de lumière d'entraînement générés par 21cmFAST.

• Données et conditionnement : Les champs cibles sont des cubes de voxels de $64 \times 64 \times 1024$ couvrant une plage de décalage vers le rouge de $z \in [7.51, 20.95]$. Les modèles sont conditionnés par deux paramètres astrophysiques : la température viriale minimale des halos ($\log_{10} T_{\text{vir}}$) et l'efficacité d'ionisation ($\zeta$).
• Architecture : Un U-Net conditionné par le contexte est employé, utilisant des pas de sous-échantillonnage/sur-échantillonnage anisotropes de $(2, 2, 4)$ pour s'adapter à la géométrie allongée du cône de lumière. Le modèle utilise 128 canaux de base avec des multiplicateurs de canaux $(1, 1, 2, 4)$ et intègre des embeddings d'étape temporelle et de condition.
• Pipeline de prétraitement : Un point central de l'étude est le pipeline de prétraitement, défini comme suit : $x_{\text{raw}} \to \text{transformation primaire} \to \text{écrasement optionnel} \to \text{diffusion} \to \text{désécrasement} \to \text{transformation inverse}$. Les auteurs comparent quatre transformations de normalisation primaires (Yeo–Johnson, Min–Max, Z-score, Arcsinh) et une étape optionnelle de compression de plage dynamique (« écrasement »).
• Conception de la validation : L'étude utilise 800 réalisations indépendantes par point de paramètre pour la validation à travers cinq ensembles de paramètres de référence distincts. Cette conception teste à la fois la tendance centrale et la diversité conditionnelle de la distribution cible.
• Métriques d'évaluation : Les échantillons générés sont évalués à l'aide d'une hiérarchie de diagnostics :
  1. Tranches d'images de température de brillance ($T_b$) (qualitatif).
  2. Signal global $\langle T_b \rangle$ en fonction du décalage vers le rouge (évolution moyenne).
  3. Spectre de puissance (structure spatiale à deux points).
  4. Coefficients de diffusion réduits (morphologie non gaussienne d'ordre supérieur).
     La performance est quantifiée à l'aide de l'erreur absolue moyenne normalisée par l'écart-type ($\text{MAE}_{\text{std}}$) comme métrique principale de classement.

Contributions et résultats clés

• Domination du prétraitement : L'étude identifie le prétraitement et le contrôle de la plage dynamique comme les facteurs dominants régissant la stabilité de l'entraînement et la fidélité physique dans les contextes 3D, surpassant la profondeur architecturale ou la durée de l'entraînement une fois l'optimisation stabilisée.
• Configuration optimale : Parmi les configurations testées, le prétraitement Yeo–Johnson combiné à une compression modérée de l'amplitude ($A=0.1$) produit le compromis le plus favorable de manière cohérente. Cette configuration surpasse les transformations Min–Max, Z-score et Arcsinh, ainsi que les configurations sans compression d'amplitude.
• Hiérarchie de fidélité : Le succès de l'émulation suit une hiérarchie claire :
  • Signal global : Le modèle atteint le meilleur accord, capturant avec précision le timing et la profondeur des creux d'absorption.
  • Spectre de puissance : L'accord est plus faible, en particulier aux modes de haut-$k$ (petites échelles), où le modèle peine à reproduire des frontières nettes.
  • Coefficients de diffusion : La morphologie non gaussienne d'ordre supérieur reste la plus difficile à reproduire. Bien que le modèle capture les motifs structurels larges, des biais mesurables persistent dans la morphologie multi-échelle.
• Modes d'échec : La limitation principale est l'incapacité du modèle à reproduire parfaitement la population de voxels fortement concentrée et proche de zéro trouvée dans 21cmFAST. Le modèle de diffusion tend à élargir ce pic, entraînant des régions ionisées « plus lisses » et de petits pixels non nuls dans des zones visuellement sombres. Ce lissage se manifeste par des écarts résiduels dans le spectre de puissance aux petites échelles et dans les statistiques d'ordre supérieur.
• Sensibilité aux hyperparamètres : L'augmentation de la profondeur du modèle (blocs résiduels) n'a pas produit d'améliorations monotones ; en fait, le modèle le plus peu profond (1 bloc résiduel) a surpassé les variantes plus profondes, probablement en raison de la dynamique d'optimisation bruyante induite par les micro-lots réduits en 3D.

Importance et affirmations
Les auteurs positionnent ce travail comme une référence au niveau de la simulation pour l'émulation 21 cm tridimensionnelle plutôt que comme une pipeline d'inférence complète de bout en bout prête pour les données de relevés.

• Insight méthodologique : L'article soutient que dans les régimes de haute dimension, une génération stable dépend principalement du prétraitement et du contrôle de la plage dynamique. La seule vraisemblance visuelle est insuffisante ; des diagnostics sensibles à la morphologie (comme les coefficients de diffusion) sont nécessaires pour détecter des biais que des résumés d'ordre inférieur pourraient manquer.
• Portée : L'étude est un test d'interpolation au sein de l'espace de paramètres d'entraînement. Elle n'aborde pas encore les défis liés aux relevés tels que la réponse instrumentale, la contamination par le fond ou l'étalonnage des incertitudes pour l'estimation des paramètres.
• Utilité future : La configuration optimale identifiée (Yeo–Johnson + $A=0.1$) fournit un point de départ robuste pour les futures études visant à intégrer des effets observationnels réalistes ou à réaliser une inférence basée sur des simulations, à condition que les limitations concernant la morphologie à petite échelle soient reconnues.