Multi-head
e a máscara causal

1 O problema de uma cabeça só

A atenção da apostila 03 funciona, mas tem uma limitação séria: cada linha do mapa de atenção soma 1. Isso significa que a atenção é um orçamento fixo. Se um token gasta 70% olhando para seu antecedente sintático, sobram 30% para todo o resto.

Só que uma palavra participa de várias relações simultâneas. Pegue "mordeu" em "o cachorro grande mordeu o carteiro ontem":

• Precisa olhar para "cachorro" — quem é o sujeito?
• Precisa olhar para "carteiro" — quem é o objeto?
• Precisa olhar para "ontem" — quando isso aconteceu?
• Precisa olhar para "o" anterior — concordância, estrutura local.

Com uma cabeça só, essas quatro necessidades competem pelo mesmo orçamento de 100%. O modelo é forçado a fazer uma média entre relações que são qualitativamente diferentes — e uma média de coisas diferentes costuma ser uma coisa ruim.

2 Como funciona: dividir, não multiplicar

O detalhe que deixa a ideia elegante: multi-head não custa mais caro. Em vez de doze cabeças de 768 dimensões cada, o modelo divide as 768 dimensões existentes em doze fatias de 64.

O procedimento completo:

A matriz W^O é essencial e frequentemente ignorada. Sem ela, as doze cabeças ficariam em compartimentos estanques, cada uma escrevendo em suas 64 dimensões e nunca conversando. É W^O (chamada c_proj no GPT-2) que permite que a informação da cabeça 3 influencie as dimensões que a cabeça 7 usa.

As dimensões, na prática

GPT-2 small, frase de 5 tokens, 12 cabecas:

h                    (5, 768)
  |
c_attn (768 -> 2304) (5, 2304)          uma matmul so
  split em 3         Q,K,V: (5, 768) cada
  reshape p/ cabecas Q,K,V: (12, 5, 64) <- 12 cabecas em paralelo
  |
atencao por cabeca   (12, 5, 64)         12 mapas 5x5 independentes
  |
concat (juntar)      (5, 768)            reagrupa as fatias
c_proj (768 -> 768)  (5, 768)            mistura as cabecas
  |
saida                (5, 768)            mesma forma da entrada

O que as cabeças de fato aprendem

Ninguém atribui papéis. As especializações emergem do treino. Estudos de interpretabilidade em GPT-2 e BERT encontraram consistentemente cabeças como:

3 A máscara causal — a linha que define tudo

Agora chegamos ao ponto mais importante do curso inteiro. Preste bastante atenção nesta seção, porque tudo na Parte III depende dela.

O problema

O GPT-2 é treinado para prever a próxima palavra. Damos "O Brasil é" e ele deve produzir "gigante". Mas, para ser eficiente, o treino não roda uma palavra de cada vez — ele processa a frase inteira de uma vez e cobra todas as previsões simultaneamente:

Entrada:  O      Brasil  é       gigante  pela
Alvo:     Brasil é       gigante pela     própria
          ^      ^       ^       ^        ^
          5 previsoes cobradas ao mesmo tempo, num unico forward pass.

E aqui está a catástrofe: a atenção da apostila 03 deixa cada token ver todos os outros — inclusive os que vêm depois. Ao tentar prever que depois de "é" vem "gigante", o token "é" está olhando diretamente para "gigante".

O modelo aprenderia a tarefa trivial de copiar a resposta. Na hora de gerar texto de verdade, onde o futuro não existe, ele não saberia fazer nada. É vazamento de rótulo — o pecado capital do aprendizado supervisionado.

A solução: proibir explicitamente

Antes do softmax, some uma matriz de máscara que coloca −∞ em toda posição futura:

Por que −∞ e não zero?

Essa é a pergunta certa, e a resposta é bonita. Zerar o score não adianta — um score 0 ainda vira e⁰ = 1 no softmax, um peso perfeitamente respeitável. E zerar o peso depois do softmax quebraria a soma-1.

O −∞ funciona porque:

Na implementação usa-se um número muito negativo (-1e9, ou torch.finfo(dtype).min) para evitar NaN em aritmética de ponto flutuante.

