注意力机制变体
核心要点:
- 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,但长上下文显存与推理速度大幅占优。
三类方法怎么选?
| 维度 | 稀疏 | 线性 | SSM | Hybrid |
|---|---|---|---|---|
| 推理:单 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 全文要点
延伸阅读
- Vaswani et al., 2017. Attention Is All You Need.
参考资料
- Yuan et al., Native Sparse Attention: Hardware-Aligned and Natively Trainable Sparse Attention (NSA), arXiv 2025. https://arxiv.org/abs/2502.11089
- Peng et al., RWKV: Reinventing RNNs for the Transformer Era, arXiv 2023. https://arxiv.org/abs/2305.13048
- Gu & Dao, Mamba: Linear-Time Sequence Modeling with Selective State Spaces, arXiv 2023. https://arxiv.org/abs/2312.00752
- Lieber et al., Jamba: A Hybrid Transformer-Mamba Language Model, arXiv 2024. https://arxiv.org/abs/2403.19887