Query, Key, Value

Passo 3 de 7. Entra: \(X\) do passo 2, \(5 \times 4\). Sai: para cada token um vetor novo que mistura os outros, calculado aqui para uma cabeça; o passo 4 acrescenta a máscara causal e a segunda cabeça. ← Embeddings · próximo passo: Máscaras e cabeças →

Um dicionário suave

Lembre a ideia da visão geral: cada token monta um vetor novo misturando os vetores dos outros, com mais peso para os que importam — "banco" puxa de "praça" numa frase e de "dinheiro" na outra.

Como decidir quanto pesa cada vizinho? A resposta vem de uma estrutura que você já conhece, o dicionário:

estoque = {"maçã": 3, "pera": 7}
estoque["pera"]     # a chave casa exatamente com uma entrada e devolve o valor dela: 7

A atenção é um dicionário suave. Em vez de exigir casamento exato, ela compara a pergunta com todas as chaves, dá a cada uma uma nota de semelhança e devolve uma média ponderada de todos os valores. Como a média é uma função suave das notas, dá para derivar — e a rede aprende, por gradiente, o que perguntar e o que anunciar.

A ferramenta para "o quanto se parecem" é o produto escalar, \(q \cdot k = \sum_i q_i k_i\): grande e positivo quando os dois vetores apontam para o mesmo lado, perto de zero quando não têm relação, negativo quando apontam para lados opostos. Tudo o que vem abaixo é essa operação, feita para cada par de tokens.

Query, Key, Value

O dicionário suave precisa de três coisas de cada token e a atenção dá um nome a cada uma:

• Query (consulta) — o que esta posição está procurando. É a pergunta.
• Key (chave) — o que cada posição anuncia sobre si mesma. É a etiqueta que será comparada com a pergunta.
• Value (valor) — o que cada posição de fato entrega, se for escolhida. É o conteúdo.

Pontue cada chave contra a query com um produto escalar, transforme os scores em uma distribuição com o softmax e devolva a mistura correspondente de valores:

Q, K e V são três projeções lineares diferentes da mesma entrada: \(Q = XW_Q\), \(K = XW_K\), \(V = XW_V\). As matrizes \(W_Q\), \(W_K\) e \(W_V\) são os únicos parâmetros aprendidos da camada. As três projeções são o motivo pelo qual um token pode pedir uma coisa, anunciar outra e entregar uma terceira — um pronome pode consultar por "o substantivo a que me refiro" sem ele próprio ser um substantivo.

O painel percorre essa cadeia uma matriz de cada vez, com as formas de uma frase de três tokens e os números do exemplo feito à mão que vem a seguir. Escolha um token para acompanhar a linha dele.

Um exemplo feito à mão

Fique com a primeira das duas frases do "banco" da visão geral e três palavras dela: "sentei · banco · praça" (o "no" e o "da" ficam de fora só para a conta caber). Vamos calcular a saída de "banco". Os vetores têm duas dimensões e foram escolhidos à mão; os valores têm um significado legível, \([\text{assento};\ \text{dinheiro}]\):

A query de "banco" é \(q = [1;\ 2]\) e \(d_k = 2\).

1. Scores. O produto escalar de \(q\) com cada chave: \(q \cdot k_\text{sentei} = 1\), \(q \cdot k_\text{banco} = 1 + 2 = 3\), \(q \cdot k_\text{praça} = 4\).
2. Escala. Dividindo por \(\sqrt{2} \approx 1{,}41\): \([0{,}71;\ 2{,}12;\ 2{,}83]\).
3. Softmax. \(e^{0{,}71} \approx 2{,}0\), \(e^{2{,}12} \approx 8{,}3\), \(e^{2{,}83} \approx 16{,}9\); somam \(27{,}3\). Dividindo cada um pela soma, os pesos são \(a = [0{,}07;\ 0{,}31;\ 0{,}62]\). Eles são positivos e somam 1.
4. Mistura. A saída é a média dos valores com esses pesos:

O "banco" que entrou em cima do muro, \([0{,}5;\ 0{,}5]\), sai da camada como \([0{,}85;\ 0{,}15]\): assento. Nenhum peso foi fixado por posição. A escolha veio de quanto a query de "banco" se parece com a chave de cada vizinho.

