Drie Matrices Per Head, Of Eén Grote Matrix

Per-Head View

Elke head heeft zijn eigen query-projectie, key-projectie en value-projectie nodig. Voor head h:

Q_h = X · W_Q^h    waarbij W_Q^h de vorm [d_model, head_dim] heeft
K_h = X · W_K^h    waarbij W_K^h de vorm [d_model, head_dim] heeft
V_h = X · W_V^h    waarbij W_V^h de vorm [d_model, head_dim] heeft

X draagt de invoervorm [batch, seq_len, d_model]. Na projectie dragen Q_h, K_h, V_h elk de vorm [batch, seq_len, head_dim].

Gefuseerde Weergave

Per-head matrices liggen naast elkaar in het geheugen. Een enkele gefuseerde matrix W_Q van vorm [d_model, d_model] produceert alle heads tegelijk:

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

De gefuseerde matmul gebruikt één BLAS-oproep in plaats van 12. CUDA tensor cores bereiken piekdoorvoer bij matmuls van deze grootte; per-head matmuls zouden de hardware onderbenutten.

Aantal parameters

Drie gefuseerde matrices W_Q, W_K, W_V, elk d_model × d_model. Plus de outputprojectie W_O, ook d_model × d_model. Voor ANDREA-120M:

params per laag's attention = 4 × 768² = 2.359.296 ≈ 2,36M
params over 12 lagen      = 12 × 2,36M ≈ 28,3M

Ongeveer een kwart van de totale parameters van ANDREA-120M bevindt zich in de attention-projecties. De resterende driekwart bevindt zich in de MLP-sublayer & embeddings.

De Projecties Benameren

Twaalf Vectoren Worden Eén

Na Elke Head Berekent

Elke head produceert een outputtensor met vorm [batch, seq_len, head_dim]. Twaalf heads produceren twaalf dergelijke tensoren. Concatenatie langs de feature-dimensie stapelt ze weer samen:

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

Concat keert de split om. De totale feature-dimensie keert terug naar d_model. Geen informatieverlies in de dimensies; het verschil zit in wat elke chunk bevat: de chunk van head 1 weerspiegelt het geleerde attention-patroon van head 1.

De Output Projectie W_O

Concatenatie alleen laat de heads geïsoleerd: de output van head 4 ligt naast de output van head 7, zonder dat ze van elkaar op de hoogte zijn. De output projectie W_O van vorm [d_model, d_model] mixt ze:

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

Na W_O draagt elke outputdimensie een geleerde lineaire combinatie van alle twaalf koppen. Informatie stroomt vrij tussen koppen via deze enkele matrixvermenigvuldiging.

Waarom koppen specialiseren

Niets in de architectuur dwingt head 4 om werkwoordtijd te leren of head 9 om gepaarde leestekens te leren. Specialisatie ontstaat door gradiëntdruk: tijdens de training ontvangen heads die redundant bijdragen kleinere gradiënten dan heads die uniek bijdragen. Over duizenden stappen vestigt elke head zich in een niche die het totale verlies het meest effectief vermindert.

Empirisch tonen getrainde transformers heads die omgaan met: positionele patronen (head kijkt naar het vorige token), syntactische patronen (head kijkt naar de bijpassende sluitende haakje), semantische patronen (head kijkt naar de meest recente benoemde entiteit). Geen labels trainen deze specialisatie. Het trainingssignaal alleen, gepropageerd door W_O, sorteert de heads.

Waarom Twaalf Heads, Niet Eén Breder Head

Hoe CUDA de Heads opslaat

Een Enkele Tensor, Herschaald

De trainingsengine van ANDREA microgpt_cuda.cu reserveert geen twaalf afzonderlijke buffers voor twaalf heads. Het reserveert één gefuseerde tensor & behandelt de head-dimensie als een stride-patroon:

// na Q = X · W_Q (één matmul, gefuseerd over heads)
// Q heeft vorm [batch, seq_len, d_model]
// reshape naar    [batch, seq_len, n_head, head_dim]
// transpose naar  [batch, n_head, seq_len, head_dim]
// elke head nu aaneengesloten in de twee binnenste dimensies

De transpose verplaatst n_head vóór seq_len. Waarom? Omdat de volgende operatie (Q_h · K_h^T) een aaneengesloten seq_len × head_dim-slice per head in het geheugen nodig heeft. CUDA matmuls werken sneller op aaneengesloten tensors.

Eén Kernel, Veel Heads

Een enkele attention CUDA-kernel voert uit over alle heads parallel. Elk thread block verwerkt één (batch, head) paar; threads binnen het block werken samen aan de seq_len × head_dim tile. De kernel weet nooit dat hij meerdere heads verwerkt; de launch grid regelt de parallelle verwerking.

Configuratie Weerspiegelt de Hardware

De keuze van ANDREA-120M voor n_head=12, head_dim=64 sluit aan bij de RTX 4090 tensor cores, die matmul-tiles prefereren in veelvouden van 16. head_dim=64 = 4 × 16 past exact bij de tile-vorm. head_dim=32 (ANDREA-12M) past ook, maar gebruikt de tile onderbenut. head_dim=72 zou niet passen & zou fallback-kernels forceren.

Finale Overzicht

Vijf stappen. Drie matmuls raken de input (Q, K, V projecties). Eén matmul raakt de geconcateneerde heads (W_O). Eén attention kernel handelt alle heads parallel. ANDREA-120M voert alle vijf stappen één keer uit per laag, twaalf lagen diep, elke forward pass.

Het lezen van de pijplijn