01 - Architettura Transformer: attention, encoder, decoder

Leggi la frase «la borsa non entrava nella valigia perché era troppo piccola». Per sapere che «era» si riferisce alla valigia e non alla borsa devi collegare due parole lontane tra loro, e lo fai in un colpo solo. Una rete ricorrente legge invece le parole una dopo l'altra e porta con sé un solo vettore di memoria: l'informazione sulla prima parola deve attraversare tutte le altre per arrivare all'ultima. Il Transformer, proposto da Vaswani e colleghi nel 2017, toglie quel passaggio obbligato: ogni posizione della sequenza guarda tutte le altre direttamente, e decide quanto peso dare a ciascuna con un calcolo che sta in due righe di algebra lineare. Questo articolo costruisce quel calcolo, lo confronta con PyTorch, e poi verifica una per una le quattro idee che lo rendono utilizzabile: il fattore di scala, le teste multiple, la codifica delle posizioni e la maschera del decoder.

Ogni listato è stato eseguito due volte con uscita identica, con numpy 2.4.6 e torch 2.12.0 su CPU, su numeri casuali a seme fisso. Non c'è nessun modello addestrato e nessun testo: l'attenzione si può studiare sui numeri, ed è il modo più rapido per capirla.

Il problema: portare l'informazione da lontano

La tabella 1 del lavoro originale (Vaswani et al., Attention Is All You Need) confronta i tipi di strato con tre misure. Con n lunghezza della sequenza e d dimensione della rappresentazione, lo strato ricorrente costa O(n·d²) operazioni, richiede O(n) passi sequenziali e ha un cammino massimo fra due posizioni di O(n). L'auto-attenzione costa O(n²·d), ma richiede O(1) passi sequenziali e il cammino massimo è O(1): due posizioni qualunque sono collegate in un solo strato. Il prezzo è il termine n², che si misurerà più avanti; il guadagno è che tutte le posizioni si calcolano insieme, quindi un acceleratore può usare tutto il suo parallelismo, mentre la ricorrenza è sequenziale per costruzione.

L'attenzione come ricerca morbida in un dizionario

Un dizionario Python restituisce il valore la cui chiave coincide con quella cercata. L'attenzione fa lo stesso ma con una corrispondenza graduale. Ogni posizione produce tre vettori: una query (che cosa sto cercando), una chiave (di che cosa parlo) e un valore (che cosa offro se mi scegli). Il prodotto scalare fra la query di una posizione e le chiavi di tutte dà un punteggio per ciascuna; il softmax trasforma i punteggi in pesi positivi che sommano a uno; l'uscita è la media dei valori pesata con quei pesi. Nella formula del paragrafo 3.2.1 del lavoro, Attention(Q, K, V) = softmax(QKᵀ/√d_k)·V, la divisione per √d_k è l'unica cosa che un dizionario non ha, e la prossima sezione spiega a che cosa serve.

Il listato calcola l'attenzione per cinque posizioni con d_k = 8 e d_v = 4. La funzione softmax sottrae il massimo di ogni riga prima di esponenziare: non cambia il risultato (il softmax è invariante a una traslazione comune), ma evita che exp vada in overflow con punteggi grandi. L'argomento mask serve più avanti. Il secondo calcolo indipendente è scaled_dot_product_attention di PyTorch, la funzione che l'API della libreria dedica a questa formula: se le due implementazioni coincidono, la formula è scritta giusta.

import numpy as np
import torch
import torch.nn.functional as F

torch.set_num_threads(2)
rng = np.random.default_rng(0)


def softmax(z):
    z = z - z.max(axis=-1, keepdims=True)      # stabilità numerica: non cambia il risultato
    e = np.exp(z)
    return e / e.sum(axis=-1, keepdims=True)


def attention(Q, K, V, mask=None):
    d_k = Q.shape[-1]
    scores = Q @ K.T / np.sqrt(d_k)            # (n_query, n_key)
    if mask is not None:
        scores = np.where(mask, scores, -np.inf)
    w = softmax(scores)
    return w @ V, w


n, d_k, d_v = 5, 8, 4
Q = rng.standard_normal((n, d_k))
K = rng.standard_normal((n, d_k))
V = rng.standard_normal((n, d_v))

