西北大学、Tilde Research与华盛顿大学联手，让AI注意力机制实现了理论上的根本性升级

这项由美国西北大学、Tilde Research和华盛顿大学联合开展的研究，于2026年5月以预印本形式发布，论文编号为arXiv:2605.29157。研究团队在这篇题为《Parallax：面向语言建模的参数化局部线性注意力》的论文中，提出了一种名为Parallax的全新注意力机制，在保持和现有高效实现兼容的前提下，从数学根基上改进了大语言模型中最核心的计算模块。

大语言模型已经深深嵌入了我们的日常生活——无论是你问AI写一首诗、让它帮你调试代码，还是用它整理会议纪要，背后都是同一套核心技术在运转。在这套技术的心脏位置，有一个叫做"注意力机制"的东西，它负责让模型在处理每个词的时候，知道应该把"注意力"集中在上下文的哪些地方。自从2017年谷歌团队在那篇划时代的论文《Attention is All You Need》中提出以来，这个机制的基本形态几乎没有发生过结构性的变化。

研究团队认为，这个核心机制存在一个可以被数学严格证明的缺陷，并且提出了一套在理论和实践两个层面都更优的替代方案。

一、注意力机制究竟在做什么，又错在哪里

要理解Parallax的意义，先要明白现有的注意力机制在做一件什么事情。可以把整个过程想象成一位图书馆员在帮你查找信息。你提出一个问题（这就是"查询向量"Q），图书馆的书架上存放着各种条目（这就是"键值对"KV），图书馆员会根据你的问题和书架条目的相关程度，决定从哪些书里摘取内容（这就是"值向量"V），最终把信息汇总给你。

现有的Softmax注意力机制在汇总信息时，本质上是在做一个"局部常数估计"——用统计学的行话说，就是以当前查询点为中心，用周围数据的加权平均值来预测答案。这种方式有一个内在的局限：当数据点不是均匀分布，而是扎堆在某一侧时（比如书架上大部分参考书都集中在左边，而你站在右边），这种简单的加权平均就会产生系统性偏差，统计学上称之为"边界偏差"。

局部线性注意力（LLA）是研究团队在此前工作中提出的一个理论概念，它把这个"局部常数估计"升级成了"局部线性估计"。继续用图书馆员打比方，普通图书馆员只会说"这一片书架的书大概都讲这个话题"，而升级版图书馆员会说"你站在书架的边缘，朝着计算机科学区走过去，越走内容越专业，我能根据这个方向性规律给你更精准的回答"——它考虑了数据的空间分布梯度，而不只是取平均值。

理论上可以严格证明，在满足一定条件下，这种局部线性估计的综合误差比局部常数估计（Softmax注意力）要小，而Softmax注意力本身又比全局线性估计（线性注意力）的误差要小。这构成了一个清晰的误差层级：全局线性估计的误差远大于局部常数估计（Softmax注意力），后者的误差又远大于局部线性估计（LLA）。

然而，LLA一直停留在理论层面，从未被真正用于大规模语言模型的预训练。原因在于它每处理一个词，都需要实时求解一个线性方程组，这不仅计算量大，还对数值精度非常敏感——在现代GPU惯用的低精度浮点数格式下很容易失稳，就像用一把误差较大的尺子去精确丈量毫米级别的距离。

二、Parallax如何把理论变成可以实际运行的系统

研究团队的核心贡献，正是把LLA那个"需要实时求解方程"的步骤彻底拆掉，换成一套可以从训练数据中学习的参数化替代方案。

在数学层面，LLA的输出可以被分解为两部分的叠加：第一部分是普通的Softmax注意力输出，第二部分是一个修正项，这个修正项依赖于键值对的"协方差结构"（可以理解为书架上各条目之间的关联模式）乘以一个"探针向量"ρ。在LLA中，这个探针向量ρ是通过求解线性方程组实时算出来的，计算昂贵且数值不稳定。

