Sparse Attention 数学学习笔记

1. Dense Attention

设输入序列为：

\[X\in\mathbb{R}^{L\times d_{model}}\]

通过投影得到：

\[Q=XW_Q,\qquad K=XW_K,\qquad V=XW_V\]

注意力分数和输出为：

\[S=\frac{QK^T}{\sqrt{d_k}}, \qquad A=\operatorname{softmax}(S), \qquad O=AV\]

对第 i 个 query：

\[o_i=\sum_{j=1}^{L}\alpha_{ij}v_j, \qquad \alpha_{ij}=\frac{e^{q_i k_j^T/\sqrt{d_k}}} {\sum_{m=1}^{L}e^{q_i k_m^T/\sqrt{d_k}}}\]

也就是说，每个 query 原则上都与全部历史 token 交互。

2. Sparse Attention 的核心：限制可见集合

对每个 query i，定义允许访问的逻辑位置集合：

\[\mathcal{S}_i\subseteq\{1,2,\ldots,L\}\]

如果是 Top-K Sparse Attention，则：

\[|\mathcal{S}_i|=k,\qquad k\ll L\]

稀疏注意力只在 𝒮_i 上归一化：

\[\alpha_{ij}^{sparse}= \frac{e^{q_i k_j^T/\sqrt{d_k}}} {\sum_{m\in\mathcal{S}_i}e^{q_i k_m^T/\sqrt{d_k}}}, \qquad j\in\mathcal{S}_i\] \[o_i^{sparse}=\sum_{j\in\mathcal{S}_i} \alpha_{ij}^{sparse}v_j\]

未选中的位置不参与当前 query 的 attention。

3. 用 mask 表示稀疏结构

定义二值 mask：

\[M_{ij}=\begin{cases} 0,&j\in\mathcal{S}_i\\ -\infty,&j\notin\mathcal{S}_i \end{cases}\]

则 Sparse Attention 可以写成标准 attention 的形式：

\[A^{sparse}=\operatorname{softmax}\left( \frac{QK^T}{\sqrt{d_k}}+M\right)\]

因此，Sparse Attention 的本质不是改变 softmax，而是把不允许访问的位置在 softmax 前屏蔽掉。

4. 选择集合从哪里来

稀疏模式的区别主要在于 𝒮_i 如何产生。

固定模式

例如滑动窗口：

\[\mathcal{S}_i=\{j\mid i-w<j\le i\}\]

其中 w 是窗口大小。它计算简单，但不能根据内容动态选择远距离 token。

内容相关的 Top-K

先计算选择分数：

\[u_{ij}=f(q_i,k_j)\]

再选择分数最高的 k 个位置：

\[\mathcal{S}_i=\operatorname{TopK}_{j\in\{1,\ldots,L\}}(u_{ij},k)\]

u_ij 不一定等于真正 attention 使用的分数。它可以由一个更小的 selector/indexer 计算，以降低选择成本。

5. Selector 分数与 Attention 分数

需要区分两种分数：

选择分数

\[u_{ij}=f_{selector}(q_i,k_j)\]

它的作用是决定哪些位置进入 𝒮_i。

注意力分数

\[s_{ij}=\frac{q_i k_j^T}{\sqrt{d_k}}\]

它的作用是在已选集合内部计算 softmax 权重。

所以完整数据流是：

query/history
      ↓
selector 分数 u_ij
      ↓
Top-K 逻辑位置 𝒮_i
      ↓
attention 分数 s_ij
      ↓
只在 𝒮_i 内 softmax

6. Top-K 对 Dense Attention 的近似

令 Dense Attention 的权重为 α_ij，选中集合的 dense 权重质量为：

\[Z_i=\sum_{j\in\mathcal{S}_i}\alpha_{ij}\]

Sparse Attention 实际上是在选中集合内重新归一化：

\[\alpha_{ij}^{sparse}=\frac{\alpha_{ij}}{Z_i}, \qquad j\in\mathcal{S}_i\]

其中 1-Z_i 是被丢弃的 dense attention mass。若所有 value 满足：

\[\|v_j\|\le V_{max}\]

则可以得到一个简单误差上界：

\[\|o_i^{dense}-o_i^{sparse}\| \le 2V_{max}(1-Z_i)\]

这说明：Top-K 的关键不是只选“数量上最少”的 token，而是尽量覆盖 dense attention 权重质量。

需要注意：

7. Decode 阶段的复杂度

对单个 decode query，Dense Attention 需要扫描 L 个历史位置：

\[T_{dense}=O\left(L(d_k+d_v)\right)\]

若 selector 已经产生了 k 个位置，真正的 Sparse Attention 只需要读取：

\[T_{attention}=O\left(k(d_k+d_v)\right)\]

但 selector 自身仍可能需要查看完整历史。例如低维 indexer 的成本可以写成：

\[T_{selector}=O(Ld_{idx})\]

因此总成本近似为：

\[T_{sparse} =O\left(Ld_{idx}+k(d_k+d_v)\right)\]

Sparse Attention 的收益来自：

8. MLA 中的 Sparse Attention

MLA 先把每个 token 的 KV 压缩到 latent：

