Résumé Technique : Patch-MoE Mamba

Énoncé du Problème

La segmentation d'images médicales repose fortement sur les réseaux de neurones convolutifs (CNN) et les Transformers. Cependant, les CNN sont intrinsèquement limités par leurs champs récepteurs locaux, peinant à modéliser les dépendances à long terme, tandis que les Transformers entraînent souvent une complexité computationnelle et mémoire quadratique. Récemment, les modèles d'espace d'état (SSM), en particulier les architectures basées sur Mamba, ont émergé comme des alternatives efficaces avec une complexité de séquence linéaire.

Malgré leur promesse, les modèles de segmentation existants basés sur Mamba font face à deux limitations critiques :

1. Perturbation de la Structure Spatiale Locale : La plupart des modèles convertissent les cartes de caractéristiques 2D en séquences 1D en utilisant un balayage raster pixel par pixel. Ce processus peut séparer des pixels spatialement adjacents dans la séquence, affaiblissant la cohérence spatiale 2D locale. Cela est préjudiciable aux tâches de prédiction dense nécessitant une localisation précise des contours et une préservation de la structure locale.
2. Fusion Directionnelle Rigide : Les méthodes actuelles fusionnent généralement les sorties de plusieurs directions de balayage par une simple sommation. Cette agrégation fixe suppose que toutes les directions sont également importantes à chaque emplacement spatial, échouant à s'adapter aux tailles, formes, textures et complexités de contours divers des cibles médicales (par exemple, de petites lésions à faible contraste versus de grandes structures irrégulières).

Méthodologie

Les auteurs proposent Patch-MoE Mamba, une architecture d'espace d'état de type Mélange d'Experts (MoE) ordonnée par patchs, conçue pour résoudre les limitations susmentionnées. L'architecture suit une conception de style U-Net, remplaçant les blocs standards Visual State Space (VSS) de VM-UNetV2 par un nouveau bloc VSS Patch-MoE.

1. Balayage Hiérarchique Ordonné par Patchs

Au lieu de balayer les pixels individuels par raster, la méthode partitionne les cartes de caractéristiques en patchs locaux et effectue un balayage directionnel de manière ordonnée par patchs.

• Mécanisme : Étant donné une carte de caractéristiques, la grille spatiale est divisée en patchs non chevauchants de taille $p \times p$. À l'intérieur de chaque patch, les pixels sont énumérés dans l'ordre ligne par ligne avant de passer au patch suivant.
• Avantage : Cela garantit que tous les pixels au sein d'un patch restent consécutifs dans la séquence, préservant les voisinages spatiaux locaux tout en permettant toujours la modélisation de dépendances à long terme à travers les patchs.
• Stratégie Hiérarchique : Pour capturer un contexte multi-échelle, le modèle utilise un ensemble de tailles de patchs $\{p^{(1)}, p^{(2)}, \dots\}$. Les petits patchs capturent les détails fins des contours, tandis que les grands patchs modélisent des structures plus grossières. Différentes directions de balayage (avant, arrière, avant largeur-hauteur, arrière largeur-hauteur) peuvent utiliser des tailles de patchs distinctes pour gérer des structures anisotropes.

2. Fusion Directionnelle Basée sur MoE

Pour remplacer la sommation fixe, les auteurs introduisent un module de fusion adaptatif traitant les sorties directionnelles comme des « experts ».

• Construction des Experts : Le module utilise cinq experts :
  • Quatre experts spécifiques à la direction ($E_1$ à $E_4$) : Les sorties normalisées des quatre scanners Mamba directionnels.
  • Un expert de concaténation ($E_{concat}$) : Formé en concaténant les quatre cartes normalisées le long de la dimension des canaux et en les projetant de nouveau via une convolution $1\times1$. Cet expert capture les corrélations inter-directionnelles et fournit une vue holistique, agnostique de la direction.
