[arXiv 2024] Stem: 重新审视稀疏注意力中的因果信息流

TL;DR

在处理超长文本时，标准的 Self-Attention 面临二次方复杂度的性能瓶颈。本文作者提出了 Stem，一种无需训练的稀疏注意力框架。它通过两个核心武器：位置衰减策略 (TPD) 保护关键的初始 Token，以及 输出感知度量 (OAM) 识别高能信号 Token。Stem 在将预填充（Pre-filling）延迟缩减为原来的 1/4 的同时，在长文本理解任务上几乎追平了全注意力（Full Attention）的精度。

痛点深挖：被忽视的递归依赖

目前的稀疏化方案（如 MInference, XAttention）往往假设：每一行（Query）的稀疏预算应该是均匀的。 然而，作者通过理论分析指出，因果注意力（Causal Attention）本质上是一个递归聚合过程。

如下图所示，第 l 层的第 n 个 Token 是由前一层第 1 到第 n 个 Token 聚合而成的。这意味着：

• 初始位置的 Token ($V_1$)：参与了后续每一个输出的计算，是信息的“树根”。一旦在这里发生剪枝误差，误差会随着层数加深被指数级放大（Global Distortion）。
• 序列末尾的 Token ($V_N$)：仅参与最后一个输出的计算，剪枝的影响非常局部。

核心方法论：Stem 的“粗到细”架构

Stem 的设计遵循两个物理直觉：位置很重要，内容更重要。

1. Token Position-Decay (TPD) 预算分配

由于初始 Token 承担了“锚点”作用，Stem 不再使用固定的 Top-k，而是采用线性衰减策略：给序列开端的 Query 分配极高的 Top-k 预算（k-start），随着位置后移，逐渐压缩预算。这在保证全局信息链完整性的同时，大幅削减了序列后半段的冗余计算。

2. Output-Aware Metric (OAM) 精准选型

传统的标准是“谁跟我相关（Score-Aware）”，但 Stem 认为“谁分量重（Signal Magnitude）”同样关键。如果一个 Token 的注意力分数很高，但其 Value 向量接近零，它对输出的影响微乎其微。 Stem 提出的度量公式为： ${\mathcal{M}}_{i,j}=Q_i{K}_j^T+\beta \cdot \max (0,\log \Vert V_j{\Vert }_2)$ 这个公式巧妙地平衡了路由相关性与信号能量值。

实验结果：以一当十的效率

作者在 Llama-3.1-8B 和 Qwen3-8B 上进行了严耕测试。

• 精度表现：在 LongBench 榜单上，Stem 以 25% 的平均计算量，在综合得分上击败了消耗更多算力的同类方法。在超复杂合成任务（SYN）中，Stem 表现出极强的鲁棒性。
• 推理加速：搭载 Triton 实现的代码块稀疏内核，Stem 在 128K 上下文下展现了惊人的提速效果。

模型兼容性：强强联手

Stem 的另一大亮点是其通用性。即使是 DeepSeek-V3.2（带有训练好的 DSA 稀疏注意力）或 MiniCPM-4.1，集成 Stem 后仍能在现有稀疏基础上继续压榨出 15% 以上的算力节省，且性能几乎无损。

深度洞察与总结

Stem 的核心贡献在于它不仅仅关注了“稀疏化”，更关注了“因果对齐”。

关键视角回顾：

1. 打破均匀假设：初始 Token 是“骨干（Stem）”，必须重金保护。
2. 引入 Value 感知：基于重建误差最小化推导出的对数模数项，为 Top-k 选择提供了物理依据。

局限性与展望： 虽然 Stem 在推理阶段非常高效，但其超参数（如衰减率 $\mu$ 和系数 $\beta$）在不同模型上可能需要微调。未来如果能将这种“位置敏感”的先验融合进大模型的预训练中，可能会产生更高效的 Native Sparse 架构。

One-liner Takeaway: Stem 证明了只要守住初始位置的“技术骨干”，即便大规模裁剪后续 Token，长文本模型的灵魂也不会丢。