ClustViT: Clustering-based Token Merging for Semantic Segmentation

1. Probleemstelling

Vision Transformers (ViTs) hebben state-of-the-art prestaties behaald in computer vision, maar lijden onder kwadratische computationele complexiteit ten opzichte van de invoergrootte. Hoewel token-pruning- en token-merging-technieken (token skimming) de kosten voor beeldclassificatie succesvol hebben verlaagd, kampen ze met problemen bij semantische segmentatie (een taak voor dichte voorspelling).

• Het Conflict: Classificatie kan "irrelevante" tokens verwijderen, maar segmentatie vereist het behoud van ruimtelijke details voor elke pixel om compatibiliteit met standaard decoderhoofden te waarborgen.
• Het Gat: Bestaande methoden vertrouwen vaak op laag-niveau token-相似heid of gestructureerde, vaste-compressie (bijvoorbeeld alleen $2\times2$ patches samenvoegen), wat niet adapteert aan de semantische inhoud van het beeld of slecht generaliseert naar taken voor dichte voorspelling.

2. Methodologie: ClustViT

De auteurs stellen ClustViT voor, een architectuur die een trainbare Cluster-module en een Regenerator-module integreert in een standaard ViT-backbone. De kerninnovatie is het gebruik van semantische informatie afgeleid van ground-truth segmentatiemaskers om token-merging te sturen, in plaats van eenvoudige pixel-nabijheid.

A. Architectuur Overzicht

1. Standaard ViT Backbone: Het proces begint zoals een standaard ViT, waarbij het beeld wordt opgesplitst in patches en deze worden geprojecteerd naar tokens.
2. Cluster-module (Injectiepunt): Op een vooraf gedefinieerd laag ($ip$) binnen de encoder voorspelt een lichtgewicht MLP cluster-toewijzingen voor elke token.
   • Pseudo-Clusters: Tijdens het trainen gebruikt het model ground-truth maskers om "pseudo-clusters" te genereren. Als alle pixels in een patch tot dezelfde semantische klasse behoren, wordt de patch toegewezen aan dat klasse-ID; anders wordt deze gemarkeerd als gemengd (0).
   • Merging: Tokens die aan hetzelfde cluster zijn toegewezen, worden samengevoegd tot één representatieve token (via mean pooling). Tokens in complexe/gemengde gebieden blijven niet-geclusterd om details te behouden.
   • Resultaat: De sequentielengte wordt verkort, wat de computationele kosten van de daaropvolgende Transformer-lagen verlaagt.
3. Regenerator-module: Na de laatste encoder-laag moet de gecomprimeerde sequentie worden hersteld naar de oorspronkelijke resolutie voor de decoder.
   • De module gebruikt de opgeslagen cluster-toewijzingen om representatieve tokens terug te uitbreiden naar hun oorspronkelijke posities.
   • Een verfijnende MLP reconstrueert de fijnkorrelige kenmerken, zodat de output-sequentie overeenkomt met de invoerlengte en compatibel is met kant-en-klare segmentatiehoofden (bijvoorbeeld Segmenter, UPerNet).

B. Trainingsstrategie

• Dual-Task Loss: Het model wordt end-to-end getraind met een gecombineerde loss-functie:
  $$L = L_{segm} + \lambda L_{clust}$$
  • $L_{segm}$: Standaard cross-entropy loss voor semantische segmentatie.
  • $L_{clust}$: Cross-entropy loss die de cluster-module dwingt de pseudo-clusters te voorspellen die zijn afgeleid van het ground-truth masker.
  • $\lambda$: Een balancerende hyperparameter ( ingesteld op 0,1 in experimenten).

3. Belangrijkste Bijdragen

