Resumo Técnico: Atenção Linear de Ordem Superior (HLA)

Declaração do Problema

O mecanismo de atenção dot-product escalada padrão nas arquiteturas Transformer incorre em complexidade computacional e de memória de $O(n^2)$ em relação ao comprimento da sequência $n$. Este custo quadrático é um gargalo primário para a escalabilidade de modelos de linguagem autoregressivos para contextos longos. Embora alternativas existentes, como Atenção Linear e Modelos de Espaço de Estados (SSMs), ofereçam complexidade por token de $O(n)$ ou $O(1)$, elas são tipicamente restritas a aproximações de primeira ordem ou métodos baseados em kernels. Essas restrições frequentemente limitam a expressividade do modelo, impedindo-o de capturar interações complexas de ordem superior entre tokens que os mecanismos de atenção padrão conseguem modelar.

Metodologia

O artigo introduz a Atenção Linear de Ordem Superior (HLA), um mecanismo causal e de streaming que realiza interações de ordem superior através de estatísticas suficientes de prefixo compactas. A ideia central é que operadores semelhantes à atenção de ordem superior podem ser fatorizados em momentos de baixa ordem (por exemplo, somas de produtos externos de chaves), permitindo streaming causal exato sem materializar matrizes de atenção $n \times n$.

1. HLA de Segunda Ordem

A variante de segunda ordem generaliza a atenção linear mantendo um estado de tamanho constante por cabeça que captura momentos de segunda ordem.

• Manutenção do Estado: Em vez de uma única soma acumulada, a HLA mantém resumos de prefixo incluindo o momento da chave $S^K_t = \sum_{i \le t} k_i k_i^\top$, o acumulador query-valor $C^{QV}_t = \sum_{i \le t} q_i v_i^\top$ e a massa da query $m^Q_t = \sum_{i \le t} q_i$.
• Causalidade Estrita: Para impor causalidade autoregressiva estrita (impedindo que tokens futuros influenciem saídas atuais) sem construir a matriz de atenção completa, os autores derivam resumos estendidos. Especificamente, eles introduzem resumos cruzados $G_t$ e $h_t$ que corrigem a inclusão de termos futuros no produto ingênuo $S^K_t C^{QV}_t$.
• Formulação da Saída: A saída não normalizada no tempo $t$ é computada como $o_t = q_t^\top (S^K_t C^{QV}_t - G_t)$. Esta formulação permite computação de $O(d^2 + d d_v)$ por token, independente do comprimento da sequência.
• Normalização: Uma variante normalizada por razão também é definida, dividindo o numerador por um denominador mascarado derivado de resumos estendidos similares.

2. Atenção Linear de Ordem Superior Assimétrica (AHLA)

O artigo propõe uma variante assimétrica complementar, a AHLA, que implementa o operador $Q(K^\top Q)(K^\top V)$. Esta estrutura pondera valores através de uma única passagem roteada por um índice de chave intermediário, diferindo da estrutura simétrica $AA^\top V$ da HLA padrão. Ela utiliza um conjunto distinto de resumos de streaming ($P^{KV}_t, m^K_t, E_t, n_t$) e alcança a mesma complexidade linear por token.

3. HLA de Terceira Ordem

O framework é estendido para interações de terceira ordem, definidas via produto matricial $AA^\top A$ onde $A = QK^\top$.

• Álgebra de Streaming: A saída mascarada de terceira ordem requer a manutenção de três momentos principais de prefixo ($S^K_t, S^Q_t, P^{KV}_t$) e seis resumos cruzados adicionais ($G^{(1)}_t, \dots, G^{(3)}_t$ e $h^{(1)}_t, \dots, h^{(3)}_t$) para impor causalidade estrita.
• Complexidade: O custo de atualização online permanece linear nas dimensões de características, embora o tamanho do estado aumente para acomodar os momentos de ordem superior.

4. Treinamento Paralelo via Varreduras Associativas

Para abordar a ineficiência do treinamento de modelos puramente recorrentes em GPUs, os autores desenvolvem um esquema de treinamento paralelo por blocos.

