Beyond Point Estimates: Benchmarking Uncertainty Quantification Methods on the AION-1 Astronomical Foundation Model

Résumé technique : Évaluation comparative de la quantification de l'incertitude sur AION-1

Énoncé du problème
Les modèles de fondation pour les relevés astronomiques, tels qu'AION-1, fournissent des représentations apprises puissantes pour les tâches de régression en aval, comme l'estimation des propriétés galactiques. Cependant, les prédictions ponctuelles seules sont insuffisantes pour l'inférence scientifique, qui nécessite une quantification de l'incertitude (UQ) fiable afin de propager les erreurs, d'identifier les entrées hors distribution et d'assurer des inférences robustes au niveau de la population. Bien que l'UQ soit critique, elle reste mal comprise pour les modèles de fondation massifs et "boîtes noires" dotés d'embeddings de haute dimension. Les méthodes existantes manquent souvent de garanties de couverture formelles ou ne parviennent pas à s'adapter aux difficultés de prédiction locales dans les données astrophysiques complexes et hétérogènes.

Méthodologie
Les auteurs mènent une comparaison directe de sept méthodes d'UQ appliquées au modèle de fondation gelé AION-1 (spécifiquement la variante AION-B avec 300 millions de paramètres). L'étude se concentre sur la régression de cinq propriétés physiques issues de la photométrie/imagerie de Legacy Survey et des spectres DESI : le décalage vers le rouge ($z$), la masse stellaire ($\log M_\star$), l'âge de la population stellaire ($t_{age}$), la métallicité de la phase gazeuse ($\log Z$) et le taux de formation stellaire spécifique ($\log sfr$). Les étiquettes sont dérivées des postulats (posteriors) de l'ajustement SED de PROVABGS.

Le dispositif expérimental comprend :

1. Architecture de base : Un MLP à trois couches projette les embeddings de 768 dimensions d'AION-1 vers des prédictions scalaires. Les poids de l'encodeur sont gelés ; seule la tête de régression est entraînée.
2. Méthodes d'UQ évaluées :
   • Méthodes conformes (garanties sans distribution) :
     • VanillaSplit : Prédiction conforme par division standard avec des quantiles globaux (intervalles de largeur constante).
     • MADSplit : Variante normalisée par MAD adaptant la largeur de l'intervalle à la difficulté locale.
     • CQR (Conformalized Quantile Regression) : Entraîne des réseaux de quantiles suivis d'une correction conforme.
     • LVD (Locally Valid and Discriminative) : Apprend un noyau dans un espace de caractéristiques réduit pour estimer les quantiles de non-conformité conditionnels. Testé dans deux espaces : sur les embeddings bruts d'AION-1 (LCAion) et sur les activations cachées de l'MLP à 256 dimensions (LCNet).
     • LVD-MAD : LVD étendu avec un prédicteur de magnitude de résidu appris pour des intervalles adaptatifs à l'échelle (LCMADAion et LCMADNet).
   • Baselines non-conformes :
     • Deep Ensembles (DE) : Cinq MLP initialisés de manière indépendante.
     • MC Dropout : Passes avant stochastiques avec le dropout activé.

L'évaluation mesure la couverture marginale, la couverture de queue (10 % supérieur/inférieur), la couverture du pire bin prédit (les 10 % de points les plus difficiles) et l'efficacité (largeur moyenne de l'intervalle de prédiction).

Résultats clés

• Couverture marginale : Toutes les méthodes conformes (variantes LVD, VanillaSplit, MADSplit, CQR) atteignent une couverture marginale à environ 1 point de pourcentage de la cible nominale de 90 % pour toutes les propriétés. En revanche, les baselines non-conformes échouent à se calibrer : les Deep Ensembles sous-estiment systématiquement la couverture (12–38 %), tandis que le MC Dropout présente une surcouverture marquée (57–99 %) avec des intervalles excessivement larges.
• Performance de queue et des pires prédictions :
  • Dans le pire bin prédit, la méthode CQR atteint la couverture la plus élevée (44–66 %), suivie par les variantes LVD-MAD (22–42 %).
  • VanillaSplit s'effondre à environ 0 % de couverture dans les régions les plus difficiles.
  • Les variantes LVD opérant sur les embeddings bruts d'AION-1 (LCAion, LCMADAion) surpassent généralement celles utilisant les activations compressées de l'MLP, suggérant que l'espace de 768 dimensions conserve une structure bénéfique pour la calibration locale.
  • Les variantes MAD (LCMADNet, LCMADAion) surpassent systématiquement leurs homologues non-MAD, soulignant l'importance de la normalisation adaptative à l'échelle pour l'hétéroscédasticité.
• Efficacité : Parmi les méthodes bien calibrées, les variantes LVD produisent les intervalles de prédiction les plus serrés tout en maintenant la couverture. La méthode CQR produit des intervalles plus larges pour garantir la validité marginale.

Contributions clés

1. Évaluation systématique : Le papier fournit la première comparaison directe de sept méthodes d'UQ sur les embeddings gelés d'AION-1 pour cinq tâches distinctes de régression de propriétés galactiques.
2. Démonstration de la validité locale : L'étude démontre que si les méthodes conformes garantissent la validité marginale, le cadre Locally Valid and Discriminative (LVD) — particulièrement lorsqu'il opère sur les embeddings de modèles de fondation — fournit une validité locale en échantillon fini. Cela permet aux intervalles de prédiction de s'adapter à la difficulté spécifique de la prédiction de chaque galaxie plutôt que de reposer uniquement sur des moyennes globales.
3. Analyse de l'espace d'embedding : Les résultats indiquent que l'utilisation des embeddings bruts de haute dimension (768 dim) plutôt que des caractéristiques intermédiaires compressées (256 dim) améliore la robustesse de la calibration locale, particulièrement dans les zones de queue.

Signification et affirmations
Les auteurs concluent que la prédiction conforme, et spécifiquement le cadre LVD, devrait être l'approche d'UQ privilégiée pour l'inférence tenant compte de l'incertitude sur les embeddings de modèles de fondation en astrophysique. Contrairement aux baselines non-conformes qui manquent de garanties formelles, et aux méthodes conformes standards qui peuvent souffrir d'une mauvaise calibration conditionnelle, LVD offre un équilibre principled entre garanties sans distribution et adaptativité locale. Ceci est crucial pour l'inférence scientifique où la fiabilité individuelle de la prédiction est aussi importante que les statistiques de population. Le papier note que bien que l'AION-B ait été utilisé, l'avantage de LVD sur les embeddings justifie des investigations plus poussées à des échelles de modèles plus larges (AION-L et AION-XL).