Fast Sampling of Cosmological Initial Conditions with Gaussian Neural Posterior Estimation

Resumo Técnico: Amostragem Rápida de Condições Iniciais Cosmológicas com Estimativa de Posterior Neural Gaussiana

Declaração do Problema
Reconstruir o campo de densidade de matéria primordial (condições iniciais cosmológicas, ou ICs) a partir de observações da estrutura em grande escala (LSS) é um desafio fundamental na cosmologia. Embora o Universo primitivo seja bem descrito por flutuações Gaussianas lineares, a teia cósmica atual é altamente não linear e não gaussiana. Reverter essa evolução é um problema mal posto: o mapeamento do início para o tempo tardio é não injetivo (diferentes ICs podem gerar o mesmo campo tardio) e a formação de estruturas em escalas de halos é caótica, amplificando o ruído de medição.

Abordagens probabilísticas tradicionais, como o Bayesian Origin Reconstruction from Galaxies (BORG), utilizam Hamiltonian Monte Carlo (HMC) para amostrar a distribuição posterior. No entanto, esses métodos exigem simuladores diferenciáveis (frequentemente dependendo de esquemas de aproximação) e são computacionalmente proibitivos, exigindo meses de tempo de CPU para convergir em espaços de parâmetros de alta dimensão ($\gtrsim 10^6$ parâmetros). Por outro lado, muitas abordagens de aprendizado de máquina neste domínio dependem de estimadores de ponto (por exemplo, U-Nets) que fornecem apenas uma reconstrução de "melhor ajuste" sem quantificar a incerteza, ou de modelos generativos (por exemplo, difusão baseada em score) que permanecem computacionalmente caros mesmo no tempo de inferência.

Metodologia
Os autores propõem um framework de Inferência Baseada em Simulação (SBI) denominado Estimativa de Posterior Neural Gaussiana (NPE). A inovação central é modelar a distribuição posterior das condições iniciais, $p(\mathbf{z}|\mathbf{x})$, como uma distribuição Gaussiana no campo inferido $\mathbf{z}$, com uma matriz de covariância diagonal no espaço de Fourier.

A metodologia procede da seguinte forma:

1. Formulação Estatística:

   • Prior: Assume-se que o prior $p(\mathbf{z})$ é um campo aleatório gaussiano definido pelo espectro de potência de matéria linear em alto redshift ($z=127$).
   • Likelihood: A verossimilhança $p(\mathbf{x}|\mathbf{z})$ é modelada como uma Gaussiana com uma média dependente da observação $\hat{\mathbf{z}}_\theta(\mathbf{x})$ e uma matriz de precisão $\mathbf{Q}^L_\theta$. Crucialmente, $\mathbf{Q}^L_\theta$ é restrita a ser diagonal no espaço de Fourier (e aproximadamente invariante à rotação), refletindo a invariância de translação da física.
   • Posterior: A combinação de prior Gaussiano e verossimilhança Gaussiana resulta em um posterior Gaussiano. A matriz de precisão do posterior é a soma das matrizes de precisão do prior e da verossimilhança ($\mathbf{Q}_\theta = \mathbf{Q}_P + \mathbf{Q}^L_\theta$).

2. Arquitetura Neural:

   • Estimador MAP: A média do posterior, $\hat{\boldsymbol{\mu}}_\theta(\mathbf{x})$, é modelada usando uma arquitetura U-Net 3D. Para melhorar o desempenho, os autores introduzem uma camada de ativação customizada que separa a reconstrução em duas escalas: um fator de escala treinável $\alpha_\theta(\mathbf{k})$ lida com a evolução linear de grande escala, enquanto a U-Net foca na reconstrução de resíduos não lineares de pequena escala.
   • Matriz de Precisão: Os elementos diagonais da matriz de precisão da verossimilhança, $D^L_\theta(k)$, são aprendidos. Os autores testaram duas parametrizações: (1) tratando os valores como parâmetros livres para cada modo de Fourier, e (2) impondo simetria rotacional via uma rede neural densa que recebe o número de onda $k$ como entrada. A abordagem de matriz "livre" mostrou maior estabilidade.

