Discrete Cosine Transform Based Decorrelated Attention for Vision Transformers

1. Enunciato del Problema

I Vision Transformers (ViT) hanno rivoluzionato la visione artificiale applicando meccanismi di auto-attenzione a patch di immagini. Tuttavia, rimangono due sfide critiche:

• Sensibilità all'Inizializzazione: L'approccio standard inizializza le matrici di proiezione per query ($W_Q$), key ($W_K$) e value ($W_V$) con numeri casuali (ad esempio, inizializzazione Xavier o He). Questo tratta i pesi come "rumore bianco", coprendo uniformemente lo spettro ma mancando di conoscenze a priori strutturali. Ciò può portare a dinamiche di addestramento instabili, specialmente su piccoli dataset, e richiede risorse computazionali significative per convergere.
• Efficienza Computazionale: I meccanismi di auto-attenzione sono computazionalmente costosi (complessità $O(N^2)$ rispetto alla lunghezza della sequenza e alti conteggi di parametri). I metodi di compressione esistenti spesso sacrificano l'accuratezza o non riescono a sfruttare le strutture di correlazione specifiche intrinseche nei dati visivi e nelle matrici di attenzione.

2. Metodologia

Gli autori propongono due metodi complementari che sfruttano la Trasformata Discreta del Coseno (DCT) per affrontare questi problemi:

A. Inizializzazione Basata su DCT

Invece dell'inizializzazione casuale, gli autori propongono di inizializzare una delle matrici di peso di proiezione ($W_Q$, $W_K$ o $W_V$) utilizzando la matrice DCT ortogonale di Tipo-II.

• Meccanismo: La matrice DCT funge da estrattore di caratteristiche strutturato. A differenza dei pesi casuali che sono rumore non correlato, i vettori di base DCT coprono l'intero spettro di frequenze in modo sistematico.
• Razionale:
  • Copertura dello Spettro: L'inizializzazione casuale inizia con tutti i vettori di peso aventi la stessa distribuzione spettrale. L'inizializzazione DCT garantisce che ogni vettore di peso copra una diversa banda di frequenze (da basse ad alte frequenze), fornendo un punto di partenza più robusto per l'apprendimento di caratteristiche diversificate.
  • Decorrelazione: I vettori di base DCT approssimano la Trasformata di Karhunen-Loève (KLT) per segnali altamente correlati (come le immagini naturali), decorrelando efficacemente le caratteristiche di input all'inizio dell'addestramento.
