Resumen Técnico: RD-ViT: Transformador de Visión con Profundidad Recurrente para Segmentación Semántica con Dependencia Reducida de Datos

Enunciado del Problema

Los Transformadores de Visión (ViT) han logrado un rendimiento de vanguardia en tareas de predicción densa como la segmentación semántica. Sin embargo, los ViT estándar sufren una alta dependencia de los datos debido a su diseño arquitectónico: cada capa en una pila profunda posee parámetros únicos que deben aprenderse de forma independiente. Esto resulta en un recuento de parámetros que escala linealmente con la profundidad ($O(L \cdot d^2)$), haciendo que los ViT sean propensos al sobreajuste cuando los datos de entrenamiento son escasos. Esta limitación es particularmente crítica en la imagenología médica, donde los conjuntos de datos etiquetados suelen ser pequeños (por ejemplo, el conjunto de datos de referencia de RM cardíaca ACDC contiene solo 100 volúmenes de pacientes) y el etiquetado es costoso y consume mucho tiempo. El artículo aborda la necesidad de arquitecturas que puedan lograr una alta precisión de segmentación con datos de entrenamiento limitados, manteniendo al mismo tiempo la eficiencia computacional.

Metodología

Los autores proponen RD-ViT (Transformador de Visión con Profundidad Recurrente), una arquitectura que adapta el concepto de Transformador con Profundidad Recurrente (RDT) del proyecto OpenMythos a tareas de predicción densa. En lugar de una pila profunda de bloques transformadores únicos, RD-ViT emplea un único bloque transformador compartido que se ejecuta en bucle $T$ veces. La arquitectura consta de cuatro componentes principales:

1. Prólogo y Epílogo: Bloques ViT estándar para la extracción inicial de características y el refinamiento final, respectivamente.
2. Bucle Recurrente: La innovación central donde se aplica iterativamente un único bloque ViT. Este bucle incorpora varios mecanismos especializados:
   • Inyección Estable LTI: Para garantizar la convergencia durante la recurrencia de largo alcance, la matriz de transición de estado $A$ se parametriza como $A = \exp(-\exp(\log d_t + \log A))$. Esta construcción garantiza que todos los valores propios se encuentren estrictamente dentro de $(0, 1)$, asegurando que la recurrencia sea una aplicación contractiva. Esto permite una extrapolación de profundidad estable (inferencia con más bucles que los utilizados en el entrenamiento).
   • Tiempo de Cálculo Adaptativo (ACT): Un mecanismo de parada aprendido asigna una probabilidad de parada a cada token de parche en cada iteración. Los tokens en regiones "fáciles" (por ejemplo, fondos homogéneos) se detienen temprano, mientras que los tokens en regiones "difíciles" (por ejemplo, bordes) sufren más iteraciones. Esto crea un mapa de cómputo espacial, asignando recursos dinámicamente.
   • LoRA por Capas: Para permitir que el bloque compartido se comporte de manera diferente en varias profundidades sin duplicar pesos de rango completo, se aplican adaptadores de Adaptación de Bajo Rango (LoRA) con escalado indexado por bucle.
   • Mezcla de Expertos (MoE): Capas MoE dispersas opcionales reemplazan la Red Neuronal de Alimentación Directa (FFN) estándar. Un enrutador dirige los tokens a sub-redes de expertos específicas, habilitando el cálculo condicional y la especialización específica de clase dentro del bloque compartido.
3. Cabeza de Segmentación: Muestra las características a nivel de parche a predicciones a nivel de píxel/vóxel utilizando convoluciones transpuestas.

La arquitectura es agnóstica a la dimensión, soportando entradas tanto 2D (mediante incrustación de parches Conv2d) como 3D (mediante Conv3d).

Contribuciones Clave

1. Adaptación Arquitectónica: Una arquitectura completa de segmentación 2D/3D que integra inyección estable LTI, ACT, LoRA por capas y MoE para predicción densa.
2. Validación de Eficiencia de Datos: Validación empírica en el conjunto de datos de RM cardíaca ACDC que muestra que RD-ViT puede superar o igualar a los ViT estándar con significativamente menos parámetros y menos datos, particularmente en configuraciones 2D.
3. Especialización Emergente: Evidencia de que los expertos de MoE se especializan espontáneamente para diferentes estructuras anatómicas (por ejemplo, Ventrículo Derecho, Miocardio, Ventrículo Izquierdo) sin supervisión de enrutamiento explícita, proporcionando un sesgo inductivo estructural para la segmentación multiclase.
4. Extrapolación de Profundidad: Demostración de que el modelo puede inferir de manera segura con más iteraciones de bucle ($T$) que las utilizadas durante el entrenamiento sin degradación del rendimiento, gracias a la garantía de estabilidad LTI.
5. Generalización: Aplicación exitosa de la misma arquitectura a una tarea fundamentalmente diferente: segmentación de instancias de dientes 3D y numeración FDI a partir de volúmenes CBCT, aprovechando la atención global para el razonamiento espacial de largo alcance.

