多头注意力机制

$$X \rightarrow Q, K, V \rightarrow QK^T \rightarrow \text{Mask} \rightarrow \text{Softmax} \rightarrow PV \rightarrow \text{Merge Heads}$$

Multi-head Attention·

设 Transformer 某一层的输入为

$$X\in\mathbb R^{B\times S\times H}$$

这里：

• $B$：batch size，一次处理多少条序列；
• $S$：sequence length，每条序列多少个 token；
• $H$：hidden size，每个 token 的隐藏向量维度。

Attention 首先对 $X$ 做三个线性变换：

$$Q = XW_Q,\qquad K = XW_K,\qquad V = XW_V.$$

如果暂时不拆 head，它们通常都是

$$Q, K, V\in\mathbb R^{B\times S\times H}.$$

然后进入 Multi-Head Attention。

设 head 数量为 $N_h$，每个 head 的维度为 $D_h$。通常有

$$H = N_hD_h.$$

原始：

$$Q\in\mathbb R^{B\times S\times H}.$$

因为 $H = N_hD_h$，先 reshape：

$$[B, S, H] \rightarrow [B, S, N_h, D_h].$$

然后为了方便矩阵乘法，把 head 维提前：

$$[B, S, N_h, D_h] \rightarrow [B, N_h, S, D_h].$$

于是：

$$Q, K, V \in \mathbb R^{B\times N_h\times S\times D_h}.$$

例如：

$$[2,32,1024,128].$$

可以理解为每个 batch 有 $N_h$ 个独立的 Attention head，每个 head 都看到 S 个 token，每个 token 在这个 head 中是 $D_h$ 维。

对于一个 head：

$$Q_h\in\mathbb R^{S\times D_h}$$

而

$$K_h\in\mathbb R^{S\times D_h}.$$

为了做矩阵乘法，对 $K$ 最后两个维度转置：

$$K_h^T \in \mathbb R^{D_h\times S}.$$

所以：

$$Q_hK_h^T : (S\times D_h)(D_h\times S)$$

得到：

$$S\times S$$

的矩阵。

因此完整 batch + multi-head 情况下：

$$QK^T \in \mathbb R^{B\times N_h\times S\times S}.$$

为了避免和 sequence length $S$ 混淆，可以把 score matrix 记为 $A$：

$$A = \frac{QK^T}{\sqrt{D_h}} \in \mathbb R^{B\times N_h\times S\times S}.$$

假设一句话有 4 个 token：

我  喜欢  深度  学习

那么 $S = 4$。每个 token 都要询问：“我应该关注哪些 token？”

所以每一个 Query 都要和所有 Key 做点积。得到：

$$A = \begin{pmatrix} a_{11}&a_{12}&a_{13}&a_{14}\\ a_{21}&a_{22}&a_{23}&a_{24}\\ a_{31}&a_{32}&a_{33}&a_{34}\\ a_{41}&a_{42}&a_{43}&a_{44} \end{pmatrix}.$$

其中：

$$a_{ij} = \frac{q_i^Tk_j}{\sqrt{D_h}}$$

表示第 $i$ 个 token 的 Query，认为第 $j$ 个 token 的 Key 有多相关。

因此：

• 行 $i$：当前 Query token；
• 列 $j$：它正在观察的 Key token。

例如 $a_{34}$ 就是：第 3 个 token 对第 4 个 token 的注意力匹配分数。

$$P, A \in \mathbb R^{B\times N_h\times S\times S}.$$

这里出现了：

$$S^2.$$

因此标准 Attention 的 score matrix 元素数量：

$$BN_hS^2.$$

当 sequence length 增大时，增长速度是平方级。这是长上下文昂贵的原因之一。

不过 Attention 不是直接计算 $QK^T$，而是：

$$\frac{QK^T}{\sqrt{D_h}}.$$

原因和数值尺度有关。假设 Query 和 Key 每一个分量大致满足均值 0、方差 1。点积

$$q^Tk = \sum_{r = 1}^{D_h}q_rk_r$$

的方差会随着 $D_h$ 增长。如果直接进入 Softmax，例如：

$$\operatorname{softmax}(100,0,-100)$$

分布会非常尖锐，容易导致梯度不好训练。除以 $\sqrt{D_h}$ 可以把点积的典型尺度稳定下来：

$$\frac{q^Tk}{\sqrt{D_h}}.$$

这就是 Scaled Dot-Product Attention 中 “Scaled” 的来源。

自回归语言模型有一个非常重要的限制：第 $i$ 个 token 不能偷看未来 token。

例如：

我 喜欢 深度 学习

预测“喜欢”时，不能提前看到后面的“深度 学习”。因此需要 causal mask。

对于 $S = 4$，mask 可以理解成：

$$M = \begin{pmatrix} 0&-\infty&-\infty&-\infty\\ 0&0&-\infty&-\infty\\ 0&0&0&-\infty\\ 0&0&0&0 \end{pmatrix}.$$

然后计算：

$$A+M.$$

第一行：

$$(a_{11},-\infty,-\infty,-\infty)$$

第二行：

$$(a_{21}, a_{22},-\infty,-\infty)$$

因为后面马上做 Softmax：

