Trois matrices par tête, ou une grande matrice

Vue par tête

Chaque tête a besoin de sa propre projection de requête, projection de clé et projection de valeur. Pour la tête h :

Q_h = X · W_Q^h    où W_Q^h a la forme [d_model, head_dim]
K_h = X · W_K^h    où W_K^h a la forme [d_model, head_dim]
V_h = X · W_V^h    où W_V^h a la forme [d_model, head_dim]

X porte la forme d'entrée [batch, seq_len, d_model]. Après projection, Q_h, K_h, V_h portent chacun la forme [batch, seq_len, head_dim].

Vue Fusionnée

Les matrices par tête sont juxtaposées les unes à côté des autres en mémoire. Une matrice fusionnée unique W_Q de forme [d_model, d_model] produit toutes les têtes en une seule fois :

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

Le matmul fusionné effectue un seul appel BLAS au lieu de 12. Les cœurs de tenseur CUDA atteignent leur débit maximal sur des matmuls de cette taille ; les matmuls par tête sous-utiliseraient le matériel.

Nombre de paramètres

Trois matrices fusionnées W_Q, W_K, W_V, chacune de taille d_model × d_model. Plus la projection de sortie W_O, également d_model × d_model. Pour ANDREA-120M :

params par couche d'attention = 4 × 768² = 2,359,296 ≈ 2,36M
params sur 12 couches      = 12 × 2,36M ≈ 28,3M

Environ un quart des paramètres totaux d'ANDREA-120M se trouvent dans les projections d'attention. Les trois quarts restants se trouvent dans la sous-couche MLP & les embeddings.

Nommer les Projections

Douze Vecteurs Deviennent Un

Après le calcul de chaque tête

Chaque tête produit un tenseur de sortie de forme [batch, seq_len, head_dim]. Douze têtes produisent douze tenseurs de ce type. La concaténation le long de la dimension des caractéristiques les réunit à nouveau :

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

Concat inverse le split. La dimension totale des caractéristiques revient à d_model. Pas de perte d'information dans les dimensions ; la différence réside dans ce que contient chaque chunk : le chunk de la tête 1 reflète le motif d'attention appris de la tête 1.

La Projection de Sortie W_O

La concaténation seule laisse les têtes isolées : la sortie de la tête 4 est juxtaposée à celle de la tête 7, sans qu'elles soient conscientes l'une de l'autre. La projection de sortie W_O de forme [d_model, d_model] les mélange :

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

Après W_O, chaque dimension de sortie transporte une combinaison linéaire apprise de toutes les douze têtes. L'information circule librement entre les têtes grâce à cette unique multiplication matricielle.

Pourquoi les têtes se spécialisent

Rien dans l'architecture n'oblige la tête 4 à apprendre le temps des verbes ou la tête 9 à apprendre la ponctuation appariée. La spécialisation émerge de la pression des gradients : pendant l'entraînement, les têtes qui contribuent de manière redondante reçoivent des gradients plus petits que celles qui contribuent de manière unique. Sur des milliers d'étapes, chaque tête se niche dans un rôle qui réduit le plus efficacement la perte totale.

Empiriquement, les transformers entraînés montrent des têtes qui gèrent : des motifs positionnels (la tête regarde le token précédent), des motifs syntaxiques (la tête regarde la parenthèse fermante correspondante), des motifs sémantiques (la tête regarde l'entité nommée la plus récente). Aucune étiquette n'entraîne cette spécialisation. Le signal d'entraînement seul, propagé à travers W_O, trie les têtes.

Pourquoi douze têtes, et non une tête plus large

Comment CUDA stocke les têtes

Un seul tenseur, remodelé

Le moteur d'entraînement d'ANDREA microgpt_cuda.cu n'alloue pas douze tampons séparés pour douze têtes. Il alloue un tenseur fusionné & traite la dimension des têtes comme un motif de stride :

// après Q = X · W_Q (une multiplication matricielle, fusionnée à travers les têtes)
// Q a la forme [batch, seq_len, d_model]
// remodeler en    [batch, seq_len, n_head, head_dim]
// transposer en  [batch, n_head, seq_len, head_dim]
// chaque tête est maintenant contiguë dans les deux dimensions internes

La transposition place n_head avant seq_len. Pourquoi ? Parce que l'opération suivante (Q_h · K_h^T) nécessite que la tranche seq_len × head_dim de chaque tête soit contiguë en mémoire. Les matmuls CUDA s'exécutent plus rapidement sur des tenseurs contigus.

Un Kernel, Beaucoup de Têtes

Un seul noyau CUDA d'attention s'exécute sur toutes les têtes en parallèle. Chaque bloc de threads gère une paire (batch, tête) ; les threads à l'intérieur du bloc coopèrent sur la tuile seq_len × head_dim. Le noyau ne sait jamais qu'il traite plusieurs têtes ; la grille de lancement gère le parallélisme.

Configuration Reflects the Hardware

Le choix de n_head=12, head_dim=64 pour ANDREA-120M s'aligne avec les cœurs tenseurs de la RTX 4090, qui préfèrent les tuiles matmul en multiples de 16. head_dim=64 = 4 × 16 correspond exactement à la forme de la tuile. head_dim=32 (ANDREA-12M) correspond aussi mais sous-utilise la tuile. head_dim=72 ne correspondrait pas & forcerait des noyaux de repli.

Final Picture

Cinq étapes. Trois multiplications matricielles touchent l'entrée (projections Q, K, V). Une multiplication matricielle touche les têtes concaténées (W_O). Un noyau d'attention gère toutes les têtes en parallèle. ANDREA-120M exécute les cinq étapes une fois par couche, douze couches de profondeur, à chaque passe avant.

Lecture du Pipeline