Resumo Técnico: Mecanismo de Atenção Bayesiana (BAM)

Declaração do Problema

Modelos de linguagem baseados em Transformer dependem de Codificação Posicional (PE) para lidar com a ordem dos tokens, no entanto, os métodos existentes enfrentam duas limitações críticas em relação à extrapolação do comprimento do contexto (a capacidade de generalizar para sequências mais longas do que aquelas vistas durante o treinamento):

1. Opacidade Teórica: Muitas técnicas de PE (por exemplo, Sinusoidal, RoPE, ALiBi, NoPE) são motivadas empiricamente com fundamentos teóricos limitados, tornando seu comportamento durante a extrapolação pouco compreendido.
2. Avaliação Inadequada: A dependência da perplexidade como métrica primária é insuficiente. Modelos podem alcançar baixa perplexidade através de padrões de atenção localizados (focando apenas em tokens recentes) sem realmente integrar informações de longo alcance, uma falha destacada por avaliações com janela deslizante.

Metodologia: O Mecanismo de Atenção Bayesiana (BAM)

Os autores propõem o Mecanismo de Atenção Bayesiana (BAM), uma estrutura teórica que reformula a autoatenção como uma expectativa sobre valores computados sob um modelo probabilístico conjunto de conteúdo e posição do token.

Estrutura Teórica

• Formulação de Probabilidade Conjunta: No BAM, o peso de atenção $p_{ij}$ (a probabilidade do token $j$ atender ao token $i$) é definido como uma distribuição de probabilidade conjunta:
  $$p_{ij} = p(f_{cont}(q_i, k_j) | g_{pos}(i, j)) \times p(g_{pos}(i, j))$$
  Aqui, $f_{cont}$ representa a similaridade de conteúdo, e $g_{pos}$ representa a relação posicional.
• Codificação Posicional como Priors: A estrutura postula que as estratégias de codificação posicional correspondem a distribuições a priori específicas $p(g_{pos}(i, j))$ sobre as posições dos tokens.
  • NoPE é derivado como um prior Uniforme (restrito pela máscara causal).
  • ALiBi é derivado como um produto de um prior Uniforme e um prior Laplaciano (decaimento exponencial baseado na distância).
• O Escalar Normalizador ($Z$): A derivação introduz um escalar $Z$ que normaliza o produto das probabilidades de conteúdo e posição, representando a energia livre de interação ou o acoplamento de informação entre conteúdo e posição.

Solução Proposta: GGD-BAM

Aproveitando essa estrutura, os autores introduzem o GGD-BAM, que utiliza uma Distribuição Gaussiana Generalizada (GGD) como prior posicional. A matriz de viés $B$ adicionada aos scores de atenção é definida como:
$$b_{ij} = -\left| \frac{j - i - \mu}{\alpha} \right|^\beta$$
Onde $\mu$, $\alpha$ e $\beta$ são parâmetros aprendíveis (ou fixos para comportamentos específicos).

• $\beta \in (0, 1)$: Produz uma distribuição de cauda pesada que decai mais lentamente do que o prior Laplaciano (ALiBi), permitindo que a atenção se estenda além do ALiBi enquanto mantém a localidade.
• $\beta < 0$: Um relaxamento novel da definição padrão de GGD. Isso cria um prior "invertido" onde a massa de probabilidade é suprimida perto do token de consulta e concentrada na extremidade distante da janela de contexto. Teoricamente, isso permite que o modelo ignore o contexto local e foque em dependências arbitrariamente longas.
• Softmax Escalável (SSMax): Os autores notam que o BAM é suscetível ao "desvanecimento da atenção" (distribuição uniforme à medida que o contexto cresce). Eles combinam o BAM com SSMax, que reescala os logits dinamicamente, para preservar a nitidez da distribuição.

Contribuições Principais

1. Estrutura Teórica Unificada: O BAM fornece uma unificação probabilística dos métodos de PE existentes, mapeando explicitamente NoPE e ALiBi para priors Uniforme e Laplaciano, respectivamente.
2. Prior Gaussiano Generalizado: A proposta do GGD-BAM introduz uma família flexível de priors posicionais. Ao aprender o parâmetro de forma $\beta$, o modelo pode alternar adaptativamente entre modos de atenção local e recuperação de longo alcance.
3. Sobrecarga Mínima: O método introduz menos de 1.000 parâmetros treináveis adicionais (especificamente 384 para um modelo de 120M parâmetros) ao aprender $\theta_\alpha$ e $\theta_\beta$ para cada cabeça de atenção, enquanto fixa $\theta_\mu = 0$.
4. Interpretabilidade: A estrutura oferece uma maneira principiante de visualizar e interpretar padrões de atenção como distribuições de probabilidade explícitas sobre posições.

Resultados Empíricos

Os experimentos foram conduzidos em modelos variando de 120M a 1,1B parâmetros, treinados no FineWeb 10B com um comprimento de contexto de 512 tokens.

• Recuperação de Passkey: O GGD-BAM (especificamente com SSMax) alcançou precisão perfeita na recuperação de um token específico de sequências de até 32.000 tokens (64x o comprimento de treinamento) e manteve precisão >80% até 512.000 tokens (1000x o comprimento de treinamento). Em contraste, baselines como RoPE, Sinusoidal e NoPE degradaram-se para precisão quase aleatória além do seu horizonte de treinamento. O ALiBi mostrou melhor extrapolação de perplexidade, mas falhou significativamente em tarefas de recuperação em comprimentos extremos.
• Benchmark RULER (NIAH): O GGD-BAM superou todas as baselines (incluindo RoPE e ALiBi) no subconjunto "Needle In A Haystack" (Agulha no Palheiro), sendo o único método a manter alta precisidade além de 6.000 tokens na tarefa Single 1.
• Perplexidade: O GGD-BAM manteve perplexidade comparável ao ALiBi e RoPE em contextos longos, demonstrando que a recuperação aprimorada não ocorre à custa da fluência da modelagem de linguagem.
• Tarefas Downstream: Nos benchmarks MMLU e ARC, o BAM SSMax superou o RoPE SSMax.
• Visualização de Atenção: As visualizações confirmaram que as cabeças com $\beta \le 0$ suprimiram efetivamente o contexto local e focaram em tokens distantes (a passkey), enquanto as cabeças com $\beta \in (0, 1)$ mantiveram atenção de cauda longa.

Significado e Alegações

O artigo alega que o BAM serve a um duplo propósito:

1. Teórico: Fornece uma fundamentação probabilística rigorosa para a codificação posicional, esclarecendo por que certos métodos (como ALiBi) funcionam para perplexidade, mas falham para recuperação (devido à sua natureza local).
2. Prático: Oferece um método altamente eficiente para extrapolação do comprimento do contexto. Ao permitir que as cabeças de atenção aprendam valores negativos de $\beta$, o modelo pode suprimir explicitamente a localidade para recuperar informações de contextos arbitrariamente longos sem exigir exposição direta a tais comprimentos durante o treinamento.

Os autores enfatizam que, embora a perplexidade seja uma métrica padrão, ela é insuficiente para avaliar capacidades verdadeiras de contexto longo; tarefas baseadas em recuperação são necessárias para verificar se os modelos estão atendendo a informações distantes em vez de apenas depender de padrões locais. O trabalho sugere que a capacidade de modelar priors posicionais explicitamente é um viés indutivo crucial para Transformers de contexto longo da próxima geração.