out, w = attention(Q, K, V)
ref = F.scaled_dot_product_attention(
    torch.from_numpy(Q)[None], torch.from_numpy(K)[None], torch.from_numpy(V)[None]
)[0].numpy()

print("forma dei pesi:", w.shape, "- somma di ogni riga:", np.round(w.sum(axis=1), 12))
print("differenza massima con torch:", f"{np.abs(out - ref).max():.2e}")
print("pesi della prima riga:", np.round(w[0], 3))

# Output:
# forma dei pesi: (5, 5) - somma di ogni riga: [1. 1. 1. 1. 1.]
# differenza massima con torch: 3.33e-16
# pesi della prima riga: [0.367 0.063 0.113 0.115 0.342]

La differenza è dell'ordine di 10^-16, cioè l'arrotondamento della doppia precisione, e ogni riga dei pesi somma a uno. Il peso è una distribuzione: la prima posizione dà 0,367 a sé stessa, 0,342 all'ultima e poco alle altre tre. L'uscita di ogni posizione è quindi una miscela di valori, e quale miscela dipende dal contenuto, non dalla distanza: per questo la prima e l'ultima posizione possono parlarsi senza intermediari.

Perché si divide per √d_k

Se le componenti di q e k sono variabili indipendenti con media 0 e varianza 1, il loro prodotto scalare ha media 0 e varianza d_k: lo dice il paragrafo 3.2.2 del lavoro, in nota. La deviazione standard dei punteggi cresce quindi come √d_k. Il softmax di numeri molto dispersi è quasi un «vincitore prende tutto»: un solo peso vicino a uno e gli altri vicino a zero, e in quella regione le derivate sono minuscole. Gli autori scrivono di sospettarlo come causa del peggior risultato dell'attenzione senza scala per d_k grande; qui lo si misura direttamente.

Il listato estrae 4000 coppie query-chiavi per ciascun d_k, con 16 chiavi per query, e misura tre cose. La deviazione standard dei punteggi grezzi si confronta con √d_k (secondo calcolo, in forma chiusa). Il peso massimo medio dice quanto il softmax è «acceso» su una sola chiave. La norma di Frobenius della matrice jacobiana del softmax, diag(p) − ppᵀ, dice quanto le uscite reagiscono a una variazione dei punteggi: se è piccola, il gradiente che attraversa lo strato è piccolo e l'apprendimento rallenta. Le ultime due colonne ripetono la misura dopo la divisione.

import numpy as np

rng = np.random.default_rng(1)
n_keys, trials = 16, 4000


def softmax(z):
    e = np.exp(z - z.max(axis=-1, keepdims=True))
    return e / e.sum(axis=-1, keepdims=True)


def jac_norm(p):
    # norma di Frobenius della matrice jacobiana del softmax: diag(p) - p p^T
    J = np.einsum("ti,ij->tij", p, np.eye(p.shape[1])) - np.einsum("ti,tj->tij", p, p)
    return np.sqrt((J ** 2).sum(axis=(1, 2)))


print(f"{'d_k':>5} {'std(q.k)':>9} {'sqrt(d_k)':>9} | {'peso max':>9} {'|J| senza':>9} | {'peso max':>9} {'|J| con':>9}")
for d_k in (4, 16, 64, 256, 1024):
    q = rng.standard_normal((trials, 1, d_k))
    k = rng.standard_normal((trials, n_keys, d_k))
    s = (q * k).sum(axis=-1)                       # prodotti scalari grezzi, forma (trials, n_keys)
    p_raw = softmax(s)
    p_scl = softmax(s / np.sqrt(d_k))
    print(f"{d_k:>5} {s.std():>9.3f} {np.sqrt(d_k):>9.3f} | "
          f"{p_raw.max(axis=1).mean():>9.3f} {jac_norm(p_raw).mean():>9.4f} | "
          f"{p_scl.max(axis=1).mean():>9.3f} {jac_norm(p_scl).mean():>9.4f}")

# Output:
#   d_k  std(q.k) sqrt(d_k) |  peso max |J| senza |  peso max   |J| con
#     4     1.986     2.000 |     0.426    0.3221 |     0.234    0.2967
#    16     3.991     4.000 |     0.681    0.2805 |     0.241    0.3024
#    64     7.973     8.000 |     0.846    0.1732 |     0.245    0.3042
#   256    15.978    16.000 |     0.923    0.0987 |     0.246    0.3049
#  1024    31.927    32.000 |     0.963    0.0504 |     0.248    0.3053

