Masks and Heads

Passo 4 de 7. Entra: \(Q\), \(K\) e \(V\) do passo 3, \(5 \times 4\) cada uma. Sai: a saída da atenção \(a\), também \(5 \times 4\), que o bloco soma de volta a \(X\). ← Query, key, value · próximo passo: O bloco →

Máscaras: causal, de padding e todas as outras

Às vezes um token não pode olhar para outro: o futuro, num modelo que gera texto; o enchimento <pad>, num batch. A máscara é somada aos scores antes do softmax, como \(-\infty\) (na prática, um número negativo grande). Como \(e^{-\infty} = 0\), os pesos proibidos ficam exatamente zero:

Ligue a máscara causal no painel da matriz da página anterior e o triângulo superior escurece. Essa única linha é toda a diferença entre um encoder e um decoder e a razão pela qual um decoder pode ser treinado em todas as posições em paralelo: com a máscara, a previsão da posição \(i\) não pode ver a própria resposta, então uma passagem direta dá \(n\) exemplos de treino em vez de um. Para a frase do exemplo com a resposta, "o gato subiu no telhado", em tokens:

Cinco exemplos, uma única passagem pela rede. Sem a máscara, a posição 1 veria ␣gato logo ao lado e aprenderia a copiar em vez de prever.

A máscara no exemplo

Na página anterior a cabeça 1 rodou sem máscara. Com ela, todo score acima da diagonal vira \(-\infty\) antes do softmax e cada linha divide o seu peso entre os tokens que sobram. A linha 2, ␣gato, tinha os scores \([-0{,}333; 0{,}471; -0{,}218; -0{,}577; -0{,}471]\); a máscara mantém os dois primeiros, \(e^{-0{,}333} = 0{,}717\) e \(e^{0{,}471} = 1{,}602\), que somam 2,319. Os pesos viram \(0{,}717/2{,}319 = 0{,}309\) e \(1{,}602/2{,}319 = 0{,}691\):

O primeiro token só pode olhar para si mesmo. A última linha não se mexe: é a linha que prevê ␣telhado e ela já podia ver tudo o que veio antes.

Multi-head e no que ela se transformou

Uma cabeça produz uma distribuição por posição — uma relação. Mas um token costuma precisar de várias ao mesmo tempo: um verbo quer achar o sujeito e olhar a palavra anterior; um pronome quer o substantivo a que se refere e a pontuação que fecha a frase. Uma única linha de pesos que soma 1 não consegue apontar forte para tudo isso. A atenção multi-cabeça roda \(h\) cabeças em paralelo, cada uma com suas próprias \(W_Q^i, W_K^i, W_V^i\) sobre fatias de dimensão \(d/h\) e concatena as saídas:

Com \(d = 4096\) e \(h = 32\), cada cabeça trabalha com vetores de 128 dimensões; as 32 saídas concatenadas voltam a ter 4096 e \(W^O\) mistura o que cada cabeça achou. As opções de cabeça do painel da matriz são exatamente esse tipo de especialização.

No código, as cabeças não são camadas separadas. Uma projeção \(d \times d\) calcula as queries de todas as cabeças de uma vez e a saída dela é cortada em \(h\) fatias de largura \(d/h\); o mesmo vale para K e V. Para um batch de \(B\) sequências:

Os parâmetros são quatro matrizes \(d \times d\) (\(W_Q\), \(W_K\), \(W_V\) e \(W^O\), mais os vieses quando o modelo os usa), qualquer que seja \(h\).

A mesma coisa, como o laboratório a escreve: separar_cabecas, a função atencao com a máscara, a concatenação e \(W^O\), linha a linha, em três tokens com \(d = 4\) e \(h = 2\). A cabeça 1 é o exemplo do banco da página anterior (a linha de "banco" termina de novo em \([0{,}85;\ 0{,}15]\)); ligue a máscara causal ou ponha \(h = 1\) e veja quais linhas se mexem:

Repare no que isso não custa: com \(d/h\) por cabeça, \(h\) cabeças custam os mesmos FLOPs que uma cabeça de largura total. Cabeças são paralelismo de graça, não capacidade extra — e é por isso que \(h\) só cresceu.

Aí veio a inferência e o desenho mudou: o custo da atenção é o porquê.

Duas cabeças no exemplo

O brinquedo tem \(d = 4\) e \(h = 2\): a cabeça 1 fica com as colunas 1–2 de \(Q\), \(K\) e \(V\) (a cabeça da página anterior), a cabeça 2 com as colunas 3–4. Para o último token a cabeça 2 tem a mesma query da cabeça 1, \(q = [u_4, u_1] = [1{,}732; -0{,}577]\), mas outras chaves, \(k = [u_1, u_4]\):

A cabeça 1 achou o verbo (␣sub e iu, 0,94); a cabeça 2 achou o sujeito, ␣gato, 0,88. Os valores dela são as colunas 3–4 de \(u\); a sua saída para o último token é \(o^{(2)} = [-0{,}521; -0{,}520]\). As duas saídas vão lado a lado e \(W^O\) — a identidade no brinquedo, para as contas ficarem curtas — dá a saída da atenção do último token:

e, para os cinco tokens com a máscara causal,

Essa matriz é o que a subcamada de atenção entrega ao bloco e é onde o capítulo de Transformers retoma o exemplo. Com uma cabeça só de largura 4 em vez de duas de largura 2, a última linha se dividiria entre ␣gato (0,39) e ␣sub (0,48): uma linha de pesos que soma 1 não consegue apontar com força para os dois — o argumento a favor do multi-head, em números.

← Query, key, value · próximo passo: O bloco → · página à parte: O custo da atenção