O painel começa exatamente nesses números. Refaça a conta acompanhando os passos ①–④, depois experimente:

• troque o valor de "praça" para \([0;\ 1]\) — é como se a frase fosse "…banco… dinheiro" — e veja o veredito mudar sem que nenhum peso de atenção se mexa;
• mova a query para \([1;\ 0]\) e veja o peso sair de "praça" e se dividir entre "sentei" e "banco";
• zere a query: todos os scores viram 0 e a atenção vira uma média simples.

O mesmo, em código

A função inteira cabe em três linhas. O script refaz a conta de cima e depois roda a mesma função para as três queries ao mesmo tempo — a linha do meio de \(A\) e de \(O\) é o "banco" que você calculou à mão:

"""O exemplo feito à mão na página, refeito em código — e depois a mesma conta para todas as linhas.

Contexto: "sentei · banco · praça". A query é a de "banco"; os values têm duas dimensões,
[assento, dinheiro]. O "banco" sozinho é ambíguo ([0.5, 0.5]); o contexto decide.
"""
import numpy as np

def softmax(s):
    e = np.exp(s - s.max(axis=-1, keepdims=True))     # subtrair o máximo não muda o resultado
    return e / e.sum(axis=-1, keepdims=True)

def attention(Q, K, V):
    S = Q @ K.T / np.sqrt(K.shape[-1])                 # (n, n): score de cada par
    A = softmax(S)                                     # (n, n): cada linha soma 1
    return A @ V, A                                    # (n, d_v): mistura de values

K = np.array([[1.0, 0.0],      # sentei
              [1.0, 1.0],      # banco
              [0.0, 2.0]])     # praça
V = np.array([[1.0, 0.0],      # sentei -> assento
              [0.5, 0.5],      # banco  -> ambíguo
              [1.0, 0.0]])     # praça  -> assento
q = np.array([1.0, 2.0])       # a query de "banco"

np.set_printoptions(precision=2, suppress=True)
s = K @ q
print("q·k         =", s)
print("q·k / √2    =", s / np.sqrt(2))
print("a = softmax =", softmax(s / np.sqrt(2)))
print("Σ a·v       =", softmax(s / np.sqrt(2)) @ V)

# A mesma função, para as três queries ao mesmo tempo: uma linha de A por token.
Q = np.array([[2.0, 0.0],      # sentei
              [1.0, 2.0],      # banco (a mesma query de cima)
              [0.0, 1.0]])     # praça
O, A = attention(Q, K, V)
print("\nA =\n", A)
print("O =\n", O)

q·k         = [1. 3. 4.]
q·k / √2    = [0.71 2.12 2.83]
a = softmax = [0.07 0.31 0.62]
Σ a·v       = [0.85 0.15]

A =
 [[0.45 0.45 0.11]
 [0.07 0.31 0.62]
 [0.14 0.28 0.58]]
O =
 [[0.78 0.22]
 [0.85 0.15]
 [0.86 0.14]]

O exemplo da página: uma cabeça, à mão

O exemplo do banco escolheu \(q\), \(k\) e \(v\) à mão. Num modelo eles são calculados a partir de \(X\) e é daí que o exemplo continua: o \(X\) da página de embeddings, cinco linhas de quatro números, uma por token de "o gato subiu no".

Passo 0 — a mesma escala para toda linha. Num bloco do GPT-2 a atenção não lê o próprio \(X\): lê uma cópia normalizada, \(u = \text{LayerNorm}(X)\): de cada linha subtraia a média \(\mu\) e depois divida pelo desvio-padrão \(\sigma\) (o porquê e o resto do bloco estão na página do bloco). Toda linha do nosso \(X\) já tem média 0; sobra só a divisão:

Para a última linha: \(\sigma = \sqrt{(0{,}25 + 0{,}25 + 0{,}25 + 2{,}25)/4} = \sqrt{0{,}75} = 0{,}866\) e \(1{,}5 / 0{,}866 = 1{,}732\).

