Joint inference of weak lensing convergence map and cosmology with diffusion models

Resumo Técnico: Inferência Conjunta de Mapas de Convergência de Lenteamento Fraco e Cosmologia com Modelos de Difusão

Definição do Problema
As análises atuais de lenteamento gravitacional fraco dependem tipicamente de estatísticas de resumo (ex: funções de correlação de dois pontos), que descartam a informação não-gaussiana presente na distribuição de matéria, particularmente em escalas pequenas e não lineares. Embora abordagens de campo total ofereçam uma maneira fundamentada de capturar toda a informação estatística, os métodos existentes enfrentam uma dicotomia:

1. Inferência Marginal: Métodos baseados em redes neurais inferem parâmetros cosmológicos ($\theta$) diretamente dos dados, mas marginalizam sobre o campo de convergência subjacente ($\kappa$), falhando em reconstruir o mapa de massa.
2. Reconstrução de Campo: Métodos de mapeamento de massa reconstroem $\kappa$ a partir de cisalhamento ruidoso ($\gamma_{obs}$), mas tratam os parâmetros cosmológicos como fixos, em vez de inferi-los conjuntamente.
3. Inferência Explícita (MCMC): Abordagens como o BORG realizam inferência conjunta ($p(\theta, \kappa | \gamma_{obs})$) usando Cadeias de Markov Monte Carlo (MCMC). No entanto, estas exigem um modelo direto diferenciável que ligue as condições iniciais aos observáveis, o que limita a complexidade da física que pode ser empregada, e sofrem com altos custos computacionais em espaços de alta dimensão.

Os autores visam desenvolver um método para inferência conjunta que reconstrça simultaneamente o campo de convergência e infira os parâmetros cosmológicos sem exigir um modelo direto diferenciável ou amostragem MCMC iterativa durante a inferência.

Metodologia: JADE
Os autores introduzem o JADE (Joint Architecture for fielD and cosmological parameter Estimation), um framework de inferência implícita baseado em modelos de difusão condicionais.

• Inferência Implícita via Modelos de Difusão: Em vez de aprender uma função de verossimilhança, o JADE aprende a distribuição posterior conjunta $p(\theta, \kappa | \gamma_{obs})$ diretamente de simulações. O método utiliza Flow Matching, um framework generativo onde uma rede neural aprende um campo de velocidade para transportar uma distribuição base (ruído Gaussiano) para a distribuição alvo. O modelo é treinado para prever a amostra "limpa" ($x_1 = [\theta, \kappa]$) a partir de um interpolante ruidoso ($x_t$) dado o campo de cisalhamento observado.
• Arquitetura Transformer: O núcleo do JADE é um Vision Transformer em espaço de pixels adaptado para dados multimodais.
  • Sequência de Tokens Unificada: O campo de convergência $\kappa$ é decomposto em patches, e os parâmetros cosmológicos $\theta$ são projetados em embeddings da mesma dimensão. Estes são concatenados em uma única sequência de tokens.
  • Condicionamento: O campo de cisalhamento ruidoso observado $\gamma_{obs}$ também é patcheado e projetado, sendo então concatenado à sequência. O mecanismo de atenção do transformer permite que os tokens alvo (campo e parâmetros) atentem para a entrada de condicionamento (cisalhamento) em todas as posições espaciais.
  • Saída: A rede produz estimativas de denoising tanto para o mapa de convergência quanto para os parâmetros cosmológicos simultaneamente.
• Inferência Amortizada: Uma vez treinado, o modelo pode gerar amostras posteriores para qualquer nova observação integrando uma Equação Diferencial Ordinária (ODE) condicionada ao $\gamma_{obs}$ específico. Este processo é "amortizado", o que significa que não é necessário retreinamento ou otimização por observação (como o MCMC).

Configuração de Simulação
O método foi validado em dados simulados de lenteamento fraco:

• Cosmologia: Um modelo $w$CDM com seis parâmetros ($\Omega_c, \Omega_b, \sigma_8, h_0, n_s, w_0$).
• Geração de Dados: Mapas de convergência foram gerados usando um modelo log-normal 2D (modelo fenomenológico para crescimento não linear) em vez de simulações completas de N-corpos, para permitir a geração rápida do grande conjunto de treinamento (100.000 amostras).
• Observáveis: Campos de cisalhamento com profundidade de LSST Ano 10 (Y10) com ruído de forma, cobrindo um patch de $5 \times 5$ deg$^2$ com $128 \times 128$ pixels.

