Resumen Técnico: BrainG3N

Planteamiento del Problema

La resonancia magnética (RM) cerebral tridimensional (3D) es crítica para la neurología clínica y la neurooncología, sin embargo, los modelos generativos capaces de operar a resolución 3D completa son computacionalmente prohibitivos en el espacio de vóxeles. En consecuencia, el campo ha convergido en la difusión latente, que depende de un tokenizador codificador-decodificador para comprimir volúmenes en un espacio latente de baja dimensión.

Existe una tensión fundamental en los procesos actuales de difusión latente radiográfica 3D:

1. Utilidad Clínica: Los embeddings del codificador deben retener información clínica específica requerida para tareas posteriores (p. ej., clasificación de enfermedades, predicción de biomarcadores).
2. Fidelidad Anatómica: El decodificador debe reconstruir volúmenes de vóxeles anatómicamente precisos.

Los enfoques actuales suelen entrenar un par codificador-decodificador único contra un objetivo de reconstrucción. Esto sesga al codificador hacia la fidelidad de los vóxeles locales, a menudo a expensas del contenido clínico global necesario para las tareas posteriores. Además, los espacios latentes existentes rara vez se evalúan sobre su capacidad para soportar tanto la generación controlable como la predicción clínica simultáneamente.

Metodología

Los autores proponen BrainG3N, un tokenizador de propósito dual y totalmente volumétrico que desacopla el codificador y el decodificador para satisfacer ambas demandas contrapuestas.

1. Corpus de Preentrenamiento

El modelo fue preentrenado con 35,309 volúmenes de RM cerebral de 17,399 sujetos únicos procedentes de 18 cohortes públicas y más de 200 sitios de adquisición. Los datos abarcan cuatro modalidades (T1, T2, FLAIR, T1c), diez categorías de enfermedades (incluyendo condiciones neurodegenerativas y tumorales) y edades que van desde los 5 hasta los 98 años. Todos los volúmenes fueron armonizados mediante registro afín ANTs al atlas SRI24, con extracción de cráneo (skull-stripping) y corrección de sesgo, resultando en volúmenes isotrópicos de 1 mm de 160×192×160.

2. Arquitectura del Tokenizador de Dos Fases

La innovación central es un codificador de Autoencoder de Máscara (MAE) congelado emparejado con un decodificador CNN dedicado mediante una proyección lineal.

• Fase 1 (Preentrenamiento): Un codificador MAE 3D (Transformer de Visión de 12 capas) es entrenado para reconstruir el 70% de los parches enmascarados aleatoriamente a partir del 30% de los parches visibles restantes. Esto obliga al codificador a modelar el contexto anatómico de largo alcance en lugar de la estadística de vóxeles locales.
• Fase 2 (Tokenización): El codificador MAE se mantiene congelado. Una proyección lineal ($P \in \mathbb{R}^{1152 \times 32}$) comprime los 1200 tokens por volumen de 1152 dimensiones a un cuello de botella de 32 canales. Un decodificador CNN 3D reconstruye entonces los vóxeles a partir de estos tokens proyectados bajo una pérdida $\ell_1$.
  • Razón de ser: Entrenar conjuntamente el codificador y el decodificador (como en los VAE) desvía el codificador hacia la fidelidad de la intensidad, mientras que el congelamiento del MAE preserva sus embeddings clínicamente informativos, mientras que el cuello de botella reduce la dimensionalidad de entrada del modelo de difusión.

3. Difusión Latente Condicional

