Resumen Técnico: Patch-MoE Mamba

Enunciado del Problema

La segmentación de imágenes médicas depende en gran medida de las Redes Neuronales Convolucionales (CNN) y los Transformadores. Sin embargo, las CNN están inherentemente limitadas por campos receptivos locales, luchando para modelar dependencias de largo alcance, mientras que los Transformadores a menudo incurren en complejidad computacional y de memoria cuadrática. Recientemente, los Modelos de Espacio de Estados (SSM), particularmente las arquitecturas basadas en Mamba, han surgido como alternativas eficientes con complejidad de secuencia lineal.

A pesar de su promesa, los modelos de segmentación existentes basados en Mamba enfrentan dos limitaciones críticas:

1. Disrupción de la Estructura Espacial Local: La mayoría de los modelos convierten mapas de características 2D en secuencias 1D utilizando un escaneo rasterizado a nivel de píxel. Este proceso puede separar píxeles espacialmente adyacentes en la secuencia, debilitando la coherencia espacial local 2D. Esto es perjudicial para tareas de predicción densa que requieren una localización precisa de los límites y la preservación de la estructura local.
2. Fusión Direccional Inflexible: Los métodos actuales suelen fusionar las salidas de múltiples direcciones de escaneo mediante una suma simple. Esta agregación fija asume que todas las direcciones son igualmente importantes en cada ubicación espacial, fallando en adaptarse a los diversos tamaños, formas, texturas y complejidades de los límites de los objetivos médicos (por ejemplo, lesiones pequeñas y de bajo contraste frente a estructuras grandes e irregulares).

Metodología

Los autores proponen Patch-MoE Mamba, una arquitectura de espacio de estados de Mezcla de Expertos (MoE) ordenada por parches diseñada para abordar las limitaciones mencionadas anteriormente. La arquitectura sigue un diseño estilo U-Net, reemplazando los bloques estándar de Espacio de Estados Visual (VSS) en VM-UNetV2 con un nuevo bloque Patch-MoE VSS.

1. Escaneo Jerárquico Ordenado por Parches

En lugar de escanear píxeles individuales de forma rasterizada, el método divide los mapas de características en parches locales y realiza un escaneo direccional de manera ordenada por parches.

• Mecanismo: Dado un mapa de características, la cuadrícula espacial se divide en parches no superpuestos de tamaño $p \times p$. Dentro de cada parche, los píxeles se enumeran en orden de fila principal antes de pasar al siguiente parche.
• Beneficio: Esto asegura que todos los píxeles dentro de un parche permanezcan consecutivos en la secuencia, preservando los vecindarios espaciales locales mientras aún permite el modelado de dependencias de largo alcance entre parches.
• Estrategia Jerárquica: Para capturar contexto multiescala, el modelo emplea un conjunto de tamaños de parche $\{p^{(1)}, p^{(2)}, \dots\}$. Los parches más pequeños capturan detalles de límites de grano fino, mientras que los parches más grandes modelan estructuras más gruesas. Diferentes direcciones de escaneo (adelante, reverso, adelante ancho-alto, reverso ancho-alto) pueden utilizar tamaños de parche distintos para manejar estructuras anisotrópicas.

2. Fusión Direccional Basada en MoE

Para reemplazar la suma fija, los autores introducen un módulo de fusión adaptativo que trata las salidas direccionales como "expertos".

• Construcción de Expertos: El módulo utiliza cinco expertos:
  • Cuatro expertos específicos de dirección ($E_1$ a $E_4$): Las salidas normalizadas de los cuatro escáneres Mamba direccionales.
  • Un experto de concatenación ($E_{concat}$): Formado concatenando los cuatro mapas normalizados a lo largo de la dimensión del canal y proyectándolos de nuevo mediante una convolución $1\times1$. Este experto captura correlaciones cruzadas entre direcciones y proporciona una visión holística agnóstica a la dirección.
