BrainAnytime: Anatomy-Aware Cross-Modal Pretraining for Brain Image Analysis with Arbitrary Modality Availability

Resumen Técnico de "BrainAnytime: Preentrenamiento Consciente de la Anatomía para Análisis de Imágenes Cerebrales con Disponibilidad Arbitraria de Modalidades"

Planteamiento del Problema

Los flujos de trabajo de diagnóstico clínico para enfermedades neurodegenerativas, particularmente la Enfermedad de Alzheimer (EA), siguen una "ruta de escalada de modalidades". Los clínicos suelen comenzar con imágenes estructurales de rutina (por ejemplo, RM ponderada en T1), añaden selectivamente secuencias (por ejemplo, FLAIR, T2) para refinar los diagnósticos diferenciales y reservan la imagen molecular (por ejemplo, PET de amiloide) para casos inciertos. En consecuencia, los datos de los pacientes en cohortes del mundo real se caracterizan por subconjuntos de modalidades heterogéneos y a menudo incompletos.

Los modelos de IA actuales asumen generalmente modalidades de entrada fijas, lo que requiere un cambio arquitectónico o un reentrenamiento cuando cambia la disponibilidad de datos. Además, los modelos existentes de preentrenamiento a gran escala para imágenes cerebrales suelen sufrir tres limitaciones en este contexto:

1. Suelen ser solo de RM, fallando en aprender puentes explícitos hacia el PET de amiloide, un biomarcador crítico.
2. Asumen combinaciones de modalidades fijas, lo que no coincide con la ausencia dependiente de la etapa de los datos del mundo real.
3. El preentrenamiento estándar con Autoencoders Enmascarados (MAE) se basa en un enmascaramiento espacialmente uniforme, ignorando el hecho de que la EA presenta vulnerabilidad regional selectiva (por ejemplo, el hipocampo) en lugar de patología uniforme.

Metodología

Los autores proponen BrainAnytime, un marco de preentrenamiento unificado diseñado para apoyar el análisis de imágenes cerebrales bajo disponibilidad arbitraria de modalidades. El marco se preentrena en 34.899 escaneos cerebrales 3D de cinco conjuntos de datos (A4, DOD-ADNI, AIBL, BraTS23, NACC) que abarcan RM de múltiples secuencias y PET de amiloide. Integra tres componentes principales dentro de una arquitectura compartida de autoencoder enmascarado 3D:

1. Multi-MAE3D (Autoencoder Enmascarado 3D Multimodal)

• Arquitectura: Un codificador Transformer compartido procesa cualquier subconjunto de modalidades disponibles (T1, T2, FLAIR, PET).
• Tokenización: Cada modalidad utiliza una incrustación de parches convolucionales 3D específica de la modalidad.
• Manejo de Datos Faltantes: Las modalidades faltantes se enmascaran completamente. Para mantener un recuento de tokens constante independientemente de la disponibilidad de entrada, el presupuesto de tokens visibles se redistribuye entre las modalidades observadas utilizando una distribución Dirichlet (Dir($\alpha$)).
• Decodificación: Cada modalidad tiene un decodificador ligero que reconstruye todos los parches. La pérdida de entrenamiento ($L_{MAE}$) se calcula solo en los parches enmascarados utilizando el error cuadrático medio con normalización por parche.

2. Destilación Recíproca Intermodal (RCMD)

• Objetivo: Aprender correspondencias entre la RM estructural y el PET molecular.
• Mecanismo: Una red maestra con Promedio Móvil Exponencial (EMA) comparte la arquitectura del codificador. Para muestras emparejadas de RM-PET, las representaciones a nivel de grupo ($z_{RM}$ y $z_{PET}$) se calculan mediante un promedio de agrupación.
• Pérdida: Dos predictores MLP simétricos ($g_{RM\to PET}$ y $g_{PET\to RM}$) se entrenan para predecir características del maestro intermodal. La pérdida ($L_{RCMD}$) maximiza la similitud del coseno entre las características predichas por el estudiante y las características del maestro en el espacio de la modalidad opuesta.

3. Enmascaramiento Curricular Consciente de la Patología (PACM)