• Implementazione: Viene inizializzata con DCT solo una matrice per evitare problemi di simmetria dei pesi (dove un'inizializzazione identica di più matrici porta a un apprendimento scarso). Le altre rimangono inizializzate casualmente.

B. Attenzione Compressa Basata su DCT

Questo metodo sfrutta il dominio della frequenza per ridurre la dimensionalità e il costo computazionale.

• Meccanismo:
  1. Trasformazione: Le patch di input vengono trasformate nel dominio della frequenza utilizzando la DCT lungo l'asse dei canali.
  2. Troncamento: Le componenti ad alta frequenza (che tipicamente corrispondono a rumore e dettagli meno critici) vengono troncate. Vengono mantenute solo la frazione superiore $\tau$ dei coefficienti (ad esempio, 25%, 50%, 75%).
  3. Proiezione: Le proiezioni Query, Key e Value vengono calcolate su questo tensore a dimensionalità ridotta ($\tau C$ invece di $C$).
  4. Ricostruzione: L'output dell'attenzione viene riempito (padded) nuovamente alla dimensione originale e trasformato indietro tramite DCT inversa (IDCT).
• Ottimizzazione: L'IDCT e il livello di output lineare finale ($L_O$) vengono fusi in un singolo livello lineare per eliminare operazioni ridondanti, poiché non vi è alcuna non-linearità tra di essi.
• Efficienza: Questo riduce la complessità computazionale dei livelli lineari da $O(C^3)$ a $O(\tau^2 C^3)$ e riduce significativamente la dimensione della matrice Softmax.

3. Contributi Chiave

1. Strategia di Inizializzazione Innovativa: Introduzione della prima inizializzazione basata su DCT per i pesi di attenzione ViT, dimostrando che le conoscenze a priori strutturate nel dominio della frequenza superano l'inizializzazione casuale.
2. Compressione Efficiente: Proposta di una tecnica di compressione basata su DCT che riduce parametri e FLOPs scartando il rumore ad alta frequenza mentre preserva le informazioni semantiche essenziali a bassa frequenza.
3. Giustificazione Teorica: Fornita una connessione matematica tra DCT e la decorrelazione di patch di immagini altamente correlate (approssimando la KLT), spiegando perché la DCT è superiore al rumore casuale o al semplice pooling per la compressione dell'attenzione.
4. Valutazione Completa: Validazione dei metodi su più architetture (Swin-T, Swin-S, ViT-B) e compiti (Classificazione, Rilevamento, Segmentazione).

4. Risultati Sperimentali

I metodi sono stati valutati sui benchmark CIFAR-10, ImageNet-1K e COCO 2017 utilizzando i Transformer Swin come backbone.

• Classificazione di Immagini (CIFAR-10 e ImageNet-1K):

  • Inizializzazione: Inizializzare $W_K$ (Key) con DCT ha migliorato costantemente l'accuratezza Top-1. Per Swin-T su ImageNet-1K, l'accuratezza è aumentata da 81,47% a 81,73%.
  • Compressione: Mantenere il 75% dei coefficienti DCT ($\tau=0,75$) ha ridotto i parametri di circa il 13% (Swin-T) e il 10% (Swin-S) mantenendo o migliorando leggermente l'accuratezza rispetto alla baseline.
  • Confronto: Il metodo proposto ha superato DCFormer (un precedente metodo basato su DCT) mantenendo accuratamente le componenti a bassa frequenza invece di mediandole, preservando così le caratteristiche discriminative.

• Rilevamento di Oggetti e Segmentazione (COCO 2017):

  • Utilizzando l'inizializzazione DCT ($W_K$) e una compressione del 75%, i modelli hanno ottenuto punteggi di Precisione Media (AP) più elevati rispetto ai Transformer Swin vanilla, riducendo contemporaneamente il conteggio dei parametri di 3M (Swin-T) e 5M (Swin-S).

• Studi di Ablazione:

  • Inizializzazione Singola vs Multipla: Inizializzare più matrici (ad esempio, sia $W_Q$ che $W_K$) con DCT ha degradato le prestazioni a causa di problemi di rottura della simmetria; inizializzare una singola matrice ha prodotto i migliori risultati.
  • Necessità della DCT: Rimuovere la trasformazione DCT dalla pipeline di compressione (lasciando solo la potatura spaziale) ha causato cali significativi di accuratezza, confermando che la decorrelazione nel dominio della frequenza è essenziale per il successo del metodo.

5. Significato

Questo paper dimostra che l'analisi nel dominio della frequenza è uno strumento potente per l'ottimizzazione dei Vision Transformers.

• Stabilità dell'Addestramento: Offre un modo principiato e non casuale per inizializzare reti di attenzione profonde, potenzialmente riducendo i tempi di addestramento e migliorando la convergenza su compiti con scarsità di dati.
• Efficienza senza Sacrifici: Dimostra che il rumore ad alta frequenza può essere scartato in sicurezza dai meccanismi di attenzione, consentendo modelli più leggeri con meno parametri e FLOPs più bassi senza compromettere (e talvolta migliorando) le prestazioni.
• Generalizzabilità: L'approccio è agnostico rispetto all'architettura (testato su Swin e ViT) e agnostico rispetto al compito (funziona per classificazione, rilevamento e segmentazione), rendendolo un'aggiunta versatile al toolkit ViT.

Il codice per l'implementazione è disponibile pubblicamente, facilitando ulteriori ricerche sui meccanismi di attenzione basati sulla frequenza.