Una guida completa a SegFormer: perché un trasformatore senza codifica posizionale funziona per la segmentazione

Prima di SegFormer, erano già stati fatti tentativi di introdurre i Transformer nella segmentazione semantica, ma molti si basavano su codifiche posizionali e decoder complessi. Di conseguenza, si trovavano ad affrontare due problemi: la precisione poteva diminuire quando la risoluzione di test differiva da quella di training, e i costi computazionali potevano diventare elevati. SegFormer (Xie, Wang, Yu, Anandkumar, Alvarez e Luo, NeurIPS 2021) risolve entrambi i problemi con un design straordinariamente semplice: un codificatore gerarchico senza codifica posizionale e un decoder All-MLP che non utilizza alcuna convoluzione. Questo articolo esamina, passo dopo passo, il motivo per cui tale design funziona, basandosi sul documento originale (arXiv:2105.15203).

Questo articolo è basato sul paper originale "SegFormer: Simple and Efficient Design for Semantic Segmentation with Transformers" (Xie et al., NeurIPS 2021, arXiv:2105.15203).

0. Cosa imparerai

• I problemi che SegFormer si propone di risolvere nei precedenti modelli di segmentazione basati su Transformer
• Come funziona l'encoder gerarchico MiT (Mix Transformer) senza codifica posizionale
• Perché un decoder composto solo da MLP, senza alcuna convoluzione, può comunque raggiungere un'elevata precisione
• mIoU misurato, numero di parametri e costo computazionale per B0–B5 su Cityscapes e ADE20K
• Come scegliere un modello nella pratica, inclusa la robustezza alla corruzione dell'immagine

1. Cos'è SegFormer?

SegFormer è un modello di segmentazione semantica che combina un encoder Transformer gerarchico (MiT) senza codifica posizionale con un decoder All-MLP leggero che non utilizza alcuna convoluzione. L'idea centrale è quella di riprogettare la divisione del lavoro tra encoder e decoder: una volta che l'encoder ha prodotto caratteristiche a diverse risoluzioni, il decoder può essere notevolmente semplificato.

2. Perché era necessario questo design?

L'utilizzo diretto di un Vision Transformer (ViT) come encoder per la segmentazione semantica crea due problemi. In primo luogo, le caratteristiche prodotte dal ViT hanno una sola risoluzione. Gli oggetti possono avere dimensioni diverse all'interno di una singola immagine e i modelli di segmentazione convenzionali basati su CNN, come FCN e U-Net, affrontano questo problema combinando caratteristiche a diverse risoluzioni. Un ViT che produce una sola risoluzione perde questa capacità di rappresentazione multiscala.

Il secondo problema è la dipendenza dalla codifica posizionale. ViT divide un'immagine in patch e le fornisce a un Transformer, ma il Transformer stesso non ha informazioni sulla posizione di ciascuna patch nell'immagine. Pertanto, è necessario aggiungere separatamente una codifica posizionale fissa o appresa. Poiché questa codifica è creata per la risoluzione di input utilizzata durante l'addestramento, un'immagine con una risoluzione diversa in fase di inferenza richiede l'interpolazione della codifica posizionale, il che può ridurre la precisione. La segmentazione spesso richiede input a risoluzione più elevata rispetto al rilevamento, rendendo questa dipendenza dalla risoluzione particolarmente problematica.

SegFormer affronta entrambi i problemi alla radice utilizzando un codificatore gerarchico ed eliminando completamente la codifica posizionale.

3. Qual è l'input?

Come per altri modelli di segmentazione, l'input è un'immagine RGB x \in \mathbb{R}^{H \times W \times 3}. SegFormer viene valutato su immagini ad alta risoluzione (1024×2048) per Cityscapes e su immagini di circa 512×512 per ADE20K. Mentre ViT effettua il downsampling con patch grossolane di 16×16 pixel, l'encoder di SegFormer inizia con patch di dimensioni ridotte, fino a 4×4 pixel, preservando maggiori dettagli all'inizio dell'estrazione delle caratteristiche.

4. Cosa predice?

