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

Resumo Técnico: Varredura de Anel Parcial: Revisitando a Ordem de Varredura em Modelos de Espaço de Estado de Visão

Definição do Problema
Modelos de Espaço de Estados (SSMs), particularmente variantes de Vision Mamba, emergiram como alternativas eficientes aos mecanismos de atenção para tarefas de visão, oferecendo processamento de sequência em tempo linear. No entanto, esses modelos exigem a serialização de imagens 2D em sequências de tokens 1D baseadas em uma ordem de varredura predefinida (ex: raster, serpentina). Os autores argumentam que essa ordem de varredura é um fator crítico, frequentemente negligenciado, que dita a adjacência espacial dentro da sequência.

As varreduras convencionais de caminho fixo (raster, bidirecional, diagonal) preservam o alinhamento espaço-sequência apenas sob transformações específicas, como reflexões (flips). Sob rotações no plano, esses caminhos sofrem com reindexação global e preenchimento (padding), o que fratura a continuidade de objetos e desalinha o vizinhança do núcleo recorrente. Isso força o SSM a gastar capacidade reparando a coerência local em vez de modelar semântica, levando à degradação do desempenho (queda de 1–2% no Top-1) e redução da robustez a transformações geométricas.

Metodologia: PRISMamba
O artigo propõe o Partial Ring Scan Mamba (PRISMamba), uma arquitetura robusta à rotação, projetada para alinhar a adjacência da sequência com a estrutura geométrica sem exigir treinamento específico para rotação. O método consiste em três componentes principais:

1. Varredura Alternada Anel por Anel:

   • Particionamento: A imagem é particionada em anéis concêntricos baseados na distância euclidiana dos pixels em relação ao centro da imagem.
   • Agregação Intra-Anel: Dentro de cada anel, os pixels são percorridos em um loop fechado (alternando entre sentido horário e anti-horário para anéis adjacentes). Um SSM seletivo leve processa este loop, e a saída é agregada em um descritor de anel via média. Esta agregação é agnóstica à ordem; uma rotação induz apenas um deslocamento cíclico no loop, o que não altera o descritor médio.
   • Propagação Inter-Radial: Os descritores de anel são processados do anel mais interno para o mais externo usando um segundo SSM curto para propagar o contexto radialmente. Este design de dois níveis mantém a complexidade de tempo linear ($O(HW)$) enquanto evita a fragilidade da serialização baseada em caminhos globais.

2. Filtragem Parcial de Canais (PCF):

   • Para melhorar a eficiência, os autores introduzem um mecanismo de roteamento rígido que filtra canais de características antes que eles entrem no caminho de anel recorrente.
   • A saliência do canal é determinada via Global Average Pooling (GAP). Canais com saliência acima do limiar médio são roteados pelo SSM de anel, enquanto os demais canais contornam esta computação pesada através de um ramo residual leve.
   • Esta abordagem reduz FLOPs e memória de ativação enquanto mantém a precisão, contrastando com mecanismos de atenção suave (ex: SE, CBAM) que processam todos os canais downstream.

3. Escrita de Volta (Write-Back) e Fusão:

   • As saídas radiais são transmitidas de volta para a grade de pixels e fundidas com o fluxo do backbone via uma projeção 1×1 e conexão residual, preservando as dimensões espaciais.

Principais Contribuições

• Análise Sistemática da Ordem de Varredura: O artigo fornece o primeiro estudo sistemático demonstrando que a ordem de varredura é um hiperparâmetro de "primeira classe". Mostra que varreduras raster e serpentina comuns fraturam a continuidade de objetos e são altamente sensíveis à rotação, enquanto varreduras baseadas em anéis mantêm a estabilidade.
• Traversal Robusto à Rotação: A introdução de um esquema de varredura baseado em anéis que agrupa pixels em anéis concêntricos e realiza agregação agnóstica à ordem. Isso preserva o núcleo de tempo linear dos SSMs enquanto elimina a necessidade de remapeamento polar ou aumento de dados específico para rotação.
• Filtragem Parcial de Canais (PCF): Um novo módulo de eficiência que roteia apenas os canais mais informativos através do caminho recorrente. Isso reduz o custo computacional e melhora o throughput enquanto aumenta ligeiramente a precisão ao focar as atualizações recorrentes em características de alto sinal.
• Arquitetura Unificada (PRISMamba): Um modelo integrado que combina esses elementos para alcançar o estado da arte em desempenho entre os Vision SSMs.

Resultados Experimentais
Os autores avaliam o PRISMamba no ImageNet-1K, MS COCO e sob testes de estresse de rotação/oclusão:

• Classificação ImageNet-1K: O PRISMamba alcança 84,5% de precisão Top-1 com 3,9G FLOPs e 3.054 img/s de throughput em uma GPU A100. Isso supera o VMamba (82,6%, 5,6G FLOPs, 1.686 img/s) em 1,9% em precisão, utilizando 30% menos FLOPs e alcançando ~1,8× maior throughput.
• Detecção/Segmentação MS COCO: Usando um cronograma de 1× (1280×800), o PRISMamba alcança 48,9 APbox e 43,2 APmask com 235G FLOPs, superando o VMamba e o GroupMamba enquanto utiliza 10–19% menos FLOPs.
• Robustez à Rotação: Sob rotações de 30° e 60°, o PRISMamba mantém uma estabilidade quase perfeita (caindo apenas 0,1–0,2%), enquanto as bases de caminho fixo (VMamba, GroupMamba, etc.) degradam em 1,0–2,0%.
• Robustez à Oclusão: Em testes de oclusão por máscara aleatória, o PRISMamba mostra quedas de desempenho significativamente menores (≤0,6%) comparado às variantes de caminho fixo (até 1,3%), atribuído à diluição do dano local via média de anel e propagação radial.
• Ablação: O módulo PCF sozinho reduz os FLOPs em ~15% e aumenta o throughput em ~40% enquanto melhora a precisão em 0,4%, demonstrando que a filtragem de canais rígida é mais eficaz do que a atenção suave para SSMs recorrentes.

Significância e Alegações
O artigo afirma que o design da ordem de varredura, aliado à alocação direcionada de canais, é um fator crucial, porém subexplorado, para precisão, eficiência e robustez em Vision SSMs. O PRISMamba demonstra que é possível resolver o descompasso entre sequência e geometria em modelos baseados em caminho sem aumentar o tamanho do modelo ou exigir treinamento especializado. Os autores posicionam seu trabalho como uma abordagem fundamentada e de baixo custo para tornar os Vision SSMs mais robustos a transformações geométricas e mais eficientes através do processamento seletivo de canais.

Limitações e Trabalho Futuro
Os autores reconhecem que sua implementação atual depende de um centro de imagem fixo e largura de anel discreta, o que pode ser suboptimal para sujeitos fora do centro ou razões de aspecto extremas. Trabalhos futuros incluem aprender origens e larguras de anéis, projetar partições adaptativas ao conteúdo e estender o conceito de travessia de anel para domínios espaço-temporais (vídeo) e volumes 3D.