From Point Estimates to Distributions: GMM Pooling for MIL in Preterm Birth Prediction

Resumen Técnico: De estimaciones puntuales a distribuciones: GMM Pooling para MIL en la predicción de parto pretérmino

Formulación del problema
El parto pretérmino (PTB, por sus siglas en inglés), definido como el parto antes de las 37 semanas de gestación, es una causa significativa de morbilidad y mortalidad neonatal. Aunque las guías clínicas recomiendan la adquisición de múltiples imágenes de ecografía transvaginal (TVUS) por paciente para compensar la variabilidad dependiente del operador en la medición de la longitud cervical, los enfoques actuales de aprendizaje automático para la predicción de PTB suelen depender de una única imagen seleccionada por paciente. Esta práctica descarta la variabilidad informativa intra-examen (por ejemplo, diferencias de vista y condiciones de adquisición) y limita el rendimiento predictivo.

Para abordar esto, los autores formulan la predicción de PTB como un problema de Aprendizaje de Instancias Múltiples (MIL, por sus siglas en inglés). En este marco, cada paciente corresponde a una "bolsa" de imágenes de TVUS de tamaño variable ($X = \{x_1, \dots, x_N\}$) asociada con una única etiqueta a nivel de bolsa ($Y$) que indica el resultado del parto (PTB vs. término). El desafío central es agregar estos números variables de características a nivel de instancia en una representación de dimensión fija, preservando la variabilidad informativa dentro de la bolsa, en lugar de colapsarla en una única estimación puntual.

Metodología: GMM Pooling
El artículo propone un módulo de GMM pooling (agregación mediante modelos de mezcla gaussiana) para reemplazar los agregadores estándar de MIL (como max, media o atención), que resumen una bolsa como una única estimación puntual en el espacio de características. El método propuesto modela explícitamente la distribución de características dentro de la bolsa para capturar la heterogeneidad. La arquitectura consta de tres etapas:

1. Extracción de características: Un codificador compartido (implementado como una arquitectura U-Net para la tarea de PTB) mapea cada imagen a un vector de características de instancia. El codificador se entrena conjuntamente para la segmentación del canal cervical (como un objetivo auxiliar) y la clasificación a nivel de bolsa para aprender representaciones con significado anatómico.
   2.Estimación de parámetros de GMM: En lugar de un promedio simple, el método estima una GMM sobre las características de las instancias. Esto implica:
   • Asignaciones suaves de instancia a componente: Un MLP ligero ($f_r$) genera una distribución de probabilidad ($r_{i,k}$) sobre $M$ componentes de la mezcla para cada instancia.
   • Ponderación de importancia de la instancia: Una red separada ($f_w$) calcula los pesos de las instancias ($w_i$) mediante softmax.
   • Combinación de atención: El peso de atención final $a_{i,k}$ combina la importancia de la instancia y la responsabilidad del componente ($a_{i,k} = w_i \cdot r_{i,k}$).
   • Aprendizaje de parámetros: Utilizando estos pesos de atención, se estiman de extremo a extremo los coeficientes de mezcla ($\pi_k$), las medias ($\mu_k$) y las varianzas diagonales ($\sigma^2_k$) de la GMM. Los autores optan por la covarianza diagonal para asegurar la estabilidad computacional y la escalabilidad.
2. Incrustación de la distribución (Distribution Embedding): Para producir una representación de bolsa de longitud fija, el método evalúa la densidad logarítmica de la GMM aprendida en un conjunto de $P$ vectores de sonda (probes) aprendibles. Los valores de densidad logarítmica resultantes forman la incrustación final de la bolsa utilizada para la predicción.

Contribuciones clave
Los autores identifican tres contribaciones principales:

1. Módulo de GMM Pooling: Un nuevo mecanismo de agregación de MIL que preserva las distribuciones de características informativas dentro de la bolsa, abordando el cuello de botella de información inherente al pooling de estimación puntual.
2. Incrustación de distribución diferenciable: Un método para construir una incrustación de longitud fija mediante la evaluación de la densidad logarítmica en vectores de sonda aprendibles, permitiendo que la representación de la distribución se aprenda de extremo a extremo.
3. Evaluación empírica exhaustiva: Una evaluación rigurosa en una cohorte privada de alto riesgo de PTB y en un benchmark público de metástasis de ganglios linfáticos, demostrando que el MIL mejora respecto a la predicción a nivel de instancia y que el GMM pooling ofrece un rendimiento competitivo en tareas de clasificación y regresión.

Resultados experimentales
El método fue evaluado en dos conjuntos de datos:

• Cohorte privada de PTB: 182 pacientes embarazadas de alto riesgo (44 casos de PTB) con 1 a 43 imágenes por paciente.
  • El MIL mejoró significativamente respecto a los modelos basados en instancias, elevando el PR-AUC de 0.44 a 0.56 (max pooling) y de 0.56 (GMM pooling).
  • El GMM pooling logró un PR-AUC de $0.56 \pm 0.03$, comparable al de max ($0.57 \pm 0.05$) y atención ($0.56 \pm 0.02$). Sin embargo, el GMM exhibió una desviación estándar menor entre ejecuciones, lo que sugiere una mayor estabilidad.
• Benchmark público de metástasis de ganglios linfáticos: 933 muestras de entrenamiento, 668 de validación y 736 de prueba.
  • El GMM pooling alcanzó un rendimiento de vanguardia con un F1-score de $0.91 \pm 0.01$ y un ROC-AUC de $0.89 \pm 0.01$ para clasificación.
  • Para regresión (porcentaje de píxeles metastásicos), logró un Error Absoluto Medio (MAE) de $0.18 \pm 0.02$, competitivo con los métodos de mejor rendimiento.

Significancia y afirmaciones
El artículo afirma que modelar las distribuciones de características mediante GMM pooling proporciona una alternativa robusta al pooling estándar de estimación puntual. Mientras que los agregadores estándar colapsan la variabilidad dentro de la bolsa, el GMM pooling captura la estructura de la distribución conjunta a través de las dimensiones de las características. Los autores señalan que, aunque el enfoque introduce hiperparámetros específicos de la distribución (número de componentes, sondas, temperatura) que requieren un ajuste dependiente del conjunto de datos, los resultados en dos tareas distintas de imágenes médicas sugieren que el modelado de distribuciones se generaliza bien. El trabajo destaca que aprovechar todas las imágenes disponibles por paciente mediante MIL es crucial para la predicción de PTB, particularmente en cohortes de alto riesgo donde la separabilidad de clases es baja.