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

Résumé Technique : Balayage par Anneaux Partiels : Reconsidérer l'Ordre de Balayage dans les Modèles d'Espace d'États de Vision

Énoncé du Problème
Les modèles d'espace d'états (SSM), en particulier les variantes de Vision Mamba, ont émergé comme des alternatives efficaces aux mécanismes d'attention pour les tâches de vision, offrant un traitement de séquence en temps linéaire. Cependant, ces modèles nécessitent de sérialiser les images 2D en séquences de tokens 1D basées sur un ordre de balayage prédéfini (ex. : raster, serpentine). Les auteurs soutiennent que cet ordre de balayage est un facteur critique, souvent négligé, qui dicte l'adjacence spatiale au sein de la séquence.

Les balayages conventionnels à chemin fixe (raster, bidirectionnel, diagonal) ne préservent l'alignement espace-séquence que sous certaines transformations spécifiques comme les retournements (flips). Sous des rotations intra-planaires, ces chemins souffrent d'un réindexage global et de remplissage (padding), ce qui fracture la continuité des objets et désaligne le voisinage du noyau récurrent. Cela force le SSM à dépenser sa capacité pour réparer la cohérence locale plutôt que pour modéliser la sémantique, entraînant une dégradation des performances (chute de 1 à 2 % de la Top-1) et une réduction de la robustesse aux transformations géométriques.

Méthodologie : PRISMamba
Le papier propose PRISMamba (Partial Ring Scan Mamba), une architecture robuste à la rotation, conçue pour aligner l'adjacence de la séquence avec la structure géométrique sans nécessiter d'entraînement spécifique à la rotation. La méthode se compose de trois composantes clés :

1. Balayage Alterné Anneau par Anneau :

   • Partitionnement : L'image est partitionnée en anneaux concentriques basés sur la distance euclidienne des pixels par rapport au centre de l'image.
   • Agrégation Intra-Anneau : Au sein de chaque anneau, les pixels sont parcourus selon une boucle fermée (alternant entre le sens horaire et antihoraire pour les anneaux adjacents). Un SSM sélectif léger traite cette boucle, et la sortie est agrégée en un descripteur d'anneau via une moyenne. Cette agrégation est indépendante de l'ordre ; une rotation induit simplement un décalage cyclique dans la boucle, ce qui ne modifie pas le descripteur moyen.
   • Propagation Inter-Radiale : Les descripteurs d'anneaux sont traités de l'anneau le plus intérieur vers l'anneau le plus extérieur à l'aide d'un second SSM court pour propager le contexte radialement. Cette conception à deux niveaux maintient une complexité en temps linéaire ($O(HW)$) tout en évitant la fragilité de la sérialisation par chemin global.

2. Filtrage de Canaux Partiel (PCF - Partial Channel Filtering) :

   • Pour améliorer l'efficacité, les auteurs introduisent un mécanisme de routage strict qui filtre les canaux de caractéristiques avant qu'ils n'entrent dans la voie récurrente de l'anneau.
   • La saillance des canaux est déterminée via un Global Average Pooling (GAP). Les canaux dont la saillance est supérieure au seuil moyen sont routés à travers le SSM par anneau, tandis que les canaux restants contournent cette computation lourde via une branche résiduelle légère.
   • Cette approche réduit les FLOPs et la mémoire d'activation tout en maintenant la précision, contrastant avec les mécanismes d'attention douce (ex. : SE, CBAM) qui traitent tous les canaux en aval.

3. Réécriture et Fusion (Write-Back and Fusion) :

   • Les sorties radiales sont redistribuées sur la grille de pixels et fusionnées avec le flux de base via une projection 1×1 et une connexion résiduelle, préservant ainsi les dimensions spatiales.

Principales Contributions

