Multi-head attention

1

Conceito

Uma única operação de attention cria uma distribuição de pontuações para cada query e retorna uma mistura ponderada de values. Ela é expressiva, mas obriga um conjunto de projeções a servir a todas as relações úteis. Multi-head attention cria várias operações em paralelo. Cada head recebe seu próprio subespaço aprendido de query, key e value, permitindo vários padrões de roteamento antes da combinação.

Se a largura do modelo é $d_{model}$ e o número de heads é $h$, um projeto comum define $d_{head}=d_{model}/h$. Implementações costumam usar três projeções lineares grandes e depois mudar o shape, em vez de criar módulos separados:

$B$ é batch e $T$ é comprimento. Attention roda de forma independente no eixo das heads:

A máscara $M$ impõe causalidade ou padding. As saídas são transpostas e concatenadas de volta para $B\times T\times d_{model}$, depois multiplicadas por $W_O$. A projeção final permite misturar features entregues por heads distintas antes de somar o resultado ao residual stream.

Por que dividir a largura? Uma query talvez precise de informação sobre pontuação próxima, um substantivo correferente e o início de uma citação. Heads separadas formam espaços de similaridade diferentes; “relevante” não precisa significar a mesma coisa em todos. Uma pode privilegiar posição relativa, outra conteúdo. Camadas posteriores compõem essas rotas.

Isso é uma capacidade, não uma promessa de interpretabilidade limpa. Pesquisas mostraram que muitas heads de certos modelos e tarefas podem ser removidas com pouco dano imediato, enquanto um subconjunto menor é importante. Redundância pode ajudar otimização ou robustez, e um cálculo significativo pode estar distribuído por heads e camadas. Batizar uma head de “head de sintaxe” com poucos exemplos é uma hipótese a testar causalmente.

Se a largura total permanece fixa e $h d_{head}=d_{model}$, as grandes projeções de $Q$, $K$, $V$ e saída têm aproximadamente a mesma contagem de parâmetros sob diferentes números de heads. Mais heads não significam automaticamente parâmetros proporcionalmente maiores; mudam a partição da largura. Ainda assim, alteram eficiência de kernels, armazenamento de matrizes e gargalos de representação.

Heads muito estreitas podem perder capacidade por rota. Heads muito largas reduzem o número de padrões independentes. Arquiteturas modernas às vezes usam quantidades diferentes de query heads e key/value heads para acelerar inference; MQA e GQA preservam várias queries enquanto compartilham keys e values no cache.

Na implementação, a ordem dos eixos é uma fonte frequente de bugs. A multiplicação deve parear cada query com keys dentro do mesmo item e head. Softmax roda sobre o eixo das posições de key. Depois da mistura dos values, o código restaura a ordem dos tokens antes da concatenação. Shapes não são burocracia em torno do algoritmo; eles expressam o algoritmo.

O modelo mental durável é um banco de canais aprendidos de comunicação. Cada head decide onde ler em seu próprio sistema de coordenadas, produz uma mensagem contextual e a entrega a um misturador comum de saída.