Arrivati alla Parte 3, avevamo a disposizione un modello in grado di apprendere: misurava il proprio margine di errore con un singolo numero e scendeva a valle finché la lingua inglese non smetteva di sorprenderlo. Tuttavia, questo apprendimento avveniva con uno svantaggio. Il modello, infatti, osservava una piccola finestra fissa di caratteri, trattandoli come un ammasso indifferenziato. Non era in grado di capire che, in una sequenza come “the cat sat on the m”, le lettere che rendono probabile la successiva “at” sono quelle che compongono “cat” qualche posizione prima, e non le parole “on the” che si trovano in mezzo.

Ciò che manca al modello è un modo affinché i caratteri precedenti giusti possano protendersi in avanti e influenzare la previsione corrente, mentre quelli irrilevanti restano in silenzio. Questo meccanismo prende il nome di attenzione (attention). Si tratta dell’idea centrale dell’articolo del 2017 che ha dato il nome all’architettura transformer, nonché la “T” di GPT; una volta che la si osserva costruita pezzo per pezzo, ogni alone di mistero svanisce. In sostanza, è una media ponderata dotata di un metodo ingegnoso per scegliere i pesi.

Il problema risolto dall’attenzione

Si pensi alla previsione del carattere successivo dopo “she opened the door and looked ”. Un’ottima ipotesi è la lettera “a” (per iniziare parole come “at” o “around”). Per formulare correttamente questa ipotesi, il modello deve sapere che c’è una persona, “she”, che sta compiendo l’azione di guardare, e che l’azione “opened the door” è già avvenuta. Queste informazioni si trovano molti caratteri più indietro. La previsione nella posizione attuale ha quindi bisogno di recuperare le posizioni precedenti rilevanti, ignorando tutto il resto.

La soluzione più ingenua, ovvero ricordare ogni cosa in una tabella di ricerca più grande, rappresenta esattamente l’ostacolo contro cui ci siamo scontrati nella Parte 2. L’attenzione agisce in modo più intelligente. Invece di memorizzare quale contesto preveda un determinato elemento, permette a ogni posizione nel testo di scorrere tutte le posizioni precedenti e decidere, al volo, quali di esse valga la pena ascoltare in quel preciso istante. Questa decisione viene appresa, non è fissa; di conseguenza, il modello può imparare che i verbi dipendono dai loro soggetti e che le virgolette aperte richiamano quelle chiuse, senza che nessuno debba programmare tali regole esplicitamente.

Queries, keys e values

Ecco il meccanismo. Il modo standard per spiegarlo ricorre a una piccola analogia che, di fatto, corrisponde esattamente a ciò che avviene a livello matematico.

Ogni posizione produce tre elementi a partire dal proprio vettore, ciascuno dei quali non è altro che una diversa combinazione appresa dei suoi numeri:

• Una query: ciò che questa posizione sta cercando. La posizione successiva a “looked” potrebbe chiedersi, grosso modo: “chi sta compiendo un’azione, e qual era l’azione?”
• Una key: ciò che questa posizione offre alle altre. La posizione che contiene “she” annuncia, in sintesi: “Sono un soggetto, una persona che fa delle cose”.
• Un value: il contributo effettivo che questa posizione fornirà se qualcuno le presta ascolto.

A questo punto avviene l’abbinamento. Per la posizione corrente, si confronta la sua query con la key di ogni posizione precedente. Quando una query e una key si allineano bene, significa che la posizione precedente è rilevante e ottiene un punteggio alto. Quando non si allineano, il punteggio è basso. Passando questi punteggi attraverso la funzione softmax vista nella Parte 3, essi si trasformano in pesi la cui somma è pari a uno: un insieme di proporzioni che indicano di “prestare il 70% dell’attenzione a questa posizione, il 20% a quell’altra e quasi zero al resto”. Il nuovo vettore della posizione diventa quindi la media ponderata dei values di tutti, utilizzando esattamente quelle proporzioni. Le posizioni rilevanti contribuiscono molto, mentre quelle irrilevanti passano quasi inosservate.

Questa è una “attention head” (testa di attenzione) nella sua interezza. Tre proiezioni apprese per creare queries, keys e values; un punteggio derivante dal confronto tra queries e keys; una softmax per trasformare i punteggi in pesi; una media ponderata dei values. Nel repository si tratta di una manciata di righe di codice, e il passaggio relativo a punteggi e pesi si legge quasi come la descrizione appena fornita:

// Weights returns the causal softmax attention weights for input x: for each
// position, how much it attends to every position it is allowed to see.
func (h *Head) Weights(x *Tensor) *Tensor {
	q := MatMul(x, h.Wq) // what each position is looking for
	k := MatMul(x, h.Wk) // what each position offers
	// score every position against every other, scaled so softmax stays sane
	scores := MulScalar(MatMul(q, Transpose(k)), 1/math.Sqrt(float64(h.headSize)))
	scores = MaskedFillCausal(scores) // no peeking at the future
	return Softmax(scores)            // scores become weights that sum to 1
}