• Puerta Consciente del Espacio: Una red enrutadora calcula pesos dependientes de la ubicación para fusionar adaptativamente estos expertos. El enrutador se impulsa mediante una mezcla de contexto local (mediante convolución profunda) y contexto global (mediante agrupación promedio adaptativa), controlado por un escalar aprendible $\alpha$.
• Estabilización Residual: Un bypass residual añade las salidas direccionales crudas de nuevo al resultado fusionado ($Z_l = \hat{Y}_l + \sum Y_l^{(i)}$). Esto asegura respuestas base robustas y preserva señales direccionales fuertes, estabilizando el entrenamiento.

3. Infusión de Semántica y Detalle (SDI)

Las salidas del codificador son procesadas por un módulo SDI (adaptado de U-Net v2) para mejorar los mapas de características infundiendo información semántica de niveles superiores e integrando detalles más finos de niveles inferiores mediante el producto de Hadamard.

Contribuciones Clave

1. Escaneo Ordenado por Parches: Un mecanismo de escaneo novedoso que reemplaza el escaneo rasterizado a nivel de píxel con un recorrido ordenado por parches, preservando efectivamente la coherencia espacial local 2D mientras mantiene la eficiencia de los modelos de espacio de estados.
2. Fusión Direccional Basada en MoE: Un módulo de fusión adaptativo que reemplaza la suma simple con una red de puerta aprendible. Pondera dinámicamente múltiples expertos direccionales y un experto de concatenación basándose en el contexto espacial, permitiendo que el modelo se adapte a diversas escalas y complejidades de lesiones.
3. Modelado Multiescala Jerárquico: La integración de tamaños de parche jerárquicos permite que la red capture simultáneamente detalles de límites de grano fino y estructuras semánticas gruesas.
4. Estabilización Residual: La inclusión de una ruta residual para las salidas direccionales crudas asegura la estabilidad del entrenamiento y previene la pérdida de señales direccionales críticas.

Resultados Experimentales

Los autores evaluaron Patch-MoE Mamba en cinco conjuntos de datos públicos de segmentación de pólipos (Kvasir-SEG, ClinicDB, ColonDB, ETIS, CVC-300) y dos conjuntos de datos de segmentación de lesiones cutáneas (ISIC 2017, ISIC 2018).

• Rendimiento: Patch-MoE Mamba logró un rendimiento de vanguardia en todos los conjuntos de datos, superando consistentemente a las baselines basadas en CNN (U-Net, U-Net v2) y a las baselines basadas en Mamba (VM-UNet, VM-UNetV2).
  • En el desafiante conjunto de datos ETIS (conocido por su bajo contraste y pólipos pequeños), el modelo mejoró la puntuación Dice del 72.56% (VM-UNetV2) al 74.04%.
  • En ISIC 2018, logró una puntuación Dice de 89.34%, superando el mejor resultado anterior de 88.36%.
• Estudios de Ablación:
  • Reemplazar el escaneo rasterizado con escaneo ordenado por parches únicamente mejoró el Dice promedio del 83.46% al 84.02%.
  • Añadir el módulo de fusión MoE aumentó aún más el promedio al 84.30%.
  • Los tamaños de parche jerárquicos (gruesos en etapas tempranas, finos en etapas más profundas) produjeron los mejores resultados.
• Complejidad Computacional: Si bien el escaneo ordenado por parches en sí no aumenta el costo computacional (la cantidad de tokens permanece sin cambios), la inclusión del experto de concatenación aumenta los parámetros a ~70M y los FLOPs a ~28G. Los autores señalan que esta sobrecarga es esperada debido a las interacciones de canales más densas requeridas para el mecanismo MoE.

Significado y Afirmaciones

El artículo afirma que Patch-MoE Mamba ofrece una solución general y efectiva para la segmentación de imágenes médicas. Al abordar las limitaciones específicas de los modelos Mamba existentes, a saber, la pérdida de estructura espacial local y la rigidez de la fusión direccional, la arquitectura propuesta demuestra una adaptabilidad superior a diversos objetivos médicos con tamaños, formas y complejidades de límites variables. Los resultados sugieren que preservar los vecindarios espaciales locales mediante el escaneo ordenado por parches, combinado con la fusión adaptativa de expertos, mejora significativamente la localización de límites y la precisión de la segmentación, particularmente en escenarios desafiantes que involucran lesiones de bajo contraste o pequeñas.