Parallax:面向语言建模的参数化局部线性注意力
阅读原文· arxiv.orgParallax是一种可扩展至大语言模型的参数化局部线性注意力机制。它消除了局部线性注意力中的数值求解器,并引入额外的查询投影器来探测KV协方差。该研究提出一种硬件感知算法,其算术强度优于FlashAttention,将注意力转向更计算密集的模式。其原型解码内核在不同批次大小和上下文长度下匹配或超越FlashAttention 2/3。在0.6B和1.7B规模的预训练中,Parallax展现出持续的困惑度改进,且该收益可迁移至下游基准测试。研究还发现Muon优化器能有效释放Parallax的性能潜力。
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.