Três Matrizes Por Cabeça, Ou Uma Grande Matriz

Visão Por Cabeça

Cada cabeça precisa de sua própria projeção de query, projeção de key e projeção de value. Para a cabeça h:

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

X carrega a forma de entrada [batch, seq_len, d_model]. Após a projeção, Q_h, K_h, V_h cada um carrega a forma [batch, seq_len, head_dim].

Visão Fundida

As matrizes por cabeça ficam lado a lado na memória. Uma única matriz fundida W_Q de forma [d_model, d_model] produz todas as cabeças de uma vez:

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

O matmul fundido envia uma única chamada BLAS em vez de 12. Os tensor cores do CUDA atingem o pico de throughput em matmuls desse tamanho; matmuls por cabeça subutilizariam o hardware.

Contagem de Parâmetros

Três matrizes fundidas W_Q, W_K, W_V, cada uma d_model × d_model. Mais a projeção de saída W_O, também d_model × d_model. Para ANDREA-120M:

parâmetros por atenção da camada = 4 × 768² = 2.359.296 ≈ 2,36M
parâmetros em 12 camadas      = 12 × 2,36M ≈ 28,3M

Aproximadamente um quarto dos parâmetros totais do ANDREA-120M está nas projeções de atenção. Os três quartos restantes estão na subcamada MLP & embeddings.

Nomeando as Projeções

Doze Vetores se Tornam Um

Após Cada Cabeça Computar

Cada cabeça produz um tensor de saída com formato [batch, seq_len, head_dim]. Doze cabeças produzem doze tensores assim. A concatenação ao longo da dimensão de características os empilha de volta juntos:

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

Concat reverte a divisão. A dimensão total de características retorna para d_model. Sem perda de informação nas dimensões; a diferença reside no que cada chunk contém: o chunk da cabeça 1 reflete o padrão de atenção aprendido da cabeça 1.

A Projeção de Saída W_O

A concatenação sozinha deixa as cabeças isoladas: a saída da cabeça 4 fica ao lado da saída da cabeça 7, sem que nenhuma esteja ciente da outra. A projeção de saída W_O de formato [d_model, d_model] as mistura:

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

Após W_O, cada dimensão de saída carrega uma combinação linear aprendida de todas as doze cabeças. A informação flui livremente entre as cabeças por meio desta única multiplicação de matriz.

Por Que as Cabeças se Especializam

Nada na arquitetura força a cabeça 4 a aprender tempo verbal ou a cabeça 9 a aprender pontuação pareada. A especialização emerge da pressão do gradiente: durante o treinamento, cabeças que contribuem de forma redundante recebem gradientes menores do que cabeças que contribuem de forma única. Ao longo de milhares de passos, cada cabeça se estabelece em um nicho que reduz a perda total de forma mais eficaz.

Empiricamente, transformadores treinados mostram cabeças que lidam com: padrões posicionais (cabeça olha para o token anterior), padrões sintáticos (cabeça olha para o colchete de fechamento correspondente), padrões semânticos (cabeça olha para a entidade nomeada mais recente). Nenhuma etiqueta treina essa especialização. O sinal de treinamento sozinho, propagado através de W_O, organiza as cabeças.

Por Que Doze Cabeças, Não Uma Cabeça Mais Larga

Como o CUDA Armazena as Cabeças

Um Único Tensor, Remodelado

O motor de treinamento da ANDREA microgpt_cuda.cu não aloca doze buffers separados para doze cabeças. Ele aloca um tensor fundido & trata a dimensão da cabeça como um padrão de stride:

// após Q = X · W_Q (uma multiplicação de matrizes, fundida entre cabeças)
// Q tem forma [batch, seq_len, d_model]
// remodelar para    [batch, seq_len, n_head, head_dim]
// transpor para  [batch, n_head, seq_len, head_dim]
// cada cabeça agora contígua nas duas dimensões internas

A transposição move n_head à frente de seq_len. Por quê? Porque a próxima operação (Q_h · K_h^T) precisa que a fatia seq_len × head_dim de cada cabeça esteja contígua na memória. Multiplicações de matrizes CUDA executam mais rápido em tensores contíguos.

Um Kernel, Muitas Cabeças

Um único kernel CUDA de atenção executa em todas as cabeças em paralelo. Cada bloco de threads lida com um par (batch, head); as threads dentro do bloco cooperam no tile seq_len × head_dim. O kernel nunca sabe que processa múltiplas cabeças; a grade de lançamento lida com o paralelismo.

Configuração Reflete o Hardware

A escolha de n_head=12, head_dim=64 do ANDREA-120M alinha-se com os tensor cores da RTX 4090, que preferem tiles de matmul em múltiplos de 16. head_dim=64 = 4 × 16 combina exatamente com a forma do tile. head_dim=32 (ANDREA-12M) também combina, mas subutiliza o tile. head_dim=72 não combinaria & forçaria kernels de fallback.

Visão Final

Cinco passos. Três matmuls tocam a entrada (projeções Q, K, V). Um matmul toca as cabeças concatenadas (W_O). Um kernel de atenção lida com todas as cabeças em paralelo. ANDREA-120M executa todos os cinco passos uma vez por camada, doze camadas de profundidade, a cada passagem forward.

Lendo o Pipeline