Learned harmonic mean estimation of the marginal likelihood for multimodal posteriors with flow matching

Resumo Técnico: Estimativa da Média Harmônica Aprendida com Flow Matching para Posteriores Multimodais

Definição do Problema
A verossimilhança marginal (ou evidência bayesiana), denotada como $z = p(d|M)$, é uma quantidade fundamental para a comparação de modelos bayesianos, permitindo uma avaliação quantitativa via fator de Bayes. No entanto, calcular essa quantidade requer a avaliação de uma integral multidimensional sobre todo o espaço de parâmetros, o que é computacionalmente proibitivo para modelos complexos e de alta dimensão. Enquanto os métodos de Cadeia de Markov Monte Carlo (MCMC) e Hamiltonian Monte Carlo (HMC) se destacam na geração de amostras posteriores em altas dimensões, eles não fornecem inerentemente a verossimilhança marginal. Por outro lado, o nested sampling fornece a evidência, mas restringe a estratégia de amostragem a uma hierarquia aninhada específica.

O "estimador de média harmônica aprendida" foi introduzido anteriormente para estimar a verossimilhança marginal usando apenas amostras da posterior, independentemente do método de amostragem utilizado para obtê-las. Essa abordagem baseia-se no aprendizado de uma densidade alvo interna $\phi(\theta)$ que está contida na posterior, mas possui caudas mais finas. Contudo, implementações anteriores utilizando fluxos normalizadores discretos ou estimadores de densidade customizados enfrentaram dificuldades com posteriores altamente multimodais. Essas arquiteturas frequentemente falhavam em capturar topologias complexas sem modificações heurísticas (por exemplo, a construção manual de misturas gaussianas para a distribuição base) ou sofriam ao atribuir massa de probabilidade não desprezável a pontes de densidade entre os modos, levando à variância infinita no estimador.

Metodologia
Este trabalho aborda as limitações de estimadores de densidade anteriores integrando o Flow Matching à estrutura da média harmônica aprendida. A metodologia procede da seguinte forma:

1. Estimador de Média Harmônica Aprendida: O estimador aproxima o recíproco da verossimilhança marginal $\rho = z^{-1}$ usando amostras da posterior $\theta_i \sim p(\theta|d)$ e uma densidade alvo aprendida $\phi(\theta)$:
   $$\hat{\rho} = \frac{1}{N} \sum_{i=1}^{N} \frac{\phi(\theta_i)}{L(\theta_i)\pi(\theta_i)}$$
   Para que o estimador tenha variância finita, $\phi(\theta)$ deve estar concentrada nas regiões de alta probabilidade da posterior e decair mais rapidamente nas caudas.

2. Arquitetura de Flow Matching: Em vez de fluxos normalizadores discretos, os autores utilizam Fluxos Normalizadores Contínuos (CNFs) treinados via flow matching.

   • Definição: Um CNF parametriza a transformação entre uma distribuição base simples (Gaussiana padrão) e a posterior alvo através de um campo de velocidade dependente do tempo $v_t(\theta)$.
   • Objetivo de Treinamento: A rede é treinada usando um objetivo de regressão livre de simulação (simulation-free). Um caminho de probabilidade condicional é definido como uma linha reta entre uma amostra de ruído $\theta_0$ e uma amostra da posterior $\theta_1$: $\theta_t = (1-t)\theta_0 + t\theta_1$. A rede minimiza o erro quadrático médio entre o campo de velocidade previsto e o verdadeiro campo vetorial condicional $(\theta_1 - \theta_0)$.
   • Vantagens: Esta abordagem permite o uso de redes neurais de alta capacidade e sem restrições arquiteturais (MLPs com ativação Swish) e evita as simulações dispendiosas de ODE exigidas em métodos de CNF anteriores durante o treinamento.

3. Escalonamento de Temperatura: Para garantir que a densidade alvo $\phi(\theta)$ possua caudas mais finas que a posterior, um parâmetro de temperatura $T \in (0, 1)$ é aplicado para escalar a variância da distribuição base do fluxo. Isso concentra a distribuição do fluxo, satisfazendo o requisito de contenção para o estimador de média harmônica.

Principais Contribuições

• Arquitetura Inovadora: A introdução de fluxos normalizadores contínuos treinados por flow matching como o estimador de densidade interna para o método da média harmônica aprendida.
• Tratamento de Multimodalidade: Demonstração de que esta arquitetura pode capturar com precisão estruturas posteriores altamente multimodais e complexas sem a necessidade de modificações heurísticas para a distribuição base ou ajustes finos.
• Escalabilidade: Provisão de um método que é agnóstico à estratégia de amostragem (compatível com MCMC, HMC, NUTS, SBI, etc.) e capaz de operar em espaços de dimensão moderada a alta.

Resultados Experimentais
Os autores validaram a abordagem em dois testes de referência desafiadores:

1. Função Rastrigin (2 Dimensões):

   • Um teste padrão com muitos modos locais.
   • O modelo de flow matching (6 camadas, 256 unidades ocultas) capturou com sucesso todos os modos com uma temperatura $T=0.98$.
   • Os resultados de 100 execuções experimentais mostraram que os valores da verossimilhança marginal estimados foram precisos e corresponderam de perto à verdade fundamental obtida via integração numérica. As estimativas de erro foram confiáveis.

2. Mistura Gaussiana de 20 Dimensões:

   • Uma mistura de cinco componentes Gaussianas em 20 dimensões com modos não sobrepostos e correlações introduzidas.
   • O modelo (10 camadas, 256 unidades ocultas) foi treinado em 40.000 amostras.
   • A distribuição alvo aprendida mostrou-se adequadamente concentrada dentro da posterior, capturando a topologia desafiadora.
   • As verossimilhanilhas marginais estimadas ao longo de 100 execuções alinharam-se com a verdade fundamental analítica, confirmando a precisão do método e a robustez da estimativa de erro mesmo neste cenário complexo de dimensão moderada.

Significância e Alegações
O artigo afirma que, ao alavancar a expressividade dos CNFs treinados por flow matching, o estimador de média harmônica aprendida pode agora estimar robustamente a evidência para posteriors altamente complexas e multimodais. A principal significância reside na remoção da dependência de ajustes heurísticos (como misturas Gaussianas escolhidas manualmente) que eram anteriormente necessários para lidar com a multimodalidade. Os autores afirmam que esta abordagem permite o cálculo preciso da evidência usando amostras de qualquer método de amostragem acelerado ou escalável, facilitando a comparação de modelos bayesianos em campos ricos em dados, como a astrofísica, sem os gargalos computacionais do nested sampling tradicional ou a instabilidade do estimador de média harmônica original.

O trabalho futuro, conforme observado pelos autores, inclui a aplicação desta metodologia a problemas do mundo real, como a astrofísica de ondas gravitacionais, e a exploração de novos refinamentos arquiteturais para tarefas de altíssima dimensão, como processamento de imagens e inferência cosmológica de nível de campo.