Resumen Técnico: UNet Condicionada por Transformador de Visión para Segmentación Semántica Adaptativa al Dominio

Enunciado del Problema

Aunque los Transformadores de Visión (ViTs) se han convertido en la columna vertebral de los modelos de segmentación más avanzados (SOTA) en dominios de imágenes naturales, su aplicación a la imagen biomédica sigue siendo un desafío. Los conjuntos de datos biomédicos suelen tener limitaciones de anotación y se caracterizan por estructuras especializadas con morfologías irregulares, dispersas, delgadas o de bajo contraste. Los autores identifican una brecha específica de rendimiento: los decodificadores de píxeles ligeros comúnmente utilizados en modelos ViT promptables (por ejemplo, cabezas lineales simples o máscaras de remuestreo) a menudo carecen del sesgo inductivo local requerido para máscaras biomédicas de alta precisión. Además, el ajuste fino de modelos grandes basados en ViT para tareas biomédicas es computacionalmente costoso y requiere grandes cantidades de datos, mientras que los enfoques híbridos existentes a menudo no logran superar consistentemente a los UNets bien ajustados (como nnU-Net) bajo validación rigurosa.

Metodología: ViTC-UNet

El artículo propone ViTC-UNet, un marco diseñado para cerrar la brecha entre los priores visuales globales de los ViT y el sesgo inductivo local de las Redes Neuronales Convolucionales (CNN). La arquitectura evita el ajuste fino de ViT de extremo a extremo, incluso en configuraciones de dominio cruzado, al condicionar un UNet sobre representaciones de ViT preentrenadas y congeladas.

Componentes Principales

1. Codificador de Espalda Congelado: El modelo utiliza una espalda de ViT congelada preentrenada en imágenes naturales a gran escala (evaluada con DINOv2 y SAM 2). Esto preserva características de alta calidad y propósito general, evitando al mismo tiempo los requisitos de datos del ajuste fino de Transformadores.
2. Decodificador de Condicionamiento: Este módulo cierra la brecha entre las incrustaciones de imágenes naturales y las representaciones espaciales biomédicas.
   • Emplea un Perceptrón Multicapa (MLP) seguido de una serie de bloques de atención bidireccional (modificados a partir de SAM).
   • Tokens de Estructura: En lugar de canales de salida fijos, el modelo utiliza "tokens de estructura" aprendibles (uno por clase objetivo) emparejados con incrustaciones posicionales espaciales compartidas.
   • Mecanismo: Las incrustaciones de imagen y los tokens de estructura se atienden mutuamente entre sí. El token de estructura agrega contexto específico de la imagen, mientras que la representación visual se modula según el objetivo solicitado.
   • Consistencia Topológica: Los autores argumentan que, dado que la espalda de ViT actúa como un homeomorfismo (una aplicación continua y biyectiva), el decodificador debe preservar esta característica topológica. Los bloques de atención bidireccional se tratan como transformaciones homeomorfas, asegurando un espacio de representación bien comportado que facilita un descodificación robusta aguas abajo.
3. Condicionamiento Latente Multi-Etapa y Multi-Escala:
   • En lugar de depender de una única representación de cuello de botella, el decodificador genera una secuencia de estados latentes condicionados por la estructura a través de bloques de atención sucesivos.
   • Estos estados proporcionan condicionamiento multi-etapa (refinando las interacciones imagen-token a través de los bloques) y condicionamiento multi-escala (fusionándose con los niveles correspondientes de la ruta de descodificación del UNet).
   • Las etapas más profundas reciben contexto global específico del objetivo, mientras que las etapas más superficiales apoyan la recuperación de límites finos y objetivos dispersos.
4. Decodificador de Píxeles UNet Condicionado por Estructura:
   • Los autores utilizan nnU-Net como decodificador de píxeles.
   • La salida del $n$-ésimo bloque de atención se concatena con la conexión de salto del $n$-ésimo bloque convolucional más superficial.
   • Salida de Canal Único: Crucialmente, el decodificador de píxeles utiliza un único canal de salida independientemente del número de estructuras objetivo. La identidad semántica se desacopla de la dimensionalidad de la salida; el modelo genera una máscara binaria condicionada por el token de estructura de entrada. La segmentación multiclase se logra repitiendo el proceso de inferencia para cada token objetivo.

