Parallax: Atención Lineal Local Parametrizada para Modelado de Lenguaje

Resumen

Los Modelos de Lenguaje de Gran Escala (LLMs) se han consolidado como el paradigma central en inteligencia artificial, aunque la primitiva computacional central de atención ha permanecido estructuralmente sin cambios. La Atención Lineal Local (LLA) es un mecanismo de atención derivado de la estadística no paramétrica en el marco de regresión en tiempo de prueba. A diferencia de investigaciones previas sobre variantes eficientes de atención, LLA mejora la estimación constante local de la atención softmax a una estimación lineal local, ofreciendo compensaciones entre sesgo y varianza demostrablemente superiores para la memoria asociativa. Sin embargo, LLA no ha sido escalada en el preentrenamiento de LLMs debido a preocupaciones de estabilidad numérica y computacional. Presentamos Parallax, una Atención Lineal Local parametrizada que es escalable para LLMs. Parallax elimina el solucionador numérico en LLA y aprende un proyector adicional similar a una consulta que sondea la covarianza KV. Ubicamos Parallax dentro de una familia de mecanismos de atención conectados por el ancho de banda, la construcción de la sonda y la estructura afín. Proponemos un algoritmo consciente del hardware que incrementa la intensidad aritmética con respecto a FlashAttention, desplazando la atención hacia un régimen más limitado por cómputo. Nuestro kernel de decodificación prototipo iguala o supera a FlashAttention 2/3 en diversos tamaños de lote y longitudes de contexto. Preentrenamos Parallax a escalas de 0.6B y 1.7B y encontramos mejoras consistentes en perplejidad a lo largo del preentrenamiento, con ganancias que se transfieren a benchmarks posteriores. La ventaja persiste tanto bajo controles emparejados por parámetros como por cómputo, demostrando una mejora de Pareto. Realizamos ablaciones cuidadosas de preentrenamiento e identificamos un fenómeno novedoso mediante el cual Muon desbloquea la capacidad de Parallax. Hasta donde sabemos, esta es la primera demostración empírica de un codiseño arquitectura-optimizador robusto para mecanismos de atención en la literatura de investigación arquitectónica.

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.