La deviazione standard dei punteggi segue √d_k a meno di una cifra decimale (31,9 contro 32 per d_k = 1024). Senza scala, passando da d_k = 4 a d_k = 1024 il peso massimo medio sale da 0,426 a 0,963 e la norma della jacobiana scende da 0,32 a 0,050: con 1024 dimensioni il softmax è quasi un massimo rigido e lascia passare circa un sesto del gradiente che passerebbe con la scala (0,050 contro 0,305). Con la scala, il peso massimo resta intorno a 0,24 e la jacobiana intorno a 0,30 qualunque sia d_k, perché la varianza dei punteggi vale uno per costruzione. La correzione quindi non rende l'attenzione «migliore»: la mantiene nello stesso regime al variare della dimensione. Resta un'avvertenza: la misura vale per vettori con componenti indipendenti e varianza uno, cioè all'inizializzazione; nei modelli addestrati le componenti non lo sono più, e questo esperimento non dice come si comportano.

Teste multiple: più attenzioni piccole invece di una grande

Una sola attenzione produce una sola miscela di valori per posizione: nella frase di apertura «era» deve guardare la valigia, ma lo stesso strato dovrebbe anche guardare il verbo. Le teste multiple (paragrafo 3.2.2) fanno girare h attenzioni in parallelo, ciascuna su una proiezione lineare diversa di dimensione d_k = d_modello/h, poi concatenano le uscite e le proiettano di nuovo con una matrice W^O. Nel modello base degli autori d_modello = 512 e h = 8, quindi d_k = 64.

Il listato controlla due cose contro torch.nn.MultiheadAttention, con d = 16 e h = 4. La prima: la mia implementazione a mano, che legge i pesi del modulo PyTorch (in_proj_weight impila le tre proiezioni Q, K, V) e assegna a ogni testa una fetta di dh colonne, riproduce uscita e pesi di attenzione di ogni testa. La seconda: il numero di parametri non dipende da h. Le teste dividono la dimensione, non la moltiplicano, quindi con 1, 4 o 16 teste i parametri sono sempre 4d² + 4d (tre proiezioni, quella di uscita, i quattro bias).

import numpy as np
import torch

torch.set_num_threads(2)
torch.manual_seed(0)
d, h, n = 16, 4, 6
dh = d // h

mha = torch.nn.MultiheadAttention(d, h, batch_first=True, dtype=torch.float64)
x = torch.randn(1, n, d, dtype=torch.float64)
with torch.no_grad():
    ref, ref_w = mha(x, x, x, need_weights=True, average_attn_weights=False)

Wi = mha.in_proj_weight.detach().numpy()       # (3d, d): le tre proiezioni impilate Q, K, V
bi = mha.in_proj_bias.detach().numpy()
Wo = mha.out_proj.weight.detach().numpy()
bo = mha.out_proj.bias.detach().numpy()
X = x[0].numpy()


def softmax(z):
    e = np.exp(z - z.max(axis=-1, keepdims=True))
    return e / e.sum(axis=-1, keepdims=True)


Q = X @ Wi[:d].T + bi[:d]
K = X @ Wi[d:2 * d].T + bi[d:2 * d]
V = X @ Wi[2 * d:].T + bi[2 * d:]
heads, weights = [], []
for i in range(h):                              # ogni testa legge la sua fetta di dh colonne
    sl = slice(i * dh, (i + 1) * dh)
    w = softmax(Q[:, sl] @ K[:, sl].T / np.sqrt(dh))
    weights.append(w)
    heads.append(w @ V[:, sl])
out = np.concatenate(heads, axis=1) @ Wo.T + bo  # concatena e proietta con W^O

print("differenza massima sull'uscita:", f"{np.abs(out - ref[0].detach().numpy()).max():.2e}")
print("differenza massima sui pesi:   ", f"{np.abs(np.stack(weights) - ref_w[0].detach().numpy()).max():.2e}")

for heads_n in (1, 4, 16):
    m = torch.nn.MultiheadAttention(d, heads_n)
    print(f"{heads_n:>2} teste -> parametri: {sum(p.numel() for p in m.parameters())}")