Parallax的做法是：直接让模型学习一个额外的投影矩阵W_R，在处理每个词的时候，用输入直接乘以这个矩阵得到探针向量ρ = W_R × x。这样，所有的"探针"信息都被编码进了可训练的参数里，完全消除了实时求解方程的需求。与此同时，研究团队还去掉了原始LLA中一个叫做"边界放大因子"的量（η），因为一旦探针向量变成参数化的，这个因子的几何意义就不复存在，如果保留反而会导致训练时数值溢出。

从使用者的角度看，Parallax就是一个比标准注意力多接受一个输入矩阵R的注意力层。除此之外，无论是架构上的位置、与其他组件的兼容性，还是对位置编码（RoPE）的支持，它都和标准Softmax注意力完全一致。研究团队特别指出，当W_R被初始化为全零矩阵时，Parallax在行为上和标准Softmax注意力完全等价——这意味着理论上可以把一个已经训练好的Transformer模型直接"插入"W_R权重并进行微调，而不需要从头开始训练。这是线性注意力家族根本做不到的事情，因为没有任何参数设置能让线性注意力精确还原成Softmax注意力。

研究团队还在一个统一的框架下，清晰地展示了Parallax与其他注意力机制之间的关系。整个注意力机制家族可以按两个维度来理解：权重是用Softmax计算还是全局均匀分布，以及探针向量ρ是零、参数化还是精确求解。沿着"把带宽h趋于无穷大"这个方向走，Softmax注意力会退化成"值平均"，Parallax会退化成一种叫做"仿射线性注意力"的机制，LLA会退化成"仿射MesaNet"。沿着"正则化λ趋于无穷大"这个方向走，Parallax和LLA都会退化回Softmax注意力，而线性注意力和MesaNet则会完全消失。这张关系图揭示了一个之前文献中从未系统梳理过的大家族结构。

三、让两倍计算量不带来额外延迟的硬件设计

Parallax相当于给每一次注意力计算都增加了一个"协方差修正"分支，从FLOPs（浮点运算次数）的角度来看，计算量大约是标准Softmax注意力的两倍。研究团队面临的工程挑战是：如何让这额外的计算量不转化成成比例的时间开销。

答案来自于"算术强度"这个概念的分析。算术强度是浮点运算量与内存访问量之比，衡量的是一个操作是被计算瓶颈还是内存读写瓶颈所限制。研究团队的分析表明，Parallax通过复用同一套键值对（KV）数据流，在不增加额外内存读写的情况下把计算量翻了一倍，从而将算术强度大幅提升——这让注意力操作从"内存受限"更多地迁移到"计算受限"状态，而现代GPU在计算受限的任务上有更大的发挥空间。

在具体的硬件实现上，研究团队利用了英伟达Hopper架构GPU的一个特性：它的张量核心指令（WGMMA）在执行矩阵乘法时，每次至少处理64行数据，而解码阶段只有一行查询向量，导致大量的硬件资源被白白浪费。Parallax的两个分支（QK和RK的矩阵乘法，以及两个PV的矩阵乘法）可以共享同一批张量核心指令，把那些原本空闲的行填满，相当于"顺手"把额外的计算做完了，不需要额外的时钟周期。

研究团队在H200 GPU上，针对从1到2048的批次大小、从128到32768的上下文长度，用BF16精度进行了全面的性能测试，覆盖了数百种不同的配置，并与FlashAttention 2和FlashAttention 3两个高度优化的工业级实现进行了对比。在"计算量匹配"设置下（Parallax使用64维，FlashAttention使用128维，使两者FLOPs相等），Parallax原型核函数在所有测试配置上均能匹配甚至超越FlashAttention 2和3的速度；在"内存访问量匹配"设置下（两者都使用128维，Parallax有更多FLOPs），Parallax同样展现出竞争力。

四、在合成测试和真实预训练中的表现

