💻 computer science

Primus: Enforcing Attention Usage for 3D Medical Image Segmentation

Este artículo presenta Primus y PrimusV2, las primeras arquitecturas centradas en Transformers competitivas para la segmentación de imágenes médicas 3D que superan las limitaciones de los modelos híbridos al maximizar el uso de la atención, logrando así un rendimiento de vanguardia que supera o iguala a los métodos basados en CNN líderes en nueve conjuntos de datos públicos.

Autores originales: Tassilo Wald, Saikat Roy, Fabian Isensee, Constantin Ulrich, Sebastian Ziegler, Dasha Trofimova, Raphael Stock, Michael Baumgartner, Gregor Köhler, Klaus Maier-Hein

Publicado 2026-05-01

📖 5 min de lectura🧠 Análisis profundo

CC BY 4.0

Autores originales: Tassilo Wald, Saikat Roy, Fabian Isensee, Constantin Ulrich, Sebastian Ziegler, Dasha Trofimova, Raphael Stock, Michael Baumgartner, Gregor Köhler, Klaus Maier-Hein

Artículo original bajo licencia CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ Esta es una explicación generada por IA del artículo a continuación. No ha sido escrita ni avalada por los autores. Para mayor precisión técnica, consulte el artículo original. Leer descargo de responsabilidad completo

El Gran Problema: Los Transformers "Impostores"

Imagina que estás intentando construir un robot que pueda mirar una imagen médica en 3D (como una tomografía computarizada de un riñón) y dibujar un contorno perfecto alrededor del tumor.

Durante mucho tiempo, los mejores robots para este trabajo fueron las CNN (Redes Neuronales Convolucionales). Piensa en una CNN como un detective muy hábil y local. Examina un pequeño vecindario de píxeles, recopila pistas y se mueve al siguiente vecindario. Es excelente para ver detalles locales, pero a veces pierde la "gran imagen" de todo el órgano.

Luego, llegaron los Transformers. Estos son como "superobservadores" que pueden mirar la imagen completa de una sola vez y entender cómo la parte superior del riñón se relaciona con la inferior. Se hicieron famosos en el procesamiento del lenguaje (como los chatbots) y en fotografías naturales. Todos pensaron: "¡Si los Transformers pueden leer un libro entero o ver un paisaje completo, deben ser perfectos para las imágenes médicas!".

Pero aquí está el truco: Cuando los investigadores intentaron usar Transformers para imágenes médicas en 3D, no funcionaron muy bien. A menudo eran más lentos y menos precisos que los viejos detectives CNN.

La Investigación: ¿Quién está haciendo el trabajo?

Los autores de este artículo decidieron investigar por qué estos Transformers estaban fallando. Examinaron nueve modelos "híbridos" populares (modelos que intentan usar tanto CNN como Transformers).

Descubrieron un secreto impactante: Los Transformers estaban mayormente dormidos.

La Analogía: Imagina una cuadrilla de construcción donde contrataste a un arquitecto famoso y costoso (el Transformer) para diseñar una casa, pero también contrataste a 100 albañiles normales (las CNN). Cuando la casa estuvo terminada, los autores se dieron cuenta de que el arquitecto en realidad no había dibujado los planos. Los albañiles construyeron toda la casa ellos mismos, y el arquitecto solo estaba allí de pie asintiendo.
La Evidencia: Cuando los investigadores eliminaron al "arquitecto" (los bloques del Transformer) de estos modelos, los modelos funcionaron casi exactamente igual. En algunos casos, eliminar el Transformer incluso los hizo más rápidos y ligeramente mejores porque el modelo dejó de desperdiciar energía en un componente inútil.

El artículo llama a esto el "Índice UNet". La mayoría de los modelos existentes tenían un índice alto, lo que significaba que en realidad eran solo CNNs disfrazadas, cargando una mochila pesada e inútil de Transformer.

La Solución: Primus y PrimusV2

