Attenzione lineare = RNN: una verifica numerica

Nel capitolo abbiamo affermato che l”attenzione lineare è, letteralmente, una rete ricorrente con uno stato-matrice di dimensione fissa. Qui lo verifichiamo con poche righe di NumPy: calcoliamo lo stesso strato in due modi diversi, la forma parallela (come un’attenzione, che costruisce una matrice \(n\times n\)) e la forma ricorrente (come una RNN, che tiene uno stato di dimensione fissa), e controlliamo che diano lo stesso risultato.

Se non programmi, questa pagina si legge lo stesso. Quello che succede qui è che facciamo lo stesso conto in due modi (tutto insieme e una parola alla volta) e andiamo a vedere se i due risultati coincidono. Ti bastano due righe: il numero stampato in fondo, che è la differenza fra i due risultati e deve essere piccolissimo, e la parola True accanto a «le due forme coincidono». Se le trovi entrambe, la promessa del capitolo è mantenuta.

Ingredienti

Generiamo query, chiavi e valori casuali e definiamo la feature map \(\phi(x)=\mathrm{elu}(x)+1\), che rende le affinità sempre positive (così la normalizzazione ha senso), come in Katharopoulos et al. (2020).

import numpy as np

rng = np.random.default_rng(0)
n, d_k, d_v = 6, 4, 5          # 6 token; chiavi in R^4, valori in R^5
Q = rng.standard_normal((n, d_k))
K = rng.standard_normal((n, d_k))
V = rng.standard_normal((n, d_v))

def phi(x):                    # feature map elu(x)+1: sempre positiva
    return np.where(x > 0, x + 1.0, np.exp(x))

PQ, PK = phi(Q), phi(K)
print("Q, K:", Q.shape, " V:", V.shape)

Q, K: (6, 4)  V: (6, 5)

Forma parallela (tipo attenzione)

Costruiamo la matrice di affinità \(\mathbf{A}=\phi(\mathbf{Q})\,\phi(\mathbf{K})^\top\), applichiamo la maschera causale e normalizziamo: l’uscita è la media pesata dei valori. Nota la dimensione di \(\mathbf{A}\): cresce come \(n^2\).

# Forma PARALLELA (tipo attenzione): costruisce la matrice di affinita n x n
A = PQ @ PK.T                              # (n, n)
A *= np.tril(np.ones((n, n)))              # maschera causale: ogni token vede solo il passato
O_par = (A @ V) / A.sum(1, keepdims=True)  # media pesata dei valori
print("matrice di attenzione:", A.shape, "-> cresce come n^2")
print("output O_par:", O_par.shape)

matrice di attenzione: (6, 6) -> cresce come n^2
output O_par: (6, 5)

Forma ricorrente (tipo RNN)

Ora lo stesso calcolo senza mai costruire la matrice \(n\times n\): scorriamo i token uno alla volta e accumuliamo nello stato \(\mathbf{S}_t = \mathbf{S}_{t-1} + \mathbf{v}_t\,\phi(\mathbf{k}_t)^\top\), leggendo \(\mathbf{o}_t = \mathbf{S}_t\,\phi(\mathbf{q}_t) \,/\, \big(\mathbf{z}_t^\top \phi(\mathbf{q}_t)\big)\), con la stessa convenzione del capitolo. Lo stato \(\mathbf{S}\) ha dimensione fissa \(d_v\times d_k\): non dipende dalla lunghezza della sequenza.

# Forma RICORRENTE (tipo RNN): uno stato-matrice di dimensione FISSA, aggiornato token per token
S = np.zeros((d_v, d_k))       # stato: NON dipende da n
z = np.zeros(d_k)              # normalizzatore
O_rec = np.zeros((n, d_v))
for t in range(n):
    S += np.outer(V[t], PK[t]) # accumula v_t phi(k_t)^T (aggiornamento di rango 1)
    z += PK[t]
    O_rec[t] = (S @ PQ[t]) / (z @ PQ[t])
print("stato S:", S.shape, "-> fisso, indipendente da n")

stato S: (5, 4) -> fisso, indipendente da n

Verifica

Le due forme devono coincidere a meno dell’errore di arrotondamento.

# Le due forme calcolano la STESSA funzione
print("max |differenza| =", np.abs(O_par - O_rec).max())
print("le due forme coincidono:", np.allclose(O_par, O_rec))

max |differenza| = 2.220446049250313e-16
le due forme coincidono: True

La differenza è dell’ordine dell’epsilon macchina (\(\approx 10^{-16}\)): lo stesso strato, calcolato in due modi. Durante l’addestramento conviene la forma parallela (sfrutta il calcolo su tutta la sequenza in una volta); in inferenza autoregressiva conviene la forma ricorrente, che aggiorna uno stato di dimensione costante: nessuna KV cache che cresce parola dopo parola. È la proprietà che accomuna tutta la famiglia di modelli di questo capitolo.