Salta al contenuto principale

95 - L'attenzione, derivata per intero: query, chiave, valore e il blocco Transformer scritto da zero

L'attenzione derivata per intero in NumPy: costo quadratico, cache chiave-valore, ablazioni su un decoder che copia e ordina, e perché i pesi non spiegano.

Cosa ne pensi di questo articolo?

95 - L'attenzione, derivata per intero: query, chiave, valore e il blocco Transformer scritto da zero
Machine Learning Engineerfedericocalo.dev

In Breve

L'attenzione derivata per intero in NumPy: costo quadratico, cache chiave-valore, ablazioni su un decoder che copia e ordina, e perché i pesi non spiegano.

Nella frase "il tecnico ha consegnato il modulo al cliente perché era incompleto" il pronome si riferisce al modulo, e chi legge lo capisce senza fatica: non rilegge dall'inizio e non si affida a un riassunto, torna indietro di sei parole, confronta "era incompleto" con i due candidati e prende da lì l'informazione. Formulare una domanda, confrontarla con tutto il disponibile, pesare il contenuto in base a quanto ogni pezzo risponde: questa è l'attenzione.

Il meccanismo nasce nel 2014 con Bahdanau, Cho e Bengio come stampella di una rete ricorrente che traduceva male le frasi lunghe (Neural Machine Translation by Jointly Learning to Align and Translate, arXiv:1409.0473); tre anni dopo Vaswani e colleghi togliono la stampella e buttano via la gamba (Attention Is All You Need, arXiv:1706.03762). Qui la derivi per intero, conti i parametri di un blocco, misuri il costo quadratico, e addestri un decoder a copiare e ordinare sequenze.

Come sono scritti gli esempi. Il codice è NumPy puro: attenzione, teste multiple, blocco, passata all'indietro scritta a mano e addestramento del decoder. Ambiente: Python 3.14.4, NumPy 2.3.5, semi fissati. Ogni listato è stato eseguito due volte con uscita identica; fanno eccezione i tempi, riportati come intervallo. I listati si leggono in ordine, perché ognuno usa i nomi definiti dai precedenti. PyTorch 2.12.0 compare in un punto solo, un confronto numerico della formula con torch.nn.functional.scaled_dot_product_attention, eseguito su CPU con l'interprete che lo contiene e senza scaricare nulla.

L'analogia: la domanda allo schedario

Immagina uno schedario. Ogni scheda ha una intestazione sul bordo — "fatture 2024", "contratti scaduti" — e dentro il contenuto. Arrivi con una domanda: quanto abbiamo speso in trasporti l'anno scorso. Non apri tutte le schede e non ne apri una sola: scorri le intestazioni, dai a ciascuna un punteggio di somiglianza, e leggi il contenuto di quelle col punteggio più alto pesando ognuna per quanto somigliava.

Il dizionario è questo: la domanda è la query, le intestazioni sono le chiavi, i contenuti sono i valori; la somiglianza è un prodotto scalare e il passaggio a pesi che sommano a uno è il softmax. Intestazione e contenuto della stessa scheda sono due cose diverse, ricavate dallo stesso foglio con criteri diversi: chi confronta direttamente i vettori d'ingresso sa dire solo "sei uguale a me".

Il muro che le reti ricorrenti hanno preso

Il modulo 06 ha smontato le ricorrenti pezzo per pezzo: la lezione sul gradiente che svanisce e che esplode nel tempo ha mostrato il prodotto di Jacobiane che si contrae, quella sulla cella LSTM con le porte ha costruito la via d'uscita, l'esperimento su più semi ha misurato quanto vale. Restano due limiti che non dipendono dall'ottimizzazione ma dalla forma del calcolo.

Il primo limite è il numero di passi sequenziali. Lo stato hₜ dipende da hₜ₋₁: per T token servono T operazioni in ordine. Un'attenzione calcola tutti i punteggi con due moltiplicazioni di matrici: passi sequenziali O(1). Il secondo limite è il cammino fra token distanti. In una ricorrente l'informazione del token 1 attraversa 99 trasformazioni per arrivare al token 100 e il gradiente rifà la strada: cammino O(T). Nell'attenzione è 1. È la tabella 1 di arXiv:1706.03762, ed è il vero argomento del paper.

Il prezzo si paga altrove. Una ricorrente costa O(T · d²) e tiene uno stato di dimensione fissa; l'attenzione costa O(T² · d) e tiene in memoria una matrice T × T. Lo misureremo con un cronometro.

Attenzione a prodotto scalato, derivata

Partiamo da X di forma (T, d): T token da d numeri. Tre matrici di parametri proiettano ogni token in tre ruoli, e i punteggi di somiglianza sono un unico prodotto di matrici:

Q = X W_Q      (T, d) · (d, d_k) → (T, d_k)      le domande
K = X W_K      (T, d) · (d, d_k) → (T, d_k)      le intestazioni
V = X W_V      (T, d) · (d, d_v) → (T, d_v)      i contenuti
S = Q Kᵀ       (T, d_k) · (d_k, T) → (T, T)      sᵢⱼ = qᵢ · kⱼ = Σ_{m=1..d_k} qᵢₘ · kⱼₘ

La riga i di Q è la domanda del token i, la riga j di K l'intestazione del token j, la riga j di V il suo contenuto; quando i tre vengono dalla stessa X si parla di auto-attenzione. La matrice S è quadrata nella lunghezza e non è simmetrica, perché domanda e intestazione passano per matrici diverse.

I punteggi diventano pesi che sommano a uno, riga per riga, con un softmax sull'asse delle chiavi; l'uscita è la media dei valori con quei pesi, e insieme danno l'equazione 1 di arXiv:1706.03762. Scrivo l'indice esplicitamente perché è l'asse che si sbaglia: se normalizzi per colonne il codice gira, le forme tornano, e nessuna eccezione ti avvisa.

pᵢⱼ = exp(sᵢⱼ / √d_k) / Σ_{j′=1..T} exp(sᵢⱼ′ / √d_k)     la somma corre su j′, cioè sulle CHIAVI, con i fissato
Y = P V        (T, T) · (T, d_v) → (T, d_v)              yᵢ = Σ_{j=1..T} pᵢⱼ · vⱼ

Attenzione(Q, K, V) = softmax( Q Kᵀ / √d_k ) V

Il listato traduce la formula riga per riga, con una dimensione in più per il lotto B. Il softmax sottrae il massimo prima di esponenziare: non cambia il risultato, perché il softmax è invariante all'aggiunta di una costante a tutti i punteggi, ma impedisce a exp di traboccare. Le matrici dei pesi sono gaussiane divise per √d, così che le proiezioni mantengano varianza vicina a uno. Le ultime righe sono i controlli: le forme delle tre matrici e la somma di ogni riga dei pesi, che deve fare uno.

import numpy as np

def softmax(s, asse=-1):
    """Softmax numericamente stabile: si sottrae il massimo prima di esponenziare."""
    s = s - s.max(axis=asse, keepdims=True)
    e = np.exp(s)
    return e / e.sum(axis=asse, keepdims=True)

def attenzione(X, W_Q, W_K, W_V, scala=True):
    Q = X @ W_Q                                   # (B, T, d_k)
    K = X @ W_K                                   # (B, T, d_k)
    V = X @ W_V                                   # (B, T, d_v)
    d_k = Q.shape[-1]
    punteggi = Q @ K.transpose(0, 2, 1)           # (B, T, T)
    if scala:
        punteggi = punteggi / np.sqrt(d_k)
    pesi = softmax(punteggi, asse=-1)             # softmax SULLE CHIAVI
    return pesi @ V, pesi, punteggi               # (B, T, d_v), (B, T, T), (B, T, T)

rng = np.random.default_rng(7)
B, T, d, d_k = 2, 5, 16, 8
X   = rng.normal(size=(B, T, d))
W_Q = rng.normal(size=(d, d_k)) / np.sqrt(d)
W_K = rng.normal(size=(d, d_k)) / np.sqrt(d)
W_V = rng.normal(size=(d, d_k)) / np.sqrt(d)

Y, P, S = attenzione(X, W_Q, W_K, W_V)
print("punteggi QK^T      ", S.shape)
print("pesi (dopo softmax)", P.shape)
print("uscita             ", Y.shape)
print("somma di ogni riga dei pesi:", np.round(P.sum(axis=-1), 12).ravel()[:5])
punteggi QK^T       (2, 5, 5)
pesi (dopo softmax) (2, 5, 5)
uscita              (2, 5, 8)
somma di ogni riga dei pesi: [1. 1. 1. 1. 1.]

Prima di fidarsi di una formula scritta a mano conviene confrontarla con una implementazione indipendente. PyTorch ha la funzione scaled_dot_product_attention, che calcola la stessa espressione; il listato la applica agli stessi Q, K, V della versione NumPy, con e senza maschera causale, e stampa lo scarto massimo. Le variabili softmax e np vengono dal listato precedente.

import torch
import torch.nn.functional as F

torch.manual_seed(0)
torch.set_num_threads(2)

rng = np.random.default_rng(7)
B, T, d_k = 2, 5, 8
Q, K, V = (rng.normal(size=(B, T, d_k)) for _ in range(3))
S = Q @ K.transpose(0, 2, 1) / np.sqrt(d_k)
mio = softmax(S) @ V
mio_causale = softmax(np.where(np.tril(np.ones((T, T), dtype=bool)), S, -np.inf)) @ V

Qt, Kt, Vt = (torch.from_numpy(a) for a in (Q, K, V))
libreria = F.scaled_dot_product_attention(Qt, Kt, Vt).numpy()
libreria_causale = F.scaled_dot_product_attention(Qt, Kt, Vt, is_causal=True).numpy()
print("torch", torch.__version__)
print("scarto massimo, senza maschera:", f"{np.abs(mio - libreria).max():.2e}")
print("scarto massimo, causale       :", f"{np.abs(mio_causale - libreria_causale).max():.2e}")
torch 2.12.0+cu130
scarto massimo, senza maschera: 2.22e-16
scarto massimo, causale       : 2.22e-16

Cosa dicono i numeri. Lo scarto è 2,22 · 10⁻¹⁶, cioè un errore di arrotondamento in doppia precisione: la formula di sopra e quella della libreria sono lo stesso calcolo. Non dimostra che l'attenzione sia la scelta giusta per il tuo problema; dimostra che il tuo codice implementa quella del paper, e che l'opzione is_causal=True equivale alla maschera con meno infinito prima del softmax.

