Diffusion Model as a Generalist Segmentation Learner

1. Probleemstelling

Huidige segmentatiesystemen zijn sterk gefragmenteerd. Modellen zijn doorgaans op maat gemaakt voor specifieke taken (semantisch, instance, panoptisch), specifieke vocabulaires (gesloten set versus open vocabulaire) of specifieke domeinen (natuurlijk, medisch, remote sensing). De overgang tussen deze instellingen vereist vaak het opnieuw ontwerpen van architecturen of het volledig opnieuw trainen van modellen.

Hoewel recent onderzoek geprobeerd heeft om voorgeprente text-to-image diffusion-modellen te hergebruiken voor segmentatie door attention-maps te extraheren, lijden deze benaderingen aan aanzienlijke beperkingen:

• Ruis en inconsistentie: Raw attention-maps zijn vaak van lage resolutie, ruisachtig en inconsistent tussen lagen.
• Afhankelijkheid van post-processing: Ze vereisen vaak zware heuristische post-processing om bruikbare maskers te produceren.
• Gebrek aan generalisatie: Bestaande methoden richten zich vaak op een enkele taak en falen in het aantonen van robuuste generalisatie over diverse domeinen en open-vocabulary queries zonder architecturale aanpassingen.

De kernvraag die dit paper adresseert is: Kan een enkel, geünificeerd op diffusion gebaseerd model dienen als een generalist segmentatieleerder die semantische, open-vocabulary en cross-domein taken aankan zonder domeinspecifieke architecturale wijzigingen?

2. Methodologie: DiGSeg Framework

De auteurs stellen DiGSeg voor, een framework dat een voorgeprent diffusion-model (specifiek Stable Diffusion v2) transformeert van een generatieve prior naar een expliciete segmentatieleerder via een geünificeerd fine-tuning-protocol.

Kernarchitectuur

Het framework bestaat uit drie hoofdcomponenten:

1. Visuele Latente Pad:

   • In plaats van de afbeelding als raw pixels te behandelen, worden de invoer RGB-afbeelding ($x$) en het ground-truth segmentatiemasker ($y$) gecodeerd naar een compacte latente ruimte met behulp van de bevroren VAE-encoder van Stable Diffusion.
   • Om compatibiliteit met de 3-kanaals VAE te behouden, wordt het enkelkanaals masker over drie kanalen gerepliceerd.
   • Dit behoudt de ruimtelijke structuur terwijl het efficiënt leren van fijnmazige correspondenties in de latente ruimte mogelijk maakt.

2. CLIP-geregistreerde Tekstconditionering:

   • Om open-vocabulary capaciteiten mogelijk te maken, genereert een bevroren CLIP-tekstencoder embeddings ($t_{clip}$) voor klassenamen of natuurlijke taalbeschrijvingen.
   • Deze embeddings worden via cross-attention op meerdere denoising-schalen ingebracht in de diffusion U-Net. Dit stelt het model in staat tekstuele queries globaal en lokaal af te stemmen op evoluerende visuele representaties.

3. Segmentatie-consistente Denoising U-Net:

   • Trainingsdoel: Het model wordt getraind om ruis te voorspellen die is toegevoegd aan de segmentatielaten ($z_y$), geconditioneerd op de afbeeldingslatent ($z_x$) en tekstembedding ($t_{clip}$).
   • Generatieve Formulering: In tegenstelling tot eerdere werken die diffusion gebruiken als feature-extractor, behandelt DiGSeg segmentatie als een conditionele denoising-generatietaken: $p_\theta(y|x)$. De U-Net leert om te denoisen van willekeurige ruis naar een segmentatie-consistente latente representatie.
   • Ruisstrategie: Er wordt een nieuwe geannealde multi-scale ruisplanning toegepast. Hoogfrequente verstoringen domineren de vroege stappen om fijne structuren te leren, terwijl laagfrequente componenten later invloedrijk worden om semantische coherentie te waarborgen.

Inferentieproces

1. De invoerafbeelding wordt gecodeerd naar een latente representatie.
2. Een segmentatielaten wordt geïnitieerd als Gaussische ruis.
3. Onder tekstconditionering denoist de U-Net de latent iteratief (met behulp van een DDIM-sampler met een "trailing" timestep-strategie voor efficiëntie).
4. De finale schone latent wordt via de bevroren VAE-decoder teruggecodeerd naar pixelruimte om het segmentatiemasker te produceren.
5. Test-Tijd Ensembling: Meerdere inferentiedoorlopen met verschillende ruiszaden worden in de latente ruimte gemiddeld om ruimtelijke consistentie te verbeteren en artefacten te verminderen.

