Interpretable Neural Marked Statistics for Cosmological Inference

Resumo Técnico: Marcas Equivariantes Interpretáveis para Inferência Cosmológica Contrastiva

Definição do Problema
Recuperar informação cosmológica além do espectro de potência é um objetivo central para os próximos levantamentos da Grande Estrutura em Larga Escala (LSS) (ex: SPHEREx, DESI, Euclid). Embora o espectro de potência seja ótimo para campos gaussianos, a densidade de matéria em tempos tardios exibe sinais não-gaussianos que as estatísticas de dois pontos sozinhas não conseguem capturar. As abordagens tradicionais para extrair essa informação enfrentam obstáculos significativos:

• Correladores de ordem superior: Subir a escada dos correladores de $n$-pontos é computacionalmente caro no espaço de configuração, e a teoria de perturbação falha no regime não-linear.
• Inferência ao nível do campo: A inferência baseada em simulações e resumos neurais ao nível do campo evita estatísticas explícitas, mas frequentemente sacrifica a interpretabilidade e exige simulações de alta fidelidade.
• Estatísticas de marcas clássicas: Métodos existentes reponderam o campo de densidade com uma marca espacial $M(x)$ para incorporar correlações de ordem superior no nível de dois pontos. No entanto, essas marcas são tipicamente projetadas manualmente com formas funcionais fixas (ex: respostas de densidade suavizadas) ou otimizadas apenas em uma única cosmologia fiducial, o que limita sua adaptabilidade e interpretabilidade através do espaço de parâmetros.

Metodologia
Os autores propõem um framework que substitui marcas projetadas manualmente por um esquema de marcação neural aprendido, cosmologicamente agnóstico e interpretável. O framework consiste em dois componentes primários:

1. Arquitetura de Marca Aprendida e Interpretável:

   • A marca $M(x)$ é construída como uma resposta escalar de translação-equivariante derivada do campo de densidade de matéria $\delta(x)$.
   • Filtragem por Harmônicos Esféricos: O campo de densidade é decomposto em canais usando filtros que consideram harmônicos esféricos ($\ell \in \{0, 1, 2\}$) no espaço de Fourier. Esses filtros produzem canais escalares invariantes de rotação:
     • $\ell=0$: Uma resposta de monopolo com sinal ($E_0$).
     • $\ell=1$: Uma resposta do tipo vetorial com amplitude invariante ($E_1$).
     • $\ell=2$: Dois invariantes escalares derivados de um tensor simétrico traço-nulo ($E_2$ para força de anisotropia, $I_3$ para distinguir formas prolatas/oblatas).
   • Processamento: Esses canais invariantes são processados por MLPs independentes e combinados aditivamente com um termo de interação cruzada. O resultado passa por uma ativação softplus para garantir a positividade.
   • Inicialização: A rede é inicializada de modo que $M(x) \approx 1$, significando que o treinamento começa a partir do campo não marcado. Isso permite que a marca aprendida seja decomposta em contribuições morfológicas distintas.

2. Objetivo de Aprendizado Contrastivo:

   • O framework emprega um objetivo de aprendizado contrastivo (InfoNCE) para alinhar resumos marcados com parâmetros cosmológicos $\theta$ em um espaço latente compartilhado.
   • Vetores de Resumo: O resumo inclui o espectro de potência do campo não marcado ($P_{\delta\delta}$) e do campo marcado ($P_{\Delta\Delta}$), juntamente com o espectro cruzado ($P_{\delta\Delta}$) e o coeficiente de correlação.
   • Residualização: Para garantir que a marca aprenda informações complementares em vez de simplesmente reproduzir o espectro de potência, o embedding do resumo marcado ($z_M$) é residualizado contra o embedding não marcado ($z_\delta$) para produzir um componente ortogonal $z_\perp$.
   • Função de Perda: A perda alinha $z_\perp$ com o embedding de parâmetro $z_\theta$ usando uma métrica de Mahalanobis aprendida, enquanto o separa de candidatos negativos amostrados do prior (negativos reais, sintéticos globais e sintéticos locais).
   • Estratégia de Treinamento: O ramo não marcado é pré-treinado e congelado; então, a rede de marcação e o embaralhador de parâmetros são ajustados usando a perda contrastiva.

Principais Contribuições

• Generalização de Estatísticas de Marcas: O artigo introduz uma arquitetícia neural que generaliza o conceito de estatísticas de marcas, indo além de formas funcionais fixas para uma transformação aprendida que permanece inspecionável no espaço de configuração.
• Interpretabilidade via Equivariância: Ao impor $\text{SO}(3)$-equivariância local e usar filtros de harmônicos esféricos, o modelo extrai características com significados físicos claros (respostas de monopolo, dipolo, quadrupolo) que podem ser diretamente visualizadas e analisadas.
• Alinhamento Contrastivo: O uso de um objetivo contrastivo com residualização garante que a marca aprendida foque especificamente em informações ortogonais ao espectro de potência padrão, quebrando degenerações que as estatísticas tradicionais não conseguem superar.

Resultados
O método foi avaliado na suíte de simulações N-body Quijote ($z=0$) com 5 parâmetros cosmológicos variáveis ($\Omega_m, \Omega_b, h, n_s, \sigma_8$).

• Restrições de Informação de Fisher: Em $k_{\max} = 0.2 \, h\text{Mpc}^{-1}$, a marca aprendida aperta as restrições marginalizadas em 2.9$\times$ para $\sigma_8$ e 1.8$\times$ para $\Omega_m$ comparado às marcas clássicas. Crucialmente, ela quebra a degenerescência canônica $\Omega_m - \sigma_8$, rotacionando o contorno do posterior conjunto para permitir a melhoria simultânea em ambos os parâmetros.
• Generalização (MSE): Através de um conjunto de teste de hipercubo latino, a marca aprendida reduz o Erro Quadrático Médio (MSE) de parâmetros em 1.45$\times$ sobre a melhor marca clássica. Demonstra uma generalização robusta, particularmente para $\sigma_8$, embora as marcas clássicas permaneçam ligeiramente mais estáveis para $\Omega_m$ em configurações específicas de $(p, \delta_s)$.
• Geometria Latente: A Análise de Componentes Principais (PCA) do embedding aprendido revela que os dois primeiros componentes principais alinham-se quase perfeitamente com as direções de $\sigma_8$ e $\Omega_m$, confirmando que o objetivo contrastivo recuperou com sucesso os eixos dominantes da informação cosmológica.
• Introspecção: A análise da marca treinada mostra que o canal isotrópico ($\ell=0$) domina a resposta, recuperando efetivamente uma reponderação não-linear da densidade, enquanto os canais anisotrópicos fornecem correções subdominantes nas fronteiras das estruturas.

Significância
O artigo afirma que esta abordagem abre as portas para estatísticas de resumo mais poderosas e interpretáveis para a inferência cosmológica. Ao substituir transformações manuais por uma arquitetura aprendida e motivada fisicamente, o método preenche a lacuna entre o alto poder de restrição das redes neurais de nível de campo e a interpretabilidade das estatísticas tradicionais. A capacidade de inspecionar a marca aprendida no espaço de configuração permite aos pesquisadores entender quais características morfológicas locais (ex: vazios, filamentos, estruturas anisotrópicas) carregam o sinal cosmológico mais rico. Os autores posicionam isso como um passo em direção a ferramentas de inferência flexíveis que podem reduzir custos de simulação e melhorar o retorno científico de levantamentos de galáxias, com potenciais extensões para outros parâmetros como a não-gaussianidade primordial.