Il viaggio di un token, passo per passo

Qui sotto ci sono sei passi, e li fai avanzare tu con i bottoni: niente scorrimento che ti porta via la pagina, niente animazione da inseguire. La frase resta sempre la stessa — «La bambina mangia la mela» — e a cambiare è soltanto il disegno.

Prima però una dichiarazione, perché conta più della grafica. I numeri di partenza — i cinque valori sotto ogni parola, le regole della batteria di rivelatori, i vettori delle parole candidate — li ho scelti io a mano, per far vedere la forma del calcolo: un modello vero ne usa migliaia per parola e se li costruisce da solo leggendo miliardi di frasi. Tutto il resto, invece, è aritmetica vera fatta dal tuo browser mentre guardi: il prodotto scalare dell'attention, la softmax, la somma sul nastro, la soglia dei rivelatori, i punteggi finali. Le prime quattro caselle e la matrice delle domande sono le stesse dichiarate nella pagina sull'attention, così le due demo tornano fra loro.

Conviene mettere i nomi giusti sulle cose che il disegno mostra. Il pezzo di testo è un token; il numero che lo identifica è la sua riga nella matrice di embedding, una tabella con una riga per ogni token del vocabolario e tante colonne quanta è la larghezza interna del modello, la d_model. Nel transformer originale di Ashish Vaswani e colleghi («Attention Is All You Need», NeurIPS 2017) quella larghezza vale 512; in GPT-3 vale 12288, e il vocabolario ha circa 50.000 righe (Brown et al., 2020). Le mie cinque caselle con un nome sopra sono una finzione utile: in un modello vero nessuna singola componente del vettore significa qualcosa da sola, e i concetti stanno distribuiti su moltissime direzioni contemporaneamente.

La cosa che vale la pena portarsi via da questa pagina, però, non è il vocabolario tecnico: è che il vettore di un token non è un'etichetta fissa. All'ingresso «mela» ha il vettore che le assegna il dizionario, uguale in qualunque frase compaia. Da lì in poi viene riscritto in funzione del contesto: dopo l'attention porta dentro di sé un pezzo di «mangia», dopo il blocco successivo un pezzo di quel che i rivelatori hanno concluso, e così via. La «mela» che esce dall'ultimo strato di questa frase è numericamente diversa dalla «mela» di «la mela cade dall'albero». È per questo che si parla di rappresentazioni contestuali, ed è la ragione per cui un transformer distingue significati che un dizionario di vettori statici — la generazione word2vec — non poteva distinguere.

Il nastro ha un nome in letteratura: flusso residuale, residual stream, e la formulazione più chiara è quella di Nelson Elhage e colleghi in «A Mathematical Framework for Transformer Circuits» (Anthropic, 2021). L'idea è questa: dentro un transformer non esiste una catena in cui ogni strato passa il testimone al successivo. Esiste un vettore per posizione — largo d_model — che attraversa tutto il modello dall'inizio alla fine, e ogni blocco legge da quel vettore e ci somma il proprio contributo. Nessun blocco cancella quello che ha scritto un altro: al massimo può scrivere qualcosa che lo contrasta. In formula è la riga x ← x + Attn(x) seguita da x ← x + MLP(x), ripetuta per ogni blocco. In GPT-3 i blocchi sono 96, quindi il vettore di ogni token viene riscritto quasi duecento volte prima di arrivare in fondo.

Quella somma — la connessione residuale, o skip connection — non è un dettaglio implementativo: è la ragione per cui reti così profonde si possono addestrare. Kaiming He e colleghi la introducono in «Deep Residual Learning for Image Recognition» (2015) per un motivo empirico spiazzante: le reti molto profonde non andavano in sovradattamento, andavano peggio anche sui dati di addestramento, e cioè non riuscivano nemmeno a imparare la funzione identità. Facendo passare l'ingresso intatto accanto al blocco, ogni strato deve imparare soltanto la correzione da aggiungere, e partire da «non faccio niente» diventa gratis. C'è anche una lettura in termini di gradiente, cioè della derivata che dice a ogni peso di quanto muoversi: la derivata di x + f(x) rispetto a x contiene un 1 che sopravvive a qualunque profondità, e il segnale di correzione arriva fino ai primi strati invece di spegnersi per strada. Con 152 strati He e colleghi vincono ImageNet 2015; due anni dopo il transformer eredita la stessa struttura. Dove le connessioni residuali vanno esattamente — prima o dopo la normalizzazione dello strato — è rimasta una questione aperta a lungo: Ruibin Xiong e colleghi hanno mostrato (ICML 2020) che mettendo la layer norm dentro il ramo residuale, e non dopo la somma, i gradienti alla partenza restano ben scalati e il riscaldamento del learning rate — che nel paper del 2017 era indispensabile — si può togliere. Quasi tutti i modelli di oggi sono pre-LN per questo motivo.