Vietato sbirciare nel futuro

C’è una regola che quella riga di codice applica silenziosamente, ed è abbastanza importante da meritare un nome: MaskedFillCausal. Un modello linguistico viene addestrato per prevedere il carattere successivo; pertanto, quando deve decidere cosa viene dopo la quinta posizione, gli deve essere permesso di guardare le posizioni dalla uno alla cinque, ma mai la sesta o le successive. Se potesse vedere il futuro, il compito sarebbe banale e il modello non imparerebbe nulla di utile, proprio come uno studente che ha accesso alle soluzioni durante un esame.

Di conseguenza, prima della softmax, si prende ogni punteggio che permetterebbe a una posizione di prestare attenzione a una posizione successiva e lo si imposta a infinito negativo. Dopo la softmax, l’infinito negativo diventa un peso pari a zero. Ogni posizione può prestare attenzione a se stessa e a tutto ciò che la precede, ma a nulla di ciò che la segue. Questo processo prende il nome di “causal masking” (mascheramento causale), ed è il motivo per cui il modello può essere addestrato su testi comuni: ogni posizione si esercita simultaneamente a prevedere il proprio successore, utilizzando esclusivamente ciò che è venuto prima.

Diversi tipi di attenzione in contemporanea

Una singola “head” apprende un solo tipo di relazione, un unico modo per decidere cosa sia rilevante. Il linguaggio, tuttavia, presenta molteplici relazioni che si sviluppano in contemporanea. Un verbo si lega al suo soggetto, un pronome al nome che sostituisce, una parentesi chiusa alla corrispondente parentesi aperta. Chiedere a una sola testa di tenere traccia di tutto questo significa pretendere troppo.

Per questo motivo, il modello esegue diverse teste in parallelo, ognuna con le proprie queries, keys e values, ed è libera di specializzarsi. Una testa potrebbe imparare a guardare il carattere precedente, un’altra l’inizio della parola corrente, un’altra ancora una virgoletta corrispondente rimasta molto indietro. I loro risultati vengono poi ricuciti insieme e mescolati. Questo meccanismo si chiama “multi-head attention” (attenzione multi-testa), ed è la ragione per cui un transformer riesce a destreggiarsi tra diversi tipi di contesto contemporaneamente senza che questi interferiscano tra loro.

Vederlo in azione

Poiché i pesi dell’attenzione sono semplici proporzioni, è possibile stamparli e osservarli. La demo della fase 4 esegue una singola testa su una breve stringa e mostra, per ogni carattere, quanta attenzione venga prestata a ciascun carattere precedente. L’output si presenta in questo modo (si tratta di una testa non addestrata, quindi le proporzioni esatte sono quasi casuali, ma è la struttura a essere fondamentale):

go run ./cmd/stage4_attention

attention weights over "hello" (row i = how char i attends):

          h    e    l    l    o
    h  1.00 0.00 0.00 0.00 0.00
    e  0.43 0.57 0.00 0.00 0.00
    l  0.31 0.32 0.36 0.00 0.00
    l  0.23 0.24 0.27 0.27 0.00
    o  0.20 0.24 0.17 0.17 0.21

Ci sono due elementi in questa griglia che racchiudono l’intera lezione. In primo luogo, tutto il triangolo in alto a destra è pari a zero: si tratta della maschera causale, per cui ogni carattere si rifiuta di guardare quelli successivi. La prima h può prestare attenzione solo a se stessa, quindi si assegna un valore di 1.00. La e può guardare la h e se stessa, dividendo l’attenzione tra le due. Quando arriviamo alla o, l’attenzione si distribuisce su tutti e cinque i caratteri. In secondo luogo, la somma di ogni riga è pari a uno, trattandosi di proporzioni. Questa testa non è addestrata, motivo per cui i suoi pesi sono vicini a una distribuzione uniforme. Dopo l’addestramento, iniziano a emergere degli schemi: le teste imparano a concentrare l’attenzione esattamente dove si trovano le informazioni utili.

Quella griglia rappresenta l’attenzione messa a nudo. Non c’è alcuna magia, solo una media ponderata, mascherata e appresa. Eppure, era proprio il tassello mancante. Grazie ad esso, ogni posizione può tornare indietro e raccogliere con precisione il contesto precedente di cui ha bisogno, imparando al contempo quale sia questo contesto.

Ora abbiamo a disposizione tutte le componenti: un tokenizer, un motore di autograd, una funzione di loss, la discesa del gradiente e l’attenzione. Nella Parte 5 le assembleremo nel transformer completo, impileremo alcuni strati e lo metteremo alla prova su un vero libro. A quel punto faremo l’unica cosa che convince davvero chiunque: osserveremo la loss scendere e leggeremo ciò che la macchina è in grado di sognare.

Il codice relativo a questa parte si trova in cmd/stage4_attention e attention.go su github.com/erubboli/go-tiny-llm.