O formato triangular

Mapa de atencao APOS a mascara (5 tokens). "." = proibido.

           O   Brasil  e   gigante  pela
   O     [0.99   .      .     .       .   ]  <- so ve a si mesmo
   Brasil[0.41  0.59    .     .       .   ]
   e     [0.12  0.55   0.33   .       .   ]
   gigante[0.08 0.31   0.19  0.42     .   ]
   pela  [0.05  0.22   0.11  0.38    0.24 ]  <- ve tudo ate aqui

Triangular inferior. Cada linha soma 1. O futuro nao existe.

A consequência prática: o cache KV

A máscara causal traz um presente de graça. Como o token 5 nunca influencia o token 3, os vetores K e V das posições anteriores nunca mudam quando um token novo é acrescentado.

Logo, ao gerar texto, não é preciso recalcular a frase inteira a cada palavra. Guarda-se K e V de tudo que já passou (o KV cache) e só se calcula o token novo. A geração cai de custo O(n²) por token para O(n).

Isso é o que torna a geração autoregressiva viável na prática — e o RoBERTa, sendo bidirecional, não pode fazer isso, porque nele acrescentar um token muda a representação de todos os outros.

4 O código completo

import torch, math
import torch.nn as nn
import torch.nn.functional as F

class AtencaoMultiCabeca(nn.Module):
    def __init__(self, d_model=768, n_cabecas=12, causal=True):
        super().__init__()
        self.h = n_cabecas
        self.d_k = d_model // n_cabecas      # 64
        self.causal = causal                  # <<< GPT=True, RoBERTa=False
        self.c_attn = nn.Linear(d_model, 3 * d_model)  # Q,K,V fundidos
        self.c_proj = nn.Linear(d_model, d_model)      # W^O

    def forward(self, x):
        B, n, d = x.shape

        # 1) uma matmul produz Q, K e V de uma vez
        q, k, v = self.c_attn(x).split(d, dim=2)

        # 2) (B, n, 768) -> (B, 12, n, 64): separa as cabecas
        q = q.view(B, n, self.h, self.d_k).transpose(1, 2)
        k = k.view(B, n, self.h, self.d_k).transpose(1, 2)
        v = v.view(B, n, self.h, self.d_k).transpose(1, 2)

        # 3) scores escalados -> (B, 12, n, n)
        att = (q @ k.transpose(-2, -1)) / math.sqrt(self.d_k)

        # 4) ===== A LINHA QUE DEFINE A ARQUITETURA =====
        if self.causal:
            futuro = torch.triu(torch.ones(n, n, dtype=torch.bool), diagonal=1)
            att = att.masked_fill(futuro, float("-inf"))
        # ==============================================

        att = F.softmax(att, dim=-1)
        y = att @ v                                    # (B, 12, n, 64)

        # 5) reagrupa as cabecas e mistura
        y = y.transpose(1, 2).contiguous().view(B, n, d)
        return self.c_proj(y)

Leia de novo o passo 4. causal=True → GPT-2. causal=False → RoBERTa. Todo o resto é idêntico.

5 No laboratório

.venv\Scripts\python.exe lab\04_multihead_mask.py

1. Imprime a máscara triangular em ASCII, para você ver a forma.
2. Roda a mesma entrada com causal=True e False, e mostra as duas matrizes de atenção lado a lado. A diferença visual é imediata.
3. Prova o vazamento: muda um token do fim da frase e mede o quanto a representação do início mudou. No modo causal: zero. No bidirecional: bastante. Este é o teste que demonstra causalidade em vez de afirmar.
4. Retrato das 12 cabeças da camada 0 do GPT-2 real: para cada cabeça, calcula a distância média a que ela olha e a entropia da sua distribuição. Você vai identificar as cabeças posicionais (distância ≈ 1) e as de "vala" (entropia baixíssima, tudo no token 0).
5. Mede o ganho do KV cache gerando 50 tokens com e sem ele.

A atenção move informação entre posições. Mas ela é essencialmente uma média ponderada — uma operação linear.
A apostila 05 mostra as peças que faltam: onde entra o pensamento não-linear, e como se empilham 48 camadas sem o treino explodir.