Steering Diffusion Models via Class-Contrastive Influence for Few-Shot Medical Classification

Resumen Técnico: Dirigiendo Modelos de Difusión mediante la Influencia de Contraste de Clase para la Clasificación Médica de Pocos Disparos (Few-Shot)

Planteamiento del Problema
En la clasificación de imágenes médicas de pocos disparos (few-shot), los datos etiquetados son escasos, lo que hace que el aumento de datos sea esencial. Aunque los modelos de difusión comerciales pueden generar muestras de entrenamiento sintéticas, los enfoques actuales se centran principalmente en mejorar el realismo, la diversidad o la adaptación de dominio de estas imágenes. Los autores argumentan que estas métricas no garantizan la utilidad para la tarea de clasificación posterior. Existe una discrepancia significativa donde las imágenes generadas con una calidad visual comparable producen desempeños de clasificación muy distintos. El problema central identificado es la falta de un criterio basado en principios para medir y optimizar la "utilidad de la tarea" de las muestras sintéticas. El papel postula que la limitación central de las prácticas actuales no es solo la calidad de la generación, sino el fraculo de generar muestras que ayuden específicamente al clasificador a refinar su frontera de decisión.

Metodología
La solución propuesta, Influencia de Contraste de Clase (C2I), desplaza el objetivo de optimización de la calidad de la imagen hacia la influencia basada en gradientes de las muestras sintéticas sobre el clasificador.

1. Criterio de Influencia de Contraste de Clase (C2I):

   • El método aprovecha la estimación de la influencia basada en gradientes para cuantificar cómo una muestra de entrenamiento afecta la pérdida de validación.
   • Los autores observan que las puntuaciones de influencia crudas y ajenas a la clase no se correlacionan con el desempeño de la clasificación porque los conjuntos de validación contienen señales conflictivas de diferentes clases.
   • En su lugar, definen la C2I basada en la brecha de influencia: una muestra se considera útil si su gradiente de pérdida se alinea positivamente con los gradientes de la misma clase en la validación y se opone (se alinea negativamente) a los gradientes de otras clases.
   • Para clasificación binaria, la C2I se formula como la diferencia al cuadrado entre la media de la influencia en los conjuntos de validación positivo y negativo, normalizada por la varianza. Para tareas multiclase, se define como la puntuación softmax de la media de la influencia en la clase correspondiente en relación con todas las clases.

2. Conexión Teórica con la Dificultad de la Muestra:

   • El análisis teórico (Teorema 1) demuestra que maximizar la C2I en un entorno de regresión logística impulsa la muestra generada hacia la media global del conjunto de validación.
   • Dado que la media global de dos grupos de clases suele situarse cerca de la frontera de decisión, maximizar la C2I selecciona inherentemente ejemplos "difíciles": muestras que son difíciles de clasificar y que se encuentran cerca de la frontera.
   • Entrenar con estos ejemplos difíciles y próximos a la frontera ayuda a refinar la frontera de decisión, mejorando la generalización y la robustez.

3. Ajuste Fino mediante Aprendizaje por Refuerzo (RL):

   • Para operativizar esto, los autores proponen un esquema de ajuste fino para modelos de difusión de Texto-a-Imagen (T2I) preentrenados (específicamente Stable Diffusion 2.1).
   • Función de Recompensa: La puntuación C2I sirve como la señal de recompensa en un marco de Aprendizaje por Refuerzo. El generador se ajusta para maximizar la recompensa esperada de C2I.
   • Optimización: Los autores utilizan la Optimización de Política de Difusión de Denotación (DDPO) para actualizar los parámetros del modelo de difusión (específicamente los pesos de LoRA) basados en la recompensa C2I. Esto dirige el proceso de generación hacia las regiones "informativas de la clase" del manifold de los datos sin requerir cambios arquitectónicos en el generador o el clasificador.
   • Flujo de Trabajo: Un clasificador (ViT-B/16) se ajusta primero en el conjunto de entrenamiento de pocos disparos. Sus gradientes en un conjunto de validación mantenido aparte se utilizan para calcular la recompensa C2I para los lotes generados, que luego guían el ajuste fino por RL del modelo de difusión.

Contribuciones Clave

• Formalización de la Efectividad: El artículo formaliza la efectividad de la muestra a través de C2I, demostrando que las muestras de alta C2I corresponden a ejemplos difíciles y próximos a la frontera, en lugar de solo muestras realistas.
• Esquema de Generación Basado en RL: Introduce un marco de RL "plug-and-play" que utiliza la C2I como recompensa para dirigir los modelos de difusión comerciales hacia la generación de muestras útiles para la tarea e informativas de la clase.
• Validación Empírica: El método se valida a través de múltiples benchmarks de imágenes médicas de pocos disparos (BreastMNIST, DermaMNIST, PneumoniaMNIST), mostrando mejoras consistentes sobre el aumento estándar y otros métodos de difusión basados en la comparación.

Resultados Experimentales
El método fue evaluado en tres benchmarks de MedMNIST bajo configuraciones de pocos disparos (16 o 32 muestras etiquetadas por clase).

• Desempeño: El generador guiado por C2I superó consistentemente a los baselines, incluyendo transformaciones estándar (RandAugment, RandomErasing, Mixup) y otros métodos de aumento basados en difusión (DataDream, Dataset Expansion, DistDiff).
  • En promedio, el método propuesto logró el AUC más alto a través de diferentes arquitecturas (ViT-B/16 y ResNet-18).
  • Notablemente, el método mejoró el desempeño en ResNet-18, indicando que las muestras sintetizadas son ampliamente informativas y se transfieren entre arquitecturas de modelos.
• Robustez: Los modelos entrenados con el aumento guiado por C2I demostraron una robustez superior bajo diversas condiciones de ruido (sal y pimienta, desenfoque gaussiano, compresión JPEG), lo que sugiere la generación de una frontera de decisión más estable.
• Correlación de Dificultad: Los experimentos confirmaron que a medida que el ajuste fino por RL progresa, la recompensa C2I aumenta y las características de las imágenes generadas se acercan a los centroides de características de la validación, confirmando la generación de ejemplos más difíciles.
• Ablación: Aumentar el número de épocas de entrenamiento por RL y el número de imágenes sintetizadas generalmente mejoró el AUC de la tarea posterior, mientras que algunos baselines (por ejemplo, Dataset Expansion) mostraron una degradación del desempeño con el aumento de datos sintéticos.

Significancia y Reivindicaciones
El artículo sostiene que el aumento sintético es más efectivo cuando se guía por la utilidad de la tarea en lugar de solo por la calidad de la imagen. Al alinear los datos sintéticos con los gradientes de validación, la C2I transforma los modelos de difusión genéricos en generadores dirigidos que fortalecen las fronteras de decisión donde los datos etiquetados son más escasos.

Los autores reconocen limitaciones, señalando que su método introduce una carga computacional adicional en comparación con el uso de los datos de entrenamiento originales, principalmente debido al proceso de ajuste fino por RL (que requiere aproximadamente 5 horas de GPU en una A100). Sin embargo, afirman que este costo se justifica por las ganancias consistentes en precisión y robustez en regímenes de pocos datos. El estudio sugiere un cambio de paradigma en el uso de modelos de difusión para el aumento: optimizar la generación para su influencia descendente en el clasificador.