Simulation-based inference for rapid Bayesian parameter estimation in epidemiological models: a comparison with MCMC

Resumen Técnico: Inferencia Basada en Simulación para la Estimación Bayesiana Rápida de Parámetros en Modelos Epidemiológicos

Planteamiento del Problema
Los modelos epidemiológicos mecanísticos son esenciales para pronosticar la dinámica de enfermedades infecciosas y guiar las decisiones de salud pública. Sin embargo, la calibración de estos modelos con datos de vigilancia mediante la inferencia bayesiana estándar, específicamente el método de Monte Carlo por Cadenas de Markov (MCMC), es computacionalmente prohibitiva para sistemas no lineales de alta dimensión. El MCMC requiere miles o millones de evaluaciones de la verosimilitud, lo que a menudo toma días o meses para análisis detallados. Esta latencia impide la recalibración repetitiva y casi en tiempo real necesaria cuando las condiciones epidémicas cambian rápidamente, como ocurrió durante la pandemia de COVID-19. Los autores investigan si la Inferencia Basada en Simulación (SBI, por sus siglas en inglés), específicamente la Estimación de la Posterior Neuronal (NPE), puede servir como una alternativa escalable y computacionalmente eficiente para la calibración bayesiana sin sacrificar la fidelidad de la distribución posterior.

Metodología
El estudio utiliza un modelo compartimental mecanístico SECIR (Susceptibles-Expuestos-Infectados-Sintomáticos-Graves-Críticos-Recuperados-Fallecidos) implementado en el marco de trabajo MEmilio para simular la dinámica del COVID-19 en Alemania durante 2020. El modelo se ajusta a los datos de ocupación de Unidades de Cuidados Intensivos (UCI), elegidos por su relevancia directa para la carga sanitaria y su menor sensibilidad a la variabilidad de las pruebas en comparación con los recuentos de casos.

• Objetivos de Inferencia: El modelo bayesiano estima cinco parámetros de progresión de la enfermedad (tiempos de transición y probabilidades) y una tasa de contacto efectiva variable en el tiempo representada por una cuadrícula de factores de reducción ($r_k$).
• Método de Referencia (MCMC): Se utiliza un algoritmo de evolución diferencial de Metropolis-Z (DE-MZ) para generar distribuciones posteriores de referencia. Este método se ejecuta enteramente en hardware de CPU y depende de evaluaciones explícitas de la verosimilitud utilizando una distribución t de Student para los errores de observación.
• Método Propuesto (SBI/NPE): Los autores emplean la Estimación de la Posterior Neuronal Amortizada.
  1. Entrenamiento: Un estimador de densidad neuronal (Flujo Autorregresivo Enmascarado, MAF) se entrena en un gran conjunto de datos de pares parámetro-trayectoria simulados $(\theta, \hat{X})$. Los parámetros se muestrean de distribuciones a priori y el simulador SECIR genera trayectorias sintéticas de UCI.
  2. Arquitectura: Una Red Neuronal Convolucional (CNN) incrusta la trayectoria de la serie temporal en una representación de características de menor dimensión, la cual condiciona al MAF para aproximar la posterior $p(\theta | X)$.
  3. Inferencia: Una vez entrenada, la red genera rápidamente muestras posteriores para los datos observados sin evaluaciones adicionales de la verosimilitud.
• Evaluación: El estudio compara SBI y MCMC a través de tres ventanas de inferencia de 31 días (que representan distintas fases epidémicas) y una ventana de reconstrucción desafiante de 201 días que involucra múltiples puntos de cambio en la transmisión. La concordancia se cuantifica mediante la distancia 1-Wasserstein y la divergencia de Kullback–Leibler (KL) entre las posteriores marginales, además de comprobaciones predictivas posteriores (RMSE e intervalos de credibilidad).

Contribuciones Clave