print("formula 4*d*d + 4*d =", 4 * d * d + 4 * d)

# Output:
# differenza massima sull'uscita: 1.11e-16
# differenza massima sui pesi:    1.11e-16
#  1 teste -> parametri: 1088
#  4 teste -> parametri: 1088
# 16 teste -> parametri: 1088
# formula 4*d*d + 4*d = 1088

L'uscita e i pesi coincidono a 10^-16 per tutte e quattro le teste. I parametri sono 1088 per qualunque numero di teste, uguali alla formula. Quindi la scelta di h non è un costo di memoria dei pesi: è una scelta su come ripartire la dimensione. La tabella 3 del lavoro originale (righe A) mostra che 8 e 16 teste danno la stessa qualità sul loro banco di prova (BLEU 25,8 su traduzione inglese-tedesco), mentre una sola testa a dimensione piena perde 0,9 punti e 32 teste scendono a 25,4: troppe teste significano chiavi troppo piccole. Sono misure di quel compito e di quel modello, non una regola per il tuo.

La posizione: l'attenzione non sa l'ordine

Una conseguenza dell'attenzione è controintuitiva e si dimostra con tre righe: se mescoli l'ordine dei token in ingresso, l'uscita viene mescolata nello stesso modo, e basta. Lo strato non vede l'ordine; vede un insieme. Per un linguaggio è un problema, perché «il cane morde l'uomo» e «l'uomo morde il cane» hanno gli stessi token. Gli autori sommano a ogni embedding una codifica posizionale (paragrafo 3.5): per la posizione pos e la coppia di dimensioni 2i, 2i+1, PE = (sin(pos/10000^2i/d), cos(pos/10000^2i/d)). Ogni coppia è una sinusoide con frequenza diversa, dalle lunghezze d'onda più brevi (2π) alle più lunghe (10000·2π). Il motivo dichiarato è che per ogni scarto fisso k, PE_pos+k è una funzione lineare di PE_pos.

Il listato verifica entrambe le cose. Prima, con MultiheadAttention senza posizioni: il confronto y_perm contro y[:, perm] dice se l'uscita del batch permutato è la permutazione dell'uscita. Poi, aggiungendo la codifica ai token e tenendo le posizioni fisse, la stessa uguaglianza deve cadere. Infine, per k = 5 costruisce a mano la matrice di rotazione 2×2 R_k con angolo ω·k e controlla che R_k·PE_pos coincida con PE_pos+k su 30 posizioni e 8 frequenze: è la dimostrazione numerica dell'affermazione degli autori, e mostra che l'operazione è una rotazione in ciascun piano.

import numpy as np
import torch

torch.set_num_threads(2)
torch.manual_seed(0)
d, h, n = 16, 4, 7
mha = torch.nn.MultiheadAttention(d, h, batch_first=True, dtype=torch.float64)


def run(x):
    with torch.no_grad():
        return mha(x, x, x, need_weights=False)[0]


