Does Your ViT Still Need U-Net for Segmentation?

Resumen Técnico: ¿Sigue necesitando su ViT una U-Net para la segmentación?

Planteamiento del problema
La segmentación de imágenes médicas ha estado dominada durante más de una década por arquitecturas codificador-decodificador de tipo U-Net. Aunque estos diseños abordan con éxito la pérdida de detalles espaciales mediante conexiones de salto (skip connections) y un decodificador simétrico, dependen de operaciones convolucionales que están limitadas por campos receptivos restringidos, lo que dificulta el modelado del contexto global. Aunque los Vision Transformers (ViTs) ofrecen mecanismos de autoatención para capturar dependencias de largo alcance, los primeros métodos de segmentación basados en ViT (por ejemplo, TransUNet, SwinUNet) conservaron decodificadores U-Net porque se consideraba que las representaciones preentrenadas de los ViT eran insuficientes para una predicción densa precisa.

Los avances recientes en el preentrenamiento a gran escala (por ejemplo, DINOv2, SigLIP) han redefinido significativamente las capacidades de representación de los ViTs, permitiéndoles aprender características visuales potentes a partir de conjuntos de datos no etiquetados masivos. Esta evolución plantea una pregunta crítica: ¿Es todavía necesario el paradigma de la U-Net para la segmentación de imágenes médicas cuando se dispone de backbones de ViT modernos y fuertemente preentrenados? Si no es así, ¿cómo debe diseñarse un marco de solo codificador (encoder-only) para predecir eficazmente las máscaras de segmentación sin un decodificador?

Metodología: EoSeg
Los autores proponen EoSeg (Encoder-only Segmentation), un marco que elimina el decodificador U-Net en favor de un diseño basado en consultas (queries) construido sobre modernos backbones de ViT preentrenados. La metodología se desarrolla mediante una exploración sistemática de decisiones arquitectónicas clave:

1. Selección del Backbone: El estudio evalúa varios backbones de ViT preentrenados (incluyendo ViT-L/14, SigLIP, MTP, DINOv2 y DINOv3) en un entorno de solo codificador. DINOv2 es seleccionado como el backbone óptimo, logrando el mayor rendimiento base (81.86% mDice en Synapse) sin ningún decodificador.
2. Generación de Máscaras:
   • Deconvolución vs. Basado en Consultas: Inicialmente, se probó un módulo de deconvolución para aumentar la resolución de las características de baja resolución de los ViT. Sin embargo, los autores descubrieron que un enfoque de predicción de máscara basado en consultas era superior. En lugar de una decodificación píxel a píxel, los embeddings de consulta aprendibles actúan como sondas semánticas globales.
   • Mecanismo: Los vectores de consulta se proyectan en un espacio de embedding de máscara e interactúan con el mapa de características globales de resolución aumentada mediante un producto punto canal a canal para generar máscaras binarias directamente.
3. Modelado de Consultas Multinivel: Inspirado en las conexiones de salto de las U-Nets, el marco aprovecha las representaciones de los últimos tres bloques Transformer en lugar de solo el bloque final. Las consultas aprendibles se inyectan antes de estos bloques y se refinan mediante autoatención con los tokens de la imagen, lo que permite al modelo agregar información semántica de múltiples etapas.
4. Fusión de Bloques Aprendible: Para combinar las predicciones de los tres bloques intermedios, se introduce un módulo de fusión aprendible. Una puerta de fusión utiliza la representación de la consulta para predecir pesos adaptativos para la salida de cada bloque, permitiendo que el modelo determine dinámicamente la contribución de cada nivel a la predicción final.
5. Configuración de Consultas: El marco adopta un paradigma de clasificación basado en consultas donde cada consulta predice una máscara y una etiqueta de clase. Los experimentos determinaron que asignar dos consultas por categoría produce un rendimiento óptimo, equilibrando la necesidad de capturar variaciones intra-clase sin sobreparametrización.
6. Estrategia de Supervisión: Contrario a las prácticas comunes en la segmentación médica, se encontró que la supervisión profunda (aplicar la pérdida a las salidas intermedias) no mejoró el rendimiento en este esquema específico de solo codificador y fue excluida del diseño final.

Contribuciones Clave

• Reevaluación Arquitectónica: El artículo proporciona una evaluación extensa demostrando que un decodificador de estilo U-Net ya no es un requisito previo para una segmentación de imágenes médicas de alta calidad cuando se aprovechan los backbones de ViT modernos preentrenados a gran escala.
• Marco EoSeg: Los autores presentan EoSeg, un marco de solo codificador basado en consultas que integra el modelado de consultas multinivel y la fusión de bloques aprendible para permitir la predicción directa de máscaras a partir de las características del codificador.
• Validación Exhaustiva: El diseño propuesto se valida en siete conjuntos de datos de referencia que abarcan cinco modalidades distintas de imágenes médicas: CT (Synapse), MRI (ACDC), Histopatología (GlaS, MoNuSeg), Endoscopia (Kvasir-Seg) y Dermoscopia (ISIC-2016, ISIC-2017).

Resultados Experimentales
EoSeg supera consistentemente a los métodos de codificador-decodificador basados en CNN, Transformer e híbridos en todas las modalidades probadas:

• Synapse (CT): Logró un 85.50% de media Dice (mDice), superando al mejor anterior (EViT-UNet con 80.87%) por 4.63 puntos. También alcanzó puntuaciones máximas para órganos específicos como la Aorta (90.80%) y el Bazo (94.13%).
• ACDC (MRI): Logró un 91.73% mDice, superando al mejor anterior por 0.24 puntos.
• GlaS (Histopatología): Logró un 93.27% mDice y 87.79% mIoU, superando al mejor anterior por 0.83 y 1.29 puntos, respectivamente.
• MoNuSeg (Histopatología): Logró un 80.51% mDice y 67.45% mIoU.
• Kvasir-Seg e ISIC: Demostraron un rendimiento superior en la segmentación de pólipos y lesiones, logrando altos puntajes de mIoU y mDice con bordes suaves y estructuras coherentes.

El análisis cualitativo confirma que EoSeg produce bordes más suaves y menos artefactos en comparación con los métodos competidores, preservando eficazmente los detalles estructurales finos sin la restauración de características basada en decodificadores.

Significado y Reivindicaciones
El artículo afirma que la dominancia del paradigma U-Net en la segmentación de imágenes médicas ya no está justificada por las limitaciones de los primeros backbones de ViT. En su lugar, las representaciones fuertes aprendidas mediante el preentrenamiento a gran escala permiten diseños de solo codificador que pueden predecir directamente máscaras de segmentación precisas.

Los autores enfatizan que el éxito de EoSeg no se debe meramente a la eliminación del decodificador, sino a la adopción de un paradigma de predicción basado en consultas apropiado combinado con el modelado multinivel y la fusión adaptativa. Estos hallazgos sugieren que los futuros marcos de segmentación de imágenes médicas construidos sobre modelos de visión preentrenados a gran escala pueden alejarse de las complejas arquitecturas de decodificador, confiando en cambio en el poder de representación intrínseco de los ViT modernos.