Riepilogo Tecnico: FTerViT: Vision Transformer Completamente Ternario

Enunciato del Problema

I Vision Transformer (ViT) sono potenti classificatori di immagini, ma generalmente non sono adatti per dispositivi di classe microcontrollore a causa del loro significativo footprint di memoria. Un ViT compatto standard come DeiT-Small richiede circa 88,3 MB in precisione FP32, mentre le tipiche unità microcontrollore (MCU) offrono solo pochi megabyte di RAM esterna. Questa discrepanza impedisce il dispiegamento di applicazioni di visione sempre attive e a basso consumo energetico su dispositivi edge senza l'offloading al cloud.

Sebbene la quantizzazione ternaria (vincolando i pesi a $\{-1, 0, +1\}$) offra una compressione teorica di 16$\times$ rispetto all'FP32, i metodi esistenti all'avanguardia ternarizzano solo i livelli dell'encoder del transformer. Mantengono l'embedding delle patch, i parametri LayerNorm e le testine di classificazione in precisione completa (FP32) o in interi a bit inferiore (INT8). Nei modelli compatti, questi componenti rimanenti in precisione completa diventano il collo di bottiglia dominante, consumando fino al 38% delle dimensioni del modello nonostante rappresentino una piccola frazione dei parametri. Inoltre, i precedenti tentativi di ternarizzare completamente questi componenti sensibili (ad esempio TerViT) hanno portato a significativi cali di accuratezza (fino a 22,4 punti percentuali), in gran parte dovuti alla nota sensibilità dei primi e ultimi livelli alla quantizzazione.

Metodologia

Gli autori propongono FTerViT, il primo Vision Transformer completamente ternarizzato in cui tutte le matrici dei pesi e i parametri di normalizzazione sono vincolati a $\{-1, 0, +1\}$. L'approccio si compone di due pilastri principali: nuovi operatori ternari e un protocollo di addestramento specializzato.

1. Primitivi Ternari

Per eliminare i colli di bottiglia FP32, gli autori sostituiscono i componenti standard con moduli ternari dedicati:

• TernaryBitLinear: Utilizzato per i livelli completamente connessi. Quantizza i pesi a $\{-1, 0, +1\}$ utilizzando una scalatura per tensor basata sulla media assoluta ($absmean$) e un arrotondamento limitato. Le attivazioni sono quantizzate a interi a 8 bit utilizzando una scalatura per token basata sul massimo assoluto.
• TernaryBitConv2d: Progettato specificamente per il livello di embedding delle patch. Applica una scalatura per canale per gestire le magnitudini eterogenee dei filtri, affrontando la sensibilità del livello di embedding delle patch che lavori precedenti hanno identificato come il componente più critico.
• TernaryLayerNorm: Ternarizza i parametri affini apprendibili ($\gamma$ e $\beta$) di LayerNorm utilizzando lo stesso schema $absmean$. Le statistiche per token (media e varianza) e i bias rimangono in FP32 poiché il loro numero di parametri è trascurabile, ma la loro ternarizzazione distruggerebbe le informazioni posizionali.

2. Protocollo di Addestramento

Il modello viene addestrato utilizzando una strategia di Distillazione della Conoscenza (KD) in due fasi:

• Fase 1 (Addestramento): Lo studente ternario viene inizializzato da un insegnante FP32 pre-addestrato e addestrato utilizzando la Distillazione Consapevole della Quantizzazione (QAD) con un tasso di apprendimento di $10^{-4}$ (decadimento coseno) per 250 epoche. La funzione di perdita utilizza solo il termine di divergenza KL in avanti tra le logit dello studente e dell'insegnante, evitando la cross-entropia che può entrare in conflitto con il KL nelle reti a basso bit.
• Fase 2 (Recupero): Una fase di fine-tuning leggera in cui il tasso di apprendimento viene riavviato a $10^{-5}$ (decadimento coseno) per 10 epoche. Questa fase permette la ricalibrazione delle scale e il recupero dell'accuratezza persa durante la saturazione iniziale della Fase 1.

