Partial Ring Scan: Revisiting Scan Order in Vision State Space Models

Resumen Técnico: Escaneo de Anillo Parcial: Revisitando el Orden de Escaneo en los Modelos de Espacio de Estados de Visión

Planteamiento del Problema
Los Modelos de Espacio de Estados (SSM), particularmente las variantes de Vision Mamba, han surgido como alternativas eficientes a los mecanismos de atención para tareas de visión, ofreciendo un procesamiento de secuencias en tiempo lineal. Sin embargo, estos modelos requieren la serialización de imágenes 2D en secuencias de tokens 1D basadas en un orden de escaneo predefinido (por ejemplo, raster, serpentina). Los autores argumentan que este orden de escaneo es un factor crítico, a menudo pasado por alto, que dicta la adyacencia espacial dentro de la secuencia.

Los escaneos convencionales de trayectoria fija (raster, bidireccional, diagonal) preservan la alineación espacio-secuencia solo bajo transformaciones específicas como las reflexiones (flips). Bajo rotaciones en el plano, estas trayectorias sufren de reindexación global y relleno (padding), lo que fractura la continuidad de los objetos y desalinea el vecindario del núcleo recurrente. Esto obliga al SSM a gastar capacidad en reparar la coherencia local en lugar de modelar la semántica, lo que conduce a una degradación del rendimiento (una caída de 1–2% en Top-1) y una reducción de la robustez ante transformaciones geométricas.

Metodología: PRISMamba
El artículo propone PRISMamba (Partial Ring Scan Mamba), una arquitectura robusta a la rotación diseñada para alinear la adyacencia de la secuencia con la estructura geométrica sin requerir un entrenamiento específico para la rotación. El método consta de tres componentes principales:

1. Escaneo Alterno Anillo por Anillo:

   • Partición: La imagen se particiona en anillos concéntricos basados en la distancia euclidiana de los píxeles desde el centro de la imagen.
   • Agregación Intra-Anillo: Dentro de cada anillo, los píxeles se recorren en un bucle cerrado (alternando el sentido horario y antihorario para los anillos adyacentes). Un SSM selectivo ligero procesa este bucle, y la salida se agrega en un descriptor de anillo mediante promediado. Esta agregación es agnóstica al orden; una rotación simplemente induce un desplazamiento cíclico en el bucle, lo cual no altera el descriptor promediado.
   • Propagación Inter-Radial: Los descriptores de los anillos se procesan desde el anillo más interno hacia el más externo utilizando un segundo SSM corto para propagar el contexto radialmente. Este diseño de dos niveles mantiene la complejidad de tiempo lineal ($O(HW)$) al tiempo que evita la fragilidad de la serialización basada en trayectorias globales.

2. Filtrado de Canales Parcial (PCF):

   • Para mejorar la eficiencia, los autores introducen un mecanismo de enrutamiento rígido que filtra los canales de características antes de que entren en la vía recurrente del anillo.
   • La prominencia del canal se determina mediante el Promedio Global de Características (GAP). Los canales con una prominencia por encima del umbral medio se enrutan a través del SSM de tipo anillo, mientras que los canales restantes evitan esta computación pesada mediante una rama residual ligera.
   • Este enfoque reduce los FLOPs y la memoria de activación manteniendo la precisión, contrastando con los mecanismos de atención suave (por ejemplo, SE, CBAM) que procesan todos los canales aguas abajo.

3. Escritura de Retorno (Write-Back) y Fusión:

   • Las salidas radiales se proyectan de nuevo a la cuadrícula de píxeles y se fusionan con la corriente del backbone mediante una proyección 1×1 y una conexión residual, preservando las dimensiones espaciales.

Contribuciones Clave

• Análisis Sistemático del Orden de Escaneo: El artículo proporciona el primer estudio sistemático que demuestra que el orden de escaneo es un hiperparámetro de "primer orden". Muestra que los escaneos comunes de tipo raster y serpentina fracturan la continuidad de los objetos y son altamente sensibles a la rotación, mientras que los escaneos basados en anillos mantienen la estabilidad.
• Trayectoria Estable a la Rotación: La introducción de un esquema de escaneo basado en anillos que agrupa píxeles en anillos concéntricos y realiza una agregación agnóstica al orden. Esto preserva el núcleo de tiempo lineal de los SSM al tiempo que elimina la necesidad de remapeo polar o aumentos de datos específicos de rotación.
• Filtrado de Canales Parcial (PCF): Un módulo de eficiencia novedoso que enruta solo los canales más informativos a través de la vía recurrente. Esto reduce el costo computacional y mejora el rendimiento (throughput) mientras aumenta ligeramente la precisión al centrar las actualizaciones recurrentes en características de alta señal.
• Arquitectura Unificada (PRISMamba): Un modelo integrado que combina estos elementos para lograr un rendimiento de vanguardia entre los SSM de visión.

Resultados Experimentales
Los autores evalúan PRISMamba en ImageNet-1K, MS COCO, y bajo pruebas de estrés de rotación/oclusión:

• Clasificación ImageNet-1K: PRISMamba logra una precisión Top-1 del 84.5% con 3.9G FLOPs y un rendimiento de 3,054 img/s en un GPU A100. Esto supera a VMamba (82.6%, 5.6G FLOPs, 1,686 img/s) por un 1.9% en precisión utilizando un 30% menos de FLOPs y logrando aproximadamente 1.8× más throughput.
• Detección/Segmentación MS COCO: Utilizando un esquema 1× (1280×800), PRISMamba alcanza 48.9 APbox y 43.2 APmask con 235G FLOPs, superando a VMamba y GroupMamba utilizando entre un 10–19% menos de FLOPs.
• Robustez a la Rotación: Bajo rotaciones de 30° y 60°, PRISMamba mantiene una estabilidad casi perfecta (cayendo solo 0.1–0.2%), mientras que las bases de trayectoria fija (VMamba, GroupMamba, etc.) se degradan entre un 1.0–2.0%.
• Robustez a la Oclusión: En pruebas de oclusión por máscara aleatoria, PRISMamba muestra caídas de rendimiento significativamente menores (≤0.6%) en comparación con las variantes de trayectoria fija (hasta 1.3%), atribuido a la dilución del daño local mediante el promedio de anillo y la propagación radial.
• Ablación: El módulo PCF por sí solo reduce los FLOPs en aproximadamente un 15% e incrementa el throughput en un ~40% mientras mejora la precisión en un 0.4%, demostrando que el filtrado de canales rígido es más efectivo que la atención suave para los SSM recurrentes.

Significado y Reivindicaciones
El artículo sostiene que el diseño del orden de escaneo, junto con la asignación de canales dirigida, es un factor crucial pero poco explorado para la precisión, la eficiencia y la robustez en los SSM de Visión. PRISMamba demuestra que es posible resolver el desajuste entre secuencia y geometría en los modelos basados en trayectorias sin aumentar el tamaño del modelo o requerir un entrenamiento especializado. Los autores posicionan su trabajo como un enfoque basado en principios y de bajo costo para hacer que los SSM de Visión sean más robustos a las transformaciones geométricas y más eficientes mediante el procesamiento selectivo de canales.

Limitaciones y Trabajo Futuro
Los autores reconocen que su implementación actual depende de un centro de imagen fijo y un ancho de anillo discreto, lo que puede ser subóptimo para sujetos descentrados o relaciones de aspecto extremas. El trabajo futuro incluye el aprendizaje de los orígenes y anchos de los anillos, el diseño de particiones adaptativas al contenido y la extensión del concepto de trayectoria de anillo a dominios espacio-temporales (video) y volúmenes 3D.