Accelerating Vision Transformers on Brain Processing Unit

Sintesi Tecnica: Accelerazione dei Vision Transformer su Brain Processing Unit

Problematica
Il documento affronta un critico disallineamento architettonico tra l'hardware specializzato per l'elaborazione neurale e i moderni modelli Vision Transformer (ViT). Nello specifico, l'Horizon Brain Processing Unit (BPU) è ottimizzato per operazioni di convoluzione quadridimensionale (quantizzate INT8) tipiche delle Reti Neurali Convoluzionali (CNN). Tuttavia, i modelli ViT, come il Data-efficient Image Transformer (DeiT), si affidano pesantemente a strati lineari e alla Layer Normalization (LayerNorm) che operano su flussi di dati tridimensionali. Questa discrepanza impedisce alla BPU di accelerare le computazioni centrali dei ViT, costringendo tali modelli a girare su CPU di uso generale e facendo perdere le capacità di accelerazione dell'hardware.

Metodologia
Per risolvere questa incompatibilità senza richiedere il riaddestramento del modello, gli autori propongono una riformulazione strutturale dell'architettura DeiT. La metodologia principale consiste nel sostituire gli operatori non convoluzionali con operatori convoluzionali matematicamente equivalenti che la BPU può accelerare:

1. Trasformazione dello Strato Lineare: Gli strati lineari, che tipicamente elaborano dati 3D $(B, N, C)$, vengono trasformati in convoluzioni point-wise. Ciò viene ottenuto riconfigurando la matrice dei pesi $W \in \mathbb{R}^{O \times I}$ in un tensore 4D $W' \in \mathbb{R}^{O \times I \times 1 \times 1}$ tramite operazioni di unsqueeze. Il tensore di input viene similmente trasposto da $(B, N, C)$ a $(B, C, 1, N)$ per allinearsi ai requisiti del flusso di dati 4D della BPU.
2. Riformulazione della LayerNorm: L'operazione di Layer Normalization viene decomposta in una sequenza di operazioni di convoluzione $1 \times 1$. La media viene calcolata utilizzando una convoluzione con pesi impostati a $1/H$ (dove $H$ è la dimensione nascosta), e la varianza viene calcolata sottraendo la media, elevando al quadrato il risultato e applicando una seconda convoluzione. I parametri apprendibili $\gamma$ e $\beta$ vengono riconfigurati per corrispondere alle dimensioni dell'input.
3. Ereditarietà dei Pesi: Un aspetto cruciale di questo approccio è che il modello ristrutturato non richiede riaddestramento o fine-tuning. I pesi pre-addestrati originali dei checkpoint standard di DeiT vengono ereditati direttamente dai nuovi strati convoluzionali ristrutturati.
4. Quantizzazione: I modelli modificati vengono elaborati attraverso la toolchain dell'algoritmo Horizon BPU, utilizzando la Post-Training Quantization (PTQ) con una strategia min-max e un dataset di calibrazione per convertire i modelli in precisione INT8.

Contributi Chiave

• Primo Deployment di Successo: Gli autori affermano che questa è la prima implementazione riuscita di Vision Transformer sulla piattaforma BPU che sfrutta appieno le sue capacità di accelerazione hardware.
• Adattamento dell'Architettura: Propongono un metodo innovativo per ristrutturare DeiT sostituendo gli strati lineari e LayerNorm con operatori convoluzionali equivalenti, consentendo l'accelerazione INT8 senza perdita architettonica.
• Deployment Senza Riaddestramento: Il metodo permette l'ereditarietà diretta dei pesi pre-addestrati, eliminando il costo computazionale e la complessità del riaddestramento o del fine-tuning dell'intero modello per la nuova struttura di operatori.
• Validazione Estesa: L'approccio è validato su due dataset distinti: il dataset su larga scala ImageNet e un dataset più piccolo di classificazione dei fiori.

Risultati Sperimentali
Il documento riporta estesi esperimenti che confrontano i modelli DeiT originali con le versioni ottimizzate per BPU e quantizzate:

• Performance su ImageNet:
  • Accuratezza: Il modello quantizzato DeiT-Base raggiunge l'80,4% di accuratezza Top-1 su ImageNet, un calo di solo l'1,4% rispetto all'81,8% originale. Le varianti distillate mostrano una resilienza ancora maggiore, con il DeiT-Base Dist. che mantiene l'81,7% di accuratezza (un calo del 1,7%).
  • Rilevamento Anomalie: Gli autori hanno osservato che sia il modello originale che quello quantizzato hanno predetto correttamente le etichette per diverse immagini nel set di validazione di ImageNet che erano note per essere etichettate erroneamente nelle annotazioni ufficiali, suggerendo che l'effettiva accuratezza dei modelli possa essere leggermente superiore alle metriche riportate.
• Performance sulla Classificazione dei Fiori:
  • Su un dataset con 5 classi di fiori, il modello quantizzato DeiT-Small ha raggiunto il 97,4% di accuratezza, con un calo di solo lo 0,3% rispetto alla versione a piena precisione.
• Accelerazione dell'Inference:
  • I modelli quantizzati hanno dimostrato significativi incrementi di velocità di inferenza sull'hardware BPU rispetto alle controparti a virgola mobile.
  • DeiT-Base: Ha ottenuto un'accelerazione di 3,82×.
  • DeiT-Small: Ha ottenuto un'accelerazione di 2,16×.
  • DeiT-Tiny: Ha ottenuto un'accelerazione di 1,30×.
  • I risultati indicano che i benefici della quantizzazione e dell'accelerazione hardware sono più pronunciati nelle architetture di modelli più grandi.

Significatività e Rivendicazioni
Il lavoro si posiziona come un passo fondamentale per colmare il divario tra i modelli di visione artificiale basati su Transformer e l'hardware specializzato embedded. Dimostrando che i ViT possono essere implementati su BPU con una degradazione minima dell'accuratezza e sostanziali miglioramenti di velocità, gli autori sostengono che il loro metodo fornisca una soluzione pratica per scenari di deployment edge nel mondo reale. La capacità di utilizzare pesi pre-addestrati senza riaddestramento è evidenziata come un grande guadagno di efficienza, rendendo la transizione verso i Transformer accelerati dall'hardware fattibile per ambienti con risorse limitate. Gli autori affermano con modestia che questo rappresenta il primo tentativo riuscito di utilizzare pienamente l'accelerazione BPU per i modelli DeiT.