🤖 machine learning

Parallax: Parameterized Local Linear Attention for Language Modeling

本文介绍了 Parallax，这是一种可扩展且数值稳定的参数化局部线性注意力机制，它通过消除局部线性注意力的计算瓶颈、优化硬件效率并展示与 Muon 优化器的新颖协同效应，在语言建模中实现了帕累托改进。

原作者： Yifei Zuo, Dhruv Pai, Zhichen Zeng, Alec Dewulf, Shuming Hu, Zhaoran Wang

发布于 2026-05-29

📖 1 分钟阅读☕ 轻松阅读

原作者： Yifei Zuo, Dhruv Pai, Zhichen Zeng, Alec Dewulf, Shuming Hu, Zhaoran Wang

原始论文采用 CC BY 4.0 许可（http://creativecommons.org/licenses/by/4.0/）。 ✨ 这是对下方论文的AI生成解释。它不是由作者撰写或认可的。如需技术准确性，请参阅原始论文。阅读完整免责声明

想象一下，你正在尝试回忆刚刚听到的一段故事。在大语言模型（LLM）的世界里，负责这种记忆的大脑部分被称为注意力（Attention）。多年来，这种注意力工作的标准方式就像查看一个单词列表，并根据每个单词的突出程度给予一个“相关性分数”。如果一个单词非常独特，它就得高分；如果它很平淡，就得低分。这被称为Softmax 注意力。

然而，这种方法有一个缺陷：它就像在看地图时只看到平坦、平均的地形，而忽略了坡度和曲线。一个名为**局部线性注意力（Local Linear Attention, LLA）**的新想法试图通过观察数据的“坡度”而不仅仅是平坦的平均值来解决这个问题。这就像意识到要预测球会滚向哪里，你不仅需要知道它在哪里，还需要知道山坡有多陡。

问题出在哪里？这种“坡度”方法的数学计算对于巨型 AI 模型在现实生活中的应用来说过于沉重且不稳定。这就像试图在土路上驾驶一级方程式赛车；引擎过于强大，无法适应这种地形。

于是，Parallax登场了。

什么是 Parallax？

Parallax 是那种“坡度”方法的一个全新、精简的版本。作者们将原始想法中复杂、沉重的数学部分进行了替换，用一种巧妙的、可学习的捷径取代了困难的部分。

可以这样理解：

旧方法（Softmax）： 你问图书管理员：“哪本书最相关？”图书管理员查看标题，挑出那个听起来最不同的。
“坡度”方法（LLA）： 图书管理员意识到，相关性不仅仅关乎标题本身，还关乎你正在寻找的那个标题周围的标题是如何变化的。但计算这种变化需要为每个单词配备一台超级计算机。
Parallax： 图书管理员学会了一种特殊技巧。他们不再每次都进行繁重的数学计算，而是随身携带一个小型、智能的记事本（一个可学习的项目器），能够根据上下文瞬间猜出“坡度”。它快速、稳定，且出奇地准确。

“魔法”成分：优化器

该论文最令人惊讶的发现之一是，Parallax 与用于训练 AI 模型的标准“引擎”（称为AdamW）配合效果不佳。这就像给汽车装上了高性能赛车引擎，却使用了错误类型的燃料；汽车会发出嘶嘶声并熄火。

论文发现，Parallax 需要一种特定的、更新型的燃料，称为Muon。当你使用 Muon 时，Parallax 运行完美，释放其全部潜力。如果没有 Muon，Parallax 仅仅比旧方法略好一点点。而有了 Muon，它就变得显著更聪明。这是一个罕见的案例，其中“引擎”（优化器）和“汽车设计”（架构）必须完美匹配才能赢得比赛。

它有多快？

作者们不仅让它变得更聪明，还让它变得更快。他们专门为现代显卡构建了一个定制的“引擎”（计算机代码内核）。

他们声称，在“解码”阶段（即 AI 正在生成下一个单词时），他们的方法与当前的行业黄金标准（FlashAttention）一样快，甚至更快，尽管它执行了更复杂的数学运算。
他们通过高效地利用内存实现了这一点，复用数据流，这样计算机就不必不断停下来获取新信息。

结果

团队在两种规模的 AI 模型（06 亿和 17 亿参数）上测试了 Parallax。

更聪明： 在测试中，Parallax 模型始终比同等规模的标准模型犯更少的错误（更低的“困惑度”），并且回答问题更好。
更好的记忆： 在旨在测试模型能否从长列表中回忆特定事实的合成测试中，Parallax 在茫茫大海中找出正确针头的能力要强得多。
效率： 即使他们调整模型以使用完全相同的计算能力或完全相同的参数数量，Parallax 仍然胜出。

结论

这篇论文介绍了Parallax，一种 AI 关注信息的新方式。它将旧的“平坦”思维方式升级为“感知坡度”的方式，但简化了数学计算，使其能够在真实的计算机上运行。关键在于，当它与一种名为Muon的特定训练工具配对时效果最佳，这种强大的组合在速度和智能方面都超越了当前的最先进模型。

作者强调，这是首次证明特定架构（Parallax）与特定优化器（Muon）之间的这种强关联能够创造出“帕累托改进”——意味着模型变得更好，而无需变得更大或更慢。

技术摘要：Parallax——用于语言建模的参数化局部线性注意力

问题陈述

大型语言模型（LLM）依赖 Softmax 注意力作为其核心计算原语，该机制自 Transformer 引入以来在结构上一直未变。尽管线性注意力和状态空间模型（SSM）等高效替代方案提供了次二次方复杂度，但在上下文信息检索任务中，它们的表现始终不及 Softmax 注意力。这种性能差距揭示了其设计中存在根本性的权衡。

