State Space Model#
C’è un’idea che l’ingegneria usa da oltre mezzo secolo per descrivere qualunque sistema che evolve nel tempo: un termostato, la traiettoria di un razzo, un circuito elettrico. Si chiama modello a spazio degli stati (state space model, in sigla SSM, che è l’abbreviazione che useremo da qui in avanti): un pugno di equazioni che riassumono tutto il passato di un segnale (una grandezza che cambia nel tempo: il livello dell’acqua in una vasca, un suono, una sequenza di parole) in uno stato interno, e da quello prevedono il futuro. È la matematica dei filtri di Kalman che portarono l’Apollo sulla Luna. Che cosa ci fa in un libro sull’intelligenza artificiale?
Ci fa che risolve il problema lasciato aperto dal capitolo sui Transformer. Là, per capire una frase, ogni parola guarda tutte le altre: raddoppiare la lunghezza del testo quadruplica il lavoro, e su un testo molto lungo il conto diventa proibitivo. Serve una macchina che legga lungo restando veloce, cioè che quando il testo raddoppia raddoppi il lavoro e basta. È questo che si intende con costo lineare (o «a tempo lineare»): il lavoro cresce di pari passo con la lunghezza, non più in fretta di lei.
La risposta arriva nell’autunno del 2021, quando Albert Gu, Karan Goel e Christopher Ré trovano il modo di far girare quelle equazioni vecchie di sessant’anni dentro uno strato di rete neurale a un costo sostenibile, e le mettono alla prova sul Long Range Arena, il banco di prova delle dipendenze a lunghissimo raggio: i legami fra parti lontane di una sequenza (in un giallo, per capire l’ultima pagina bisogna ricordare il nome che compariva alla prima). Il loro modello, S4 [GGRe22], riesce là dove Transformer e reti ricorrenti si arrendevano, ed è il primo a risolvere il compito di quel banco in cui il legame da riconoscere è lungo sedicimila passi, dove un «passo» è un elemento della sequenza (una parola, un campione audio, un pixel). È l’atto di nascita di una seconda strada verso il modello di sequenze a costo lineare: non quella dell’attenzione resa economica del capitolo precedente, ma quella, apparentemente lontana, dei sistemi dinamici. Alla fine, scopriremo, le due strade portano allo stesso posto.
Un sistema che riassume il passato#
Segui il livello dell’acqua in una vasca mentre entra ed esce di continuo. Non ti serve ricordare ogni singola goccia, ti basta un numero, quanta acqua c’è adesso. Quel numero si porta dietro tutta la storia, ed è lo stato.
Le cose che lo muovono restano sempre le stesse. Dal rubinetto entra acqua e il livello sale. Dallo scarico socchiuso ne esce, e più acqua c’è più in fretta cala, così di quello che c’era ogni minuto ne resta una frazione. Sul fianco della vasca un galleggiante muove un ago su una scala graduata, e quello che l’ago segna è la tua risposta a chi chiede. Che cosa entra, quanto resta di ciò che c’era, che cosa se ne legge, e la regola non ha altri pezzi.
Metti che dal rubinetto entrino due litri al minuto e che lo scarico porti via ogni minuto metà dell’acqua che trova. Parti da vasca vuota e segui l’ago: zero, due, tre (metà di due, più i due che entrano), 3,5, 3,75, 3,875. Il livello sale sempre più piano e si assesta sui quattro litri. Ogni getto intanto sbiadisce, e di quello entrato cinque minuti fa resta un trentaduesimo, di quello di mezz’ora fa meno di un miliardesimo.
Resta da scegliere ogni quanto guardare. L’acqua scorre senza interruzione, tu l’ago lo segni una volta al minuto, e di quel che è successo fra una segnatura e l’altra hai soltanto quello che riesci a ricostruire. Segnando fitto ti sfugge poco di ciò che è entrato. Segnando di rado può passare in mezzo un getto intero che non hai visto, e quello che ricostruisci esce grossolano.
Uno state space model fa questo con una sequenza. Dove la vasca ha un livello solo, il modello ne tiene molte, ciascuna con il suo rubinetto e il suo scarico, e a ogni parola in arrivo si aggiornano tutte insieme, ognuna al ritmo che le tocca (e la parola, nel modello, è una fila di numeri, non una parola scritta); quante siano si decide prima di aprire il rubinetto e non cambia più, per lungo che sia il testo. È lo stesso spirito della rete ricorrente, con rubinetto, scarico e ago presi dai sistemi che evolvono nel tempo.
Proprio perché la regola non cambia mai, a quel 3,875 ci si arriva per due strade. Passo dopo passo, dal livello di prima a quello di adesso, una parola alla volta. Oppure tutto insieme, sommando i getti entrati fin qui, ciascuno sbiadito secondo quanto tempo fa è entrato: 2 + 1 + 0,5 + 0,25 + 0,125 fa lo stesso numero. Quella fila di sbiadimenti è un filtro che si fa scorrere in un colpo solo sull’intera sequenza. Sono la stessa identica cosa vista da due lati, ed è la doppia natura. Si addestra il modello nel secondo modo, veloce perché fa tutti i conti in una volta, e lo si usa nel primo, economico perché a ogni parola gli basta il riassunto di prima.
Tutto questo sta in piedi finché nessuno tocca il rubinetto e lo scarico. Se qualcuno stesse alla vasca a girarli minuto per minuto, regolandoli in base all’acqua in arrivo, non ci sarebbe più una sola fila di sbiadimenti buona per l’intera storia, e quella non si potrebbe più far scorrere in un colpo solo. Modi di fare i conti tutti insieme ne restano, ma vanno ritrovati da capo, ed è metà del lavoro che il capitolo racconta.
Il mattone è un sistema lineare a tempo continuo che mappa un ingresso \(u(t)\) in un’uscita \(y(t)\) attraverso uno stato latente \(\mathbf{h}(t)\):
La matrice \(\mathbf{A}\) governa la dinamica interna (come lo stato evolve da solo), \(\mathbf{B}\) come l’ingresso vi entra, \(\mathbf{C}\) come se ne legge l’uscita. Per usarlo su una sequenza discreta lo si discretizza con un passo \(\Delta\), ottenendo una ricorrenza \(\mathbf{h}_t = \bar{\mathbf{A}}\, \mathbf{h}_{t-1} + \bar{\mathbf{B}}\, x_t\), dove \(x_t\) è l’ingresso campionato al passo \(t\) e \(\bar{\mathbf{A}}, \bar{\mathbf{B}}\) sono le versioni discrete di \(\mathbf{A}\) e \(\mathbf{B}\). E qui sta la ricchezza: finché i parametri sono costanti nel tempo, questa ricorrenza ha una doppia natura; si può calcolare passo per passo come una RNN (inferenza a costo costante per token) oppure, a stato iniziale nullo, tutta in una volta come una convoluzione (addestramento parallelo). È la stessa doppia natura parallelo/ricorrente che muove il capitolo sull’attenzione lineare, raggiunta però dalla teoria dei segnali.
Due strade, una meta#
L’attenzione lineare del capitolo precedente e gli state space model di questo nascono da mondi diversi, e arrivano alla stessa macchina. È quella appena descritta: tiene un riassunto di taglia sempre uguale (lo stato) e a ogni parola lo aggiorna con una regola semplice, in cui il nuovo riassunto è il vecchio, un po’ sbiadito, più ciò che entra adesso. Si addestra lavorando su tutta la sequenza in una volta sola, e poi genera una parola alla volta senza che la memoria cresca mai.
Il nome per esteso di quella macchina è rete ricorrente lineare a stato di dimensione fissa, ed è quello con cui il capitolo precedente l’ha già chiamata: «ricorrente» perché ogni passo riparte dal risultato del passo precedente, «lineare» perché il riassunto nuovo si ricava dal vecchio con sole moltiplicazioni e somme (è un altro mestiere della stessa parola: qui non dice quanto costa il conto, dice come è fatto), «a stato di dimensione fissa» perché il riassunto non si allarga mai. L’attenzione lineare ci arriva dal meccanismo di attenzione, gli state space model dai sistemi dinamici; la macchina, alla fine, è la stessa.
Fig. 17.1 Due modi di portarsi dietro il passato. L’attenzione lo tiene tutto e lo riguarda; la ricorrenza lo riassume in uno stato di taglia fissa e ci scrive sopra, decidendo di volta in volta che cosa vale la pena scrivere.#
Il confronto di Fig. 17.1 mostra anche dove sta il prezzo. Un riassunto di taglia fissa deve, prima o poi, dimenticare qualcosa. Il primo dei due fili che attraversano il capitolo è la selettività (il rombo sul lato destro della figura): decidere cosa scrivere nello stato, e cosa lasciar cadere, in funzione di ciò che sta arrivando. È facile chiedere troppo alla selettività: cambia come si usa lo spazio del riassunto, non lo allarga. Il tetto di un riassunto di taglia fissa resta, ed è l’argomento di Panorama e limiti.
Il secondo filo è una parentela. Alla fine, con Mamba-2 [DG24] (che nel capitolo precedente era una riga della tabella delle ricorrenze), vedremo che non è una somiglianza vaga: un state space model di forma opportuna è un’attenzione mascherata, cioè un’attenzione che guarda solo all’indietro, in cui il confronto fra due parole è pesato da quanto della prima è sopravvissuto nel frattempo. Le due famiglie che raccontiamo in due capitoli sono, in fondo, due viste dello stesso disegno.
Ma prima c’è una tensione da sciogliere. La forma a convoluzione, cioè il filtro unico che si fa scorrere sull’intera sequenza, vale solo se il sistema è invariante nel tempo: le stesse regole a ogni passo. Ed è proprio questa rigidità che rompe Mamba, il modello con cui Albert Gu e Tri Dao chiudono il 2023 [GD24]. Mamba rende il sistema selettivo, cioè capace di decidere in base al contenuto che cosa ricordare e che cosa dimenticare, e in cambio perde il filtro unico. Il modo «tutto insieme» non se ne va con lui, ma va ricostruito su un’altra strada.
Dai sistemi dinamici a Mamba#
Quattro tappe, dall’idea di base alla frontiera.
Dai sistemi dinamici a S4: che cos’è una macchina che riassume il passato in un pugno di numeri, come si adatta a una sequenza fatta di passi separati, e come si fa a darle una memoria lunga (sono HiPPO e S4).
Mamba: come si insegna alla macchina a scegliere, invece di trattare tutte le parole allo stesso modo; che cosa costa quella scelta (si perde la convoluzione), e con quali due mosse si recupera la velocità perduta: una proprietà della ricorrenza che permette di raggruppare i passi a piacere, e la forma di una scheda grafica.
La dualità: la scoperta che questa macchina, scritta in un altro modo, è un’attenzione mascherata, cioè la formula dei Transformer con la softmax tolta di mezzo, e che riscriverla così la fa girare molto più in fretta. Poi le tre messe a punto di Mamba-3.
Panorama e limiti: una mappa che tiene insieme questo capitolo e il precedente, dove un riassunto di taglia fissa diventa un collo di bottiglia, e le architetture ibride che mettono insieme il meglio delle due strade.
Da ricordare
Un modello a spazio degli stati (in sigla SSM) riassume tutto quello che ha letto in un riassunto di dimensione sempre uguale, e a ogni parola lo aggiorna. Sono le stesse equazioni con cui l’ingegneria descrive un termostato o la traiettoria di un razzo; S4 [GGRe22] trova il modo di farle girare dentro una rete neurale a un costo sostenibile, ed è il primo a risolvere il compito del Long Range Arena in cui il legame da riconoscere è lungo sedicimila passi.
Il problema che vengono a risolvere: far guardare ogni parola a tutte le altre costa al quadrato (testo doppio, lavoro quadruplo). Qui il costo cresce di pari passo con la lunghezza, ed è ciò che nel libro si chiama costo lineare.
Finché le regole non cambiano da un passo all’altro, lo stesso calcolo si può fare in due modi: passo dopo passo (economico per generare) oppure con un filtro solo, fatto scorrere sull’intera sequenza (veloce per addestrare). È la doppia natura, la stessa già vista con l’attenzione lineare.
Mamba [GD24], il modello che Albert Gu e Tri Dao presentano alla fine del 2023, rompe quella regola fissa: lascia decidere alla parola in arrivo quanto scrivere e quanto dimenticare (è la selettività), e in cambio rinuncia al filtro unico. Mamba-2 [DG24] mostra poi che, nella sua versione più semplice, questa macchina è un’attenzione che guarda solo all’indietro: le due famiglie si incontrano su quel gradino.
Il percorso: dai sistemi dinamici a S4 → Mamba (scegliere, e restare veloci) → la dualità (Mamba-2 e Mamba-3) → panorama, limiti e ibridi.
Da ricordare
Uno state space model riassume il passato in uno stato di dimensione fissa, con equazioni che l’ingegneria usa da decenni per i sistemi dinamici; S4 [GGRe22] ne trova la parametrizzazione che le rende calcolabili, e con essa risolve per primo Path-X, il compito del Long Range Arena a \(16\,384\) passi su cui tutti i lavori precedenti fallivano.
Discretizzato, un SSM invariante nel tempo ha una doppia natura: ricorrente (inferenza a costo costante per token) e convoluzionale, a stato iniziale nullo (addestramento parallelo). È la stessa doppia natura dell’attenzione lineare, da un’altra strada.
Mamba [GD24] rompe l’invarianza temporale con la selettività, e con essa la forma convoluzionale; il parallelismo si riconquista per un’altra via, lo scan. Mamba-2 [DG24] mostra che un SSM di forma opportuna è un’attenzione mascherata: le due famiglie si incontrano su quel gradino.
Il percorso: dai sistemi dinamici a S4 → Mamba (selezione e scan) → la dualità (Mamba-2 e Mamba-3) → panorama, limiti e ibridi.