CardioComposer: Leveraging Differentiable Geometry for Compositional Control of Anatomical Diffusion Models

1. El Problema

La generación de modelos anatómicos 3D realistas es crucial para la investigación clínica, la evaluación de dispositivos médicos y las simulaciones biofísicas (como ensayos clínicos virtuales). Sin embargo, los modelos generativos actuales enfrentan una compensación fundamental (trade-off) entre:

• Realismo: La capacidad de generar estructuras anatómicamente plausibles y complejas.
• Control Geométrico: La capacidad de manipular atributos específicos (tamaño, forma, posición) de forma interpretable y precisa.

Los enfoques existentes tienen limitaciones:

• Los primitivos geométricos simples (cilindros, elipsoides) ofrecen control pero carecen de realismo anatómico.
• Los modelos basados en datos (autoencoders, difusión condicional) a menudo requieren reentrenamiento para nuevos atributos o no permiten un control independiente (desenredado) de múltiples subestructuras.
• Las variaciones geométricas menores (del orden de milímetros) pueden alterar drásticamente los resultados fisiológicos en simulaciones, por lo que el control preciso es crítico.

2. Metodología: CardioComposer

El autores proponen CardioComposer, un marco de trabajo programable que opera en tiempo de inferencia (sin reentrenar el modelo) para controlar modelos de difusión latente incondicionales mediante geometría diferenciable.

Conceptos Clave:

1. Primitivas Elipsoidales Interpretables:
   En lugar de controlar píxeles directamente, el método utiliza elipsoides para representar subestructuras anatómicas discretas. Estos elipsoides codifican atributos geométricos: tamaño, forma y posición.

2. Funciones de Medición Diferenciables (Momentos Geométricos):
   El núcleo de la innovación es el cálculo de momentos geométricos a partir de mapas de etiquetas de voxels (segmentaciones) que son totalmente diferenciables. Esto permite calcular gradientes para guiar el proceso de difusión:

   • Tamaño (Mass): Calculado mediante momentos de orden cero (volumen).
   • Posición (Centroid): Calculado mediante momentos de primer orden.
   • Forma (Covarianza): Calculado mediante momentos de segundo orden normalizados (matriz de covarianza), lo que captura la orientación y la relación de aspecto.

3. Guía Basada en Energía (Energy-Based Guidance):
   Durante el proceso de muestreo inverso (denoising) del modelo de difusión, se aplica una corrección basada en la pérdida geométrica:

   • Se genera una predicción limpia $\hat{x}_0$ a partir del latente ruidoso.
   • Se extraen los momentos geométricos medidos ($G$) de las subestructuras seleccionadas.
   • Se compara con los momentos objetivo ($\bar{G}$) para calcular una pérdida ($L_{geom}$).
   • El gradiente de esta pérdida se utiliza para corregir el paso de denoising, guiando la generación hacia la geometría deseada sin modificar los pesos del modelo base.

   La actualización se formula como:
   $$D^w_\theta(z_\sigma; \sigma) = D_\theta(z_\sigma; \sigma) - \sigma^2 \cdot w \cdot \nabla_{z_\sigma} L_{geom}$$
   Donde $w$ es el peso de la guía.

4. Control Composicional:
   El marco permite combinar múltiples funciones de pérdida (tamaño, posición, forma) para diferentes subestructuras simultáneamente, permitiendo la generación de anatomías complejas con múltiples componentes restringidos.

3. Contribuciones Clave

• Geometría Diferenciable para Caracterización Anatómica: Introducen un conjunto de funciones de medición geométrica basadas en momentos que permiten calcular atributos fisiológicamente relevantes (volumen, centroide, forma) directamente desde mapas de etiquetas, habilitando la optimización basada en gradientes.
• Guía en Tiempo de Inferencia: Demuestran que se puede controlar modelos de difusión latente incondicionales para generar segmentaciones multi-clase con atributos específicos sin necesidad de entrenamiento condicional, evitando el sesgo de entrenamiento y permitiendo flexibilidad.
• Control Desenredado (Disentangled Control): Validan que es posible controlar independientemente el tamaño, la posición y la forma de una subestructura sin afectar significativamente a las otras, algo difícil de lograr con métodos condicionales tradicionales.
• Control Composicional Multi-Parte: El método escala para controlar múltiples subestructuras (desde 1 hasta 6 tejidos diferentes) simultáneamente, incluyendo estructuras no convexas y ramificadas (como vasos sanguíneos).

4. Resultados Experimentales

Los autores evaluaron su método en el conjunto de datos TotalSegmentator (corazón y grandes vasos) y lo extendieron a la aorta, columna vertebral y rodilla.

• Calidad de Generación vs. Fidelidad Condicional:
  • CardioComposer supera a los métodos de entrenamiento condicional (concatenación explícita/implícita, atención cruzada) en términos de fidelidad geométrica (precisión en tamaño, posición y forma).
  • Mantiene métricas de calidad morfológica y de nubes de puntos (MMD, Cobertura, 1-NN) comparables a las del modelo incondicional, demostrando que la guía no degrada el realismo anatómico.
• Desenredamiento:
  • Los experimentos muestran que aplicar una pérdida solo de "masa" ajusta el volumen sin alterar significativamente la posición o la forma, y viceversa.
  • Las distribuciones marginales de las métricas morfológicas se estrechan alrededor del valor objetivo solo para los atributos controlados.
• Composicionalidad:
  • Al restringir múltiples subestructuras (ej. ventrículo derecho, aurícula derecha, arteria pulmonar), el método logra generar anatomías coherentes donde todas las restricciones se cumplen simultáneamente, algo que los baselines condicionales fallan al hacer debido a problemas de distribución de datos durante el entrenamiento (dropout).
• Aplicaciones en Simulación:
  • Se demostró la capacidad de editar anatomías de pacientes específicos para simulaciones biomecánicas (ej. aumentar/disminuir el tamaño del ventrículo derecho y simular la presión biventricular), mostrando cambios predecibles en el desplazamiento de la pared cardíaca.
• Generalidad:
  • El método funcionó exitosamente en sistemas anatómicos diversos (cardíaco, vascular, esquelético) y con geometrías complejas (no convexas, ramificadas).

5. Significado e Impacto

El trabajo de CardioComposer representa un avance significativo en la intersección entre la inteligencia artificial generativa y la simulación médica:

1. Puente entre IA y Física: Permite generar datos sintéticos que no solo son visualmente realistas, sino que obedecen restricciones geométricas precisas necesarias para simulaciones físicas (FEM, CFD).
2. Eliminación de la Dependencia del Entrenamiento: Al operar en tiempo de inferencia, el método es más flexible y robusto que los modelos condicionales, permitiendo a los investigadores definir nuevas restricciones sin volver a entrenar modelos costosos.
3. Interpretabilidad Clínica: El uso de primitivas elipsoidales y momentos geométricos ofrece una interfaz intuitiva para clínicos e ingenieros, traduciendo conceptos anatómicos (volumen, orientación) directamente en parámetros de generación.
4. Habilitador de Ensayos Clínicos Virtuales: Facilita la creación de cohortes virtuales con variaciones anatómicas específicas y raras, esenciales para probar dispositivos médicos en escenarios que podrían ser difíciles o éticamente complejos de obtener en pacientes reales.

En resumen, CardioComposer ofrece un marco robusto para el control composicional de la anatomía 3D, resolviendo el dilema entre el realismo generativo y la precisión geométrica requerida para aplicaciones médicas críticas.