\[c_j^{KV}=W^{DKV}h_j, \qquad d_c\ll d_k+d_v\]

之后从 latent 恢复用于 attention 的表示：

\[k_j=W^{UK}c_j^{KV}, \qquad v_j=W^{UV}c_j^{KV}\]

实际模型还可能把 content 部分和 RoPE 部分分开；这里保留抽象形式。

Sparse MLA 的计算可以写成：

\[\mathcal{S}_i=\operatorname{TopK}_j(u_{ij},k)\] \[o_i^{Sparse\text{-}MLA} =\sum_{j\in\mathcal{S}_i} \operatorname{softmax}_{j\in\mathcal{S}_i}(s_{ij}) \left(W^{UV}c_j^{KV}\right)\]

这里有两种不同的缓存数据：

数据 数学作用
Indexer key 计算 u_ij，产生 𝒮_i
MLA latent c_j^KV 在 𝒮_i 中恢复并计算真正 attention

因此，Indexer 决定“看哪些位置”，MLA latent 决定“这些位置提供什么 value”。

两套缓存的生命周期

对历史 token j，可以把数据明确写成两条独立的缓存：

缓存 每个历史 token 保存的内容 当前 query 时做什么
Indexer KV cache k_j^I，低维 indexer key 与 q_t^I 比较，扫描历史并产生 index score
MLA attention KV cache c_j^KV，MLA 压缩 latent Top-K 确定后，只取选中位置做真正 attention

它们虽然都按 token 位置 j 对齐，但不是同一份数据：

\[k_j^I\in\mathbb{R}^{d^I}, \qquad c_j^{KV}\in\mathbb{R}^{d_c}, \qquad k_j^I\ne c_j^{KV}\]

一个 decode query 的实际数学顺序是：

\[h_t\rightarrow(q_t^I,w_t^I) \rightarrow \{I_{t,j}\}_{j=1}^{L} \rightarrow \mathcal{S}_t \rightarrow \{c_j^{KV}\mid j\in\mathcal{S}_t\} \rightarrow o_t\]

所以不能把“index KV”理解成主 attention 使用的 K/V：

9. DSA 作为一个具体例子

DSA 的 indexer 并不是把 MLA 的 c_j^KV 直接拿来做 Top-K，而是为选择任务单独生成 indexer key：

\[\mathbf{k}_s^I=f_K(\mathbf{h}_s), \qquad \mathbf{q}_{t,j}^I=f_{Q,j}(\mathbf{h}_t)\]

对于历史位置 s，lightning indexer 计算：

\[I_{t,s}= \sum_{j=1}^{H^I} w_{t,j}^{I}\cdot \operatorname{ReLU}\left( \mathbf{q}_{t,j}^{I}\cdot\mathbf{k}_{s}^{I} \right)\]

再按 index score 选择：

\[\mathcal{S}_t= \left\{ s\mid I_{t,s}\in\operatorname{TopK}(I_{t,:},k) \right\}\]

最后才读取另一条缓存中的 attention 数据：

\[\left\{c_s^{KV}\mid s\in\mathcal{S}_t\right\}\]

并将它们送入 MLA attention。整个过程可以概括为：

\[\underbrace{\{\mathbf{k}_s^I\}_{s=1}^{L}}_{\text{Indexer KV cache}} \xrightarrow{\text{score + Top-K}} \underbrace{\mathcal{S}_t}_{\text{逻辑位置}} \xrightarrow{\text{gather}} \underbrace{\{c_s^{KV}\mid s\in\mathcal{S}_t\}}_{\text{MLA attention KV cache}}\]

因此，indexer KV 的维度较小，适合扫描长历史；MLA KV 的数据量更大，只对 Top-K 位置读取。

10. 稀疏 Attention 的数学数据流

h_1,...,h_L                         当前 h_t
      │                                  │
      ├─ 生成 k_s^I                   ├─ 生成 q_t^I,w_t^I
      │                                  │
      ├─ Indexer KV cache               └─ 与全部 k_s^I 算 I_t,s
      │                                            │
      └─ 生成 c_s^KV                              ↓
           MLA attention KV cache            Top-K 集合 𝒮_t
                                                     │
                                                     ↓
                                    只取 {c_s^KV | s ∈ 𝒮_t}
                                                     │
                                                     ↓
                                            Sparse MLA attention
                                                     │
                                                     ↓
                                                     o_t

最重要的关系是：

\[\text{选择集合 }\mathcal{S}_i \quad\rightarrow\quad \text{参与 attention 的数据} \quad\rightarrow\quad \text{输出 }o_i\]

11. 与 HiSparse 的关系

Sparse Attention 的数学算法决定每一步只访问 𝒮_i；HiSparse 进一步利用这一点，把未被当前 step 访问的 KV 放到其他存储层。

因此：

\[\text{Sparse Attention：决定访问谁}\] \[\text{HiSparse：决定这些数据暂时放在哪里}\]

前者改变计算集合，后者只改变 KV placement。只要 𝒮_i 中的 KV 在计算前仍能被准确取回，HiSparse 不会改变 Sparse Attention 的数学输出。

参考