Mixed neural posterior estimation for simulators with discrete and continuous parameters

Resumen Técnico: Estimación Posterior Neural Mixta para Simuladores con Parámetros Discretos y Continuos

Planteamiento del Problema
La Estimación Posterior Neural (NPE) se ha convertido en una herramienta estándar para la Inferencia Basada en Simulación (SBI), permitiendo una inferencia de parámetros rápida y amortizada para simuladores complejos con verosimilitudes intratables. Sin embargo, los métodos estándar de NPE asumen espacios de parámetros continuos, modelando típicamente la posterior como una densidad de probabilidad continua mediante flujos normalizadores o modelos de difusión. Esta suposición falla para una amplia clase de simuladores científicos donde el espacio de parámetros es mixto, conteniendo tanto dimensiones continuas (por ejemplo, tasas, conductancias) como dimensiones discretas (por ejemplo, variables categóricas, índices de selección de modelos, ubicaciones de puntos de cambio).

Los enfoques existentes para la inferencia mixta enfrentan limitaciones significativas. Los métodos estándar de Cadena de Markov Monte Carlo (MCMC) requieren un ajuste cuidadoso, no amortizan los costos de inferencia entre observaciones y a menudo luchan con variables discretas (por ejemplo, requiriendo marginalización analítica o dinámicas de Hamiltoniano discontinuas especializadas). Los métodos recientes de SBI basados en difusión intentan incrustar parámetros discretos en espacios continuos o tratarlos como índices de modelo, pero estos pueden ser computacionalmente pesados o requerir incrustaciones riemannianas complejas. En consecuencia, los simuladores con espacios de parámetros mixtos permanecen insuficientemente respaldados por las cajas de herramientas actuales de SBI.

Metodología: Estimación Posterior Neural Mixta (MNPE)
Los autores proponen la Estimación Posterior Neural Mixta (MNPE), una extensión de NPE que maneja conjuntamente parámetros discretos y continuos a través de una red de inferencia factorizada.

1. Estrategia de Factorización:
   La posterior conjunta $p(\theta_d, \theta_c | x)$ sobre parámetros discretos $\theta_d$ y parámetros continuos $\theta_c$ se factoriza como:
   $$p(\theta_d, \theta_c | x) = p(\theta_d | x) \cdot p(\theta_c | \theta_d, x)$$
   Esto separa el problema en dos sub-tareas: clasificación sobre el espacio discreto y estimación de densidad condicional en el espacio continuo, condicionada al estado discreto.

2. Arquitectura de la Red:

   • Componente Discreto: El término $p(\theta_d | x)$ se modela utilizando un Estimador de Densidad Autoregresivo Enmascarado (MADE). Esta red autoregresiva genera probabilidades de clase para cada dimensión discreta, condicionada a las dimensiones discretas precedentes y a la observación $x$. Esto permite manejar números arbitrarios de parámetros discretos con conteos de clases variables.
   • Componente Continuo: El término $p(\theta_c | \theta_d, x)$ se modela utilizando un modelo generativo condicional estándar (por ejemplo, un flujo de spline neuronal o un modelo de difusión) condicionado tanto a la observación $x$ como a los parámetros discretos $\theta_d$.
   • Objetivo de Entrenamiento: La red se entrena conjuntamente para minimizar el logaritmo negativo de la probabilidad de los parámetros mixtos:
     $$\mathcal{L} = -\log q(\theta_d, \theta_c | x) = -\log q(\theta_c | \theta_d, x) - \log q(\theta_d | x)$$
     La parte continua utiliza la verosimilitud estándar del flujo, mientras que la parte discreta se reduce a una pérdida de entropía cruzada.

3. Diagnósticos de Calibración:
   Para evaluar la calidad de las posteriores mixtas sin soluciones de referencia, los autores proponen un marco diagnóstico combinado:

   • Dimensiones Continuas: Utiliza Calibración Basada en Simulación (SBC) mediante estadísticas de rango para verificar la uniformidad en las distribuciones marginales de rango.
   • Dimensiones Discretas: Utiliza el Error de Calibración Esperado (ECE) y diagramas de fiabilidad, tratando la inferencia discreta como un problema de clasificación.
   • Líneas Base: Los autores derivan líneas base empíricas de muestras finitas para ambas métricas para distinguir la mala calibración real del ruido estadístico, reconociendo que la calibración perfecta es inalcanzable con muestras finitas.

Resultados Clave
Los autores validan MNPE en tres simuladores de complejidad creciente:

1. Ejemplo de Juguete Gaussiano Tractable:

   • Un simulador con una variable latente continua y una variable latente discreta acopladas a través de una observación compartida.
   • MNPE recupera con precisión la posterior analítica de verdad fundamental, incluidos los regímenes bimodales donde el estado discreto es incierto.
   • La puntuación de la Prueba de Dos Muestras del Clasificador (C2ST) converge al nivel de azar (~0.5) con aproximadamente 1.000 simulaciones, indicando indistinguibilidad de la posterior verdadera.

2. Simulador de Cola en Tándem:

   • Un simulador con tres tasas continuas y dos conteos de servidores discretos ($|D|=25$).
   • MNPE se compara con referencias MCMC (PyMC y NumPyro) donde las dimensiones discretas se marginalizan analíticamente.
   • Con 100.000 simulaciones, las posteriores de MNPE coinciden con la referencia MCMC tanto en dimensiones continuas como discretas.
   • Los diagnósticos de calibración (histogramas de rango SBC y diagramas de fiabilidad ECE) confirman que la posterior está bien calibrada, con dimensiones continuas acercándose a la línea base de muestras finitas y dimensiones discretas coincidiendo con el error esperado de un clasificador perfectamente calibrado.

3. Simulador Biofísico Hodgkin–Huxley:

   • Un simulador de "caja negra" para la actividad neuronal con verosimilitudes intratables, que involucra conductancias continuas e identidades de canales discretas (por ejemplo, presencia de canales iónicos específicos).
   • MNPE infiere con éxito posteriores conjuntas para conductancias continuas y tipos de canales categóricos.
   • El método produce posteriores bien calibradas a través de presupuestos de entrenamiento, permitiendo el análisis de mecanismos compensatorios entre canales iónicos (por ejemplo, la relación entre los canales de calcio tipo L y la conductancia de potasio).

Significado y Afirmaciones
El artículo afirma que MNPE llena una brecha crítica en SBI al proporcionar un marco unificado y amortizado para simuladores con espacios de parámetros mixtos. Sus contribuciones principales son:

• Generalizabilidad: Soporta números arbitrarios de parámetros discretos y continuos y es compatible con cualquier red generativa continua (flujos, difusión).
• Eficiencia: A diferencia de MCMC, MNPE amortiza el costo de la simulación y el entrenamiento, permitiendo una inferencia rápida en nuevas observaciones mediante una sola pasada hacia adelante.
• Calibración: Proporciona un método principiado para evaluar la calibración en configuraciones mixtas, combinando métricas basadas en rangos y clasificación con líneas base empíricas.
• Implementación: El método está implementado en el paquete Python sbi disponible públicamente, haciéndolo accesible a científicos de dominio sin requerir un ajuste extenso.

Los autores señalan limitaciones inherentes a la familia NPE, como la necesidad de grandes conjuntos de datos de entrenamiento para posteriores complejas y la susceptibilidad a la especificación incorrecta del modelo. También aclaran que su procedimiento de calibración evalúa la calibración marginal por dimensión en lugar de la calibración conjunta completa, lo cual sigue siendo un desafío para espacios mixtos.