Parallax: Parameterized Local Linear Attention for Language Modeling
Yifei Zuo, Dhruv Pai, Zhichen Zeng, Alec Dewulf, Shuming Hu, Zhaoran Wang
cs.LG, cs.AI, cs.CL
2026-05-28
把局部线性注意力的逐 query 求解器换成可学协方差支路,0.6B/1.7B 预训练困惑度稳定低于 Transformer,解码核不输 FlashAttention 2/3。
Transformer 的注意力主干从 2017 年沿用至今。替代路线一直有:Linear Attention 家族(DeltaNet 等)和 SSM(Mamba)用常数大小的循环状态换来亚二次复杂度,代价是上下文检索一致输给 softmax 注意力。速度省下来,记忆能力赔进去。
测试时回归框架把各种注意力统一解释成「在上下文的 KV 对上做回归」:softmax 注意力对应非参数统计里的 Nadaraya-Watson 局部常数估计,查询落在注意力分布边界时有消除不掉的偏差。前作 Local Linear Attention(LLA)把估计升级为局部线性,理论上积分 MSE 严格更小,却从未进入真实的 LLM 预训练:它要对每个 query 用共轭梯度解线性方程组,前向访存比 softmax 注意力高一个量级,正则系数 λ 调大退化、调小病态,共轭梯度还与 BF16 低精度训练不兼容。
Parallax 先把 LLA 改写成「softmax 输出减去一个协方差修正项」:修正项是 softmax 加权的 KV 协方差 ΣKV 乘一个探针向量 ρ。原版要对每个 query 解方程求 ρ,Parallax 砍掉求解器,学一个投影矩阵 WR,ρ = WR·x 从层输入一步算出,与 Q/K/V 并列多出一个 R 矩阵(实现上对 ρ 加 RMSNorm,套 RoPE 还有额外收益)。同时删掉原式的边界放大系数 η,参数化后它取值无界,会引发训练不稳定。最终输出就是 softmax 注意力输出减 ΣKV·ρ,形式干净。WR 置零时与 softmax 注意力严格等价,这给后处理迁移留了门。
同一坐标系装得下一整族机制:按带宽 h、探针来源(置零/可学/求解)、有无截距分类,h→∞ 时 Parallax 退化成带截距的线性注意力,λ→∞ 时退化回 softmax。顺带澄清一个概念:Linear Attention 里传统叫 query 的向量,在这个视角下其实是探针 ρ。
工程上,两条支路共享同一条 KV 流、同一个在线最大值与缩放因子,流式计算不增加 HBM 访存,算术强度(每字节显存搬运对应的浮点运算量)比 FlashAttention 约翻倍。解码还有个巧活:Hopper 的 WGMMA 指令最少按 64 行 tile 计,而解码每步只有 1 行 query,标准注意力里累加器大部分行闲置;Parallax 把 QK 和 RK 两个矩阵乘塞进同一条张量核指令,修正支路的计算近乎白捡。
合成任务 MAD(六项,全部用 Muon 训练;ICR 为上下文内检索,SC 为选择性复制):
| 模型 | 平均 | ICR | SC |
| Parallax | 0.716 | 0.951 | 0.988 |
| Softmax Attention | 0.672 | 0.803 | 0.950 |
| GDN | 0.665 | 0.920 | 0.939 |
| Mesa | 0.659 | 0.998 | 0.568 |
| Mamba | 0.631 | 0.756 | 0.950 |
难度拉到词表 512、上下文 2048 后,基线大幅掉点,Parallax 基本保住精度。解码核在 H200、BF16 下扫遍 batch 1–2048 × 上下文 128–32768,对 FA2/FA3 的较好者全配置持平或更快,算力对齐设置最高 1.54×,I/O 对齐设置最高 1.14×。
预训练(Qwen-3 架构,Ultra-FineWeb,4096 上下文):
| 模型(0.6B,78.6B tokens) | LMB ppl ↓ | Wiki ppl ↓ | 下游平均 ↑ |
| Transformer | 22.15 | 23.43 | 54.54 |
| 参数对齐 Transformer† | 22.35 | 23.36 | 54.90 |
| 算力对齐 Parallax† | 20.29 | 22.49 | 55.79 |
| Parallax | 18.56 | 22.25 | 55.99 |
1.7B(157.2B tokens)同样成立:LAMBADA 困惑度 13.07 → 10.80,下游平均 61.43 → 62.45。0.6B 下 KDA(52.73)和 GDN(53.67)都跑不过 Transformer 基线,Parallax 赢了;参数对齐(给 Transformer 加等量 query 头)只追回一小截差距,算力对齐(头维减半)依然明显赢,增益来自机制本身。
换 AdamW 训练,故事就变了:0.6B 下游平均,AdamW+WSD 两组 52.61 对 52.68 打平;AdamW+Cosine 下 Parallax 51.45,反而低于 Transformer 的 52.34。机制分析给出解释:修正项与输出的强度比在 Muon 深层超过 8,AdamW 下不到 4;给探针加一个可学门控,AdamW 把门压到 0.26 附近,等于主动关闭修正支路,Muon 则把门打开。权重谱定位到瓶颈:新增投影 WR 在 AdamW 下稳定秩塌得最狠。另外它的合成注意力权重可以为负、绝对值可远超 1(深层 ±40),attention sink(softmax 注意力把过多概率质量压在首 token 上的现象)明显减弱,基座 softmax 熵更高,细粒度筛 token 的活挪给了修正支路。
这是 LLA 理论路线第一次跑进 LLM 预训练,并且在参数、算力双重对齐下都赢,属于 Pareto 改进。设计可搬运:不动 softmax 主体,只加一条支路;WR 置零等价于 softmax,意味着现有预训练 checkpoint 理论上能加一个 WR 微调迁移,比换线性注意力(无法精确还原 softmax)便宜得多,这条路论文只提未验。
对更多人的提醒在优化器这边:同一架构,AdamW 下优势归零甚至反噬,Muon 下才兑现,默认 AdamW 的架构评测可能系统性漏掉这类机制。论文自称这是注意力架构研究里第一次实证的架构-优化器强协同,按给出的证据,这个定位站得住。
也要诚实:下游平均分差 1 分出头,主要增益在困惑度(尤其 LAMBADA),不是下游任务的碾压。
作者自列:只验证到 1.7B、157B tokens、4096 上下文,MoE、长上下文、稀疏变体都没做;Muon 依赖只有实证刻画没有理论;checkpoint 微调迁移路线未验证。
再补几点。卖点是检索与记忆,但训练评测上下文只有 4096,LLA 理论针对的联想记忆(associative memory)主战场在长上下文,论文没碰。WikiText 困惑度差距(22.25 对 23.43)远小于 LAMBADA(18.56 对 22.15),增益分布不均。门控实验证明 AdamW 会主动关掉修正支路,这既是发现也是脆弱性:优化器配方一换,机制表现要重测,论文也承认 Muon+WSD 未必最优。解码对比是自研 CuTeDSL 原型对 FA2/FA3(H200 单卡、BF16),训练侧端到端墙钟时间未报告。