Parallax: Parameterized Local Linear Attention for Language Modeling

大语言模型（LLMs）已成为人工智能领域的核心范式，然而其底层计算原语——注意力机制——在结构上却始终未发生根本性变革。局部线性注意力（Local Linear Attention, LLA）是一种源于测试时回归框架下非参数统计理论的注意力机制。与以往针对高效注意力变体的研究不同，LLA 将传统 Softmax 注意力中采用的局部常数估计，升级为局部线性估计，从而在关联记忆任务中实现了可被严格证明更优的偏差-方差权衡。然而，由于计算开销与数值稳定性方面的挑战，LLA 尚未在大语言模型的预训练中实现规模化应用。为此，我们提出 Parallax——一种可扩展用于大语言模型的、带参数化的局部线性注意力机制。Parallax 摒弃了 LLA 中依赖的数值求解器，并额外引入一个类查询（query-like）的可学习投影器，用以探测键值（KV）对的协方差结构。我们将 Parallax 置于一个由带宽（bandwidth）、探针（probe）构造方式以及仿射结构（affine structure）所定义的注意力机制家族之中加以系统定位。我们进一步设计了一种面向硬件特性的算法，显著提升了算术强度（arithmetic intensity），使其超越 FlashAttention，从而将注意力计算从访存受限（memory-bound）转向更偏计算受限（compute-bound）的运行状态。我们在原型解码核（decode kernel）上的实验表明：Parallax 在各类批量大小（batch size）与上下文长度（context length）设置下，性能均达到或优于 FlashAttention 2/3。我们在 0.6B 和 1.7B 参数量级上完成了 Parallax 的完整预训练，结果表明其困惑度（perplexity）在整个训练过程中持续优于基线模型，且该优势能稳定迁移到下游各项评测基准中。值得注意的是，这一优势在参数量匹配（parameter-matched）与计算量匹配（compute-matched）两种严格控制条件下均保持显著，体现出一种帕累托意义上的实质性提升（Pareto improvement）。我们还开展了细致的预训练消融实验，并首次发现一种新现象：优化器 Muon 能够充分释放 Parallax 的建模潜力。据我们所知，这是架构研究文献中首次通过实证方式，清晰揭示注意力机制领域中“架构–优化器协同设计”（architecture-optimizer co-design）所能带来的强大增益。

传统注意力机制（如softmax attention）在大语言模型中计算开销大、难以扩展，而现有高效注意力变体多聚焦于近似或稀疏化，未从统计建模本质改进其偏差-方差权衡；LLA虽理论上更优（局部线性估计优于局部常数估计），但因需数值求解、稳定性差、不可微且难以扩展，尚未应用于LLM预训练。本文旨在验证：能否设计一种可扩展、稳定、硬件友好的参数化局部线性注意力，在真实LLM预训练中实现一致且Pareto最优的性能提升？

提出Parallax——首个为LLM预训练量身定制的参数化局部线性注意力：1）摒弃LLA中不稳定的数值求解器，用可学习的query-like协方差探针（KV covariance projector）替代；2）将局部线性回归显式嵌入注意力计算流，通过带宽、探针构造与仿射结构三要素统一建模注意力族；3）设计硬件感知高算术强度算法，使注意力从内存受限转向计算受限，显著提升GPU利用率。

• 首个在0.6B/1.7B规模LLM上完成端到端预训练验证的局部线性注意力，全阶段困惑度下降且下游任务泛化提升；• 严格控制实验：在参数量匹配（parameter-matched）和计算量匹配（compute-matched）下均显著优于基线，证实Pareto改进；• 发现关键协同现象：Muon优化器（而非AdamW）是释放Parallax容量的必要条件，首次实证揭示注意力架构与优化器的强耦合设计空间；• 开源高性能decode kernel（原型），在多种batch size与context length下持平或超越FlashAttention 2/3；• 算法完全兼容标准Transformer接口，无需修改训练框架。

Local Linear Attention (LLA, ICML 2024); FlashAttention / FlashAttention-2 (NeurIPS 2022/2023); Linformer (ICLR 2021); Performer (NeurIPS 2020); Nyströmformer (ICLR 2021); RingAttention (NeurIPS 2023); Mamba (ICML 2024); RetNet (ACL 2023)