Un Transformer de Difusión de Coincidencia de Flujo (DiT) es entrenado sobre las secuencias de tokens proyectados ($z' \in \mathbb{R}^{1200 \times 32}$).

• Condicionamiento: El DiT acepta seis variables de condicionamiento: enfermedad (8 vías), sexo, modalidad (4 vías), sitio de adquisición (19 vías), edad (continua) e estado de mutación IDH1 (binario).
• Mecanismo: Las condiciones se canalizan a través de la modulación adaLN-Zero. La Guía de Clasificador Libre (CFG) se implementa mediante el descarte aleatorio de condiciones (probabilidad $p=0.1$) durante el entrenamiento, excepto para la modalidad, que siempre se especifica.
• Variante Longitudinal: Una versión especializada reemplaza el interpolante de ruido-a-datos estándar con un puente latente de línea base-a-seguimiento, permitiendo el pronóstico específico de pacientes sobre horizontes temporales solicitados ($\Delta t$).

Contribuciones Clave

1. Espacio de Embedding de Propósito Dual: Los autores demuestran que un único codificador MAE congelado puede servir como tokenizador tanto para tareas clínicas posteriores como para la generación controlable.
2. Desacoplamiento Arquitectónico: Al congelar el codificador MAE y utilizar una proyección lineal + decodificador CNN, el método evita el compromiso entre el contenido clínico y la fidelidad de reconstrucción inherente al entrenamiento conjunto.
3. Preentrenamiento Escalable: El modelo fue preentrenado en un corpus masivo, multimodal y multisitio (más de 35k volúmenes), significativamente mayor que los conjuntos de datos generativos médicos 3D típicos.

Resultados Experimentales

1. Representación Clínica (Probing Lineal)

El codificador congelado fue evaluado en un benchmark de 23 tareas de probing lineal (sin ajuste fino del codificador).

• Rendimiento: El modelo superó o igualó a tres modelos fundacionales de vanguardia (BrainIAC, BrainSegFounder, MedicalNet) en 21 de las 23 tareas.
• Métricas Clave:
  • Estado de Mutación IDH1: AUC 0.937 (frente a 0.805 para BrainIAC).
  • Grado de Tumor: AUC 0.921.
  • Regresión de Edad Cerebral: Error Absoluto Medio (MAE) 4.43 años.
  • Predicción de Sexo: AUC 0.967.
• Comparación: El enfoque MAE congelado superó a un modelo basado en cortes 2D (BrainDINO) en tareas volumétricas, posicionando el enfoque de auto-supervisión totalmente volumétrico como superior para datos 3D.

2. Generación Controlable

Un DiT entrenado en estos embeddings fue evaluado por su capacidad para generar muestras condicionadas a variables específicas.

• Controlabilidad: Sondas de datos reales aplicadas a las muestras generadas confirmaron que el modelo recupera con éxito los atributos solicitados (p. ej., estado de enfermedad, sexo, edad) bajo Guía de Clasificador Libre (CFG).
  • Ejemplo: El condicionamiento de edad mostró una correlación de Pearson de $r=0.93$ entre la edad solicitada y la generada.
  • Contrafácticos: Mantener constantes las semillas de ruido mientras se variaban las condiciones produjo cambios visibles y anatómicamente consistentes (p. ej., agrandamiento ventricular en casos de Alzheimer) preservando la identidad.
• Fidelidad: El FID 3D agrupado entre las muestras generadas y los volúmenes reales fue de 107.3, competitivo con el trabajo generativo médico 3D previo. La brecha entre las muestras generadas y las reales se atribuyó principalmente a la compresión por parte del tokenizador de los detalles de alta frecuencia del escáner, más que a un fallo del generador.

3. Pronóstico Longitudinal

El modelo fue probado en la extrapolación temporal intra-paciente utilizando datos de ADNI.

• Configuración: Dado un latente de línea base y un horizonte temporal ($\Delta t$), el modelo pronostica un escaneo de seguimiento.
• Resultados: El modelo recuperó aproximadamente el 27% de la magnitud real del envejecimiento (pendiente 0.268) con una correlación de Pearson de $r=0.716$ entre el $\Delta t$ predicho y el solicitado.
• Localización: Los cambios anatómicos se localizaron correctamente en los ventrículos y surcos, consistentes con patrones conocidos de neurodegeneración, aunque la magnitud se vio atenuada debido al regularizador de puente estocástico que interpola hacia la media condicional.

Significado y Reivindicaciones

El artículo establece que un único espacio de embedding de RM cerebral 3D puede soportar simultáneamente:

1. Tareas clínicas de alto rendimiento (superando a otros modelos fundacionales en el probing lineal).
2. Modelado generativo controlable (permitiendo la generación condicional a través de seis variables y el pronóstico longitudinal específico de paciente).

Los autores argumentan que este enfoque valida el uso de codificadores MAE congelados como tokenizadores para la difusión radiográfica 3D, superando las limitaciones de los tokenizadores impulsados por la reconstrucción que sacrifican el contenido clínico en favor de la fidelidad de los vóxeles. El trabajo permite aplicaciones tales como la generación de datos sintéticos para cohortes subrepresentadas, la simulación de trayectorias de enfermedad contrafácticas y la provisión de lecturas basadas solo en imágenes para biomarcadores clínicos (p. ej., estado de IDH1, edad cerebral) sin requerir procedimientos invasivos.

Limitaciones señaladas por los autores:

• Los hallazgos son específicos de la RM cerebral; su transferencia a otras modalidades (TC, RM de cuerpo completo) no ha sido probada.
• El pronóstico longitudinal recupera la dirección y la localización del envejecimiento, pero solo el ~27% de la magnitud.
• La comparación con otros modelos fundacionales implica posibles desajustes de distribución debido a sus diferentes corpora de preentrenamiento.
• Los factores de confusión de la adquisición (fabricante del escáner, fuerza de campo) son detectables en los embeddings, lo que requiere una evaluación emparejada por cohorte para aplicaciones posteriores.