Atenção: o que a fórmula não conta sobre custo

A fórmula em uma frase

Para cada token da sequência, a atenção calcula o quanto ele deve "olhar" para cada outro token e mistura as informações na proporção desses pesos. Cada token vira três vetores: uma query (o que estou procurando), uma key (o que eu ofereço) e um value (o que eu carrego). O produto entre queries e keys dá os pesos, o softmax normaliza, e o resultado é uma soma ponderada dos values.

Em código, cabe em seis linhas, e vale ver uma vez para o mistério acabar:

import numpy as np

def attention(Q, K, V):
    d_k = Q.shape[-1]
    scores = Q @ K.transpose(0, 2, 1) / np.sqrt(d_k)   # (batch, n, n)
    weights = np.exp(scores - scores.max(-1, keepdims=True))
    weights /= weights.sum(-1, keepdims=True)          # softmax
    return weights @ V                                 # (batch, n, d_v)

Foi isso que substituiu as redes recorrentes: em vez de empurrar um estado oculto token a token, todos os tokens se enxergam de uma vez, em paralelo. O ganho em paralelismo e em dependências de longa distância é o que tornou possível treinar modelos do tamanho dos atuais. Até aqui, é o que todo material conta.

O que a fórmula não conta

Olhe de novo para a dimensão de scores: (batch, n, n). Para uma sequência de n tokens, a atenção calcula e guarda n² pesos, por cabeça, por camada. Dobrar o contexto quadruplica esse custo. É a única linha do modelo que cresce assim, e é ela que dita três decisões de engenharia que aparecem depois que o modelo funciona.

Decisão 1: contexto longo ou recuperação

Modelos com janelas de centenas de milhares de tokens criam a tentação de "jogar tudo no prompt". Funciona, e custa quadraticamente em computação e linearmente em dinheiro a cada chamada, porque você paga por todos os tokens de entrada toda vez. A alternativa é recuperar só os trechos relevantes com busca vetorial e mandar um contexto curto. Perde-se a garantia de que a informação certa está lá; ganha-se latência, custo e a possibilidade de citar a fonte.

Decisão 2: memória de inferência e o cache de K e V

Na geração, cada novo token precisa olhar para as keys e values de todos os anteriores. Recalcular isso a cada passo seria absurdo, então os frameworks guardam K e V em cache. Esse cache cresce com o comprimento da conversa e com o número de cabeças e camadas, e é ele, não os pesos, que estoura a memória da GPU em produção com muitos usuários simultâneos. É por isso que variantes como multi-query e grouped-query attention existem: compartilham K e V entre cabeças para encolher o cache, aceitando uma pequena perda de qualidade em troca de servir mais requisições por placa.

Quando coloco um modelo aberto para servir, essa é a conta que faço antes de escolher a instância: tamanho dos pesos (fixo) mais cache de KV por usuário vezes usuários simultâneos vezes comprimento médio de contexto. O segundo termo costuma surpreender.

Decisão 3: multi-head e o que cada cabeça aprende

Na prática ninguém usa uma atenção só. O modelo roda várias "cabeças" em paralelo, cada uma com suas projeções de Q, K e V, e concatena os resultados. Cabeças diferentes aprendem relações diferentes: uma acompanha concordância sintática, outra correferência, outra posição. Em PyTorch isso é uma linha:

import torch, torch.nn as nn

mha = nn.MultiheadAttention(embed_dim=512, num_heads=8, batch_first=True)
x = torch.rand(1, 128, 512)          # (batch, tokens, embedding)
out, weights = mha(x, x, x)            # self-attention
weights.shape                          # (1, 128, 128): os n² pesos, já na média das cabeças

A decisão aqui é menos glamourosa do que parece: mais cabeças não é melhor por si só. Com a dimensão total fixa, mais cabeças significa cada cabeça com menos dimensões. O que importa é o produto, e os valores que funcionam já estão bem mapeados pelos modelos públicos. Onde a escolha de fato pesa é em interpretabilidade: olhar os pesos de atenção ajuda a entender por que o modelo errou, mas não é explicação causal. Já vi time concluir que o modelo "olhou para o token certo" e mesmo assim respondeu errado.

Onde isso me pegou de verdade

Em previsão de séries financeiras com Transformers, a janela de entrada é uma decisão de modelagem e de custo ao mesmo tempo. Janela maior captura regimes mais longos e custa quadraticamente no treino, que já é o gargalo quando se roda grid search sobre dezenas de combinações. No Nahas, parte do trabalho foi justamente escolher onde gastar: janela, número de cabeças, ou mais rodadas de busca de hiperparâmetros. Com orçamento fixo de GPU, cada um desses tira do outro.

Resumo da conversa

• Atenção é uma soma ponderada; o segredo está em quem define os pesos. A fórmula cabe em seis linhas.
• O custo é quadrático no comprimento da sequência, e essa linha dita contexto, memória e dinheiro.
• Contexto longo e recuperação são escolhas com perdas diferentes. Meça a fração do contexto que o modelo usa.
• Em produção, o cache de K e V, não os pesos, costuma ser o que estoura a memória.
• Pesos de atenção ajudam a investigar erros, mas não são explicação causal.

Referências que valem a leitura: o artigo original Attention Is All You Need, o Illustrated Transformer de Jay Alammar e a documentação de MultiheadAttention no PyTorch.