Passo 1 — três projeções. \(Q = uW_Q\), \(K = uW_K\) e \(V = uW_V\), cada uma \(5 \times 4\). As matrizes de um modelo treinado são densas; as do brinquedo só escolhem colunas de \(u\), para que cada número possa ser rastreado:

A barra separa as duas cabeças: as colunas 1–2 são a cabeça 1, a desta página; as colunas 3–4 são a cabeça 2, da próxima página. Para o último token, ␣no, \(u = [-0{,}577; -0{,}577; -0{,}577; 1{,}732]\), então a sua query na cabeça 1 é \(q = [u_4, u_2] = [1{,}732; -0{,}577]\). As chaves \(k = [u_2, u_4]\) e os valores \(v = [u_1, u_2]\) da cabeça 1 para os cinco tokens:

Passos 2 a 4 — as quatro contas do exemplo do banco, de novo. \(q = [1{,}732; -0{,}577]\) e \(d_k = 2\):

④ A saída mistura os valores com esses pesos:

O último token põe 94% da sua atenção em ␣sub e iu: ele olha para o verbo, que o tokenizador tinha cortado em dois. Ninguém disse onde está o verbo; os pesos vieram de quanto a query dele se parece com cada chave e isso veio de \(W_Q\) e \(W_K\).

O painel roda a mesma cadeia para cada linha. As linhas 1 a 4 olham para tokens que vêm depois delas — o, na posição 1, põe 0,66 em ␣sub —, o que um modelo que gera texto não pode fazer; a próxima página proíbe isso com a máscara. A última linha não tem nada depois dela; por isso os seus números sobrevivem à máscara sem mudar.

As formas de cada matriz

Para uma frase de \(n\) tokens, cada um com \(d\) dimensões, é assim que as formas se encadeiam. Metade dos bugs de uma implementação nova aparece como um erro de forma nesta tabela:

No exemplo, \(n = 5\), \(d = 4\) e \(d_k = d_v = 2\) por cabeça: \(S\) e \(A\) são \(5 \times 5\) e \(O\) é \(5 \times 2\). No GPT-2 small, \(d = 768\) e \(d_k = d_v = 64\) por cabeça.

A matriz inteira e por que \(\sqrt{d_k}\) não é cosmético

Uma query contra três chaves é o mecanismo. O que uma camada de fato computa é a matriz \(n \times n\) completa: cada posição consulta cada posição. O painel abaixo a constrói para uma frase real. Leia por linha: a linha de um token mostra para onde ele olha e cada linha soma 1.

Cada opção em cabeça (head) é uma função de score diferente. As três primeiras descrevem padrões que realmente ocorrem em modelos treinados (uma cabeça que casa por conteúdo, uma que apenas olha um token atrás, uma que decai com a distância); a quarta mistura as três. Uma camada roda muitas cabeças ao mesmo tempo.

Agora desligue dividir por \(\sqrt{d_k}\) e arraste \(d_k\) para cima. A matriz vira binária.

A razão é um argumento de duas linhas. Se \(q\) e \(k\) têm componentes independentes com média 0 e variância 1, então

porque é uma soma de \(d_k\) termos independentes, cada um com variância 1. Os scores têm então desvio padrão \(\sqrt{d_k}\) — com \(d_k = 128\), cerca de 11. Scores rotineiramente diferem por 20. Um softmax sobre scores que diferem por 20 é um argmax duro: \(e^{-20} \approx 2 \times 10^{-9}\), então um peso vale 1, o resto vale \(10^{-9}\) e o gradiente que passa pelo softmax é zero. Dividir por \(\sqrt{d_k}\) restaura variância unitária e mantém a camada no regime em que ela ainda consegue aprender.

No exemplo da página \(d_k = 2\); o efeito é pequeno mas visível: sem a divisão os pesos da última linha vão de \([0{,}011; 0{,}035; 0{,}622; 0{,}319; 0{,}014]\) para \([0{,}002; 0{,}012; 0{,}707; 0{,}275; 0{,}003]\), já mais concentrados. Desmarque "dividir por √dₖ" no painel do exemplo, acima, para ver.

