Atenção como aposta: e se cada cabeça do Transformer fosse uma crença sobre o mundo?

Modelos de linguagem treinados em sequências de algumas centenas de tokens precisam, na prática, lidar com documentos inteiros — relatórios médicos, contratos, livros, bases de código. Essa habilidade chama-se context length extrapolation: a capacidade de generalizar para sequências muito mais longas do que as vistas durante o treino. É uma das competências mais valiosas — e mais frustrantes — dos LLMs modernos.

A engrenagem responsável por isso é o positional encoding (PE), que diz ao modelo onde cada token está, já que a atenção sozinha é cega a ordem. Vários métodos foram propostos: Sinusoidal, RoPE, ALiBi, e até a opção radical de não usar nada (NoPE). Todos funcionam em algum sentido, mas dois problemas persistem. Primeiro, faltam fundamentos teóricos sólidos — boa parte dessas técnicas surgiu de intuições empíricas. Segundo, a métrica mais usada para avaliá-las, a perplexidade, pode ser enganosa: um modelo consegue mantê-la baixa apenas prestando atenção nos vizinhos imediatos, sem realmente acessar contexto longo.

É contra esse pano de fundo que Arthur S. Bianchessi, Yasmin C. Aguirre, Rodrigo C. Barros e Lucas S. Kupssinskü — pesquisadores do laboratório MALTA da PUCRS — propõem o Bayesian Attention Mechanism (BAM): um arcabouço probabilístico que reformula a atenção como uma esperança sobre valores, calculada a partir de uma distribuição conjunta sobre conteúdo e posição dos tokens. Sob essa lente, positional encoding deixa de ser um truque de engenharia e passa a ser, literalmente, uma crença a priori sobre quais posições importam.

O reenquadre bayesiano

A definição central do BAM é elegante. A clássica self-attention é reescrita como uma esperança: o valor que sai da camada de atenção é a média dos vetores de valor, ponderada pela probabilidade de o token j ser relevante para a query i. O salto vem quando os autores mostram que essa probabilidade pode ser fatorada em uma componente de conteúdo (similaridade entre query e key) e uma componente posicional (a crença a priori sobre a relevância da posição).

Quando a função de score é aditiva — algo verdadeiro para boa parte dos PEs existentes — essa fatoração emerge naturalmente. Com isso, o paper mostra que NoPE equivale a usar uma distribuição Uniforme como prior sobre posições, e ALiBi corresponde a uma distribuição de Laplace. Métodos que pareciam ad-hoc passam a ter nome, forma e família.

A Gaussiana Generalizada

Se ALiBi é Laplace, e Laplace é um caso particular de uma família maior de distribuições, por que não usar a família inteira? É o que os autores fazem ao propor o GGD-BAM, onde o prior posicional é uma Distribuição Gaussiana Generalizada (GGD), parametrizada por forma β, escala α e localização μ.

Com β = 1, o GGD-BAM recupera ALiBi. Mas com β entre 0 e 1, a cauda da distribuição fica mais pesada e a atenção passa a contemplar tokens distantes com probabilidade não desprezível — algo que ALiBi não consegue, já que sua atenção decai exponencialmente e vira efetivamente local em contextos longos.

O lance mais ousado: e se β < 0? Tecnicamente, isso quebra a definição estrita de uma GGD. Mas o resultado é fascinante — o modelo passa a ignorar o contexto local e a focar exclusivamente em tokens distantes. Surgem cabeças de atenção especializadas em recuperação de informação remota: verdadeiras "cabeças farejadoras" de longo alcance.

Três tipos de cabeça, um ecossistema

Depois do treino, as cabeças de atenção se organizam em três clusters no espaço de parâmetros. O primeiro, com β > 0, comporta-se como ALiBi clássico — bom para dependências locais e perplexidade baixa. O segundo, com −0,6 ≤ β ≤ 0, abriga as cabeças que funcionam como retrievers, capazes de pescar informação ao longo de toda a sequência. O terceiro, menor e mais raro, com β < −0,6, reúne cabeças totalmente cegas ao local e dedicadas ao distante.

Esses três agrupamentos aparecem de forma estável tanto em diferentes tamanhos de modelo (120M, 432M, 1.1B parâmetros) quanto em diferentes tamanhos de contexto de treino (512, 1024, 2048 tokens), sugerindo que essa estrutura é uma propriedade emergente do problema, não um acidente.

A prova de fogo: 500× o contexto de treino