• Gating Sensible à l'Espace : Un réseau routeur calcule des poids dépendants de la localisation pour fusionner adaptativement ces experts. Le routeur est piloté par un mélange de contexte local (via une convolution par profondeur) et de contexte global (via un pooling moyen adaptatif), contrôlé par un scalaire apprenable $\alpha$.
• Stabilisation Résiduelle : Une déviation résiduelle ajoute les sorties directionnelles brutes au résultat fusionné ($Z_l = \hat{Y}_l + \sum Y_l^{(i)}$). Cela assure des réponses de base robustes et préserve les signaux directionnels forts, stabilisant l'entraînement.

3. Infusion de Sémantique et de Détails (SDI)

Les sorties de l'encodeur sont traitées par un module SDI (adapté de U-Net v2) pour améliorer les cartes de caractéristiques en infusant des informations sémantiques des niveaux supérieurs et en intégrant des détails plus fins des niveaux inférieurs via le produit de Hadamard.

Contributions Clés

1. Balayage Ordonné par Patchs : Un mécanisme de balayage novateur qui remplace le balayage raster pixel par pixel par une traversal ordonnée par patchs, préservant efficacement la cohérence spatiale 2D locale tout en maintenant l'efficacité des modèles d'espace d'état.
2. Fusion Directionnelle Basée sur MoE : Un module de fusion adaptatif qui remplace la simple sommation par un réseau de gating apprenable. Il pondère dynamiquement plusieurs experts directionnels et un expert de concaténation en fonction du contexte spatial, permettant au modèle de s'adapter à des échelles et complexités de lésions variables.
3. Modélisation Hiérarchique Multi-échelle : L'intégration de tailles de patchs hiérarchiques permet au réseau de capturer simultanément les détails fins des contours et les structures sémantiques grossières.
4. Stabilisation Résiduelle : L'inclusion d'un chemin résiduel pour les sorties directionnelles brutes assure la stabilité de l'entraînement et empêche la perte de signaux directionnels critiques.

Résultats Expérimentaux

Les auteurs ont évalué Patch-MoE Mamba sur cinq ensembles de données publics de segmentation de polypes (Kvasir-SEG, ClinicDB, ColonDB, ETIS, CVC-300) et deux ensembles de données de segmentation de lésions cutanées (ISIC 2017, ISIC 2018).

• Performance : Patch-MoE Mamba a atteint des performances de pointe sur tous les ensembles de données, surpassant constamment les bases CNN (U-Net, U-Net v2) et Mamba (VM-UNet, VM-UNetV2).
  • Sur l'ensemble de données ETIS (connu pour son faible contraste et ses petits polypes), le modèle a amélioré le score Dice de 72,56 % (VM-UNetV2) à 74,04 %.
  • Sur ISIC 2018, il a obtenu un score Dice de 89,34 %, dépassant le meilleur précédent de 88,36 %.
• Études d'Ablation :
  • Le remplacement du balayage raster par un balayage ordonné par patchs seul a amélioré le Dice moyen de 83,46 % à 84,02 %.
  • L'ajout du module de fusion MoE a encore porté la moyenne à 84,30 %.
  • Des tailles de patchs hiérarchiques (grossières dans les premières étapes, fines dans les étapes plus profondes) ont produit les meilleurs résultats.
• Complexité Computationnelle : Bien que le balayage ordonné par patchs lui-même n'augmente pas le coût computationnel (le nombre de tokens restant inchangé), l'inclusion de l'expert de concaténation augmente les paramètres à ~70M et les FLOPs à ~28G. Les auteurs notent que cette surcharge est attendue en raison des interactions de canaux plus denses requises pour le mécanisme MoE.

Importance et Revendications

L'article revendique que Patch-MoE Mamba offre une solution générale et efficace pour la segmentation d'images médicales. En abordant les limitations spécifiques des modèles Mamba existants, à savoir la perte de structure spatiale locale et la rigidité de la fusion directionnelle, l'architecture proposée démontre une adaptabilité supérieure à des cibles médicales diverses avec des tailles, formes et complexités de contours variables. Les résultats suggèrent que la préservation des voisinages spatiaux locaux grâce au balayage ordonné par patchs, combinée à une fusion adaptative d'experts, améliore considérablement la localisation des contours et la précision de la segmentation, en particulier dans des scénarios difficiles impliquant des lésions à faible contraste ou de petite taille.