A volta por uma cabeça

Esta subseção é a mais técnica do capítulo. Na primeira leitura, basta a conclusão em negrito do terceiro parágrafo.

A frase acima — o gradiente que passa pelo softmax é zero — sai da conta, que custa quatro linhas. Com \(S = QK^\top/\sqrt{d_k}\), \(A = \text{softmax}(S)\) e \(O = AV\):

O termo do meio é a jacobiana do softmax e é onde está a lição. Repare no fator \(A \odot\): se uma linha de \(A\) é quase um one-hot — um peso valendo 1 e o resto valendo \(10^{-9}\) —, então quase todo o parêntese é multiplicado por \(10^{-9}\); e na única posição onde \(A\) não é minúsculo, o parêntese em si é praticamente zero, porque a média ponderada que ele subtrai é o próprio elemento. Uma atenção saturada devolve gradiente numericamente nulo em toda a linha. É disso que o \(\sqrt{d_k}\) protege.

A mesma fórmula dá uma invariante barata para conferir implementação: cada linha de \(\partial L/\partial S\) soma exatamente zero, porque somar uma constante a uma linha de scores não muda o softmax dela.

"""A passagem reversa de uma cabeça de atenção, conferida por diferenças finitas.

O termo que importa é o do softmax: ele não é um fator por elemento, é uma jacobiana que
subtrai a média ponderada da linha — e é ela que zera o gradiente quando a atenção satura.
"""
import numpy as np

def softmax(s):
    e = np.exp(s - s.max(axis=-1, keepdims=True))
    return e / e.sum(axis=-1, keepdims=True)

def forward(Q, K, V):
    d = Q.shape[-1]
    S = Q @ K.T / np.sqrt(d)
    A = softmax(S)
    return A @ V, A

def backward(dO, Q, K, V, A):
    d = Q.shape[-1]
    dV = A.T @ dO
    dA = dO @ V.T
    dS = A * (dA - (dA * A).sum(axis=-1, keepdims=True))     # jacobiana do softmax, linha a linha
    return dS @ K / np.sqrt(d), dS.T @ Q / np.sqrt(d), dV, dS

rng = np.random.default_rng(0)
Q, K, V = (rng.normal(size=(3, 4)) for _ in range(3))
dO = rng.normal(size=(3, 4))

O, A = forward(Q, K, V)
dQ, dK, dV, dS = backward(dO, Q, K, V, A)

def num_grad(M):
    g, h = np.zeros_like(M), 1e-6
    for i in range(M.shape[0]):
        for j in range(M.shape[1]):
            old = M[i, j]
            M[i, j] = old + h; up = (dO * forward(Q, K, V)[0]).sum()
            M[i, j] = old - h; dn = (dO * forward(Q, K, V)[0]).sum()
            M[i, j] = old; g[i, j] = (up - dn) / (2 * h)
    return g

np.set_printoptions(precision=3, suppress=True)
print("pesos de atenção A (cada linha soma 1):\n", A)
print("dQ:\n", dQ)
print("erro máximo contra diferenças finitas — dQ:", f"{np.abs(num_grad(Q) - dQ).max():.1e}",
      " dK:", f"{np.abs(num_grad(K) - dK).max():.1e}",
      " dV:", f"{np.abs(num_grad(V) - dV).max():.1e}")
print("soma de cada linha de dS:", np.round(dS.sum(axis=-1), 12),
      "— sempre zero: somar a mesma constante a uma linha de scores não muda o softmax")

pesos de atenção A (cada linha soma 1):
 [[0.226 0.475 0.298]
 [0.144 0.598 0.259]
 [0.668 0.235 0.097]]
dQ:
 [[0.456 0.354 0.099 0.209]
 [0.407 0.404 0.013 0.129]
 [0.345 0.228 0.109 0.184]]
erro máximo contra diferenças finitas — dQ: 3.2e-10  dK: 6.1e-10  dV: 3.4e-10
soma de cada linha de dS: [ 0. -0. -0.] — sempre zero: somar a mesma constante a uma linha de scores não muda o softmax

← Embeddings · próximo passo: Máscaras e cabeças →