Decoupling the "What" and "Where" With Polar Coordinate Positional Embeddings

Resumo Técnico: Desacoplando o "O quê" e o "Onde" com o Positional Embedding em Coordenadas Polares

Problema
Nas arquiteturas Transformer, o mecanismo de atenção combina chaves (keys) com consultas (queries) com base em dois fatores: conteúdo ("o quê") e posição na sequência ("onde"). Os autores argumentam que o Rotary Position Embedding (RoPE), amplamente adotado em modelos de fronteira como o Llama 3 e o DeepSeek, entrelaça esses dois fatores. No RoPE, a rotação aplicada aos vetores de chave e consulta depende tanto das magnitudes de seu conteúdo quanto de suas posições relativas. Esse entrelaçamento cria um confundimento "o quê-onde", prejudicando o desempenho do modelo quando tarefas exigem correspondência independente apenas de conteúdo ou de posição. Especificamente, o termo de interação no escore de atenção do RoPE faz com que a fase da rotação seja influenciada pelas características de conteúdo, tornando difícil para o modelo aprender regras que dependam estritamente de indexação posicional ou correspondência de conteúdo.

Metodologia: Polar Coordinate Positional Embedding (PoPE)
Os autores propõem o PoPE, uma modificação do RoPE que desacopla conteúdo e posição ao reformular as representações de chave e consulta em coordenadas polares.

• Reformulação: No RoPE, os componentes são rotacionados com base na posição, mas a fase inicial é determinada pelo conteúdo. No PoPE, os autores transformam os vetores de chave e consulta de valores reais em vetores complexos onde:
  • A magnitude ($\mu$) é derivada exclusivamente das características de conteúdo usando uma função de ativação softplus ($\sigma(x) = \ln(1 + e^x)$) para garantir a não negatividade.
  • A fase ($\phi$) é determinada estritamente pela posição da sequência ($s$ ou $t$) e por uma frequência específica de componente ($\theta_c$).
• Escore de Atenção: O escore de atenção é computado como a parte real do produto escalar entre a transposta conjugada da consulta e a chave. Matematicamente, isso resulta em:
  $$a^{PoPE}_{ts} = \sum_{c=1}^{d} \mu_{\tilde{q}_{tc}} \mu_{\tilde{k}_{sc}} \cos((s - t)\theta_c + \delta_c)$$
  onde $\delta_c$ é um termo de viés (bias) aprendível.
• Principais Diferenças do RoPE:
  1. Desacoplamento: O termo de interação ($\phi_{k_{sc}} - \phi_{q_{tc}}$) que mistura conteúdo e posição no RoPE é eliminado. A magnitude (conteúdo) e a fase (posição) são estritamente separadas.
  2. Resolução de Frequência: O PoPE opera em $d$ componentes (elementos individuais) em vez de $d/2$ pares, efetivamente dobrando o número de canais de frequência.
  3. Viés Aprendível: Um viés $\delta_c$ fixo, porém aprendível, pode ser introduzido para ajustar o deslocamento relativo ideal para cada frequência sem reintroduzir o confundimento conteúdo-posição.

Principais Contribuições

1. Tarefa de Diagnóstico (Indexação Indireta): Os autores introduzem uma tarefa que exige que os modelos identifiquem um caractere alvo com base em um deslocamento relativo de um caractere de origem (ex: "encontre o caractere 3 posições à esquerda de 'z'"). Esta tarefa requer a manipulação independente de conteúdo e posição.
2. Insight Teórico: O artigo fornece uma derivação algébrica mostrando como o RoPE confunde conteúdo e posição, enquanto o PoPE permite que o mecanismo de atenção implemente regras como uma conjunção de correspondências independentes de "o quê" e "onde".
3. Implementação Eficiente: Uma implementação eficiente usando Triton e Flash Attention 2 é fornecida, que computa a parte real do produto escalar complexo sem materializar a matriz complexa completa, adicionando apenas um único overhead de multiplicação em comparação com a atenção padrão.

Resultados

• Indexação Indireta: Na tarefa de diagnóstico, o RoPE alcançou uma precisão média de apenas 11,16%, lutando para aprender a correspondência independente necessária. Em contraste, o PoPE alcançou 94,82% de precisão, demonstrando sua capacidade de desacoplar os fatores.
• Modelagem de Sequência (Música e Genômica):
  • Música (JSB e MAESTRO): O PoPE superou o RoPE em log-verossimilhança negativa (NLL) em ambos os conjuntos de dados, alcançando menor perplexidade.
  • Genômica (Genoma de Referência Humana): O PoPE mostrou uma queda significativa em NLL comparado ao baseline de RoPE.
• Modelagem de Linguagem (OpenWebText): Em diversas escalas de modelo (124M, 253M e 777M parâmetros), o PoPE consistentemente alcançou menor perplexidade do que o RoPE. A lacuna de desempenho permaneceu estável ou aumentou com o tamanho do modelo.
• Tarefas de Downstream: Avaliações zero-shot em seis tarefas de downstream (LAMBADA, BLiMP, CBT, HellaSwag, PIQA, ARC-E) mostraram que os modelos baseados em PoPE alcançaram maior precisão média em todos os tamanhos de modelo comparados aos modelos baseados em RoPE.
• Extrapolação de Comprimento:
  • Zero-Shot: O PoPE demonstrou fortes capacidades de extrapolação de comprimento zero-shot em sequências até 10x mais longas que o contexto de treinamento (1024 vs 10240 tokens) sem qualquer ajuste fino ou interpolação de frequência.
  • Comparação com YaRN: Enquanto o YaRN (um método projetado para extrapolação) mitiga a degradação do RoPE até seu limite de ajuste fino (4096 tokens), ele degrada significativamente além disso. O PoPE manteve o desempenho em sequências mais longas sem exigir as técnicas de interpolação ou o ajuste fino extensivo que o YaRN necessita.
  • Análise de Uso de Frequência: A análise das normas de características revelou que os modelos RoPE tendem a suprimir canais de alta frequência para evitar ruído, enquanto os modelos PoPE utilizam uma gama de frequências mais distribuída, incluindo altas frequências, através das camadas.

Significância e Alegações
O artigo afirma que o PoPE oferece uma solução fundamentada para o entrelaçamento "o quê-onde" inerente ao RoPE. Ao desacoplar conteúdo e posição, o PoPE melhora a eficiência dos dados, a precisão assintótica e, notavelmente, a generalização do comprimento do contexto. Os autores enfatizam que esses ganhos são alcançados com uma modificação menor à arquitetura existente, não exigindo mudanças arquitetônicas complexas.

O trabalho é apresentado como pesquisa fundamental sobre modelagem de sequências. Os autores observam que, embora o método melhore as capacidades em domínios como música, genômica e linguagem natural (potencialmente auxiliando aplicações em dobramento de proteínas ou descoberta de fármacos), ele carrega riscos de duplo uso comuns a modelos de fundação, como a geração de conteúdo nocivo ou deepfakes sofisticados. A principal contribuição é a demonstração de que separar informações de conteúdo e posicionais leva a modelos Transformer mais robustos e generalizáveis.