2、自注意力机制回顾：QKV矩阵计算、注意力分数计算、Softmax与加权求和

好，咱们正式开始聊自注意力机制。说实话，这个机制是整个Transformer的基石，也是KV Cache能发挥作用的前提。如果你对这部分已经烂熟于心，可以快速过一遍；但如果你还有点模糊，我建议你跟着我的思路走一遍。

自注意力机制，说白了就是让模型在生成每个词的时候，能“回头看”一下已经处理过的所有词。它不像RNN那样一步步传递隐状态，而是直接计算当前位置和所有位置之间的关联度。嗯，这个设计很暴力，也很优雅。

2.1 QKV矩阵计算：从输入到查询、键、值

我们先从最基础的开始。假设输入是一个序列，每个位置有一个向量表示。比如输入序列长度是 \( L \)，每个词的维度是 \( d_{model} \)。

那么，Q、K、V是怎么来的？

我个人习惯把这一步理解为“三个不同的投影”。输入向量分别乘以三个不同的权重矩阵：

• Q（Query）：查询向量，代表“当前词想找谁”
• K（Key）：键向量，代表“当前词能被谁找到”
• V（Value）：值向量，代表“当前词真正要贡献的信息”

用公式表示就是：

Q = X * W_Q
K = X * W_K
V = X * W_V

其中 \( X \) 是输入矩阵（形状 \( L \times d_{model} \)），\( W_Q, W_K, W_V \) 都是可训练的权重矩阵（形状 \( d_{model} \times d_k \)）。注意，\( d_k \) 通常等于 \( d_{model} / h \)，其中 \( h \) 是注意力头的数量。

我在项目中遇到过一个问题：刚开始做多头注意力时，我忘了把Q、K、V的维度除以头数，结果显存直接爆了。你想想看，如果 \( d_{model}=1024 \)，头数 \( h=16 \)，那么每个头的维度应该是64。如果不分头，直接做1024维的注意力，计算量会大16倍。嗯，这个坑我踩过，你别再踩了。

2.2 注意力分数计算：点积与缩放

有了Q、K、V之后，下一步就是计算注意力分数。这一步的核心是：用Q去“匹配”所有的K。

具体来说，对于第 \( i \) 个位置的查询向量 \( q_i \)，我们计算它与所有位置键向量 \( k_j \) 的点积：

score(i, j) = q_i · k_j

把所有分数拼起来，就得到了一个 \( L \times L \) 的注意力分数矩阵。这个矩阵的第 \( i \) 行第 \( j \) 列，代表位置 \( i \) 对位置 \( j \) 的关注程度。

但是，这里有个问题。如果 \( d_k \) 比较大，点积的结果会很大。比如 \( d_k=64 \) 时，点积的方差大约是64。这会导致softmax之后，概率分布过于尖锐（接近one-hot），梯度容易消失。

解决办法很简单——缩放。把点积结果除以 \( \sqrt{d_k} \)：

score(i, j) = (q_i · k_j) / sqrt(d_k)

这就是“Scaled Dot-Product Attention”名字的由来。我建议你把这个缩放因子记牢，因为很多面试官会问“为什么除以根号d_k”。

scores = torch.matmul(Q, K.transpose(-2, -1)) / math.sqrt(d_k)

2.3 Softmax与加权求和：注意力机制的输出

分数算完了，接下来就是两步：归一化和加权求和。

第一步：Softmax归一化

对每个位置 \( i \)，把它的分数向量（即分数矩阵的第 \( i \) 行）做softmax，得到注意力权重：

attention_weights(i, :) = softmax(scores(i, :))

softmax的作用是让所有权重之和为1，并且大的分数变得更大，小的变得更小。嗯，这里要注意，softmax之前通常要加一个mask，把未来位置（在自回归生成中）的分数设为负无穷，这样softmax之后权重就是0。这就是所谓的“因果掩码”（causal mask）。

第二步：加权求和

有了注意力权重之后，用它们对V进行加权求和：

output(i) = sum_j attention_weights(i, j) * V(j)

这一步的物理意义很直观：每个位置的输出，是所有位置的值向量的加权平均，权重由注意力分数决定。如果位置 \( i \) 对位置 \( j \) 的关注度很高，那么 \( V(j) \) 就会更多地贡献到 \( output(i) \) 中。

用矩阵形式表示就是：

Attention(Q, K, V) = softmax(Q * K^T / sqrt(d_k)) * V

这就是完整的自注意力计算公式。简洁、优雅、强大。

2.4 完整流程的SVG示意图

下面我用一张图把整个流程串起来。这张图是我自己画的，你可以把它当作“自注意力机制速查表”。

2.5 为什么KV Cache能从这里切入？

好了，现在自注意力的完整流程我们已经过了一遍。你可能会问：这跟KV Cache有什么关系？

关系大了。你想想看，在自回归生成中，我们是一个词一个词地生成的。假设已经生成了 \( t \) 个词，现在要生成第 \( t+1 \) 个词。按照标准流程，我们需要重新计算所有 \( t+1 \) 个位置的Q、K、V，然后做注意力。

但仔细想想，前 \( t \) 个词的K和V其实已经算过了。每次生成新词时，它们并没有变化。那为什么还要重复计算呢？

这就是KV Cache的切入点——把之前算好的K和V缓存起来，每次只计算新词的Q、K、V，然后用新词的Q去匹配所有缓存的K，最后用缓存的V做加权求和。

嗯，这个优化思路很直接，但效果惊人。我后面会详细展开，这里先留个印象。

好，自注意力机制的回顾就到这里。记住三个核心步骤：QKV投影、缩放点积注意力、softmax加权求和。这是后面所有内容的基础。