Resumen Técnico: Mamba Goes HoME

Planteamiento del Problema

La segmentación de imágenes médicas tridimensionales (3D) es crítica para el diagnóstico y la planificación del tratamiento en diversas modalidades como la Tomografía Computarizada (TC), la Resonancia Magnética (RM) y la Ecografía (US). Sin embargo, las arquitecturas existentes enfrentan compromisos significativos:

• Las Redes Neuronales Convolucionales (CNN) ofrecen una complejidad computacional lineal pero sufren de campos receptivos limitados, lo que dificulta la captura de patrones espaciales globales esenciales para comprender la relación jerárquica entre lesiones locales y estructuras anatómicas globales.
• Los Vision Transformers (ViTs) utilizan la atención global para modelar dependencias de largo alcance, pero incurren en una complejidad computacional cuadrática con respecto al número de tokens. Esto los hace computacionalmente prohibitivos para datos volumétricos 3D de alta resolución.
• Los Modelos de Espacio de Estados (SSM) Selectivos, como Mamba, proporcionan una complejidad lineal para el modelado de secuencias largas, ofreciendo una alternativa eficiente a los ViTs. Sin embargo, los SSM estándar no están diseñados intrínsecamente para manejar de forma adaptativa los diversos patrones locales presentes en los complejos datos médicos multiescala.
• Los marcos de Trabajo de Mezcla de Expertos (MoE) pueden enrutar dinámicamente las características hacia subredes especializadas para manejar patrones locales, pero la combinación de la eficiencia global de los SSM con la adaptabilidad localizada de MoE en la imagen médica 3D permanece en gran medida inexplorada.

Metodología: Mamba-HoME

Los autores proponen Mamba-HoME, una arquitectura novedosa que integra la columna vertebral del Modelo de Espacio de Estados (SSM) Mamba Selectivo con una capa de Mezcla de Expertos Suave Jerárquica (HoME). El diseño tiene como objetivo modelar eficientemente las jerarquías de patrones de local a global con una complejidad computacional lineal.

1. Mezcla de Expertos Suave Jerárquica (HoME)

La innovación central es una capa de enrutamiento de tokens de dos niveles diseñada para equilibrar la especialización local con la integración del contexto global:

• Nivel 1: Agrupación y Enrutamiento Local. Las secuencias de entrada se particionan en grupos locales. Una capa de Mezcla de Expertos Suave (SMoE) enruta los tokens dentro de estos grupos hacia expertos locales especializados. Esto permite la extracción de características localizadas mientras preserva la localidad espacial y reduce la sobrecarga de memoria en comparación con el enrutamiento global.
• Nivel 2: Agregación Global. Las salidas del primer nivel se agregan y se pasan a través de una segunda capa SMoE global. Esto facilita la fusión de información entre grupos y el refinamiento del contexto global.
• Escalabilidad: Al realizar el enrutamiento dentro de los grupos antes de la agregación global, el método evita los costos de memoria prohibitivos del enrutamiento global en datos 3D de alta resolución (donde el recuento de tokens puede alcanzar millones), manteniendo una complejidad computacional de $O(Nd)$.

2. Bloque Mamba-HoME

La arquitectura incrusta la capa HoME en una red codificadora-decodificadora en forma de U (inspirada en SegMamba):

• Columna Vertebral (Backbone): Utiliza capas Mamba para un modelado eficiente de dependencias de largo alcance con complejidad de tiempo lineal.
• Convolución Espacial con Compuerta (GSC): Integrada para capturar prioridades espaciales locales antes de aplanar los datos en secuencias 1D para el procesamiento de Mamba.
• Tangente Hiperbólica Dinámica (DyT): Reemplaza la Normalización de Capa estándar con una no linealidad punto a punto ($f(x) = w \cdot \tanh(\alpha \cdot x) + b$). Esto estabiliza los gradientes y reduce la sobrecarga computacional al evitar los cálculos de media y varianza, acelerando el entrenamiento y la inferencia.
• Progresión Jerárquica: El número de expertos aumenta monótonamente con la profundidad ($E_1 < E_2 < \dots$), mientras que los tamaños de grupo disminuyen ($K_1 > K_2 > \dots$), permitiendo un procesamiento progresivamente más fino a medida que la red se profundiza.

Principales Contribuciones

