稀疏注意力#

复杂度瓶颈#

Transformer 架构的核心是缩放点积注意力 (Scaled Dot-Product Attention)：

$$o_i = \sum_{j=1}^{n} \alpha_{ij} \cdot v_j, \quad \alpha_{ij} = \operatorname{softmax}\left(\frac{q_i^T k_j}{\sqrt{d_k}}\right)_j$$

其中 $Q, K, V \in \mathbb{R}^{n \times d}$ 分别是查询、键、值矩阵。

$QK^T$ 矩阵的形状为 $n \times n$，计算复杂度 $O(n^2 d)$，空间复杂度 $O(n^2)$。对于长序列（$n=128\text{K}$，$d=4096$，FP16），仅注意力矩阵就需要 $128\text{K} \times 128\text{K} \times 2\text{B} \approx 32\text{ GB}$ 显存——单张 GPU 无法承受。

注意力矩阵的稀疏性来源于一个直觉观察：并非所有 token 对都同等重要。大多数 token 对之间的注意力权重接近零，仅有少数 token 对承载了主要的信息交互。

这种稀疏性有两层含义：

• 自然的局部性：相邻 token 之间的依赖远强于远距离 token
• 内容相关性：语义相关的 token 即使距离远也需要交互

稀疏注意力方法试图在信息完整性和计算效率之间寻找最优平衡。

MHA（多头注意力）#

$$Attention(Q, K, V) = softmax(Q·K^T / \sqrt{d_h}) · V$$

$$q_t = W^Q h_t$$

$$k_t = W^K h_t$$

$$v_t = W^V h_t$$

• $d$: Embedding维度；
• $d_h$：每一个头的维度；
• $n_h$: Attention头数；
• $h_t ∈ ℝ^{d}$： Attention单层中第$t$ 个token的输入；
• $W^Q, W^K, W^V ∈ ℝ^{d_hn_h × d}$ ；
• $q_t,k_t,v_t ∈ ℝ^{d_hn_h}$ ； $$[q_{𝑡,1};q_{𝑡,2}; ...; q_{𝑡,𝑛_ℎ} ] = q_t$$ $$[k_{𝑡,1};k_{𝑡,2}; ...; k_{𝑡,𝑛_ℎ} ] = k_t$$ $$[v_{𝑡,1};v_{𝑡,2}; ...; v_{𝑡,𝑛_ℎ} ] = v_t$$ $$o_{t,j} = \sum_{j=1}^t{softmax_j(q_{t,i}^T·k_{j,i} / \sqrt{d_h})} · v_{j,i}$$ $$u_𝑡 = 𝑊^𝑂 [o_{𝑡,1}; o_{𝑡,2}; ...; o_{𝑡,𝑛_ℎ} ]$$
• $q_{t,i},k_{t,i},v_{t,i} ∈ ℝ^{d_h}$ : 第$i$个Attention头；
• $W^O ∈ ℝ^{d × d_hn_h}$ ：输出投影矩阵；
• 每个 token 单层需要缓存的 cache 的是 $n_h$ 套 ($k,v$) — 共 $2n_hd_h$

MQA（多查询注意力）#

所有 $n_h$ 个 Q head 共享同一套 K、V，只剩 1 套，更快的推理速度：由于内存和缓存需求大幅下降，内存带宽开销也基本消除。虽然训练效率基本保持不变，但与 MHA 相比，共享 K/V 对可能会导致模型质量和输出精度略有下降