ZeroS: Zero-Sum Linear Attention for Efficient Transformers¶

会议: NeurIPS 2025
arXiv: 2602.05230
代码: 有
领域: LLM效率
关键词: 零和注意力, 线性注意力, softmax分解, 径向角向解耦, O(N)复杂度

一句话总结¶

通过移除 softmax 的零阶均匀项 \(1/t\)，构建零和权重的线性注意力机制 ZeroS，突破凸组合只能做加法混合的限制，支持单层内的差分/对比操作，在保持 \(O(Nd^2)\) 线性复杂度的同时，在多个序列建模基准上匹配甚至超越标准 softmax 注意力。

研究背景与动机¶

领域现状：线性注意力通过核特征映射 \(\phi(q)\phi(k)^\top\) 近似 softmax 注意力，将复杂度从 \(O(N^2)\) 降至 \(O(N)\)。代表性方法包括 Performer、GLA、Mamba 等。但这些方法普遍不如标准 softmax 注意力。
现有痛点：
凸组合瓶颈：softmax 注意力产生的是非负权重，只能对 value 向量做"加性混合"。线性注意力为保证数值稳定也限制权重为正。但这意味着单层注意力无法表达差分或对比操作——即使只有两个 token，也不能用一层做 \(v_1 - v_2\)。
均匀权重偏差：softmax 的 Taylor 展开中，零阶项 \(1/t\) 产生了持续的平均效应。在长序列中，这个均匀成分稀释了聚焦注意力，导致"注意力退化"。
核心矛盾：要实现线性复杂度就需要可分解的核形式，但已有的核方法都试图模拟 softmax 的非负性，这恰恰是限制表达力的根源。
本文要解决什么：(a) 能否允许注意力权重取负值同时保持数值稳定？(b) 能否让线性注意力在单层内就支持对比操作，弥合与 softmax 的性能差距？
切入角度：从 softmax 的 Taylor 展开入手——将 \(\text{softmax}(s_i)\) 分解为零阶（\(1/t\)）、一阶（\(\delta_i/t\)）、高阶残差（\(\varepsilon_i\)）。零阶项只提供均匀平均，对 token 间交互无贡献。移除它后得到自然的零和权重 \(\sum_i w_i = 0\)，既有正又有负，且范数稳定性不依赖序列长度。
核心 idea：与其设计复杂核来近似 softmax 同时保留非负性，不如反过来——去掉 softmax 中造成非负约束的零阶项，用零和残差权重构建更有表达力的线性注意力。

方法详解¶

整体框架¶

ZeroS 直接替换 Transformer 中的多头注意力模块，保留其他组件（MLP、LN、残差连接）不变。每个 ZeroS 层包含三个核心操作：(1) 计算偏差 logits \(s_i\)（仅依赖步骤 \(i\)，不依赖当前步 \(t\)），(2) 计算重加权零和 softmax 权重 \(w_{t,i}\)，(3) 将径向权重与角向余弦分量相乘得到最终注意力输出。整体保持 \(O(Nd^2)\) 时间和 \(O(d^2)\) 内存。

关键设计¶

