From Sparsity to Simplicity: Enabling Simpler Sequential Replacements via Sparse Attention Distillation

Riepilogo Tecnico: Dalla Sparsità alla Semplicità (S2S)

Enunciato del Problema
I meccanismi di self-attention costituiscono la base del preaddestramento di transformer su larga scala, abilitando un mixing flessibile dei token e una modellazione delle dipendenze a lungo raggio. Tuttavia, la loro complessità quadratica rispetto alla lunghezza della sequenza comporta un'elevata latenza di inferenza e un sovraccarico di memoria. Sebbene la sostituzione dell'attenzione con moduli sequenziali più semplici (ad esempio, RNN, Modelli di Spazio di Stato) rappresenti una via attraente per l'efficienza, una sostituzione ingenua spesso porta a un significativo degrado della precisione, specialmente su scale maggiori. Gli approcci esistenti per l'efficienza modificano spesso il calcolo dell'attenzione stesso o potano i token, ma mantengono l'attenzione come operatore fondamentale. Questo lavoro affronta la sfida di sostituire completamente i livelli di attenzione con moduli sequenziali senza una sostanziale perdita di prestazioni.

Metodologia
Il lavoro propone un framework chiamato Sparsity to Simplicity (S2S), che postula che i transformer preaddestrati decompongano le complesse dipendenze tra i token in mappature blocchiali di complessità effettiva variabile. Gli autori ipotizzano che i livelli che esibiscono pattern di attenzione sparsi siano intrinsecamente più facili da approssimare con moduli sequenziali più semplici.

La metodologia è composta da tre componenti fondamentali:

1. Sostituzione Unificata dei Livelli di Attenzione:

   • Gli autori introducono un framework plug-and-play in cui specifici livelli di attenzione in un Vision Transformer (ViT) preaddestrato vengono sostituiti con mixer di token sequenziali (nello specifico BiMamba e BiLSTM), preservando l'interfaccia originale del livello (LayerNorm, connessioni residue e MLP).
   • I moduli sequenziali adottano una struttura bidirezionale multi-testa per corrispondere alla natura non causale e a contesto completo dell'attenzione ViT.
   • Viene utilizzata una supervisione allineata ai livelli durante la distillazione: l'output del mixer sequenziale sostituito viene confrontato direttamente con l'output del livello di attenzione originale alla stessa profondità, utilizzando una funzione di perdita di errore quadratico medio (MSE).

2. Distillazione Guidata dalla Sparsità:

   • Per controllare e studiare l'impatto della sparsità, gli autori integrano A-ViT (Adaptive ViT), che utilizza punteggi di arresto per generare maschere binarie di ritenzione dei token.
   • Fase 1 (Distillazione): Il modello studente (con sostituzioni sequenziali) viene addestrato sotto le stesse maschere di token a livello di livello generate dal teacher (A-ViT). Ciò garantisce che studente e teacher operino su token attivi identici, isolando la difficoltà di approssimare la funzione di mixing dei token dalle differenze nel routing dei token. La perdita include una perdita di similarità, un termine di allineamento dell'arresto e una perdita di classificazione.
   • Fase 2 (Fine-tuning): Lo studente viene sottoposto a fine-tuning utilizzando i propri moduli di arresto appresi per generare maschere, permettendogli di operare in autonomia mantenendo le prestazioni nel compito.

3. Strategia di Sostituzione Sparsa Consapevole del Livello:

   • Invece di una sostituzione uniforme, la strategia mira a specifici gruppi di livelli basati sui loro profili di sparsità naturale. Gli autori sostituiscono solo i livelli più profondi (che esibiscono la sparsità naturale più elevata) e applicano il framework di distillazione guidato dalla sparsità a questi specifici blocchi.

Principali Contributi

• L'Ipotesi S2S: Il lavoro stabilisce che la sparsità dell'attenzione è un segnale della complessità della mappatura. I blocchi sparsi nei transformer preaddestrati possono essere sostituiti da moduli sequenziali più semplici con perdita minima, mentre i blocchi densi sono più difficili da approssimare.
• Framework di Sostituzione Controllata: Un metodo unificato per scambiare l'attenzione con moduli sequenziali (Mamba/LSTM) mantenendo la compatibilità dell'interfaccia e abilitando un'analisi selettiva per livello.
• Controllo Esplicito della Sparsità: L'introduzione di una pipeline di distillazione in cui la sparsità indotta dal teacher (tramite maschere A-ViT) è esplicitamente allineata allo studente, dimostrando che l'aumento della sparsità del teacher riduce costantemente il divario di prestazioni tra studente e teacher.
• Trade-off Ottimizzato: Una strategia che combina la sparsità naturale a livello di livello con una sparsificazione esplicita per ottenere migliori compromessi precisione-efficienza rispetto alle baseline di attenzione sparsa.

Risultati Sperimentali
Gli esperimenti sono stati condotti su ImageNet utilizzando backbones DeiT (Tiny, Small, Base) con BiMamba e BiLSTM come sostituti.

• Sparsità Naturale e Sostituibilità: Le visualizzazioni delle mappe di interazione tra i token mostrano che, mentre i livelli superficiali esibiscono un mixing globale denso, i livelli più profondi diventano marcatamente più sparsi e selettivi. Esperimenti controllati che sostituiscono diversi gruppi di livelli rivelano che sostituire livelli più profondi (più sparsi) comporta cadute di precisione significativamente minori rispetto alla sostituzione di livelli più superficiali (più densi).
• Effetto della Sparsità Esplicita: Quando il rapporto di ritenzione dei token viene ridotto tramite A-ViT, il divario di precisione tra il teacher basato su attenzione e lo studente sequenziale si riduce costantemente. Ciò conferma che la sparsificazione esplicita semplifica la mappatura target, rendendo più facile per i moduli sequenziali apprendere.
• Trade-off Precisione-Efficienza: La sostituzione sparsa consapevole del livello proposta (sostituendo solo i 3 livelli più profondi con BiMamba sotto budget di sparsità corrispondenti) ottiene:
  • Precisione paragonabile o leggermente superiore alle baseline di attenzione sparsa (A-ViT) a tassi di ritenzione di token simili.
  • Guadagni significativi in termini di efficienza: fino a 1,71× di accelerazione stimata a una risoluzione di 384² rispetto alla baseline densa, con una maggiore capacità di throughput di mixing dei token rispetto ai modelli di attenzione sparsa a conteggi di token più elevati.

Significato e Affermazioni
Il lavoro afferma che la sostituzione dell'attenzione dovrebbe essere trattata come un problema strutturato e dipendente dal livello piuttosto che come uno scambio uniforme di architettura. Il significato del lavoro risiede nel dimostrare che:

1. La sparsità è un segnale di sostituibilità: I modelli preaddestrati decompongono naturalmente la complessità in modo tale che i livelli sparsi siano i più suscettibili alla semplificazione.
2. La sparsificazione esplicita aiuta la distillazione: Imporre la sparsità al teacher rende il compito di apprendimento per lo studente sequenziale più semplice, riducendo il divario di prestazioni.
3. La sostituzione strategica supera gli approcci uniformi: Concentrando la sostituzione su livelli naturalmente sparsi e guidando il processo con sparsità controllata, è possibile raggiungere un equilibrio superiore tra precisione e latenza rispetto alla semplice potatura dell'attenzione o alla sostituzione dell'intero modello.

Gli autori concludono che questa prospettiva incoraggia la futura ricerca su modelli ibridi che sfruttano la conoscenza dei transformer preaddestrati integrando al contempo moduli sequenziali più semplici ed efficienti laddove la struttura interna del modello lo consente.