L'output è una mappa di segmentazione a 1/4 della risoluzione di input, con la classe di ciascun pixel rappresentata come \hat{y} \in \mathbb{R}^{\frac{H}{4} \times \frac{W}{4} \times N_{cls}} (N_{cls} è il numero di classi). Viene quindi riportata alla risoluzione originale utilizzando l'interpolazione del vicino più prossimo o un metodo simile. L'idea centrale è che, se l'encoder crea caratteristiche multiscala sufficientemente ricche, il decoder necessita solo di una combinazione di semplici layer lineari per integrarle e convertirle in una maschera accurata.

5. Architettura di base

SegFormer è composto da due blocchi principali: un encoder MiT (Mix Transformer) che estrae le caratteristiche a quattro risoluzioni e un decoder All-MLP che combina queste caratteristiche e produce una maschera di segmentazione.

Figura 1 — L'encoder MiT produce caratteristiche a quattro risoluzioni (1/4, 1/8, 1/16 e 1/32) e il decoder All-MLP le combina utilizzando solo layer lineari per creare una maschera di segmentazione. L'unica convoluzione è presente nel Mix-FFN dell'encoder; Il decodificatore è letteralmente composto interamente da strati lineari.

Poiché l'encoder crea le caratteristiche a quattro risoluzioni in anticipo, il decoder può ottenere una maschera accurata con un'elaborazione pressoché minima: allinea le caratteristiche e le combina. Questa divisione del lavoro è il motivo principale per cui SegFormer mantiene basso il numero totale di parametri.

6. Dettagli tecnici dei componenti

Unione di patch sovrapposte: una gerarchia a quattro fasi

Anziché dividere l'immagine in patch una sola volta come fa ViT, l'encoder MiT riduce gradualmente la risoluzione in quattro fasi. Le risoluzioni di output sono 1/4, 1/8, 1/16 e 1/32 dell'input, riproducendo la struttura gerarchica utilizzata dalle reti neurali convoluzionali (CNN) come ResNet.

Anche il metodo di partizionamento delle patch differisce dalle patch non sovrapposte di ViT: le patch adiacenti si sovrappongono quando vengono unite. La prima fase utilizza una dimensione del kernel di K=7, uno stride di S=4 e un padding di P=3; le fasi successive utilizzano K=3, S=2, P=1. La sovrapposizione preserva la continuità locale attraverso i confini delle patch, ovvero le informazioni condivise dalle patch adiacenti, mentre la mappa delle caratteristiche viene ridotta.

Autoattenzione efficiente

L'autoattenzione multi-testa ordinaria richiede O(N^2) calcoli per una sequenza di lunghezza N. Poiché N diventa molto grande per la segmentazione ad alta risoluzione, questo rappresenta un serio collo di bottiglia. SegFormer introduce la riduzione della sequenza, che comprime le sequenze di chiave e valore con un rapporto di riduzione di R prima di calcolare l'attenzione.

Le sequenze K, V vengono ridotte dalla lunghezza N alla lunghezza N/R prima di essere utilizzate. Il fattore di riduzione R per le fasi 1-4 è impostato a [64, 16, 4, 1]: le fasi superficiali hanno alta risoluzione e sequenze lunghe, quindi vengono compresse in modo aggressivo; le fasi profonde hanno risoluzione inferiore e sequenze più corte, quindi non vengono compresse affatto. Ciò riduce il costo computazionale da O(N^2) a O(N^2/R) e rende l'autoattenzione praticabile anche per input ad alta risoluzione.

Mix-FFN: come viene eliminata la codifica posizionale

SegFormer può evitare completamente la codifica posizionale perché integra una convoluzione 3×3 nella rete feed-forward (FFN), creando Mix-FFN.

La convoluzione 3×3 utilizza un padding di zeri al bordo dell'immagine. Questo padding fornisce indirettamente un indizio sulla "posizione nell'immagine" all'interno della convoluzione, in modo che le informazioni posizionali vengano incorporate nelle caratteristiche semplicemente passando attraverso Mix-FFN, senza una codifica posizionale esplicita. Di conseguenza, il test a una risoluzione diversa da quella di addestramento non richiede l'interpolazione di una codifica posizionale, che altrimenti potrebbe ridurre la precisione.

Decodificatore All-MLP: perché la convoluzione è superflua

Il decodificatore è costituito dai seguenti quattro passaggi, utilizzando letteralmente solo layer lineari (MLP).