Los autores se preguntaron: "¿Qué pasaría si construyéramos un modelo donde el Transformer tenga que hacer el trabajo?". Crearon dos nuevas arquitecturas llamadas Primus (latín para "Primero") y PrimusV2.

Así es como forzaron al Transformer a despertar y hacer su trabajo:

1. No Aplastar los Detalles (El Tokenizador)

Los Transformers estándar a menudo cortan la imagen 3D en trozos enormes (como cortar un pastel en rebanadas gigantes y gruesas) para convertirlos en tokens de datos. Esto tira información diminuta e importante, como un tumor pequeño o un vaso sanguíneo fino.

La Solución: Primus utiliza un "tokenizador de alta resolución". En lugar de rebanadas gigantes, usa rebanadas más pequeñas y finas (8x8x8 vóxeles).
La Analogía: En lugar de mirar un mapa de una ciudad donde cada calle es solo una línea borrosa, Primus mira un mapa donde puedes ver casas individuales. Esto le da al Transformer información más detallada con la que trabajar.

2. El Enfoque "Iterativo" (PrimusV2)

Incluso con rebanadas más pequeñas, el Transformer a veces luchaba por entender las formas complejas en 3D de los órganos de inmediato.

La Solución: PrimusV2 utiliza una "Incrustación de Parches Iterativa". En lugar de intentar entender todo el trozo de una sola vez, procesa la imagen en etapas, como pelar una cebolla o refinar un boceto. Utiliza unos pocos pasos convolucionales pequeños e inteligentes para preparar los datos antes de que el Transformer los vea.
La Analogía: Imagina intentar resolver un rompecabezas complejo. Primus no simplemente tira todas las piezas sobre la mesa. Primero las ordena por color y piezas de borde (los pasos iterativos), lo que hace mucho más fácil para el "superobservador" (el Transformer) ver la imagen final.

3. Dándole un GPS (3D RoPE)

Los Transformers son naturalmente "ciegos" al espacio; no saben que "izquierda" es diferente de "derecha" a menos que se lo digas.

La Solución: Los autores añadieron una "Incrustación de Posición Rotatoria" (RoPE) 3D especial.
La Analogía: Es como darle al Transformer un sistema GPS que entiende la profundidad, el ancho y la altura simultáneamente. Sabe exactamente dónde está ubicado un tumor en el espacio 3D en relación con el resto del órgano.

Los Resultados: El Transformer Finalmente Gana

Después de estos cambios, los resultados fueron impresionantes:

Primus se convirtió en el primer modelo basado en Transformers que podía competir con la CNN estándar de "estándar de oro" (nnU-Net).
PrimusV2 no solo compitió; superó a la CNN estándar en la mayoría de las pruebas e igualó a las CNN más complejas y mejores disponibles hoy en día.

La Conclusión Clave:
El artículo demuestra que los Transformers pueden ser la mejor herramienta para la imagenología médica en 3D, pero solo si dejas de tratarlos como un ayudante de una CNN. Tienes que diseñar el sistema para que el Transformer sea el personaje principal, alimentado con datos de alta resolución y equipado con las herramientas adecuadas para entender el espacio 3D.

Resumen en una Frase

Los autores construyeron un nuevo modelo de IA llamado Primus que finalmente fuerza al Transformer "superobservador" a hacer el trabajo pesado en la imagenología médica, demostrando que, cuando se diseñan correctamente, los Transformers pueden superar a las CNN tradicionales de "detective local".

1. Planteamiento del Problema

A pesar del éxito de los Transformers en el procesamiento del lenguaje natural y la visión por computadora de imágenes naturales, su impacto en la segmentación de imágenes médicas 3D ha sido limitado. Los benchmarks actuales siguen dominados por las Redes Neuronales Convolucionales (CNN), particularmente el marco nnU-Net. Los autores identifican dos deficiencias críticas en los enfoques basados en Transformers existentes para la imagen médica:

