Modality-Aware and Anatomical Vector-Quantized Autoencoding for Multimodal Brain MRI

1. Problema e Motivação

A ressonância magnética (RM) estrutural multimodal (ex: T1, T2, FLAIR) é fundamental para o diagnóstico de doenças cerebrais, pois diferentes contrastes fornecem informações complementares sobre anatomia, integridade tecidual e patologia. No entanto, adquirir RM multimodal de alta qualidade para todos os sujeitos é desafiador devido a artefatos de movimento, restrições de tempo e custos, resultando frequentemente em conjuntos de dados limitados ou desbalanceados.

Para contornar isso, a síntese de RM cerebral (gerar imagens de uma modalidade a partir de outra) tornou-se uma direção promissora. A maioria das abordagens modernas utiliza Modelos de Difusão Latente (LDM), que dependem de um Autoencoder Variacional (VAE) para comprimir volumes 3D de alta dimensão em um espaço latente compacto.

Limitações das abordagens atuais:

• Foco em Modalidade Única: Muitos VAEs existentes são treinados apenas para T1, ignorando o valor diagnóstico de outras modalidades.
• Abordagens 2D vs. 3D: Alguns métodos usam VAEs 2D pré-treinados em imagens naturais (como o Stable Diffusion) aplicados fatia por fatia. Isso falha em capturar o contexto volumétrico, resultando em inconsistências anatômicas entre fatias e reconstruções enviesadas pela modalidade.
• VAEs Universais Genéricos: Modelos como MedVAE ou MediTok, treinados em diversos tipos de exames médicos, muitas vezes tratam as sequências de RM (T1, T2) como entradas independentes ou usam um único backbone sem modelar explicitamente as relações complementares entre as modalidades cerebrais.
• Custo de Retreinamento: VAEs específicos de conjunto de dados exigem retreinamento pesado para cada novo estudo, limitando a generalização.

O objetivo é criar um VAE 3D unificado que modele conjuntamente a estrutura anatômica compartilhada e a aparência específica de cada modalidade, preservando a coerência volumétrica e permitindo generalização escalável.

---

2. Metodologia: NeuroQuant

Os autores propõem o NeuroQuant, um framework de autoencoder vetorial quantizado (VQ-VAE) 3D, consciente de modalidade e fundamentado na anatomia. A arquitetura é projetada para desentrelaçar (disentangle) a informação anatômica compartilhada das características de aparência dependentes da modalidade.

Arquitetura Principal

O modelo utiliza uma arquitetura de duplo fluxo (dual-stream) dentro de um espaço latente unificado:

1. Codificador de Duplo Fluxo:

   • Fluxo Anatômico: Codifica informações estruturais espaciais consistentes (invariantes à modalidade).
   • Fluxo de Modalidade: Captura variações de intensidade e contraste específicas de cada modalidade (T1, T2, etc.).
   • Atenção Multi-eixo Fatorada: Para capturar dependências de longo alcance entre regiões cerebrais distantes sem o custo computacional proibitivo da atenção 3D completa ($O((DHW)^2)$), o modelo utiliza blocos de atenção que operam sequencialmente nos eixos axial, coronal e sagital. Isso reduz a complexidade para $O(DHW(D+H+W))$.

2. Quantização Vetorial (VQ) Anatômica:

   • O latente anatômico ($z_{anat}$) é discretizado usando um livro de códigos (codebook) compartilhado entre todas as modalidades. Isso cria um "vocabulário" compacto de tokens anatômicos, garantindo que a estrutura básica seja representada de forma consistente, independentemente da modalidade de entrada.

3. Decodificador Modulado por FiLM:

   • O latente de modalidade ($z_{mod}$) é usado para gerar parâmetros de Modulação Linear Específica de Recursos (FiLM) ($\gamma, \beta$).
   • Esses parâmetros modulam as ativações do decodificador, permitindo que o modelo ajuste dinamicamente o contraste e a intensidade da imagem reconstruída enquanto mantém a estrutura anatômica fixa fornecida pelo latente quantizado.

Estratégia de Treinamento (2D/3D Joint)

Para lidar com a aquisição baseada em fatias de dados 3D e melhorar a escalabilidade, o NeuroQuant emprega uma estratégia de treinamento conjunta 2D/3D:

• O modelo é treinado alternando entre volumes 3D completos e fatias 2D aleatórias (axial, coronal ou sagital).
• Isso permite que o modelo aprenda tanto a coerência volumétrica global quanto detalhes de alta frequência em nível de fatia, sem a necessidade de codificadores extras.

Funções de Perda e Desentrelaçamento