Il quinto passo è la parte del transformer di cui si parla meno e che pesa di più. Ogni blocco contiene, dopo l'attenzione, un piccolo percettrone applicato a ciascuna posizione indipendentemente: il feed-forward, o MLP. Fa due sole cose: allarga il vettore a quattro volte la propria larghezza, applica una non linearità — ReLU nel paper del 2017, GELU in GPT-2 e successori — e lo riporta alla larghezza di prima. Nel transformer originale il passaggio è 512 → 2048 → 512; in GPT-3 è 12288 → 49152 → 12288. Contando i parametri di un blocco, l'attenzione ne usa 4·d² (le quattro matrici Q, K, V e quella di uscita) e il feed-forward 8·d²: due terzi dei parametri di un blocco stanno lì dentro.

Che cosa ci facciano è la domanda a cui rispondono Mor Geva, Roei Schuster, Jonathan Berant e Omer Levy in «Transformer Feed-Forward Layers Are Key-Value Memories» (EMNLP 2021): la prima matrice funziona da insieme di chiavi, ciascuna delle quali si accende su un certo schema del testo — un pattern lessicale, una costruzione sintattica, un argomento — e la seconda da insieme di valori, cioè di distribuzioni sul vocabolario che quella chiave, quando è accesa, aggiunge al flusso residuale. È esattamente la struttura della demo, in miniatura: quattro chiavi con un nome, una soglia, e una firma che chi si accende deposita sul nastro. E si vede bene una proprietà che nella prosa divulgativa si perde: nel mio giocattolo la casella «qualità» non esisteva nei dati d'ingresso. Nessuna parola nasce con quella proprietà; la crea il feed-forward combinando due caselle che c'erano già. Togli quel blocco e le parole «rossa» e «matura» del passo 6 prendono punteggio zero, perché niente nel nastro le sostiene.

All'uscita dell'ultimo blocco il vettore viene moltiplicato per una matrice che ha una riga per ogni token del vocabolario: il risultato è un punteggio grezzo per parola, il logit. In GPT-2 quella matrice non è nemmeno nuova — è la matrice di embedding usata al contrario, un trucco che Ofir Press e Lior Wolf hanno proposto (EACL 2017) e che risparmia decine di milioni di parametri migliorando la qualità. La softmax trasforma poi i logit in una distribuzione di probabilità: esponenziale di ciascuno, diviso per la somma di tutti. Durante l'addestramento quella distribuzione viene confrontata con il token che compare davvero, e la misura dello scarto è la cross-entropy, cioè quanta sorpresa il modello prova davanti alla parola giusta; è quel numero che la discesa del gradiente cerca di abbassare, come racconta la pagina sui pesi. In generazione, invece, dalla distribuzione si pesca, ed è lì che entrano temperatura, top-k e top-p: la pagina sulla risposta li smonta uno per uno.

Dove il mio giocattolo smette di valere, e conviene saperlo. Primo, non c'è nessuna codifica di posizione: «La» e «la» hanno lo stesso vettore e ricevono lo stesso peso, mentre un modello vero somma al token un segnale posizionale — sinusoidale nel 2017, oggi quasi sempre rotatorio — senza il quale l'ordine delle parole sarebbe invisibile. Secondo, manca la layer norm, che a ogni blocco riporta il vettore a una scala confrontabile. Terzo, c'è una testa sola e un blocco solo, contro i 96 strati per 96 teste di GPT-3: novemiladuecentosedici pattern di attenzione diversi sulla stessa frase. Quarto, ed è il limite più serio: nella demo la parola vincente vince perché somiglia al vettore sul nastro, mentre in un modello vero l'ultimo vettore non assomiglia affatto al token corrente — gli strati lo hanno progressivamente ruotato verso ciò che viene dopo. Che questa rotazione avvenga davvero si può guardare con la logit lens, la tecnica proposta da nostalgebraist nel 2020: si applica la matrice di uscita al flusso residuale di uno strato intermedio e si legge che cosa il modello predirebbe se si fermasse lì. Le previsioni dei primi strati sono confuse e si affilano man mano che si sale — ma la stessa tecnica ha mostrato che il quadro cambia da modello a modello, e va usata come indizio, non come radiografia.