Spatio-temporal stochastic graph-based learning for infectious disease forecasting

Resumen Técnico: Aprendizaje Basado en Grafos Estocásticos Espacio-Temporales para la Previsión de Enfermedades Infecciosas

Planteamiento del Problema
Los modelos existentes basados en grafos espacio-temporales para la previsión de enfermedades infecciosas (por ejemplo, para el COVID-19 y la varicela) han investigado insuficientemente la integración del modelado estocástico en sus procesos de aprendizaje. Además, la mayoría de los trabajos previos se centran en redes geográficas de tamaño pequeño a mediano (por ejemplo, 20–400 regiones) o exhiben una variabilidad limitada en sus trayectorias de prueba, las cuales a menudo no logran reflejar las fluctuaciones complejas, heterogéneas y de alta frecuencia observadas en la propagación de enfermedades en el mundo real a través de naciones extensas. En consecuencia, no está claro si las arquitecturas actuales pueden pronosticar con precisión las trayectorias de enfermedades en redes epidemiológicas de gran escala y altamente variables.

Metodología
Los autores proponen una Arquitectura de Aprendizaje Basada en Grafos Estocásticos Espacio-Temporales diseñada para pronosticar nuevos casos de enfermedades infecciosas mediante la integración de la formulación estocástica y la aproximación de la incertidumbre. La arquitectura consta de tres módulos secuenciales:

1. Codificación de Secuencias Multivariantes: Este módulo estandariza los datos temporales de incidencia de enfermedades y las conexiones geográficas. Crea muestras de secuencias con retardo ($X(t)$) donde los datos temporales se escalan al intervalo $[0, 1]$ y se organizan en un dominio de grafo espacial estático. La entrada incluye un tensor espacio-temporal que representa los casos de la enfermedad a través de $G_l$ ubicaciones durante una ventana $W$, junto con tensores de índice de aristas y atributos que representan la topología geográfica.
2. Procesamiento Espacio-Temporal: Este módulo extrae características latentes utilizando dos variantes propuestas:
   • PrGCNrmsle: Basado en una capa de Red Convolucional de Grafos (GCN). Agrega información de los nodos vecinos utilizando normalización simétrica para aprender correlaciones espaciales invariantes a la densidad de conectividad local.
   • PrGLSTMrmsle: Extiende la GCN con una capa de Memoria de Corto Plazo de Longitud de la LSTM (LSTM). Esta versión procesa las salidas de la GCN para capturar dependencias temporales y la evolución del estado oculto, filtrando el ruido de corto plazo mientras retiene la memoria de la dinámica de múltiples olas.
3. Método Estocástico: Una capa totalmente conectada (MLP) mapea las características latentes a un descriptor de 2 dimensiones por ubicación, que representa la media ($\mu$) y el logaritmo de la varianza ($\ln \sigma$) de una distribución normal. Utilizando un truco de reparametrización, el modelo muestrea de esta distribución para generar predicciones estocásticas ($\hat{y} = \mu + \sigma \odot \epsilon$). Esto induce perturbaciones mínimas para simular la variabilidad interna sin reintroducir prioris estructurales que pudieran suavizar excesivamente la representación.

Entrenamiento y Evaluación

• Función de Pérdida: Los modelos se entrenan de extremo a extremo utilizando el Error Cuadrático Medio Logarítmico (RMSLE). Esta elección penaliza más fuertemente la subestimación de los casos de enfermedad que la sobreestimación y trata los errores relativos de manera equitativa entre ubicaciones de diferentes tamaños, enfocando el modelo en las características dinámicas globales en lugar de las fluctuaciones locales.
• Estimación de la Incertidumbre: La incertidumbre epistémica se aproxima mediante un método de ensamble. El modelo se entrena 100 veces con diferentes inicializaciones (valores de semilla) para generar una distribución de predicciones, a partir de la cual se deriva un intervalo de variabilidad del 95%.
• Conjuntos de Datos:
  • COVID-19 (EE. UU.): 3,218 condados de los EE. UU. (2020–2022), cubriendo los periodos de entrenamiento (enero 2020–diciembre 2021) y prueba (diciembre 2021–octubre 2022) con patrones de olas altamente variables.
  • Varicela (Hungría): 20 condados (2005–2014), cubriendo los periodos de entrenamiento (2005–2012) y prueba (2012–2014) con patrones de olas estacionales más consistentes.
