Parallax：面向语言建模的参数化局部线性注意力
机器学习
2026-05-29 v1 人工智能
计算与语言
摘要
大型语言模型 (LLM) 已成为人工智能中的核心范式，但注意力的核心计算原语结构基本保持不变。局部线性注意力 (Local Linear Attention, LLA) 是一种从非参数统计学推导的注意力机制，基于测试时回归框架。在与先前高效注意力变体相关的研究不同，LLA 将softmax注意力中的局部常数估计升级为局部线性估计，为关联记忆提供了可证明的偏差-方差收益 tradeoff。然而，由于计算和数值稳定性问题，LLA 尚未在 LLM 预训练中实现规模化。我们引入Parallax，这是一种可规模化用于 LLM 的参数化局部线性注意力。Parallax 消除了 LLA 中的数值求解器，学习一个额外的类似查询的投影器，以探测 KV 协方差。我们将 Parallax 放置在由带宽、探针构建和仿射结构连接的注意力机制家族中。我们提出一种硬件感知算法，提高了相对于 FlashAttention 的算术强度，将注意力机制置于更计算受限的 regime。我们的原型解码内核在各种 batch 大小和上下文长度下可匹配或超越 FlashAttention 2/3。在 0.6B 和 1.7B 规模上进行预训练，我们发现预训练期间一致的困惑度改进可传递到下游基准测试。优势在参数匹配和计算匹配控制下均持续存在，表明其具有 Pareto 改进性。我们进行了严谨的预训练消融实验，识别出一种新现象，即 Muon 能否解锁 Parallax 的容量。据我们了解，这是架构研究文献中首次展示注意力机制在架构-优化器 codesign 中具有强大效果的实证结果。
引用
@article{arxiv.2605.29157,
title = {Parallax: Parameterized Local Linear Attention for Language Modeling},
author = {Yifei Zuo and Dhruv Pai and Zhichen Zeng and Alec Dewulf and Shuming Hu and Zhaoran Wang},
journal= {arXiv preprint arXiv:2605.29157},
year = {2026}
}