DAMamba-UNet3D: A Parameter-Efficient Mamba State Space U-Net with Dynamic Adaptive Scan for 3D Medical Image Segmentation

Resumen Técnico: DAMamba-UNet3D

Planteamiento del Problema
La segmentación de imágenes médicas 3D requiere equilibrar la delimitación precisa de los bordes locales con un contexto espacial global coherente, todo ello cumpliendo con presupuestos estrictos de computación y parámetros. Las arquitecturas existentes enfrentan distintos compromisos:

• U-Nets Convolucionales: Ofrecen eficiencia de parámetros y mezcla local $O(n)$, pero carecen de un contexto global explícito, lo que requiere el escalado de la profundidad para capturar dependencias de largo alcance.
• Vision Transformers: Proporcionan razonamiento global mediante atención propia (self-attention), pero incurren en costos prohibitivos de memoria y cómputo $O(n^2)$ para volúmenes 3D de gran tamaño (por ejemplo, un parche de $128^3$ contiene aproximadamente 2.1 millones de tokens).
• Modelos de Espacio de Estados (SSMs/Mamba): Propagan información globalmente en tiempo y memoria $O(n)$. Sin embargo, los segmentadores médicos de SSM existentes (por ejemplo, SegMamba) dependen de patrones de escaneo fijos y cuentan con presupuestos de parámetros elevados ($\sim 70$ millones de parámetros), careciendo de adaptabilidad dependiente de los datos.

Aunque se ha introducido el Escaneo Adaptativo Dinámico (DAS) para el reconocimiento de imágenes naturales 2D para aprender el reordenamiento dependiente de los datos antes del escaneo selectivo, este no se ha aplicado a las imágenes médicas ni se ha extendido a volúmenes 3D.

Metodología
Los autores proponen DAMamba-UNet3D, una arquitectura híbrida codificador-decodificador que integra bloques Tri-Plane 3D-DAS dentro de un marco de trabajo U-Net.

1. Diseño de la Arquitectura:

   • Estructura Híbrida: El modelo conserva bloques convolucionales para el tallo (E1), el cuello de botella (bottleneck) y todo el decodificador para preservar los detalles locales.
   • Integración del Codificador: Los bloques Tri-plane DAMamba se insertan exclusivamente en las etapas del codificador E2–E4.
   • Variante Compacta: El modelo principal utiliza un ancho base de $C=16$, lo que resulta en aproximadamente 5.3 millones de parámetros ($\sim 13\times$ menos que SegMamba).
   • Variante de Gran Escala (DAMamba-L): Una variante más ancha ($\sim 70$ millones de parámetros) con bloques DAMamba apilados en el codificador y un cuello de botella convolucional, diseñada para igualar la capacidad de SegMamba para una comparación directa.

2. Escaneo Adaptativo Dinámico Tri-Plano (Tri-Plane DAS):

   • Mecanismo: A diferencia de los escaneos fijos (orden de filas, zigzag o múltiples orientaciones), el DAS aprende el reordenamiento espacial dependiente de los datos. Una Red de Predicción de Desplazamiento (OPN) predice desplazamientos 2D a partir del mapa de características, los cuales se utilizan para remuestrear las características mediante interpolación bilineal antes del escaneo selectivo.
   • Extensión 3D: El operador DAS 2D se extiende a volúmenes aplicando este a través de tres planos anatómicos: Axial ($H \times W$), Coronal ($D \times W$) y Sagital ($D \times H$). Las salidas se promedian para formar el bloque 3D DAS:
     $$DAS3D(X) = \frac{1}{3}(DAS_{ax} + DAS_{cor} + DAS_{sag})$$
   • Escaneo Selectivo: El mapa de características reordenado es procesado por un bloque de escaneo selectivo Mamba estándar, permitiendo la propagación del contexto global con un costo de $O(N)$ sin mapas de atención cuadráticos.