• Líneas Base (Baselines): Los modelos propuestos fueron comparados con cuatro modelos basados en grafos espacio-temporales (EvolveGCNH, GCLSTM, GConvLSTM, MPNNLSTM) y varias variantes de ablación (eliminación de la estocasticidad, cambio de funciones de pérdida o eliminación de módulos espaciales).

Resultos Clave

• Desempeño: El modelo propuesto PrGCNrmsle logró un desempeño semanal general competitivo en todos los 3,218 condados de EE. UU. y los 20 condados de Hungría. Demostró una capacidad mejorada para pronosticar la primera ola de COVID-19 de 2022 en los EE. UU. y las olas comparativas de la varicela en Hungría (2012–2014).
• Estabilidad: El enfoque propuesto exhibió una menor sensibilidad a la variabilidad de alta frecuencia y baja amplitud en comparación con las líneas base. Mostró intervalos de variabilidad más bajos en sus predicciones a través de 100 ejecuciones, lo que indica una mayor estabilidad.
• Limitaciones Observadas:
  • Retraso Temporal: Se observó un retraso constante de un paso (una semana) en las predicciones respecto a las observaciones. El modelo tiende a seguir la dirección de la tendencia pero presenta un desfase al capturar cambios abruptos o el momento de los picos.
  • Estimación de Picos: Aunque PrGCNrmsle aproximó bien la evolución temporal general, en ocasiones sobreestimó los picos para olas menores o subestimó los picos rápidos en datos de alta frecuencia.
  • Variante LSTM: La variante PrGLSTMrmsle (GCN + LSTM) mostró un aumento en el error durante las predicciones de los picos invernales para el COVID-19 y requirió significativamente más tiempo de entrenamiento (~4.5x) en comparación con PrGCNrmsle, lo que sugiere que la capa LSTM puede estar sobreajustándose a las olas de entrenamiento más pequeñas o teniendo dificultades con los picos más grandes y no vistos.
• Eficiencia Computacional: PrGCNrmsle demostró una convergencia más rápida y menos parámetros entrenables en comparación con las líneas base basadas en RNN, lo que lo hace más eficiente computacionalmente para redes a gran escala.

Significancia y Reivindicaciones
El artículo sostiene que la arquitectura propuesta logra adaptarse para codificar tanto redes geográficas de población grande (condados de EE. UU.) como pequeñas (condados de Hungría) dentro de una única estructura de modelo. Al incorporar una formulación estocástica, el modelo representa la progresión epidémica general en relación con las líneas base mientras mantiene la integridad estructural y minimiza la propagación de la incertidumbre.

Los autores posicionan este trabajo como un paso hacia un pronóstico de epidemias más robusto en entornos heterogéneos. Argumentan que, si bien el modelo captura las tendencias dominantes y proporciona un monitoreo útil de las trayectorias de las enfermedades, su aplicación a las políticas de salud pública requiere mejoras metodológicas adicionales para abordar el retraso de un paso y gestionar mejor los cambios rápidos y no lineales en la incidencia. El estudio destaca que los modelos actuales basados en grafos a menudo luchan con el desajuste entre las trayectorias de entrenamiento (patrones consistentes) y las trayectorias de prueba (patrones altamente variables), y que su enfoque estocástico ofrece un mecanismo para representar mejor el rango de posibilidades en la propagación de enfermedades sin sacrificar la eficiencia computacional.