视差：语言建模的参数化局部线性注意力

大语言模型（LLM）已成为人工智能的中心范例，但注意力的核心计算原语在结构上保持不变。局部线性注意力（LLA）是一种源自测试时回归框架中的非参数统计的注意力机制。与之前关于有效注意力变体的研究相比，LLA 将 softmax 注意力中的局部常数估计升级为局部线性估计，从而为联想记忆带来了可证明优越的偏差-方差权衡。然而，由于计算和数值稳定性问题，LLA 尚未在 LLM 预训练中进行缩放。我们引入了 Parallax，这是一种可针对 LLM 进行扩展的参数化局部线性注意力。 Parallax 消除了 LLA 中的数值求解器，并学习了一个额外的类似查询的投影仪来探测 KV 协方差。我们将视差置于由带宽、探针结构和仿射结构连接的一系列注意机制中。我们提出了一种硬件感知算法，可以增加 FlashAttention 的算术强度，将注意力转移到更受计算限制的区域。我们的原型解码内核在不同的批量大小和上下文长度上匹配或优于 FlashAttention 2/3。我们在 0.6B 和 1.7B 尺度上预训练视差，并在整个预训练过程中发现了一致的困惑度改进，并将收益转移到下游基准。在参数匹配和计算匹配控制下，优势仍然存在，证明了帕累托改进。我们进行了仔细的预训练消融，并发现了一种新现象，即 Muon 释放了视差的能力。据我们所知，这是架构研究文献中针对注意力机制的强大架构优化器协同设计的第一个实证演示。