3. Belangrijkste Bijdragen

• Geünificeerd Framework: DiGSeg is het eerste framework dat een voorgeprente diffusion-ruggengraat succesvol hergebruikt als een generalist segmentatieleerder die gesloten-vocabulaire, open-vocabulaire en cross-domein taken aankan zonder taakspecifieke architecturale wijzigingen.
• Direct Latent Leren: In plaats van te vertrouwen op ruisachtige attention-maps, fine-tunt de methode de denoising U-Net om direct segmentatie-consistente latents te genereren, waardoor hoogwaardige maskers worden geproduceerd zonder zware post-processing.
• Multi-Scale Tekstconditionering: Introduceert een mechanisme om taalfeatures over meerdere denoising-schalen in te brengen, wat robuuste afstemming mogelijk maakt tussen willekeurige tekstprompts en visuele features.
• Nieuwe Ruisstrategie: Stelt een geannealde multi-scale ruisplanning voor die fijnmazige grenslering afweegt tegen globale semantische consistentie.

4. Experimentele Resultaten

De auteurs hebben DiGSeg geëvalueerd op standaard benchmarks en gespecialiseerde domeinen, waarbij State-of-the-Art (SOTA) prestaties werden aangetoond.

• Open-Vocabulary Segmentatie:
  • Presteerde beter dan 13 baselines (waaronder ODISE, OVSeg, CAT-Seg) op A-847, PC-459, A-150, PC-59 en Cityscapes.
  • Bereikte aanzienlijke winsten (bijv. +2,8 mIoU op PC-59, +2,2 op PC-459) ten opzichte van de op één na beste methoden.
• Standaard Semantische Segmentatie:
  • Bereikte SOTA op COCO (50,8 mIoU) en ADE20K (58,6 mIoU), en sloot sterke transformer-gebaseerde modellen zoals Mask2Former en OneFormer voorbij.
• Cross-Domein Transfer:
  • Toonde sterke zero-shot transfer naar medische (REFUGE-2), landbouw (PhenoBench) en remote sensing (DeepGlobe) datasets zonder architecturale modificatie.
  • Opmerking: De prestaties in medische domeinen waren lager dan in natuurlijke scènes, toegeschreven aan de domeinkloof in de pretraining-data van CLIP, maar waren nog steeds concurrerend.
• Efficiëntie & Gebruik van Data:
  • Het model is zeer data-efficiënt en bereikt bijna identieke prestaties met slechts 50% van de trainingsdata.
  • De inferentiesnelheid is verbeterd via DDIM-trailing (bereikt verzadiging in 1 stap), hoewel het door de aard van diffusion trager blijft dan feed-forward modellen.

5. Betekenis

Dit werk overbrugt fundamenteel de kloof tussen visuele generatie en visueel begrip.

• Paradigmaverschuiving: Het daagt de notie uit dat diffusion-modellen uitsluitend voor synthese zijn, en bewijst dat ze krachtige, generalistische discriminatieve leermiddelen kunnen zijn wanneer ze correct geconditioneerd zijn.
• Generalisatie: Het toont aan dat de rijke ruimtelijke en semantische priors die zijn gecodeerd in grootschalige generatieve modellen, overdraagbaar zijn naar dichte predictietaken over diverse domeinen (natuurlijk, medisch, landbouw).
• Eenvoud: Door complexe multi-stage pipelines of gespecialiseerde heads te vermijden, biedt DiGSeg een gestroomlijnde, geünificeerde interface voor segmentatie, wat wijst op een toekomst waarin een enkele generatieve ruggengraat een breed scala aan computer vision-taken kan afhandelen.

Concluderend stelt DiGSeg vast dat moderne diffusion-ruggengraten, wanneer ze worden gefine-tuned met een latent-geconditioneerd protocol en tekstafstemming, gespecialiseerde segmentatiemodellen kunnen overtreffen terwijl ze superieure flexibiliteit bieden in open-vocabulaire en cross-domein scenario's.