研究团队首先在一个叫做MAD基准测试的合成任务集上验证了Parallax的能力。这个测试集包含六种任务，可以粗略分为两类：一类是考察模型在上下文中精确检索信息的能力，包括上下文召回（ICR）、模糊上下文召回（FCR）、噪声上下文召回（NCR）和选择性复制（SC）；另一类是考察模型压缩和记忆训练数据的能力（CMP和MEM）。

在这六项任务中，Parallax在召回类任务上明显优于标准注意力、Mamba、Gated DeltaNet和MesaNet，同时在压缩和记忆类任务上保持了竞争力，综合平均准确率最高。研究团队还专门设计了更难版本的测试，把词汇量和序列长度都大幅提升，结果发现随着难度增加，其他方法的准确率急剧下降，而Parallax的表现基本保持稳定，尤其在最长上下文的选择性复制任务上优势最为突出。

在大规模语言模型预训练实验中，研究团队采用了Qwen-3架构作为基础，在一个叫做Ultra-FineWeb的大规模文本数据集上进行预训练，上下文长度为4096个词元。他们训练了0.6B（6亿参数）和1.7B（17亿参数）两个规模的模型，并与标准Transformer基线进行了详细对比。

为了确保比较的公平性，研究团队还特别设计了两个对照实验。一个是"参数量匹配的Transformer"，它通过增加查询头的数量，使参数总量与Parallax相同，以排除"Parallax胜出只是因为参数更多"的可能性。另一个是"计算量匹配的Parallax"，它把头维度减半以使注意力层的FLOPs与标准注意力相同，减少的参数量由扩大前馈网络来补偿，以排除"Parallax胜出只是因为用了更多计算"的可能性。

结果显示，在使用Muon优化器的情况下，完整的Parallax（0.6B规模）在LAMBADA困惑度测试上达到了18.56，WikiText困惑度达到了22.25，均低于标准Transformer（分别为22.15和23.43）以及参数量匹配的Transformer（22.35和23.36），而且在HellaSwag、PIQA、ARC等一系列下游推理和常识问答基准测试上的平均准确率也是所有方案中最高的（55.99%对比标准Transformer的54.54%）。计算量匹配的Parallax也明显优于标准Transformer和参数匹配的Transformer，排除了额外计算量是主要贡献来源的可能性。这两项控制实验共同证明，性能提升来自于机制本身，而不是更多的参数或更多的计算。

这一优势在1.7B规模上同样成立，LAMBADA困惑度从13.07下降到10.80，WikiText困惑度从18.11下降到17.08，下游任务平均准确率从61.43%提升到62.45%。

五、一个意外发现：优化器和架构之间存在深度耦合

这项研究中最出乎意料的发现，是Parallax的性能提升与所使用的优化器高度相关。研究团队发现，在使用AdamW（目前最常用的语言模型训练优化器）时，Parallax和标准Transformer的表现几乎没有差别。但换成Muon优化器后，差距就变得非常显著。

为了理解这背后的原因，研究团队设计了一套量化诊断工具。核心指标叫做"修正-输出比"（COR），衡量的是Parallax中那个协方差修正项的强度相对于基础Softmax注意力输出的比值——这个比值越高，说明Parallax新增的修正机制在实际发挥的作用越大。

测量结果揭示了一个清晰的模式：在Muon训练下，COR随着网络深度增加而持续攀升，在最深层达到8以上；在AdamW训练下，COR几乎维持在4以下。换句话说，用AdamW训练时，Parallax的修正分支虽然存在，但基本处于"摆设"状态，对最终输出贡献甚微。

研究团队进一步把COR的差距分解成两个来源。一个是"方向性"因素：用来调制协方差修正的探针向量ρ，是否指向了KV协方差矩阵最重要的方向。这用"协方差-探针对齐度"（CPA）来衡量。另一个是"幅度"因素：探针向量本身的大小。结果发现，Muon训练下不仅探针向量的范数更大，而且对齐度也更高——这说明差距不仅仅是一个缩放系数的问题，而是Muon帮助模型学到了在结构上更有意义的键值关联，并且让探针向量与这些关联方向更好地对齐。