3. Treinamento e Dados:

   • O modelo é treinado em 2.000 pares de campos de sobredensidade de matéria escura inicial ($z=127$) e final ($z=0$) da suíte de simulações N-body Quijote.
   • O modelo de propagação (forward model) utiliza simulações N-body completas e não diferenciáveis (Gadget-III), evitando a necessidade de aproximações diferenciáveis exigidas pelos métodos HMC.
   • O treinamento minimiza o log-posterior negativo (Eq. 7) ao longo de 30 épocas em uma única GPU Nvidia A100, levando aproximadamente 1,5 hora.

Principais Contribuições

• Amostragem Eficiente: Ao assumir um posterior Gaussiano com uma matriz de precisão diagonal no espaço de Fourier, os autores permitem a geração de milhares de amostras do posterior em segundos (aproximadamente 3 segundos para 1.000 amostras em uma única GPU), ordens de magnitude mais rápido que os métodos HMC ou de difusão existentes.
• Simuladores Não Diferenciáveis: O método não exige que o modelo de propagação seja diferenciável, permitindo o uso de simulações N-body de última geração e completas sem esquemas de aproximação.
• Ajuste Analítico para Amostragem Rápida: Os autores demonstram que os valores diagonais da matriz de precisão aprendida $D_\theta(k)$ podem ser ajustados com precisão através de uma função analítica simples (uma Gaussiana mais uma constante) até $k \approx 0,2 \, h/\text{Mpc}$. Isso permite que qualquer estimador de ponto (MAP) treinado seja convertido em um amostrador de posterior rápido sem necessidade de novo treinamento da matriz de precisão completa.
• Quantificação de Incerteza: Ao contrário dos estimadores de ponto padrão, este método fornece uma distribuição posterior completa, permitindo uma quantificação rigorosa de incerteza para a reconstrução de campos.

Resultos
Os autores validam o método usando uma simulação fiducial não presente no conjunto de treinamento:

• Consistência Visual e Estatística: As amostras geradas assemelham-se visualmente ao valor real (ground truth). Estatísticas de resumo (espectro de potência, função de transferência, correlação cruzada) mostram alta concordância com as condições iniciais reais. O espectro de potência coincide com o real dentro de 1–2% até a escala Nyquist ($k_{Nyq} \approx 0,4 \, h/\text{Mpc}$).
• Correlação Cruzada: A correlação cruzada entre as amostras geradas e o valor real permanece alta ($\gtrsim 70\%$) até $k \approx 0,3 \, h/\text{Mpc}$.
• Consistência Bayesiana: Um teste de consistência dependente de escala confirma que as diferenças normalizadas entre os modos reais e a média do posterior seguem uma distribuição normal padrão ($N(0,1)$) através de vários bins de $k$. Isso verifica que as incertezas inferidas estão bem calibradas.
• Robustez: O método mantém o desempenho mesmo quando ruído Gaussiano é adicionado ao campo final observado.

Significância e Alegações
O artigo afirma que esta abordagem fornece um caminho eficiente em termos de simulação para SBI sequencial em cosmologia. Ao reduzir a complexidade do posterior para uma forma Gaussiana com uma matriz de precisão diagonal, os autores alcançam um equilíbrio entre precisão física e velocidade computacional que era anteriormente inalcançável para a reconstrução de campos de alta dimensão.

Os autores observam modestamente que a aproximação Gaussiana de Fourier-diagonal é suficiente para o intervalo de escala selecionado ($k \lesssim 0,4 \, h/\text{Mpc}$), mas pode exigir matrizes de precisão mais sofisticadas (por exemplo, baseadas em convolução ou wavelets) para escalas menores ou quando efeitos observacionais quebram a invariância de translação. Eles enfatizam que, embora o trabalho atual assuma parâmetros cosmológicos fixos e observações idealizadas de matéria escura, o framework foi desenhado para ser estendido a cosmologias variáveis e levantamentos de galáxias realistas em trabalhos futuros. A principal significância reside em permitir a geração de milhares de realizações contidas em segundos, abrindo caminho para análises estatísticas robustas da LSS que eram anteriormente inviáveis computacionalmente.