CLIMB: Controllable Longitudinal Brain Image Generation using Mamba-based Latent Diffusion Model and Gaussian-aligned Autoencoder

Resumo Técnico: CLIMB

1. Problema e Motivação

A geração de imagens médicas longitudinais (previsão de imagens de acompanhamento com base em uma imagem basal) é crucial para modelar a progressão de doenças neurodegenerativas, como o Alzheimer, permitindo intervenções precoces e planejamento de tratamento.
No entanto, os métodos existentes enfrentam desafios significativos:

• Custo Computacional: Modelos de Difusão (DM) tradicionais operam no espaço de pixels, exigindo recursos intensivos e inferência lenta devido ao processo iterativo de remoção de ruído.
• Limitações dos Autoencoders Variacionais (VAE): A maioria dos Modelos de Difusão Latente (LDM) utiliza VAEs para comprimir dados. O VAE introduz ruído estocástico durante o treinamento (amostragem de uma distribuição Gaussiana), o que frequentemente resulta em reconstruções de imagens desfocadas ou sem detalhes nítidos.
• Ineficiência de Arquiteturas Baseadas em Transformers: As abordagens atuais de LDM dependem de mecanismos de Self-Attention (atenção) para capturar contexto global. Embora eficazes, eles possuem complexidade quadrática, tornando-se proibitivos para imagens médicas 3D de alta dimensão e exigem grandes conjuntos de dados.
• Falta de Controle Clínico: Muitos modelos não conseguem integrar eficazmente múltiplas variáveis clínicas (idade projetada, gênero, status da doença, volumes cerebrais) para controlar a evolução temporal da anatomia.

2. Metodologia Proposta: CLIMB

O CLIMB (Controllable Longitudinal brain Image generation via MamBa-based latent diffusion model) é um framework inovador que combina três componentes principais para superar as limitações acima:

A. Autoencoder Alinhado a Gaussiana (GATE)

• Em vez de um VAE tradicional, os autores propõem o GATE (Gaussian-aligned Autoencoder).
• Mecanismo: O GATE utiliza uma estratégia determinística de alinhamento de distribuição baseada na Distância de Wasserstein Sliced (Sliced Wasserstein Distance).
• Objetivo: Forçar as representações latentes a seguirem uma distribuição Gaussiana padrão sem a necessidade de amostragem estocástica (ruído) durante o treinamento. Isso preserva a fidelidade e a nitidez das imagens reconstruídas, eliminando o efeito de "borramento" comum em VAEs.
• Funções de Perda: O treinamento combina perda de reconstrução, alinhamento de CDF (Função de Distribuição Cumulativa) cortada, perda perceptual (baseada em redes pré-treinadas como VGG/CLIP) e perda adversarial para garantir realismo fotográfico.

B. Rede de Remoção de Ruído Baseada em Mamba (State Space Models)

• Arquitetura: O núcleo do modelo de difusão latente substitui os mecanismos de Self-Attention (Transformers) pela arquitetura Mamba (Modelos de Espaço de Estado - SSMs).
• Vantagem: Os SSMs modelam dependências de longo alcance com complexidade linear (tempo O(n)), ao contrário da complexidade quadrática dos Transformers. Isso permite processamento eficiente de dados 3D e funciona bem mesmo com conjuntos de dados menores.
• Estratégia de Treinamento em Duas Etapas:
  1. Treinamento inicial apenas com variáveis projetadas (idade, volumes, etc.) para aprender a dinâmica latente.
  2. Congelamento desses parâmetros e ajuste fino da rede de processamento de características da imagem, utilizando camadas de convolução inicializadas com zeros para garantir estabilidade.

C. Processo de Inferência e Controle

• O modelo aceita uma MRI basal e sua idade de aquisição como entrada.
• Utiliza um modelo pré-existente (IRLSTM) para prever dados futuros (status da doença, volumes de estruturas cerebrais) baseados na idade projetada, gênero e genética.
• Essas variáveis condicionais são injetadas na rede de difusão via operações de Cross-Attention, permitindo a geração controlada de uma MRI de acompanhamento futura.
• Utiliza o Denoising Diffusion Implicit Model (DDIM) para acelerar a amostragem (25 passos) e emprega amostragem média latente (10 amostras por execução) para robustez.

3. Principais Contribuições

1. Modelo CLIMB: Um framework de geração de imagens longitudinais que sintetiza mudanças anatômicas usando uma MRI basal e dados projetados, integrando múltiplos fatores condicionais clínicos.
2. GATE (Autoencoder): Introdução de um autoencoder que alinha as características latentes a uma distribuição Gaussiana sem ruído de amostragem, utilizando a distância de Wasserstein Sliced, resultando em reconstruções mais nítidas.
3. Arquitetura Mamba: Substituição do mecanismo de atenção por Mamba no modelo de difusão latente, reduzindo drasticamente a complexidade computacional e o uso de memória, mantendo a capacidade de capturar contexto global e local.
4. Integração de Previsão: Uso do modelo IRLSTM para prever variáveis de progressão da doença que servem como condições para a geração da imagem futura.

4. Resultados Experimentais

O modelo foi avaliado no conjunto de dados ADNI (Alzheimer's Disease Neuroimaging Initiative), contendo 6.306 exames de MRI de 1.390 participantes.

• Métricas Quantitativas: O CLIMB superou os métodos de última geração (SADM e BrLP) em todas as métricas:
  • SSIM (Similaridade Estrutural): 0.9433 (o mais alto, indicando melhor preservação estrutural).
  • PSNR: 27.82 (melhor qualidade de sinal).
  • MSE (Erro Quadrático Médio): 2.01e-3 (menor erro).
  • LPIPS (Similaridade Perceptual): 0.0587 (o mais baixo, indicando maior fidelidade perceptual e menos artefatos).
• Eficiência:
  • Tempo de Inferência: 2.92 segundos por imagem (mais rápido que os baselines).
  • Uso de Memória: 5.699 MB (significativamente menor que os modelos baseados em Transformers).
• Estudos de Ablação:
  • A substituição do VAE pelo GATE reduziu o LPIPS de 0.1080 para 0.0587.
  • A substituição do Self-Attention pelo Mamba melhorou a qualidade da reconstrução e a eficiência computacional.

5. Significado e Conclusão

O CLIMB representa um avanço significativo na geração de imagens médicas longitudinais. Ao combinar a eficiência dos Modelos de Espaço de Estado (Mamba) com a fidelidade de um Autoencoder alinhado a Gaussiana (GATE), o modelo consegue simular a evolução temporal do cérebro com alta precisão e baixo custo computacional.

Impacto Potencial:

• Clínico: Pode auxiliar no planejamento de tratamento, previsão de progressão da doença e redução de custos ao simular exames futuros sem a necessidade de novos escaneamentos reais.
• Pesquisa: Oferece uma ferramenta robusta para gerar dados sintéticos de alta qualidade para treinar outros modelos de IA, especialmente em cenários onde dados longitudinais reais são escassos.

O artigo conclui que, embora as métricas numéricas sejam excelentes, a validação final por especialistas médicos (radiologistas e neurologistas) é essencial para garantir a utilidade clínica real das imagens geradas.