Diffusion Model as a Generalist Segmentation Learner

1. Énoncé du problème

Les systèmes de segmentation actuels sont fortement fragmentés. Les modèles sont généralement conçus pour des tâches spécifiques (sémantique, instance, panoptique), des vocabulaires spécifiques (ensemble fermé vs vocabulaire ouvert) ou des domaines spécifiques (naturel, médical, télédétection). Le passage entre ces configurations nécessite souvent une refonte des architectures ou un réentraînement complet.

Bien que des recherches récentes aient tenté de réutiliser des modèles de diffusion texte-vers-image préentraînés pour la segmentation en extrayant des cartes d'attention, ces approches souffrent de limitations significatives :

• Bruit et incohérence : Les cartes d'attention brutes sont souvent de faible résolution, bruyantes et incohérentes entre les couches.
• Dépendance au post-traitement : Elles nécessitent fréquemment un post-traitement heuristique lourd pour produire des masques utilisables.
• Manque de généralisation : Les méthodes existantes se concentrent souvent sur une seule tâche et échouent à démontrer une généralisation robuste à travers des domaines divers et des requêtes à vocabulaire ouvert sans personnalisation architecturale.

La question centrale abordée par ce papier est : Un modèle unique et unifié basé sur la diffusion peut-il servir d'apprenant généraliste en segmentation, capable de gérer des tâches sémantiques, à vocabulaire ouvert et inter-domaines sans modifications architecturales spécifiques au domaine ?

2. Méthodologie : Cadre DiGSeg

Les auteurs proposent DiGSeg, un cadre qui transforme un modèle de diffusion préentraîné (spécifiquement Stable Diffusion v2) d'une prière générative en un apprenant de segmentation explicite via un protocole de fine-tuning unifié.

Architecture de base

Le cadre se compose de trois composants principaux :

1. Voie latente visuelle :

   • Au lieu de traiter l'image comme des pixels bruts, l'image RVB d'entrée ($x$) et le masque de segmentation vérité terrain ($y$) sont encodés dans un espace latent compact en utilisant l'encodeur VAE figé de Stable Diffusion.
   • Pour maintenir la compatibilité avec le VAE à 3 canaux, le masque à canal unique est répliqué sur trois canaux.
   • Cela préserve la structure spatiale tout en permettant un apprentissage efficace des correspondances fines dans l'espace latent.

2. Conditionnement textuel aligné CLIP :

   • Pour permettre des capacités à vocabulaire ouvert, un encodeur de texte CLIP figé génère des embeddings ($t_{clip}$) pour les noms de classes ou les descriptions en langage naturel.
   • Ces embeddings sont injectés dans le U-Net de diffusion via une attention croisée à plusieurs échelles de débruitage. Cela permet au modèle d'aligner les requêtes textuelles avec les représentations visuelles en évolution, globalement et localement.

3. U-Net de débruitage cohérent avec la segmentation :

   • Objectif d'entraînement : Le modèle est entraîné à prédire le bruit ajouté au latent de segmentation ($z_y$), conditionné par le latent d'image ($z_x$) et l'embedding textuel ($t_{clip}$).
   • Formulation générative : Contrairement aux travaux précédents qui utilisent la diffusion comme extracteur de caractéristiques, DiGSeg traite la segmentation comme une tâche de génération de débruitage conditionnelle : $p_\theta(y|x)$. Le U-Net apprend à débruiter depuis un bruit aléatoire vers une représentation latente cohérente avec la segmentation.
   • Stratégie de bruit : Un calendrier de bruit multi-échelle atténué novateur est employé. Les perturbations haute fréquence dominent les premières étapes pour apprendre les structures fines, tandis que les composantes basse fréquence deviennent influentes plus tard pour assurer la cohérence sémantique.

Processus d'inférence

