Raddoppia la lunghezza del testo che dai in pasto a un Transformer e il lavoro dell’attention non raddoppia: quadruplica. Ogni parola, per decidere a cosa prestare attenzione, guarda tutte le altre; con mille parole sono un milione di confronti, con duemila quattro milioni. È il prezzo nascosto dietro la potenza dell’architettura che regge quasi tutti i modelli di oggi, e diventa proibitivo esattamente dove servirebbe di più: documenti lunghi, genomi, tracce audio, segnali che si estendono per centinaia di migliaia di passi.
La selettività S6, lo scan hardware-aware e il blocco Mamba sono spiegati nel libro.
Nel dicembre 2023 Albert Gu e Tri Dao pubblicano un paper che riprende un filone considerato perdente — i modelli a spazio di stato — e lo modifica quel tanto che basta per farlo competere con l’attention sul suo terreno più difficile, il linguaggio. Il modello si chiama Mamba, e il suo interesse non sta nell’aver «battuto» i Transformer (non lo ha fatto), ma nell’aver dimostrato che l’attention non era l’unica strada possibile.
Il costo nascosto dell’attention
La forza dell’attention è che ogni token può, in linea di principio, attingere a qualunque altro token della sequenza, vicino o lontanissimo. Nessuna informazione va persa per strada: tutto resta accessibile a tutto. Questa è anche la sua condanna. Tenere ogni coppia di token in relazione significa un costo che cresce con il quadrato della lunghezza, in tempo e in memoria. Per le sequenze corte è un affare; per quelle lunghe è un muro.
Negli anni erano nate molte proposte per aggirarlo — attention lineari, convoluzioni con gate, modelli ricorrenti, gli state space model strutturati. Tutte più economiche, tutte con lo stesso difetto: sul linguaggio non reggevano il confronto. Gu e Dao mettono il dito sul perché.
Una memoria che scorre: gli state space model
Un modello a spazio di stato non guarda indietro a tutti i token: mantiene uno stato, un riassunto compresso di ciò che ha letto finora, e lo aggiorna un passo alla volta man mano che la sequenza avanza. L’analogia è il modo in cui leggi un romanzo. Non rileggi da capo tutte le pagine precedenti ogni volta che incontri una parola nuova: porti in testa un riassunto — chi sono i personaggi, dove siamo arrivati — e lo aggiorni pagina dopo pagina. Il costo di leggere una parola in più è costante, non cresce con la lunghezza del libro. È lineare, non quadratico.
Il difetto delle versioni precedenti stava in come aggiornavano quel riassunto. Lo facevano con una regola fissa, uguale per ogni parola: la stessa compressione applicata meccanicamente a tutto, che si trattasse del nome del protagonista o di una congiunzione. Va benissimo per segnali regolari come l’audio, dove ogni istante conta più o meno quanto gli altri. È un disastro per il linguaggio, dove il valore di una parola dipende disperatamente da quale parola è.
La mossa di Mamba: uno stato che sceglie
Il contributo centrale del paper ha un nome preciso: selettività. Gu e Dao rendono i parametri che governano l’aggiornamento dello stato dipendenti dall’input. In pratica il modello, a ogni token, decide quanto quel token debba entrare nel riassunto e quanto del passato debba conservare o dimenticare. Non una regola fissa, ma un filtro che reagisce al contenuto: trattieni il nome, lascia scorrere il «che».
L’attention non dimentica nulla, e paga tutto. Mamba sceglie cosa dimenticare, e proprio per questo può permettersi di essere lineare.
È il ponte che mancava tra l’economia dei modelli ricorrenti e la capacità dell’attention di fare ragionamento basato sul contenuto. Con questa modifica gli autori costruiscono un’architettura essenziale, senza blocchi di attention e senza i consueti strati completamente connessi: un unico tipo di blocco selettivo, ripetuto.
A sinistra l’attention: ogni token è collegato a tutti gli altri, e il numero di connessioni cresce col quadrato della lunghezza. A destra lo state space selettivo: uno stato (i quadrati teal) scorre in avanti di token in token, e per ciascuno un filtro (i rombi ocra) decide quanta informazione far entrare e quanta lasciar dimenticare. Il costo cresce in modo lineare.
Cosa hanno misurato
La selettività ha un prezzo tecnico: rendere i parametri dipendenti dall’input impedisce le scorciatoie di calcolo (le convoluzioni) su cui gli state space model precedenti facevano leva per essere veloci. Gu e Dao lo pagano progettando un algoritmo che tiene conto di come è fatta davvero una GPU — di quale memoria è veloce e quale lenta — così che il calcolo ricorrente resti efficiente senza materializzare lo stato dove costerebbe troppo. È metà del contributo: l’idea da sola non sarebbe bastata senza l’ingegneria che la rende praticabile.
I risultati riportati nel paper sono concreti. In inferenza Mamba dichiara un throughput cinque volte superiore a quello di un Transformer, e un costo che scala linearmente con la lunghezza della sequenza — con prestazioni che continuano a migliorare su dati reali fino a sequenze lunghe un milione di passi, il territorio dove l’attention soffoca. Sul linguaggio, il modello Mamba da 3 miliardi di parametri supera i Transformer di pari dimensione e regge il confronto con Transformer grandi il doppio, sia in pre-addestramento sia nelle valutazioni a valle. E lo stesso schema, senza modifiche di sostanza, ottiene ottimi risultati anche su audio e genomica: è un’architettura generale, non una scorciatoia buona per un solo dominio.
Perché ti riguarda
Sarebbe fuori luogo raccontare Mamba come il successore dei Transformer. A due anni e mezzo di distanza l’attention è ancora l’architettura dominante, e la strada più battuta non è stata la sostituzione ma il compromesso: modelli ibridi che alternano blocchi di attention e blocchi in stile Mamba, per prendere il ragionamento sul contenuto dei primi e l’economia sulle sequenze lunghe dei secondi. Il dibattito su quanto e dove convenga davvero è tutt’altro che chiuso.
Il motivo per cui vale la pena conoscerlo è un altro. Per anni «modello linguistico» ha significato, di fatto, «Transformer»: l’attention era diventata talmente ovvia da sparire come ipotesi. Mamba l’ha rimessa sul tavolo come una scelta, con dei costi e delle alternative. Se costruisci o valuti sistemi che devono digerire contesti lunghissimi — interi archivi, log, sequenze biologiche — sapere che il conto quadratico dell’attention non è una legge di natura, ma una fra le opzioni possibili, cambia le domande che poni prima di scegliere l’architettura. E imparare a chiedersi cosa un modello può permettersi di dimenticare è, spesso, più utile che pretendere che ricordi tutto.
Il paper originale
Albert Gu, Tri Dao, Mamba: Linear-Time Sequence Modeling with Selective State Spaces, 2023. Disponibile su arXiv: arxiv.org/abs/2312.00752.
I commenti sono riservati agli iscritti.
Accedi per commentare