Para garantir que a estrutura e a aparência sejam realmente separadas, o modelo utiliza:

• Perda de Reconstrução Cruzada (Cross-Modal): Tenta reconstruir uma imagem T2 usando o código anatômico de uma T1 e os parâmetros FiLM da T2 (e vice-versa).
• Perda Adversarial de Modalidade: Um classificador tenta prever a modalidade a partir do latente anatômico. O codificador é treinado para enganar esse classificador (via reversão de gradiente), forçando o latente anatômico a ser invariante à modalidade.

---

3. Contribuições Principais

1. NeuroQuant: Um framework VQ-VAE 3D unificado para RM cerebral multimodal que captura estrutura anatômica e aparência específica da modalidade em um espaço latente compartilhado.
2. Arquitetura Inovadora: Introdução de um codificador de duplo fluxo com tokens anatômicos vetorialmente quantizados e um decodificador modulado por FiLM, reforçado por atenção multi-eixo fatorada para consistência estrutural.
3. Desempenho Superior: Demonstração experimental de que a abordagem atinge fidelidade de reconstrução state-of-the-art (SOTA), consistência estrutural mais forte e representações latentes mais discriminativas em comparação com VAEs existentes.

---

4. Resultados Experimentais

O NeuroQuant foi avaliado em dois grandes conjuntos de dados multimodais: NCANDA e ABCD.

Métricas de Avaliação

• Fidelidade em Nível de Voxel: PSNR e SSIM 3D.
• Consistência Anatômica: Pontuação Dice baseada na segmentação SynthSeg (avalia volumes regionais e formas).
• Utilidade do Espaço Latente: Precisão na classificação de sexo usando um classificador treinado nos embeddings latentes.

Desempenho Quantitativo

• O NeuroQuant superou consistentemente os métodos SOTA (VQGAN, SD-VAE, MediTok, MedVAE) em todas as métricas.
• Exemplo (Dataset NCANDA, T1):
  • PSNR: 28.89 dB (NeuroQuant) vs. 28.26 dB (Melhor baseline MedVAE).
  • SSIM: 95.27 vs. 94.38.
  • Dice (Segmentação): 95.25 vs. 94.83.
  • Precisão de Classificação de Sexo: 79.09% vs. 78.12%.
• O modelo alcançou uma taxa de compressão espacial de 16x (vs. 8x em outros), mantendo uma qualidade superior, o que é crucial para modelos de difusão latente.

Análise Qualitativa

• Inconsistências 2D: Modelos baseados em 2D (como SD-VAE) apresentaram artefatos visíveis, como fronteiras corticais descontínuas e distorções em "escada" entre fatias, especialmente nos planos coronal e sagital.
• Suavização Excessiva: Modelos baseados em KL (como MedVAE) tendiam a suavizar detalhes anatômicos finos (dobras corticais, limites ventriculares).
• NeuroQuant: Produziu bordas mais nítidas, transições inter-fatias suaves e estruturas anatômicas coerentes em todas as orientações, preservando detalhes de alta frequência sem perder a coerência volumétrica.

Estudo de Ablação

• A remoção do treinamento conjunto 2D/3D degradou significativamente a qualidade (bordas mais borradas, menor consistência).
• A adição da atenção multi-eixo fatorada (FMA) e do decodificador FiLM melhorou substancialmente a modelagem estrutural de longo alcance e a fidelidade de textura.
• As perdas auxiliares (cruzada e adversarial) foram essenciais para o desentrelaçamento limpo entre estrutura e aparência.

---

5. Significado e Conclusão

O NeuroQuant representa um avanço significativo na modelagem de RM cerebral. Ao explicitamente separar a anatomia invariante da modalidade da aparência dependente da modalidade, o modelo resolve problemas fundamentais de coerência volumétrica e generalização que afetam os VAEs atuais.

Impactos Principais:

• Fundação Escalável: Fornece uma base robusta para modelagem generativa (como LDMs) em neuroimagem, permitindo a síntese de dados de alta qualidade para aumento de dados e exploração de mecanismos de doenças.
• Análise Transversal: A capacidade de gerar imagens de uma modalidade a partir de outra (síntese cross-modal) facilita a análise em estudos onde certas modalidades estão ausentes.
• Eficiência: A compressão 16x com alta fidelidade torna o processamento de grandes conjuntos de dados 3D mais viável para modelos de difusão e autoregressivos.

Em suma, o trabalho demonstra que a modelagem unificada multimodal, combinada com técnicas de quantização vetorial e atenção eficiente, é superior às abordagens fragmentadas ou baseadas em 2D para a representação de estruturas cerebrais complexas.