Parallax: Atenção Linear Local Parametrizada para Modelagem de Linguagem

Resumo

Modelos de Linguagem de Grande Escala (LLMs) tornaram-se o paradigma central em inteligência artificial, mas a primitiva computacional central da atenção permaneceu estruturalmente inalterada. A Atenção Linear Local (LLA) é um mecanismo de atenção derivado da estatística não paramétrica no arcabouço de regressão em tempo de teste. Em contraste com pesquisas anteriores sobre variantes eficientes de atenção, a LLA atualiza a estimativa constante local da atenção softmax para uma estimativa linear local, proporcionando compensações viés-variância comprovadamente superiores para memória associativa. No entanto, a LLA não foi escalada no pré-treinamento de LLMs devido a problemas de estabilidade computacional e numérica. Apresentamos Parallax, uma Atenção Linear Local parametrizada que é escalável para LLMs. Parallax elimina o solucionador numérico da LLA e aprende um projetor adicional semelhante a uma query que sonda a covariância KV. Posicionamos Parallax dentro de uma família de mecanismos de atenção conectados pela largura de banda, pela construção da sonda e pela estrutura afim. Propomos um algoritmo ciente de hardware que aumenta a intensidade aritmética em relação ao FlashAttention, deslocando a atenção para um regime mais limitado por computação. Nosso kernel de decodificação protótipo iguala ou supera o FlashAttention 2/3 em diversos tamanhos de lote e comprimentos de contexto. Pré-treinamos Parallax nas escalas de 0,6B e 1,7B e encontramos melhorias consistentes de perplexidade ao longo do pré-treinamento, com ganhos que se transferem para benchmarks downstream. A vantagem persiste sob controles pareados tanto por parâmetros quanto por computação, demonstrando uma melhoria de Pareto. Realizamos ablações cuidadosas de pré-treinamento e identificamos um fenômeno novo pelo qual o Muon desbloqueia a capacidade do Parallax. Até onde sabemos, esta é a primeira demonstração empírica de um forte codesign entre arquitetura e otimizador para mecanismos de atenção na literatura de pesquisa de arquitetura.

English

Large Language Models (LLMs) have become the central paradigm in artificial intelligence, yet the core computational primitive of attention has remained structurally unchanged. Local Linear Attention (LLA) is an attention mechanism derived from nonparametric statistics in the test-time regression framework. In contrast to prior research on efficient attention variants, LLA upgrades the local constant estimate in softmax attention to a local linear estimate, yielding provably superior bias-variance tradeoffs for associative memory. However, LLA has not been scaled in LLM pretraining due to computational and numerical stability concerns. We introduce Parallax, a parameterized Local Linear Attention that is scalable for LLMs. Parallax eliminates the numerical solver in LLA and learns an extra query-like projector that probes the KV covariance. We place Parallax within a family of attention mechanisms connected by the bandwidth, the probe construction and the affine structure. We propose a hardware-aware algorithm that increases the arithmetic intensity over FlashAttention, shifting attention into a more compute bound regime. Our prototype decode kernel matches or outperforms FlashAttention 2/3 across diverse batch sizes and context lengths. We pretrain Parallax at 0.6B and 1.7B scales and find consistent perplexity improvements throughout pretraining with gains that transfer to downstream benchmarks. The advantage persists under both parameter-matched and compute-matched controls, demonstrating a Pareto improvement. We perform careful pretraining ablations and identify a novel phenomenon whereby Muon unlocks the capacity of Parallax. To our knowledge, this is the first empirical demonstration of strong architecture-optimizer codesign for attention mechanisms in the architecture research literature.