Synthetic Data Generation for Long-Tail Medical Image Classification: A Case Study in Skin Lesions

Resumen Técnico: Generación de Datos Sintéticos para la Clasificación de Imágenes Médicas de Cola Larga

Enunciado del Problema
Las distribuciones de clases de cola larga son omnipresentes en conjuntos de datos médicos multiclase, planteando desafíos significativos para los modelos de aprendizaje profundo. En estos escenarios, unas pocas clases "cabeza" dominantes contienen la mayoría de las muestras, mientras que numerosas clases "cola" clínicamente críticas (que a menudo corresponden a enfermedades raras y de alto riesgo) están severamente subrepresentadas. Los modelos entrenados con datos tan sesgados suelen lograr una precisión general competitiva, pero fallan al generalizar a las clases de cola. Las soluciones existentes, categorizadas en diseño de funciones de pérdida, mejoras en la arquitectura del modelo y reequilibrio de clases, ofrecen solo mejoras marginales. Estos métodos a menudo dependen de un ajuste meticuloso de hiperparámetros, muestran una generalización limitada entre diferentes modelos y tienen dificultades para escalar debido a la naturaleza determinista de la augmentación de datos manual. Además, aunque los modelos de difusión han demostrado éxito en la generación de imágenes naturales, su aplicación a la imagenología médica específica de dominio y de cola larga sigue siendo poco explorada, particularmente en lo que respecta a la generación de muestras diversas y clínicamente significativas.

Metodología
Los autores proponen una tubería de augmentación de datos sintéticos impulsada por modelos de difusión, adaptada para la Clasificación de Imágenes Médicas de Cola Larga (LTMIC), demostrada específicamente en la clasificación de lesiones cutáneas. El enfoque consta de tres etapas principales:

1. Arquitectura del Modelo de Difusión de Inpainting:
   El núcleo del método es un novedoso modelo de difusión de inpainting. A diferencia de los modelos de incrustación de texto estándar, que no son adecuados para descripciones médicas altamente especializadas, los autores entrenan una red simple de incrustación de clases totalmente conectada. El modelo toma cuatro entradas: la imagen original de dermatoscopia ($x$), una máscara binaria de área de interés ($b$), el fondo de la imagen de dermatoscopia ($d$) y la etiqueta de clase ($y$). La imagen y el fondo se codifican en un espacio latente mediante un Autoencoder Variacional (VAE) preentrenado. La máscara binaria se reescala para coincidir con las dimensiones latentes. Estos componentes se concatenan y se introducen en una UNet parametrizada por $\theta_d$ para estimar el ruido añadido durante el proceso de difusión hacia adelante. Una red de incrustación de clases parametrizada por $\theta_c$ proporciona la condición de la etiqueta.

2. Ajuste Fino con Selección de Muestras:
   Para asegurar que el modelo de difusión aprenda la distribución de cada clase de manera igualmente efectiva, se emplea una estrategia de sobre-muestreo durante el ajuste fino. Los autores aprovechan hallazgos de trabajos anteriores que sugieren que las muestras "aproximadamente dentro de la distribución" (AID) (aquellas ligeramente alejadas del centro de la distribución) son las más importantes para entrenar modelos generativos. En consecuencia, se utiliza un método de detección de anomalías médicas auto-supervisado para calcular puntuaciones de anomalía para cada muestra. Se asignan pesos basados en la distancia de una muestra respecto a la puntuación de anomalía mediana de su clase; las muestras más cercanas a la región AID reciben pesos más altos y se repiten para equilibrar el conjunto de datos de ajuste fino entre todas las clases.

