Resumo Técnico: CONFLUX

Definição do Problema

Modelos generativos controláveis para imagens médicas 3D visam sintetizar volumes com atributos clínicos específicos para tarefas como aumento de coorte, compartilhamento de dados preservando a privacidade e design de estudos contrafatuais. No entanto, alcançar uma síntese 3D nativa de alta fidelidade que seja simultaneamente fiel ao condicionamento solicitado continua sendo um desafio significativo. Os modelos existentes frequentemente têm dificuldade em garantir que amostras individuais geradas exibam os achados clínicos específicos solicitados, pois os objetivos de treinamento padrão (como o flow matching) otimizam a verossimilhança da distribuição agregada em vez da realização de atributos por amostra. Essa lacuna é particularmente aguda em TC de tórax, que exige alta resolução e coerência volumétrica, sendo descrita por metadados radiológicos estruturados.

Metodologia

O CONFLUX é um framework de difusão latente de três estágios projetado para a síntese de TC de tórax, integrando um autoencoder variacional (VAE), um transformer de fluxo retificado (rectified-flow transformer) e um estágio de pós-treinamento por aprendizado por reforço (RL).

1. Autoencoder Latente (Estágio 1)

Para mitigar o custo computacional da geração direta no espaço de voxels, um VAE convolucional 3D comprime volumes de TC pré-processados ($216 \times 176 \times 200$) em um espaço latente de menor resolução ($16 \times 27 \times 22 \times 25$).

• Arquitetura: O codificador e o decodificador utilizam blocos residuais de normalização de grupo com atenção autoatencional 3D na resolução mais baixa.
• Treinamento: O modelo minimiza uma perda de reconstrução $\ell_1$, um regularizador Kullback–Leibler e uma perda perceptual LPIPS triplanar.
• Normalização: Pós-treinamento, o espaço latente é normalizado para uma variância próxima à unidade para facilitar o flow matching.

2. Transformer de Fluxo Retificado Condicional (Estágio 2)

Um transformer de fluxo único gera embeddings latentes a partir de ruído Gaussiano dentro do espaço latente aprendido.

• Flow Matching: O modelo é treinado sob um objetivo de fluxo retificado, aprendendo a prever o campo de velocidade constante que conecta o ruído aos latentes dos dados ao longo de caminhos de linha reta.
• Condicionamento: Metadados radiológicos estruturados (18 achados de anormalidades binárias, sexo, década de idade e kernel de reconstrução) são concatenados em um vetor de 42 dimensões. Este vetor modula o transformer via normalização de camada adaptativa (adaLN-zero).
• Amostragem: A geração envolve a integração da ODE de fluxo de probabilidade de $t=1$ até $t=0$ usando uma grade Euler com deslocamento de tempo, seguida pela decodificação com o decodificador VAE congelado.

3. Pós-treinamento por Aprendizado por Reforço (Estágio 3)

Para abordar a desconexão entre a correspondência de distribuição agregada e a fidelidade de atributos por amostra, o CONFLUX emprega a Otimização de Política Relativa ao Grupo (GRPO).

• Motivação: O flow matching padrão não recompensa explicitamente se um achado específico solicitado aparece em um volume gerado individual.
• Mecanismo: O amostrador de ODE determinístico é substituído por um amostrador estocástico (SDE) para fornecer log-probabilidades tratáveis para atualizações de gradiente de política.
• Sinal de Recompensa: Um classificador independente e congelado, treinado em representações latentes reais, pontua os volumes gerados com base no acordo entre os achados solicitados e os realizados. A recompensa é definida como a entropia cruzada ponderada negativa entre as predições do classificador e o vetor de condicionamento.
• Otimização: O modelo é ajustado para maximizar as vantagens relativas ao grupo (subtraindo a recompensa média de um grupo de execrações por prompt e dividindo pelo desvio padrão). As atualizações usam um substituto PPO com recorte (clipped PPO) com uma penalidade de divergência KL contra a política de referência pré-RL para evitar o desvio (drift).

Principais Contribuições

1. Modelo de Fluxo Retificado Latente 3D Controlável: Uma arquitetura 3D nativa para TC de tórax que condiciona diretamente em metadados radiológicos estruturados via adaLN, alcançando qualidade de síntese competitiva com fortes baselines de difusão 3D.
2. Pós-treinamento por RL para Fidelidade: A primeira aplicação de GRPO a um modelo de fluxo médico 3D. Este estágio otimiza explicitamente a confiabilidade dos achados solicitados nos volumes gerados, verificada por um modelo julgador independente.
3. Lançamento Aberto: O lançamento do modelo treinado e de um conjunto de dados de TC de tórax sintético de aproximadamente 200.000 volumes. Este conjunto de dados é condicionado às proporções de metadados reais, abrangendo uma ampla variedade de achados clínicos, destinado ao aumento de coorte e design de estudos controlados.

Resultados

Qualidade de Síntese

Avaliado em um conjunto de validação mantido fora (held-out) ($N=501$) usando FID triplanar, o CONFLUX supera significativamente os baselines existentes:

• FID Triplanar: 32,3 (CONFLUX) vs. 74,6 (MAISI) e 145,4 (GenerateCT).
• Cobertura e Densidade: As amostras do CONFLUX são 3,5$\times$ mais densas e cobrem 6$\times$ mais do manifold real em comparação ao MAISI.
• Diversidade: O modelo mantém alta diversidade (MS-SSIM 0,507), comparável aos dados reais (0,41), evitando o colapso de modo (mode collapse).

Fidelidade de Condicionamento

O pós-treinamento com GRPO melhora o acordo entre os atributos solicitados e realizados, conforme medido por um classificador independente no espaço da imagem:

• Ganho de Desempenho: A Precisão Média (AP) aumentou de 0,330 para 0,344, e a AUROC de 0,684 para 0,699.
• Recuperação de Lacuna: Embora os ganhos absolutos pareçam modestos, eles recuperam 47% da lacuna de desempenho entre o modelo base e o teto dos dados reais (AP 0,360).
• Especificidade: As melhorias foram específicas para o grupo de achados otimizados; grupos não ponderados (sexo, idade, kernel) permaneceram inalterados dentro do ruído.

Desempenho do Classificador de Recompensa

O classificador usado para o sinal de recompensa de RL alcançou um macro AUROC de 0,793 no split de validação real, superando variantes do CT-CLIP e o CT-Net no benchmark CT-RATE, validando sua utilidade como sinal de recompensa, apesar de operar em latentes comprimidos.

Significância e Alegações

O artigo afirma que o CONFLUX representa um passo significativo à frente na síntese de imagens médicas 3D controláveis ao desacoplar a geração de alta fidelidade do controle de atributos. Ao integrar o pós-treinamento por RL, os autores demonstram que é possível otimizar diretamente a "fidelidade" de atributos clínicos em amostras individuais, uma propriedade não garantida pelos objetivos padrão de flow-matching. O lançamento de um conjunto de dados sintético condicionalmente diverso visa facilitar a pesquisa em cenários onde os dados reais são escassos ou possuem restrições de privacidade. O trabalho posiciona o CONFLUX como a primeira adaptação de GRPO para modelos de fluxo médico 3D, oferecendo um caminho para ferramentas generativas mais confiáveis e controláveis para o design de estudos clínicos.