1. Unificazione dei canali: mappare ogni caratteristica di fase F_i allo stesso numero di canali C con un layer lineare indipendente.

2. Upsampling e concatenazione: eseguire l'upsampling di ogni caratteristica a 1/4 della risoluzione e concatenarla lungo la dimensione dei canali.

3. Fusione: mappare la caratteristica concatenata a 4C dimensioni a C con un layer lineare.

4. Previsione: generare una maschera di segmentazione con un canale per ogni classe utilizzando un layer lineare finale.

L'articolo originale supporta questo progetto con un'analisi del campo recettivo effettivo (ERF). I decoder basati su CNN come DeepLabv3+ necessitano di meccanismi complessi, come convoluzioni dilatate e pooling multiscala, per ottenere un ampio campo recettivo e raggiungere un'elevata precisione. Al contrario, l'encoder MiT acquisisce naturalmente ERF locali nelle fasi superficiali e ERF non locali che coprono l'intera immagine nelle fasi profonde attraverso l'autoattenzione, senza meccanismi aggiuntivi. Poiché l'encoder stesso crea già caratteristiche contenenti informazioni da scale locali a globali, il decoder non necessita di un complesso meccanismo di espansione del campo recettivo. Questa è la ragione teorica per cui una semplice combinazione di layer lineari è sufficiente.

Metodo di addestramento

L'articolo originale utilizza una semplice configurazione di addestramento senza particolari accorgimenti. L'encoder MiT è pre-addestrato su ImageNet-1K, mentre il decoder è addestrato a partire da un'inizializzazione casuale. Per l'ottimizzazione viene utilizzato AdamW, con un tasso di apprendimento che diminuisce da un valore iniziale di 0,00006 secondo uno schema polinomiale (decadimento a legge di potenza con coefficiente 1,0). Il numero di iterazioni è di 160.000 su ADE20K e Cityscapes e di 80.000 su COCO-Stuff. L'aumento dei dati di training è limitato al ridimensionamento casuale con un rapporto di 0,5-2,0, al ribaltamento orizzontale casuale e al ritaglio casuale (512×512 per ADE20K, 1024×1024 per Cityscapes e 640×640 per B5 su ADE20K). L'articolo afferma esplicitamente di non utilizzare tecniche ampiamente adottate come OHEM (Online Hard Example Mining), perdite ausiliarie o perdite bilanciate per classe. L'obiettivo è dimostrare che l'encoder MiT e il decoder All-MLP possono raggiungere un elevato mIoU senza tali aggiunte.

7. Confronto tra le varianti del modello

SegFormer è disponibile in sei dimensioni, da MiT-B0 a MiT-B5. Di seguito sono riportati i risultati misurati per Cityscapes e ADE20K, descritti nella Tabella 2 dell'articolo originale.

Fonte: Tabella 2 dell'articolo originale (Xie et al., NeurIPS 2021). Il mIoU di Cityscapes è misurato con test multi-scala e di inversione sinistra-destra (MS). Il conteggio dei parametri include l'encoder e il decoder completi; anche per B5, il decoder rappresenta solo circa il 4% del totale. A titolo di riferimento, B5 raggiunge 82,4 mIoU con test a scala singola (SS) e 84,0 con test MS.

L'articolo confronta questi risultati con lo stato dell'arte dell'epoca e riporta che SegFormer-B4 raggiunge il 50,3% di mIoU su ADE20K con 64,1 milioni di parametri, superando il miglior metodo di quel periodo di 2,2 punti con un modello di dimensioni quasi un quinto inferiori. B0 mantiene un mIoU pratico su Cityscapes del 76,2% con soli 3,8 milioni di parametri, e i numeri supportano la sua progettazione orientata ai bordi.

8. Le sue difficoltà

La maggior parte delle debolezze di SegFormer sono comuni ai modelli basati su Transformer in generale. In primo luogo, i modelli più grandi (B3 e superiori) hanno costi di addestramento e inferenza elevati. I FLOP di Cityscapes aumentano da 125,5 GB per B0 a 1460,4 GB per B5, un aumento di oltre dieci volte. Per la segmentazione ad alta risoluzione, questa crescita può rendere difficile il funzionamento in tempo reale.

