Mixed neural posterior estimation for simulators with discrete and continuous parameters

Resumo Técnico: Estimação Neural de Posterior Mista para Simuladores com Parâmetros Discretos e Contínuos

Declaração do Problema
A Estimação Neural de Posterior (NPE) tornou-se uma ferramenta padrão para Inferência Baseada em Simulação (SBI), permitindo inferência paramétrica rápida e amortizada para simuladores complexos com verossimilhanças intratáveis. No entanto, os métodos padrão de NPE assumem espaços paramétricos contínuos, modelando tipicamente o posterior como uma densidade de probabilidade contínua usando fluxos normalizadores ou modelos de difusão. Essa suposição falha para uma ampla classe de simuladores científicos onde o espaço paramétrico é misto, contendo tanto dimensões contínuas (por exemplo, taxas, condutâncias) quanto dimensões discretas (por exemplo, variáveis categóricas, índices de seleção de modelos, localizações de pontos de mudança).

Abordagens existentes para inferência mista enfrentam limitações significativas. Os métodos padrão de Cadeia de Markov Monte Carlo (MCMC) exigem ajuste cuidadoso, não amortizam os custos de inferência entre observações e frequentemente lutam com variáveis discretas (por exemplo, exigindo marginalização analítica ou dinâmicas hamiltonianas descontínuas especializadas). Métodos recentes de SBI baseados em difusão tentam incorporar parâmetros discretos em espaços contínuos ou tratá-los como índices de modelo, mas estes podem ser computacionalmente pesados ou exigir embebidamentos riemannianos complexos. Consequentemente, simuladores com espaços paramétricos mistos permanecem insuficientemente suportados pelas atuais caixas de ferramentas de SBI.

Metodologia: Estimação Neural de Posterior Mista (MNPE)
Os autores propõem a Estimação Neural de Posterior Mista (MNPE), uma extensão da NPE que trata conjuntamente parâmetros discretos e contínuos por meio de uma rede de inferência fatorada.

1. Estratégia de Fatorização:
   O posterior conjunto $p(\theta_d, \theta_c | x)$ sobre parâmetros discretos $\theta_d$ e parâmetros contínuos $\theta_c$ é fatorado como:
   $$p(\theta_d, \theta_c | x) = p(\theta_d | x) \cdot p(\theta_c | \theta_d, x)$$
   Isso separa o problema em duas subtarefas: classificação sobre o espaço discreto e estimação de densidade condicional no espaço contínuo, condicionada ao estado discreto.

2. Arquitetura da Rede:

   • Componente Discreto: O termo $p(\theta_d | x)$ é modelado usando um Estimador de Densidade Autoregressivo Mascado (MADE). Esta rede autoregressiva produz probabilidades de classe para cada dimensão discreta, condicionada às dimensões discretas precedentes e à observação $x$. Isso permite lidar com números arbitrários de parâmetros discretos com contagens de classes variáveis.
   • Componente Contínuo: O termo $p(\theta_c | \theta_d, x)$ é modelado usando um modelo generativo condicional padrão (por exemplo, um fluxo de spline neural ou modelo de difusão) condicionado tanto à observação $x$ quanto aos parâmetros discretos $\theta_d$.
   • Objetivo de Treinamento: A rede é treinada conjuntamente para minimizar o log-probabilidade negativa dos parâmetros mistos:
     $$\mathcal{L} = -\log q(\theta_d, \theta_c | x) = -\log q(\theta_c | \theta_d, x) - \log q(\theta_d | x)$$
     A parte contínua usa a verossimilhança padrão de fluxo, enquanto a parte discreta reduz-se a uma perda de entropia cruzada.

3. Diagnósticos de Calibração:
   Para avaliar a qualidade de posterores mistos sem soluções de referência, os autores propõem um quadro diagnóstico combinado:

   • Dimensões Contínuas: Usa Calibração Baseada em Simulação (SBC) via estatísticas de rank para verificar a uniformidade nas distribuições marginais de rank.
   • Dimensões Discretas: Usa Erro de Calibração Esperado (ECE) e diagramas de confiabilidade, tratando a inferência discreta como um problema de classificação.
   • Linhas de Base: Os autores derivam linhas de base empíricas de amostra finita para ambas as métricas para distinguir má calibração real de ruído estatístico, reconhecendo que a calibração perfeita é inatingível com amostras finitas.

Resultados Chave
Os autores validam a MNPE em três simuladores de complexidade crescente:

1. Exemplo Toy Gaussiano Tratável:

   • Um simulador com uma variável latente contínua e uma variável latente discreta acopladas através de uma observação compartilhada.
   • A MNPE recupera com precisão o posterior analítico de verdade, incluindo regimes bimodais onde o estado discreto é incerto.
   • A pontuação do Teste de Duas Amostras com Classificador (C2ST) converge para o nível de acaso (~0,5) com aproximadamente 1.000 simulações, indicando indistinguibilidade do posterior verdadeiro.

2. Simulador de Fila em Tandem:

   • Um simulador com três taxas contínuas e duas contagens de servidores discretas ($|D|=25$).
   • A MNPE é comparada contra referências MCMC (PyMC e NumPyro) onde as dimensões discretas são marginalizadas analiticamente.
   • Com 100.000 simulações, os posteriores da MNPE correspondem à referência MCMC em ambas as dimensões contínuas e discretas.
   • Diagnósticos de calibração (histogramas de rank SBC e diagramas de confiabilidade ECE) confirmam que o posterior está bem calibrado, com dimensões contínuas aproximando-se da linha de base de amostra finita e dimensões discretas correspondendo ao erro esperado de um classificador perfeitamente calibrado.

3. Simulador Biofísico de Hodgkin–Huxley:

   • Um simulador "caixa preta" para atividade neural com verossimilhanças intratáveis, envolvendo condutâncias contínuas e identidades de canais discretas (por exemplo, presença de canais iônicos específicos).
   • A MNPE infere com sucesso posterores conjuntos para condutâncias contínuas e tipos de canais categóricos.
   • O método produz posterores bem calibrados em diferentes orçamentos de treinamento, permitindo a análise de mecanismos compensatórios entre canais iônicos (por exemplo, a relação entre canais de cálcio do tipo L e condutância de potássio).

Significado e Alegações
O artigo alega que a MNPE preenche uma lacuna crítica na SBI ao fornecer um quadro unificado e amortizado para simuladores com espaços paramétricos mistos. Suas principais contribuições são:

• Generalidade: Suporta números arbitrários de parâmetros discretos e contínuos e é compatível com qualquer rede generativa contínua (fluxos, difusão).
• Eficiência: Diferentemente do MCMC, a MNPE amortiza o custo de simulação e treinamento, permitindo inferência rápida em novas observações através de uma única passagem direta.
• Calibração: Fornece um método principiado para avaliar calibração em configurações mistas, combinando métricas baseadas em rank e baseadas em classificação com linhas de base empíricas.
• Implementação: O método é implementado no pacote Python sbi publicamente disponível, tornando-o acessível a cientistas de domínio sem exigir ajuste extensivo.

Os autores notam limitações inerentes à família NPE, como a necessidade de grandes conjuntos de dados de treinamento para posterores complexos e suscetibilidade a má especificação do modelo. Eles também esclarecem que seu procedimento de calibração avalia calibração marginal por dimensão em vez de calibração conjunta completa, o que permanece desafiador para espaços mistos.