研究团队还做了一个"门控"实验来进一步验证：他们在探针向量前面加了一个可学习的Sigmoid门控，允许模型自主决定是否激活修正分支。在Muon训练下，门控值随训练进行逐渐升高，最终稳定在接近完全开放的状态，修正分支也确实带来了与无门控版本相当的最终性能提升。在AdamW训练下，门控值则持续下降，最终稳定在约0.26，说明模型"主动学会了"压制这个修正分支，最终性能与普通Transformer相当。

从权重矩阵的谱结构来看，Muon训练下所有投影矩阵的"稳定秩"（一种衡量矩阵有效维度的指标）都远高于AdamW，而W_R在这一点上尤为突出——AdamW下的稳定秩只有9.3，Muon下高达134.0，是差距最大的一个矩阵。这解释了为什么AdamW下的探针向量对齐度更低：如果W_R的有效秩很低，那么它能产生的探针方向也极为有限，自然难以与高维的KV协方差矩阵的主方向对齐。

此外，研究团队还发现了一个纯粹由架构本身带来的效果：无论使用哪种优化器，Parallax模型中的W_V、W_O和OV电路的稳定秩都一致高于标准Transformer，说明Parallax的修正结构为值-输出路径提供了更丰富的信息。

研究团队坦诚地指出，目前还没有数学理论能完整解释为什么Muon对Parallax有如此特殊的作用。从已有的理解来看，Muon的更新总是正交化的（条件数精确为1），这让权重矩阵保持高秩、避免谱坍缩；而AdamW在这方面的保障更弱，容易出现有效秩快速萎缩的现象。但精确的机制仍然是一个开放的研究问题。

六、Parallax的得分模式与注意力分布

研究团队还发现，Parallax产生的权重分布与标准Softmax注意力在本质上存在结构性差异。在标准Softmax注意力中，每个位置对上下文各位置的权重都是正数且加和为1（就像100%的预算分配给不同词元）。在Parallax中，修正分支的加入使得单个位置的权重可以是负数，绝对值也可以远大于1。这种"可以是负数"的权重赋予了模型一种新能力：它可以主动"减去"某些词元的影响，而不只是降低它们的权重。

研究发现，Muon训练下的Parallax权重范围随网络深度增加而扩大，在最深层中最小和最大权重之间的跨度可以达到约±40，远超AdamW下的Parallax（跨度接近0）。这与COR随深度增大的模式完全一致。

另一个值得注意的现象是注意力"沉没"（attention sink）问题的缓解。标准Softmax注意力有一个众所周知的怪癖：它倾向于把大量的概率质量集中分配给序列的第一个词元（通常是特殊的起始符号），无论这个词元是否与当前预测相关，形成一个"权重沉没"。Parallax大幅减少了这种现象——无论是看基础Softmax分量还是看综合权重，第一个词元的权重都明显低于标准Transformer，暗示Parallax的修正分支可能承担了一部分原本被外包给第一个词元的"路由"功能。

与此同时，Parallax基础Softmax分量的信息熵（衡量分布均匀程度的指标）也一致高于标准Transformer，说明注意力分布更加分散、更加"平等"地覆盖上下文——Parallax用Softmax做更广泛的上下文聚合，把细粒度的词元分辨任务交给修正分支来完成，形成了一种分工结构。

七、局限性与未来方向

研究团队在论文中坦诚地列出了多个尚未解决的问题，并指出了值得探索的方向。

目前的预训练实验最大到1.7B参数量，在更大规模（比如70B以上）、更长上下文（比如128k以上）以及结合专家混合（MoE）架构等更复杂设置下，Parallax的表现是否依然稳定，还有待验证。由于Parallax将算术强度翻倍，这也为调整头维度、头数量和注意力与前馈网络的比例提供了新的灵活性空间，最优的配比还需要系统性的实验来确定。