理论框架，特别是测试时回归框架，将注意力解释为一种上下文回归求解器。在此框架内，**局部线性注意力（LLA）**被提出作为优于 Softmax 注意力的替代方案。LLA 用局部线性估计器取代了 Softmax 注意力中使用的局部常数估计器（Nadaraya-Watson），理论上能产生更优的偏差 - 方差权衡及联想记忆容量。然而，由于三个主要挑战，LLA 尚未成功扩展到 LLM 预训练：

计算开销：精确的 LLA 需要为每个令牌通过并行共轭梯度（CG）求解器求解线性系统，引入了显著的 I/O 和计算成本。
数值不稳定性：CG 求解器对精度格式敏感，使其难以在低精度硬件环境中管理。
正则化权衡：在实际预训练设置中，平衡正则化参数 $\lambda$ 以确保数值稳定性而不使模型退化为 Softmax 注意力并非易事。

方法论：Parallax

作者引入了Parallax，这是一种参数化的局部线性注意力版本，旨在为 LLM 预训练提供可扩展、高效且数值稳定的方案。

核心机制

Parallax 将 LLA 重新表述为对 Softmax 注意力的加性修正。Parallax 不再为每个查询求解最优修正向量 $\rho^*_i$ （这需要求逆协方差矩阵），而是学习从层输入 $x_i$ 到探针向量 $\rho_i = W_R x_i$ 的直接映射。
输出计算如下：
$o^{PLX}_i = o^{SA}_i - \Sigma^{(i)}_{KV} \rho_i$
其中 $o^{SA}_i$ 是标准 Softmax 注意力输出， $\Sigma^{(i)}_{KV}$ 是 KV 协方差矩阵。

关键设计选择包括：

消除求解器：通过对 $\rho_i$ 进行参数化，消除了每个令牌共轭梯度求解的需求。
移除边界放大：作者将边界放大因子 $\eta_i$ 设为零。在精确 LLA 中， $\eta_i$ 用于补偿边界偏差，但在参数化设置中，它可能导致训练不稳定。移除它简化了前向传播，同时保留了协方差修正能力。
统一框架：Parallax 被定位在一个由带宽、探针构建和仿射结构连接的注意力机制家族中。它弥合了非参数方法（LLA）与参数方法（线性注意力、MesaNet）之间的差距。

硬件感知算法

为解决效率问题，作者开发了一种与 FlashAttention (FA) 兼容的流式算法。

算术强度：该算法通过大致加倍计算量并复用相同的 KV 流，将算术强度（每字节内存流量的浮点运算次数）相比 FlashAttention 提高了约一倍。这将操作转移到了更符合现代 GPU 架构的计算受限区域。
内核优化：作者在 CuTeDSL 中为 NVIDIA Hopper GPU 实现了自定义解码内核。它利用了 Softmax 和协方差两个分支共享相同 KV 流的事实，允许在相同的 Tensor Core 指令（WGMMA）内联合计算 $QK $和$ RK$ 乘积，从而最小化 HBM 流量。

主要贡献

架构：推导了 Parallax，通过消除每令牌求解器并引入可学习探针，将局部线性注意力的理论优势扩展到大规模预训练。
效率：开发了一种硬件感知的流式算法和自定义解码内核，在多样化的批量大小和上下文长度下，其性能匹配或优于 FlashAttention 2/3。
优化器 - 架构交互：识别了 Parallax 架构与Muon优化器之间的关键交互。论文表明，Parallax 的修正分支高度依赖优化器的几何特性；它在 Muon 下显示出显著增益，但在 AdamW 下表现与 Softmax 注意力相当。

实验结果

作者通过合成基准测试以及在 0.6B 和 1.7B 参数规模下的 LLM 预训练验证了 Parallax。

合成基准测试：在 MAD-Benchmark 上，Parallax 在面向召回的任务（上下文召回、模糊召回、噪声召回、选择性复制）中始终优于 Softmax 注意力、Mamba 和 Gated DeltaNet，实现了最高的整体准确率。
语言建模：
- 困惑度：使用 Muon 训练的 Parallax 模型在 LAMBADA 和 WikiText 上的困惑度低于 Transformer 基线。
- 下游任务：Parallax 在零样本问答和常识推理基准测试中表现出一致的提升。
- 控制实验：增益在参数匹配（增加基线中的查询头）和计算匹配（减少 Parallax 中的头维度）的控制下依然存在，证实改进源于机制本身而非额外的参数或计算量。
优化器依赖性：一个重要的发现是 Parallax 的优势取决于优化器。在 Muon 下，修正与输出的比率（COR）随层深增长，模型学会利用协方差修正。在 AdamW 下，模型学会抑制修正分支（门控值 $\approx 0.26$ ），导致性能与 Transformer 基线相当。
谱分析：对权重矩阵的分析显示，与 AdamW 相比，Muon 在投影权重（特别是 $W_R$ 和 $W_R K$ 电路）中保持了更高的稳定秩，使得探针能够有效对齐主要的协方差方向。

意义与主张

该论文声称提供了文献中关于注意力机制架构 - 优化器协同设计的首个实证演示。它强调，像 Parallax 这样的高级注意力机制的成功不仅仅是架构设计的函数，而是与优化景观深度交织的。

作者将 Parallax 定位为帕累托改进：只要使用正确的优化器（Muon），它就能在不牺牲效率的情况下提供优于 Softmax 注意力的性能。他们还指出，Parallax 允许从预训练的 Transformer 进行训练后适应（通过初始化 $W_R=0$ ），这是线性注意力家族所不具备的特性。

这项工作表明，未来关于注意力机制的研究必须将“优化器 - 架构交互”视为主要设计约束，特别是对于仿射结构，其中学习到的探针的幅度和对齐至关重要。