Attention 机制与 QKV 矩阵

Attention 机制与 QKV 矩阵的手推过程与直觉解释。

从零手推:大语言模型核心——Attention 机制与 QKV 矩阵

前置知识: 序列模型、线性层与矩阵乘法

引言:为什么我们需要 Attention?

在学习《序列模型》(如 RNN / LSTM)时,我们知道处理“我 爱 你”这句话,模型是按顺序吃的:先吃“我”,再吃“爱”,最后吃“你”。 RNN 的致命弱点:

  1. 遗忘: 句子一长,看到后面就忘了前面。
  2. 慢: 必须串行计算,无法发挥 GPU 强大的并行矩阵运算能力。

2017 年,Google 提出了 Self-Attention(自注意力机制),彻底掀翻了 RNN。它的核心思想极其霸道:不排队了,句子里的所有词同时互相看,一次性算清楚谁和谁关系最铁!

下面,我们就用“我 爱 你”这三个词,手工推演这个“互相看”的过程。

第一步:词的数字化与 QKV 的诞生 (回顾)

假设我们的词向量维度 $D=4$,序列长度 $L=3$。 输入矩阵 $X$(包含了“我”、“爱”、“你”三个词的特征):

$$X = \begin{bmatrix} 1 & 0 & 1 & 0 \ 0 & 1 & 0 & 1 \ 1 & 1 & 0 & 0 \end{bmatrix} \quad \text{(维度: } 3 \times 4\text{)}$$

输入矩阵 $X$ 分别乘以三个权重矩阵 $W_Q, W_K, W_V$(这三个矩阵是模型在训练中学习到的,维度设为 $4 \times 4$),我们得到了 Q(查询)、K(键/标签)、V(值/本体)三个矩阵。

这里假设算出来的 $Q$ 和 $K$ 矩阵如下:

$$Q = \begin{bmatrix} 1 & 0 & 1 & 1 \ 0 & 1 & 0 & 1 \ 1 & 1 & 0 & 0 \end{bmatrix} \begin{matrix} \leftarrow \text{我的查询} \ \leftarrow \text{爱的查询} \ \leftarrow \text{你的查询} \end{matrix}$$

$$K = \begin{bmatrix} 1 & 0 & 0 & 1 \ 0 & 1 & 1 & 0 \ 1 & 1 & 0 & 0 \end{bmatrix} \begin{matrix} \leftarrow \text{我的标签} \ \leftarrow \text{爱的标签} \ \leftarrow \text{你的标签} \end{matrix}$$

第二步:灵魂碰撞 —— 计算打分矩阵 ($Q \times K^T$)

Attention 的核心思想是:拿每个词的 Query,去和其他所有词的 Key 算相似度。 在数学中,计算两个向量相似度最直接的方法就是点积(Dot Product)。点积越大,说明两个向量越相似,注意力应该分配得越多。

为了一次性算出所有词两两之间的相似度,我们用矩阵乘法:$Q$ 乘以 $K$ 的转置 ($K^T$)

首先,把 $K$ 躺平变成 $K^T$(维度变为 $4 \times 3$):

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

现在计算 $Scores = Q \times K^T$:

$$\begin{bmatrix} 1 & 0 & 1 & 1 \ 0 & 1 & 0 & 1 \ 1 & 1 & 0 & 0 \end{bmatrix}{3 \times 4} \times \begin{bmatrix} 1 & 0 & 1 \ 0 & 1 & 1 \ 0 & 1 & 0 \ 1 & 0 & 0 \end{bmatrix}{4 \times 3} = \begin{bmatrix} 2 & 1 & 1 \ 1 & 1 & 1 \ 1 & 1 & 2 \end{bmatrix}_{3 \times 3}$$

💡 极其重要的物理意义解析(面试必考): 看看算出来的这个 $3 \times 3$ 的矩阵,这就是注意力打分矩阵(Attention Score Matrix)

  • 第一行 [2, 1, 1]:代表“我”这个词的 Q,分别去匹配“我”、“爱”、“你”的 K。发现“我”对“我”自己的注意力得分是 2,对“爱”和“你”是 1。
  • 第二行 [1, 1, 1]:代表“爱”去匹配所有词的得分。
  • 第三行 [1, 1, 2]:代表“你”对“我”、“爱”、“你”的打分。

到这一步,模型已经算出了任意两个词之间的关联程度。没有任何先后顺序,全部是一次矩阵乘法完成的。这就是 Transformer 比 RNN 快的根本原因!

第三步:缩放与概率化 (Scale & Softmax)

得出了打分矩阵,我们还要做两层处理。

1. 缩放 (Scale):除以 $\sqrt{d_k}$

公式里有一步是 $\frac{Q K^T}{\sqrt{d_k}}$。这里的 $d_k$ 就是词向量的维度(我们这里是 4,所以除以 $\sqrt{4} = 2$)。

🔥为什么一定要除以 $\sqrt{d_k}$?不除会怎样? 标准答案: 为了防止梯度消失。当词向量维度 $d_k$ 很大时(比如 768),点积算出来的数值会非常大。如果直接把这么大的数值扔给后面的 Softmax 函数,Softmax 会被“撑爆”(推向饱和区),导致输出的概率变成一个极端的 One-hot 向量(如 [0.999, 0.001, 0]),这时候梯度接近于 0,模型参数就无法更新学习了。除以根号 $d_k$ 是一种方差控制手段,把数值拉回到合理的范围内。

在我们这个简化的例子里,除以 2 之后,打分矩阵变成了:

$$\text{Scaled Scores} = \begin{bmatrix} 1.0 & 0.5 & 0.5 \ 0.5 & 0.5 & 0.5 \ 0.5 & 0.5 & 1.0 \end{bmatrix}$$

2. 概率化 (Softmax)

我们要把分数变成加起来等于 1 的百分比权重。沿着每一行做 Softmax。 (为了直观,这里直接给出近似的概率结果)

$$\text{Attention Weights} = \text{Softmax}(\text{Scaled Scores}) \approx \begin{bmatrix} 0.45 & 0.27 & 0.27 \ 0.33 & 0.33 & 0.33 \ 0.27 & 0.27 & 0.45 \end{bmatrix}$$

看第一行:对于“我”这个词,它分配了 45% 的注意力给自己,27% 给“爱”,27% 给“你”。

第四步:提取精华 —— 乘以 $V$ 矩阵

现在我们有了每个词对其他词的“关注度百分比”,最后一步,就是拿这个百分比,去把 $V$ 矩阵(包含真实内容)里的特征按比例提取出来。

最终结果 = $Attention Weights \times V$

$$\begin{bmatrix} 0.45 & 0.27 & 0.27 \ 0.33 & 0.33 & 0.33 \ 0.27 & 0.27 & 0.45 \end{bmatrix}{3 \times 3} \times \begin{bmatrix} v{我} \ v_{爱} \ v_{你} \end{bmatrix}_{3 \times 4}$$

以第一行(更新后的“我”的特征)为例: 新“我” = $0.45 \times v_{我} + 0.27 \times v_{爱} + 0.27 \times v_{你}$

结论: 经过这套 Attention 运算后,输出的仍然是一个 $3 \times 4$ 的矩阵。但是,此时的“我”,已经不再是原本孤立的“我”了。它融合了 27% “爱”的语境和 27% “你”的语境。词语之间产生了灵魂的交融,模型就这样“理解”了上下文!

下一篇预告: 当有了 Attention 之后,为什么原版代码里还需要加上位置编码 (Rotary Embedding) 和 GQA 机制?我们将在下篇继续手撕大模型源码!

留言讨论

0 条留言

正在加载留言...