Mamba-2 (2024): SSM e attention sono la stessa cosa

Prendete due algoritmi che sembrano non avere niente in comune. Il primo è una ricorrenza: legge una sequenza un elemento alla volta, aggiorna uno stato compatto, va avanti, non torna mai indietro — è il cuore dei modelli a spazio degli stati, gli SSM, di cui Mamba è l’incarnazione recente. Il secondo è l’attenzione dei Transformer: costruisce una matrice che confronta ogni posizione con tutte le altre, un reticolo quadratico di interazioni. Uno scorre in linea, l’altro guarda tutto insieme. Nel maggio 2024 Tri Dao e Albert Gu — gli stessi autori dietro FlashAttention e Mamba — pubblicano un paper il cui titolo è già una tesi: Transformers are SSMs. I due algoritmi calcolano la stessa cosa.

Due modi di guardare la stessa matrice

Il ponte tra i due mondi è un oggetto matematico con un nome intimidatorio e un’idea semplice: la matrice semiseparabile. Immaginate un registro enorme in cui, però, ogni voce si può ricostruire da un riassunto compatto: non serve conservare tutto il reticolo, basta una descrizione a basso rango della sua struttura. Il paper dimostra che l’intera trasformazione operata da un SSM su una sequenza si può scrivere come moltiplicazione per una matrice di questo tipo, triangolare (il presente non dipende dal futuro) e semiseparabile.

Ecco lo snodo. Una stessa matrice semiseparabile si può calcolare in due modi opposti, che danno il medesimo risultato. Il primo la applica come farebbe l’attenzione: costruisce esplicitamente le interazioni tra le posizioni e le combina — costo quadratico nella lunghezza della sequenza, ma parallelizzabile e amico delle GPU. Il secondo la applica come farebbe un SSM: scorre in avanti mantenendo uno stato riassuntivo — costo lineare, ideale per generare un token dopo l’altro. Gli autori chiamano questa equivalenza State Space Duality, dualità dello spazio degli stati. Non è un’analogia suggestiva: è un’identità algebrica.

La dualità in un’immagine: la stessa matrice semiseparabile M si calcola come un reticolo di interazioni in stile attenzione (a sinistra, quadratico e parallelo) oppure come una catena di stati in stile ricorrenza (a destra, lineare e sequenziale). Stesso risultato, due algoritmi.

Il trucco a blocchi

Un’identità matematica sarebbe poca cosa se non producesse un algoritmo migliore. Qui la produce, e nasce dal domandarsi: perché scegliere? Se le due forme calcolano la stessa matrice, si possono anche mescolare. Il paper divide la sequenza in blocchi contigui. Dentro ogni blocco usa la forma quadratica, in stile attenzione: i blocchi sono corti, il costo quadratico su pochi elementi è trascurabile e sfrutta a pieno le moltiplicazioni matriciali per cui le GPU sono ottimizzate. Tra un blocco e l’altro usa la forma lineare: passa avanti solo lo stato riassuntivo, senza ricostruire tutte le interazioni.

È la strategia di chi organizza un lungo viaggio a tappe: dentro ogni tappa cammina e vede ogni cosa nel dettaglio; tra una tappa e l’altra non rifà il percorso, si porta appresso solo un riassunto di dove è arrivato. Il risultato combina il meglio dei due regimi — la velocità hardware del calcolo denso a blocchi e la scalabilità lineare del passaggio di stato — ed è la ragione tecnica per cui il nuovo strato risulta molto più rapido del precedente.

La domanda non era «ricorrenza o attenzione», ma «in quale porzione del calcolo conviene l’una e in quale l’altra».

Cosa cambia in Mamba-2

Da questa impalcatura teorica gli autori derivano un’architettura, Mamba-2, che è una revisione dello strato SSM selettivo di Mamba. Le differenze non sono cosmetiche. Per rendere il calcolo esprimibile come moltiplicazione matriciale, Mamba-2 impone alla dinamica dello stato una struttura più semplice rispetto a Mamba-1, rinunciando a un po’ di libertà nella ricorrenza in cambio della possibilità di usare le unità matriciali della GPU. In compenso può permettersi uno stato interno molto più grande — l’ordine di grandezza cresce in modo netto rispetto ai valori tipici di Mamba-1 — perché la forma a blocchi rende quel calcolo comunque efficiente. Lo strato acquisisce inoltre una struttura a «teste» parallele che ricalca da vicino l’organizzazione a testa multipla dell’attenzione, chiudendo il cerchio con l’architettura da cui il paper prende le distanze.

Il conto pratico, riportato dagli autori: il nuovo strato SSM selettivo risulta da 2 a 8 volte più veloce di quello di Mamba-1, restando competitivo con i Transformer nei compiti di modellazione del linguaggio. La dualità, insomma, non resta sulla lavagna: si traduce in un’architettura che si addestra più in fretta e regge sequenze più lunghe con uno stato più capiente.

Quello che la dualità non cancella

Vale la pena non sovrainterpretare. «SSM e attention sono la stessa cosa» è vero nel senso preciso e ristretto del paper: esiste una classe di attenzione strutturata — con una particolare mascheratura — che coincide con la forma matriciale di questi SSM. Non significa che qualunque Transformer sia un SSM travestito, né che l’attenzione piena, quella che confronta liberamente ogni coppia di posizioni, sparisca in una ricorrenza economica. Rispetto a quest’ultima gli SSM comprimono il passato in uno stato di dimensione fissa, e restano prove — soprattutto su compiti di recupero esatto di informazioni dal contesto — in cui questa compressione costa qualcosa che l’attenzione piena non paga. La dualità non promette che una famiglia soppianti l’altra; illumina il terreno comune, e spiega perché nella pratica gli ibridi che alternano strati SSM e strati di attenzione funzionino così bene.

Perché ti riguarda

Se costruisci sistemi con LLM, il valore di questo paper è meno nel singolo modello e più nel modo di pensare. Le due grandi famiglie di architetture per sequenze — quella dominante dei Transformer e quella emergente degli SSM — non sono tribù rivali da tifare, ma casi particolari di una stessa struttura algebrica, con profili di costo diversi da spendere dove servono: attenzione dove conta il recupero fedele dal contesto, ricorrenza lineare dove contano sequenze lunghissime e generazione economica. È il tipo di unificazione che raramente fa titolo ma cambia il campo in profondità: quando due cose che sembravano alternative si rivelano due viste dello stesso oggetto, smetti di scegliere per fede e cominci a scegliere per ingegneria. E, come spesso capita, l’algoritmo più veloce è arrivato non da un’ottimizzazione in più, ma da una comprensione più chiara di cosa si stava calcolando davvero.

Il paper originale

Tri Dao, Albert Gu, Transformers are SSMs: Generalized Models and Efficient Algorithms Through Structured State Space Duality, 2024 — arXiv:2405.21060.