50 lines
2.0 KiB
Markdown
50 lines
2.0 KiB
Markdown
---
|
||
title: "Stem Sparse Attention"
|
||
created: 2026-06-05
|
||
updated: 2026-06-05
|
||
type: concept
|
||
tags: [sparse-attention, causal, training-free, pre-filling]
|
||
sources: [[niu-stem-causal-sparse-attention]]
|
||
---
|
||
|
||
# Stem Sparse Attention
|
||
|
||
**Stem** 是 Niu et al. (2026) 提出的 training-free 稀疏注意力框架，从 [[causal-information-flow|因果信息流]] 的角度重新思考稀疏化策略。名称 Stem（茎）隐喻早期 token 作为信息流的结构性"主干"。
|
||
|
||
## 核心思想
|
||
|
||
> In causal attention, V₁ participates in every output; V_N only in the last. Uniform sparsity ignores this fundamental asymmetry.
|
||
|
||
Stem 将因果注意力中的 token 分为两类：
|
||
- **茎 Token（Stem Tokens）**：初始位置的 token，参与所有后续 token 的聚合 → 需要更多预算保护
|
||
- **叶 Token（Leaf Tokens）**：尾部位置的 token，仅影响局部 → 可以激进稀疏化
|
||
|
||
## 双组件架构
|
||
|
||
```
|
||
Stem
|
||
├── TPD (Token Position-Decay)
|
||
│ └── 位置依赖的 top-k：前面多，后面少
|
||
└── OAM (Output-Aware Metric)
|
||
└── 基于输出幅度选 token（不止看注意力分数）
|
||
```
|
||
|
||
## 与现有方法的根本差异
|
||
|
||
| 维度 | 传统稀疏方法 | Stem |
|
||
|------|------------|------|
|
||
| 预算分配 | 统一 top-k（所有位置相同） | 位置衰减（前多后少） |
|
||
| 选择标准 | 仅注意力分数 | 注意力分数 + Value 信息（OAM） |
|
||
| 信息流意识 | 无（逐层独立决策） | 有（保护递归依赖链） |
|
||
| 训练需求 | 部分需要微调 | Training-free，即插即用 |
|
||
|
||
## 性能
|
||
|
||
在 H20 GPU 上，Stem 在 128K 上下文中将 attention kernel 时间从 dense 的 401ms 压缩到 90ms（约 4.5× 加速），同时保持 lossless 精度。还可以作为插件叠加到 [[native-sparse-attention|NSA]] 等训练型稀疏模型上进一步压缩。
|
||
|
||
## 技术栈
|
||
|
||
- 底层实现：[[block-sparse-attention|Block Sparse Attention]] 内核（Triton）
|
||
- 目标阶段：pre-filling（首次 token 生成前的批量编码）
|
||
- 适用模型：Llama 系列、Qwen 系列、DeepSeek-V3.2、MiniCPM-4.1
|