Resumo Técnico: Atenção Interdomínio

Declaração do Problema

Transformers e Modelos de Espaço de Estado (SSMs) profundos representam dois paradigmas de design opostos para modelagem de sequências. Mecanismos de atenção padrão (Vaswani et al., 2017) utilizam um cache crescente de chave-valor (KV), roteando consultas por meio de correspondência baseada em conteúdo. Embora isso proporcione uma recuperação precisa, incorre em custos computacionais quadráticos ($O(N_q N)$) e crescimento linear de memória ($O(N)$) em relação ao comprimento da sequência. Por outro lado, SSMs profundos (por exemplo, S4, Mamba) comprimem todo o contexto em um estado recorrente de tamanho fixo, atualizando em $O(1)$ tempo por passo. No entanto, essa compressão ocorre à custa da perda de correspondência direta consulta-chave; o estado fixo não é diretamente endereçável por consultas específicas, frequentemente levando a um desempenho mais fraco em tarefas que exigem recuperação precisa ou atenção baseada em conteúdo.

Arquiteturas híbridas tentaram intercalar essas abordagens, e variantes de atenção linear buscaram reduzir diretamente os custos de atenção. Este artigo propõe uma rota diferente: integrar um SSM diretamente no mecanismo de atenção para manter a eficiência de um estado de tamanho fixo, enquanto permite que cada consulta atenda ao histórico comprimido.

Metodologia: Atenção Interdomínio

A contribuição central é a Atenção Interdomínio, um misturador de tokens que unifica atenção por kernel e modelos de espaço de estado. O método é derivado de uma visão de regressão por kernel de atenção e de uma projeção de base estilo HiPPO.

Derivação Teórica

1. Mapa de Características de Kernel: O kernel de atenção $K(q, k)$ é aproximado por um mapa de características finito $\xi(q)^\top \xi(k)$.
2. Projeção de Base: Em vez de armazenar todas as chaves e valores, o método projeta as características de chave $\xi(k)$ e os valores $v$ sobre um conjunto compartilhado de funções de base mantido por uma única recorrência SSM (especificamente estilo HiPPO). Isso resulta em um estado de tamanho fixo independente do comprimento da sequência.
3. Coeficientes Comprimidos: As integrais dessas projeções geram três componentes de estado:
   • $U$: Coeficientes de projeção de base para chaves (variáveis indutoras).
   • $\Gamma$: Projeções de valores.
   • $\eta$: Constantes de normalização.
4. Leitura Condicionada à Consulta: Em cada posição, a consulta $q$ é mapeada para o espaço de características $\xi(q)$. A saída é computada atendendo aos coeficientes comprimidos em vez de tokens individuais:
   $$\tilde{o} = \frac{\sum_m (\xi(q)^\top u_m) \gamma_m}{\sum_m (\xi(q)^\top u_m) \eta_m}$$
   Na forma matricial: $\tilde{O} = \frac{F_q U^\top \Gamma}{F_q U^\top \eta}$, onde $F_q$ é a matriz de características da consulta.

Implementação Escalável

A implementação prática é um relaxamento aprendido da derivação teórica, projetado para estabilidade e eficiência:

• Mapa de Características: Utiliza um mapa de características neural estilo DeltaNet $\xi(x) = \text{SiLU}(\text{SC}(x)) / \|\text{SiLU}(\text{SC}(x))\|^2$, onde SC é uma convolução profunda causal.
• Núcleo Recorrente: Aproxima a projeção HiPPO usando uma recorrência S4D com diagonal complexa. A entrada do SSM concatena a característica de chave $\xi(k)$ e o valor $v$.
• Normalização: Aplica RMSNorm de entrada com vieses por cabeça aprendíveis às metades de chave e valor antes de entrar no SSM.
• Leitura sem Denominador: Para evitar a mistura de escalas incompatíveis causadas pelo redimensionamento de entrada, o termo de normalização $\eta$ é descartado, resultando em uma leitura de atenção linear: $\tilde{O} = F_q U^\top \Gamma$.
• Arquitetura: Incorporada em um decodificador estilo Llama com pré-normalização, feedforwards SwiGLU e incorporações de posição rotacionais (RoPE).

Contribuições Principais

1. Construção: Um misturador de tokens de estado fixo condicionado à consulta derivado de regressão por kernel e projeção de base HiPPO, com uma distinção clara entre a derivação teórica ideal e a implementação escalável.
2. Decomposição do Mecanismo: Um estudo de ablação em 125M de parâmetros identifica a projeção condicionada à consulta como o fator dominante que impulsiona os ganhos de desempenho sobre misturadores S4D padrão. A entrada dupla de chave/valor contribui com um ganho menor e secundário.
3. Estudo de Escala: Uma avaliação abrangente no FineWeb-Edu de 125M a 1,3B de parâmetros sob um orçamento de estado recorrente correspondido.

Resultados Experimentais

O estudo compara a Atenção Interdomínio contra um misturador de tokens apenas S4D e uma linha de base de atenção softmax padrão.

• Desempenho vs. S4D: A Atenção Interdomínio melhora o misturador de tokens S4D em todas as escalas (125M–1,3B), alcançando 13–16% de menor perplexidade de validação.
• Desempenho vs. Softmax: Em 1,3B de parâmetros, a Atenção Interdomínio supera a linha de base softmax da mesma receita na perplexidade de validação (7,98 vs. 8,63) e em um conjunto de oito tarefas de senso comum (+3,03 pontos percentuais de precisão média). Em 125M, está aproximadamente empatada com o softmax.
• Extrapolação de Comprimento: Aproveitando seu núcleo de estado fixo, a Atenção Interdomínio exibe comportamento "plano em comprimento". Enquanto a perplexidade da linha de base softmax degrada significativamente além do contexto de treinamento (4K), a Interdomínio mantém desempenho estável até 14K (3,5× o contexto de treinamento).
• Limitações de Recuperação: O artigo observa que a Atenção Interdomínio, como outros modelos de estado fixo, luta com recuperação associativa de strings exatas (por exemplo, tarefas de agulha no palheiro) em comparação com a atenção softmax, identificando isso como uma limitação estrutural da compressão de estado fixo e não como uma falha específica do mecanismo Interdomínio.

Significado e Alegações

O artigo alega que a Atenção Interdomínio conecta com sucesso a lacuna entre a eficiência dos SSMs de estado fixo e as capacidades de recuperação baseadas em conteúdo da atenção. Ao projetar características sobre uma base SSM, permite que consultas atendam a um histórico comprimido sem o custo quadrático ou o crescimento linear de memória da atenção padrão.

O significado principal reside em demonstrar que uma projeção condicionada à consulta sobre um estado recorrente fixo pode superar tanto SSMs puros quanto, em escalas maiores, a atenção softmax padrão, enquanto herda as propriedades desejáveis de extrapolação de comprimento dos SSMs. Os autores sugerem que essa abordagem oferece um caminho viável para modelagem de linguagem de longo contexto, onde eficiência de memória e desempenho consistente em contextos variáveis são críticos. Direções de trabalho futuro mencionadas incluem a integração de núcleos de estado fixo mais fortes (por exemplo, Mamba-3) ou a combinação da leitura com regras de atualização de pesos rápidos.