1. Semantisch-Gestuurde Token Merging: In tegenstelling tot eerdere methoden die mergen op basis van visuele similariteit of vaste roosters, voegt ClustViT tokens samen op basis van semantische inhoud (afgeleid van segmentatiemaskers), wat ongestructureerde, inhouds-adaptieve compressie mogelijk maakt.
2. End-to-End Trainbare Backbone: Het clustering-component is direct geïntegreerd in de ViT-backbone en wordt gezamenlijk getraind met de segmentatietask, waardoor de noodzaak voor aparte beleidsnetwerken of voor-trainingsstappen wordt geëlimineerd.
3. Regenerator voor Dichte Voorspelling: De introductie van een Regenerator-module lost het compatibiliteitsprobleem op, waardoor gecomprimeerde token-representaties volledig kunnen worden hersteld voor dichte voorspellinghoofden zonder verlies van ruimtelijke resolutie.
4. Efficiëntie-Nauwkeurigheids Trade-off: De methode vermindert de computationele kosten (GFLOPs) en de inferentietijd aanzienlijk terwijl concurrerende nauwkeurigheid wordt behouden, met name in scenario's met grote achtergrondgebieden die gebruikelijk zijn in robotica.

4. Experimentele Resultaten

De auteurs evalueerden ClustViT op drie datasets: ADE20K (algemene scènes), SUIM (onderwaterrobotica) en RumexWeeds (landbouwrobotica). Ze vergeleken met een standaard ViT en de state-of-the-art Content-aware Token Sharing (CTS) methode.

• Prestatiemetrics:
  • Computationele Efficiëntie: Bereikte tot 2,18× minder GFLOPs en 1,64× snellere inferentie (afbeeldingen per seconde) in vergelijking met de baseline ViT.
  • Nauwkeurigheid: Behield vergelijkbare segmentatienauwkeurigheid (mIoU). Op ADE20K bereikte de beste configuratie (ClustViT-bk3,ip3) een 1,47× snelheidswinst met slechts een daling van ongeveer 1,66 mIoU-punten.
  • Robotica-scenario's: De methode excelleerde op datasets met grote achtergrondgebieden (SUIM en RumexWeeds). Bijvoorbeeld, op RumexWeeds behaalde het 1,54× FPS en 2,12× GFLOP winsten terwijl het de mIoU licht verbeterde ten opzichte van de baseline.
• Ablatiestudies:
  • Aantal Clusters ($k$): Het verhogen van $k$ (meer clusters) verbetert de snelheid maar vermindert de nauwkeurigheid omdat meer details worden samengevoegd. $k=3$ bood de beste balans.
  • Injectiepunt ($ip$): Het comprimeren eerder in het netwerk ($ip=2$) maximaliseerde de snelheid maar schaade de nauwkeurigheid. Een middelijk injectiepunt ($ip=3$ of $4$) bood de optimale trade-off.
  • Modelgrootte: De efficiëntiewinsten waren het meest uitgesproken bij grotere modellen (Base/Large), waarbij de overhead van de clustering-MLP verwaarloosbaar was ten opzichte van de besparingen in attention-berekening.

5. Betekenis en Impact

• Robotica-toepasbaarheid: Het artikel adresseert een kritieke bottleneck voor het implementeren van ViTs op echte robotische systemen (bijvoorbeeld autonome landbouw, onderwaterverkenning) waar computationele middelen beperkt zijn en scènes vaak grote, uniforme achtergronden bevatten.
• Paradigmaverschuiving: Het toont aan dat token-compressie in dichte voorspelling geen inboet aan ruimtelijke fideliteit vereist als de merging wordt gestuurd door semantisch begrip in plaats van alleen visuele similariteit.
• Generaliseerbaarheid: Door standaard decoderhoofden (Segmenter en UPerNet) te ondersteunen, biedt ClustViT een drop-in vervanging voor bestaande ViT-backbones, wat het zeer praktisch maakt voor de bredere computer vision-gemeenschap.

Concluderend sluit ClustViT succesvol de kloof tussen de hoge nauwkeurigheid van Vision Transformers en de efficiëntie-eisen van real-time robotische perceptie door een semantisch bewust, trainbaar token-merging-mechanisme in te voeren.