Atenção Linear de Ordem Superior

Resumo

O custo quadrático da atenção por produto escalar dimensionado é um obstáculo central para a escalabilidade de modelos de linguagem autorregressivos em contextos longos. Atenção em tempo linear e Modelos de Espaço de Estados (SSMs) oferecem alternativas escaláveis, mas são tipicamente restritos a aproximações de primeira ordem ou baseadas em kernels, o que pode limitar a expressividade. Apresentamos a Atenção Linear de Ordem Superior (HLA), um mecanismo causal e contínuo que realiza interações de ordem superior por meio de estatísticas suficientes compactas de prefixo. No caso de segunda ordem, a HLA mantém um estado de tamanho constante e calcula saídas por token em tempo linear sem materializar quaisquer matrizes n por n. Fornecemos identidades de streaming em forma fechada, uma variante mascarada estritamente causal usando dois resumos adicionais e um esquema de treinamento paralelo em blocos baseado em varreduras associativas que reproduz exatamente as ativações de uma recorrência serial. Esboçamos ainda extensões para terceira ordem e ordens superiores. Coletivamente, esses resultados posicionam a HLA como um bloco de construção escalável e fundamentado que combina a mistura dependente de dados, similar à atenção, com a eficiência das arquiteturas recorrentes modernas. Página do Projeto: https://github.com/yifanzhang-pro/HLA.

English

The quadratic cost of scaled dot-product attention is a central obstacle to scaling autoregressive language models to long contexts. Linear-time attention and State Space Models (SSMs) provide scalable alternatives but are typically restricted to first-order or kernel-based approximations, which can limit expressivity. We introduce Higher-order Linear Attention (HLA), a causal, streaming mechanism that realizes higher interactions via compact prefix sufficient statistics. In the second-order case, HLA maintains a constant-size state and computes per-token outputs in linear time without materializing any n times n matrices. We give closed-form streaming identities, a strictly causal masked variant using two additional summaries, and a chunk-parallel training scheme based on associative scans that reproduces the activations of a serial recurrence exactly. We further outline extensions to third and higher orders. Collectively, these results position HLA as a principled, scalable building block that combines attention-like, data-dependent mixing with the efficiency of modern recurrent architectures. Project Page: https://github.com/yifanzhang-pro/HLA.