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

1. El Problema

La resonancia magnética (MRI) estructural multimodal (ej. T1, T2, FLAIR) es fundamental para el diagnóstico de enfermedades cerebrales, ya que los diferentes contrastes ofrecen información complementaria sobre la anatomía y la patología. Sin embargo, existen desafíos significativos:

• Limitaciones en la adquisición: Obtener MRI de alta calidad en múltiples modalidades para cada sujeto es costoso, requiere mucho tiempo y es propenso a artefactos por movimiento, lo que resulta en conjuntos de datos desequilibrados o incompletos.
• Deficiencias de los modelos actuales:
  • Los enfoques basados en VAE (Autoencoders Variacionales) existentes suelen centrarse en una sola modalidad (generalmente T1) o tratan las modalidades como entradas independientes, ignorando la correspondencia anatómica compartida.
  • Los modelos preentrenados en imágenes naturales (como los de Stable Diffusion) tienen una gran brecha de dominio con las imágenes médicas, produciendo reconstrucciones con inconsistencias anatómicas y sesgos de modalidad.
  • Los enfoques basados en 2D (que procesan corte por corte) fallan en capturar la coherencia volumétrica 3D, generando artefactos en las fronteras anatómicas entre cortes.
  • Los VAEs "universales" médicos a menudo no modelan explícitamente los contrastes complementarios dentro de la MRI cerebral multimodal.

Objetivo: Desarrollar un VAE 3D unificado que pueda aprender una representación latente compartida que disocié la estructura anatómica (invariante a la modalidad) de la apariencia dependiente de la modalidad (contraste e intensidad), permitiendo una síntesis y análisis de MRI cerebral de alta fidelidad.

2. Metodología: NeuroQuant

NeuroQuant es un Autoencoder Vectorial Cuantizado (VQ-VAE) 3D unificado diseñado específicamente para MRI cerebral multimodal. Su arquitectura se basa en tres pilares principales:

A. Arquitectura de Doble Flujo (Dual-Stream)

En lugar de usar codificadores separados para cada modalidad, NeuroQuant utiliza un codificador unificado que bifurca la información en dos ramas paralelas:

1. Codificación Anatómica ($z_{anat}$): Captura la estructura espacial compartida entre modalidades. Esta representación se discretiza utilizando un libro de códigos (codebook) vectorial compartido, creando un vocabulario compacto de "tokens anatómicos".
2. Codificación de Modalidad ($z_{mod}$): Captura las variaciones de intensidad y contraste específicas de cada modalidad (T1, T2, etc.).

B. Atención Multi-Eje Factorizada

Para manejar la alta resolución de los volúmenes 3D sin incurrir en el costo computacional prohibitivo de la atención 3D completa ($O((DHW)^2)$), el modelo emplea Atención Multi-Eje Factorizada.

• Descompone la atención global en operaciones secuenciales a lo largo de los ejes axial, coronal y sagital.
• Esto reduce la complejidad a $O(DHW(D+H+W))$, permitiendo modelar dependencias de largo alcance entre regiones cerebrales distantes manteniendo la coherencia espacial.

C. Decodificador con FiLM (Modulación Lineal por Características)

El decodificador reconstruye el volumen 3D utilizando:

• Los tokens anatómicos discretizados ($\tilde{z}_{anat}$) como base estructural.
• Los rasgos de modalidad ($z_{mod}$) para generar parámetros de FiLM ($\gamma, \beta$) mediante una red MLP.
• Estos parámetros FiLM modulan las activaciones del decodificador, adaptando dinámicamente el contraste y la intensidad sin alterar la estructura anatómica subyacente.

D. Estrategia de Entrenamiento 2D/3D Conjunta

Para abordar la adquisición basada en cortes de los datos 3D y mejorar la escalabilidad:

