Resumen Técnico: MIS-HCC para la Segmentación Eficiente de Imágenes Médicas

Declaración del Problema

El despliegue de redes neuronales profundas complejas (p. ej., U-Net, Transformers) para la segmentación de imágenes médicas enfrenta un cuello de botella crítico: el compromiso entre la alta precisión y las restricciones computacionales/de parámetros de las plataformas médicas con recursos limitados (p. ej., dispositivos móviles, nodos de borde). Aunque existen técnicas de compresión de modelos existentes como la poda de redes, la cuantización y la destilación de conocimiento, estas suelen fallar al no abordar las características específicas de las imágenes médicas.

Los métodos actuales de poda estructurada (p. ej., HRank, CHEX) típicamente evalúan las unidades de la red de forma aislada, pasando por alto las correlaciones intrínsecas entre canales y la alta similitud de características inducida por las estructuras anatómicas. En consecuencia, estos métodos corren el riesgo de descartar semánticas finas esenciales al simplemente eliminar filtros "redundantes", en lugar de preservar el poder representacional colectivo de las características correlacionadas. Existe la necesidad de un marco de compresión que integre la similitud de representación para fusionar características correlacionadas en lugar de meramente eliminarlas.

Metodología: MIS-HCC

Los autores proponen MIS-HCC (Clustering Jerárquico de Canales para la Segmentación Eficiente de Imágenes Médicas), un marco de compresión estructurada diseñado para reemplazar las decisiones binarias de "mantener o descartar" con una fusión inteligente de parámetros. El método opera en tres etapas principales:

1. Similitud de Representación mediante la Distancia de Wasserstein:
   En lugar de utilizar métricas estándar como la similitud de coseno sobre vectores brutos, MIS-HCC emplea la distancia de Wasserstein para cuantificar la similitud entre las distribuciones de los canales dentro de una capa de la red preentrenada. Esta métrica se elige por su capacidad para capturar discrepancias de distribución y su robustez ante desplazos de masa local, lo que la hace adecuada para las características de los canales agregadas sobre dimensiones espaciales, temporales o de kernel.

   • Para una capa con canales de entrada, el tensor de pesos se aplana para formar una matriz de características de canal.
   • Se calcula una matriz de distancia de canales por pares utilizando la distancia de Wasserstein.
   • Esta matriz de distancia se normaliza y se convierte en una matriz de similitud ($M_2$) donde valores más altos indican una mayor similitud.

2. Clustering Jerárquico:
   Para tener en cuenta la interdependencia de los canales, el método aplica clustering jerárquico aglomerativo con Vínculo Promedio (Average Linkage) sobre la matriz de similitud derivada.

   • Los canales se tratan inicialmente como clusters individuales.
   • El algoritmo fusiona iterativamente los clusters basándose en la distancia promedio entre sus miembros.
   • Este proceso particiona los $N_{in}$ canales de entrada en $K$ clusters ($K \leq N_{in}$), determinados por una tasa de compresión deseada o un umbral de corte del dendrograma. Este enfoque asegura que los canales funcionalmente superpuestos se agrupen.

3. Fusión de Parámetros:
   Una vez que los canales se han agrupado, el método fusiona los parámetros dentro de cada cluster para crear una red compacta.

   • Para cada cluster $c$ con conjunto de índice $I_c$, los pesos se fusionan mediante el promedio aritmético a lo largo de la dimensión de canal de entrada.
   • Esto resulta en un nuevo tensor de pesos con $K$ canales de entrada, reduciendo efectivamente el tamaño del modelo y los FLOPs por un factor de $K/N_{in}$.
   • Crucialmente, esta fusión es una operación de forma cerrada y libre de entrenamiento que mantiene la topología original de la red, asegurando la compatibilidad con hardware de propósito general y evitando la necesidad de kernels de inferencia especializados o reentrenamiento.

El marco está diseñado para ser aplicable tanto a arquitecturas CNN como Transformer para datos 2D y 3D.

Principales Contribuciones

El artículo describe tres contribuciones principales:

1. Marco MIS-HCC: Un método de compresión estructurada que introduce el clustering jerárquico en la poda de modelos médicos, cambiando de la simple eliminación de filtros a la fusión inteligente de parámetros.
2. Similitud Basada en Wasserstein: La utilización de la distancia de Wasserstein para construir una matriz de similitud de canales robusta, modelando explícitamente las correlaciones entre canales para preservar características visuales críticas durante la compresión.
3. Validación Empírica: Experimentos extensos en múltiples benchmarks de segmentación médica demostrando que MIS-HCC logra compromisos de compresión-precisión superiores comparado con métodos de estado del arte.

Resultados Experimentales

El método fue evaluado en tres conjuntos de datos 2D (BUSI, DSB, ISIC) utilizando backbones U-Net y U-Net++, y un conjunto de datos 3D (FLARE 2021) utilizando el backbone MedSAM2. Se aplicó una tasa de poda uniforme del 87.5% para la comparación.

• Segmentación 2D: En los conjuntos de datos BUSI, DSB e ISIC, MIS-HCC superó a los métodos existentes (RP, PFEC, HRank, CHEX, TPP) tanto en puntuaciones de Dice como de IoU. Notablemente, mientras que otros métodos a menudo mostraron caídas de precisión respecto al baseline, MIS-HCC frecuentemente logró puntuaciones ligeramente superiores al baseline sin comprimir o mantuvo un rendimiento casi sin pérdidas. Por ejemplo, en el conjunto de datos BUSI con U-Net, MIS-HCC logró un Dice de 0.7841 comparado con el 0.7781 del siguiente mejor método (TPP).
• Segmentación 3D: En el conjunto de datos FLARE 2021 con MedSAM2, MIS-HCC superó significativamente a otras estrategias. Mientras que la poda aleatoria (RP) causó una caída del 10.35% en Dice y TPP causó una caída del 1.61%, MIS-HCC redujo el Dice en solo un 0.14% (0.8471 vs. 0.8483 del baseline).
• Eficiencia de Compresión: Con altas tasas de poda (87.5%), MIS-HCC demostró una reducción superior en FLOPs y parámetros en comparación con TPP manteniendo una mayor precisión.
• Estudios de Ablación: Los experimentos confirmaron que la combinación de la distancia de Wasserstein con el Clustering Jerárquico produce el mejor rendimiento, superando a las variantes que utilizan solo uno de estos componentes o métricas estándar como la similitud de Coseno.
• Análisis Visual: Los resultados cualitativos mostraron que MIS-HCC produjo mapas de segmentación con contornos de lesiones más nítidos y una mejor preservación de detalles finos en comparación con otros métodos, que a menudo exhibían bordes borrosos o falsos positivos.

Significado y Reivindicaciones

El artículo afirma que MIS-HCC ofrece una solución efectiva para desplegar modelos de segmentación de imágenes médicas en plataformas con recursos limitados. Al aprovechar la similitud estructural intrínseca prevalente en las imágenes médicas, el método equilibra con éxito la compresión del modelo con la preservación del poder representacional crítico.

Los autores aseguran que su enfoque es particularmente efectivo para tareas de segmentación 3D a gran escala donde la diversidad de características es crucial. Posicionan a MIS-HCC como un marco robusto y generalizable aplicable a diversas arquitecturas (CNNs y Transformers) y modalidades de datos (2D y 3D), proporcionando un camino prometedor para el despliegue clínico práctico sin las pesadas cargas computacionales de los modelos actuales de vanguardia. El trabajo concluye sugiriendo que estos principios podrían extenderse a otros escenarios del mundo real para mejorar aún más la eficiencia de los modelos.