提出可扩展的局部线性注意力机制,显著提升语言模型性能。
Parallax: Parameterized Local Linear Attention for Language Modeling

- 用可学习投影替代数值求解,实现高效且稳定的线性注意力
- 在0.6B和1.7B规模预训练中持续降低困惑度,下游任务表现更优
- 首次实证架构与优化器协同设计对注意力机制的有效性
大型语言模型(LLM)已成为人工智能的核心范式,但注意力机制的结构长期未变。局部线性注意力(LLA)源自测试时回归框架的非参数统计方法,将softmax注意力中的局部常数估计升级为局部线性估计,理论上具有更优的偏差-方差权衡。然而,由于计算与数值稳定性问题,LLA尚未在大规模语言模型预训练中应用。本文提出Parallax,一种可参数化的局部线性注意力机制,通过消除LLA中的数值求解器,并引入一个类似查询的投影器来探测键值协方差,实现了对LLM的可扩展性。我们建立了一个以带宽、探针构造和仿射结构为连接的注意力家族。提出一种硬件感知算法,提升算术强度,使注意力从内存密集转向计算密集。原型解码内核在多种批大小和上下文长度下表现优于或匹配FlashAttention 2/3。我们在0.6B和1.7B规模进行预训练,发现整个训练过程中的困惑度持续下降,且优势在下游基准上迁移有效。该优势在参数匹配和计算匹配条件下均存在,证明了帕累托改进。通过细致的预训练消融实验,我们发现一种新现象:Muon能解锁Parallax的潜力。据我们所知,这是架构研究文献中首个关于注意力机制的强架构-优化器协同设计的实证演示。
原文摘要 · Abstract (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.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。