Drei Matrizen pro Head oder eine große Matrix

Sicht pro Head

Jeder Head benötigt seine eigene Query-Projektion, Key-Projektion und Value-Projektion. Für Head h:

Q_h = X · W_Q^h    wobei W_Q^h die Form [d_model, head_dim] hat
K_h = X · W_K^h    wobei W_K^h die Form [d_model, head_dim] hat
V_h = X · W_V^h    wobei W_V^h die Form [d_model, head_dim] hat

X trägt die Eingabeform [batch, seq_len, d_model]. Nach der Projektion haben Q_h, K_h, V_h jeweils die Form [batch, seq_len, head_dim].

Geführte Ansicht

Die pro-Head-Matrizen liegen nebeneinander im Speicher. Eine einzige geführte Matrix W_Q mit der Form [d_model, d_model] erzeugt alle Heads auf einmal:

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

Die fused Matmul ersetzt 12 BLAS-Aufrufe durch einen einzigen. CUDA Tensor Cores erreichen ihre maximale Durchsatzleistung bei Matmuls dieser Größe; pro-Head-Matmuls würden die Hardware unterauslasten.

Parameteranzahl

Drei fusionierten Matrizen W_Q, W_K, W_V, jeweils d_model × d_model. Plus die Ausgabprojektion W_O, ebenfalls d_model × d_model. Für ANDREA-120M:

params pro Layer-Attention = 4 × 768² = 2.359.296 ≈ 2,36M
params über 12 Layer      = 12 × 2,36M ≈ 28,3M

Ungefähr ein Viertel der Gesamtparameter von ANDREA-120M befindet sich in den Attention-Projektionen. Die restlichen drei Viertel befinden sich im MLP-Unterschicht & Embeddings.

Benennung der Projektionen

Zwölf Vektoren werden zu einem

Nachdem jeder Head berechnet hat

Jeder Head erzeugt ein Ausgabetensor der Form [batch, seq_len, head_dim]. Zwölf Heads erzeugen zwölf solcher Tensoren. Die Konkatenation entlang der Feature-Dimension stapelt sie wieder zusammen:

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 kehrt das Split um. Die Gesamtdimension der Merkmale kehrt zu d_model zurück. Kein Informationsverlust in den Dimensionen; der Unterschied liegt darin, was jeder Chunk enthält: Der Chunk von Head 1 spiegelt das gelernte Attention-Muster von Head 1 wider.

Die Output-Projektion W_O

Die reine Konkatenation lässt die Heads isoliert: Die Ausgabe von Head 4 liegt neben der Ausgabe von Head 7, ohne voneinander zu wissen. Die Output-Projektion W_O der Form [d_model, d_model] mischt sie:

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

Nach W_O trägt jede Ausgabedimension eine gelernte lineare Kombination aller zwölf Köpfe. Informationen fließen frei zwischen den Köpfen durch diese einzige Matrixmultiplikation.

Warum sich Köpfe spezialisieren

Nichts in der Architektur zwingt Kopf 4, die Verbzeit zu lernen, oder Kopf 9, gepaarte Satzzeichen zu lernen. Die Spezialisierung entsteht durch Gradientendruck: Während des Trainings erhalten Köpfe, die redundant beitragen, kleinere Gradienten als Köpfe, die einzigartig beitragen. Über Tausende von Schritten siedelt sich jeder Kopf in eine Nische an, die den Gesamtverlust am effektivsten reduziert.

Empirisch zeigen trainierte Transformer Köpfe, die folgendes handhaben: positionsbezogene Muster (Kopf schaut auf das vorherige Token), syntaktische Muster (Kopf schaut auf die passende schließende Klammer), semantische Muster (Kopf schaut auf die neueste benannte Entität). Keine Labels trainieren diese Spezialisierung. Das Trainingssignal allein, propagiert durch W_O, sortiert die Köpfe.

Warum zwölf Köpfe, nicht ein breiterer Kopf

Wie CUDA die Heads speichert

Ein einzelnes Tensor, umgeformt

ANDREA's Training-Engine microgpt_cuda.cu reserviert keine zwölf separaten Puffer für zwölf Heads. Sie reserviert ein fusioniertes Tensor & behandelt die Head-Dimension als Stride-Muster:

// nach Q = X · W_Q (eine Matmul, fusioniert über Heads)
// Q hat die Form [batch, seq_len, d_model]
// umformen zu    [batch, seq_len, n_head, head_dim]
// transponieren zu  [batch, n_head, seq_len, head_dim]
// jeder Head jetzt kontinuierlich in den inneren zwei Dimensionen

Die Transposition verschiebt n_head vor seq_len. Warum? Weil die nächste Operation (Q_h · K_h^T) den seq_len × head_dim-Ausschnitt jedes Heads kontinuierlich im Speicher benötigt. CUDA-MatMuls laufen schneller auf kontinuierlichen Tensoren.

Ein Kernel, viele Heads

Ein einzelner Attention-CUDA-Kernel läuft parallel über alle Heads. Jeder Thread-Block bearbeitet ein (Batch, Head)-Paar; Threads innerhalb des Blocks kooperieren bei der seq_len × head_dim-Kachel. Der Kernel weiß nie, dass er mehrere Heads verarbeitet; das Launch-Grid kümmert sich um die Parallelität.

Konfiguration spiegelt die Hardware wider

Die Wahl von n_head=12, head_dim=64 bei ANDREA-120M passt zu den Tensor-Cores der RTX 4090, die Matmul-Kacheln in Vielfachen von 16 bevorzugen. head_dim=64 = 4 × 16 passt exakt zur Kachelform. head_dim=32 (ANDREA-12M) passt ebenfalls, nutzt die Kachel aber unteraus. head_dim=72 würde nicht passen & würde Fallback-Kernel erzwingen.

Gesamtbild

Fünf Schritte. Drei Matmuls berühren den Input (Q, K, V Projektionen). Ein Matmul berührt die verketteten Heads (W_O). Ein Attention-Kernel verarbeitet alle Heads parallel. ANDREA-120M führt alle fünf Schritte einmal pro Schicht aus, zwölf Schichten tief, bei jedem Forward-Pass.

Lesen des Pipelines