1. L'image d'entrée est encodée en une représentation latente.
2. Un latent de segmentation est initialisé comme du bruit gaussien.
3. Sous conditionnement textuel, le U-Net débruite itérativement le latent (en utilisant un échantillonneur DDIM avec une stratégie de pas de temps "trailing" pour l'efficacité).
4. Le latent propre final est décodé dans l'espace des pixels en utilisant le décodeur VAE figé pour produire le masque de segmentation.
5. Ensemble au moment du test : Plusieurs passes d'inférence avec différents seeds de bruit sont moyennées dans l'espace latent pour améliorer la cohérence spatiale et réduire les artefacts.

3. Contributions clés

• Cadre unifié : DiGSeg est le premier cadre à réussir la réutilisation d'une colonne vertébrale de diffusion préentraînée en un apprenant généraliste en segmentation capable de gérer des tâches à vocabulaire fermé, à vocabulaire ouvert et inter-domaines sans modifications architecturales spécifiques à la tâche.
• Apprentissage latent direct : Au lieu de s'appuyer sur des cartes d'attention bruyantes, la méthode affine le U-Net de débruitage pour générer directement des latents cohérents avec la segmentation, produisant des masques de haute qualité sans post-traitement lourd.
• Conditionnement textuel multi-échelle : Introduit un mécanisme pour injecter des caractéristiques linguistiques à travers plusieurs échelles de débruitage, permettant un alignement robuste entre des invites textuelles arbitraires et des caractéristiques visuelles.
• Stratégie de bruit novatrice : Propose un calendrier de bruit multi-échelle atténué qui équilibre l'apprentissage des contours fins avec la cohérence sémantique globale.

4. Résultats expérimentaux

Les auteurs ont évalué DiGSeg sur des benchmarks standards et des domaines spécialisés, démontrant des performances à l'état de l'art (SOTA).

• Segmentation à vocabulaire ouvert :
  • A surpassé 13 lignes de base (y compris ODISE, OVSeg, CAT-Seg) sur A-847, PC-459, A-150, PC-59 et Cityscapes.
  • A réalisé des gains significatifs (par exemple, +2,8 mIoU sur PC-59, +2,2 sur PC-459) par rapport aux deuxièmes meilleures méthodes.
• Segmentation sémantique standard :
  • A atteint le SOTA sur COCO (50,8 mIoU) et ADE20K (58,6 mIoU), surpassant de puissants modèles basés sur des transformateurs comme Mask2Former et OneFormer.
• Transfert inter-domaines :
  • A démontré un fort transfert zero-shot vers des données médicales (REFUGE-2), agricoles (PhenoBench) et de télédétection (DeepGlobe) sans modification architecturale.
  • Note : Les performances dans les domaines médicaux étaient inférieures à celles des scènes naturelles, attribuées à l'écart de domaine dans les données de préentraînement de CLIP, mais restaient compétitives.
• Efficacité et utilisation des données :
  • Le modèle est très efficace en termes de données, atteignant des performances presque identiques avec seulement 50 % des données d'entraînement.
  • La vitesse d'inférence est améliorée grâce à DDIM-trailing (atteignant la saturation en 1 étape), bien qu'elle reste plus lente que les modèles feed-forward en raison de la nature de la diffusion.

5. Importance

Ce travail comble fondamentalement le fossé entre la génération visuelle et la compréhension visuelle.

• Changement de paradigme : Il remet en question l'idée que les modèles de diffusion sont uniquement destinés à la synthèse, prouvant qu'ils peuvent servir d'apprenants discriminatifs puissants et généralistes lorsqu'ils sont correctement conditionnés.
• Généralisation : Il démontre que les priors spatiaux et sémantiques riches encodés dans les modèles génératifs à grande échelle sont transférables aux tâches de prédiction dense à travers des domaines divers (naturel, médical, agricole).
• Simplicité : En évitant les pipelines multi-étapes complexes ou les têtes spécialisées, DiGSeg offre une interface unifiée et rationalisée pour la segmentation, suggérant un avenir où une seule colonne vertébrale générative peut gérer une large gamme de tâches de vision par ordinateur.

En conclusion, DiGSeg établit que les architectures de diffusion modernes, lorsqu'elles sont affinées avec un protocole conditionné par le latent et un alignement textuel, peuvent surpasser les modèles de segmentation spécialisés tout en offrant une flexibilité supérieure dans les scénarios à vocabulaire ouvert et inter-domaines.