Sobredependencia de las Convoluciones: La mayoría de las arquitecturas "Transformer" son en realidad modelos híbridos donde la mayoría de los parámetros y la capacidad computacional residen en bloques convolucionales (codificadores/decodificadores), convirtiendo a los componentes Transformer en meros auxiliares.
Ineficacia de la Atención: En muchos modelos híbridos, eliminar los bloques Transformer resulta en caídas de rendimiento neglijibles, lo que sugiere que los modelos están aprendiendo representaciones únicamente mediante CNNs. Además, los modelos Transformer puros o casi puros existentes a menudo tienen un rendimiento inferior al de las CNNs base sólidas, probablemente debido a la dificultad de entrenar Transformers desde cero en conjuntos de datos médicos pequeños y escasamente etiquetados.

2. Metodología

Los autores proponen un enfoque de dos frentes: primero, una deconstrucción rigurosa de los modelos existentes para demostrar sus limitaciones, y segundo, el diseño de nuevas arquitecturas verdaderamente centradas en Transformers.

A. Análisis de Arquitecturas Existentes

Los autores analizaron nueve modelos populares basados en Transformers (por ejemplo, TransUNet, SwinUNETR, UNETR) utilizando una métrica llamada Índice UNet.

Definición del Índice UNet: Una relación que mide la cantidad de parámetros convolucionales fuera de los bloques Transformer en relación con un U-Net estándar.
Hallazgos: La mayoría de los modelos tienen un Índice UNet alto (gran dependencia de CNNs). Cuando los bloques Transformer fueron reemplazados por mapeos de identidad, 6 de cada 9 modelos mostraron una caída de menos del 3% en el Coeficiente de Similitud Dice (DSC), demostrando que el Transformer no era el motor principal del rendimiento.

B. Las Arquitecturas Primus y PrimusV2

Para abordar estos problemas, los autores diseñaron Primus y PrimusV2, las primeras arquitecturas centradas en Transformers competitivas para la imagen médica 3D. Estos modelos minimizan los parámetros convolucionales, forzando a la red a depender de los mecanismos de atención.

Innovaciones Clave de Diseño:

Tokenización de Alta Resolución:
- A diferencia de los ViT estándar que utilizan parches grandes de $16\times16\times16$ (que pierden detalle local), Primus utiliza un tamaño de parche de $8\times8\times8$ para preservar estructuras anatómicas de fino detalle.
- PrimusV2 introduce una Codificación de Parches Iterativa. En lugar de una sola proyección, utiliza una cadena de bloques convolucionales residuales para submuestrear y extraer características semánticas iterativamente antes de la tokenización. Esto crea tokens más ricos sin depender de una columna vertebral CNN pesada.
Embebidos Posicionales Rotatorios 3D (RoPE):
- Los embebidos posicionales absolutos estándar luchan con los desplazamientos espaciales comunes en los datos médicos 3D. Los autores adaptaron los Embebidos Posicionales Rotatorios Axiales (RoPE) a 3D, codificando distancias y orientaciones relativas directamente en el mecanismo de atención para capturar mejor las correlaciones espaciales volumétricas.
Componentes Modernos de Transformer:
- Bloques MLP: Adopción de SwiGLU (Unidades Lineales Puertas activadas por Swish) con Normalización de Capa, probada como efectiva en Transformers de imágenes naturales (por ejemplo, EVA-02).
- Mecanismos de Estabilidad: Uso de LayerScale y Normalización Post-Atención para estabilizar el entrenamiento en modelos grandes (Primus-L) y prevenir la divergencia.
- Decodificador Ligero: Un decodificador compuesto por convoluciones transpuestas que submuestrea tokens de forma independiente, minimizando la influencia convolucional y manteniendo bajo el Índice UNet.

3. Contribuciones Clave

