Resumen Técnico: CONFLUX

Planteamiento del Problema

Los modelos generativos controlables para imágenes médicas 3D buscan sintetizar volúmenes con atributos clínicos específicos para tareas como el aumento de cohortes, el intercambio de datos preservando la privacidad y el diseño de estudios contrafácticos. Sin embargo, lograr una síntesis 3D nativa de alta fidelidad que sea simultáneamente fiel a las condiciones solicitadas representa un desafío significativo. Los modelos existentes suelen tener dificultades para garantizar que los volúmenes generados individualmente exhiban los hallazgos clínicos específicos solicitados, ya que los objetivos de entrenamiento estándar (como el flow matching) optimizan el realismo de la distribución agregada en lugar de la realización de atributos por muestra. Esta brecha es particularmente aguda en la TC de tórax, que requiere alta resolución y coherencia volumétrica, además de estar descrita mediante metadatos radiológicos estructurados.

Metodología

CONFLUX es un marco de difusión latente de tres etapas diseñado para la síntesis de TC de tórax, que integra un autoencoder variacional (VAE), un transformador de flujo rectificado (rectified-flow transformer) y una etapa de post-entrenamiento mediante aprendizaje por refuerzo (RL).

1. Autoencoder Latente (Etapa 1)

Para mitigar el costo computacional de la generación directa en el espacio de vóxeles, un VAE convolucional 3D comprime volúmenes de TC preprocesados ($216 \times 176 \times 200$) en un espacio latente de menor resolución ($16 \times 27 \times 22 \times 25$).

• Arquitectura: El codificador y el decodificador utilizan bloques residuales de normalización de grupo con atención propia (self-attention) 3D en la resolución más baja.
• Entrenamiento: El modelo minimiza una pérdida de reconstrucción $\ell_1$, un regularizador de Kullback–Leibler y una pérdida perceptual LPIPS triplanar.
• Normalización: Tras el entrenamiento, el espacio latente se normaliza a una varianza cercana a la unidad para facilitar el flow matching.

2. Transformador de Flujo Rectificado Condicional (Etapa 2)

Un transformador de flujo único genera incrustaciones (embeddings) latentes a partir de ruido gaussiano dentro del espacio latente aprendido.

• Flow Matching: El modelo se entrena bajo un objetivo de flujo rectificado, aprendiendo a predecir el campo de vectores de velocidad constante que conecta el ruido con los latentes de los datos a lo largo de trayectorias de línea recta.
• Condicionamiento: Los metadatos radiológicos estructurados (18 hallazgos de anomalías binarias, sexo, década de edad y núcleo de reconstrucción) se concatenan en un vector de 42 dimensiones. Este vector modula el transformador mediante normalización de capa adaptativa (adaLN-zero).
• Muestreo: La generación implica la integración de la ODE de flujo de probabilidad desde $t=1$ hasta $t=0$ utilizando una rejilla de Euler con desplazamiento temporal, seguido de la decodificación con el decodificador VAE congelado.

3. Post-entrenamiento mediante Aprendizaje por Refuerzo (Etapa 3)

Para abordar la desconexión entre el emparejamiento de la distribución agregada y la fidelidad de los atributos por muestra, CONFLUX emplea la Optimización de Política Relativa de Grupo (GRPO).

• Motivación: El flow matching estándar no recompensa explícitamente si un hallazgo específico solicitado aparece en un volumen generado individual.
• Mecanismo: El muestreador de ODE determinista se reemplaza con un muestreador estocástico (SDE) para proporcionar log-probabilidades tratables para las actualizaciones de gradiente de política.
• Señal de Recompensa: Un clasificador independiente y congelado, entrenado en representaciones latentes reales, califica los volúmenes generados basándose en el acuerdo entre los hallazgos solicitados y los realizados. La recompensa se define como la entropía cruzada ponderada negativa entre las predicciones del clasificador y el vector de condicionamiento.
• Optimización: El modelo se ajusta para maximizar las ventajas relativas de grupo (restando la media de la recompensa de un grupo de ejecuciones por prompt y dividiendo por la desviación estándar). Las actualizaciones utilizan un sustituto PPO recortado con una penalización de divergencia KL contra la política de referencia pre-RL para evitar el desvío (drift).