Softmax 零阶移除与零和权重构建
做什么：将 softmax 分解为零阶（\(1/t\)）、一阶偏移（\(\delta_{t,i}/t\)）和高阶残差（\(\varepsilon_{t,i}\)），去掉零阶项后用可学习门控混合剩余项
核心思路：\(w_{t,i} = \sigma_t^1 \frac{\delta_{t,i}}{t} + \sigma_t^h \varepsilon_{t,i}\)，其中 \(\sigma_t^1, \sigma_t^h\) 是依赖当前步 \(t\) 的 sigmoid 门控。所有权重满足 \(\sum_i w_{t,i} = 0\)，自然支持正负值
理论支撑：Proposition 3.1 证明零和权重的可达集严格包含凸组合的可达集（只要 value 向量不完全相同），即零和注意力比标准注意力表达力更强
设计动机：零阶项 \(1/t\) 只贡献均匀平均（对所有 \((t,i)\) 对无区分），是表达力的"最低成本基"。移除后不失一般性（第一层可选保留零阶项以覆盖仿射包），后续层通过残差连接恢复均匀方向的信息
径向-角向解耦（Radial-Angular Decoupling）
做什么：将最终注意力权重分解为径向部分 \(r_{t,i}\)（来自零和 softmax）和角向部分 \(\cos\theta\)（query-key 方向相似度）
核心思路：\(o_t = \sum_{i=1}^{t} r_{t,i} \cos\theta \cdot v_i\)，其中 \(\cos\theta = \hat{q}_t \hat{k}_i^\top\)（query/key 归一化后的内积）。由于零和权重本身已经数值稳定，可以直接将 \(\cos\theta\) 乘进来而无需正性约束
设计动机：标准 softmax 中 \(\exp(\|q\|\|k\|\cos\theta)\) 将幅值和方向耦合在一起——当 \(\cos\theta\) 变号时，大正值变成极小值，这种"翻转效应"是 softmax 表达力的关键。线性注意力用 \(\phi(q)\phi(k)\) 替代，但正映射 \(\phi\) 限制了角度范围（\(<90°\)），丢失了翻转效应。ZeroS 通过解耦明确恢复这一能力
与 RoPE 的兼容：\(\cos\theta' = \hat{q}_t R_{t-i} \hat{k}_i^\top\)，旋转位置编码自然融入角向分量
偏差 logits 设计
做什么：为每个步骤 \(i\) 计算一个标量 logit \(s_i\)（不依赖当前步 \(t\)）
核心思路：\(s_i = -\frac{1}{\sqrt{d}} u_i \bar{u}_i^\top\)，其中 \(u_i = x_i W_u\)，\(\bar{u}_i = \frac{e^\tau \mu + \sum_{j=1}^i u_j}{e^\tau + i}\) 是带可学习先验的累积平均。\(s_i\) 度量步骤 \(i\) 相对历史的偏离程度
设计动机：logits 仅依赖 \(i\) 而非 \((t,i)\)，才能通过前缀和实现线性时间计算。步骤 \(t\) 的影响通过门控 \(\sigma_t^1, \sigma_t^h\) 注入，实现不同阶零和分量的动态控制
线性时间前缀扫描
做什么：通过维护 5 个前缀和状态实现 \(O(d^2)\) 每步计算
核心思路：\(E_t = \sum e^{s_i}\)，\(P_t = \sum s_i\)，\(F_t = \sum e^{s_i} \hat{k}_i^\top v_i\)，\(G_t = \sum s_i \hat{k}_i^\top v_i\)，\(H_t = \sum \hat{k}_i^\top v_i\)。最终输出 \(o_t = \hat{q}_t(\alpha_t F_t + \beta_t G_t + \gamma_t H_t)\)，其中 \(\alpha_t, \beta_t, \gamma_t\) 是由门控和前缀和标量组合得到的系数
设计动机：所有状态矩阵大小为 \(d \times d\)，更新复杂度 \(O(d^2)\)，总复杂度 \(O(Nd^2)\)，内存 \(O(d^2)\)——与其他线性注意力相同

损失函数 / 训练策略¶

直接替换标准 Transformer 的多头注意力，保留原有 MLP/嵌入/超参数/训练配置
默认不在第一层添加零阶项（实验中影响极小）
因果（decoder）和非因果（encoder）两种模式均支持

实验关键数据¶

主实验：MAD 基准（上下文学习）¶

模型	Compress	Fuzzy	In-Context	Memorize	Noisy	Sel.Copy	平均
Transformer	51.6	29.8	94.1	85.2	86.8	99.6	74.5
Mamba	52.7	6.70	90.4	89.5	90.1	86.3	69.3
DeltaNet	42.2	35.7	100	52.8	100	100	71.8
LinAttn	31.1	8.15	91.0	74.9	75.6	93.1	62.3
ZeroS	44.0	14.9	99.9	88.1	96.1	97.8	73.5
ZeroS-SM	45.2	28.0	100	84.3	96.6	98.5	75.4

消融实验（WikiText-103 + MAD）¶

配置	WikiText PPL↓	MAD 平均
ZeroS（完整）	24.61	73.5
+ 保留零阶项	24.74 (+0.13)	69.4 (-4.1)
替换为标准 softmax（无 RWSM）	24.97 (+0.36)	67.6 (-5.9)
去掉门控	—	70.8 (-2.7)
去掉 LayerNorm	—	69.4 (-4.1)

关键发现¶

零和权重对上下文学习至关重要：In-Context Recall 从 91.4（含零阶）提升到 99.9（零和），验证了凸组合限制算法推理能力的假设
WikiText 上超越 Transformer：ZeroS PPL 24.61 vs Transformer 24.78，线性注意力首次在此基准上击败标准注意力
ImageNet 分类也有效：DeiT-Tiny 上 ZeroS 75.51% vs DeiT 72.20%，在视觉任务上同样适用
时间序列全面领先：在 Weather/Solar/ETT 等数据集上优于 GLA、AFT 和领域方法 iTransformer
记忆能力未受影响：尽管使用负权重，Memorize 任务得分（88.1）与 Transformer（85.2）和 Mamba（89.5）可比，说明零阶移除不损害序列记忆

亮点与洞察¶

"去掉它"的简洁哲学：不是设计更复杂的核来近似 softmax，而是大胆去掉造成限制的零阶项。简单到令人惊讶——居然没人早this么做
凸组合→零和的理论飞跃：Proposition 3.1 严格证明零和可达集严格包含凸组合可达集。这是对线性注意力表达力瓶颈的精确诊断
数值稳定性的优雅保证：Lemma 3.4 证明 \(\|\sum_i w_{t,i} v_i\| = O(B)\)（不依赖 \(t\)），说明即使权重有负值，输出范数仍然有界——不需要凸组合的范数保证
径向-角向解耦恢复"翻转效应"：这是对 softmax 注意力核心表达力来源的深刻理解——不是 \(\exp\) 本身重要，而是幅值-方向的耦合交互重要

局限性 / 可改进方向¶

未提供 GPU 加速实现（如 CUDA kernel），仅有算法层面的线性复杂度；实际速度可能不如工程优化的 Mamba/GLA
受资源限制未做大规模 LLM 预训练（7B+）的验证，所有实验在小模型（<50M 参数）上进行
主要评估自回归任务，非因果场景（如 BERT-style）的探索不足
偏差 logits 的设计（负内积 + 累积平均）虽然合理但不唯一，可能存在更好的 logit 函数设计

评分¶

新颖性: ⭐⭐⭐⭐⭐ 零和权重的理论视角全新且深刻，"去掉零阶项"的切入既简洁又有力
实验充分度: ⭐⭐⭐⭐ MAD/WikiText/ImageNet/时间序列多任务验证，消融详尽，但大规模 LLM 实验缺失
写作质量: ⭐⭐⭐⭐⭐ Softmax 分解的推导链条清晰，理论命题与实验结论一一对应
价值: ⭐⭐⭐⭐⭐ 对线性注意力的理论和实践都有重要突破，可能改变后续线性注意力的设计范式