Perché si divide per la radice di d_k, misurato

La radice quadrata è l'unico pezzo di arXiv:1706.03762 relegato in una nota a piè di pagina pur essendo decisivo. Il punteggio è sᵢⱼ = Σ_{m=1..d_k} qᵢₘ · kⱼₘ; se le componenti di q e k sono indipendenti con media 0 e varianza 1, ogni addendo ha media 0 e varianza 1, e la somma di d_k addendi indipendenti ha varianza d_k:

Var(qᵢ · kⱼ) = d_k          sd(qᵢ · kⱼ) = √d_k          sd( (qᵢ · kⱼ) / √d_k ) = 1

In parole: più grande è la dimensione del confronto, più larghi diventano i punteggi. Senza la divisione, a d_k = 512 i punteggi oscillano su decine e il softmax concentra quasi tutto il peso su una chiave sola. La derivata del softmax è ∂pᵢ/∂sⱼ = pᵢ(δᵢⱼ − pⱼ) e la sua traccia Σⱼ pⱼ(1 − pⱼ) va a zero quando un peso vale quasi uno: il gradiente che insegna dove guardare non arriva.

def misura(d_k, scala, seme, T=64, ripetizioni=400):
    """Statistiche dei punteggi e dei pesi, con e senza la divisione per sqrt(d_k)."""
    rng = np.random.default_rng(seme)
    Q = rng.normal(size=(ripetizioni, T, d_k))        # componenti indipendenti, media 0, varianza 1
    K = rng.normal(size=(ripetizioni, T, d_k))
    grezzi = Q @ K.transpose(0, 2, 1)                 # (r, T, T): i prodotti scalari
    punteggi = grezzi / np.sqrt(d_k) if scala else grezzi
    P = softmax(punteggi, asse=-1)
    sensibilita = (P * (1.0 - P)).sum(axis=-1)        # traccia della jacobiana del softmax
    entropia = -(P * np.log(P + 1e-300)).sum(axis=-1)
    return (grezzi.var(), punteggi.std(), P.max(axis=-1).mean(),
            entropia.mean(), sensibilita.mean())

print("d_k  diviso  Var(q.k)  sd punteggi  peso max  entropia  sensibilita")
for d_k in (8, 64, 512):
    for scala in (False, True):
        v, sd, pm, h, s = misura(d_k, scala, 100 + d_k)
        print(f"{d_k:4d}  {'si    ' if scala else 'no    '} {v:8.2f} {sd:12.3f}"
              f" {pm:9.4f} {h:9.4f} {s:12.4f}")
print("entropia massima su 64 chiavi = ln(64) =", round(float(np.log(64)), 4),
      "| sensibilita massima = 1 - 1/64 =", round(1 - 1 / 64, 4))
d_k  diviso  Var(q.k)  sd punteggi  peso max  entropia  sensibilita
   8  no         7.98        2.825    0.4240    2.0727       0.7203
   8  si         7.98        0.999    0.1063    3.6936       0.9596
  64  no        63.99        8.000    0.7913    0.5835       0.2908
  64  si        63.99        1.000    0.1072    3.6861       0.9610
 512  no       512.59       22.641    0.9303    0.1760       0.0998
 512  si       512.59        1.001    0.1075    3.6848       0.9611
entropia massima su 64 chiavi = ln(64) = 4.1589 | sensibilita massima = 1 - 1/64 = 0.9844

Cosa dicono i numeri. La colonna Var(q.k) conferma la derivazione — 7,98, 63,99, 512,59 per d_k 8, 64, 512 — e la deviazione dei punteggi passa da 2,825 a 22,641 senza la divisione, mentre con la divisione resta a 1,00. A d_k = 512 senza scala il peso massimo di una riga vale in media 0,9303 su 64 chiavi e la sensibilità scende a 0,0998 contro 0,9611: un fattore 9,6 di gradiente perduto prima di cominciare.

Esercizio 1

La deviazione che decide il softmax

Un'attenzione lavora con d_k = 64 e vettori le cui componenti sono indipendenti, con media 0 e varianza 1. Carta e penna, nessun calcolatore oltre a e ≈ 2,718.

  1. Quanto valgono varianza e deviazione standard del prodotto scalare q · k?
  2. Due chiavi ottengono punteggi che distano due deviazioni standard. Qual è il rapporto fra i loro pesi dopo il softmax, senza la divisione per √d_k? E con la divisione?
  3. Perché il rapporto non dipende dalle altre 62 chiavi?

Le tre domande sono in sequenza: la prima fissa la scala dei punteggi, la seconda la usa per confrontare due pesi con e senza divisione, la terza spiega perché il confronto è legittimo. Prova a rispondere con il solo ordine di grandezza prima di fare i conti.

Mostra la soluzione

1. Varianza d_k = 64, deviazione √64 = 8: la varianza è la somma di 64 varianze unitarie indipendenti.

2. Due deviazioni sono Δ = 16 sui punteggi grezzi, quindi il rapporto dei pesi vale e¹⁶ ≈ 8,9 · 10⁶: una chiave si prende quasi tutto. Con la divisione per 8 diventa Δ = 2 e il rapporto scende a e² ≈ 7,39: una preferenza, non un'esclusione.

3. Il denominatore del softmax è lo stesso per i due pesi e si semplifica: pᵢ/pⱼ = exp(sᵢ − sⱼ). Il rapporto dipende solo dalla differenza, ed è per questo che √d_k decide tutto.

La maschera causale: meno infinito prima del softmax

Un decoder genera un token alla volta e non può guardare quelli che non ha ancora prodotto, altrimenti impara a copiare la risposta dal futuro. Il punto delicato è dove si applica il divieto: i punteggi proibiti vanno a −∞ prima del softmax, non si azzerano i pesi dopo. Con −∞ prima, quel termine non entra nel denominatore; azzerando dopo, il denominatore ha già incluso le chiavi proibite e la riga non è più una media.

np.set_printoptions(precision=3, suppress=True)
T = 5
rng = np.random.default_rng(3)
punteggi = rng.normal(size=(T, T))
causale = np.tril(np.ones((T, T), dtype=bool))            # True = posso guardare

giusto = softmax(np.where(causale, punteggi, -np.inf), asse=-1)
sbagliato = softmax(punteggi, asse=-1) * causale          # azzerare DOPO il softmax
print("maschera causale (1 = visibile):"); print(causale.astype(int))
print("PRIMA del softmax - somme di riga:", np.round(giusto.sum(-1), 12)); print(giusto)
print("DOPO il softmax   - somme di riga:", np.round(sbagliato.sum(-1), 4))
print("scarto massimo fra le due matrici:", round(float(np.abs(giusto - sbagliato).max()), 4))

lunghezze = np.array([5, 3])                              # la sequenza 2 ha 3 token reali
valido = np.arange(T)[None, :] < lunghezze[:, None]       # (B, T) chiavi reali
combinata = causale[None, :, :] & valido[:, None, :]      # (B, T, T)
print("chiavi valide nel lotto:"); print(valido.astype(int))
print("maschera combinata della sequenza 2 (causale AND padding):")
print(combinata[1].astype(int))
maschera causale (1 = visibile):
[[1 0 0 0 0]
 [1 1 0 0 0]
 [1 1 1 0 0]
 [1 1 1 1 0]
 [1 1 1 1 1]]
PRIMA del softmax - somme di riga: [1. 1. 1. 1. 1.]
[[1.    0.    0.    0.    0.   ]
 [0.859 0.141 0.    0.    0.   ]
 [0.462 0.259 0.278 0.    0.   ]
 [0.119 0.285 0.138 0.458 0.   ]
 [0.115 0.529 0.194 0.068 0.094]]
DOPO il softmax   - somme di riga: [0.733 0.031 0.759 0.874 1.   ]
scarto massimo fra le due matrici: 0.8317
chiavi valide nel lotto:
[[1 1 1 1 1]
 [1 1 1 0 0]]
maschera combinata della sequenza 2 (causale AND padding):
[[1 0 0 0 0]
 [1 1 0 0 0]
 [1 1 1 0 0]
 [1 1 1 0 0]
 [1 1 1 0 0]]

Cosa fa il codice. np.tril costruisce la triangolare inferiore: la riga i vede se stessa e il passato. La seconda riga della matrice sbagliata somma a 0,031 invece di 1 — quel token avrebbe dato il 97% del peso al futuro, gettato via senza ridistribuirlo — e lo scarto massimo fra le due matrici è 0,8317. Solo l'ultima riga coincide, perché per lei non esiste futuro.

Maschera causale e maschera di padding sono due cose diverse, e confonderle è un errore silenzioso. La causale dipende solo dalle posizioni ed è uguale per tutto il lotto: (T, T). Quella di padding dipende dal contenuto — token reali contro riempimento — e cambia da sequenza a sequenza: (B, T) sulle chiavi. Si combinano in AND, e se dimentichi la seconda niente si rompe: i risultati peggiorano in modo che dipende da come hai ordinato il lotto.

Teste multiple: perché h teste da d/h e non una da d

Una sola attenzione produce, per ogni token, una sola combinazione convessa dei valori: un token che deve guardare il soggetto della frase e l'ultimo numero citato ha un unico insieme di pesi per due compiti. L'attenzione a teste multiple (§3.2.2 di arXiv:1706.03762) esegue h attenzioni in parallelo su sottospazi distinti e concatena i risultati.

La scelta che sorprende è la dimensione: ogni testa lavora in d/h dimensioni, non in d. Con una testa piena servono W_Q, W_K, W_V, W_O di forma (d, d): 4d². Con h teste ciascuna ha matrici (d, d/h), le h di una famiglia si impilano in un'unica (d, d), e anche W_O resta (d, d). Totale 4d², identico: le teste multiple sono gratis in parametri.

def multi_testa(X, W_Q, W_K, W_V, W_O, h, maschera=None):
    B, T, d = X.shape
    d_h = d // h
    def dividi(M):                       # (B, T, d) -> (B, h, T, d/h)
        return M.reshape(B, T, h, d_h).transpose(0, 2, 1, 3)
    Q, K, V = dividi(X @ W_Q), dividi(X @ W_K), dividi(X @ W_V)
    punteggi = Q @ K.transpose(0, 1, 3, 2) / np.sqrt(d_h)
    if maschera is not None:
        punteggi = np.where(maschera, punteggi, -np.inf)
    pesi = softmax(punteggi, asse=-1)                       # softmax SULLE CHIAVI
    ctx = pesi @ V                                          # (B, h, T, d/h)
    ctx = ctx.transpose(0, 2, 1, 3).reshape(B, T, d)        # ricuce le teste
    return ctx @ W_O, pesi