Contribuciones Clave

1. Modelo de Flujo Rectificado Latente 3D Controlable: Una arquitectura 3D nativa para TC de tórax que se condiciona directamente en metadatos radiológicos estructurados mediante adaLN, logrando una calidad de síntesis competitiva con bases de comparación de difusión 3D robustas.
2. Post-entrenamiento con RL para la Fidelidad: La primera aplicación de GRPO a un modelo de flujo médico 3D. Esta etapa optimiza explícitamente la fiabilidad de los hallazgos solicitados en los volúmenes generados, verificada por un modelo juez independiente.
3. Lanzamiento Abierto: El lanzamiento del modelo entrenado y de un conjunto de datos de TC de tóxica sintética de aproximadamente 200,000 volúmenes. Este conjunto de datos está condicionado en proporciones de metadatos reales, abarcando una amplia variedad de hallazgos clínicos, destinado al aumento de cohortes y al diseño de estudios controlados.

Resultados

Calidad de Síntesis

Evaluado en un conjunto de validación no utilizado ($N=501$) mediante la Distancia de Inception de Fréchet (FID) triplanar, CONFLUX supera significativamente a las bases de comparación existentes:

• FID Triplanar: 32.3 (CONFLOW) frente a 74.6 (MAISI) y 145.4 (GenerateCT).
• Cobertura y Densidad: Las muestras de CONFLUX son 3.5 veces más densas y cubren 6 veces más del espectro de la distribución real en comparación con MAISI.
• Diversidad: El modelo mantiene una alta diversidad (MS-SSIM 0.507), comparable a los datos reales (0.41), evitando el colapso de modo.

Fidelidad del Condicionamiento

El post-entrenamiento con GRPO mejora el acuerdo entre los atributos solicitados y los realizados, medido por un clasificador independiente en el espacio de la imagen:

• Ganancia de Desempeño: La Precisión Media (AP) aumentó de 0.330 a 0.344, y el AUROC de 0.684 a 0.699.
• Recuperación de la Brecha: Aunque las ganancias absoldas parecen modestas, recuperan el 47% de la brecha de rendimiento entre el modelo base y el techo de los datos reales (AP 0.360).
• Especificidad: Las mejoras fueron específicas para el grupo de hallazgos optimizados; los grupos no ponderados (sexo, edad, núcleo) permanecieron sin cambios dentro del ruido.

Desempeño del Clasificador de Recompensa

El clasificador utilizado para la señal de recompensa de RL alcanzó un macro AUROC de 0.793 en el conjunto de validación real, superando a las variantes de CT-CLIP y CT-Net en el benchmark CT-RATE, validando su utilidad como señal de recompensa a pesar de operar en latentes comprimidos.

Significancia y Reclamaciones

El artículo afirma que CONFLUX representa un paso significativo hacia adelante en la síntesis de imágenes médicas 3D controlables al desacoplar la generación de alta fidelidad del control de atributos. Al integrar el post-entrenamiento con RL, los autores demuestran que es posible optimizar directamente la "fidelidad" de los atributos clínicos en muestras individuales, una propiedad que no está garantizada por los objetivos estándar de flow-matching. El lanzamiento de un conjunto de datos sintéticos condicionalmente diversos y a gran escala tiene como objetivo facilitar la investigación en escenarios donde los datos reales son escasos o están sujetos a restricciones de privacidad. El trabajo posiciona a CONFLUX como el primero en adaptar GRPO para modelos de flujo médico 3D, ofreciendo una vía para herramientas generativas más fiables y controlables para el diseño de estudios clínicos.