def pe(n, d):
    pos = np.arange(n)[:, None]
    i = np.arange(d // 2)[None, :]
    ang = pos / 10000 ** (2 * i / d)
    out = np.zeros((n, d))
    out[:, 0::2] = np.sin(ang)
    out[:, 1::2] = np.cos(ang)
    return out


x = torch.randn(1, n, d, dtype=torch.float64)
perm = torch.randperm(n, generator=torch.Generator().manual_seed(3))

y = run(x)
y_perm = run(x[:, perm])
print("senza posizione: uscita(permutato) == permutata(uscita)?",
      torch.allclose(y_perm, y[:, perm]), f"(errore {float((y_perm - y[:, perm]).abs().max()):.1e})")

P = torch.from_numpy(pe(n, d))[None]
z = run(x + P)
z_perm = run(x[:, perm] + P)                     # stessi token, ordine diverso, posizioni fisse
print("con codifica posizionale: stessa uguaglianza?",
      torch.allclose(z_perm, z[:, perm]), f"(errore {float((z_perm - z[:, perm]).abs().max()):.1e})")

# PE(pos + k) = R_k PE(pos): ogni coppia (sin, cos) ruota di un angolo che dipende solo da k
k = 5
big = pe(40, d)
worst = 0.0
for pos in range(0, 30):
    for i in range(d // 2):
        w = 1 / 10000 ** (2 * i / d)
        R = np.array([[np.cos(w * k), np.sin(w * k)], [-np.sin(w * k), np.cos(w * k)]])
        pred = R @ big[pos, 2 * i:2 * i + 2]
        worst = max(worst, np.abs(pred - big[pos + k, 2 * i:2 * i + 2]).max())
print(f"PE(pos+{k}) = R_{k} PE(pos): errore massimo su 30 posizioni e {d // 2} frequenze = {worst:.1e}")

# Output:
# senza posizione: uscita(permutato) == permutata(uscita)? True (errore 1.1e-16)
# con codifica posizionale: stessa uguaglianza? False (errore 2.1e-01)
# PE(pos+5) = R_5 PE(pos): errore massimo su 30 posizioni e 8 frequenze = 8.9e-16

Senza posizioni l'uguaglianza regge a 10^-16: lo strato è equivariante alle permutazioni. Con la codifica cade, con una differenza di 0,21, perché lo stesso token in posizioni diverse porta ora un contributo diverso. L'ultima riga conferma il perché della forma sinusoidale: lo spostamento di 5 posizioni è la stessa rotazione per ogni pos, a 10^-15. Il lavoro originale ha provato anche embedding posizionali appresi e ha trovato risultati quasi identici (tabella 3, riga E); ha scelto la versione sinusoidale perché potrebbe estrapolare a sequenze più lunghe di quelle viste, un'ipotesi che qui non si è messa alla prova. Altri schemi di posizione sono stati proposti dopo il 2017, ma questo articolo non li confronta.

Encoder, decoder e la maschera causale

Lo stesso strato di attenzione compare in tre ruoli, descritti nei paragrafi 3.1 e 3.2.3 del lavoro:

• auto-attenzione dell'encoder: query, chiavi e valori vengono dall'uscita dello strato precedente, e ogni posizione guarda tutte le altre;
• auto-attenzione mascherata del decoder: ogni posizione guarda solo sé stessa e le precedenti;
• attenzione encoder-decoder: le query vengono dal decoder, chiavi e valori dall'uscita dell'encoder, e ogni posizione del decoder può guardare tutto l'ingresso.

Cambia solo da dove arrivano Q, K, V e quali connessioni sono permesse; la formula è sempre la stessa. L'encoder del lavoro è una pila di N = 6 strati identici, ciascuno con attenzione a più teste e una rete feed-forward per posizione, con una connessione residua intorno a ciascuna sotto-parte seguita da normalizzazione di strato: l'uscita è LayerNorm(x + Sottostrato(x)). Il decoder ha gli stessi strati più il terzo ruolo della lista, e l'uscita è spostata di una posizione perché la previsione in i deve dipendere solo dalle uscite note prima di i. La residua impone che tutte le sotto-parti producano la stessa dimensione d_modello (512 nel lavoro).

La maschera si implementa nella formula stessa: i punteggi illegali si portano a −∞ prima del softmax, e il peso diventa esattamente zero. Il listato la verifica nel modo più diretto: modifica solo le posizioni 5, 6 e 7 dell'ingresso, e misura quanto cambia l'uscita in ciascuna posizione, con e senza maschera triangolare. Se la maschera funziona, le posizioni 0-4 non devono accorgersi della modifica; se non c'è, la modifica deve propagarsi a tutte.

import torch

torch.set_num_threads(2)
torch.manual_seed(0)
d, h, n = 16, 4, 8
mha = torch.nn.MultiheadAttention(d, h, batch_first=True, dtype=torch.float64)
causal = torch.triu(torch.ones(n, n, dtype=torch.bool), diagonal=1)   # True = vietato guardare


def run(x, mask):
    with torch.no_grad():
        return mha(x, x, x, attn_mask=mask, need_weights=False)[0]


x = torch.randn(1, n, d, dtype=torch.float64)
x2 = x.clone()
x2[:, 5:] = torch.randn(1, n - 5, d, dtype=torch.float64)     # cambio solo le posizioni 5, 6, 7

for nome, mask in (("senza maschera", None), ("con maschera causale", causal)):
    diff = (run(x, mask) - run(x2, mask)).abs().amax(dim=-1)[0]
    print(f"{nome:>22}: variazione dell'uscita per posizione", [f"{v:.1e}" for v in diff.tolist()])

# Output:
#         senza maschera: variazione dell'uscita per posizione ['2.3e-01', '3.4e-01', '3.3e-01', '2.9e-01', '2.5e-01', '3.4e-01', '2.8e-01', '2.8e-01']
#   con maschera causale: variazione dell'uscita per posizione ['0.0e+00', '0.0e+00', '0.0e+00', '0.0e+00', '0.0e+00', '4.1e-01', '2.0e-01', '2.8e-01']

Senza maschera la modifica cambia l'uscita di ogni posizione, con scarti fra 0,23 e 0,34. Con la maschera causale le posizioni 0-4 restano a zero esatto, e cambiano solo le posizioni 5-7, che «vedono» il token modificato. La conseguenza pratica è che in addestramento il decoder può calcolare in un solo passaggio la previsione di tutte le posizioni senza barare, e in generazione deve invece procedere un token alla volta, perché i token successivi non esistono ancora.

Quanto costa il termine n²

Il punteggio di ogni coppia di posizioni occupa una cella: una matrice n × n per testa e per strato. Il listato misura sulla CPU, con d = 64 e float32, la dimensione di quella matrice e il tempo medio di cinque esecuzioni di un singolo calcolo di attenzione, per n da 512 a 4096.

import time
import numpy as np

rng = np.random.default_rng(0)
d = 64
print(f"{'n':>5} {'entrate n*n':>12} {'MiB (float32)':>14} {'tempo medio (s)':>16}")
for n in (512, 1024, 2048, 4096):
    Q = rng.standard_normal((n, d)).astype(np.float32)
    K = rng.standard_normal((n, d)).astype(np.float32)
    V = rng.standard_normal((n, d)).astype(np.float32)
    t = []
    for _ in range(5):
        t0 = time.perf_counter()
        S = Q @ K.T / np.sqrt(np.float32(d))
        S -= S.max(axis=1, keepdims=True)
        np.exp(S, out=S)
        S /= S.sum(axis=1, keepdims=True)
        out = S @ V
        t.append(time.perf_counter() - t0)
    print(f"{n:>5} {n * n:>12} {n * n * 4 / 2**20:>14.1f} {np.mean(t):>16.3f}")

# Output:
#     n  entrate n*n  MiB (float32)  tempo medio (s)
#   512       262144            1.0            0.001
#  1024      1048576            4.0            0.004
#  2048      4194304           16.0            0.017
#  4096     16777216           64.0            0.068

Ogni raddoppio di n quadruplica le celle (1, 4, 16, 64 MiB) e il tempo cresce di circa quattro volte per raddoppio: da 512 a 4096 le celle crescono di 64 volte, e il tempo, in tre esecuzioni, è passato da circa un millisecondo a 0,07-0,09 s. I tempi cambiano da una macchina e da un'esecuzione all'altra (e con la libreria di algebra lineare usata), quindi vale l'ordine di grandezza e non il valore. Con una sola testa e uno strato, 4096 posizioni costano 64 MiB di punteggi; con 8 teste e 6 strati, se si conservassero tutti per l'addestramento, sarebbero 3 GiB per ogni esempio del batch. Per questo gli strati ottimizzati calcolano il softmax a blocchi senza materializzare la matrice intera, ma il conto quadratico delle operazioni resta.

Che cosa sai fare adesso

Sai scrivere l'attenzione a scala di prodotto scalare in poche righe di NumPy e verificarla contro PyTorch, sai misurare perché serve il fattore 1/√d_k, sai ricostruire a mano un'attenzione a più teste da un modulo di libreria, sai dimostrare che le posizioni vanno aggiunte e che la codifica sinusoidale è una rotazione, e sai controllare che una maschera causale impedisca davvero di guardare avanti. Quanto di tutto questo regga a scala di miliardi di parametri non è argomento di questo articolo: il successivo, 02 - LLM scaling laws e capacità emergenti, studia come l'errore cambia con la dimensione del modello e dei dati. Per chi viene dal machine learning classico, l'articolo sulla scomposizione bias e varianza è un buon prerequisito, e l'indice completo è nella pagina della serie.

Più avanti nella serie trovi anche 08 - Inferenza on-prem con Ollama, vLLM, llama.cpp, che riprende il costo dell'attenzione dal lato dell'inferenza, con la memoria della cache chiave-valore.