Re-mixing Embeddings for Patient Augmentation in Data Scarce Multiple Instance Learning

Resumen Técnico: Re-mezcla de Embeddings para la Aumentación de Pacientes en el Aprendizaje de Múltiples Instancias con Escasez de Datos

Planteamiento del Problema
La escasez de datos sigue siendo un cuello de botella crítico en el Aprendizaje de Múltiples Instancias (MIL, por sus siglas en inglés) médico, particularmente para enfermedades raras, modalidades especializadas costosas y escenarios donde la recolección de datos invasivos es limitada. En el MIL, una "bolsa" (por ejemplo, un paciente) consiste en múltiples instancias (por ejemplo, parches de imágenes, células individuales o características genómicas), y el modelo predice la etiqueta a nivel de bolsa sin anotaciones a nivel de instancia. Las técnicas de aumentación tradicionales (rotaciones, giros) operan a nivel de instancia y no logran generar variaciones realistas a nivel de bolsa. Además, las estrategias de mezcla a nivel de bolsa existentes (por ejemplo, PseMix, ReMix) sufren de tres limitaciones principales:

1. Operan exclusivamente de forma online, creando bolsas sintéticas efímeras que no pueden ser validadas ni reutilizadas.
2. Su alcance de mezcla local no logra capturar las estadísticas globales de la población.
3. Requieren ejemplos reales de cada clase, lo que las hace ineficaces cuando falta una clase entera (por ejemplo, controles sanos) del conjunto de datos.

Metodología: RECIPE
Los autores presentan RECIPE (Re-mixing Embeddings via Clustering for Patient augmEntation), un marco de trabajo offline y estadísticamente fundamentado para generar nuevas bolsas de pacientes en el espacio de embeddings. El flujo de trabajo consta de dos etapas principales:

1. Muestreo No Supervisado Basado en Clustering (Aprendizaje de la Receta):

   • Agrupación de Embeddings (Embedding Pooling): Los embeddings de las instancias de todos los pacientes de entrenamiento se agrupan.
   • Modelado de Mezclas Gaussianas (GMM): Se ajusta un GMM a los embeddings agrupados para definir clusters no supervisados.
   • Derivación de la Receta: Para cada clase de enfermedad, el método calcula una "receta" —una distribución estadística de instancias a través de los clusters aprendidos. Esto se define mediante la media ($\mu_{c,k}$) y la desviación estándar ($\sigma_{c,k}$) del número de instancias asignadas a cada cluster $k$ para la clase $c$.
   • Generación de Pacientes: Se generan nuevos pacientes muestreando el número de instancias para cada cluster basándose en la receta aprendida ($n_{c,k} \sim \mathcal{N}(\mu_{c,k}, \sigma^2_{c,k})$) y luego extrayendo esa cantidad de embeddings del grupo de instancias reales asignadas a esos clusters. Esto permite la generación de pacientes incluso cuando no existen ejemplos reales de una clase específica, siempre que la receta se derive de un conjunto de datos externo.

2. Selección de Pacientes Basada en la Incertidumbre:

   • Para asegurar que los datos generados mejoren el rendimiento del modelo en lugar de introducir ruido, el método emplea la cuantificación de la incertidumbre.
   • Un modelo MIL pre-entrenado evalúa las bolsas generadas utilizando Monte Carlo (MC) dropout durante la inferencia.
   • La Entropía Predictiva se calcula para cuantificar la incertidumbre ($U(x) = H[\bar{p}(x)]$).
   • Solo los pacientes generados más inciertos (aquellos con alta entropía) son seleccionados para aumentar el conjunto de entrenamiento, bajo la hipótesis de que estas muestras proporcionan la señal de aprendizaje más informativa.

Contribuciones Clave

• Generación Offline y Estadísticamente Fundamentada: A diferencia de los métodos de mezcla online, RECIPE aprende "recetas" distributivas globales para crear cohortes sintéticas persistentes y analizables.
• Gestión de Clases Ausentes: RECIPE es el primer marco de MIL capaz de generar clases de pacientes enteramente ausentes (por ejemplo, controles sanos) mediante la transferencia de recetas estadísticas de conjuntos de datos externos sin compartir datos sensibles de pacientes.
• Aplicabilidad Multimodal: El método se demuestra efectivo en diversas modalidades de datos, incluyendo imágenes histopatológicas, secuenciación de ARN de célula única (scRNA-seq) y citometría de flujo.
• Agnóstico a la Arquitectura: El marco es compatible con varios agregadores de MIL, incluyendo ABMIL (basado en atención), DSMIL y Transformers.

Resultados Experimentales
Los autores evaluaron RECIPE a través de tres escenarios de escasez clínicamente relevantes:

1. Transferencia entre Datasets (Clase Ausente):

   • Escenario: Se eliminaron todos los controles sanos del conjunto de datos AML-Hehr. Una "receta sana" se derivó del dataset externo cAItomorph y se aplicó al pool de embeddings de AML-Hehr.
   • Resultado: Los pacientes sanos generados coincidieron estrechamente con las composiciones reales de tipos celulares (Divergencia de Jensen-Shannon = 0.12). Entrenar con estos pacientes generados logró un rendimiento comparable al entrenamiento con el dataset completo (0.78 vs. 0.83 de precisión balanceada) y superó significativamente la transferencia directa de datos externos (0.66).

2. Regímenes de Bajos Datos:

   • Escenario: El número de controles sanos en AML-Hehr se redujo a tan solo $n=1$.
   • Resultado: Con un solo control sano real, el baseline alcanzó una precisión balanceada de 0.26. RECIPE logró 0.68, superando sustancialmente a MixUp (0.60) y PseMix (0.64).

3. Modalidades No de Imagen (scRNA-seq y Citometría de Flujo):

   • Escenario: Aumentación en datasets de severidad de COVID-19 con tamaños de muestra naturalmente limitados.
   • Resultado: RECIPE mejoró la precisión balanceada en un +14% (embeddings de PCA) y un +26% (embeddings de scVI) en datos de scRNA-seq, y en un +22% en datos de citometría de flujo, superando todas las estrategias de aumentación de base.

Significancia y Reivindicaciones
El artículo afirma que RECIPE ofrece una solución novedosa al cuello de botella de la escasez de datos en el MIL médico al permitir la generación de cohortes de pacientes realistas mediante la re-mezcla de embeddings agrupados basados en distribuciones estadísticas aprendidas. Su principal significancia radica en su capacidad para:

• Abordar el problema de la "clase ausente", permitiendo la síntesis de controles sanos o fenotipos de enfermedades raras utilizando estadísticas externas sin comprometer la privacidad del paciente.
• Mejorar el rendimiento del MIL en regímenes de datos extremadamente bajos donde la aumentación tradicional falla.
• Proporcionar un marco generalizable aplicable más allá de las imágenes a datos de célula única y citometría.

Los autores enfatizan que su enfoque va más allá de la mezcla local y online para modelar estadísticas poblacionales globales, permitiendo la creación de datos sintéticos validados y reutilizables que pueden cerrar brechas en la investigación de enfermedades raras y en la aumentación de datos preservando la privacidad.