LaplacianFormer:Rethinking Linear Attention with Laplacian Kernel

Titolo: LaplacianFormer: Ripensare l'Attenzione Lineare con il Kernel di Laplace

1. Il Problema

L'architettura Transformer, fondamentale per il modellamento delle sequenze e la visione artificiale, soffre di una complessità computazionale e spaziale quadratica ($O(N^2)$) dovuta al meccanismo di attenzione softmax. Questo limita la scalabilità a sequenze lunghe o immagini ad alta risoluzione.
Per mitigare questo problema, sono state proposte varianti di "attenzione lineare" che sostituiscono il softmax con approssimazioni basate su kernel (spesso kernel Gaussiani). Tuttavia, il paper identifica due criticità principali nelle soluzioni esistenti:

• Mancanza di fondamento teorico: L'uso diffuso del kernel Gaussiano è spesso una convenzione di default piuttosto che una scelta teoricamente giustificata per le interazioni Query-Key nella visione artificiale.
• Soppressione eccessiva delle interazioni: Il kernel Gaussiano, basato sulla distanza $\ell_2^2$, decade rapidamente. Le distribuzioni empiriche delle distanze tra vettori Query e Key nei Transformer visivi (come DeiT, PVT, Swin) mostrano code pesanti (outliers). Il kernel Gaussiano amplifica l'effetto di queste code, portando a una soppressione eccessiva delle interazioni a media distanza e a gradienti che svaniscono, compromettendo l'espressività dell'attenzione e la stabilità dell'ottimizzazione.

2. Metodologia

Gli autori propongono LaplacianFormer, una variante del Transformer che sostituisce il kernel Gaussiano con un kernel di Laplace, motivata da analisi empiriche e teoriche.

• Kernel di Laplace: Invece della distanza euclidea quadrata, il modello utilizza la distanza $\ell_1$ (Manhattan) nella funzione di similarità:
  $$k(x, y) = \exp\left(-\frac{\|x - y\|_1}{\lambda}\right)$$
  Il kernel di Laplace decade più lentamente rispetto a quello Gaussiano, preservando meglio le interazioni a media distanza e fornendo gradienti non nulli anche quando i vettori sono molto simili (grazie alla natura lineare a tratti della norma $\ell_1$), migliorando la stabilità dell'addestramento.

• Mappatura delle Feature Iniettiva: Per evitare la degradazione dell'espressività tipica delle approssimazioni a basso rango, viene introdotta una mappatura delle feature provabilmente iniettiva. Questa trasformazione normalizza i vettori di similarità query-key utilizzando una matrice di sbiancamento (whitening) approssimata in modo diagonale, garantendo che query distinte producano output distinti e preservando le informazioni token a grana fine.

• Approssimazione Nyström e Newton-Schulz:

  • Per calcolare efficientemente l'attenzione basata su kernel, viene utilizzata l'approssimazione di Nyström, che seleziona un sottoinsieme di "chiavi landmark" per ridurre la complessità della matrice del kernel.
  • Per invertire la matrice del kernel dei landmark (necessaria per la normalizzazione), invece di usare costose decomposizioni SVD o inversioni dirette, viene adottata l'iterazione di Newton-Schulz. Questo metodo si basa solo su moltiplicazioni di matrici, è adatto per le GPU e garantisce convergenza stabile anche per matrici mal condizionate (grazie a una piccola perturbazione diagonale).

• Implementazione CUDA: Sono stati sviluppati kernel CUDA personalizzati per il calcolo del kernel di Laplace e per l'iterazione di Newton-Schulz, ottimizzando l'accesso alla memoria globale e il riutilizzo dei registri per abilitare passaggi forward e backward ad alto throughput, adatti anche al deployment su dispositivi edge.

3. Contributi Chiave

1. LaplacianFormer: Un nuovo framework di attenzione lineare basato sul kernel di Laplace che migliora la modellazione delle dipendenze a lungo raggio mantenendo la scalabilità.
2. Analisi Teorica ed Empirica: Dimostrazione che il kernel Gaussiano è subottimale per la distribuzione reale delle distanze Query-Key nella visione, e che il kernel di Laplace offre un profilo di decadimento più adatto e gradienti più stabili.
3. Ottimizzazione Efficiente: Sviluppo di un'implementazione CUDA ottimizzata che integra l'attenzione di Laplace con un risolutore inverso basato su Newton-Schulz, eliminando la necessità di inversioni di matrice costose.
4. Validazione Estesa: Sperimentazione su ImageNet-1k e compiti di visione downstream (rilevamento oggetti e segmentazione di istanze), dimostrando prestazioni competitive.

4. Risultati

Gli esperimenti condotti su ImageNet-1k e altri benchmark mostrano:

• Precisione-Efficienza: LaplacianFormer supera i modelli esistenti (come SOFT++, Swin, PVT, AgentFormer) in tutti i range di FLOPs, offrendo un miglior compromesso tra accuratezza e costo computazionale. Ad esempio, la versione Large raggiunge l'85.6% di Top-1 accuracy con 11.2G FLOPs, superando modelli simili.
• Scalabilità: Il modello mostra una scalabilità lineare della memoria rispetto alla lunghezza della sequenza, superando significativamente il Transformer standard e competendo con altre soluzioni lineari efficienti.
• Task Downstream: Nei compiti di rilevamento oggetti (Mask R-CNN) e segmentazione (RetinaNet), LaplacianFormer stabilisce nuovi stati dell'arte (SOTA) per diverse dimensioni di modello (es. 48.0 APb per la versione Medium su Mask R-CNN).
• Convergenza: L'uso dell'iterazione di Newton-Schulz dimostra una convergenza più robusta rispetto al metodo del gradiente coniugato (CG), specialmente in regimi mal condizionati, portando a una maggiore accuratezza finale.

5. Significato

Questo lavoro sfida la convenzione consolidata nell'uso del kernel Gaussiano per l'attenzione lineare nella visione artificiale. Sostituendolo con un kernel di Laplace, gli autori non solo migliorano le prestazioni pratiche, ma forniscono anche una giustificazione teorica basata sulla distribuzione statistica delle interazioni token. L'introduzione di un risolutore inverso efficiente (Newton-Schulz) e di implementazioni CUDA dedicate rende il modello non solo teoricamente solido, ma anche praticamente deployabile su hardware con risorse limitate, aprendo la strada a Transformer più scalabili ed espressivi per applicazioni reali ad alta risoluzione.