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

Resumo Técnico: CardioComposer

1. Problema e Motivação

Modelos generativos de anatomia cardiovascular 3D têm o potencial de sintetizar estruturas informativas para pesquisa clínica e avaliação de dispositivos médicos. No entanto, existe um compromisso (trade-off) crítico entre controlabilidade geométrica e realismo anatômico.

• Desafios Específicos:
  • Escala Crítica: Variações geométricas de milímetros podem alterar drasticamente o comportamento fisiológico em simulações.
  • Efeitos Específicos de Atributos: Tamanho, posição e forma desempenham papéis distintos nos resultados biofísicos.
  • Interação Composicional: A geometria de múltiplas subestruturas interage de forma complexa, onde o resultado depende do arranjo coletivo.
  • Limitações Atuais: Métodos existentes baseados em condicionamento durante o treinamento (ex: concatenar vetores de tamanho) são rígidos, exigem retreinamento para novas restrições e muitas vezes falham em controlar atributos além do volume (como orientação ou forma detalhada). Métodos de edição (inpainting) modificam apenas estruturas existentes, não gerando novas variações plausíveis do zero.

O objetivo é criar um framework que permita o controle interpretável e composicional de atributos geométricos (tamanho, forma, posição) em modelos de difusão não condicionados (unconditional), sem a necessidade de retreinamento.

2. Metodologia: CardioComposer

O CardioComposer é um framework de guia baseado em energia (energy-based guidance) aplicado no momento da inferência (inference-time). Ele utiliza um modelo de difusão latente não condicionado e aplica correções baseadas em gradientes derivados de funções de perda geométrica diferenciáveis.

Componentes Principais:

1. Primitivas Elipsoidais Interpretáveis:

   • As restrições são definidas por elipsoides 3D que representam subestruturas anatômicas (ex: ventrículos, artérias).
   • Os atributos controláveis são: Tamanho (Massa/Volume), Posição (Centróide) e Forma (Covariância/Eixo de orientação).

2. Medidas Geométricas Diferenciáveis (Momentos Geométricos):

   • O sistema calcula momentos geométricos a partir de mapas de rótulos (segmentações) em nível de voxel:
     • Massa ($M$): Momento de ordem zero (volume).
     • Centróide ($C$): Momento de primeira ordem (posição).
     • Covariância ($S$): Momento de segunda ordem normalizado (forma e orientação).
   • A covariância é normalizada para ser invariante à escala, permitindo o controle independente da forma.

3. Guia de Amostragem (Sampling Guidance):

   • Durante o processo de difusão reversa (denoising), o modelo gera uma previsão limpa $\hat{x}_0$.
   • Calcula-se uma Perda Geométrica ($L_{geom}$) comparando os momentos medidos na amostra atual com os momentos alvo ($\bar{G}$).
   • A perda é uma soma ponderada: $L_{geom} = \lambda_0 L_{size} + \lambda_1 L_{pos} + \lambda_2 L_{shape}$.
   • O gradiente dessa perda em relação ao latente ruidoso ($z_\sigma$) é usado para corrigir o passo de denoising:
     $$D^w_\theta(z_\sigma; \sigma) = D_\theta(z_\sigma; \sigma) - \sigma^2 \cdot w \cdot \nabla_{z_\sigma} L_{geom}$$
   • Isso permite direcionar a geração para atender às restrições geométricas sem alterar os pesos do modelo treinado.

4. Controle Composicional:

   • O framework permite aplicar perdas simultâneas a múltiplas subestruturas (ex: controlar o ventrículo direito e a aorta ao mesmo tempo), permitindo a geração de anatomias complexas e não convexas.

3. Contribuições Chave

• Geometria Diferenciável para Caracterização Anatômica: Introdução de funções de medição baseadas em momentos geométricos voxel a voxel que permitem calcular tamanho, posição e forma de forma diferenciável para guiar modelos de difusão.
• Guia no Momento da Inferência: Demonstração de que gradientes simples derivados de perdas geométricas podem controlar modelos de difusão latente não condicionados, permitindo o controle independente ou conjunto de atributos sem retreinamento.
• Controle Composicional Complexo: Validação de que múltiplas perdas geométricas específicas de subestruturas podem ser compostas para criar restrições anatômicas complexas, aplicáveis a estruturas não convexas (ramificadas ou curvas) e uniões de classes de tecidos.
• Generalidade: O método foi validado em diversos sistemas anatômicos (cardíaco, vascular e esquelético), demonstrando robustez além do coração.

4. Resultados Experimentais

Os experimentos foram realizados utilizando o dataset TotalSegmentator (anatomia cardíaca) e datasets adicionais de aorta, coluna vertebral e joelho.

• Fidelidade Condicional vs. Realismo:
  • O CardioComposer superou métodos de condicionamento explícito e implícito (treinados com dropout) na fidelidade condicional (precisão em atingir o tamanho, posição e forma alvo).
  • Ao contrário dos métodos condicionados, o CardioComposer manteve a qualidade morfológica e a distribuição de dados (métricas de nuvem de pontos como MMD, Coverage e 1-NNA) mesmo com pesos de guia altos, evitando o colapso de modo ou perda de realismo.
• Desacoplamento (Disentanglement):
  • O sistema permitiu o controle independente de atributos. Por exemplo, ao aplicar apenas a perda de "massa", o volume era alterado drasticamente enquanto a posição e a forma permaneciam distribuídas livremente, e vice-versa.
• Geração Composicional Multi-Parte:
  • O método conseguiu gerar anatomias com múltiplas subestruturas restritas simultaneamente (ex: 6 estruturas cardíacas ao mesmo tempo). Métodos baseados em condicionamento com dropout falharam em gerar anatomias de alta qualidade quando todas as condições estavam presentes, devido ao desequilíbrio no treinamento.
• Aplicações em Simulação Biofísica:
  • O framework foi usado para editar anatomias de pacientes específicos (inpainting geométrico) para simulações de pressão biventricular. A alteração controlada do volume do ventrículo direito resultou em mudanças previsíveis e fisicamente plausíveis no deslocamento da parede cardíaca, demonstrando utilidade para "gêmeos digitais" e ensaios clínicos virtuais.
• Generalidade: O método funcionou eficazmente em estruturas complexas como a aorta ascendente ramificada, vértebras da coluna e fêmur, confirmando que não se limita a formas convexas simples.

5. Significado e Impacto

O CardioComposer representa um avanço significativo na interseção entre aprendizado de máquina generativo e engenharia biomédica:

• Interface Intuitiva: Permite que engenheiros e clínicos especifiquem requisitos anatômicos usando primitivas geométricas simples (elipsoides), em vez de vetores latentes abstratos.
• Flexibilidade: Elimina a necessidade de retreinar modelos para cada nova configuração experimental ou variação anatômica, acelerando a criação de coortes virtuais para testes de dispositivos médicos.
• Validação para Simulação: Ao garantir que as variações anatômicas geradas respeitem restrições geométricas precisas, o método viabiliza o uso de modelos generativos em simulações físicas rigorosas (CFD, FEM), onde pequenas variações geométricas têm grandes impactos funcionais.

Em resumo, o trabalho oferece uma solução robusta para o controle composicional de modelos de difusão 3D, tornando a geração de anatomias sintéticas mais precisa, controlável e pronta para uso em simulações biofísicas de alta fidelidade.