Rhamba: Region-Aware Hybrid Attention-Mamba Framework for Self-Supervised Learning in Resting-State fMRI

Resumen Técnico: Rhamba – Marco de Atención Híbrida Consciente de la Región basado en Mamba

Enunciado del Problema
El aprendizaje auto-supervisado (SSL) ofrece una vía prometedora para el análisis de neuroimagen a gran escala, sin embargo, los enfoques existentes enfrentan limitaciones significativas. Los métodos actuales a menudo carecen de una integración explícita de la estructura neuroanatómica, dependiendo de enmascaramiento aleatorio de vóxeles o estructuras de grafos genéricas que pueden sacrificar información espacial de alta resolución. Además, las decisiones arquitectónicas para modelar datos de resonancia magnética funcional en estado de reposo (rs-fMRI) presentan una compensación: los modelos basados en Transformers capturan interacciones globales pero sufren de complejidad cuadrática con secuencias largas, mientras que los Modelos de Espacio de Estados (SSM) como Mamba ofrecen eficiencia en tiempo lineal pero pueden carecer de la flexibilidad para modelar interacciones globales complejas. Además, la relación entre los objetivos de reconstrucción del preentrenamiento y el rendimiento en tareas posteriores permanece poco explorada, particularmente en lo que respecta a cómo las estrategias de enmascaramiento conscientes de la región influyen en la calidad de la representación.

Metodología
Los autores proponen Rhamba, un marco de preentrenamiento consciente de la región diseñado para el análisis de rs-fMRI a gran escala. La metodología se basa en tres principios de diseño fundamentales:

1. Tokenización y Enmascaramiento Conscientes de la Región:

   • Alineación Anatómica: El marco utiliza el atlas de Etiquetado Anatómico Automatizado (AAL3) para alinear volúmenes fMRI 4D (normalizados al espacio MNI152) con Regiones de Interés (ROI) definidas anatómicamente.
   • Generación de Parches: Los volúmenes se dividen en parches de vóxeles no superpuestos de $6 \times 6 \times 6$.
   • Estrategias de Enmascaramiento: Se implementan tres estrategias de enmascaramiento distintas basadas en el grado de homogeneidad anatómica dentro de un parche:
     • Cualquiera (Any): Un parche se enmascara si contiene al menos un vóxel perteneciente a una macro-región específica.
     • Mayoría: Un parche se enmascara solo si más del 50% de sus vóxeles pertenecen a la macro-región objetivo.
     • Puro: Un parche se enmascara solo si la etiqueta anatómica dominante representa al menos el 70% de los vóxeles etiquetados, asegurando una alta homogeneidad espacial.

2. Arquitectura Híbrida de Atención–Mamba:

   • La columna vertebral se basa en NeuroSTORM, un modelo fundamental que utiliza bloques de Mamba con Ventana Desplazada (SWM).
   • Los autores experimentan con cuatro variantes arquitectónicas para determinar la integración óptima de mecanismos de atención y espacio de estados:
     • Solo Mamba: Modelado puro de espacio de estados.
     • Alternado: Bloques de Mamba y Atención intercalados.
     • Atención–Mamba (AM): Bloques de Atención en el codificador, Mamba en el decodificador.
     • Mamba–Atención (MA): Bloques de Mamba en el codificador, Atención en el decodificador.
   • El modelo emplea un objetivo de autoencoder enmascarado (MAE), reconstruyendo parches espacio-temporales faltantes durante el preentrenamiento.

3. Pipeline de Entrenamiento y Evaluación:

   • Preentrenamiento: Realizado en el conjunto de datos ABIDE (1.041 sujetos) utilizando las tres estrategias de enmascaramiento y las cuatro variantes arquitectónicas.
   • Ajuste Fino: Los modelos preentrenados se adaptaron para tareas de clasificación binaria (Esquizofrenia vs. Control; TDAH vs. Control) utilizando los conjuntos de datos COBRE y ADHD-200.
   • Interpretabilidad: Se utilizaron Gradientes Integrados con SmoothGrad al cuadrado (IG-SQ) para generar mapas de atribución, identificando las regiones cerebrales que contribuyen a las predicciones del modelo.

Contribuciones Clave