rng = np.random.default_rng(11)
B, T, d = 1, 6, 64
X = rng.normal(size=(B, T, d))
W = [rng.normal(size=(d, d)) / np.sqrt(d) for _ in range(4)]
for h in (1, 2, 4, 8):
    Y, pesi = multi_testa(X, *W, h=h)
    print(f"h={h:2d}  d_testa={d//h:3d}  pesi {str(pesi.shape):16s} uscita {Y.shape}"
          f"  4*d^2 = {4*d*d}")

perm = rng.permutation(T)                                   # nessuna maschera, nessuna posizione
Y, _ = multi_testa(X, *W, h=4)
Yp, _ = multi_testa(X[:, perm, :], *W, h=4)
print("max |attn(PX) - P attn(X)| =", f"{np.abs(Yp - Y[:, perm, :]).max():.3e}")
print("max |attn(PX) -   attn(X)| =", f"{np.abs(Yp - Y).max():.3e}")
h= 1  d_testa= 64  pesi (1, 1, 6, 6)     uscita (1, 6, 64)  4*d^2 = 16384
h= 2  d_testa= 32  pesi (1, 2, 6, 6)     uscita (1, 6, 64)  4*d^2 = 16384
h= 4  d_testa= 16  pesi (1, 4, 6, 6)     uscita (1, 6, 64)  4*d^2 = 16384
h= 8  d_testa=  8  pesi (1, 8, 6, 6)     uscita (1, 6, 64)  4*d^2 = 16384
max |attn(PX) - P attn(X)| = 4.441e-16
max |attn(PX) -   attn(X)| = 1.849e+00

Cosa fa il codice. reshape(B, T, h, d_h) spezza l'ultimo asse senza spostare un numero in memoria; transpose(0, 2, 1, 3) porta la testa davanti al tempo. Al ritorno si percorre il cammino inverso nello stesso ordine: se inverti i due passi ottieni un tensore della forma giusta con i numeri mescolati fra teste diverse, un bug che non solleva nessuna eccezione.

Dove sta la posizione

C'è un buco grosso nella formula, e si vede in due righe. Sia Π una permutazione delle righe: allora Q → ΠQ, K → ΠK, V → ΠV, quindi S → Π S Πᵀ e P → Π P Πᵀ, perché il softmax agisce riga per riga; l'uscita è P V → Π P Πᵀ Π V = Π P V, perché Πᵀ Π = I. Permutare l'ingresso permuta l'uscita e nient'altro: l'attenzione è equivariante, e per lei una frase è un sacchetto di parole. ∎

Posizione sinusoidale. La scelta del paper originale (§3.5): all'embedding si somma un vettore fisso di seni e coseni a frequenze geometricamente scalate, PE(pos, 2i) = sin(pos / 10000^{2i/d}) e PE(pos, 2i+1) = cos(pos / 10000^{2i/d}), dove le onde lente distinguono le posizioni lontane. Zero parametri, comportamento definito oltre la lunghezza vista. Posizione appresa. Una tabella (T_max, d) addestrata come qualunque embedding, che lo stesso paper dà "quasi identica": è quella dell'esperimento, costa T_max · d parametri, e alla posizione T_max + 1 non hai niente.

RoPE, posizione rotatoria. Invece di sommare, ruota: le coppie di componenti di q e k girano di un angolo proporzionale alla posizione, prima del prodotto scalare (Su e colleghi, RoFormer: Enhanced Transformer with Rotary Position Embedding, arXiv:2104.09864). Il prodotto scalare fra query in m e chiave in n dipende così solo da m − n, cioè dalla distanza relativa: è la scelta di fatto dei modelli recenti.

Il blocco completo, e il conto dei suoi parametri

L'attenzione da sola non basta: posizione per posizione è una trasformazione lineare dei valori. Serve una non linearità che lavori dentro ogni posizione, servono le connessioni residue perché la pila sia profonda, e serve la normalizzazione per strato (Ba, Kiros e Hinton, Layer Normalization, arXiv:1607.06450). Insieme sono il blocco Transformer.

Il blocco Transformer pre-LN, con le forme e i parametri UN BLOCCO, VARIANTE PRE-LN ingresso x layer norm ATTENZIONE MULTI-TESTA Q, K, V in h teste da d/h, poi W_O somma residua layer norm RETE FEED-FORWARD d allarga a 4d, ReLU, 4d torna a d somma residua uscita, al blocco successivo forma (B, T, d) qui e in ogni altro punto del blocco 2d parametri (gamma, beta) 4d² + 4d parametri il numero di teste non compare 0 parametri: e una somma 2d parametri (gamma, beta) 8d² + 5d parametri due terzi del blocco stanno qui 0 parametri TOTALE 12d² + 13d
Il grosso dei parametri sta nella rete feed-forward, non nell'attenzione: due terzi contro un terzo. Le due vie tratteggiate sono le connessioni residue, le sole che garantiscono un cammino diretto dal gradiente all'ingresso.

La rete feed-forward è una rete a due strati applicata indipendentemente a ogni posizione: espande da d a 4d, applica una non linearità, torna a d; il fattore quattro è la scelta del paper (d = 512, d_ff = 2048). Le connessioni residue rendono l'uscita di ogni sotto-strato x + f(x): il gradiente trova sempre un cammino che non passa per f, e una pila di 24 blocchi resta addestrabile.

Qui si apre la scelta che distingue le due famiglie. Nel post-LN del paper originale si calcola x_nuovo = LayerNorm(x + f(x)), con la normalizzazione dopo la somma; nel pre-LN si calcola x_nuovo = x + f(LayerNorm(x)) e il cammino residuo resta pulito. Oggi si usa il pre-LN, e la ragione è documentata: Xiong e colleghi (arXiv:2002.04745) mostrano che nel post-LN i gradienti presso gli strati vicini all'uscita crescono con la profondità, e servono un passo piccolo e un riscaldamento.

Contiamo i parametri in forma chiusa: quattro matrici (d, d) coi bias danno 4d² + 4d, la feed-forward 8d² + 5d, le due layer norm 4d. Totale 12d² + 13d, e il numero di teste non compare. Per verificarlo contro l'inventario dei tensori serve il modello vero, che serve anche all'esperimento più avanti. Lo scrivo qui per intero, in quattro pezzi: parametri e layer norm, attenzione con la sua passata all'indietro, blocco e perdita, ottimizzatore.

Parametri, layer norm, teste

inizializza costruisce un dizionario nome → array, con i nomi che il conteggio di sotto userà (b0.Wq, b0.g1, e così via). Le matrici sono gaussiane divise per la radice del numero di righe, i bias partono da zero, γ da uno; embedding e posizioni hanno scala 0,1, e la posizione è una tabella appresa (t_max, d), quella dell'esperimento. ln_indietro usa la forma chiusa del gradiente della layer norm, dx = inv · (dx̂ − media(dx̂) − x̂ · media(dx̂ · x̂)): non l'ho derivata nel testo, e la prova che sia giusta è il controllo numerico più sotto. teste e unisci sono i due passi del reshape e transpose visti prima, uno l'inverso dell'altro.

def inizializza(vocab, d, h, n_strati, d_ff, t_max, rng):
    """Tutti i tensori del decoder in un dizionario: nome -> array."""
    n = lambda *forma: rng.normal(size=forma) / np.sqrt(forma[0])
    p = {"E": 0.1 * rng.normal(size=(vocab, d)),      # embedding dei token
         "P": 0.1 * rng.normal(size=(t_max, d))}      # posizione appresa
    for l in range(n_strati):
        q = f"b{l}."
        for nome in ("Wq", "Wk", "Wv", "Wo"):
            p[q + nome] = n(d, d)
            p[q + "b" + nome[1].lower()] = np.zeros(d)
        p[q + "g1"], p[q + "be1"] = np.ones(d), np.zeros(d)
        p[q + "g2"], p[q + "be2"] = np.ones(d), np.zeros(d)
        p[q + "W1"], p[q + "b1"] = n(d, d_ff), np.zeros(d_ff)
        p[q + "W2"], p[q + "b2"] = n(d_ff, d), np.zeros(d)
    p["gf"], p["bef"] = np.ones(d), np.zeros(d)
    p["Wout"], p["bout"] = n(d, vocab), np.zeros(vocab)
    return p

def ln_avanti(x, g, b, eps=1e-5):
    inv = 1.0 / np.sqrt(x.var(-1, keepdims=True) + eps)
    xh = (x - x.mean(-1, keepdims=True)) * inv
    return g * xh + b, (xh, inv, g)

def ln_indietro(dy, c):
    xh, inv, g = c
    dxh = dy * g
    dx = inv * (dxh - dxh.mean(-1, keepdims=True) - xh * (dxh * xh).mean(-1, keepdims=True))
    return dx, (dy * xh).sum((0, 1)), dy.sum((0, 1))

