Increasing the Precision of Surrogate Models for Weak Lensing Mass Maps with Flow Matching

Resumo Técnico: Aumento da Precisão de Modelos Surrogados para Mapas de Massa de Lente Gravitacional Fraca com Flow Matching

Declaração do Problema
A lente gravitacional fraca fornece uma sonda direta da estrutura em grande escala do cosmos, codificando informações não gaussianas em mapas de convergência (massa) que são inacessíveis às estatísticas de dois pontos tradicionais. A modelagem direta precisa desses mapas é essencial para a inferência cosmológica, dependendo tipicamente de simulações de $N$-corpos computacionalmente caras. Embora emuladores de aprendizado profundo, particularmente Redes Adversariais Generativas (GANs), tenham sido propostos para acelerar esse processo, os métodos existentes sofrem de limitações significativas. Surrogados baseados em GANs anteriores frequentemente falham em corresponder a estatísticas de resumo com precisão suficiente (tipicamente erros de 5–10%), produzem mapas excessivamente suavizados, lutam para capturar a distribuição correta de conjuntos de mapas gerados (falhando em reproduzir a distribuição probabilística completa em vez de apenas o "modo") e exibem instabilidade no treinamento. Além disso, a geração condicional através de parâmetros cosmológicos contínuos (por exemplo, densidade de matéria $\Omega_m$ e amplitude de aglomeração $\sigma_8$) permanece desafiadora devido a dificuldades de aprendizado heterogêneas através do espaço de rótulos.

Metodologia
Os autores propõem uma estrutura de Flow Matching Condicional de Rótulo Residual (RLCFM) para gerar mapas de convergência de alta fidelidade condicionados a parâmetros cosmológicos. A abordagem baseia-se em Fluxos Normalizantes Contínuos (CNF) e Flow Matching (FM), que aprendem um campo vetorial dependente do tempo para transportar uma distribuição a priori para uma distribuição de dados alvo.

Componentes metodológicos chave incluem:

1. Flow Matching Condicional de Rótulo Residual: Em vez de treinar um fluxo condicional de rótulo padrão, os autores introduzem uma reparametrização em um espaço residual. Eles subtraem a média específica do rótulo ($\mu_0(y)$) dos mapas alvo antes do treinamento. Isso remove deslocamentos de média triviais dependentes do rótulo, permitindo que a rede se concentre em aprender a estrutura residual transferível. Este design visa estabilizar a interpolação através de parâmetros cosmológicos.
2. Priori Gaussiana Específica do Rótulo: O modelo emprega uma priori gaussiana específica do rótulo, orientada pelos dados, $p(x|y)$, em vez de uma única priori compartilhada. A média e a variância da priori são derivadas das estatísticas da distribuição condicional $q(x|y)$ em $t=0$, transportando essa priori específica para a distribuição alvo em $t=1$.
3. Objetivo de Treinamento: A rede é treinada para minimizar a perda de Flow Matching, regressando o campo vetorial dependente do tempo que transporta a priori específica do rótulo para os mapas de convergência alvo. O treinamento é realizado em um conjunto de dados de 57 modelos cosmológicos ($\Omega_m, \sigma_8$) dentro de um framework $\Lambda$CDM plano, compreendendo 684.000 mapas de $128 \times 128$ pixels.
4. Pré-processamento de Dados: Mapas de convergência são transformados via logaritmo (aproximando um campo lognormal) e padronizados. O modelo aprende a gerar o campo residual $r = \tilde{x} - \mu_0(y)$, que é então reconstruído na inferência adicionando a média específica do rótulo de volta.

Contribuições Principais

• Arquitetura Inovadora: Introdução da arquitetura RLCFM, que combina priores específicos do rótulo com treinamento em espaço residual para abordar a instabilidade e os desafios de interpolação da modelagem generativa condicional em cosmologia.
• Fidelidade Aprimorada: O método alcança precisão significativamente maior do que benchmarks anteriores de GANs. Reduz erros típicos em estatísticas básicas (distribuições de pixels/picos, espectro de potência) de $<10\%$ para $<1\%$, e em estatísticas de ordem superior (bispectro, funcionais de Minkowski) de $<20\%$ para $<5\%$.
• Precisão ao Nível de Distribuição: Diferentemente das GANs, que frequentemente ajustam o "modo" dos dados, o modelo RLCFM reproduz com sucesso a distribuição completa de estatísticas de resumo geradas a partir de condições iniciais aleatórias, correspondendo à variabilidade das simulações de $N$-corpos.
• Interpolação Robusta: O modelo demonstra desempenho estável através do espaço de parâmetros cosmológicos, incluindo em cosmologias de teste retidas e nas bordas da grade, com erros aumentando apenas moderadamente nas fronteiras em comparação com o centro.

Resultados
Os autores avaliaram o modelo RLCFM contra simulações de $N$-corpos usando uma suite abrangente de métricas:

• Qualidade Visual: Mapas gerados são visualmente indistinguíveis das referências de $N$-corpos, com evolução morfológica consistente à medida que $\Omega_m$ e $\sigma_8$ variam.
• Estatísticas de Resumo:
  • Distribuições de Pixels e Picos: Distâncias de Wasserstein-1 normalizadas são tipicamente $<0,01$.
  • Espectro de Potência ($C_\ell$): Diferenças fracionárias são tipicamente $<1\%$.
  • Bispectro ($B_\ell$): Diferenças fracionárias são tipicamente $<5\%$.
  • Funcionais de Minkowski ($V_0, V_1, V_2$): Diferenças fracionárias são geralmente abaixo de 3%.
  • Matrizes de Correlação: O contraste da norma de Frobenius entre as matrizes de correlação RLCFM e de $N$-corpos é tipicamente $<3\%$.
  • Semelhança Estrutural (MS-SSIM): A significância da diferença nas pontuações de semelhança estrutural é baixa, indicando que o modelo não sofre de colapso de modo e preserva o realismo estrutural multiescala.
• Eficiência: O modelo foi treinado em aproximadamente 7 horas (65.000 passos) em uma única GPU NVIDIA A100, exigindo significativamente menos épocas para alcançar desempenho estável em comparação com as bases de GAN.

Significado e Afirmações
O artigo afirma que o flow matching condicional, quando aumentado com uma formulação em espaço residual e uma priori orientada pelos dados, fornece uma estrutura poderosa e eficiente para gerar mapas de massa cosmológicos. O significado principal reside na usabilidade prática do emulador: ele captura o sinal cosmológico com alta fidelidade através de estatísticas de baixa e alta ordem, enquanto suporta múltiplas formas de análise de dados. Os autores afirmam que este trabalho aproxima o campo de um emulador de mapas de massa prático capaz de suportar a modelagem direta para futuras pesquisas de campo amplo (por exemplo, Euclid, LSST) sem o custo computacional de executar simulações completas de $N$-corpos para cada ponto de parâmetro cosmológico. O método é apresentado como um passo em direção a pipelines de inferência cosmológica totalmente generativos e baseados em simulação.