• Marco Rhamba: Introducción de una arquitectura híbrida consciente de la región que integra un enmascaramiento guiado anatómicamente con columnas vertebrales híbridas de Atención–Mamba para rs-fMRI.
• Enmascaramiento Guiado Anatómicamente: Propuesta de un enfoque sistemático para la tokenización y el enmascaramiento utilizando parcelaciones AAL3, avanzando más allá del enmascaramiento aleatorio de vóxeles hacia perturbaciones estructuradas y alineadas con regiones.
• Análisis Arquitectónico: Demostración de que intercalar mecanismos de espacio de estados y atención permite un modelado eficiente de secuencias largas mientras se preservan las interacciones funcionales globales.
• Evaluación Comparativa: Evaluación exhaustiva frente a modelos fundamentales de última generación, incluyendo SwiFT (basado en Transformers) y NeuroSTORM (basado en Mamba).

Resultados

• Estrategia de Enmascaramiento vs. Reconstrucción: Hubo un ordenamiento claro en la pérdida de reconstrucción basado en la especificidad del enmascaramiento: Cualquiera > Mayoría > Puro. Un enmascaramiento más estricto (Puro) resultó en una pérdida de reconstrucción más baja, indicando que los parches homogéneos son más fáciles de reconstruir.
• Estrategia de Enmascaramiento vs. Rendimiento Posterior: La tendencia de la pérdida de reconstrucción no se tradujo directamente en el rendimiento de clasificación posterior. Las diferencias entre las estrategias de enmascaramiento fueron modestas y dependientes del conjunto de datos. Por ejemplo, en el conjunto de datos ADHD-200, las estrategias "Cualquiera" y "Puro" produjeron un rendimiento comparable, mientras que en COBRE, "Cualquiera" y "Mayoría" fueron similares.
• Rendimiento Arquitectónico: La configuración Mamba–Atención (MA) logró consistentemente el AUROC promedio más alto en ambos conjuntos de datos.
  • COBRE: Rhamba-MA (Mayoría Occipital) logró un 71.32% de Precisión y un 68.44% de AUROC, superando a SwiFT (58.38% / 61.30%) y a las variantes de NeuroSTORM.
  • ADHD-200: Rhamba-MA (Mayoría Límbica) logró un 71.44% de AUROC, superando a SwiFT (69.06%) y a NeuroSTORM, aunque la precisión fue comparable a la de SwiFT.
• Análisis por Región: El rendimiento máximo dependió de la interacción entre la estrategia de enmascaramiento y la arquitectura en lugar de una única región dominante. Por ejemplo, la región Temporal tuvo el mejor rendimiento con MA en COBRE, mientras que la región Límbica mostró un rendimiento máximo con MA en ADHD-200.
• Interpretabilidad: Los mapas de atribución destacaron regiones consistentes con los perfiles neurobiológicos conocidos para los respectivos trastornos (por ejemplo, redes DMN y frontoparietales para TDAH; redes DMN y tálamo-corticales para Esquizofrenia).

Significado y Afirmaciones
El artículo afirma que Rhamba ofrece un marco flexible para equilibrar la interpretabilidad, la escalabilidad y el rendimiento en el aprendizaje de representaciones de fMRI a gran escala. Los hallazgos clave incluyen:

• Reconstrucción vs. Representación: La dificultad de reconstrucción no es un proxy confiable para la calidad de la representación. El enmascaramiento heterogéneo (Cualquiera), a pesar de una mayor pérdida de reconstrucción, promueve la emergencia de características más transferibles al fomentar la integración inter-regional.
• Superioridad Híbrida: La configuración MA subraya la importancia de combinar el modelado de contexto global (mediante Atención) con dinámicas de secuencia eficientes (mediante Mamba) para capturar tanto las interacciones de redes distribuidas como las dependencias temporales de largo alcance.
• Robustez: El preentrenamiento consciente de la región sigue siendo efectivo incluso en conjuntos de datos heterogéneos y multi-sitio (como ADHD-200) donde la alineación anatómica varía, lo que sugiere que el marco captura patrones funcionalmente relevantes más allá de la correspondencia anatómica estricta.
• Dirección Futura: El trabajo sugiere que los objetivos de preentrenamiento y el diseño arquitectónico deben considerarse conjuntamente. Propone que la auto-supervisión informada anatómicamente sirve como una herramienta principista para sondear la organización cerebral funcionalmente relevante, siempre que exista una distinción clara entre el comportamiento de aprendizaje de representaciones y los mecanismos neurobiológicos subyacentes.