Tre Matrici Per Testa, O Una Grande Matrice

Visione Per-Testa

Ogni testa ha bisogno della propria proiezione query, proiezione key e proiezione value. Per la testa h:

Q_h = X · W_Q^h    dove W_Q^h ha forma [d_model, head_dim]
K_h = X · W_K^h    dove W_K^h ha forma [d_model, head_dim]
V_h = X · W_V^h    dove W_V^h ha forma [d_model, head_dim]

X porta la forma di input [batch, seq_len, d_model]. Dopo la proiezione, Q_h, K_h, V_h portano ciascuna la forma [batch, seq_len, head_dim].

Vista Fusa

Le matrici per testa sono disposte una accanto all'altra in memoria. Una singola matrice fusa W_Q di forma [d_model, d_model] produce tutte le teste contemporaneamente:

Q_fused = X · W_Q              # [batch, seq_len, d_model]
Q_per_head = reshape(Q_fused)  # [batch, n_head, seq_len, head_dim]

Il matmul fuso esegue una sola chiamata BLAS invece di 12. I tensor core di CUDA raggiungono il picco di throughput su matmul di questa dimensione; i matmul per testa sottoutilizzerebbero l'hardware.

Conteggio Parametri

Tre matrici fuse W_Q, W_K, W_V, ognuna d_model × d_model. Più la proiezione di output W_O, anch'essa d_model × d_model. Per ANDREA-120M:

parametri per attention di un layer = 4 × 768² = 2,359,296 ≈ 2.36M
parametri su 12 layer             = 12 × 2.36M ≈ 28.3M

Circa un quarto dei parametri totali di ANDREA-120M risiede nelle proiezioni di attenzione. I restanti tre quarti risiedono nel sottolayer MLP ed embeddings.

Nominare le Proiezioni

Dodici Vettori Diventano Uno

Dopo che Ogni Head Ha Calcolato

Ogni head produce un tensore di output di forma [batch, seq_len, head_dim]. Dodici head producono dodici tensori di questo tipo. La concatenazione lungo la dimensione delle feature li rimette insieme:

concat_output = concat(head_1, head_2, ..., head_12)
shape         = [batch, seq_len, n_head × head_dim]
= [batch, seq_len, 768]    # for ANDREA-120M

Concat inverte lo split. La dimensione totale delle feature ritorna a d_model. Nessuna perdita di informazione nelle dimensioni; la differenza risiede in ciò che contiene ogni chunk: il chunk della testa 1 riflette il pattern di attenzione appreso dalla testa 1.

La Proiezione di Output W_O

La semplice concatenazione lascia le teste isolate: l'output della testa 4 si trova accanto all'output della testa 7, senza che nessuna sia consapevole dell'altra. La proiezione di output W_O di forma [d_model, d_model] le mescola:

attention_output = concat_output · W_O
shape            = [batch, seq_len, d_model]

Dopo W_O, ogni dimensione di output trasporta una combinazione lineare appresa di tutte e dodici le teste. Le informazioni fluiscono liberamente tra le teste attraverso questa singola moltiplicazione di matrici.

Perché le Teste si Specializzano

Niente nell'architettura forza la testa 4 a imparare il tempo verbale o la testa 9 a imparare la punteggiatura accoppiata. La specializzazione emerge dalla pressione del gradiente: durante l'addestramento, le teste che contribuiscono in modo ridondante ricevono gradienti più piccoli rispetto a quelle che contribuiscono in modo unico. In migliaia di passi, ogni testa si stabilizza in una nicchia che riduce la perdita totale nel modo più efficace.

Empiricamente, i transformer addestrati mostrano teste che gestiscono: pattern posizionali (la testa guarda il token precedente), pattern sintattici (la testa guarda la parentesi chiusa corrispondente), pattern semantici (la testa guarda l'entità nominata più recente). Nessuna etichetta addestra questa specializzazione. Il solo segnale di addestramento, propagato attraverso W_O, ordina le teste.

Perché Dodici Teste, Non Una Testa Più Larga

Come CUDA Memorizza le Teste

Un Singolo Tensore, Rimodellato

Il motore di addestramento di ANDREA microgpt_cuda.cu non alloca dodici buffer separati per dodici teste. Alloca un tensore fuso unico & tratta la dimensione della testa come un pattern di stride:

// dopo Q = X · W_Q (un matmul, fuso tra le teste)
// Q ha forma [batch, seq_len, d_model]
// ridimensiona a    [batch, seq_len, n_head, head_dim]
// trasponi a  [batch, n_head, seq_len, head_dim]
// ogni testa ora è contigua nelle due dimensioni interne

Il trasponi sposta n_head prima di seq_len. Perché? Perché l'operazione successiva (Q_h · K_h^T) richiede che la fetta seq_len × head_dim di ogni testa sia contigua in memoria. I matmul CUDA sono più veloci sui tensori contigui.

Un Kernel, Molte Teste

Un singolo kernel CUDA di attention viene eseguito su tutte le teste in parallelo. Ogni thread block gestisce una coppia (batch, head); i thread all'interno del block collaborano sul tile seq_len × head_dim. Il kernel non sa mai di processare multiple teste; la griglia di lancio gestisce il parallelismo.

La Configurazione Riflette l'Hardware

La scelta di ANDREA-120M di n_head=12, head_dim=64 si allinea con i tensor core della RTX 4090, che preferiscono tile matmul in multipli di 16. head_dim=64 = 4 × 16 corrisponde esattamente alla forma del tile. head_dim=32 (ANDREA-12M) corrisponde anch'esso ma sottoutilizza il tile. head_dim=72 non corrisponderebbe e forzerebbe kernel di fallback.

Immagine Finale

Cinque passaggi. Tre matmul toccano l'input (proiezioni Q, K, V). Una matmul tocca le teste concatenate (W_O). Un kernel di attention gestisce tutte le teste in parallelo. ANDREA-120M esegue tutti e cinque i passaggi una volta per layer, dodici layer in profondità, ad ogni forward pass.

Lettura del Pipeline