跳到主要内容

注意力机制变体

核心要点

  • dense attention $O(n^2)$ 是根本算力瓶颈
  • 三个降算力方向:稀疏 / 线性 / 状态空间
  • 稀疏检索强、工业最常用;线性 / SSM 质量略低
  • Hybrid(少量 attention + 多数高效层)是工程主流
  • 各方向的实现机制见 06 注意力机制实现卷

本文回答 02-第一性挑战 中 attention $O(n^2)$ 算力的算法层缓解方案,是问题视角的方向导览——讲长上下文有哪些降算力的路、各缓解什么、怎么选;具体算子怎么实现见 06-大模型解构 / 注意力机制实现卷。模型变体降算力,与位置编码外推(03-位置编码与外推)、KV 架构压缩(05-kv-cache架构压缩)三者正交,工业模型常同时用。

为什么 attention 是 $O(n^2)$

原版 attention 算每个 query 与每个 key 的内积,共 $n^2$ 对。复杂度推导见 02-第一性挑战 § 挑战 1

降算力的根本只有三条路——少算(稀疏)、换计算顺序(线性)、换算子(状态空间)。下面逐条讲思路与缓解的问题,机制实现链到 06。

三个降算力方向缓解什么?

方向思路缓解什么代价实现详见
稀疏每 query 只算部分 key(按位置或按内容选)算力 $O(n^2)\to O(n\cdot k)$,检索强KV 显存仍 $O(n)$;稀疏 kernel 工程难07 / 08
线性用核函数重排计算顺序,先算 KV 侧$O(n d^2)$ 线性、无 KV cache失去 softmax 尖锐选择性,质量弱11
状态空间 (SSM)用递归隐状态代替全 attention推理 $O(1)$/token、无 KV cache精确检索差(隐状态固定)11

@tbl-longctx-attn-three-paths 注意力降算力的三个方向

稀疏:少算 query-key 对

稀疏注意力让每个 query 只看一部分 key,把 $O(n^2)$ 降到 $O(n\cdot k)$,且检索能力强(看到的位置是精确的)。按"怎么选 key"分两支:按位置固定选(滑窗 SWA、Longformer、BigBird)和按内容动态选(NSA、DSA 的 indexer top-k)。NSA[1] 等原生稀疏在 2025 年让稀疏 attention 重回工业视野。机制实现见 07-滑窗与固定稀疏08-动态稀疏选择

线性:换计算顺序

线性注意力用核函数近似 softmax,靠矩阵结合律先算 KV 侧,把复杂度降到 $O(nd^2)$、推理可 RNN 化无 KV cache。代表 Performer / RWKV[2]。代价是失去 softmax 的尖锐选择性——纯线性在 1B+ 模型上质量弱于 dense,故多作 hybrid 的长距离分量。机制实现见 11-线性注意力与 SSM

状态空间:换算子

SSM 用递归隐状态(借控制论状态空间方程)代替全 attention,推理 $O(1)$/token、无 KV cache、长度可无限扩展。Mamba[3] 的选择性 SSM 让它 GPU 友好,Mamba-2 的 SSD 进一步用矩阵乘加速。代价是精确检索差(隐状态固定,无法保留所有 token)。机制实现见 11-线性注意力与 SSM

Hybrid 为什么是工程主流?

纯 dense、纯 SSM、纯线性都有短板,hybrid 按层混合取长补短——少量 attention 层负责精确检索,多数高效层负责序列处理。这是当前长上下文的工程主流。

  • Jamba[4]:Mamba : Transformer = 7 : 1(每 8 层 1 层 attention),256K 上下文,已部署。
  • Hymba(NVIDIA):每层内并行组合 attention head 与 Mamba head,输出融合,让"短距离精确 + 长距离平滑"逐层对齐。
  • Falcon-Mamba(TII):纯 Mamba 通过架构调整接近 Transformer,证明 SSM 也可单独承担。

hybrid 在多数任务上质量接近 dense Transformer,但长上下文显存与推理速度大幅占优。

三类方法怎么选?

维度稀疏线性SSMHybrid
推理:单 token$O(k)$$O(d^2)$$O(1)$取决于混合比
KV cache$O(n)$仅 attention 层有
检索能力强(保留 attention 层)
长依赖取决于窗口 / 选择
工业成熟度高(NSA/SWA)中(RWKV)中(Mamba 研究为主)高(Jamba 已部署)

@tbl-longctx-attn-selection 注意力机制变体选型对照

当前实践共识

  • 保守路线:dense attention + GQA/MQA/MLA 压缩 KV(→ 05-kv-cache架构压缩)+ FlashAttention + 长上下文外推(→ 03-位置编码与外推
  • 激进路线:Hybrid(Jamba / Hymba),少量 attention 锚定 + 多数 SSM 高效
  • 新兴方向:原生稀疏(NSA),训练即稀疏,保留精确性同时降算力

Takeaway

知识点核心结论
三个方向稀疏(少算)/ 线性(换顺序)/ SSM(换算子),各缓解算力但有不同代价
稀疏检索强、工业最常用;NSA 训练即稀疏(实现见 07/08)
线性 / SSM$O(n)$ 但纯用质量弱于 dense,宜作 hybrid 分量(实现见 11)
Hybrid少量 attention 锚定 + 多数高效层,质量与效率兼得,激进路线首选
与本篇分工本篇讲方向与选型,机制实现全部在 06 注意力机制实现卷

@tbl-longctx-attn-takeaway 全文要点

延伸阅读

参考资料

  1. Yuan et al., Native Sparse Attention: Hardware-Aligned and Natively Trainable Sparse Attention (NSA), arXiv 2025. https://arxiv.org/abs/2502.11089
  2. Peng et al., RWKV: Reinventing RNNs for the Transformer Era, arXiv 2023. https://arxiv.org/abs/2305.13048
  3. Gu & Dao, Mamba: Linear-Time Sequence Modeling with Selective State Spaces, arXiv 2023. https://arxiv.org/abs/2312.00752
  4. Lieber et al., Jamba: A Hybrid Transformer-Mamba Language Model, arXiv 2024. https://arxiv.org/abs/2403.19887