• El modelo se entrena mezclando volúmenes 3D completos y cortes 2D aleatorios (axiales, coronales o sagitales).
• Esto actúa como una augmentación de datos que mejora la generalización y la recuperación de detalles de alta frecuencia sin requerir codificadores adicionales.

E. Pérdidas de Desacoplamiento

Para asegurar que la estructura y la apariencia estén separadas, se utilizan pérdidas auxiliares:

• Pérdida de Reconstrucción Cruzada: Intercambia los parámetros FiLM entre modalidades para un mismo sujeto, forzando al código anatómico a ser invariante a la modalidad.
• Pérdida Adversarial: Un clasificador intenta predecir la modalidad a partir de $z_{anat}$; el codificador se entrena para engañar a este clasificador (reversión de gradiente), eliminando así la información de la modalidad de la representación anatómica.

3. Contribuciones Clave

1. NeuroQuant: Un marco VQ-VAE 3D unificado que modela conjuntamente la estructura anatómica y la apariencia específica de la modalidad en un espacio latente compartido.
2. Diseño Innovador: Introducción de un codificador de doble flujo con tokens anatómicos VQ y un decodificador FiLM, reforzado por atención multi-eje factorizada para una codificación estructural consistente.
3. Rendimiento Superior: Demostración experimental de que el enfoque logra una fidelidad de reconstrucción superior, mayor consistencia estructural y representaciones latentes más discriminativas en comparación con los VAEs existentes.

4. Resultados Experimentales

El modelo se evaluó en dos grandes conjuntos de datos multimodales: NCANDA y ABCD, comparándolo con VQGAN, SD-VAE, MediTok y MedVAE.

• Métricas Cuantitativas:
  • NeuroQuant superó consistentemente a los modelos baselines en PSNR (calidad de píxel), SSIM (similitud estructural) y Dice Score (consistencia anatómica tras segmentación con SynthSeg).
  • Logró una compresión espacial de 16x (vs. 8x en otros modelos) manteniendo una fidelidad superior.
  • En la clasificación de sexo basada en el espacio latente, NeuroQuant obtuvo la mayor precisión, indicando que sus representaciones latentes son semánticamente ricas y compactas.
• Análisis Cualitativo:
  • Los modelos basados en 2D mostraron inconsistencias estructurales y artefactos de "escalera" entre cortes.
  • Los modelos 3D basados en KL (como MedVAE) tendían a suavizar excesivamente los detalles finos (pliegues corticales).
  • NeuroQuant recuperó bordes más nítidos, transiciones inter-corte suaves y estructuras anatómicas coherentes en todas las orientaciones.
• Estudios de Ablación:
  • La eliminación del entrenamiento conjunto 2D/3D degradó significativamente la nitidez y la coherencia volumétrica.
  • La atención multi-eje (FMA) y el decodificador FiLM fueron críticos para mejorar la fidelidad y la discriminabilidad.
  • Las pérdidas de desacoplamiento (cruzada y adversaria) mejoraron la pureza de los códigos anatómicos.

5. Significado e Impacto

NeuroQuant representa un avance significativo en el modelado de neuroimagen:

• Fundamento Escalable: Proporciona una base sólida para el modelado generativo de MRI cerebral, permitiendo la síntesis de escaneos de alta calidad para aumentación de datos y análisis de contrafactuales.
• Eficiencia y Generalización: Al aprender una representación anatómica compartida y desacoplada, el modelo es más eficiente y generaliza mejor entre estudios y protocolos de adquisición diferentes.
• Aplicaciones Futuras: Facilita tareas downstream como la detección de enfermedades, la exploración de mecanismos de enfermedad y el análisis causal, al ofrecer representaciones latentes que separan claramente la anatomía de las variaciones de contraste.

En resumen, NeuroQuant resuelve el problema de la inconsistencia anatómica y la falta de coherencia volumétrica en los VAEs actuales mediante una arquitectura 3D consciente de la modalidad, estableciendo un nuevo estado del arte en la síntesis y análisis de MRI cerebral multimodal.