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

Resumo Técnico: Benchmarking de Quantificação de Incerteza em AION-1

Definição do Problema
Modelos de fundação para levantamentos astronômicos, como o AION-1, fornecem representações aprendidas poderosas para tarefas de regressão downstream, como a estimação de propriedades de galáxias. No entanto, previsões pontuais isoladas são insuficientes para a inferência científica, que requer Quantificação de Incerteza (UQ) confiável para propagar erros, identificar entradas fora da distribuição e garantir inferências robustas em nível populacional. Embora a UQ seja crítica, ela permanece mal compreendida para modelos de fundação de grande escala e "caixa-preta" com embeddings de alta dimensão. Os métodos existentes frequentemente carecem de garantias formais de cobertura ou falham em se adaptar às dificuldades locais de predição em dados astrofísicos complexos e heterogêneos.

Metodologia
Os autores realizam uma comparação direta de sete métodos de UQ aplicados ao modelo de fundação congelado AION-1 (especificamente a variante AION-B com 300 milhões de parâmetros). O estudo foca na regressão de cinco propriedades físicas da fotometria/imagem do Legacy Survey e espectros do DESI: redshift ($z$), massa estelar ($\log M_\star$), idade da população estelar ($t_{age}$), metalicidade da fase gasosa ($\log Z$) e taxa específica de formação estelar ($\log sSFR$). Os rótulos são derivados dos posteriors de ajuste de SED do PROVABGS.

A configuração experimental envolve:

1. Arquitetura Base: Um MLP de três camadas mapeia os embeddings de 768 dimensões do AION-1 para previsões escalares. Os pesos do codificador estão congelados; apenas a cabeça de regressão é treinada.
2. Métodos de UQ Avaliados:
   • Métodos Conformes (Garantias independentes de distribuição):
     • VanillaSplit: Predição conforme de divisão padrão com quantis globais (intervalos de largura constante).
     • MADSplit: Variante com normalização MAD que adapta a largura do intervalo à dificuldade local.
     • CQR (Regressão de Quantis Conformalizada): Treina redes de quantis seguidas de correção conformal.
     • LVD (Locally Valid and Discriminative): Aprende um kernel em um espaço de características reduzido para estimar quantis de não conformidade condicional. Testado em dois espaços: nos embeddings brutos do AION-1 (LCAion) e nas ativações ocultas do MLP de 256 dimensões (LCNet).
     • LVD-MAD: LVD estendido com um preditor de magnitude de resíduo aprendido para intervalos adaptativos de escala (LCMADAion e LCMADNet).
   • Baselines Não-Conformes:
     • Deep Ensembles (DE): Cinco MLPs inicializados de forma independente.
     • MC Dropout: Passagens diretas estocásticas com dropout habilitado.

A avaliação mede cobertura marginal, cobertura de cauda (10% superior/inferior), cobertura no pior bin previsto (10% dos pontos mais difíceis) e eficiência (largura média do intervalo de predição).

Principais Resultados

• Cobertura Marginal: Todos os métodos conformes (variantes LVD, VanillaSplit, MADSplit, CQR) alcançam cobertura marginal dentro de $\sim$1 ponto percentual do alvo nominal de 90% para todas as propriedades. Em contraste, as baselines não-conformes falham em calibrar: Deep Ensembles subcobertura sistematicamente (12–38%), enquanto o MC Dropout exibe forte sobrecobertura (57–99%) com intervalos excessivamente largos.
• Desempenho de Cauda e de Pior Predição:
  • No pior bin previsto, o CQR alcança a maior cobertura (44–66%), seguido pelas variantes LVD-MAD (22–42%).
  • O VanillaSplit colapsa para $\sim$0% de cobertura nas regiões mais difíceis.
  • As variantes LVD operando nos embeddings brutos do AION-1 (LCAion, LCMADAion) geralmente superam aquelas que utilizam ativações comprimidas do MLP, sugerindo que o espaço de 768 dimensões retém estrutura benéfica para calibração local.
  • As variantes MAD (LCMADNet, LCMADAion) superam consistentemente suas contrapartes não-MAD, destacando a importância da normalização adaptativa de escala para heterocedasticidade.
• Eficiência: Entre os métodos bem calibrados, as variantes LVD produzem os intervalos de predição mais estreitos enquanto mantêm a cobertura. O CQR produz intervalos mais largos para garantir validade marginal.

Principais Contribuições

1. Benchmarking Sistemático: O artigo fornece a primeira comparação direta de sete métodos de UQ em embeddings congelados do AION-1 para cinco tarefas distintas de regressão de propriedades galácticas.
2. Demonstração de Validade Local: O estudo demonstra que, embora os métodos conformes garantam validade marginal, o framework Locally Valid and Discriminative (LVD) — particularmente quando opera em embeddings de modelos de fundação — fornece validade local de amostra finita. Isso permite que os intervalos de predição se adaptem à dificuldade específica da predição de cada galáxia, em vez de depender apenas de médias globais.
3. Análise do Espaço de Embedding: Os resultados indicam que utilizar os embeddings bros de alta dimensão (768-dim) em vez de características intermediárias comprimidas (256-dim) aumenta a robustez da calibração local, particularmente em regiões de cauda.

Significância e Alegações
Os autores concluem que a predição conforme, e especificamente o framework LVD, deve ser a abordagem de UQ preferencial para inferência consciente de incerteza em embeddings de modelos de fundação na astrofísica. Diferente das baselines não-conformes, que carecem de garantias formais, e dos métodos conformes padrão, que podem sofrer de má calibração condicional, o LVD oferece um equilíbrio principiado entre garantias independentes de distribuição e adaptabilidade local. Isso é crucial para a inferência científica, onde a confiabilidade da predição individual é tão importante quanto as estatísticas populacionais. O artigo observa que, embora o AION-B tenha sido utilizado, a vantagem do LVD em embeddings justifica investigações adicionais em escalas de modelos maiores (AION-L e AION-XL).