Dalla softmax alla ricorrenza#
Il capitolo sui Transformer si chiude con un paradosso. Il meccanismo che ha vinto la partita (l’attenzione, che lascia ogni parola libera di guardare tutte le altre) vince proprio perché guarda tutte le altre; ed è esattamente questo il suo conto da pagare. I conti sono due. Il primo è il lavoro, che cresce col quadrato della lunghezza (gli informatici lo scrivono \(O(n^2)\), che è solo un modo compatto di dirlo). Il secondo è l’archivio di appunti che il modello si tiene mentre scrive, con dentro le etichette e le informazioni già calcolate per non rifare ogni volta gli stessi conti: quell’archivio si allunga a ogni parola generata, si chiama KV cache, e la sezione sull’attenzione in pratica lo smonta pezzo per pezzo.
Le reti ricorrenti che abbiamo studiato prima dei Transformer non avevano nessuno di questi due problemi: leggevano in fila, a costo lineare, con uno stato di dimensione fissa che non cresceva mai. Le avevamo abbandonate per un difetto altrettanto grave, la sequenzialità: ogni passo deve aspettare quello prima, e le schede grafiche (le GPU), che sono fatte per macinare montagne di conti tutti insieme, restano a guardare. La domanda è se si possano avere entrambe le cose: distribuire il lavoro su tanti conti simultanei quando il modello impara, e pagare poco, sempre lo stesso, quando scrive. La risposta parte da un’osservazione che si fa alle medie, e la vediamo subito.
Il trucco del kernel: spezzare la softmax#
L’ostacolo ha un nome preciso, ed è la softmax. Per calcolare l’attenzione bisogna prima misurare quanto ogni parola somiglia a ogni altra, e riempire con quei numeri una tabella che ha una riga e una colonna per ciascuna parola: mille parole, un milione di caselle. La tabella grande serve perché la softmax mescola, cioè spartisce fra tutte: quanto la parola numero uno dà retta alla numero due dipende anche da tutte le altre, e finché è così nessuna casella si può calcolare per conto suo.
La via d’uscita è cambiare il modo di misurare la somiglianza. Serve una misura che, invece di nascere dal confronto di una coppia, si spezzi in due pezzi indipendenti: uno che riguarda solo chi fa la domanda, uno che riguarda solo chi risponde. (Ogni parola fa tutte e due le parti: interroga le altre con la propria domanda, e alle domande delle altre risponde con la propria etichetta.) Se la somiglianza si spezza così, i conti si possono riordinare, e riordinarli li fa crollare di numero.
Mille persone in una sala, una per ogni parola della frase, e ognuna ha qualcosa da dire alle altre. Farle parlare tutte con tutte sono \(499\,500\) conversazioni, contando una volta sola ogni coppia; la tabella per segnarle ha un milione di caselle, perché ogni persona ha la sua riga e la sua colonna, sé stessa compresa. Il costo esplode perché quanto io do retta a te nasce dall’incontro fra me e te: quel numero non esiste finché non ci siamo visti.
Cambiamo la regola. All’ingresso ognuno riceve un cartellino con un numero, che è la sua etichetta ridotta a un numero solo, e da lì in poi quanto io do retta a te è il mio numero per il tuo. Il peso si spezza in due, un pezzo mio e uno tuo.
Tre persone hanno da consegnare \(2\), \(5\) e \(9\). Prima chi voleva farsi un’idea doveva incontrarle tutte e tre, e ne usciva con \(1\times 2 + 6\times 5 + 8\times 9\): tre pesi diversi, tre incontri, tre caselle della tabella grande. Adesso invece nel numero che ciascuno consegna il proprio cartellino è già conteggiato, e a chi ascolta resta soltanto il numero suo, sempre quello: \(3\times 2 + 3\times 5 + 3\times 9 = 3\times(2+5+9) = 3\times 16\). È il raccoglimento a fattor comune di seconda media. E quel \(16\) non dipende da chi ascolta: la persona dopo farà \(7\times 16\), quella dopo ancora \(4\times 16\).
Allora il \(16\) si calcola una volta sola. All’ingresso c’è un registro: chi entra ci somma il proprio contributo, sempre nelle stesse caselle, e chi vuole farsi un’idea non gira più per la sala, legge il totale e lo moltiplica per il proprio cartellino. Le persone restano mille, il registro è uno solo, e lo si aggiorna una volta per ciascuna: mille aggiornamenti invece di mezzo milione di conversazioni. Scrivere sul registro non è gratis, costa quante sono le sue caselle. Ma quel costo resta lo stesso man mano che la sala si riempie: raddoppia la gente e raddoppiano gli aggiornamenti, mentre le conversazioni diventerebbero quattro volte tante.
Il numero sul cartellino può essere piccolissimo, mai negativo: se no chi entra cancellerebbe dal registro quello che gli altri ci hanno scritto, invece di aggiungercisi. Accanto al registro sta poi un secondo foglio, molto più corto, dove si segna quanto pesano in tutto le risposte raccolte: alla fine si divide per quel totale, così quello che esce è una media dei contributi e non una somma che si gonfia man mano che entra gente. Chi si presenta col cartellino azzerato trova un totale nullo, non ha niente per cui dividere, ed è lì che il meccanismo si pianta.
Un modo di misurare la somiglianza che si spezza in due pezzi, in matematica, si chiama kernel.
Nell’attenzione dei Transformer [VSP+17] l’uscita (non normalizzata) per la query \(i\) è una somma dei value pesati dalla somiglianza esponenziale \(\exp(\mathbf{q}_i^\top \mathbf{k}_j)\) (il fattore di scala \(1/\sqrt{d}\) lo consideriamo assorbito in query e chiavi). Il guaio è che quell’esponenziale non si spezza: non esiste una fattorizzazione esatta a dimensione finita in un prodotto di una funzione della sola \(\mathbf{q}_i\) per una funzione della sola \(\mathbf{k}_j\). Una feature map che lo riproduca c’è, ma con infinite componenti, e in pratica la si può solo approssimare: è la strada delle random features del Performer [CLD+21]. Restando esatti, quindi, l’esponenziale va valutato per ogni coppia \((i,j)\): la matrice \(n \times n\).
L’idea di Katharopoulos e colleghi [KVPF20] è sostituire la somiglianza con una che si spezza, cioè un prodotto scalare fra versioni trasformate di query e key:
dove \(\phi\) è una feature map applicata a ciascun vettore, riga per riga sulle matrici di query e chiavi. In generale è una \(\phi: \mathbb{R}^d \to \mathbb{R}^C\), che può cambiare dimensione: \(C\) è la dimensione dello spazio di feature ed è un parametro libero (il Performer, per dire, ne sceglie 256 di random features), e da essa dipende la taglia della memoria. Con questa scelta l’uscita per la query \(i\) diventa
dove i passaggi sono tre, e conviene separarli perché è su questi che si regge il capitolo: il prodotto scalare è simmetrico, quindi \(\phi(\mathbf{q}_i)^\top \phi(\mathbf{k}_j) = \phi(\mathbf{k}_j)^\top \phi(\mathbf{q}_i)\); uno scalare si può spostare a destra del vettore \(\mathbf{v}_j\); e a quel punto è l’associatività del prodotto a permettere di raccogliere \(\phi(\mathbf{q}_i)\) fuori dalla somma su \(j\), che si stacca dalla query e si condensa in un’unica matrice \(\mathbf{S} = \sum_j \mathbf{v}_j\, \phi(\mathbf{k}_j)^\top \in \mathbb{R}^{d_v\times C}\) (il «registro» chiave→valore). Qui \(\mathbf{q}_i, \mathbf{k}_j, \mathbf{v}_j\) sono i vettori query, key e value del token e \(d_v\) è la dimensione del value. Da qui in avanti, per semplicità, assumiamo key, query e value della stessa dimensione \(d\) e una \(\phi\) che non cambia dimensione (\(C = d\)), così che lo stato sia una \(d \times d\): è il caso della feature map che sceglieremo fra poco, ma non il caso generale.
Le due strade hanno costi diversissimi. Calcolare tutti i prodotti \(\phi(\mathbf{q}_i)^\top \phi(\mathbf{k}_j)\) è la matrice \(n \times n\), costo \(O(n^2 d)\); costruire \(\mathbf{S}\) una volta e applicarla a ogni query costa \(O(n d^2)\): una matrice \(d \times d\) al posto di una \(n \times n\). Quando \(n \gg d\) (sequenze lunghe), la seconda vince nettamente, ed è il passaggio da \(O(n^2 d)\) a \(O(n d^2)\), cioè da quadratico a lineare nella lunghezza (lo stesso \(O(n d^2)\) delle ricorrenti che avevamo incontrato nel confronto fra Transformer e RNN). In forma matriciale compatta, per tutte le query insieme, è l’identità
la stessa \(\text{softmax}(\mathbf{Q}\mathbf{K}^\top)\mathbf{V}\) del capitolo sui Transformer con la softmax tolta di mezzo: senza di lei il prodotto si può ri-associare, e conviene fare prima \(\phi(\mathbf{K})^\top \mathbf{V}\), la matrice piccola. Un’avvertenza sulla convenzione, che in un capitolo dedicato a da che parte moltiplica un fattore non è pedanteria: qui i token stanno sulle righe di \(\mathbf{Q}\), \(\mathbf{K}\) e \(\mathbf{V}\), quindi la matrice piccola \(\phi(\mathbf{K})^\top \mathbf{V}\) è \(d \times d_v\), cioè \(\mathbf{S}^\top\) e non \(\mathbf{S}\).
Resta da scegliere \(\phi\). Serve una feature map che dia somiglianze positive (perché i pesi si comportino come quelli di una media), e Katharopoulos et al. propongono la più semplice che funzioni:
applicata componente per componente e sempre maggiore di zero (in aritmetica
esatta: per \(x<0\) vale \(\exp(x)\), ma il calcolo passa da \(-1 + \exp(x)\), e in
float32 quella somma arrotonda a \(-1\) appena \(\exp(x)\) scende sotto
\(2^{-25}\), cioè da \(x = -25\ln 2 \approx -17{,}33\) in giù; da lì \(\phi\) è zero esatto, e un
denominatore che si azzera è il modo tipico in cui questa implementazione si
rompe). La softmax aveva anche un denominatore che
normalizzava i pesi: lo si conserva come un secondo accumulatore
\(\mathbf{z} = \sum_j \phi(\mathbf{k}_j)\), e la lettura normalizzata diventa
\(\mathbf{o}_i = \mathbf{S}\,\phi(\mathbf{q}_i) \,/\, \big(\mathbf{z}^\top \phi(\mathbf{q}_i)\big)\).
L’attenzione lineare è una RNN#
Fin qui abbiamo ragionato come se tutte le parole fossero disponibili insieme. Ma quando il modello scrive vale una regola in più: ogni parola può guardare solo quelle che la precedono, non quelle che verranno, e questo si dice causale. Allora il registro non è più uno solo, compilato alla fine: ce n’è una versione a ogni passo, e ogni parola nuova aggiunge il suo contributo a quello che c’è già scritto. Attenzione a non fraintendere: a cambiare sono i numeri scritti nelle caselle, non il numero di caselle, che resta quello.
Un registro che si aggiorna così, un passo alla volta e a partire da come era al passo prima, ha un nome: è una ricorrenza. Ed è il nome che portavano le reti che i Transformer avevano mandato in pensione.
A una riunione lunga si prendono appunti in due modi.
Il primo è la KV cache dei Transformer: ogni volta che qualcuno parla, aggiungi uno scontrino alla pila. Non butti via niente, e questo è comodo (hai tutto) ma la pila cresce, e a fine giornata occupa mezzo tavolo. Ogni parola nuova ne aggiunge un’altra.
Il secondo è un foglio-registro di dimensione fissa. Non aggiungi scontrini: aggiorni le stesse caselle. Quando qualcuno parla, sommi il suo contributo a ciò che c’è già scritto, e il foglio resta un foglio: sempre lo stesso, che tu sia alla decima o alla decimillesima parola. Per rispondere a una domanda leggi il foglio, non rovisti nella pila.
L’attenzione lineare tiene esattamente questo foglio. Ecco perché la memoria non cresce: qualunque sia la lunghezza del testo, il foglio (negli articoli si chiama stato) ha sempre le stesse dimensioni. È il ritorno, sotto mentite spoglie, della vecchia idea delle reti ricorrenti.
Basta riscrivere \(\mathbf{S}\) come somma cumulativa fino al passo \(t\):
e la lettura al passo \(t\) usa lo stato corrente:
Qui \(\mathbf{S}_t \in \mathbb{R}^{d\times d}\) è lo stato, una memoria chiave→valore che a ogni passo incassa il prodotto esterno \(\mathbf{v}_t\, \phi(\mathbf{k}_t)^\top\) (scrivi il value \(\mathbf{v}_t\) sotto l’etichetta \(\phi(\mathbf{k}_t)\)); \(\mathbf{z}_t\) è il normalizzatore che accumula le key trasformate; \(\mathbf{o}_t\) è l’uscita. Leggere con \(\phi(\mathbf{q}_t)\) significa \(\mathbf{S}_t \phi(\mathbf{q}_t) = \sum_{i\le t} \big(\phi(\mathbf{k}_i)^\top\phi(\mathbf{q}_t)\big)\, \mathbf{v}_i\): la query ripesca dai value in proporzione a quanto la sua etichetta somiglia a ciascuna key già scritta.
Una cosa sul modo di scrivere il meccanismo, che non lo cambia di una virgola: Katharopoulos tiene il normalizzatore \(\mathbf{z}_t\), e la lettura è una media pesata di cui \(\mathbf{z}_t^\top \phi(\mathbf{q}_t)\) è il denominatore. Già i lavori sui fast weight di poco successivi lo abbandonano, giudicandolo instabile (quell’accumulatore può crescere senza controllo), e normalizzano invece chiavi e query trasformate; le varianti più recenti vi rinunciano del tutto, aggiungendo una layer normalization in uscita e, dove la transizione lo richiede (DeltaNet), riportando le key a norma unitaria. Sono tre impostazioni diverse, che è meglio non mescolare: nelle prossime sezioni le teniamo distinte e, parlando delle architetture moderne, usiamo l’ultima, cioè feature map \(\phi\) posta all’identità e nessun \(\mathbf{z}_t\) nelle formule.
Guardiamo bene questa ricorrenza. È esattamente una RNN, ma con due differenze rispetto alle celle dei modelli di sequenza. La prima: lo stato non è un vettore \(\mathbf{h}_t\) ma una matrice \(\mathbf{S}_t\) (una memoria molto più capiente). La seconda, decisiva: la transizione di stato è lineare, anzi è l’identità (\(\mathbf{S}_{t-1}\) passa intatto, gli si somma soltanto un termine nuovo). Non c’è nessuna \(\tanh\) o non-linearità sullo stato, come invece in \(\mathbf{h}_t = \tanh(\mathbf{W}_{hh} \mathbf{h}_{t-1} + \dots)\). L’aggiornamento costa \(O(d^2)\) per token e la memoria è costante: la matrice \(d \times d\) non cambia dimensione, che siamo al decimo o al milionesimo token.
È il senso, volutamente ironico, del titolo del paper di Katharopoulos et al., Transformers are RNNs [KVPF20]: sotto una certa scelta della somiglianza, un Transformer è una rete ricorrente; solo che se n’era dimenticato.
Fig. 16.2 Due modi di ricordare. A sinistra la memoria dei Transformer, che cresce di una coppia etichetta-informazione a ogni parola. A destra l’attenzione lineare: un’unica tabella di numeri, sempre della stessa taglia, in cui ogni parola somma la propria informazione sotto la propria etichetta, e che si rilegge facendole una domanda.#
Come mostra Fig. 16.2, i due schemi raccolgono la stessa informazione in modi opposti: la KV cache la conserva tutta e paga in memoria che cresce; il registro la comprime in un foglio di taglia fissa.
Addestrare in parallelo, generare in ricorrenza#
Fin qui il registro lo abbiamo visto riempirsi una parola alla volta. Ma quando il testo c’è già tutto non c’è nessun bisogno di aspettare: lo stesso identico risultato si può ottenere anche in un colpo solo, spartendo il lavoro fra tante unità di calcolo che macinano insieme. Sono due volti dello stesso calcolo. Non è una coincidenza tecnica: è la proprietà che rende interessante tutta questa famiglia di modelli, e che ritroveremo (con parole diverse) nel capitolo sugli State Space Model, i modelli che descrivono una sequenza come un sistema che evolve nel tempo.
Fig. 16.3 La memoria che si aggiorna parola per parola: a ogni passo i numeri nelle caselle cambiano, perché ci si somma sopra il contributo della parola appena letta, ma le caselle restano quelle. Alla quinta parola la tabella è grande esattamente come alla prima.#
La Fig. 16.3 mostra la parte «economica» del patto: la memoria non cresce. Ma mostra anche, senza dirlo, il prezzo: dentro ogni cella i contributi di parole diverse finiscono sommati fra loro, e una volta sommati non si possono più separare. È il tema di Il limite dell’accumulo.
Studiare un libro che hai già tutto in mano è un mestiere; raccontare a voce una storia che stai inventando adesso è un altro.
Nel primo caso (è l’addestramento, quando il modello impara) il testo esiste già per intero: puoi aprirlo a metà, dare un capitolo a testa a dieci persone e finire in un decimo del tempo. È la forma «tutta insieme», quella che tiene occupata tutta la scheda grafica.
Nel secondo caso (è la generazione, quando il modello scrive) procedi parola per parola, perché il seguito non esiste ancora: lo stai inventando. Qui l’unica cosa che ti porti dietro è il foglio-registro, che aggiorni a ogni parola e che non cresce mai.
Le due forme danno lo stesso risultato, e allora si usa ciascuna dove conviene: il modello si allena con quella che sfrutta tutte le unità di calcolo insieme, e scrive con quella che occupa sempre la stessa memoria. La KV cache dei Transformer, al contrario, obbliga a trascinare una pila che si allunga a ogni parola generata.
Un’avvertenza, perché il conto non sembri troppo bello: la forma «tutta insieme», fatta nel modo più ovvio, si ritrova per le mani la tabella grande di prima.
C’è una seconda strada, e sembra risolvere tutto: dare a ognuna delle dieci persone il registro già aggiornato fino al punto in cui comincia il suo capitolo, così nessuno aspetta nessuno. Quei registri intermedi però bisogna prima costruirli, e il modo di costruirli senza mettersi in fila ne produce uno per ogni parola: con un testo di ottomila parole sono ottomila fogli invece di uno. I conti da fare restano proporzionali alla lunghezza, ma il tempo che si perde a scrivere, tenere e rileggere quella pila si mangia il guadagno di lavorare in dieci.
La strada che si prende davvero sta in mezzo, e si lavora a blocchi: si taglia il testo in pezzi, e dentro ogni pezzo si fa tutto insieme, dove la tabella è piccola perché le parole sono poche; poi il registro passa da un pezzo al successivo. In fila vanno i pezzi, che sono pochi e grossi; il grosso del lavoro, quello dentro ai pezzi, resta in parallelo. Il risultato non cambia, e i blocchi torneranno in tutto il resto del capitolo.
E c’è una seconda avvertenza, sulle promesse: meno conti non vuol dire subito più veloce. Su un testo corto un’attenzione dei Transformer scritta con cura vince ancora, perché è stata limata per anni; il guadagno si raccoglie quando il testo si allunga, ed è lì che il conto proporzionale stacca quello a valanga.
In addestramento si usa la forma parallela, e qui serve un momento di onestà sui costi. La strada più diretta è il prodotto fra matrici mascherato dalla causalità, \(\big(\phi(\mathbf{Q})\,\phi(\mathbf{K})^\top \odot \mathbf{M}\big)\mathbf{V}\) con \(\mathbf{M}\) la maschera triangolare: parallelo esattamente come un Transformer, ma la maschera impedisce di ri-associare il prodotto e si torna a pagare \(O(n^2 d)\). L’alternativa è srotolare la somma cumulativa \(\mathbf{S}_t = \sum_{i\le t} \mathbf{v}_i\,\phi(\mathbf{k}_i)^\top\) come prefix sum (una somma progressiva): l’operazione è associativa, quindi si calcola con uno scan parallelo a costo \(O(n d^2)\), lineare. L’ostacolo qui non è il numero di operazioni ma la memoria: uno scan pretende di materializzare tutti gli \(n\) stati intermedi \(d \times d\), cioè \(O(n d^2)\) di memoria contro gli \(O(d^2)\) della forma ricorrente, e il traffico da e verso la memoria della GPU si mangia il guadagno del parallelismo (con \(n = 8192\) e \(d = 64\) per testa sono più di 33 milioni di valori per testa e per strato, contro i 4096 dello stato). Nessuna delle due forme dà insieme le due cose; la conciliazione usata in pratica è il calcolo a blocchi (chunkwise): parallelo dentro ogni blocco, ricorrente fra un blocco e l’altro, costo \(O(nBd + nd^2)\) con blocchi di ampiezza \(B\), cioè lineare in \(n\). Lo ritroveremo, formalizzato, in RetNet e DeltaNet.
Un’ultima onestà sul lato addestramento, perché «lineare» non vuol dire «subito più veloce»: l’attenzione softmax ha implementazioni curatissime nel traffico di memoria, e sotto qualche migliaio di token un calcolo a blocchi scritto in modo ingenuo perde il confronto. Il vantaggio si raccoglie quando il contesto cresce, ed è il motivo per cui buona parte del lavoro su queste architetture è lavoro di implementazione, non di formule.
In inferenza autoregressiva si usa la forma ricorrente: si aggiorna \(\mathbf{S}_t\) sul posto e si legge \(\mathbf{o}_t\), con costo \(O(d^2)\) per token e memoria \(O(d^2)\) costante. Nessuna KV cache che si allunga: lo stato è sempre la stessa matrice \(d \times d\). In un Transformer, ricordiamo, la cache cresce di una coppia \((\mathbf{k}_t, \mathbf{v}_t)\) per token e per strato, e il costo di generare l’\(n\)-esimo token sale con la lunghezza del prefisso; qui resta piatto.
È da questo contrasto che nasce l’accelerazione più spettacolare riportata da Katharopoulos et al.: fino a circa quattromila volte più veloce nella generazione autoregressiva di sequenze molto lunghe. Va letta con onestà (è un caso limite, riguarda solo l’inferenza e non l’addestramento) ma misura bene il vantaggio della memoria costante quando \(n\) diventa enorme.
Il passo con cui il modello genera una parola sta in poche righe, e la cosa che
conta si legge a colpo d’occhio: S entra ed esce dalla funzione sempre della
stessa taglia, e in tutto il codice non c’è nessuna lista che si allunga.
import torch
import torch.nn.functional as F
phi = lambda x: F.elu(x) + 1.0 # feature map: sempre positiva
def genera_passo(S, z, q_t, k_t, v_t):
"""Un passo di attenzione lineare, forma ricorrente.
S: stato d x d (fisso) z: normalizzatore d q_t, k_t, v_t: vettori d."""
pk = phi(k_t)
S = S + torch.outer(v_t, pk) # scrivi: S += v_t phi(k_t)^T
z = z + pk # aggiorna il normalizzatore
pq = phi(q_t)
o_t = (S @ pq) / (z @ pq) # leggi: o_t = S phi(q_t) / (z . phi(q_t))
return o_t, S, z
La memoria occupata da quella tabella non dipende da quanti token abbiamo già
generato: è il cuore del vantaggio. (Accanto a S viaggia un secondo
accumulatore, z, molto più piccolo, che serve solo a tenere i numeri in
scala; anche lui è di taglia fissa.)
Il limite dell’accumulo#
Tanta eleganza ha un prezzo, ed è bene dirlo subito e senza sconti. Il registro dell’attenzione lineare sa fare una cosa sola: sommare. Non cancella e non corregge. Ogni parola che passa lascia la sua traccia sul foglio, sommata a tutte le altre, e lì resta per sempre.
Un foglio-registro su cui continui a sommare senza mai cancellare niente, prima o poi diventa illeggibile: le scritte si sovrappongono, e quando cerchi un’informazione precisa ti ritrovi un pasticcio di tracce che si confondono.
Se due parole diverse hanno etichette simili, i loro contributi si mescolano, e rileggendo non capisci più bene quale informazione appartenesse a quale etichetta. Il disturbo non comincia di colpo a una certa soglia: comincia subito, piano, e cresce con quante cose hai scritto.
E il foglio si riempie molto prima di quanto lascino sperare le sue caselle, perché un’informazione non occupa una casella: quando la scrivi si spalma su tutto il foglio, e la successiva si spalma sopra di lei. Prendi un foglio piccolo, trentadue righe per trentadue colonne (nei modelli veri il lato è di sessantaquattro o centoventotto): di caselle ne ha più di mille, ma le etichette che riesce a tenere separate sono trentadue, cioè il suo lato, tante quante ne stanno una accanto all’altra senza pestarsi i piedi.
E trentadue è già il caso migliore, quello che nella pratica non capita: le etichette non le sceglie nessuno apposta distinte, le assegna il modello, e due parole qualunque finiscono per somigliarsi un po’. Infatti già a otto informazioni la risposta torna sbagliata di circa metà del suo valore, e a trentadue lo sbaglio è grande quanto la risposta. Da lì in poi ritrovare il dettaglio giusto («di che colore era il cappotto citato venti pagine fa?») è impossibile. L’attenzione dei Transformer, quella che tiene tutti gli scontrini, quel dettaglio ce l’ha ancora; il registro riassuntivo può averlo perso.
E la domanda ovvia (perché non prendersi un foglio più grande?) ha una risposta altrettanto ovvia: si può, ed è una delle manopole di chi progetta il modello, ma si paga. Un foglio più largo vuol dire più caselle da aggiornare e da rileggere a ogni parola, quindi più conti e più memoria: allargandolo abbastanza si torna a spendere quanto un Transformer, e il vantaggio che eravamo venuti a cercare svanisce. Il gioco di tutto il resto del capitolo è un altro: tenere il foglio piccolo e imparare a scriverci meglio.
Il limite è di capacità, ed è una conseguenza della dimensione finita. Lo stato \(\mathbf{S}\) è una matrice \(d \times d\): in uno spazio di dimensione \(d\) non esistono più di \(d\) vettori mutuamente ortogonali. Se le key \(\phi(\mathbf{k}_1), \dots\) fossero esattamente ortogonali e di norma uno, leggere con \(\phi(\mathbf{q})\) recupererebbe il value giusto pulito fino a \(d\) associazioni; ma le key le produce una proiezione lineare, non un’ortogonalizzazione, e con chiavi casuali il crosstalk (le briciole degli altri value che il retrieval raccoglie insieme a quello cercato) non compare a una soglia: cresce da subito come \(\sqrt{N/d}\) con il numero \(N\) di associazioni scritte. A \(N \approx d\) l’interferenza vale ormai quanto il valore cercato. In \(d=32\), con chiavi gaussiane riportate a norma unitaria e value gaussiani, l’errore relativo medio del richiamo (media su tutte le chiavi scritte e su duemila estrazioni) è \(0{,}99\) a \(N=d\) e \(0{,}46\) a \(N=d/4\), in accordo con l’andamento atteso \(\sqrt{(N-1)/d}\), cioè \(\sqrt{N/d}\) a meno del contributo della chiave che si sta interrogando. Il richiamo pulito vuole quindi \(N\) ben minore di \(d\), non \(N \le d\). Ed essendo la transizione l’identità, non c’è modo di dimenticare: una scrittura spuria fatta all’inizio resta a disturbare per sempre. È la ragione per cui l’attenzione lineare pura, così com’è, resta indietro rispetto all’attenzione softmax proprio sui compiti di richiamo preciso.
Il conto si rifà in poche righe: si riempie una memoria \(32\times32\) con \(N\) associazioni, la si rilegge con le stesse chiavi con cui è stata scritta, e si guarda di quanto la risposta si discosta dal valore giusto.
import numpy as np
def errore_richiamo(N, d=32, prove=2000, seme=0):
# errore relativo medio nel rileggere N associazioni da una memoria d x d
rng = np.random.default_rng(seme)
errori = []
for _ in range(prove):
K = rng.standard_normal((N, d))
K /= np.linalg.norm(K, axis=1, keepdims=True) # chiavi a norma unitaria
V = rng.standard_normal((N, d))
S = V.T @ K # la memoria dopo N scritture
letto = (S @ K.T).T # rileggo con le stesse chiavi
errori.append((np.linalg.norm(letto - V, axis=1)
/ np.linalg.norm(V, axis=1)).mean())
return float(np.mean(errori))
for N in (8, 16, 32):
print(f"N={N:2d} errore relativo medio = {errore_richiamo(N):.2f}"
f" (andamento atteso ~{np.sqrt((N - 1) / 32):.2f})")
N= 8 errore relativo medio = 0.46 (andamento atteso ~0.47)
N=16 errore relativo medio = 0.68 (andamento atteso ~0.68)
N=32 errore relativo medio = 0.99 (andamento atteso ~0.98)
Da qui in avanti l’intero capitolo è un tentativo di curare questi due mali tenendo però il dono della memoria che non cresce. Servono due ingredienti: un modo per dimenticare, cioè un interruttore (in inglese gate) che lasci sbiadire quello che è vecchio invece di tenerlo eterno; e un modo per correggere invece di sommare alla cieca, cioè andare a vedere che cosa il foglio risponde già a quella etichetta e scriverci sopra solo la differenza (è la regola delta). Sono esattamente i due fili della prossima sezione.
Da ricordare
Far parlare ogni parola con tutte le altre è un’esplosione di conversazioni: mille parole sono quasi mezzo milione di scambi. Sostituire quel confronto a due a due con un registro riassuntivo, che ciascuno aggiorna una volta sola, porta il conto a mille aggiornamenti, ciascuno grande quanto il registro e sempre uguale a sé: da esplosivo a proporzionale alla lunghezza del testo, ed è esattamente ciò che si intende per costo lineare.
Il registro è un foglio di dimensione fissa: ogni parola ci scrive la propria informazione sotto la propria etichetta, e per rispondere a una domanda lo si rilegge, invece di rovistare nella pila degli scontrini.
Letto parola per parola, quel foglio è il riassunto di una vecchia rete ricorrente: si aggiorna sommando, costa sempre lo stesso a ogni parola e non cresce mai, che si sia alla decima o alla milionesima. È l’ironia del titolo Transformers are RNNs (Katharopoulos e colleghi, 2020).
Stesso risultato, due modi di ottenerlo: tutto insieme quando il testo c’è già (studiare un libro che si ha in mano, cioè addestrare, e in pratica si fa a blocchi) e una parola alla volta quando il testo si sta inventando (raccontarlo a voce, cioè generare), senza la pila che si allunga. La stessa doppia natura tornerà con gli State Space Model.
Il difetto del registro: somma e basta, non cancella e non corregge. Le scritte si sovrappongono un po’ fin dalla prima riga, e il foglio si riempie molto prima di quanto lascino sperare le sue caselle, perché ogni informazione si spalma su tutto il foglio: uno da trentadue righe per trentadue colonne ne ha più di mille, ma già a otto informazioni la risposta torna sbagliata di circa metà del suo valore, e a trentadue lo sbaglio è grande quanto la risposta. Prendere un foglio più grande si può, ma costa conti e memoria a ogni parola, e a quel punto tanto vale un Transformer.
I due rimedi delle sezioni successive: un modo per sbiadire ciò che è vecchio e un modo per correggere invece di sommare alla cieca.
Da ricordare
La softmax dei Transformer costa \(O(n^2 d)\) perché non si fattorizza in dimensione finita; sostituirla con una somiglianza \(\text{sim}(\mathbf{q},\mathbf{k})=\phi(\mathbf{q})^\top\phi(\mathbf{k})\) e ri-associare il prodotto porta il costo a \(O(n d^2)\), cioè lineare nella lunghezza.
Il calcolo si condensa in uno stato-matrice \(\mathbf{S} = \sum_j \mathbf{v}_j\,\phi(\mathbf{k}_j)^\top\) di dimensione fissa \(d \times d\): una memoria chiave→valore che si legge con \(\mathbf{S}\,\phi(\mathbf{q})\).
In forma causale è una ricorrenza, \(\mathbf{S}_t = \mathbf{S}_{t-1} + \mathbf{v}_t\,\phi(\mathbf{k}_t)^\top\): cioè una RNN a stato matriciale con transizione lineare (l’identità), aggiornamento \(O(d^2)\) per token e memoria costante; da cui l’ironia di Transformers are RNNs (Katharopoulos et al., 2020).
Stessa funzione, due forme: parallela per addestrare, ricorrente per generare a memoria costante, senza la KV cache che invece cresce. La stessa dualità tornerà per gli State Space Model. Lineare non vuol dire però subito più veloce: sotto qualche migliaio di token una implementazione ingenua perde contro un’attenzione softmax ben scritta.
Il difetto dell’accumulo puro: non dimentica e non corregge. Con chiavi casuali l’interferenza (crosstalk) cresce come \(\sqrt{N/d}\) fin da subito e a \(N \approx d\) pareggia il segnale: il richiamo pulito vuole \(N\) ben minore di \(d\).
I due rimedi (un gate che dimentica e una regola delta che corregge invece di sommare) sono il filo delle sezioni successive.