1. Capa HoME: Introducción de una capa de enrutamiento de tokens MoE de dos niveles que captura eficientemente las jerarquías de patrones de local a global mediante la agrupación de tokens para el enrutamiento de expertos locales seguido del refinamiento de expertos globales.
2. Arquitectura Unificada: Diseño del Bloque Mamba-HoME, que integra los SSM de Mamba con HoME, combinando el procesamiento de secuencias largas con eficiencia de memoria y el enrutamiento de expertos jerárquicos.
3. Red de Segmentación de Múltiples Etapas: Embebimiento de estos componentes en una arquitectura en forma de U diseñada específicamente para la segmentación de imágenes médicas 3D, mejorando la representación contextual a través de las dimensiones espaciales y de profundidad.

Resultados Experimentales

Los autores evaluaron Mamba-HoME en cuatro conjuntos de datos públicos (PANORAMA, AMOS, FeTA 2022, MVSeg) y un conjunto de datos de TC propio, cubriendo las modalidades de TC, RM y US.

• Desempeño: Mamba-HoME logró resultados de vanguardia (SOTA) en precisión de segmentación (Coeficiente de Similitud de Dice - DSC) y precisión de contorno (Distancia de Hausdorff 95% - HD95) en todas las modalidades probadas.
  • En el conjunto de datos PANORAMA, alcanzó un DSC de 56.7% (pre-entrenado) y 54.8% (desde cero), superando a modelos base como SegMamba y Swin UNETR.
  • En el conjunto de datos de RM fetal FeTA 2022, alcanzó un DSC de 87.7% (pre-entrenado).
  • En el conjunto de datos de US 3D MVSeg, alcanzó un DSC de 85.0% (pre-entrenado).
• Eficiencia: A pesar de tener un recuento de parámetros mayor (170.1M) en comparación con algunos competidores, el modelo demostró un bajo uso de memoria de GPU durante la inferencia, una ventaja crítica para datos 3D de alta resolución. Aunque la inferencia fue aproximadamente un 30% más lenta que algunos modelos base, los autores señalan que este es un compromiso atractivo por las ganancias significativas en precisión.
• Generalizabilidad: El modelo demostró una fuerte generalización cross-modal. Un modelo pre-entrenado en datos de TC y RM se adaptó con éxito a tareas de ecografía 3D con un ajuste fino mínimo, superando a los métodos SOTA tanto en métricas DSC como en HD95.
• Estudios de Ablación:
  • Se encontró el desempeño óptimo con conteos de expertos específicos ($E_1 \in \{4, 8, 12, 16\}$, $E_2 \in \{8, 16, 24, 32\}$) y tamaños de grupo ($K \in \{2048, 1024, 512, 256\}$).
  • La Tangente Hiperbólica Dinámica (DyT) mejoró la velocidad de entrenamiento e inferencia en un ~6% en comparación con la Normalización de Capa sin sacrificar la precisión.
  • El pre-entrenamiento supervisado en conjuntos de datos de TC y RM a gran escala mejoró significamente el desempeño y la robustez a través de diversas regiones anatómicas y modalidades.

Significado y Afirmaciones

El artículo afirma que Mamba-HoME representa un avance significativo en la segmentación de imágenes médicas 3D al cerrar eficazmente la brecha entre la eficiencia de los SSM y la adaptabilidad de los marcos de trabajo MoE. Los autores sostienen que su diseño jerárquico permite al modelo capturar las jerarquías espaciales inherentes de local a global de los datos médicos de manera más efectiva que los enfoques previos basados en CNN o Transformers.

El trabajo destaca el potencial de las representaciones de características agnósticas a la modalidad, demostrando que un modelo pre-entrenado en TC y RM puede generalizarse efectivamente a tareas de ecografía 3D, abordando el desafío de la variabilidad de los datos entre diferentes modalidades de imagen. Los autores posicionan a Mamba-HoME como una solución escalable para datos 3D de alta resolución que mantiene una eficiencia competitiva de memoria y cómputo mientras entrega una precisión de segmentación superior.

Limitaciones Reconocidas: Los autores señalan que la escalabilidad a conjuntos de datos extremadamente grandes (por ejemplo, >10,000 escaneos) y el comportamiento bajo aprendizaje autosupervisado a gran escala (>200,000 escaneos) aún no han sido explorados, identificando estos puntos como direcciones para trabajos futuros.