Anche la dipendenza da dati di pre-addestramento su larga scala è rilevante. L'encoder MiT è pre-addestrato su ImageNet-1K, quindi la qualità e la quantità dei dati di fine-tuning influenzano fortemente il trasferimento a domini che differiscono sostanzialmente dai dati di pre-addestramento, come le immagini mediche o satellitari.

Anche oggetti molto piccoli e strutture sottili, ad esempio cavi e cartelli stradali, non sono intrinsecamente facili da gestire per SegFormer. La fusione di patch sovrapposte (Overlapped Patch Merging) effettua un downsampling di 4×4 anche nella prima fase, quindi strutture larghe solo pochi pixel possono perdere informazioni prima di raggiungere le fasi successive.

Allo stesso tempo, la valutazione della robustezza su Cityscapes-C riportata nell'articolo originale rappresenta un chiaro punto di forza. Per immagini corrotte con rumore gaussiano aggiunto, SegFormer-B5 registra un miglioramento relativo massimo dell'mIoU del 588% rispetto a DeepLabv3+ con un backbone Xception-71; per la corruzione simile alla neve, il miglioramento relativo massimo è del 295%. Queste misurazioni mostrano una maggiore tolleranza alla corruzione delle immagini rispetto ai modelli basati su CNN tradizionali. Si ritiene che questa robustezza derivi dal fatto che l'autoattenzione è meno facilmente dominata dal rumore locale.

9. Come scegliere un modello nella pratica

Per dispositivi edge ed elaborazione in tempo reale, SegFormer-B0 o B1 rappresentano una scelta realistica. B0 bilancia precisione e compattezza con 3,8 milioni di parametri e 76,2 mIoU su Cityscapes, rendendo plausibile l'implementazione su dispositivi embedded e droni.

Per applicazioni come il riconoscimento dell'area percorribile nella guida autonoma, dove sono richieste immagini ad alta risoluzione e precisione costante, B2 o B3 sono spesso un compromesso ragionevole. B4 e B5 offrono la massima precisione, ma i loro FLOP su Cityscapes superano i 1200G, quindi l'inferenza in tempo reale su una GPU di bordo richiede generalmente ottimizzazioni come la quantizzazione o la conversione TensorRT.

Per l'analisi di precisione offline di immagini mediche o satellitari, la precisione ha la priorità sulle prestazioni in tempo reale. B4 o B5 sono appropriati, con una messa a punto su dati specifici del dominio, se necessario.

Per ambienti esterni con frequenti condizioni meteorologiche avverse o variazioni di illuminazione, come robot agricoli e telecamere di sorveglianza per esterni, la robustezza misurata su Cityscapes-C rappresenta un vantaggio pratico. Rispetto ai modelli basati su CNN con una precisione simile, ci si può aspettare che SegFormer perda meno precisione in presenza di immagini corrotte.

Per apprendere il rilevamento di oggetti e la segmentazione semantica dalle basi, consultare "Introduzione al rilevamento di oggetti e alla segmentazione semantica". Per le tendenze recenti nel campo della segmentazione, consultare "Tendenze tecnologiche nella segmentazione semantica".

10. Riepilogo in tre righe

• SegFormer combina un codificatore gerarchico (MiT) che produce caratteristiche multiscala senza codifica posizionale con un decodificatore All-MLP che non utilizza convoluzioni.

• La convoluzione 3×3 in Mix-FFN rilascia implicitamente informazioni posizionali, eliminando la necessità di codifica posizionale, mentre l'ampio campo recettivo effettivo del Transformer rende possibile un semplice decodificatore MLP.

• La famiglia spazia da B0 (3,8 milioni di parametri, 76,2 mIoU su Cityscapes) a B5 (84,7 milioni di parametri, 84,0 mIoU) e i risultati misurati ne confermano la robustezza alla corruzione delle immagini.

Riferimenti

• SegFormer: Progettazione semplice ed efficiente per la segmentazione semantica con Transformer (NeurIPS 2021, arXiv:2105.15203)
• Articolo su SegFormer su ar5iv (testo completo)
• Implementazione ufficiale di SegFormer (NVlabs/SegFormer, GitHub)
• Atti di NeurIPS 2021: SegFormer