3. Entrenamiento y Evaluación:

   • Conjunto de Datos: Conjunto de entrenamiento BraTS 2020 (369 sujetos, 4 modalidades).
   • Pérdida: Modulated Combo Loss ($L = (L_{CE} + \delta) \cdot (L_{mod}^{Dice} + \delta)$) para enfocarse en vóxeles difíciles y bordes.
   • Protocolo: Validación cruzada de cinco pliegues (five-fold) a nivel de sujeto con inferencia de volumen completo (ventanas deslizantes, parches de $128^3$). Las métricas incluyen el Dice medio y HD95 para el Tumor Enriquecido (ET), Núcleo Tumoral (TC) y Tumor Completo (WT).

Contribuciones Clave

1. Primera Aplicación Médica de DAS: Transfiere el Escaneo Adaptativo Dinámico del reconocimiento de clasificación de imágenes naturales 2D a la segmentación densa de imágenes médicas 3D.
2. Extensión de 2D a 3D: Eleva el operador DAS a datos volumétricos mediante el escaneo adaptativo tri-plano, ofreciendo una alternativa aprendible al aplanamiento de múltiples orientaciones diseñado a mano.
3. Eficiencia de Parámetros: Demuestra un U-Net híbrido con bloques DAS tri-plano solo en el codificador, logrando una paridad cercana con modelos mucho más grandes con un costo de parámetros $\sim 13\times$ menor.
4. Validación de Dos Escalas: Valida el enfoque tanto en una escala compacta ($\sim 5.3$M) como en una escala grande ($\sim 70$M), mostrando que el DAS aprendido puede superar a los patrones de escaneo fijos cuando la colocación arquitectónica está optimizada.

Resultados

• Modelo Compacto (DAMamba-UNet3D, $\sim 5.3$M): Logró un Dice medio de 0.815 $\pm$ 0.013 en BraTS 2020. Esto es casi par con el SegMamba reentrenado (0.824 $\pm$ 0.014) a pesar de tener $\sim 13\times$ menos parámetros.
• Modelo de Gran Escala (DAMamba-L, $\sim 70$M): Superó al SegMamba reentrenado con un Dice medio de 0.829 $\pm$ 0.012 (una mejora de +0.5 puntos), observándose la mayor ganancia en la clase de Núcleo Tumoral (TC) (0.842 vs. 0.827).
• Estudios de Ablación:
  • Colocación: Colocar los bloques DAS solo en el codificador (E2–E4) produjo los mejores resultados (0.815). Añadir DAS al cuello de botella o al decodificador redujo el rendimiento (0.807 y 0.805, respectivamente), indicando que la colocación solo en el codificador es crítica.
  • Patrones de Escaneo: El DAS tri-plano aprendido superó a los órdenes de escaneo fijos (Raster, Snake, Cross, Tri-orientated) cuando la colocación en el codificador fue emparejada.

Significado y Reivindicaciones
El artículo afirma que el DAS tri-plano aprendido en un U-Net híbrido es una alternativa competitiva al escaneo fijo de Mamba de Orientación Tripartita (ToM) para la segmentación de tumores cerebrales 3D.

• Eficiencia: El modelo compacto demuestra que se puede lograr un alto rendimiento con recuentos de parámetros significativamente reducidos aprovechando el escaneo adaptativo en lugar de la capacidad bruta.
• Rendimiento: A gran escala, el mecanismo adaptativo aprendido (DAS) puede mejorar sobre las trayectorias de escaneo fijas (ToM) bajo protocolos idénticos de entrenamiento y evaluación.
• Modestia: Los autores advierten que, aunque los resultados en BraTS 2020 son prometedores, el estudio se limita a un solo conjunto de datos. Declaran que "se necesita validación en múltiples conjuntos de datos y pruebas estadísticas antes de que se puedan hacer afirmaciones generales sobre DAS frente a ToM". Su principal contribución es la extensión de DAS a volúmenes médicos 3D y la demostración de su viabilidad en una arquitectura híbrida eficiente en parámetros.