Resultados Experimentales

Segmentación Cardíaca (Conjunto de Datos ACDC)

• Volumétrico 3D:
  • RD-ViT Tiny + MoE (3.0M parámetros) logró una puntuación Dice de 0.812, alcanzando el 99.4% de la línea base ViT estándar (0.817, 5.7M parámetros) mientras utilizaba solo el 53% de los parámetros.
  • Ablación: MoE proporcionó la mayor ganancia (+0.8 puntos porcentuales sobre RD-ViT Tiny). ACT y LoRA proporcionaron mejoras menores pero consistentes (+0.3 y +0.2 pp, respectivamente).
  • Extrapolación de Profundidad: El rendimiento se mantuvo estable al aumentar los bucles de inferencia de $T=8$ (entrenamiento) a $T=16$, confirmando la estabilidad de la inyección LTI.
• Nivel de Rebanada 2D:
  • RD-ViT superó a ViT estándar en todas las fracciones de datos. Con el 10% de datos de entrenamiento, RD-ViT logró 0.774 Dice frente a 0.762 para ViT estándar. Con el 100% de datos, logró 0.882 frente a 0.872.
  • Los autores atribuyen esto al contexto espacial más rico en las rebanadas 2D, lo que permite que el refinamiento iterativo sea más efectivo que en la configuración 3D de menor resolución.

Segmentación Dental (Conjunto de Datos ToothFairy2)

• Aplicado a la segmentación de instancias de dientes 3D y numeración FDI (33 clases).
• Un modelo de fase única RD-ViT Tiny + MoE (5.4M parámetros) logró:
  • 0.797 Dice de Semilla (diente vs. fondo).
  • 89.1% de precisión en la clasificación FDI.
• Los resultados sugieren que la autoatención global del transformador captura efectivamente el contexto de cuadrante de mandíbula de largo alcance requerido para la numeración de dientes, una tarea donde los campos receptivos locales de las CNN a menudo luchan.

Eficiencia Computacional

• Dinámica de ACT: El tiempo de ponderación medio (iteraciones efectivas) disminuyó de ~2.6 a ~1.4 durante el entrenamiento, lo que indica que el modelo aprendió a procesar las regiones fáciles más rápido.
• Asignación de Cómputo: Los mapas de parada de ACT mostraron una mayor asignación de cómputo en los límites cardíacos y una asignación menor en regiones homogéneas.

Significado y Afirmaciones

El artículo afirma que RD-ViT ofrece una estrategia arquitectónica viable para la segmentación eficiente en parámetros en dominios médicos con datos limitados. Al combinar el compartir pesos (recurrencia) con el cálculo condicional (MoE), el modelo logra un rendimiento cercano al de vanguardia con aproximadamente la mitad de los parámetros de un ViT estándar.

Los autores enfatizan que la inyección estable LTI es crítica para habilitar la extrapolación de profundidad, proporcionando un "control de tiempo de inferencia gratuito" para intercambiar velocidad y precisión sin reentrenar. Además, la especialización espontánea de los expertos de MoE sugiere que las arquitecturas de profundidad recurrente pueden descomponer naturalmente problemas complejos de segmentación en subtareas específicas de clase.

El trabajo destaca una dicotomía entre el rendimiento 2D y 3D: mientras que RD-ViT sobresale en 2D debido al rico contexto espacial, el rendimiento 3D está actualmente más limitado por la capacidad del modelo que por la propia arquitectura recurrente. Sin embargo, la transferencia exitosa a la imagenología dental demuestra el potencial agnóstico al dominio de la arquitectura.

El artículo concluye que, aunque RD-ViT aún no supera a las pipelines altamente optimizadas y con muchos datos en todos los escenarios 3D, reduce significativamente la dependencia de los datos y ofrece un nuevo paradigma para la segmentación de imágenes médicas eficiente, interpretable y adaptable. Todo el código, las notas y los resultados se han publicado públicamente.