1. Marco Comparativo: El artículo proporciona una comparación rigurosa, cara a cara, entre SBI y MCMC para un modelo epidemiológico complejo y de alta dimensión, yendo más allá de la simple precisión predictiva para evaluar la concordancia estructural de las distribuciones posteriores inferidas.
2. Demostración de Escalabilidad: Los autores demuestran que la SBI puede reducir el tiempo de ejecución de la inferencia aproximadamente un orden de magnitud para ventanas cortas y más de dos órdenes de magnitud para ventanas largas y de alta dimensión en comparación con el MCMC basado en CPU.
3. Análisis de Sensibilidad de la Prior: El estudio destaca la importancia crítica de la especificación de la distribución a priori en la SBI. Para el problema de 201 días, las distribuciones a priori amplias sobre los factores de reducción variables en el tiempo condujeron a dinámicas epidémicas implausibles en los datos de entrenamiento, lo que perjudicó el aprendizaje. Los autores muestran que estrechar los límites de la a priori basándose en comprobaciones predictivas a priori es esencial para una SBI exitosa en entornos de alta dimensión.
4. Optimización de Hiperparámetros: El trabajo identifica que las configuraciones óptimas de SBI (presupuesto de simulación y tamaño de lote) difieren según el horizonte de inferencia y la dimensionalidad, lo que sugiere que la selección de hiperparámetros no puede transferirse ciegamente entre ventanas de inferencia cortas y largas.

Resultados

• Concordancia de la Posterior: En todos los entornos, la SBI recuperó distribuciones posteriores ampliamente consistentes con la referencia de MCMC. Aunque las marginales de la SBI fueron generalmente más suaves y ligeramente más amplias (lo que indica mayor incertidumbre), los modos dominantes de la posterior y las tendencias de los parámetros se preservaron.
  • Para las ventanas de 31 días, la distancia de Wasserstein entre las posteriores de SBI y MCMC osciló entre 0.63 y 2.47, con divergencias KL simétricas generalmente por debajo de 100.
  • Para la ventana de 201 días, a pesar del aumento en la dimensionalidad y la incertidumbre, la SBI preservó la estructura principal de la posterior, con una distancia de Wasserstein de 0.35 y una KL simétrica de 72.14.
• Desempeño Predictivo: La SBI reprodujo con precisión las trayectorias observadas de UCI. Para las ventanas de 31 días, los valores de RMSE fueron comparables a los de MCMC (por ejemplo, 66.41 frente a 31.8 para el desfase 0). En el entorno de 201 días, el RMSE de la SBI fue de 325.72 frente a los 23.1 de MCMC; sin embargo, ambos métodos capturaron las dinámicas temporales dominantes, incluyendo el descenso tras la primera ola y el aumento de otoño. El error más alto en la ventana larga se atribuye a la dificultad de reconstruir múltiples puntos de cambio simultáneamente.
• Eficiencia de Tiempo de Ejecución:
  • Ventanas de 31 días: MCMC requirió ~1,077 segundos por ejecución. La SBI requirió 60–71 segundos en una única GPU NVIDIA A100.
  • Ventana de 201 días: MCMC requirió ~19,102 segundos y no logró alcanzar la convergencia para todos los parámetros (máximo $\hat{R} = 1.605$). La SBI requirió ~157 segundos y mantuvo una amplia concordancia con la estructura de la posterior de MCMC.
• Convergencia: El MCMC tuvo dificultades para converger en el entorno de alta dimensión de 201 días, mientras que la SBI proporcionó estimaciones de la posterior estables en una fracción de ese tiempo.

Significancia y Reivindicaciones
El artículo afirma que la SBI con NPE proporciona un marco rápido y computacionalmente eficiente para la calibración bayesiana de modelos epidemiológicos mecanísticos. La principal significancia radica en permitir la inferencia repetitiva casi en tiempo real y el análisis rápido de brotes, los cuales son críticos para la toma de decisiones en salud pública pero actualmente están limitados por el costo computacional del MCMC.

Los autores señalan modestamente que, aunque la SBI preserva la estructura dominante de la posterior, puede producir estimaciones de incertidumbre más amplias que el MCMC, potencialmente debido a la falta de un modelado explícito del ruido de observación en la fase de entrenamiento. Enfatizan que la SBI es un método amortizado: sus ventajas son más pronunciadas en entornos que requieren inferencia repetida bajo el mismo modelo y especificación de la a priori, donde los costos iniciales de simulación y entrenamiento se distribuyen a través de múltiples análisis. El estudio concluye que la SBI es un enfoque prometedor para la calibración rápida, siempre que se integren procesos rigurosos de comprobación predictiva a priori y validación de la posterior para asegurar la identificabilidad de los parámetros y la plausibilidad de las trayectorias simuladas.