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

Resumen Técnico: Evaluación de la Cuantificación de la Incertidumbre en AION-1

Planteamiento del Problema
Los modelos fundacionales para sondeos astronómicos, como AION-1, proporcionan representaciones aprendidas potentes para tareas de regresión descendentes, como la estimación de propiedades galácticas. Sin embargo, las predicciones puntuales por sí solas son insuficientes para la inferencia científica, la cual requiere una Cuantificación de la Incertidumbre (UQ, por sus siglas en inglés) fiable para propagar errores, identificar entradas fuera de la distribución e asegurar inferencias robustas a nivel poblacional. Aunque la UQ es crítica, sigue estando poco comprendida para los modelos fundacionales de gran escala y tipo "black-box" con embeddings de alta dimensión. Los métodos existentes a menudo carecen de garantías formales de cobertura o no logran adaptarse a las dificultades locales de predicción en datos astrofísicos complejos y heterogéneos.

Metodología
Los autores llevan a cabo una comparación directa de siete métodos de UQ aplicados al modelo fundacional congelado AION-1 (específicamente la variante AION-B con 300 millones de parámetros). El estudio se centra en la regresión de cinco propiedades físicas de la fotometría/imágenes de Legacy Survey y espectros de DESI: desplazamiento al rojo ($z$), masa estelar ($\log M_\star$), edad de la población estelar ($t_{age}$), metalicidad de fase gaseosa ($\log Z$) y tasa específica de formación estelar ($\log sSFR$). Las etiquetas se derivan de los posteriors de ajuste de SED de PROVABGS.

La configuración experimental involucra:

1. Arquitectura Base: Un MLP de tres capas mapea los embeddings de 768 dimensiones de AION-1 hacia predicciones escalares. Los pesos del codificador permanecen congelados; solo se entrena la cabeza de regresión.
2. Métodos de UQ Evaluados:
   • Métodos Conformes (Garantías libres de distribución):
     • VanillaSplit: Predicción conforme de división estándar con cuantiles globales (intervalos de ancho constante).
     • MADSplit: Variante de MAD-normalizada que adapta el ancho del intervalo a la dificultad local.
     • CQR (Regresión de Cuantiles Conformalizada): Entrena redes de cuantiles seguidas de una corrección conforme.
     • LVD (Localmente Válido y Discriminativo): Aprende un kernel en un espacio de características reducido para estimar los cuantiles de no conformidad condicional. Se probó en dos espacios: sobre los embeddings brutos de AION-1 (LCAion) y sobre las activaciones ocultas del MLP de 256 dimensiones (LCNet).
     • LVD-MAD: LVD extendido con un predictor de magnitud de residuo aprendido para intervalos adaptativos de escala (LCMADAion y LCMADNet).
   • Baselines No Conformes:
     • Deep Ensembles (DE): Cinco MLPs inicializados de forma independiente.
     • MC Dropout: Pasadas hacia adelante estocásticas con el dropout habilitado.

La evaluación mide la cobertura marginal, la cobertura de cola (10% superior/inferior), la cobertura en el bin peor predicho (10% de los puntos más difíciles) y la eficiencia (ancho medio del intervalo de predicción).

Resultados Clave

• Cobertura Marginal: Todos los métodos conformes (variantes de LVD, VanillaSplit, MADSplit, CQR) logran una cobertura marginal dentro de $\sim$1 punto porcentual del objetivo nominal del 90% para todas las propiedades. En contraste, los baselines no conformes fallan en la calibración: Deep Ensembles subcalibra sistemáticamente (12–38%), mientras que MC Dropout exhibe una sobrecobertura fuerte (57–99%) con intervalos excesivamente anchos.
• Desempeño en Cola y Peor Predicción:
  • En el bin peor predicho, CQR logra la mayor cobertura (44–66%), seguido de las variantes LVD-MAD (22–42%).
  • VanillaSplit colapsa a $\sim$0% de cobertura en las regiones más difíciles.
  • Las variantes de LVD que operan sobre los embeddings brutos de AION-1 (LCAion, LCMADAion) generalmente superan a aquellas que utilizan activaciones intermedias comprimidas del MLP, lo que sugiere que el espacio de 768 dimensiones retiene la estructura beneficiosa para la calibración local.
  • Las variantes MAD (LCMADNet, LCMADAion) superan consistentemente a sus contrapartes sin MAD, resaltando la importancia de la normalización adaptativa de escala para la heterocedasticidad.
• Eficiencia: Entre los métodos bien calibrados, las variantes de LVD producen los intervalos de predicción más estrechos manteniendo la cobertura. CQR produce intervalos más anchos para garantizar la validez marginal.

Contribuciones Clave

1. Benchmarking Sistemático: El artículo proporciona la primera comparación directa de siete métodos de UQ sobre los embeddings congelados de AION-1 para cinco tareas distintas de regresión de propiedades galácticas.
2. Demostración de Validez Local: El estudio demuestra que, si bien los métodos conformes garantizan la validez marginal, el marco de la Predicción Localmente Válida y Discriminativa (LVD) —particularmente cuando opera sobre los embeddings del modelo fundacional— proporciona validez local de muestra finita. Esto permite que los intervalos de predicción se adapten a la dificultad específica de la predicción de cada galaxia en lugar de depender únicamente de promedios globales.
3. Análisis del Espacio de Embeddings: Los resultados indican que utilizar los embeddings brutos de alta dimensión (768-dim) en lugar de las características intermedias comprimidas (256-dim) mejora la robustez de la calibración local, particularmente en las regiones de cola.

Significancia y Afirmaciones
Los autores concluyen que la predicción conforme, y específicamente el marco LVD, debería ser el enfoque de UQ preferido para la inferencia consciente de la incertidumbre sobre los embeddings de modelos fundacionales en astrofísica. A diferencia de los baselines no conformes que carecen de garantías formales, y los métodos conformes estándar que pueden sufrir de descalibración condicional, LVD ofrece un equilibrio principista entre las garantías libres de distribución y la adaptabilidad local. Esto es crucial para la inferencia científica donde la fiabilidad de la predicción individual es tan importante como las estadísticas de la población. El artículo señala que, aunque se utilizó AION-B, la ventaja de LVD en los embeddings justifica una mayor investigación en escalas de modelos más grandes (AION-L y AION-XL).