• Operadores Associativos: Eles definem operadores associativos (especificamente produtos semidiretos) para concatenar segmentos de tokens. Esses operadores accountam pelos termos cruzados entre segmentos (por exemplo, $S^K_B C^{QV}_A$) que surgem ao combinar blocos.
• Equivalência de Varredura: Ao aplicar uma varredura exclusiva de Blelloch sobre esses resumos de segmento, o método produz ativações forward que são exatamente idênticas às de uma recorrência serial. Isso permite treinamento paralelo sem recorrer a retropropagação aproximada através do tempo.
• Varreduras de Terceira Ordem: Para o caso de terceira ordem, o operador de varredura é aumentado com mapas lineares de nível de segmento ($M^{KQP}, M^{KQm}$) para lidar com a composição de estados corrigidos de terceira ordem de forma exata.

Contribuições Principais

1. Streaming Mascarado Exato em Segunda Ordem: O artigo fornece uma álgebra completa de resumos estendidos que produz HLA de segunda ordem estritamente causal com custo constante por token. Inclui provas formais estabelecendo identidades de streaming mascarado e atualizações online, distinguindo entre variantes não normalizadas e normalizadas por razão.
2. Varreduras Associativas Correspondendo a Ativações Seriais: Os autores definem operadores associativos para configurações mascaradas e não mascaradas (com e sem decaimento exponencial). Eles provam que varreduras exclusivas padrão sobre esses operadores produzem ativações forward idênticas a recorrências seriais, permitindo treinamento paralelo eficiente.
3. Extensão de Terceira Ordem: O trabalho apresenta o estado completo mascarado de terceira ordem, atualizações online, um kernel de streaming estritamente causal e um algoritmo exato paralelo por blocos. Isso inclui a derivação de mapas de segmento necessários para compor estados de terceira ordem de forma exata.
4. Variante Assimétrica (AHLA): Um novo mecanismo de atenção de ordem superior assimétrico é introduzido, oferecendo um viés indutivo diferente enquanto mantém as mesmas propriedades de streaming e complexidade.

Resultados e Afirmações

O artigo foca primariamente na estrutura algorítmica e propriedades teóricas em vez de benchmarks empíricos em conjuntos de dados específicos. As afirmações são fundamentadas em derivação matemática:

• Complexidade: A HLA alcança tempo e memória de $O(d^2 + d d_v)$ por token para inferência, onde $d$ é a dimensão da query/chave e $d_v$ é a dimensão do valor. Isso é independente do comprimento da sequência $n$.
• Exatidão: O esquema de treinamento paralelo por blocos é provado reproduzir as ativações de uma recorrência serial exatamente, evitando erros de aproximação comuns em outros métodos de treinamento de atenção linear.
• Causalidade: As formulações mascaradas propostas impõem estritamente a causalidade autoregressiva, garantindo que a saída no tempo $t$ dependa apenas de tokens $j \le t$.
• Escalabilidade: Ao evitar a materialização de matrizes $n \times n$ e utilizar varreduras associativas, a HLA é posicionada como um bloco de construção escalável para modelos de contexto longo.

Significado

O artigo posiciona a HLA como um "bloco de construção principiado e escalável" que preenche a lacuna entre mistura dependente de dados semelhante à atenção e a eficiência de arquiteturas recorrentes modernas.

• Expressividade vs. Eficiência: Diferentemente da atenção linear de primeira ordem que depende de mapas de características fixos, a HLA aprende uma métrica dependente de dados (via $S^K_t$) e captura interações de ordem superior (polinômios de grau 2 ou grau 3) enquanto retém o paradigma de atualização de estado $O(1)$ por token.
• Implementação: O trabalho enfatiza que a HLA pode servir como um substituto direto para a subcamada de atenção padrão em blocos Transformer, exigindo apenas uma troca de kernel enquanto preserva codificações posicionais e lógica de mascaramento.
• Fundação Teórica: Ao fornecer identidades de streaming exatas e operadores de varredura associativa para termos de ordem superior, o artigo estabelece um framework algébrico rigoroso para estender a atenção linear além da primeira ordem sem sacrificar restrições causais ou eficiência de treinamento.