Strong Gravitational Lensing Posterior Sampling in Pixel-Space Using Diffusion Models and Recurrent Inference Machines

Resumo Técnico: Amostragem de Posterior de Lente Gravitacional Forte em Espaço de Pixels utilizando Modelos de Difusão e Máquinas de Inferência Recorrentes

Definição do Problema
Modelar lentes gravitacionais fortes de galáxia-galáxia para inferir a distribuição de brilho de uma galáxia fonte e a distribuição de massa de uma galáxia em primeiro plano é um problema inverso computacionalmente exigente. Este desafio é exacerbado no regime de observações de alta resolução e alta razão sinal-ruído (SNR), onde suposições paramétricas de baixa dimensão (ex: leis de potência elípticas simples) são insuficientes para capturar a complexidade dos dados. Para modelar observações até o nível do ruído, são necessárias representações pixeladas de alta dimensão tanto para a fonte quanto para a massa em primeiro plano. No entanto, representar a lente e a fonte como imagens pixeladas torna o modelo direto (forward model) não linear em relação à massa em primeiro plano, criando um problema inverso de alta dimensão e não linear, caracterizado por distribuições posteriores complexas e multimodais e severas degenerações de parâmetros. Algoritmos de amostragem tradicionais e abordagens padrão de aprendizado de máquina frequentemente têm dificuldade em amostrar essas distribuições.

Metodologia: O Framework DiRIM
Os autores propõem o Diffusion Recurrent Inference Machine (DiRIM), um método Bayesiano projetado para gerar amostras conjuntas da posterior da luminosidade da galáxia fonte ($s$) e da distribuição de massa em primeiro plano (representada como um mapa de convergência, $\kappa$) condicionadas às observações ($y$). O framework integra dois paradigmas distintos de aprendizado de máquina:

1. Modelagem Generativa Baseada em Difusão: O método utiliza modelos de difusão baseados em score para amostrar de distribuições de probabilidade complexas. O processo envolve uma equação diferencial estocástica (SDE) direta que adiciona ruído gradualmente aos dados, e uma SDE reversa que reconstrói os dados a partir do ruído. Para amostrar da posterior $p(x|y)$, a SDE reversa requer o score da posterior $\nabla_x \log p_t(x|y)$.
2. Máquinas de Inferência Recorrentes (RIMs): As RIMs são empregadas para resolver o problema inverso ao aprender um algoritmo de inferência iterativo orientado por dados. Em vez de produzir uma única estimativa de Máxima Posterior (MAP), o framework RIM é adaptado aqui para refinar iterativamente a estimativa do score da posterior.

A Arquitetura DiRIM:

• Entrada: A rede neural $g_\theta$ (uma arquitetura U-Net) recebe como entrada os parâmetros do modelo ruidosos no tempo $t$ ($x_t$), a observação $y$, o passo de tempo $t$, os gradientes de verossimilhança ($\nabla L$) e os resíduos da imagem lensada normalizados ($\delta y / \sigma_y$).
• Refinamento Iterativo: Ao contrário dos modelos de difusão padrão que realizam um único passo de denoising por passo de tempo, o DiRIM usa o framework RIM para atualizar iterativamente sua estimativa da fonte e do mapa de convergência denoisados ($\hat{x}_0$) dentro de cada passo de tempo de difusão.
• Estimativa de Score: A rede é treinada para minimizar uma perda de correspondência de score de denoising condicional. A saída final das iterações da RIM ($\hat{x}_0^{(M)}$) é usada para calcular o score da posterior necessário para a SDE reversa:
  $$\nabla_{x_t} \log p_t(x_t|y) = \frac{\hat{x}_0^{(M)} - x_t}{\sigma^2(t)}$$
• Dados de Treinamento: O modelo é treinado em observações fictícias realistas geradas a partir de simulações magnetohidrodinâmicas cosmológicas (IllustrisTNG) e perfis analíticos. As imagens fonte são extraídas do conjunto de dados SKIRT-TNG (processado por transferência radiativa de poeira), e os mapas de convergência são derivados de halos de simulação suavizados ou renderizados de perfis analíticos (Lei de Potência Elíptica + multipolos + subhalos NFW).

Principais Contribuições

• Primeira Amostragem de Posterior Conjunta em Espaço de Pixels: Este trabalho apresenta o primeiro método capaz de amostrar a posterior conjunta da luminosidade da fonte e da massa em primeiro plano como imagens pixeladas para lentes gravitacionais fortes, indo além de estimativas pontuais ou ajustes paramétricos de baixa dimensão.
• Integração de RIMs e Difusão: O artigo introduz uma arquitetura inovadora onde uma RIM é incorporada dentro de um modelo de difusão condicional para refinar recursivamente os scores da posterior, abordando a não linearidade e a alta dimensão do problema de lensing.
• Modelagem ao Nível do Ruído: O método demonstra a capacidade de modelar observações até o nível do ruído, uma capacidade frequentemente perdida quando se utilizam modelos paramétricos simplificados.

Resultados
Os autores treinaram e avaliaram o DiRIM em dois conjuntos de dados: um utilizando mapas de convergência de simulações cosmológicas e outro utilizando perfis analíticos.

• Calibração: Testes de Accuracy with Random Points (TARP) indicam que as amostras da posterior estão quase perfeitamente calibradas, sugerindo que o estimador é não enviesado.
• Bondade de Ajuste (Goodness of Fit): A distribuição dos valores de $\chi^2$ para as imagens reconstruídas aproxima-se estreitamente da distribuição teórica de $\chi^2$, confirmando que os modelos capturam os dados até o nível do ruído.
• Tratamento de Complexidade: O método modela com sucesso sistemas com distribuições de massa complexas, incluindo múltiplos defletores principais e subhalos. Em casos onde modelos paramétricos analíticos falharam em ajustar os dados ao nível do ruído devido à falta de flexibilidade, o DiRIM obteu sucesso.
• Detecção de Subhalos: Ao ser aplicado a mapas analíticos contendo subhalos NFW, o método detectou os subhalos em mais de 98% das amostras da posterior, alcançando uma significância de detecção de $>3\sigma$.
• Amostragem de Degenerescência: O método amostra efetivamente direções degeneradas no espaço de parâmetros, como a correlação entre o tamanho da fonte e o declínio radial do mapa de convergência.
• Robustez (OOD): O modelo demonstrou robustez a dados fora da distribuição (OOD), inferindo com sucesso a morfologia de fontes e mapas de convergência que diferiam significamente da distribuição de treinamento (ex: galáxias em forma de "7" geradas por IA e mapas de convergência somados).
• Ablação: Remover o componente RIM (definindo as iterações $M=1$) resultou em inferências enviesadas e resíduos com estrutura significativa, confirmando que o refinamento iterativo dos scores é essencial para o sucesso do método.

Significância e Alegações
O artigo afirma que o DiRIM constitui uma "primeira prova de conceito" da capacidade de modelos generativos baseados em difusão combinados com máquinas de inferência recorrentes para resolver problemas inversos não lineares de alta dimensão na astrofísica. Os autores enfatizam que seu método aproveita representações flexíveis e de alta dimensão para extrair informações de dados que são inacessíveis à modelagem paramétrica tradicional. Eles posicionam este trabalho como um passo necessário para realizar o potencial científico de futuros levantamentos em larga escala (como LSST e Euclid), que devem descobrir mais de 100.000 lentes gravitacionais fortes. Os autores notam modestamente que, embora o método seja eficaz, pequenos vieses permanecem em sistemas complexos (ex: múltiplos defletores), e trabalhos futuros abordarão estes pontos e explorarão formulações de espaço latente para reduzir custos computacionais.