MHMamba: Multi-Head Mamba for 3D Brain Tumor Segmentation

Resumen Técnico: MHMamba para Segmentación de Tumores Cerebrales en 3D

Enunciado del Problema
La segmentación de tumores cerebrales en resonancia magnética (RM) 3D enfrenta desafíos significativos debido a la alta heterogeneidad de la morfología tumoral y el contraste multimodal. Si bien las Redes Neuronales Convolucionales (CNN) destacan en el modelado de texturas y límites locales, luchan con las dependencias de largo alcance debido a campos receptivos limitados. Por el contrario, los modelos basados en Transformers pueden capturar el contexto global, pero sufren una complejidad computacional y de memoria cuadrática ($O(N^2)$), lo que dificulta su aplicación a volúmenes 3D de alta resolución sin un submuestreo o ventaneo agresivo, lo cual puede provocar incoherencia contextual e inconsistencias en los límites. Además, los enfoques híbridos existentes a menudo padecen desajustes estadísticos entre las ramas convolucionales y de atención (por ejemplo, Normalización por Lotes/Instancia frente a Normalización por Capa), lo que conduce a un entrenamiento inestable y problemas de fusión de características.

Metodología
El artículo propone MHMamba (Mamba de Múltiples Cabezas), una arquitectura codificador-decodidor en forma de U diseñada para modelar características globales con complejidad lineal mientras preserva detalles de alta resolución. El marco consta de tres componentes principales:

1. Módulo Multi-Head Mamba (MHM):

   • Reemplaza los bloques estándar en el codificador con un Modelo de Espacio de Estados (SSM) de múltiples cabezas.
   • La dimensión de canales se divide en cabezas SSM paralelas ($N_h$), permitiendo que diferentes cabezas aprendan patrones de dependencia de largo alcance complementarios a partir de entradas multimodales sin asignaciones de modalidad fijas.
   • Cada cabeza serializa las características espaciales 3D en una secuencia, aplica actualizaciones de estado selectivas (Mamba) y fusiona las salidas mediante concatenación y proyección.
   • Este diseño mantiene la complejidad lineal mientras mejora las capacidades de modelado global y la estabilidad del entrenamiento bajo condiciones multimodales.

2. Módulo de Atención de Calibración Canal-Espacial (CSCA):

   • Diseñado para alinear las estadísticas de características y mejorar las respuestas de lesiones después del procesamiento de Mamba de múltiples cabezas.
   • Emplea un mecanismo de doble vía:
     • Vía de Canal: Utiliza agrupación global promedio y máxima seguida de una MLP para generar pesos de canal.
     • Vía Espacial: Integra innovadoramente cuatro estadísticas espaciales (media, desviación estándar, máximo y mínimo) para generar un mapa de pesos espaciales, capturando la heterogeneidad y los picos salientes.
   • Un mecanismo de puerta fusiona adaptativamente estas ramas, seguido de una conexión residual para preservar las respuestas de características originales.

3. Módulo de Fusión de Puerta Adaptativa (AGF):

   • Ubicado en las conexiones de salto entre el codificador y el decodificador.
   • Aborda el conflicto entre la semántica de alto nivel y los detalles de bajo nivel fusionando dinámicamente características multiescala.
   • Las características de entrada se dividen en subgrupos a lo largo de la dimensión de canales. Para cada subgrupo, el módulo aprende pesos de fusión independientes para equilibrar las características del codificador y del decodificador, evitando conflictos de características que a menudo conducen a predicciones de límites inconsistentes.

Contribuciones Clave

• Arquitectura Multi-Head Mamba: El artículo introduce un enfoque de SSM de múltiples cabezas para RM 3D que descompone los canales en cabezas paralelas. Esto mejora el modelado de largo alcance y la adaptabilidad multimodal mientras retiene la ventaja de complejidad lineal de Mamba.
• Calibración y Fusión de Características: La integración del módulo CSCA para unificar escalas y estadísticas de características, y del módulo AGF para fusionar dinámicamente características multiescala en las conexiones de salto, mejora significativamente la consistencia de los límites y la detección de lesiones de pequeño volumen.
• Validación Sistemática: El método fue evaluado rigurosamente en los conjuntos de datos BraTS2021 y BraTS2023, demostrando mejoras estables sobre bases sólidas (incluyendo nnU-Net, SegMamba y modelos basados en Transformers) en términos de puntuaciones Dice y Distancia de Hausdorff (HD95).

Resultados Experimentales
Los experimentos realizados en los conjuntos de datos BraTS2021 y BraTS2023 muestran que MHMamba logra un rendimiento de vanguardia o cercano a la vanguardia:

• BraTS2021: Logró una puntuación Dice media del 91.02% y un HD95 promedio de 3.38, superando a bases como SegMamba (90.32% Dice, 4.18 HD95) y nnU-Net (89.40% Dice, 5.58 HD95).
• BraTS2023: Logró una puntuación Dice media del 90.23% y un HD95 promedio de 3.67.
• Estudios de Ablación:
  • La adición individual de CSCA, AGF y el módulo Multi-Head Mamba contribuyó a ganancias de rendimiento en todas las subregiones tumorales (Tumor Completo, Núcleo Tumoral y Tumor Realzante).
  • La configuración óptima resultó ser 4 cabezas, lo que equilibró la capacidad representacional y la eficiencia computacional; aumentar a 8 cabezas condujo a una degradación del rendimiento debido a una sobre-partición y ruido de optimización en el entrenamiento 3D con tamaños de lote pequeños.
  • Los resultados cualitativos indicaron que MHMamba produce predicciones menos fragmentadas con límites más ajustados, particularmente en regiones tumorales heterogéneas y áreas realzantes pequeñas.

Significado y Afirmaciones
El artículo afirma que MHMamba aborda eficazmente la compensación entre el modelado de contexto global y la preservación de detalles locales en la segmentación de tumores cerebrales en 3D. Al aprovechar la complejidad lineal de los Modelos de Espacio de Estados a través de una arquitectura de múltiples cabezas, el método evita la sobrecarga de memoria cuadrática de los Transformers mientras supera las limitaciones del campo receptivo local de las CNN. Los módulos propuestos (CSCA y AGF) se dirigen específicamente a los desajustes estadísticos y la incoherencia contextual que a menudo se encuentran en la segmentación 3D multimodal, resultando en un marco robusto, eficiente y preciso adecuado para aplicaciones clínicas como la planificación de radioterapia y la evaluación preoperatoria. Los autores enfatizan que el método mantiene las ventajas de eficiencia del modelado basado en Mamba mientras mejora significativamente la sensibilidad a los núcleos tumorales y las lesiones de pequeño volumen.