Riepilogo Tecnico: Vanilla ViT per la Segmentazione Semantica di Nuvole di Punti Automotive (VaViT)

Problema

Sebbene i plain Transformer siano diventati l'architettura de-facto per l'elaborazione di testo, audio e immagini grazie al loro design unificato e alle superiori proprietà di scalabilità, la loro applicazione alla segmentazione semantica di nuvole di punti 3D rimane limitata. Gli attuali metodi allo stato dell'arte per la segmentazione LiDAR automotive si affidano prevalentemente ad architetture ibride U-Net che intercalano convoluzioni sparse con attenzione locale o a finestre (ad es., PTv3, LitePT). Alternativamente, alcuni approcci proiettano i dati 3D in viste di range 2D (ad es., RangeViT), il che introduce distorsioni geometriche e scarta la struttura spaziale.

Gli autori pongono una domanda fondamentale: Un Vision Transformer (ViT) semplice e non gerarchico può operare direttamente su nuvole di punti 3D grezze per ottenere prestazioni competitive per la segmentazione semantica LiDAR su larga scala? La sfida consiste nel colmare il divario tra i dati irregolari delle nuvole di punti e la rappresentazione basata su token richiesta dai ViT senza fare affidamento sugli bias induttivi delle convoluzioni o sulle distorsioni geometriche delle proiezioni 2D.

Metodologia

Il framework proposto, VaViT (Vanilla ViT), elabora punti 3D grezzi utilizzando un backbone ViT a scala singola con auto-attenzione globale, evitando stadi gerarchici o blocchi convoluzionali. La pipeline consiste di quattro componenti chiave:

1. Strategia di Tokenizzazione

Invece di convertire l'intero scan in un'immagine 2D, VaViT estrae embedding punto-specifici direttamente nello spazio 3D e li proietta su una griglia Bird's Eye View (BEV) grossolana per formare i token.

• Rete di Embedding: I punti grezzi (coordinate, range, intensità) passano attraverso una rete di embedding multi-strato ispirata a PointNet++ e DGCNN. Questa rete include:
  • Un ramo globale che elabora le feature dei punti.
  • Un ramo locale che calcola le rappresentazioni relative con $k$-vicini più prossimi, potenziate da un meccanismo di gating ($f_{gate}$) per filtrare le feature spurie.
  • Una concatenazione di feature globali e locali.
• Tokenizzazione a Pillar: La scena è divisa in pillar quadrati non sovrapposti. Per ogni pillar occupato, le embedding di tutti i punti contenuti vengono raccolte e un'operazione di max-pooling genera un singolo token del pillar.
• Codifica Posizionale: I token sono associati a coordinate BEV 2D, codificate utilizzando Rotary Positional Embeddings (RoPE) ripetuti su ogni asse.

2. Backbone Vanilla ViT

La sequenza di token dei pillar viene elaborata da un ViT standard, non gerarchico.

• L'architettura segue il design standard dei blocchi Transformer: LayerNorm (LN) $\rightarrow$ Multi-Head Self-Attention (MSA) $\rightarrow$ Residuo $\rightarrow$ LN $\rightarrow$ MLP $\rightarrow$ Residuo.
• Per gestire lunghezze di sequenza variabili tra i batch, vengono utilizzati token di padding per equalizzare le lunghezze.
• Il backbone si affida interamente all'auto-attenzione globale, consentendo al modello di catturare dipendenze spaziali a lungo raggio in tutta la scena senza vincoli di finestre locali.

3. Testa di Segmentazione Leggera

Per recuperare predizioni a livello di punto fini dai grossolani token dei pillar:

• Lifting: I token finali del transformer vengono sollevati dalla griglia BEV nuovamente ai punti 3D originali. Tutti i punti all'interno di un determinato pillar ricevono il corrispondente token del pillar.
• Fusione delle Feature: Le feature dei token sollevati vengono concatenate con le embedding dei punti ad alta risoluzione originali calcolate durante la tokenizzazione.
• Gating e Classificazione: Un meccanismo di gating ($g_{gate}$) ricalibra il contributo delle feature del ViT rispetto alle embedding dei punti originali. Le feature combinate passano attraverso un livello lineare finale per la classificazione semantica.

4. Data Augmentation: PillarMix+

