ClustViT: Clustering-based Token Merging for Semantic Segmentation

1. Enunciato del Problema

I Vision Transformers (ViT) hanno raggiunto prestazioni all'avanguardia nella visione artificiale, ma soffrono di complessità computazionale quadratica rispetto alla dimensione dell'input. Sebbene le tecniche di pruning e merging dei token (token skimming) abbiano ridotto con successo i costi per la classificazione di immagini, faticano nella segmentazione semantica (un compito di previsione densa).

• Il Conflitto: La classificazione può scartare token "irrilevanti", ma la segmentazione richiede la preservazione dei dettagli spaziali per ogni pixel per garantire la compatibilità con le testate di decodifica standard.
• Il Divario: I metodi esistenti spesso si basano sulla similarità a basso livello dei token o su una compressione strutturata a dimensione fissa (ad esempio, unendo solo patch $2\times2$), il che fallisce nell'adattarsi al contenuto semantico dell'immagine o generalizza male ai compiti di previsione densa.

2. Metodologia: ClustViT

Gli autori propongono ClustViT, un'architettura che integra un modulo Cluster addestrabile e un modulo Regenerator in un backbone ViT standard. L'innovazione principale consiste nell'utilizzare informazioni semantiche derivate dalle maschere di segmentazione ground-truth per guidare il merging dei token, piuttosto che la semplice prossimità dei pixel.

A. Panoramica dell'Architettura

1. Backbone ViT Standard: Il processo inizia come un ViT standard, dividendo l'immagine in patch e proiettandole in token.
2. Modulo Cluster (Punto di Iniezione): A un livello predefinito ($ip$) all'interno dell'encoder, un MLP leggero prevede le assegnazioni di cluster per ogni token.
   • Pseudo-Cluster: Durante l'addestramento, il modello utilizza le maschere ground-truth per generare "pseudo-cluster". Se tutti i pixel di una patch appartengono alla stessa classe semantica, la patch viene assegnata a quell'ID di classe; altrimenti, viene contrassegnata come mista (0).
   • Merging: I token assegnati allo stesso cluster vengono aggregati in un singolo token rappresentativo (tramite media pooling). I token in regioni complesse/miste rimangono non clusterizzati per preservare i dettagli.
   • Risultato: La lunghezza della sequenza viene ridotta, abbassando il costo computazionale dei successivi livelli Transformer.
3. Modulo Regenerator: Dopo il livello finale dell'encoder, la sequenza compressa deve essere ripristinata alla risoluzione originale per il decodificatore.
   • Il modulo utilizza le assegnazioni di cluster memorizzate per espandere i token rappresentativi nelle loro posizioni originali.
   • Un MLP di raffinamento ricostruisce le caratteristiche fini, garantendo che la sequenza di output corrisponda alla lunghezza di input e sia compatibile con testate di segmentazione pronte all'uso (ad esempio, Segmenter, UPerNet).

B. Strategia di Addestramento

• Funzione di Perdita a Doppio Compito: Il modello è addestrato end-to-end con una funzione di perdita combinata:
  $$L = L_{segm} + \lambda L_{clust}$$
  • $L_{segm}$: Perdita standard cross-entropy per la segmentazione semantica.
  • $L_{clust}$: Perdita cross-entropy che forza il modulo cluster a prevedere i pseudo-cluster derivati dalla maschera ground-truth.
  • $\lambda$: Un iperparametro di bilanciamento (impostato a 0,1 negli esperimenti).

3. Contributi Chiave

1. Merging dei Token Guidato dalla Semantica: A differenza dei metodi precedenti che uniscono basandosi sulla similarità visiva o su griglie fisse, ClustViT unisce i token basandosi sul contenuto semantico (derivato dalle maschere di segmentazione), permettendo una compressione non strutturata e adattiva al contenuto.
2. Backbone Addestrabile End-to-End: Il componente di clustering è integrato direttamente nel backbone ViT e addestrato congiuntamente al compito di segmentazione, eliminando la necessità di reti di policy separate o passaggi di pre-addestramento.
3. Regenerator per Previsione Densa: L'introduzione di un modulo Regenerator risolve il problema di compatibilità, permettendo alle rappresentazioni di token compresse di essere completamente ripristinate per le testate di previsione densa senza perdere la risoluzione spaziale.
4. Compromesso Efficienza-Accuratezza: Il metodo riduce significativamente il costo computazionale (GFLOPs) e il tempo di inferenza mantenendo un'accuratezza competitiva, in particolare in scenari con grandi aree di sfondo comuni nella robotica.

4. Risultati Sperimentali

Gli autori hanno valutato ClustViT su tre dataset: ADE20K (scene generali), SUIM (robotica subacquea) e RumexWeeds (robotica agricola). Hanno confrontato i risultati con ViT standard e il metodo all'avanguardia Content-aware Token Sharing (CTS).

• Metriche di Prestazione:
  • Efficienza Computazionale: Ha raggiunto fino a 2,18× meno GFLOPs e 1,64× inferenza più veloce (immagini al secondo) rispetto al ViT di base.
  • Accuratezza: Ha mantenuto un'accuratezza di segmentazione comparabile (mIoU). Su ADE20K, la configurazione migliore (ClustViT-bk3,ip3) ha ottenuto un speedup di 1,47× con una diminuzione di soli ~1,66 punti mIoU.
  • Scenari Robotici: Il metodo ha eccellito su dataset con grandi aree di sfondo (SUIM e RumexWeeds). Ad esempio, su RumexWeeds, ha ottenuto guadagni di 1,54× FPS e 2,12× GFLOP migliorando leggermente il mIoU rispetto alla baseline.
• Studi di Ablazione:
  • Numero di Cluster ($k$): Aumentare $k$ (più cluster) migliora la velocità ma riduce l'accuratezza man mano che vengono uniti più dettagli. $k=3$ ha offerto il miglior equilibrio.
  • Punto di Iniezione ($ip$): Comprimere più presto nella rete ($ip=2$) massimizza la velocità ma danneggia l'accuratezza. Un punto di iniezione intermedio ($ip=3$o$4$) ha fornito il compromesso ottimale.
  • Scala del Modello: I guadagni di efficienza sono stati più pronunciati nei modelli più grandi (Base/Large), dove l'overhead del MLP di clustering era trascurabile rispetto ai risparmi nel calcolo dell'attenzione.

5. Significato e Impatto

• Applicabilità Robotica: Il paper affronta un collo di bottiglia critico per il dispiegamento di ViT su sistemi robotici reali (ad esempio, agricoltura autonoma, esplorazione sottomarina) dove le risorse computazionali sono limitate e le scene spesso contengono grandi sfondi uniformi.
• Cambiamento di Paradigma: Dimostra che la compressione dei token nella previsione densa non richiede di sacrificare la fedeltà spaziale se il merging è guidato dalla comprensione semantica piuttosto che dalla semplice similarità visiva.
• Generalizzabilità: Supportando testate di decodifica standard (Segmenter e UPerNet), ClustViT offre una sostituzione diretta per i backbone ViT esistenti, rendendolo altamente pratico per la più ampia comunità di visione artificiale.

In conclusione, ClustViT colma con successo il divario tra l'alta accuratezza dei Vision Transformers e i requisiti di efficienza della percezione robotica in tempo reale, introducendo un meccanismo di merging dei token consapevole della semantica e addestrabile.