Enhancing Layer Interaction Using Key-Correlated Layer Attention

Sintesi Tecnica: Miglioramento dell'Interazione tra i Layer tramite la Key-Correlated Layer Attention

Definizione del Problema
I recenti progressi nella progettazione di reti neurali profonde (DNN) hanno utilizzato la Layer Attention per migliorare le interazioni tra i layer, trattando gli output dei layer come token di una sequenza. Tuttavia, la standard layer attention soffre di una complessità computazionale quadratica ($O(D^2)$) rispetto alla profondità della rete $D$, insieme a una crescita significativa della memoria dovuta alla necessità di memorizzare (cache) gli output di tutti i layer precedenti.

Per affrontare queste problematiche, lavori precedenti hanno introdotto la Recurrent Layer Attention (RLA) e approssimazioni lineari come MRLA-L. Sebbene la RLA riduca l'overhead computazionale memorizzando le rappresentazioni Key ($K$) e Value ($V$), essa funziona come un meccanismo statico in cui le caratteristiche dei primi layer rimangono fisse una volta generate, limitando gli aggiornamenti dinamici delle informazioni. Inoltre, le approssimazioni lineari come MRLA-L tentano di ridurre la complessità approssimando linearmente l'output della attention. Gli autori osservano che questo approccio causa un attenuazione esponenziale dei contributi dei layer precedenti (come mostrato nella loro analisi dei coefficienti di decadimento $\beta_l$), rendendo difficile stabilire dipendenze cross-layer a lungo raggio efficaci. Inoltre, i metodi di linearizzazione esistenti spesso non affrontano il collo di bottiglia della memoria nel caching degli output dei layer precedenti o le sfide specifiche nell'applicare la linear attention a strutture DNN gerarchiche.

Metodologia: Key-Correlated Layer Attention (KCLA)
La proposta Key-Correlated Layer Attention (KCLA) è un nuovo metodo di layer attention linearizzato, progettato per raggiungere una complessità computazionale lineare preservando al contempo gli aggiornamenti dinamici e la modellazione delle dipendenze a lungo raggio. La metodologia si basa su tre intuizioni e meccanismi chiave:

1. Osservazione della Correlazione delle Key: L'analisi empirica rivela che le rappresentazioni Key attraverso diversi layer nella standard layer attention esibiscono un'alta similarità del coseno. Sfruttando ciò, KCLA approssima l'output della standard layer attention fattorizzando la trasformazione lineare del Value ($V$) e approssimando le rappresentazioni Key come direzionalmente correlate. Ciò consente l'aggregazione degli output dei layer in uno stato ricorrente ($U^t$) senza memorizzare l'intera cronologia delle matrici $K$ e $V$.
2. Restaurazione Dinamica della Softmax: Una sfida critica nella linearizzazione della attention è la perdita dell'operazione di softmax, essenziale per normalizzare i punteggi di attention e consentire aggiustamenti dinamici della distribuzione. KCLA introduce un meccanismo di multi-head attention modificato per ripristinare questa capacità:
   • Utilizza più teste con diversi fattori di scala (temperature) derivati da una sequenza logaritmicamente uniforme.
   • Utilizza una trasformazione lineare apprendibile $\Phi$ per mappare la norma della Query ($\|Q^t\|$) e l'angolo di correlazione ($\cos \theta$) in pesi specifici per ogni testa.
   • Ciò consente al modello di regolare dinamicamente la distribuzione dell'attention ad ogni layer, simulando efficacemente il ruolo del parametro di temperatura della softmax senza calcolare l'intera matrice $N \times N$.
3. Adattamento Cross-Stage: Per adattarsi a DNN con diverse dimensioni delle mappe di caratteristiche e dimensioni dei canali attraverso gli stage, KCLA implementa una specifica procedura di adattamento (Algoritmo 1). Questa prevede la normalizzazione delle caratteristiche accumulate, l'applicazione di una convoluzione $1 \times 1$ con stride per il downsampling e il ricalcolo dei fattori di normalizzazione per colmare le disparità dimensionali tra gli stage.

Contributi Chiave

• Osservazione e Dimostrazione: Gli autori dimostrano che le rappresentazioni Key attraverso diversi layer nella standard layer attention sono altamente correlate in direzione, fornendo la base teorica per l'approssimazione lineare.
• Nuova Linearizzazione: KCLA è proposto come un metodo di layer attention linearizzato che evita la natura statica della RLA e l'attenuazione esponenziale di MRLA-L. Raggiunge una complessità computazionale lineare e una complessità spaziale fissa indipendente dalla profondità della rete.
• Approssimazione della Softmax Dinamica: Viene introdotto un nuovo metodo per approssimare l'operazione di softmax utilizzando lo scaling multi-head e trasformazioni apprendibili, consentendo un aggiustamento dinamico e adattivo delle distribuzioni di attention ad ogni layer.

Risultati Sperimentali
Gli autori hanno valutato KCLA su tre compiti fondamentali di visione: classificazione di immagini, rilevamento di oggetti e segmentazione di immagini mediche.

• Classificazione di Immagini: Su CIFAR-100, KCLA ha ottenuto le migliori prestazioni su profondità ResNet di 20, 56 e 110, superando MRLA-L, DLA-L e la standard layer attention. Su ImageNet-1K, l'integrazione di KCLA in ResNet-50 e ResNet-101 ha migliorato l'accuratezza Top-1 rispettivamente del +1,7% e del +1,4% rispetto al baseline. Ha mantenuto un vantaggio lieve ma costante su MRLA-L utilizzando significativamente meno parametri rispetto al metodo dinamico DLA-L.
• Rilevamento di Oggetti: Sul dataset COCO2017 utilizzando Faster R-CNN e Mask R-CNN, KCLA ha dimostrato prestazioni competitive. Con ResNet-101, ha superato l'architettura più complessa DLA-L in diverse metriche (AP$_{75}$, AP$_S$, AP$_M$) utilizzando 2,2 milioni di parametri in meno.
• Segmentazione di Immagini Mediche: Sui dataset ISIC2017, ISIC2018 e Kvasir-SEG, KCLA ha costantemente superato MRLA-L e il baseline ResNet-UNet, ottenendo punteggi mIoU e DSC superiori con un overhead di parametri minimo.
• Analisi della Complessità: A differenza di MRLA-B, la cui memoria e FLOPs crescono quadraticamente con la profondità, KCLA mostra una crescita lineare sia nel carico computazionale che nel costo di memoria.

Significatività e Rivendicazioni
Il paper sostiene che KCLA colma con successo il divario tra l'efficienza della linear attention e il potere rappresentativo della standard attention. Approssimando direttamente il meccanismo della standard layer attention basandosi sulle correlazioni delle Key, KCLA evita i problemi di rappresentazione statica delle feature della RLA e la degradazione delle dipendenze a lungo raggio di MRLA-L.

Gli autori sottolineano che il loro metodo mantiene la capacità di modellare efficientemente le dipendenze cross-layer a lungo raggio, offrendo un equilibrio desiderabile tra efficienza e prestazioni. Notano che il metodo proposto per ripristinare la funzione softmax è un fattore critico per il raggiungimento di questi risultati. Sebbene l'attuale implementazione si basi su parametri di temperatura pre-impostati, gli autori identificano la progettazione di una layer attention multi-head con temperatura auto-adattiva e dinamica come una direzione per l'ottimizzazione futura. Il codice è pubblicamente disponibile per facilitare la riproducibilità.