Resultos Principais

1. Reconstrução Conjunta: O modelo gerou com sucesso amostras posteriores de ambos $\kappa$ e $\theta$ condicionados ao cisalhamento ruidoso. A inspeção visual e testes estatísticos confirmaram que os mapas de convergência reconstruídos recuperam estruturas da verdade fundamental (ground-truth), enquanto refletem corretamente a incerteza residual em escalas pequenas.
2. Consistência Estatística:
   • Espectros de Potência: Os espectros de potência auto e cruzados das amostras posteriores coincidiram com os de simulações independentes extraídas nas cosmologias amostradas, com um erro relativo de ~6,7% em média sobre bins e escalas.
   • PDFs de um ponto: As distribuições recuperadas exibiram a forma log-normal esperada, confirmando a captura de estatísticas não-Gaussianas.
   • Resíduos: A estatística de qui-quadrado reduzido ($\chi^2_{red} \approx 1,002$) confirmou que os resíduos entre observações e amostras posteriores foram dominados pelo ruído de forma injetado, como esperado.
3. Restrições Cosmológicas: As posteriors marginais cosmológicas ($p(\theta | \gamma_{obs})$) derivadas do JADE mostraram excelente concordância com cadeias MCMC tradicionais NUTS (No-U-Turn Sampler) em termos de regiões de credibilidade e direções de degenerescência.
4. Calibração:
   • Pontuação MIRA: A posterior conjunta alcançou uma pontuação de calibração MIRA (Mass In Random Areas) de $0,635 \pm 0,017$ (onde 0,667 é o ideal), indicando boa calibração com leve excesso de confiança. A posterior cosmológica marginal pontuou $0,659 \pm 0,0139$.
   • TARP: Testes de Precisão com Pontos Aleatórios confirmaram que as regiões de credibilidade atingiram a cobertura frequentista correta através do espaço de parâmetros.
5. Eficiência: Gerar uma única amostra da posterior conjunta leva aproximadamente 0,2 segundos (256 passos de ODE). Extrair 1.000 amostras leva ~3 minutos, comparado ao custo computacional significativamente maior de rodar cadeias MCMC (que exigiram ~2,18 milhões de avaliações de gradiente para o benchmark).

Significância e Alegações
O artigo afirma estabelecer a inferência implícita baseada em difusão como uma rota viável para análises de campo total de nível de lenteamento. A principal significância reside em unificar o mapeamento de massa e a estimação de parâmetros cosmológicos dentro de um único framework não-diferenciável.

As principais alegações sobre as vantagens do método incluem:

• Não Requer Modelo Direto Diferenciável: Ao contrário da inferência de nível de campo baseada em MCMC, o JADE não exige gradientes analíticos do modelo direto, permitindo teoricamente o uso de simuladores arbitrariamente complexos (ex: códigos de N-corpos completos ou hidrodinâmicos) uma vez que os dados de treinamento estejam disponíveis.
• Amortização: O modelo é treinado uma única vez e aplicado a qualquer nova observação instantaneamente, evitando a necessidade de amostragem por observação.
• Inferência Conjunta: É o primeiro método apresentado que infere conjuntamente a posterior completa $p(\theta, \kappa | \gamma_{obs})$ usando um único modelo de difusão condicional, em vez de marginalizar sobre o campo ou inferir parâmetros separadamente.

Limitações e Trabalhos Futuros
Os autores observam que a demonstração atual depende de simulações log-normais, que são mais rápidas que códigos de N-corpos, mas menos completas fisicamente. Escalar para simuladores realistas e caros apresenta desafios em relação à densidade do conjunto de treinamento no espaço de parâmetros cosmológicos. Além disso, como um método amortizado, o JADE deve ser retreinado se o modelo de dados (ex: máscara do levantamento, modelo de ruído) mudar, ao contrário de métodos de inferência explícita que podem se adaptar no tempo de inferência. Os autores propõem como trabalho futuro estender esta arquitetura para a inferência conjunta de condições iniciais e parâmetros cosmológicos, e explorar abordagens de difusão latente para lidar com campos de maior resolução de forma mais eficiente.