def teste(M, h):                      # (B, T, d) -> (B, h, T, d/h)
    B, T, d = M.shape
    return M.reshape(B, T, h, d // h).transpose(0, 2, 1, 3)

def unisci(M):                        # (B, h, T, d/h) -> (B, T, d)
    B, h, T, dh = M.shape
    return M.transpose(0, 2, 1, 3).reshape(B, T, h * dh)

Attenzione, avanti e indietro

La passata in avanti conserva ciò che serve all'indietro: l'ingresso, Q, K, V, i pesi P e il contesto concatenato. La maschera è applicata con −∞ prima del softmax, come nella sezione dedicata. All'indietro la riga che conta è dS = P · (dP − Σ dP · P): è la derivata del softmax riga per riga, la stessa pᵢ(δᵢⱼ − pⱼ) della sezione sulla scala, contratta con il gradiente dP che arriva dal prodotto con V. Dove la maschera ha messo P = 0 anche dS vale zero, senza lavoro in più: una posizione che non ha influito sul risultato non riceve gradiente.

def att_avanti(x, p, q, h, scala):
    Q, K, V = (teste(x @ p[q + w] + p[q + "b" + w[1].lower()], h) for w in ("Wq", "Wk", "Wv"))
    T = x.shape[1]
    S = Q @ K.transpose(0, 1, 3, 2)
    if scala:
        S = S / np.sqrt(Q.shape[-1])
    S = np.where(np.tril(np.ones((T, T), dtype=bool)), S, -np.inf)   # -inf PRIMA del softmax
    P = softmax(S, -1)
    C = unisci(P @ V)
    return C @ p[q + "Wo"] + p[q + "bo"], (x, Q, K, V, P, C)

def att_indietro(dy, c, p, q, h, g, scala):
    x, Q, K, V, P, C = c
    flat = lambda a: a.reshape(-1, a.shape[-1])
    g[q + "Wo"] += flat(C).T @ flat(dy); g[q + "bo"] += dy.sum((0, 1))
    dC = teste(dy @ p[q + "Wo"].T, h)
    dP = dC @ V.transpose(0, 1, 3, 2)
    dV = P.transpose(0, 1, 3, 2) @ dC
    dS = P * (dP - (dP * P).sum(-1, keepdims=True))
    if scala:
        dS = dS / np.sqrt(Q.shape[-1])
    dQ, dK = dS @ K, dS.transpose(0, 1, 3, 2) @ Q
    dx = 0
    for w, dM in (("Wq", dQ), ("Wk", dK), ("Wv", dV)):
        dM = unisci(dM)
        g[q + w] += flat(x).T @ flat(dM); g[q + "b" + w[1].lower()] += dM.sum((0, 1))
        dx = dx + dM @ p[q + w].T
    return dx

Blocco, perdita, gradiente

avanti costruisce il decoder, nelle due varianti: post_ln falso dà il pre-LN, vero il post-LN, e in entrambe c'è una layer norm finale prima della proiezione sul vocabolario, così il numero di parametri resta lo stesso. perdita_e_gradienti è l'entropia incrociata pesata da pesa, divisa per la somma dei pesi: il gradiente iniziale sui logit è (P − uno-caldo) · peso / somma dei pesi, e poi si percorre il blocco in ordine inverso. Per l'embedding si usa np.add.at e non un'assegnazione, perché lo stesso token può comparire più volte nel lotto e i contributi si sommano.

def avanti(p, tok, n_strati, h, scala=True, post_ln=False):
    B, T = tok.shape
    x = p["E"][tok] + p["P"][:T]
    cache = []
    for l in range(n_strati):
        q = f"b{l}."
        if post_ln:
            a, ca = att_avanti(x, p, q, h, scala)
            x1, cn1 = ln_avanti(x + a, p[q + "g1"], p[q + "be1"])
            z = x1 @ p[q + "W1"] + p[q + "b1"]; f = np.maximum(z, 0) @ p[q + "W2"] + p[q + "b2"]
            x2, cn2 = ln_avanti(x1 + f, p[q + "g2"], p[q + "be2"])
            cache.append((ca, cn1, x1, z, cn2))
        else:
            n1, cn1 = ln_avanti(x, p[q + "g1"], p[q + "be1"])
            a, ca = att_avanti(n1, p, q, h, scala)
            x1 = x + a
            n2, cn2 = ln_avanti(x1, p[q + "g2"], p[q + "be2"])
            z = n2 @ p[q + "W1"] + p[q + "b1"]; f = np.maximum(z, 0) @ p[q + "W2"] + p[q + "b2"]
            x2 = x1 + f
            cache.append((ca, cn1, n2, z, cn2))
        x = x2
    xf, cf = ln_avanti(x, p["gf"], p["bef"])
    return xf @ p["Wout"] + p["bout"], (tok, cache, xf, cf)

def perdita_e_gradienti(p, tok, ber, pesa, n_strati, h, scala=True, post_ln=False):
    logit, (tok, cache, xf, cf) = avanti(p, tok, n_strati, h, scala, post_ln)
    P = softmax(logit, -1)
    B, T, V = P.shape
    pr = np.take_along_axis(P, ber[..., None], -1)[..., 0]
    perdita = -(pesa * np.log(pr + 1e-12)).sum() / pesa.sum()
    dl = P.copy(); np.put_along_axis(dl, ber[..., None], np.take_along_axis(dl, ber[..., None], -1) - 1, -1)
    dl *= (pesa / pesa.sum())[..., None]
    g = {k: np.zeros_like(v) for k, v in p.items()}
    flat = lambda a: a.reshape(-1, a.shape[-1])
    g["Wout"] += flat(xf).T @ flat(dl); g["bout"] += dl.sum((0, 1))
    dx, g["gf"], g["bef"] = ln_indietro(dl @ p["Wout"].T, cf)
    for l in reversed(range(n_strati)):
        q = f"b{l}."
        ca, cn1, nn, z, cn2 = cache[l]
        if post_ln:
            dx2, g[q + "g2"], g[q + "be2"] = ln_indietro(dx, cn2)       # x2 = LN(x1 + f)
            x1 = nn
            g[q + "W2"] += flat(np.maximum(z, 0)).T @ flat(dx2); g[q + "b2"] += dx2.sum((0, 1))
            dz = (dx2 @ p[q + "W2"].T) * (z > 0)
            g[q + "W1"] += flat(x1).T @ flat(dz); g[q + "b1"] += dz.sum((0, 1))
            dx1 = dx2 + dz @ p[q + "W1"].T
            dsum, g[q + "g1"], g[q + "be1"] = ln_indietro(dx1, cn1)     # x1 = LN(x + a)
            dx = dsum + att_indietro(dsum, ca, p, q, h, g, scala)
        else:
            g[q + "W2"] += flat(np.maximum(z, 0)).T @ flat(dx); g[q + "b2"] += dx.sum((0, 1))
            dz = (dx @ p[q + "W2"].T) * (z > 0)
            g[q + "W1"] += flat(nn).T @ flat(dz); g[q + "b1"] += dz.sum((0, 1))
            dn2, g[q + "g2"], g[q + "be2"] = ln_indietro(dz @ p[q + "W1"].T, cn2)
            dx1 = dx + dn2
            dn1 = att_indietro(dx1, ca, p, q, h, g, scala)
            dn1, g[q + "g1"], g[q + "be1"] = ln_indietro(dn1, cn1)
            dx = dx1 + dn1
    np.add.at(g["E"], tok, dx)
    g["P"][:T] += dx.sum(0)
    return perdita, g

Ottimizzatore

Adam con i valori consueti del paper originale dell'ottimizzatore (β₁ = 0,9, β₂ = 0,999, ε = 10⁻⁸): una stima della media e una del quadrato del gradiente, entrambe corrette dal bias iniziale. È l'unico pezzo che non è specifico del Transformer.

def stato_adam(p):
    return {k: [np.zeros_like(v), np.zeros_like(v)] for k, v in p.items()}

def adam(p, g, stato, lr, t, b1=0.9, b2=0.999, eps=1e-8):
    for k in p:
        m, v = stato[k]
        m *= b1; m += (1 - b1) * g[k]
        v *= b2; v += (1 - b2) * g[k] ** 2
        p[k] -= lr * (m / (1 - b1 ** t)) / (np.sqrt(v / (1 - b2 ** t)) + eps)

Con il modello in mano si può fare il conto che la formula prometteva. Il listato crea un blocco per tre larghezze (64, 256, 512), somma le dimensioni di ogni tensore il cui nome comincia per b0. e le confronta con 12d² + 13d; poi, per d = 512, scompone il totale voce per voce, in modo da vedere dove stanno i parametri.

def parametri_blocco(p, l):
    """Somma le dimensioni effettive di tutti i tensori del blocco l."""
    q = f"b{l}."
    return sum(v.size for k, v in p.items() if k.startswith(q))

rng = np.random.default_rng(1)
print(" d    d_ff   formula 12d^2+13d   conteggio reale")
for d in (64, 256, 512):
    p = inizializza(vocab=32, d=d, h=4, n_strati=1, d_ff=4 * d, t_max=16, rng=rng)
    print(f"{d:4d}  {4*d:6d}   {12*d*d + 13*d:17d}   {parametri_blocco(p, 0):15d}")

d = 512
voci = [("W_Q, W_K, W_V, W_O", 4 * d * d), ("bias delle quattro proiezioni", 4 * d),
        ("W_1 (d -> 4d)", 4 * d * d), ("b_1", 4 * d),
        ("W_2 (4d -> d)", 4 * d * d), ("b_2", d),
        ("due layer norm (gamma e beta)", 4 * d)]
print("scomposizione per d = 512:")
for nome, n in voci:
    print(f"  {nome:32s} {n:10d}")
print(f"  {'TOTALE':32s} {sum(n for _, n in voci):10d}")
 d    d_ff   formula 12d^2+13d   conteggio reale
  64     256               49984             49984
 256    1024              789760            789760
 512    2048             3152384           3152384
scomposizione per d = 512:
  W_Q, W_K, W_V, W_O                  1048576
  bias delle quattro proiezioni          2048
  W_1 (d -> 4d)                       1048576
  b_1                                    2048
  W_2 (4d -> d)                       1048576
  b_2                                     512
  due layer norm (gamma e beta)          2048
  TOTALE                              3152384

Cosa dicono i numeri. Conto chiuso e inventario coincidono alla cifra: 49 984, 789 760, 3 152 384. La scomposizione per d = 512 mostra tre blocchi da 1 048 576 parametri — attenzione, W_1, W_2 — più 6 656 parametri di bias e scale, lo 0,21% del totale. Ecco perché nelle stime a mente si scrive 12d² e si ignora il resto.

La regola della catena l'hai già derivata nella lezione sulla backpropagation; qui conta come si verifica un gradiente scritto a mano. Si perturba un parametro di ±ε e si confronta (L(θ+ε) − L(θ−ε)) / 2ε con il gradiente analitico. Il listato lo fa su un modello minuscolo (d = 16, 2 strati, 3 sequenze), sulle prime sei coordinate di ogni tensore dove il gradiente supera 10⁻⁶, e riporta lo scarto relativo peggiore.

def scarto_massimo(post_ln, eps, seme=0):
    """Confronta il gradiente analitico con la differenza finita centrata, tensore per tensore."""
    rng = np.random.default_rng(seme)
    p = inizializza(21, 16, 4, 2, 64, 13, rng)
    tok = rng.integers(0, 21, size=(3, 12))
    ber = rng.integers(0, 21, size=(3, 12))
    pesa = np.zeros((3, 12)); pesa[:, 6:] = 1.0
    _, g = perdita_e_gradienti(p, tok, ber, pesa, 2, 4, True, post_ln)
    peggiore = 0.0
    for nome in p:
        for i in np.argwhere(np.abs(g[nome]) > 1e-6)[:6]:       # solo dove il gradiente e' misurabile
            i = tuple(i); v0 = p[nome][i]
            p[nome][i] = v0 + eps; piu = perdita_e_gradienti(p, tok, ber, pesa, 2, 4, True, post_ln)[0]
            p[nome][i] = v0 - eps; meno = perdita_e_gradienti(p, tok, ber, pesa, 2, 4, True, post_ln)[0]
            p[nome][i] = v0
            numerico = (piu - meno) / (2 * eps)
            peggiore = max(peggiore, abs(numerico - g[nome][i]) / max(abs(numerico), abs(g[nome][i])))
    return peggiore

for post_ln in (False, True):
    for eps in (1e-6, 1e-4):
        print(f"{'post-LN' if post_ln else 'pre-LN '}  eps={eps:.0e}  scarto relativo massimo {scarto_massimo(post_ln, eps):.2e}")
pre-LN   eps=1e-06  scarto relativo massimo 2.56e-06
pre-LN   eps=1e-04  scarto relativo massimo 9.94e-02
post-LN  eps=1e-06  scarto relativo massimo 4.90e-05
post-LN  eps=1e-04  scarto relativo massimo 6.60e-03

Cosa dicono i numeri. Con ε = 10⁻⁶ lo scarto relativo massimo è 2,56 · 10⁻⁶ per il pre-LN e 4,90 · 10⁻⁵ per il post-LN: il gradiente scritto a mano è giusto in entrambe le varianti, nei limiti di un campione di coordinate. Con ε = 10⁻⁴ lo scarto sale a 9,94 · 10⁻² e 6,60 · 10⁻³. Se il controllo peggiora quando aumenti ε, il problema è ε: la differenza finita è un'approssimazione di una derivata, e se la perturbazione attraversa uno spigolo della ReLU smette di misurarla. Quest'ultima è la spiegazione coerente con il comportamento, non l'ho verificata coordinata per coordinata.

Il conto che decide tutto: O(T²)

La matrice dei punteggi ha T² elementi, e questa frase spiega la maggior parte dell'ingegneria dei modelli linguistici recenti. Il tempo di Q Kᵀ e di P V è O(T² · d), e la memoria per tenere P nella passata in avanti — serve in quella all'indietro — è O(T²) per testa e per elemento del lotto.

import numpy as np, time

def attenzione_causale(Q, K, V):
    T, d_k = Q.shape
    S = (Q @ K.T) / np.sqrt(d_k)
    S = np.where(np.tril(np.ones((T, T), dtype=bool)), S, -np.inf)
    return softmax(S, asse=-1) @ V

def tempo_minimo(Q, K, V, prove):
    """Il MINIMO su piu' prove: il carico della macchina puo' allungare un tempo,
    non accorciarlo, quindi il minimo e' la stima meno inquinata."""
    attenzione_causale(Q, K, V)                        # giro di riscaldamento
    misure = []
    for _ in range(prove):
        t0 = time.perf_counter()
        attenzione_causale(Q, K, V)
        misure.append((time.perf_counter() - t0) * 1000)
    return min(misure)

rng = np.random.default_rng(0)
d_k = 64
print(" T     prove   tempo (ms)   matrice TxT (MiB)   Q,K,V insieme (MiB)")
tempi, lunghezze, prove = [], (128, 512, 2048), (400, 60, 12)
for T, n in zip(lunghezze, prove):
    Q, K, V = (rng.normal(size=(T, d_k)) for _ in range(3))
    ms = tempo_minimo(Q, K, V, n)
    tempi.append(ms)
    print(f"{T:5d} {n:7d} {ms:12.2f} {T*T*8/2**20:18.2f} {3*T*d_k*8/2**20:21.2f}")

pendenza = np.polyfit(np.log(np.array(lunghezze, float)), np.log(np.array(tempi)), 1)[0]
print(f"pendenza di log(tempo) su log(T): {pendenza:.2f}  (una legge T^2 da' 2)")
for i in range(1, 3):
    print(f"  da T={lunghezze[i-1]} a T={lunghezze[i]}: lunghezza x4, "
          f"tempo x{tempi[i]/tempi[i-1]:.1f}, memoria della matrice x16")
 T     prove   tempo (ms)   matrice TxT (MiB)   Q,K,V insieme (MiB)
  128     400         1.46               0.12                  0.19
  512      60        23.47               2.00                  0.75
 2048      12       339.52              32.00                  3.00
pendenza di log(tempo) su log(T): 1.97  (una legge T^2 da' 2)
  da T=128 a T=512: lunghezza x4, tempo x16.1, memoria della matrice x16
  da T=512 a T=2048: lunghezza x4, tempo x14.5, memoria della matrice x16

Cosa dicono i numeri. L'uscita riportata è una di tre esecuzioni consecutive, su una macchina condivisa con altri processi, e per questo i tempi sono intervalli: 1,22-1,48 ms a T = 128, 23,5-25,3 ms a 512, 340-404 ms a 2048; la pendenza nel piano logaritmico è fra 1,97 e 2,06, contro il 2 esatto di una legge quadratica. Quadruplicare la lunghezza moltiplica il tempo per 16,1-20,8 nel primo salto e per 14,5-16,0 nel secondo. Che il secondo resti sotto 16 può dipendere dalla cache del processore, ma non l'ho verificato. L'ultima colonna è la memoria di Q, K e V insieme, 3,00 MiB: la matrice dei punteggi pesa più di dieci volte i tensori che l'hanno generata.

I numeri sono in float64; in float16 si dividono per quattro, ma la dipendenza resta quadratica — a T = 32 768 una sola testa occupa 2 GiB per la matrice. Per questo esistono implementazioni che non materializzano mai P ma la calcolano a blocchi, tenendo solo le somme parziali del softmax: è ciò che fa torch.nn.functional.scaled_dot_product_attention col suo nucleo fuso, come descrive la documentazione di PyTorch. La formula è la stessa (il confronto numerico è più sopra); cambia dove vivono i numeri intermedi.

La cache chiave-valore: cosa si conserva, e quanto pesa

In generazione il modello produce un token per volta e rifà la passata in avanti su tutto quello che ha scritto, ma chiavi e valori delle posizioni già viste non cambiano: la k₇ e la v₇ del token in posizione 7 sono le stesse quando il modello sta generando il token 300. L'unica cosa nuova al passo t è la query qₜ.

Da qui la formula della memoria della cache chiave-valore, derivata da cosa si conserva davvero: due matrici, da cui il fattore 2; per ogni strato, perché ogni strato ha le sue proiezioni; per ogni testa di chiave-valore; d_testa componenti per ciascuna delle T posizioni, per il lotto e per i byte del tipo. Nota cosa non c'è: non le query, non i pesi di attenzione, non le attivazioni della feed-forward.

Il listato costruisce un modello giocattolo a quattro strati, solo attenzione e residuo, con pesi casuali: serve a misurare il tempo, non la qualità. Le due funzioni producono lo stesso risultato, e lo dimostra lo scarto massimo stampato. rinormalizza compare soltanto perché, riportando in ingresso l'uscita al posto di un token, le somme residue crescerebbero a ogni passo: non fa parte del meccanismo della cache. Il prompt è di 32 righe di rumore, e si generano 128 vettori.

import time
import numpy as np

rng = np.random.default_rng(5)
d, h, n_strati = 256, 8, 4
d_h = d // h
Ws = [tuple(rng.normal(size=(d, d)) / np.sqrt(d) for _ in range(4)) for _ in range(n_strati)]

def dividi(M):                                      # (T, d) -> (h, T, d_h)
    return M.reshape(M.shape[0], h, d_h).transpose(1, 0, 2)

def prefisso_intero(X):
    """Nessuna cache: a ogni token si ricalcolano K e V di tutto il prefisso."""
    for Wq, Wk, Wv, Wo in Ws:
        T = X.shape[0]
        Q, K, V = dividi(X @ Wq), dividi(X @ Wk), dividi(X @ Wv)
        S = (Q @ K.transpose(0, 2, 1)) / np.sqrt(d_h)
        S = np.where(np.tril(np.ones((T, T), dtype=bool)), S, -np.inf)
        C = (softmax(S, -1) @ V).transpose(1, 0, 2).reshape(T, d)
        X = X + C @ Wo
    return X

def un_token(x, cache):
    """Con cache: del token nuovo si calcolano q, k, v; k e v si appendono."""
    for l, (Wq, Wk, Wv, Wo) in enumerate(Ws):
        q = (x @ Wq).reshape(h, 1, d_h)
        k = (x @ Wk).reshape(h, 1, d_h)
        v = (x @ Wv).reshape(h, 1, d_h)
        K = k if cache[l][0] is None else np.concatenate([cache[l][0], k], axis=1)
        V = v if cache[l][1] is None else np.concatenate([cache[l][1], v], axis=1)
        cache[l] = [K, V]
        S = (q @ K.transpose(0, 2, 1)) / np.sqrt(d_h)   # (h, 1, t): nessuna maschera serve
        c = (softmax(S, -1) @ V).reshape(1, d)
        x = x + c @ Wo
    return x

def memoria_cache(n_strati, n_teste_kv, d_testa, T, batch, byte):
    return 2 * n_strati * n_teste_kv * d_testa * T * batch * byte

def rinormalizza(x):
    """Tiene il giocattolo in una scala numerica stabile: senza, le somme residue crescono."""
    return x / np.linalg.norm(x, axis=-1, keepdims=True) * np.sqrt(d)

def genera_senza_cache(prompt, n_nuovi):
    X = prompt.copy()
    for _ in range(n_nuovi):
        X = np.vstack([X, rinormalizza(prefisso_intero(X)[-1])])
    return X

def genera_con_cache(prompt, n_nuovi):
    cache = [[None, None] for _ in Ws]
    for riga in prompt[:-1]:                         # riempie la cache con il prompt
        un_token(riga, cache)
    x, righe = prompt[-1], list(prompt)
    for _ in range(n_nuovi):
        x = rinormalizza(un_token(x, cache)[0])
        righe.append(x)
    return np.array(righe)

prompt = rng.normal(size=(32, d))
t0 = time.perf_counter(); senza = genera_senza_cache(prompt, 128); t_senza = time.perf_counter() - t0
t0 = time.perf_counter(); con = genera_con_cache(prompt, 128); t_con = time.perf_counter() - t0
print(f"{n_strati} strati, {h} teste, d = {d}, prompt 32 token, 128 token generati")
print(f"  scarto massimo fra le due generazioni: {np.abs(senza - con).max():.1e}")
print(f"  senza cache: {t_senza:6.2f} s")
print(f"  con cache:   {t_con:6.2f} s")
print(f"  rapporto:    {t_senza / t_con:6.1f}x")
print("memoria = 2 * n_strati * n_teste_kv * d_testa * T * batch * byte")
print(f"  questo modello, T = 160, batch 1, float64: {memoria_cache(4, 8, 32, 160, 1, 8) / 2**20:.2f} MiB")
for kv in (32, 1):
    print(f"  32 strati, d_testa 128, T = 4096, batch 1, float16, {kv:2d} coppie K,V: "
          f"{memoria_cache(32, kv, 128, 4096, 1, 2) / 2**30:.3f} GiB")
4 strati, 8 teste, d = 256, prompt 32 token, 128 token generati
  scarto massimo fra le due generazioni: 6.9e-15
  senza cache:  14.29 s
  con cache:     0.19 s
  rapporto:      77.2x
memoria = 2 * n_strati * n_teste_kv * d_testa * T * batch * byte
  questo modello, T = 160, batch 1, float64: 2.50 MiB
  32 strati, d_testa 128, T = 4096, batch 1, float16, 32 coppie K,V: 2.000 GiB
  32 strati, d_testa 128, T = 4096, batch 1, float16,  1 coppie K,V: 0.062 GiB

Cosa dicono i numeri. Lo scarto massimo fra le due generazioni è 6,9 · 10⁻¹⁵: la cache non cambia il risultato, cambia il costo. Sul modello giocattolo generare 128 vettori dopo un prompt di 32 costa 10,0-17,6 s senza cache e 0,15-0,26 s con la cache, in nove esecuzioni su una macchina condivisa: il rapporto va da circa 40 a circa 110 (la riga qui sopra ne mostra una, 77). Cresce con la lunghezza generata, perché senza cache a ogni passo si ricalcolano chiavi e valori di tutto il prefisso. Su una configurazione realistica (32 strati, 32 coppie di chiave-valore, d_testa = 128, T = 4096, lotto 1, float16) la formula dà 2,000 GiB; con una sola coppia K, V condivisa, l'attenzione multi-query di Shazeer (arXiv:1911.02150), 0,062 GiB. Sono formule, non misure su un modello vero.

Esercizio 2

Il conto che decide il budget

Stai dimensionando un decoder con d = 512, h = 8 teste, 24 blocchi, d_testa = 64. Carta e penna; le potenze di due aiutano.

  1. Quanti parametri ha un blocco? E tutti i 24 blocchi insieme?
  2. In generazione con T = 8192, lotto di 4 richieste e float16, quanto occupa la cache chiave-valore?
  3. Confronta le due cifre in byte, assumendo anche i pesi in float16. Che cosa decide la memoria della tua macchina?

Anche qui l'ordine conta: il conto dei parametri della prima domanda entra nel confronto in byte della terza, e la cache della seconda dipende dal contesto e dal lotto, che i parametri non conoscono. Scrivi sempre le unità accanto ai numeri, o perderai un fattore due fra byte e parametri.

Mostra la soluzione

1. 12 · 512² + 13 · 512 = 3 145 728 + 6 656 = 3 152 384 parametri per blocco; per 24 blocchi 75 657 216, cioè 144,3 MiB in float16.

2. 2 · 24 · 8 · 64 · 8192 · 4 · 2 = 1 610 612 736 byte, cioè 1,5 GiB esatti.

3. La cache pesa 10,6 volte i pesi dei blocchi: quando pianifichi la memoria di un servizio di generazione, il modello è la voce piccola e il contesto per le richieste simultanee è quella grande. È anche il motivo dell'attenzione multi-query (arXiv:1911.02150), che taglia quel numero per otto: 0,188 GiB.

L'esperimento: copiare e ordinare sequenze

Per vedere l'attenzione all'opera serve un compito scelto bene, e la scelta è metodologica. Niente testo: la perplessità scende e non sai dire se il modello ha imparato la struttura o la frequenza delle parole. Niente immagini: non c'è un ordine da ricostruire. Il compito giusto è sintetico e verificabile a vista: dato un alfabeto di 20 simboli e un separatore, copiare la sequenza dopo il separatore oppure ordinarla. Richiede esattamente quello che l'attenzione sa fare: prendere un valore in una posizione precisa.

I simboli sono già numeri interi, cioè già token: da dove vengono lo ha mostrato la lezione sulla tokenizzazione. La sequenza data al modello è [s₁ … s_L, SEP, risposta₁ … risposta_L], lunga 2L + 1, e il modello predice il token successivo a ogni posizione. La funzione di costo è l'entropia incrociata della lezione sull'entropia incrociata e la divergenza KL, calcolata solo sulle L posizioni della risposta.

import numpy as np

ALFABETO = 20
SEP = ALFABETO            # il 21-esimo simbolo separa domanda e risposta
VOCAB = ALFABETO + 1

def lotto(rng, B, L, compito):
    s = rng.integers(0, ALFABETO, size=(B, L))
    risposta = np.sort(s, axis=1) if compito == "ordina" else s
    seq = np.concatenate([s, np.full((B, 1), SEP), risposta], axis=1)   # (B, 2L+1)
    tok = seq[:, :-1]                      # ingresso:  posizioni 0 .. 2L-1
    ber = seq[:, 1:]                       # bersaglio: posizioni 1 .. 2L
    pesa = np.zeros_like(tok, dtype=float)
    pesa[:, L:] = 1.0                      # la perdita conta solo la risposta, L token
    return tok, ber, pesa

def addestra(compito, L, passi=1500, B=32, d=32, h=4, n_strati=2, lr=3e-3, seme=0,
             scala=True, post_ln=False, posizione=True):
    rng = np.random.default_rng(seme)
    p = inizializza(VOCAB, d, h, n_strati, 4 * d, 2 * L + 1, rng)
    if not posizione:
        p["P"][:] = 0.0                    # tabella a zero...
    stato = stato_adam(p)
    for t in range(1, passi + 1):
        tok, ber, pesa = lotto(rng, B, L, compito)
        perdita, g = perdita_e_gradienti(p, tok, ber, pesa, n_strati, h, scala, post_ln)
        if not posizione:
            g["P"][:] = 0.0                # ...e mai aggiornata
        adam(p, g, stato, lr, t)
    return p, perdita

def valuta(p, compito, L, n=500, h=4, n_strati=2, seme=999, scala=True, post_ln=False):
    """Generazione vera: la risposta si costruisce un token alla volta, senza bersagli."""
    rng = np.random.default_rng(seme)
    s = rng.integers(0, ALFABETO, size=(n, L))
    atteso = np.sort(s, axis=1) if compito == "ordina" else s
    corrente = np.concatenate([s, np.full((n, 1), SEP)], axis=1)
    for _ in range(L):
        logit, _ = avanti(p, corrente, n_strati, h, scala, post_ln)
        corrente = np.concatenate([corrente, logit[:, -1, :].argmax(-1)[:, None]], axis=1)
    prodotto = corrente[:, L + 1:]
    return (prodotto == atteso).mean(), (prodotto == atteso).all(axis=1).mean()

Il compito è dichiarato, il modello è quello scritto sopra; resta il lavoro di contabilità. Il secondo listato lancia i 23 addestramenti (due compiti, tre lunghezze, tre semi, più cinque ablazioni) in processi separati, un thread ciascuno, e stampa le tabelle di questa e delle prossime due sezioni. Ho verificato su un caso che il risultato di un addestramento non dipende dal modo in cui lo si lancia: la perdita finale di copia con L = 6 e seme 0 è identica, alla cifra, lanciata da sola e dentro il pool. Da solo, un addestramento dura circa 58 s.

import os
os.environ["OMP_NUM_THREADS"] = "1"              # un processo per addestramento, un thread ciascuno
os.environ["OPENBLAS_NUM_THREADS"] = "1"
from multiprocessing import Pool

SEQUENZA_MAPPA = np.array([[17, 3, 0, 12, 7, 9, SEP, 17, 3, 0, 12, 7]])

def prova(spec):
    nome, compito, L, seme, opzioni = spec
    p, perdita = addestra(compito, L, seme=seme, **opzioni)
    ok = {k: opzioni[k] for k in ("scala", "post_ln") if k in opzioni}
    per_token, per_sequenza = valuta(p, compito, L, **ok)
    esito = {"perdita": perdita, "token": per_token, "sequenza": per_sequenza}
    if nome in ("completo", "senza posizione") and compito == "copia" and L == 6:
        rng = np.random.default_rng(999)         # le stesse 500 sequenze di valuta()
        s = rng.integers(0, ALFABETO, size=(500, L))
        corrente = np.concatenate([s, np.full((500, 1), SEP)], axis=1)
        for _ in range(L):
            logit, _ = avanti(p, corrente, 2, 4)
            corrente = np.concatenate([corrente, logit[:, -1, :].argmax(-1)[:, None]], axis=1)
        giusta = (corrente[:, L + 1:] == s).all(axis=1)
        distinti = np.array([len(set(r)) == L for r in s])
        esito["distinti"] = (int(distinti.sum()), giusta[distinti].mean(), giusta[~distinti].mean())
    if nome == "completo" and compito == "copia" and L == 6 and seme == 0:
        _, (_, cache, _, _) = avanti(p, SEQUENZA_MAPPA, 2, 4)
        esito["attenzione"] = [c[0][4][0] for c in cache]          # per strato: (teste, T, T)
    return (nome, compito, L, seme), esito

if __name__ == "__main__":
    lavori = [("completo", c, L, seme, {}) for c in ("copia", "ordina") for L in (4, 6, 8) for seme in (0, 1, 2)]
    lavori += [("senza scala", "copia", 6, 0, {"scala": False}),
               ("senza posizione", "copia", 6, 0, {"posizione": False}),
               ("post-LN", "copia", 6, 0, {"post_ln": True}),
               ("pre-LN, lr alto", "copia", 6, 0, {"lr": 1e-2}),
               ("post-LN, lr alto", "copia", 6, 0, {"lr": 1e-2, "post_ln": True})]
    with Pool(16) as pool:
        r = dict(pool.map(prova, lavori, chunksize=1))

    n_par = sum(v.size for v in inizializza(VOCAB, 32, 4, 2, 128, 13, np.random.default_rng(0)).values())
    print("Configurazione: d=32, 4 teste (d_testa=8), 2 strati, d_ff=128, batch 32,")
    print("Adam lr=3e-3, 1500 passi, alfabeto 20 simboli + separatore (vocabolario 21).")
    print(f"Parametri del modello: {n_par}\n")
    print("compito   L   accuratezza per token      accuratezza per sequenza")
    print("              (media su 3 semi)          (media su 3 semi)")
    for compito in ("copia", "ordina"):
        for L in (4, 6, 8):
            e = [r[("completo", compito, L, s)] for s in range(3)]
            t, q = [x["token"] for x in e], [x["sequenza"] for x in e]
            print(f"{compito:<8} {L:>2}   {np.mean(t):.3f}  (min {min(t):.3f})        "
                  f"{np.mean(q):.3f}  (min {min(q):.3f})")

    print("\nAblazioni sul compito di copia, L = 6, seme 0, stessi iperparametri salvo dove indicato:")
    for nome in ("completo", "senza scala", "senza posizione", "post-LN", "pre-LN, lr alto", "post-LN, lr alto"):
        e = r[(nome, "copia", 6, 0)]
        print(f"  {nome:<18} perdita finale {e['perdita']:.4f}   token {e['token']:.3f}   sequenza {e['sequenza']:.3f}")

    print("\nCopia senza posizione, 500 sequenze di prova:")
    n_dist, a_dist, a_rip = r[("senza posizione", "copia", 6, 0)]["distinti"]
    print(f"  sulle sequenze a simboli tutti distinti:       {a_dist:.3f}   ({n_dist} casi)")
    print(f"  sulle sequenze con almeno un simbolo ripetuto: {a_rip:.3f}   ({500 - n_dist} casi)")
    print(f"  quota di sequenze a simboli distinti: misurata {n_dist / 500:.3f}, "
          f"attesa 20*19*18*17*16*15/20^6 = {20*19*18*17*16*15 / 20**6:.3f}")

    print("\nsequenza: [17 3 0 12 7 9 SEP 17 3 0 12 7]")
    print("le 6 posizioni di uscita sono 6..11 e devono emettere i simboli che stanno alle chiavi 0..5\n")
    print("strato  testa   chiave piu' pesata per le uscite 6..11   scarto             peso medio")
    for l, pesi in enumerate(r[("completo", "copia", 6, 0)]["attenzione"], start=1):
        for t in range(4):
            righe = pesi[t][6:]
            scelta = righe.argmax(-1)
            scarto = scelta - np.arange(6, 12)
            fisso = f"{scarto[0]:+d} costante" if (scarto == scarto[0]).all() else "variabile"
            print(f"{l:>4} {t:>6}     {str(scelta):<36} {fisso:<18} {righe.max(-1).mean():.3f}")
Configurazione: d=32, 4 teste (d_testa=8), 2 strati, d_ff=128, batch 32,
Adam lr=3e-3, 1500 passi, alfabeto 20 simboli + separatore (vocabolario 21).
Parametri del modello: 27253

compito   L   accuratezza per token      accuratezza per sequenza
              (media su 3 semi)          (media su 3 semi)
copia     4   1.000  (min 1.000)        1.000  (min 1.000)
copia     6   1.000  (min 1.000)        1.000  (min 1.000)
copia     8   1.000  (min 1.000)        1.000  (min 1.000)
ordina    4   0.995  (min 0.989)        0.981  (min 0.958)
ordina    6   0.993  (min 0.989)        0.970  (min 0.958)
ordina    8   0.987  (min 0.979)        0.951  (min 0.918)

Cosa fa il codice. Il tensore pesa è la maschera della perdita, non una maschera di attenzione. L'addestramento genera un lotto nuovo a ogni passo, quindi non c'è un insieme da sovradattare: lo spazio delle sequenze è 20⁶ ≈ 6,4 · 10⁷ e il modello ne vede 48 000, lo 0,075%. La valutazione è generazione vera, con il modello che riceve la propria uscita precedente.

Cosa dicono i numeri. Sulla copia il modello è perfetto a tutte e tre le lunghezze, 1,000 per token e per sequenza su tre semi: 27 253 parametri risolvono esattamente il compito. Sull'ordinamento l'accuratezza per token resta fra 0,987 e 0,995, ma quella per sequenza scende da 0,981 a 0,951, col minimo a 0,918: ordinare è un confronto fra tutte le coppie, e un errore su sei token basta a sbagliare la sequenza. 0,987 e 0,951 sono lo stesso modello a L = 8: è anche una lezione sulle metriche.

La scala è ridotta, e la dichiaro. L'addestramento gira in NumPy su CPU con la passata all'indietro scritta a mano: non è un modello, è un banco di prova. d = 32, 4 teste da 8, 2 blocchi, d_ff = 128, lotti da 32, 1500 passi. Tre semi sono pochi: il minimo mostrato accanto alla media serve a ricordarlo. Meglio un esperimento piccolo eseguito che uno grande descritto.

Tre ablazioni, e due che non mostrano niente

Un esperimento serve soprattutto a togliere pezzi. Rilanciando lo stesso addestramento sulla copia con L = 6, seme 0 e tutto il resto invariato, un pezzo alla volta:

L'uscita delle ablazioni è nello stesso listato della sezione precedente, e la riporto qui. Ogni riga è un addestramento con seme 0; «senza posizione» azzera la tabella delle posizioni e non la aggiorna mai.

Ablazioni sul compito di copia, L = 6, seme 0, stessi iperparametri salvo dove indicato:
  completo           perdita finale 0.0002   token 1.000   sequenza 1.000
  senza scala        perdita finale 0.0002   token 1.000   sequenza 1.000
  senza posizione    perdita finale 0.7342   token 0.656   sequenza 0.204
  post-LN            perdita finale 0.0002   token 1.000   sequenza 1.000
  pre-LN, lr alto    perdita finale 0.0000   token 1.000   sequenza 1.000
  post-LN, lr alto   perdita finale 0.0000   token 1.000   sequenza 1.000

Copia senza posizione, 500 sequenze di prova:
  sulle sequenze a simboli tutti distinti:       0.246   (224 casi)
  sulle sequenze con almeno un simbolo ripetuto: 0.170   (276 casi)
  quota di sequenze a simboli distinti: misurata 0.448, attesa 20*19*18*17*16*15/20^6 = 0.436

Togliere la posizione rompe il compito. La perdita passa da 0,0002 a 0,7342 e l'accuratezza per sequenza da 1,000 a 0,204. Resta però 0,656 per token, ben sopra 1/20 = 0,05: il modello non è completamente cieco alla posizione. Una spiegazione plausibile è che la maschera causale la lasci filtrare, perché la riga t vede esattamente t + 1 chiavi, e per i decoder senza codifica posizionale esiste letteratura in questo senso (Haviv e colleghi, arXiv:2203.16634). Non ho messo alla prova questa spiegazione sul mio banco.

Un'indicazione, non un risultato. Sulle sequenze a simboli tutti distinti il modello senza posizione indovina 0,246 contro 0,170 con una ripetizione (224 e 276 casi). La direzione è quella attesa, perché un simbolo ripetuto è un punto in cui il contenuto da solo non basta a decidere il posto, ma è un solo seme e due sottoinsiemi dello stesso campione: non lo uso per sostenere niente. La quota misurata di sequenze a simboli distinti, 0,448, è quella attesa a meno del campionamento (0,436).

Le altre due ablazioni non mostrano niente, e questo è un risultato da riportare. Togliere la divisione per √d_k lascia l'accuratezza a 1,000: non contraddice la sezione sulla scala, contraddice l'idea che il problema si veda a qualunque dimensione. Qui d_testa = 8 e la deviazione dei punteggi grezzi è √8 ≈ 2,83: a d_k = 8 la sensibilità scende da 0,96 a 0,72, non di un ordine di grandezza. Il regime in cui la scala decide comincia a d_k = 64, come nella tabella di sopra (sensibilità 0,29 contro 0,96), e il mio banco non ci arriva.

Lo stesso vale per pre-LN contro post-LN: indistinguibili a lr = 3 · 10⁻³ e a lr = 10⁻² senza riscaldamento, perché due blocchi non sono una pila profonda e quello di arXiv:2002.04745 è un problema di profondità. Quindi riporto la fonte e non un numero mio: una misura su due blocchi non sostiene una raccomandazione su ventiquattro.

La mappa di attenzione, e perché non è una spiegazione

Sul compito di copia possiamo guardare direttamente la matrice dei pesi: se il modello ha imparato il meccanismo giusto, la posizione che emette il k-esimo simbolo della risposta dovrebbe pesare la posizione che contiene il k-esimo simbolo della domanda, a distanza L indietro.

La tabella è l'ultima parte dell'uscita del listato dell'esperimento: per ognuna delle otto teste, la chiave con il peso più alto in ciascuna delle sei posizioni di uscita, lo scarto dalla posizione attesa (le uscite 6..11 devono guardare le chiavi 0..5, cioè sei indietro) e il peso massimo medio.

sequenza: [17 3 0 12 7 9 SEP 17 3 0 12 7]
le 6 posizioni di uscita sono 6..11 e devono emettere i simboli che stanno alle chiavi 0..5

strato  testa   chiave piu' pesata per le uscite 6..11   scarto             peso medio
   1      0     [0 1 2 3 4 5]                        -6 costante        0.835
   1      1     [0 3 6 6 6 5]                        variabile          0.416
   1      2     [0 1 2 3 4 5]                        -6 costante        0.996
   1      3     [0 1 2 3 4 5]                        -6 costante        0.767
   2      0     [ 0  1  0  4 10 11]                  variabile          0.355
   2      1     [ 0  7  8  5  0 11]                  variabile          0.369
   2      2     [ 0  7  8  4 10 11]                  variabile          0.365
   2      3     [ 0  1  0  4 10 11]                  variabile          0.391
I pesi di attenzione sul compito di copia, primo strato, testa 2 PESI DI ATTENZIONE - strato 1, testa 2, compito di copia (L = 6) chiavi (j): posizione e simbolo 0 17 1 3 2 0 3 12 4 7 5 9 6 SEP 7 17 8 3 9 0 10 12 11 7 query (i) 0 (17) 1 (3) 2 (0) 3 (12) 4 (7) 5 (9) 6 (SEP) 7 (17) 8 (3) 9 (0) 10 (12) 11 (7) separatore Righe 0-5: il modello legge la sequenza. Righe 6-11: il modello risponde, e ogni riga mette da 0,99 a 1,00 di peso sulla chiave esattamente sei posizioni prima. Il massimo medio, sulle sei righe di risposta, vale 0,996: questa testa ha imparato «guarda sei indietro», e nient'altro. peso: da 0 a 1
La diagonale spostata di sei posizioni e il meccanismo della copia, visibile a occhio: la posizione che deve emettere il k-esimo simbolo della risposta pesa quasi tutto sulla posizione che contiene il k-esimo simbolo della domanda.

Cosa dicono i numeri. Tre delle quattro teste del primo strato (0, 2 e 3) hanno imparato la cosa giusta: per ognuna delle sei posizioni di uscita il peso massimo cade esattamente sei posizioni indietro, scarto −6 costante, con peso medio 0,835, 0,996 e 0,767. La testa 2, quella disegnata, è un puntatore quasi puro; la testa 1 non lo è (0,416, scarto variabile). Le teste del secondo strato sono diffuse, con massimo medio fra 0,355 e 0,391: il compito, in questo addestramento, lo risolve il primo strato. Un altro seme potrebbe distribuire il lavoro in un altro modo, ed è una ragione in più per non generalizzare dalla figura.

E qui arriva l'avvertenza onesta. Una mappa leggibile su un compito artificiale non autorizza a leggere le mappe come spiegazioni su un compito reale. Jain e Wallace (Attention is not Explanation, arXiv:1902.10186) costruiscono distribuzioni molto diverse da quella appresa che producono la stessa predizione: se due mappe contraddittorie portano allo stesso risultato, nessuna delle due è la ragione del risultato. Wiegreffe e Pinter (Attention is not not Explanation, arXiv:1908.04626) replicano distinguendo spiegazione plausibile e fedele.

La regola da portare via. I pesi di attenzione sono una variabile interna del calcolo, non una motivazione. Qui la mappa è leggibile perché il compito ha una sola soluzione meccanica; su un compito reale la stessa immagine può essere una storia convincente e sbagliata. Per sostenere perché un modello ha deciso servono interventi sull'ingresso, non una mappa di calore.

Trade-off

  • Cammino corto contro costo quadratico. Due token qualunque a distanza 1 e passata parallelizzata, ma T² in tempo e memoria: pendenza circa 2 misurata (1,97-2,06). Mitigazione: finestre locali, attenzione a blocchi, nuclei che non materializzano la matrice — tutte rinunciano a qualche coppia.
  • Contesto lungo contro memoria del servizio. Raddoppiare T quadruplica la matrice e raddoppia la cache. Nell'esercizio 2 la cache pesa 10,6 volte i pesi: è il contesto a decidere quante richieste simultanee puoi servire.
  • Molte teste contro teste ricche. A parità di parametri, h grande dà molti schemi in sottospazi piccoli e h piccolo pochi schemi in sottospazi ricchi; l'equilibrio empirico è d/h fra 32 e 128.
  • Multi-query contro capacità. Condividere una sola coppia K, V divide la cache per il numero di teste — sopra, da 2,000 a 0,062 GiB — ma le teste perdono i sottospazi distinti. Si adotta quando il collo di bottiglia è la memoria.
  • Posizione appresa contro calcolata. La tabella appresa è una riga di codice, ma oltre T_max non esiste; sinusoidale e RoPE non costano parametri e restano definite oltre la lunghezza vista.
  • Pre-LN contro post-LN. Il pre-LN tollera passi più grandi senza riscaldamento, il post-LN normalizza l'uscita di ogni blocco. Sul mio banco a due blocchi non c'è differenza misurabile: è il limite della mia misura, non l'assenza del fenomeno.

I sei compromessi hanno due direzioni sole: quelli sul costo (cammino contro T², contesto contro memoria, multi-query) li misuri con un cronometro o con una formula, e vanno decisi sul tuo carico; quelli sulla qualità (numero di teste, tipo di posizione, posizione della normalizzazione) dipendono dai dati e dalla profondità, e un banco piccolo come il mio non li decide. Chi sceglie un'architettura dovrebbe sapere in quale dei due gruppi cade ogni dubbio.

Le trappole in cui si cade davvero

  1. Dimenticare la divisione per √d_k. Sintomo: la perdita scende lentissima e i pesi sono quasi tutti 0 o 1. Causa: il softmax satura e perde la derivata rispetto ai punteggi (sensibilità da 0,96 a 0,10 a d_k = 512). Rimedio: dividere, e stampare il peso massimo medio come diagnostica.
  2. Applicare la maschera dopo il softmax. Sintomo: nessuno. Causa: azzerare i pesi non rinormalizza e le righe sommano a meno di uno — qui una riga somma a 0,031. Rimedio: −∞ prima del softmax, più un assert sulle somme di riga.
  3. Sbagliare l'ordine degli assi nel ricucire le teste. Sintomo: nessuna eccezione, perché le forme tornano. Causa: reshape prima di transpose mescola componenti di teste diverse. Rimedio: cammino inverso nello stesso ordine, e con h = 1 l'identità.
  4. Includere il riempimento nel softmax. Sintomo: risultati che cambiano se riordini il lotto. Causa: manca la maschera di padding e i vettori che rappresentano "niente" prendono peso. Rimedio: maschera sulle chiavi, in AND con quella causale.
  5. Confondere maschera causale e di padding. Sintomo: funziona con sequenze della stessa lunghezza e si rompe in produzione. Causa: la prima è (T, T) e dipende dalle posizioni, la seconda è (B, T) e dipende dal contenuto. Rimedio: due nomi diversi, combinati in un solo punto.
  6. Normalizzare sull'asse delle query. Sintomo: l'addestramento procede ma non converge mai bene. Causa: axis=-2 invece di axis=-1, cioè somma a uno per colonne. Rimedio: il controllo delle somme di riga.
  7. Dimenticare la posizione. Sintomo: accuratezza per sequenza molto più bassa, 0,204 contro 1,000, con un andamento da sacchetto di parole. Causa: l'attenzione è equivariante per permutazione. Rimedio: embedding posizionale, sinusoide o RoPE; come prova, l'uscita non deve più permutarsi.
  8. Usare un ε troppo grande nel controllo del gradiente. Sintomo: il controllo peggiora quando aumenti ε. Causa: la differenza finita attraversa uno spigolo della ReLU. Rimedio: ε = 10⁻⁶ e solo coordinate con gradiente misurabile.
  9. Presentare i pesi di attenzione come la spiegazione di una decisione. Sintomo: una mappa di calore in una presentazione, e nessuno che la contesta. Causa: plausibilità confusa con fedeltà (arXiv:1902.10186, arXiv:1908.04626). Rimedio: intervieni sull'ingresso e misura l'uscita.

Le nove trappole hanno una proprietà comune: nessuna solleva un'eccezione. Le forme tornano, la perdita scende, il modello produce un testo; il difetto si vede soltanto in una misura che qualcuno ha deciso di fare. Per questo, accanto a ciascuna, c'è un controllo che si può scrivere in una riga: le somme di riga, l'identità con h = 1, il test di permutazione, il gradiente per differenze finite.

Quattro esercizi da completare (varianza del prodotto scalare e saturazione del softmax, maschera causale prima e dopo la normalizzazione, parametri di un blocco in forma chiusa, memoria della cache chiave-valore) stanno in esercizi.py, Python 3.12+ e NumPy. Le soluzioni commentate sono in fondo al file. Non l'ho eseguito per questa revisione: i listati dell'articolo sì, il file scaricabile no.

Cosa portarti via

L'attenzione è una domanda confrontata con delle intestazioni per pesare dei contenuti, e tutta la formula softmax(Q Kᵀ / √d_k) V è quella frase in algebra. I pezzi che la rendono utilizzabile non sono accessori: la divisione per √d_k tiene in vita la derivata del softmax, la maschera causale va prima della normalizzazione o le righe non sono più medie, la posizione va iniettata da fuori, e h teste costano quanto una. Il blocco è 12d² + 13d parametri; il costo è T², pendenza misurata fra 1,97 e 2,06; la cache vale un fattore fra 40 e 110 in generazione sul mio giocattolo.

Due cose da tenere oltre la formula. La prima è che un compito artificiale insegna molto proprio perché è artificiale: un decoder da 27 253 parametri copia sequenze in modo perfetto e le ordina per intero nel 95-98% dei casi, e da quelle due cifre si capisce la differenza fra un recupero posizionale e un confronto fra coppie. La seconda è che due delle mie tre ablazioni non hanno mostrato nulla, e la risposta giusta non era insistere fino a trovare un numero favorevole. Un meccanismo lo capisci quando sai anche in quali condizioni il suo effetto non si vede. La prossima lezione usa questi vettori per cercare per significato: la lezione su embedding e ricerca semantica.

Nella stessa serie trovi anche Far girare un modello in locale: quantizzazione, memoria e latenza misurate.

Fonti

Le sintesi degli articoli sono quelle dell'originale e le ho controllate sugli abstract, non sui testi integrali. Tutti i numeri degli esperimenti sono miei: NumPy 2.3.5 e Python 3.14.4 per i listati, PyTorch 2.12.0 solo per il confronto numerico; il modello, il compito e i tempi sono stati ripetuti due volte, con tempi in intervallo. Non ho confrontato il mio decoder con una implementazione di riferimento a livello di addestramento, solo di attenzione.

Commenti

Caricamento commenti...

Accedi per commentare

Discuti questo articolo nella community

Unisciti al gruppo Telegram per discutere con altri sviluppatori, fare domande e condividere le tue esperienze.

Ti è piaciuto questo articolo?

Esplora altri contenuti sul blog o scopri i miei progetti