Self Attention 计算#

Step 0：输入#

假设经过线性投影后，Q、K、V 矩阵为：

$Q = \begin{bmatrix} 1 & 0 & 1 & 0 \\ 0 & 1 & 0 & 1 \\ 1 & 1 & 0 & 0 \end{bmatrix}, \quad K = \begin{bmatrix} 0 & 1 & 1 & 0 \\ 1 & 0 & 0 & 1 \\ 1 & 1 & 0 & 0 \end{bmatrix}, \quad V = \begin{bmatrix} 1 & 2 & 3 & 4 \\ 5 & 6 & 7 & 8 \\ 9 & 10 & 11 & 12 \end{bmatrix}$

Step 1：计算 $QK^T$（相似度矩阵）#

$QK^T = \begin{bmatrix} 1 & 0 & 1 & 0 \\ 0 & 1 & 0 & 1 \\ 1 & 1 & 0 & 0 \end{bmatrix} \begin{bmatrix} 0 & 1 & 1 \\ 1 & 0 & 1 \\ 1 & 0 & 0 \\ 0 & 1 & 0 \end{bmatrix} = \begin{bmatrix} 1 & 1 & 1 \\ 1 & 1 & 1 \\ 1 & 1 & 2 \end{bmatrix}$

矩阵中第 $i$ 行第 $j$ 列 = 词 $i$ 的 Q 与词 $j$ 的 K 的点积 = 它们的相关度。

Step 2：缩放 $\div \sqrt{d_k}$#

$\frac{QK^T}{\sqrt{4}} = \frac{1}{2}\begin{bmatrix} 1 & 1 & 1 \\ 1 & 1 & 1 \\ 1 & 1 & 2 \end{bmatrix} = \begin{bmatrix} 0.5 & 0.5 & 0.5 \\ 0.5 & 0.5 & 0.5 \\ 0.5 & 0.5 & 1.0 \end{bmatrix}$

为什么要除以 ？

当 较大时，Q 和 K 的点积结果方差约为 ，值会很大。大值送入 Softmax 后会产生接近 one-hot 的分布（一个接近 1，其余接近 0），梯度几乎为零——训练卡住。除以 将方差稳定在 1 附近。

数学证明：若 $q_i, k_j \sim \mathcal{N}(0, 1)$，则 $q^T k = \sum_{i=1}^{d_k} q_i k_i$ 的方差 = $d_k$，除以 $\sqrt{d_k}$ 后方差 = 1。

Step 3：Softmax 归一化#

对每一行做 Softmax（每个词对所有词的注意力权重归一化为概率分布）：

$\text{softmax}(\begin{bmatrix} 0.5 & 0.5 & 0.5 \end{bmatrix}) = \begin{bmatrix} 0.333 & 0.333 & 0.333 \end{bmatrix}$

$\text{softmax}(\begin{bmatrix} 0.5 & 0.5 & 1.0 \end{bmatrix}) = \begin{bmatrix} 0.262 & 0.262 & 0.476 \end{bmatrix}$

完整的注意力权重矩阵 $A$：

$A = \begin{bmatrix} 0.333 & 0.333 & 0.333 \\ 0.333 & 0.333 & 0.333 \\ 0.262 & 0.262 & 0.476 \end{bmatrix}$

Step 4：加权求和 $A \times V$#

$\text{Output} = A \cdot V = \begin{bmatrix} 0.333 & 0.333 & 0.333 \\ 0.333 & 0.333 & 0.333 \\ 0.262 & 0.262 & 0.476 \end{bmatrix} \begin{bmatrix} 1 & 2 & 3 & 4 \\ 5 & 6 & 7 & 8 \\ 9 & 10 & 11 & 12 \end{bmatrix}$

$= \begin{bmatrix} 5.0 & 6.0 & 7.0 & 8.0 \\ 5.0 & 6.0 & 7.0 & 8.0 \\ 5.86 & 6.86 & 7.86 & 8.86 \end{bmatrix}$

解读：词 1 和词 2 均匀关注所有词（权重相等），词 3 更偏向自己（权重 0.476 > 0.262）。