Contributi Chiave

1. Primo ViT Completamente Ternario: Gli autori dimostrano che i componenti più fragili dei ViT (embedding delle patch, LayerNorm e testina di classificazione) possono essere ternarizzati con successo a $\{-1, 0, +1\}$ quando addestrati con distillazione della conoscenza, spingendo i limiti di quanto fosse precedentemente considerato quantizzabile.
2. Operatori Novelli: Introduzione di TernaryBitConv2d con scalatura per canale e TernaryLayerNorm per gestire la specifica sensibilità dei livelli non-encoder.
3. Strategia di Addestramento: Una semplice ma efficace strategia di addestramento in due fasi (addestramento con alto LR seguito da recupero con basso LR) che permette una convergenza più rapida e un'accuratezza finale significativamente più alta rispetto ai precedenti approcci ViT ternari.
4. Implementazione su Dispositivo: La prima implementazione C autonoma di un Vision Transformer ternario su un microcontrollore ESP32-S3 dual-core con risorse limitate (8 MB di PSRAM).

Risultati Sperimentali

Il metodo è stato valutato su ImageNet-1K e altri benchmark:

• Accuratezza e Compressione:

  • FTerViT-DeiT-III-S (384$\times$384): Raggiunge un'accuratezza top-1 del 82,43% a 6,09 MB (compressione 14,6$\times$), perdendo solo 2,42 punti percentuali rispetto al suo insegnante FP32. Questo supera i precedenti metodi ViT ternari (ad esempio TerViT) fino a 8 punti percentuali.
  • FTerViT-DeiT-III-S (224$\times$224): Raggiunge un 79,64% top-1 a 5,81 MB.
  • FTerViT-DeiT-S (224$\times$224): Raggiunge un 77,47% top-1 a 5,81 MB, superando TerViT-DeiT-S di 3,27 punti percentuali.
  • Benchmark CIFAR: Le varianti DeiT-Tiny raggiungono il 97,43% (CIFAR-10) e l'86,01% (CIFAR-100) con solo 1,53 MB di archiviazione.

• Analisi della Sensibilità dei Componenti:

  • Le analisi del primo ordine di Taylor e della traccia di Hessian confermano che LayerNorm e embedding delle patch dominano l'importanza del modello (34–39% e 4,6% rispettivamente) nonostante occupino <0,2% dei parametri. Questo spiega perché i metodi precedenti li lasciavano in precisione più alta e perché la loro ternarizzazione di successo è critica.

• Dispiegamento su Dispositivo (ESP32-S3):

  • Il modello da 5,81 MB si adatta interamente agli 8 MB di PSRAM dell'ESP32-S3, lasciando circa 4,5 MB liberi per i buffer della fotocamera e dell'LCD.
  • Prestazioni: Un passaggio in avanti richiede 21,06 secondi.
  • Efficienza: Rispetto all'FP32, l'inferenza ternaria a 2 bit compressa è 1,71$\times$ più veloce (804 ms vs 1376 ms per un livello specifico) e riduce il consumo energetico sopra l'idle del 55% (59 mJ vs 130 mJ).

Significato

Il documento afferma che FTerViT riduce sostanzialmente il divario tra l'accuratezza standard dei ViT e i vincoli di memoria rigorosi del dispiegamento su scala microcontrollore. Ternarizzando con successo l'intera architettura, inclusi gli embedding delle patch e i livelli di normalizzazione precedentemente "intoccabili", gli autori abilitano il dispiegamento di modelli di visione ad alta accuratezza su MCU commerciali a basso costo (10$). Il lavoro valida che, con la corretta strategia di addestramento (KD con stessa architettura e recupero in due fasi), la sensibilità di questi componenti può essere gestita, rendendo i ViT completamente ternari una realtà pratica per l'AI edge. Gli autori notano che il divario di accuratezza rispetto all'FP32 scala inversamente con la capacità del modello, con modelli più grandi (DeiT-III-S384) che subiscono una perdita di accuratezza relativa inferiore rispetto a quelli più piccoli.