在效率优化方面，Parallax继承了Softmax注意力的流式计算结构，理论上所有适用于Softmax注意力的稀疏化技术（滑窗、空洞、块稀疏等）都可以直接移植到Parallax。它与多头潜在注意力（MLA）等压缩KV缓存的技术也具有结构上的兼容性。目前的解码核心还只是一个原型，完整的训练核心（包括反向传播）尚未以同等精细的程度优化。

在微调层面，把已有的Transformer检查点转化为Parallax模型是一个值得探索的方向。理论上只需加入W_R并进行微调即可，因为W_R=0时两者行为完全等价，但具体效果如何、不同优化器设置下的微调效果是否有差异，都是有趣的未决问题。

在理论层面，优化器与架构之间的耦合机制为什么存在、是否也在其他类似的"仿射"结构（比如仿射线性注意力、仿射DeltaNet）中出现，都是值得深入研究的问题。研究团队在论文中建立的统一框架还没有纳入DeltaNet，推导DeltaNet的非参数化对应形式将是对理论体系的自然延伸。

说到底，Parallax这项工作最有意思的地方，不在于它把某个基准测试的数字又往上推了一点，而在于它打通了两件原本看起来互不相关的事情：一是来自非参数统计学的严谨理论，告诉我们现有注意力机制存在可以被数学证明的内在局限；二是来自深度学习优化领域的实践经验，关于不同优化器如何塑造模型内部的权重结构。把这两件事联系起来的，是对注意力机制"究竟在做什么"这个问题的一种新的理解方式——它不是一个固定的计算公式，而是一个在测试时执行的回归估计问题，不同的设计选择对应着不同的统计假设，而这些假设的优劣可以被精确地分析和比较。

有兴趣深入了解技术细节的读者，可以通过arXiv编号2605.29157查询完整论文，代码也已经在GitHub上以yifei-zuo/Parallax的路径公开发布。

Q&A

Q1：Parallax注意力机制和标准Softmax注意力相比，多了哪些计算开销？

A：Parallax在标准注意力的基础上增加了一个协方差修正分支，需要额外的矩阵投影W_R和一次键值协方差加权计算，FLOPs大约翻倍。但由于两个分支共享同一套KV数据流，内存访问量几乎不增加，算术强度大幅提升。研究团队在H200 GPU上的测试显示，原型解码核函数在计算量匹配和内存访问量匹配两种设置下均能匹配甚至超越FlashAttention 2和3的速度。

Q2：为什么Parallax配合Muon优化器才有明显效果，换成AdamW就几乎没有提升？

A：根本原因在于Parallax的修正分支需要探针向量ρ具有足够大的范数并且与KV协方差矩阵的主方向对齐。Muon的梯度更新始终是正交化的，能让权重矩阵维持高稳定秩，W_R的稳定秩在Muon下高达134，AdamW下只有9.3。低秩的W_R只能产生极有限的探针方向，无法有效对齐高维协方差结构，导致修正分支实际上处于近乎关闭的状态，模型在有门控实验中甚至主动学会了压制这个分支。

Q3：局部线性注意力LLA和Parallax有什么区别，为什么LLA没法直接用于大模型训练？

A：LLA是严格按照局部线性回归的最优解来计算探针向量的，每处理一个词元都需要实时求解一个线性方程组，迭代次数最多可达头维度次，内存访问量是Softmax注意力的T倍（T为迭代次数）。此外，正则化参数λ难以平衡——太大会退化回Softmax注意力，太小又容易数值不稳定，在低精度浮点数格式下尤为严重。Parallax把这个实时求解过程替换成可训练的投影矩阵W_R，消除了所有这些问题，同时保留了局部线性估计的核心思想。