3. Generación Sintética y Filtrado OOD:
   El modelo ajustado fino genera imágenes sintéticas tomando las imágenes de fondo y las etiquetas de clase del conjunto de datos de entrenamiento original como entrada. Debido a la naturaleza estocástica de los modelos de difusión, el conjunto de datos sintético inicial ($D_{syn}$) contiene muestras "ruidosas", incluidas imágenes Fuera de la Distribución (OOD) que pueden no ser clínicamente válidas. Para abordar esto, se aplica un mecanismo de filtrado médico OOD a $D_{syn}$ para producir un conjunto de datos sintético limpio ($D_{clean}^{syn}$). El conjunto de datos de entrenamiento final se forma combinando el conjunto de datos de cola larga original ($D_{LT}$) con $D_{clean}^{syn}$. Un hiperparámetro $\gamma$ controla el porcentaje de muestras limpias retenidas del conjunto sintético para cada clase.

Contribuciones Clave
El artículo describe tres contribuciones principales:

1. Tubería Novel: Una tubería de generación de datos sintéticos que integra un modelo de difusión de inpainting con filtrado OOD, diseñada específicamente para abordar el problema de LTMIC.
2. Arquitectura Especializada: Una nueva arquitectura de modelo de difusión de inpainting adaptada para la generación de imágenes médicas, que utiliza una red de incrustación de clases dedicada y condicionamiento de fondo en lugar de incrustaciones de texto genéricas.
3. Estrategia de Muestreo: Una estrategia de muestreo novedosa basada en OOD para el ajuste fino de modelos de difusión en conjuntos de datos médicos de cola larga, que prioriza las muestras AID para mejorar la calidad de los datos generados para las clases de cola.

Resultados Experimentales
El método se evaluó en el conjunto de datos de clasificación de lesiones cutáneas ISIC2019, un referente que contiene 25.331 imágenes en 8 clases con un desequilibrio severo (que va desde 12.875 muestras para Nevus hasta solo 62 para Lesión Vascular). Los experimentos utilizaron validación cruzada de 5 pliegues y un clasificador simple basado en CNN (sin ensamblaje ni datos externos) para aislar el impacto del método de generación de datos.

• Rendimiento General: El enfoque propuesto logró una Precisión Multiclase Balanceada (BMA) de 0.802, superando significativamente a los métodos actuales del Estado del Arte (SOTA), incluidos los ensamblajes de múltiples modelos (MRE) y otros enfoques basados en difusión (por ejemplo, CBDM, Dreambooth). Cabe destacar que superó al método MRE en aproximadamente un 4% en BMA sin requerir datos externos.
• Mejora en Clases de Cola: Las ganancias más significativas se observaron en las clases de cola. Para la clase con menos muestras (Dermatofibroma, DF), el método produjo una mejora en la precisión de más del 28% en comparación con la línea base sin datos sintéticos.
• Impacto del Filtrado OOD: El estudio encontró que el filtrado OOD es crucial para equilibrar el rendimiento. Si bien los datos sintéticos sin filtrado mejoraron significativamente las clases de cola, introdujeron ruido que degradó ligeramente el rendimiento de las clases cabeza. El proceso de filtrado restauró el rendimiento de las clases cabeza mientras mantenía las ganancias en las clases de cola.
• Sensibilidad a Hiperparámetros: Un estudio de ablación sobre $\gamma$ (la tasa de retención de muestras sintéticas limpias) reveló que el rendimiento mejora bruscamente incluso con una pequeña cantidad de datos sintéticos, pero se estabiliza o disminuye a medida que $\gamma$ se acerca a 1 (sin filtrado). Se identificó el rango óptimo para $\gamma$ entre 0.2 y 0.6.

Significado y Afirmaciones
Los autores afirman que sus resultados demuestran la efectividad de la augmentación basada en difusión para mitigar el desequilibrio de cola larga y mejorar la robustez de la clasificación médica. Al generar muestras sintéticas diversas, realistas y clínicamente significativas, el método aborda la escasez de datos de enfermedades raras sin depender de ensamblajes de modelos complejos o funciones de pérdida especializadas que son sensibles a los hiperparámetros. El artículo concluye que, aunque el enfoque actual trata la generación de datos y la clasificación como procesos independientes, el trabajo futuro se centrará en la optimización conjunta de estos componentes para mejorar aún más el rendimiento de LTMIC. El método se presenta como agnóstico al modelo, capaz de integrarse con cualquier arquitectura de clasificación.