• Analyse Systématique de l'Ordre de Balayage : Le papier fournit la première étude systématique démontrant que l'ordre de balayage est un hyperparamètre de « premier rang ». Il montre que les balayages raster et serpentine courants fracturent la continuité des objets et sont très sensibles à la rotation, tandis que les balayages par anneaux maintiennent la stabilité.
• Traversée Stable à la Rotation : L'introduction d'un schéma de balayage par anneaux qui groupe les pixels en anneaux concentriques et effectue une agrégation indépendante de l'ordre. Cela préserve le cœur en temps linéaire des SSM tout en éliminant le besoin de remappage polaire ou d'augmentation de données spécifique à la rotation.
• Filtrage de Canaux Partiel (PCF) : Un nouveau module d'efficacité qui route uniquement les canaux les plus informatifs à travers la voie récurrente. Cela réduit le coût computationnel et améliore le débit tout en améliant légèrement la précision en concentrant les mises à jour récurrentes sur les caractéristiques à haut signal.
• Architecture Unifiée (PRISMamba) : Un modèle intégré qui combine ces éléments pour atteindre des performances de pointe parmi les Vision SSMs.

Résultats Expérimentaux
Les auteurs évaluent PRISMamba sur ImageNet-1K, MS COCO, et sous des tests de stress de rotation/occlusion :

• Classification ImageNet-1K : PRISMamba atteint une précision Top-1 de 84,5 % avec 3,9G FLOPs et un débit de 3 054 img/s sur un GPU A100. Cela surpasse VMamba (82,6 %, 5,6G FLOPs, 1 686 img/s) de 1,9 % en précision tout en utilisant 30 % de FLOPs en moins et en atteignant un débit environ 1,8× plus élevé.
• Détection/Segmentation MS COCO : En utilisant un planning 1× (1280×800), PRISMamba atteint 48,9 APbox et 43,2 APmask avec 235G FLOPs, surpassant VMamba et GroupMamba tout en utilisant 10–19 % de FLOPs en moins.
• Robustesse à la Rotation : Sous des rotations de 30° et 60°, PRISMamba maintient une stabilité quasi parfaite (chute de seulement 0,1–0,2 %), alors que les bases à chemin fixe (VMamba, GroupMamba, etc.) se dégradent de 1,0–2,0 %.
• Robustesse à l'Occlusion : Dans les tests d'occlusion par masquage aléatoire, PRISMamba montre des chutes de performance nettement plus faibles (≤0,6 %) par rapport aux variantes à chemin fixe (jusqu'à 1,3 %), attribuées à la dilution des dommages locaux via la moyenne des anneaux et la propagation radiale.
• Ablation : Le module PCF seul réduit les FLOPs d'environ 15 % et augmente le débit de ~40 % tout en améliorant la précision de 0,4 %, démontrant que le filtrage de canaux strict est plus efficace que l'attention douce pour les SSM récurrents.

Signification et Revendications
Le papier affirme que la conception de l'ordre de balayage, couplée à une allocation de canaux ciblée, est un facteur crucial mais sous-exploré pour la précision, l'efficacité et la robustesse dans les Vision SSMs. PRISMamba démontre qu'il est possible de résoudre le décalage séquence-géométrie dans les modèles par chemin sans augmenter la taille du modèle ni nécessiter un entraînement spécialisé. Les auteurs positionnent leur travail comme une approche fondée et à faible coût pour rendre les Vision SSMs plus robustes aux transformations géométriques et plus efficaces grâce au traitement sélectif des canaux.

Limites et Travaux Futurs
Les auteurs reconnaissent que leur implémentation actuelle repose sur un centre d'image fixe et une largeur d'anneau discrète, ce qui peut être sous-optimal pour les sujets décentrés ou les rapports d'aspect extrêmes. Les travaux futurs incluent l'apprentissage des origines et des largeurs d'anneaux, la conception de partitions adaptatives au contenu, et l'extension du concept de traversée d'anneau aux domaines spatio-temporels (vidéo) et aux volumes 3D.