Contribuciones Clave

• Arquitectura Novel: Introducción de ViTC-UNet, que condiciona un UNet sobre representaciones de ViT congeladas mediante tokens aprendibles y un decodificador de atención bidireccional, combinando priores globales con sesgo inductivo local.
• Eficiencia y Adaptabilidad: El método logra un rendimiento SOTA sin ajuste fino de ViT de extremo a extremo, lo que lo hace adecuado para dominios biomédicos con escasez de datos.
• Flexibilidad del Espacio de Etiquetas: Al desacoplar la identidad semántica de la capa de salida, la arquitectura permite la adición incremental de nuevas estructuras objetivo mediante tokens aprendibles sin modificar la arquitectura de la red ni reentrenar todo el modelo.
• Perspectiva Teórica: El artículo proporciona una justificación teórica para el diseño del decodificador basada en el homeomorfismo, argumentando que preservar las propiedades topológicas del espacio de incrustación de ViT es esencial para la fidelidad de la segmentación.

Resultados

El modelo fue evaluado en una compilación de conjuntos de datos de CT y MRI (incluyendo Abdomen1K, Aorta, AirwayTree, AMOS, BraTS, etc.) curados por BiomedParse.

• Rendimiento: El ViTC-UNet basado en DINOv2 logró un mIoU promedio de fondo de 0.90 en los conjuntos de datos probados, superando la línea base de nnU-Net (0.79) y el resultado agregado de BiomedParse (0.88). Igualó o superó a nnU-Net en 14 de 15 puntos de referencia.
• Comparación con Híbridos: Los resultados demuestran que simplemente concatenar incrustaciones de ViT en un cuello de botella de UNet (un enfoque híbrido anterior) es insuficiente. El condicionamiento multi-etapa y multi-escala propuesto es necesario para delinear correctamente las morfologías anatómicas, particularmente para objetivos dispersos y de estructura fina.
• Aprendizaje Incremental: En un experimento de entrenamiento en dos etapas, el modelo incorporó con éxito clases nuevas sin degradar el rendimiento en estructuras previamente aprendidas, confirmando su capacidad para la expansión del espacio de etiquetas.
• Sensibilidad de la Espalda: El rendimiento dependía altamente de la geometría del espacio latente de la espalda. DINOv2 (auto-supervisado) superó a SAM 2 (supervisado), lo que sugiere que los subespacios lineales ortogonales de DINOv2 son más propicios para las transformaciones homeomorfas modeladas en el decodificador.

Significado y Afirmaciones

El artículo afirma que ViTC-UNet ofrece una ruta computacionalmente eficiente y eficiente en muestras para adaptar priores visuales a gran escala a la segmentación biomédica de alta complejidad. Al tratar el UNet como un decodificador de píxeles promptable y condicionado por estructura, el trabajo demuestra que las representaciones visuales de propósito general pueden complementar eficazmente las pipelines de segmentación biomédica robustas.

Los autores afirman que este enfoque cierra la brecha entre los modelos de visión a gran escala y las aplicaciones clínicas especializadas, particularmente en escenarios con bajas relaciones señal-ruido y escasez de etiquetas. Destacan que la arquitectura permite una descodificación de alta fidelidad sin el costo computacional prohibitivo de ajustar fino grandes Transformadores, manteniendo al mismo tiempo la capacidad de escalar a nuevas estructuras anatómicas mediante condicionamiento basado en tokens en lugar de cambios arquitectónicos.

Limitaciones

Los autores reconocen varias limitaciones:

• Procesamiento 2D: Debido a que los codificadores ViT están preentrenados en imágenes naturales 2D, la pipeline procesa los datos por rebanadas, sin aprovechar completamente el contexto volumétrico 3D.
• Tokens Discretos: El uso de tokens de estructura discretos y finitos limita la generalización de vocabulario abierto.
• Geometría Latente: El rendimiento del modelo depende de la ortogonalidad del espacio latente del codificador; la adaptación a codominios con geometrías diferentes requiere mayor investigación.
• Tipos de Prompt: El trabajo actual se centra en tokens de estructura semántica y aún no incorpora prompts visuales interactivos (por ejemplo, puntos o cuadros delimitadores).