• Objetivo: Priorizar la anatomía vulnerable a la enfermedad durante la reconstrucción en lugar de enmascarar uniformemente.
• Puntuación de Importancia: Se deriva una puntuación estática a partir de priores neuropatológicos (estadificación de Braak) mapeados mediante el atlas AAL116. Una puntuación dinámica se obtiene agregando mapas de atención de un maestro de anatomía EMA. Estas se combinan para generar puntuaciones de importancia de parches.
• Programa Curricular: Se emplea una estrategia de tres fases:
  1. Etapa temprana: Enmascaramiento aleatorio uniforme.
  2. Etapa media: El parámetro de temperatura ($\tau$) se recocido mediante decaimiento coseno, afilando progresivamente la distribución de la máscara para favorecer regiones de alta importancia para el enmascaramiento (es decir, forzando al modelo a reconstruir regiones críticas).
  3. Etapa tardía: Temperatura fija.
• Selección: Los parches se seleccionan utilizando el truco Gumbel-top-k, asegurando que los parches de alta importancia tengan más probabilidades de ser enmascarados y, por lo tanto, reconstruidos.

Entrenamiento y Ajuste Fino:
La pérdida total es $L = L_{MAE} + \lambda \cdot \mathbb{1}_{[emparejado]} \cdot L_{RCMD}$. Para tareas aguas abajo, solo se utiliza el codificador estudiante con una cabeza de tarea ligera. Las modalidades faltantes permanecen en cero y bloqueadas por atención durante la inferencia, asegurando que el modelo maneje subconjuntos arbitrarios sin cambios arquitectónicos.

Contribuciones Clave

1. Marco Unificado: Un único modelo capaz de aceptar cualquier subconjunto de modalidades (desde una sola exploración T1 hasta un estudio multimodal completo) sin reentrenamiento ni cambio arquitectónico.
2. Aprendizaje Intermodal: Aprende explícitamente correspondencias estructura-molecular entre RM y PET mediante destilación recíproca, cerrando la brecha entre la evidencia estructural y patológica.
3. Preentrenamiento Consciente de la Anatomía: Introduce una estrategia de enmascaramiento curricular que dirige el esfuerzo reconstructivo del modelo hacia la neuroanatomía saliente para la EA, abordando las limitaciones del enmascaramiento uniforme en los MAE estándar.
4. Preentrenamiento Escalable: Se preentrenó con éxito en un conjunto de datos heterogéneo a gran escala (34.899 escaneos) que imita los patrones de ausencia clínica.

Resultados Experimentales

El modelo se evaluó en cuatro tareas aguas abajo (CN vs. EA, CN vs. ECI, regresión MMSE, predicción de edad) en cinco configuraciones de modalidades motivadas clínicamente (desde solo T1 hasta T1+T2+FLAIR+PET).

• Rendimiento: BrainAnytime superó a modelos específicos de modalidad (por ejemplo, ResNet 3D, M3T), métodos de fusión de doble modalidad, modelos fundacionales cerebrales a gran escala (BrainIAC, SAM-Brain3D) y enfoques de aprendizaje multimodal incompleto (Flex-MoE, FuseMoE).
• Métricas Clave:
  • CN vs. EA: Logró una precisión promedio del 89,4% (frente al 84,2% del siguiente mejor punto de referencia, Flex-MoE), lo que representa una mejora relativa del 6,2% sobre el punto de referencia de modalidad faltante más fuerte.
  • CN vs. ECI: Logró una precisión promedio del 67,2% (frente al 60,6% de MoE-retriever), una mejora relativa del 7,0%.
  • Regresión: Demostró un rendimiento superior en la predicción de MMSE y edad en todas las configuraciones de modalidades.
• Robustez: Un estudio de robustez intra-sujeto confirmó que el rendimiento mejoró consistentemente a medida que se añadían más modalidades (por ejemplo, la F1 de CN vs. EA aumentó del 65,0% con solo T1 al 80,4% con entrada completa), demostrando la capacidad del modelo para aprovechar eficazmente los datos disponibles.
• Ablación: Eliminar RCMD o PACM individualmente resultó en caídas de rendimiento, confirmando la necesidad de tanto la destilación intermodal como el enmascaramiento consciente de la anatomía.

Significado

El artículo afirma que BrainAnytime aborda una brecha crítica entre la realidad clínica (adquisición de datos escalonada e incompleta) y los paradigmas actuales de IA (modelos de entrada fija). Al permitir que un único modelo acepte flexiblemente cualquier imagen disponible, se alinea con la ruta de escalada de modalidades clínica. El marco demuestra que el preentrenamiento con destilación intermodal y enmascaramiento consciente de la patología produce un modelo fundacional robusto que supera significativamente a los puntos de referencia existentes, particularmente en escenarios donde los datos están incompletos o faltan modalidades. Este enfoque ofrece una solución práctica para implementar la IA en entornos clínicos reales donde los datos multimodales completos rara vez están disponibles para cada paciente.