Como medir se um modelo realmente acessa contexto longo, e não apenas mantém perplexidade baixa via atenção local? Os autores recorrem ao Passkey Retrieval: enterra-se uma chave numérica de cinco dígitos em algum ponto de uma sequência longa, preenchida com texto irrelevante, e mede-se se o modelo consegue recuperá-la. É uma versão minimalista — e cruel — do clássico "agulha no palheiro".

Os números são impressionantes. Modelos de 120M parâmetros treinados com contexto de 512 tokens foram avaliados em sequências de até 32 mil tokens — 64× o contexto de treino. Sinusoidal, RoPE e NoPE colapsam para acurácia próxima de zero quase imediatamente. ALiBi também se sai mal: sua acurácia despenca para cerca de 10% logo no início da extrapolação. BAM com Scalable Softmax mantém 100% de acurácia em todo o intervalo testado.

Nos modelos maiores (1.1B parâmetros), o passkey foi recuperado com sucesso em sequências de até 512 mil tokens — cerca de 500× o contexto de treino. Os autores comentam, com humor, que não puderam testar em contextos ainda maiores porque a VRAM acabou.

Além do passkey: RULER e LongBench

Para evitar a suspeita de overfit na tarefa específica, os autores avaliam o método também no benchmark RULER (subtarefa Needle-in-a-Haystack) e no LongBench v2, que mede compreensão real de documentos longos. Em ambos, BAM SSMax supera todas as alternativas.

No LongBench v2, com modelos de 1B parâmetros, BAM atinge score geral de 28,6 contra 24,2 do RoPE SSMax, com vantagens em compreensão de repositórios de código (41,7 vs 25,0), QA de documento único (26,5 vs 18,8) e in-context learning longo (36,4 vs 30,3). Em downstream tasks padrão como MMLU e ARC, BAM também supera o RoPE — os ganhos não vêm às custas das capacidades fundamentais do modelo.

384 parâmetros adicionais

Talvez o dado mais surpreendente seja o custo do GGD-BAM. Para o modelo de 120M, são exatamente 384 parâmetros aprendíveis a mais — 0,00032% do total. Em troca, ganha-se a capacidade de extrapolar para contextos centenas de vezes maiores que o treino, sem impacto perceptível no tempo de inferência. É um custo praticamente nulo para um benefício enorme.

Scalable Softmax como parceiro natural

Um ingrediente importante é o Scalable Softmax (SSMax). Ele resolve um problema complementar: o attention fading. Conforme o contexto cresce, o denominador da softmax cresce também, e a distribuição de atenção tende a se aplainar. SSMax aplica um fator de reescala que mantém a "afiação" da distribuição independentemente do tamanho da entrada.

BAM e SSMax se combinam naturalmente: o primeiro cuida do prior posicional, o segundo cuida da normalização. Curiosamente, o parâmetro aprendível do SSMax tem papel matemático muito parecido com o fator de normalização que aparece na derivação bayesiana do BAM — como se duas linhas de pesquisa, partindo de motivações diferentes, tivessem chegado ao mesmo lugar.

Interpretabilidade como bônus

As curvas de prior do BAM não são heurísticas — são distribuições de probabilidade reais e visualizáveis. É possível olhar para uma cabeça de atenção treinada e dizer, com precisão matemática, qual é a sua "filosofia": localista, distribuída ou especialista em recuperação remota? O paper inclui visualizações de pesos de atenção em diferentes regimes de β, mostrando empiricamente o que a teoria previa. Isso transforma positional encoding em algo que pode ser inspecionado, ajustado e projetado deliberadamente — um avanço em uma área que costuma operar como caixa-preta.

O que fica no horizonte

Permanecem em aberto questões relevantes: como BAM se comporta após instruction tuning ou alinhamento por preferências? Os ganhos se preservam em cenários multimodais? E como ele escala para modelos com dezenas ou centenas de bilhões de parâmetros?

Mas talvez a pergunta mais instigante seja filosófica. O BAM mostra que é possível reformular um componente fundamental dos Transformers como uma escolha bayesiana explícita. Quantas outras peças desses modelos, que hoje tratamos como engenharia, estão esperando uma releitura probabilística? Atenção pode ser uma. Normalização pode ser outra.

A história do BAM é, no fundo, sobre o que acontece quando se olha com calma para algo que parecia resolvido. Ao perguntar "o que isso é, em termos probabilísticos?", surge um framework que unifica métodos existentes, permite criar novos e melhora a generalização para contextos longos em ordem de grandeza. É um lembrete de que rigor teórico não é o oposto de impacto prático — pelo contrário, costuma ser o caminho mais curto para chegar lá.