Tre Matriser Per Huvud, Eller En Stor Matris

Vy Per Huvud

Varje huvud behöver sin egen query-projektion, key-projektion och value-projektion. För huvud h:

Q_h = X · W_Q^h    där W_Q^h har form [d_model, head_dim]
K_h = X · W_K^h    där W_K^h har formen [d_model, head_dim]
V_h = X · W_V^h    där W_V^h har formen [d_model, head_dim]

X bär på indatans form [batch, seq_len, d_model]. Efter projektion bär Q_h, K_h, V_h vardera formen [batch, seq_len, head_dim].

Sammanslagen vy

Per-head-matriserna ligger bredvid varandra i minnet. En enda sammanslagen matris W_Q med formen [d_model, d_model] producerar alla huvuden på en gång:

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

Den fusade matmul skickar ett BLAS-anrop istället för 12. CUDA tensor cores når toppgenomströmning på matmuls av denna storlek; per-head matmuls skulle underutnyttja hårdvaran.

Antal parametrar

Tre sammansmälta matriser W_Q, W_K, W_V, varje d_model × d_model. Plus utdataprojektionen W_O, också d_model × d_model. För ANDREA-120M:

parametrar per lager's attention = 4 × 768² = 2,359,296 ≈ 2.36M
parametrar över 12 lager      = 12 × 2.36M ≈ 28.3M

Ungefär en fjärdedel av ANDREA-120M:s totala parametrar finns i attention-projektionerna. De återstående tre fjärdedelarna finns i MLP-underskiktet & inbäddningarna.

Benämning av projektionerna

Tolv vektorer blir en

Efter att varje huvud har beräknat

Varje huvud producerar ett utdatatensor av formen [batch, seq_len, head_dim]. Tolv huvuden producerar tolv sådana tensorer. Konkatenering längs funktionsdimensionen staplar dem tillbaka tillsammans:

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

Concat vänder upp splitningen. Den totala funktionsdimensionen återgår till d_model. Ingen informationsförlust i dimensionerna; skillnaden ligger i vad varje chunk innehåller: chunk för head 1 reflekterar head 1:s inlärda uppmärksamhetsmönster.

Utdata-projektionen W_O

Endast konkatenering lämnar huvudena isolerade: head 4:s utdata sitter bredvid head 7:s utdata, utan vetskap om varandra. Utdata-projektionen W_O med formen [d_model, d_model] blandar dem:

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

Efter W_O bär varje utmatningsdimension en inlärd linjär kombination av alla tolv huvuden. Information flödar fritt mellan huvudena genom denna enda matris-multiplikation.

Varför huvuden specialiserar sig

Ingen del av arkitekturen tvingar huvud 4 att lära sig verbtidsformer eller huvud 9 att lära sig matchad parentes. Specialisering uppstår från gradienttryck: under träning får huvuden som bidrar redundantly mindre gradienter än huvuden som bidrar unikt. Över tusentals steg slår sig varje huvud ner i en nisch som minskar totalförlusten mest effektivt.

Empiriskt visar tränade transformatorer huvuden som hanterar: positionsmönster (huvud tittar på föregående token), syntaktiska mönster (huvud tittar på matchande slutparentes), semantiska mönster (huvud tittar på den senaste namngivna enheten). Inga etiketter tränar denna specialisering. Träningsignalen ensam, propagerad genom W_O, sorterar huvudena.

Varför Tolv Huvuden, Inte Ett Bredare Huvud

Hur CUDA lagrar huvudena

En enda tensor, omformad

ANDREA:s träningsmotor microgpt_cuda.cu allokerar inte tolv separata buffertar för tolv huvuden. Den allokerar en sammanslagen tensor & behandlar huvudimensionen som ett stegmönster:

// efter Q = X · W_Q (en matmul, sammanslagen över huvuden)
// Q har form [batch, seq_len, d_model]
// omforma till    [batch, seq_len, n_head, head_dim]
// transposera till  [batch, n_head, seq_len, head_dim]
// varje huvud är nu kontinuerligt i de inre två dimensionerna

Transponeringen flyttar n_head framför seq_len. Varför? Eftersom nästa operation (Q_h · K_h^T) behöver varje heads seq_len × head_dim-skiva kontinuerlig i minnet. CUDA-matmul-operationer körs snabbare på kontinuerliga tensorer.

En Kernel, Många Huvuden

En enda attention CUDA-kärna körs över alla huvuden parallellt. Varje trådblock hanterar ett (batch, head)-par; trådar inuti blocket samarbetar på seq_len × head_dim-plattan. Kärnan vet aldrig att den processar flera huvuden; lanseringsgaller hanterar parallellismen.

Konfigurationen speglar hårdvaran

ANDREA-120M:s val av n_head=12, head_dim=64 stämmer överens med RTX 4090 tensor cores, som föredrar matmul-plattor i multipler av 16. head_dim=64 = 4 × 16 matchar plattformen exakt. head_dim=32 (ANDREA-12M) matchar också men underutnyttjar plattformen. head_dim=72 skulle inte matcha & skulle tvinga fallback-kärnor.

Slutliga bilden

Fem steg. Tre matmuls berör input (Q, K, V-projektioner). En matmul berör de sammanslagna headsen (W_O). En attention-kärna hanterar alla heads parallellt. ANDREA-120M kör alla fem steg en gång per lager, tolv lager djupt, varje forward-pass.

Läsa pipelinen