$$P =\operatorname{softmax}(A+M).$$

而 $e^{-\infty}=0$，所以一个被 mask 的位置最终概率就是 0。

例如：

$$\operatorname{softmax}(3,1,-\infty)$$

第三项概率：

$$\frac{e^{-\infty}} {e^3+e^1+e^{-\infty}} = 0.$$

因此 causal mask 最终实现：

$$\text{未来 token 的 Attention 权重}= 0.$$

---

现在：

$$A\in \mathbb R^{B\times N_h\times S\times S}.$$

对于每一个：

• batch $b$
• head $h$
• query token $i$

都有一整行：

$$A_{b, h, i,:}.$$

Softmax 对最后一个维度，也就是 Key/token 维度 进行：

$$P_{b, h, i,:} = \operatorname{softmax} (A_{b, h, i,:}).$$

所以：

$$\sum_{j = 1}^S P_{b, h, i, j}= 1.$$

每一个 Query token 都得到一组对所有 Key token 的概率权重。例如：

$$P_{i,:} = (0.1,0.2,0.6,0.1).$$

表示 token $i$ 把 10% 注意力放第 1 个 token，20% 放第 2 个，60% 放第 3 个，10% 放第 4 个。

得到 Attention probability：

$$P\in \mathbb R^{B\times N_h\times S\times S}.$$

而：

$$V\in \mathbb R^{B\times N_h\times S\times D_h}.$$

计算：

$$O = PV.$$

观察最后两个矩阵维度：

$$(S\times S)(S\times D_h)$$

得到：

$$S\times D_h.$$

所以：

$$O \in \mathbb R^{B\times N_h\times S\times D_h}.$$

跟原来的 $Q, K, V$ shape 一样。

对于 token $i$：

$$o_i = \sum_{j = 1}^S P_{ij}v_j.$$

例如 Attention 权重：

$$P_{i,:} = (0.1,0.2,0.6,0.1).$$

那么：

$$o_i = 0.1v_1 + 0.2v_2 + 0.6v_3 + 0.1v_4.$$

所以 Attention 的真正含义是 根据 Query-Key 相似度，对不同 Value 做加权求和。

因此：

• $Q$：我在寻找什么？
• $K$：我有什么特征可供匹配？
• $V$：如果你关注我，你真正获得什么信息？

---

现在：

$$O \in \mathbb R^{B\times N_h\times S\times D_h}.$$

先 transpose：

$$[B, N_h, S, D_h] \rightarrow [B, S, N_h, D_h].$$

然后把最后两个维度拼起来：

$$N_hD_h = H.$$

所以：

$$[B, S, N_h, D_h] \rightarrow [B, S, H].$$

最终：

$$O\in\mathbb R^{B\times S\times H}.$$

和原始输入

$$X\in\mathbb R^{B\times S\times H}$$

shape 完全一致。

Flash Attention·

Flash Attention 的核心不是改变 Attention 数学公式，而是不把完整 $S\times S$ 中间矩阵写回 HBM。

普通实现逻辑上可能类似：

$$Q, K \rightarrow QK^T.$$

把巨大的 $S\times S$ 矩阵写入 GPU HBM。然后再：

$$\text{HBM} \rightarrow \text{Softmax}$$

再写回。然后：

$$P, V \rightarrow PV.$$

大量时间浪费在：

$$\text{HBM 读写}.$$

而不仅仅是 FLOPs。Flash Attention 不一次生成整个：

$$S\times S.$$

而是把 $Q, K, V$ 切成小块。例如概念上：

Q block
   ↓
K block → 算局部 QKᵀ
   ↓
局部 Softmax
   ↓
和 V block 累积

只在 GPU 更快但容量更小的片上存储中保留当前 block 所需数据，例如：

• registers；
• shared memory / SRAM。

算完这个块后立即把贡献累积到输出。因此不需要：

$$\text{完整 materialize 一个 }S\times S\text{ score matrix}.$$

普通 Softmax：

$$p_j = \frac{e^{x_j}} {\sum_k e^{x_k}}$$

通常需要知道整行：

$$x_1,\dots, x_S.$$

而数值稳定版本还需要：

$$m =\max_jx_j.$$

然后：

$$p_j = \frac{e^{x_j-m}} {\sum_ke^{x_k-m}}.$$

看起来必须一次拿到所有 $S$ 个元素。

Flash Attention 使用一种 online softmax 思想分块处理时持续维护当前行的最大值和归一化因子。

假设已经处理了一部分，保存：

$$m_{\text{old}}$$

和：

$$l_{\text{old}} = \sum e^{x_i-m_{\text{old}}}.$$

新 block 出现新的最大值之后，可以重新缩放旧累积值，从而得到正确的新 normalization。因此无需保存整行 score，也能精确完成稳定 Softmax。事实上，Flash Attention 没有消除标准 Attention 的平方计算复杂度。Flash Attention 主要是 IO-aware exact attention，显著降低了 HBM IO 带宽从而降低显存压力。

标准 self-attention 仍然需要比较：

$$S\times S$$

个 Query-Key 对。所以主要计算复杂度仍然近似：

$$O(S^2D_h).$$

多头总体近似：

$$O(S^2H).$$