Riconoscendo che i dataset LiDAR sono più piccoli dei dataset di immagini utilizzati per il pretraining dei ViT, gli autori introducono PillarMix+, una strategia di aumento geometrico potenziata.

• Meccanismo: A differenza del PillarMix standard che mescola due scan in un pattern a scacchiera, PillarMix+ mescola $N$ nuvole di punti (tipicamente 3).
• Processo: Per ogni pillar occupato in uno scan di riferimento, se la stessa posizione è occupata in uno degli altri $N-1$ scan, il pillar viene casualmente sostituito da uno dei pillar occupati appartenenti al set.
• Vantaggi: Questo preserva il pattern di occupazione e la forma generale della scena pur introducendo variazioni di contesto significative. Evita l'effetto "ritaglio stand-alone" del mixing a scacchiera e mantiene un numero costante di token, diversamente dai metodi che aumentano il numero di token.

Contributi Chiave

1. Validazione delle Prestazioni: Ha dimostrato che un Vanilla ViT accuratamente progettato può eguagliare o superare i metodi allo stato dell'arte (inclusi gli ibridi U-Net e i modelli di proiezione 2D) per la segmentazione LiDAR.
2. Tokenizzazione Specializzata: Ha introdotto uno stem di tokenizzazione dedicato che estrae informative embedding di punti 3D e le proietta su una griglia BEV grossolana, preservando informazioni geometriche sufficienti per i task downstream senza modifiche architettoniche al ViT.
3. Testa Leggera: Ha progettato una testa di segmentazione che combina efficacemente le feature contestuali globali del ViT con le embedding dei punti ad alta risoluzione tramite un meccanismo di gating.
4. PillarMix+: Ha proposto una robusta strategia di aumento geometrico che migliora significativamente la generalizzazione per l'addestramento dei ViT sui dataset di guida autonoma, mescolando più scan pur preservando la coerenza della scena.

Risultati Sperimentali

Gli autori hanno valutato VaViT su tre benchmark principali: nuScenes, SemanticKITTI e Waymo Open Dataset (WOD).

• nuScenes: VaViT-B ha raggiunto un mIoU (mean Intersection over Union) dell'81.3%, superando il precedente stato dell'arte LitePT (80.7%). Notevolmente, si è classificato al primo posto su oggetti piccoli come coni stradali e biciclette nonostante la griglia di tokenizzazione grossolana.
• Waymo Open Dataset: VaViT-B ha raggiunto il 70.5% di mIoU, superando SphereFormer (69.9%) e PTv3 (70.2%).
• SemanticKITTI: VaViT-B ha raggiunto il 67.6% di mIoU (senza Test-Time Augmentation), competendo da vicino con PTv3 (66.2%) e SphereFormer (67.8%).
• Robustezza: Un risultato significativo è che i modelli VaViT sono intrinsecamente robusti alle trasformazioni globali (rotazione, scala, flipping). Mentre altri metodi (PTv3, LitePT) guadagnano circa 2 punti mIoU con il Test-Time Augmentation (TTA), le prestazioni di VaViT rimangono quasi invariate, suggerendo che il meccanismo di auto-attenzione globale ha già appreso queste invarianze.

Significato e Rivendicazioni

Il paper sostiene che i design ibridi specializzati (che combinano convoluzioni e attenzione locale) potrebbero non essere strettamente necessari per la segmentazione semantica 3D. Invere, i plain Transformer costituiscono un'alternativa valida e più semplice che offre un backbone unificato attraverso i task di percezione.

Il successo di VaViT suggerisce che il "lavoro pesante" nella segmentazione di nuvole di punti può essere svolto dall'auto-attenzione globale se la rappresentazione dell'input (tokenizzazione) e la strategia di addestramento (augmentation) sono accuratamente calibrate. Questo lavoro apre la porta a:

• Design di architetture multimodali più semplici (ad es., mescolare nuvole di punti e immagini).
• L'applicazione diretta delle tecniche di pretraining auto-supervisionato per i transformer nel dominio 3D.
• Una riduzione della complessità architettonica eliminando la necessità di gerarchie stile U-Net e convoluzioni sparse.

Gli autori sottolineano che i loro risultati validano l'efficienza di un approccio ViT unificato, a condizione che l'interfaccia tra i dati 3D grezzi e il transformer sia meticolosamente ingegnerizzata.