Deconstrucción Empírica: Proporcionó evidencia cuantitativa (mediante el Índice UNet y estudios de ablación) de que los modelos médicos actuales "basados en Transformers" son efectivamente CNNs con adiciones menores de atención.
Primeros Modelos Centrados en Transformers Competitivos: Introdujo Primus y PrimusV2, que logran un rendimiento de última generación mientras mantienen la vasta mayoría de los parámetros y FLOPs dentro de los bloques Transformer.
Innovaciones Arquitectónicas: Demostró que combinar tokenización de alta resolución ( $8\times8\times8$ ), codificación de parches iterativa y RoPE 3D permite a los Transformers superar la escasez de datos y los requisitos de características locales de la imagen médica.
Eficiencia en Aprendizaje Auto-supervisado: Destacó que los diseños centrados en Transformers son más eficientes para el pre-entrenamiento con Modelado de Imágenes Enmascaradas (MIM), utilizando significativamente menos VRAM (40% menos) en comparación con arquitecturas pesadas en CNN al enmascarar tokens.

4. Resultados Experimentales

Los modelos fueron evaluados en nueve conjuntos de datos públicos (incluyendo AMOS22, KiTS23, TotalSegmentator-BTCV, SBM, etc.) que cubren diversas modalidades (TC, RMN) y órganos.

Rendimiento vs. Líneas Base:
- Primus: Superó los métodos actuales basados en Transformers y compitió directamente con el nnU-Net predeterminado, superándolo en 4 de 9 conjuntos de datos.
- PrimusV2: Superó significativamente a todas las demás líneas base de Transformers (incluyendo UNETR++, SwinUNETRv2) y alcanzó paridad con las CNNs de última generación como ResEnc-L y MedNeXt.
- Reducción de la Brecha: PrimusV2 redujo la brecha de rendimiento con las mejores CNNs a dentro de 0.05–0.29 puntos DSC, una mejora masiva sobre los intentos anteriores de Transformers que se rezagaban entre 8–11 puntos DSC.
Insights de Ablación:
- Eliminar el Transformer de PrimusV2 causó una caída de ~5.0 DSC, confirmando que el modelo depende fuertemente de sus mecanismos de atención (a diferencia de los híbridos anteriores donde la caída era cercana a 0).
- El Tokenizador Iterativo en PrimusV2 proporcionó una ganancia de ~2 DSC sobre la proyección estándar, mejorando particularmente el rendimiento en lesiones pequeñas (por ejemplo, metástasis cerebrales en el conjunto de datos SBM).
Escalado: Aunque escalar el tamaño del modelo (de S a L) mostró rendimientos decrecientes en conjuntos de datos muy pequeños (debido al sobreajuste), PrimusV2-L logró las puntuaciones generales más altas en conjuntos de datos más grandes.

5. Significado

Este trabajo marca un cambio de paradigma en la segmentación de imágenes médicas 3D:

Viabilidad de Transformers Puros: Demuestra que los Transformers pueden ser el motor principal del rendimiento en la imagen médica, no solo un añadido a las CNNs.
Preparación para el Futuro: Al establecer una línea base fuerte centrada en Transformers, el artículo permite la adopción de aprendizaje auto-supervisado (SSL) y integración multi-modal (por ejemplo, combinar imágenes con informes de texto) en la IA médica, áreas donde los Transformers sobresalen pero las CNNs luchan.
Aplicabilidad Clínica: El análisis del tiempo de inferencia muestra que incluso los modelos PrimusV2 más grandes caben dentro de las ventanas de decisión clínica (por ejemplo, detección de accidentes cerebrovasculares), haciéndolos viables para su implementación en el mundo real.

En conclusión, Primus y PrimusV2 demuestran que con las adaptaciones arquitectónicas adecuadas (tokens de alta resolución, RoPE 3D, codificación iterativa), los Transformers pueden igualar o superar a las mejores CNNs en la segmentación médica 3D, abriendo la puerta a pipelines de IA médica de próxima generación.