Flow Matching with Optimized Subclass Priors for Medical Image Augmentation

Resumen Técnico: Flow Matching con Priors de Subclase Optimizados para Aumento de Imágenes Médicas

Enunciado del Problema
Los conjuntos de datos de imágenes médicas se caracterizan por desequilibrios severos de cola larga, donde las condiciones comunes dominan mientras que las patologías raras y clínicamente críticas (por ejemplo, neumotórax pequeños, bronquiectasias) están subrepresentadas. Este desequilibrio provoca que los clasificadores discriminativos se sobreajusten a las clases principales y fallen en condiciones raras. Si bien el aumento de datos generativo ofrece una solución potencial, el flow matching (FM) condicional estándar y los modelos de difusión enfrentan dos dificultades compounding cuando se entrenan con etiquetas de enfermedades gruesas:

1. Agregación Multimodal: Una única etiqueta diagnóstica a menudo agrega subtipos diversos, protocolos de adquisición y comorbilidades. Esto crea condicionales multimodales que sesgan a los generadores hacia los submodos dominantes, truncando efectivamente variaciones raras pero clínicamente válidas.
2. Geometría de Transporte Ineficiente: Extraer todas las trayectorias de generación desde una fuente gaussiana estándar compartida fuerza a subpoblaciones específicas a través de caminos de transporte desproporcionadamente largos. Esto infla la dificultad de regresión y degrada la fidelidad de la generación de clases de cola.

Metodología
Los autores proponen una estrategia offline que introduce priores informativos en dos niveles para refinar el proceso generativo sin modificar el bucle de entrenamiento central del modelo de flow matching.

1. Inducción de Subclase mediante Modelado de Mezclas Gaussianas (GMM):

   • El método divide cada etiqueta gruesa $c$ en submodos coherentes (subclases $k$) dentro del espacio latente del modelo generativo.
   • Utilizando un VAE congelado (FLUX1.dev), se calculan latentes deterministas. Para cada clase $c$, se ajusta un Modelo de Mezcla Gaussiana diagonal a los latentes residuales (después de restar el centro de clase $\mu_c$).
   • El número de componentes $K_c$ se selecciona utilizando el Criterio de Información Bayesiano Extendido (EBIC). Las muestras reciben asignaciones duras a subclases, creando una partición con menor varianza condicional.
   • Teóricamente, los autores demuestran que condicionar sobre estas etiquetas refinadas no puede aumentar el riesgo óptimo de Bayes del flow matching; la mejora es igual a la variabilidad esperada entre subclases del desplazamiento medio condicional.

2. Optimización de Fuente Dependiente de Subclase:

   • En lugar de una fuente gaussiana compartida, el método aprende una distribución de fuente específica $\pi_0(\cdot | c, k)$ para cada subclase. Estas fuentes se parametrizan como Gaussianas aprendibles $N(\mu_{c,k}, \text{diag}(\sigma^2_{c,k}))$ que re-centran y re-escalan la distribución de inicio.
   • Esto acorta las trayectorias de transporte típicas y reduce la dispersión dentro de la subclase.
   • Control Geométrico: Para prevenir soluciones degeneradas (como inflación de norma o caminos excesivamente concentrados), los autores optimizan un objetivo proxy offline que impone:
     • Compacidad Direccional: Fomentar que las direcciones de desplazamiento normalizadas se alineen moderadamente con prototipos aprendibles.
     • Limitación de Longitud de Trayectoria: Imponer una penalización fija a los valores atípicos utilizando un límite específico de subclase ($r_{cap}$), definido como el percentil 99 de las distancias al centro de clase. Esto evita que el modelo eluda la pérdida inflando las normas de desplazamiento.
     • Regularización de Escala: Regularizar las escalas logarítmicas de las distribuciones de fuente hacia una referencia unitaria.

Contribuciones Clave

1. Formulación Condicional de Subclase: Un enfoque novedoso para la generación de imágenes médicas de cola larga que divide condicionales multimodales mediante asignaciones duras inducidas por GMM, reduciendo efectivamente la ambigüedad condicional.
2. Fuentes Aprendidas Dependientes de Subclase: La introducción de priores informativos que remodelan el transporte condicional al re-centrar y re-escalar la distribución de fuente por subclase, reduciendo así la dispersión dentro de la subclase.
3. Objetivo Regularizado Geométricamente: Un objetivo proxy práctico que equilibra una compacidad angular moderada con un control explícito de valores atípicos en la longitud de trayectoria, evitando atajos de inflación de norma y colapso direccional.
4. Validación Empírica: Demostración de una mejor cobertura de cola y rendimiento de clasificación de cola larga aguas abajo en los puntos de referencia de rayos X de tórax (MIMIC-LT, NIH-LT) y TC (CT-RATE).

Resultados Experimentales
El método se evaluó en tres puntos de referencia de imágenes médicas de cola larga:

• Clasificación Aguas Abajo: El método propuesto logró consistentemente la mayor precisión equilibrada (bAcc) y puntuaciones macro-F1 en los tres conjuntos de datos en comparación con los baselines sin aumento y otras estrategias de aumento (FM Vanilla, CBFM, CPD). Cabe destacar que fue la única estrategia que mejoró de manera confiable el rendimiento en ambas modalidades de rayos X y TC.
• Calidad Generativa: En MIMIC-LT, el método logró la mejor Distancia Fréchet Inception (FID) y Puntuación de Ratio Inception (IRS), indicando ganancias simultáneas en fidelidad y diversidad. En CT-RATE, aunque el FID fue ligeramente superior al de un baseline, el método mantuvo un IRS alto y una clasificación aguas abajo superior, sugiriendo una mejor cobertura distribucional en lugar de mera memorización.
• Estudios de Ablación: Los resultados confirmaron que ambas etapas del pipeline (subclasificación y optimización de fuente) son necesarias. La subclasificación por sí sola mejoró el rendimiento, pero la combinación arrojó los mejores resultados. La optimización sin subclasificación a veces degradó el rendimiento al promediar sobre modos conflictivos.
• Validez de Subclase: El análisis mostró que las subclases inducidas capturan una estructura relevante para la patología en lugar de solo metadatos de adquisición. Las sondas de confusores indicaron que los metadatos eran detectables pero no dominantes, y la pureza kNN de las subclases aprendidas superó significativamente a las particiones aleatorias.

Significado y Afirmaciones
El artículo afirma que el diseño de fuentes y el refinamiento condicional son palancas efectivas para el aumento generativo robusto en entornos clínicos de cola larga. Al dividir etiquetas heterogéneas y aprender fuentes dependientes de subclase, el método reduce la ambigüedad condicional y estabiliza las trayectorias de transporte a un costo computacional insignificante (todo el proceso es offline y toma menos de una hora de GPU por conjunto de datos). Los autores lo posicionan como una estrategia principista que mejora la fidelidad y diversidad de la generación de clases de cola, conduciendo a clasificadores aguas abajo más confiables para condiciones médicas raras. El trabajo destaca que la generación condicional ingenua a menudo falla en subpoblaciones raras debido a problemas de geometría de transporte, los cuales pueden mitigarse mediante priores estructurados.