Transformer 模型中的 QKV 机制是如何运作的

通过一个简化的例子来展示 Transformer 模型中的 QKV 机制是如何运作的。假设我们正在处理一个包含两个词“你好”和“世界”的序列，并且为了简单起见，我们使用一个非常小的嵌入维度来说明这个概念。

步骤 1: 嵌入和投影到 QKV 空间

首先，我们将输入的词转换为嵌入向量。假设“你好”和“世界”的嵌入向量如下：

• Embedding(你好) = [1, 2, 3]
• Embedding(世界) = [4, 5, 6]
  然后，这些嵌入向量被投影到 QKV 空间，使用三个不同的投影矩阵（W_Q, W_K, W_V）。为了简化，我们假设这些矩阵是单位矩阵，所以投影后的向量就是嵌入向量本身：
• Q(你好) = [1, 2, 3]
• K(你好) = [1, 2, 3]
• V(你好) = [1, 2, 3]
• Q(世界) = [4, 5, 6]
• K(世界) = [4, 5, 6]
• V(世界) = [4, 5, 6]

步骤 2: 计算注意力分数

接下来，我们通过将每个查询向量与所有键向量进行点积来计算注意力分数：

• “你好”与“你好”的注意力分数：Q(你好) · K(你好) = 1✖1 + 2✖2 + 3✖3 = 14
• “你好”与“世界”的注意力分数：Q(你好) · K(世界) = 1✖4 + 2✖5 + 3✖6 = 32
• “世界”与“你好”的注意力分数：Q(世界) · K(你好) = 4✖1 + 5✖2 + 6✖3 = 32
• “世界”与“世界”的注意力分数：Q(世界) · K(世界) = 4✖4 + 5✖5 + 6✖6 = 77

步骤 3: 应用 Softmax

现在我们对注意力分数应用 softmax 函数来获得注意力概率：

• “你好”的注意力概率：[e^14 / (e^14 + e^32), e^32 / (e^14 + e^32)] ≈ [0.0024, 0.9976]
• “世界”的注意力概率：[e^32 / (e^32 + e^77), e^77 / (e^32 + e^77)] ≈ [0.0024, 0.9976]

步骤 4: 计算输出

最后，我们通过将值向量与注意力概率加权求和来计算自注意力层的输出：

• “你好”的输出：0.0024 * V(你好) + 0.9976 * V(世界) ≈ 0.0024*[1, 2, 3] + 0.9976*[4, 5, 6] ≈ [3.9952, 4.9960, 5.9968]
• “世界”的输出：0.0024 * V(你好) + 0.9976 * V(世界) ≈ 0.0024*[1, 2, 3] + 0.9976*[4, 5, 6] ≈ [3.9952, 4.9960, 5.9968]
  在这个例子中，“你好”和“世界”的输出几乎是相同的，因为注意力机制已经确定在这个上下文中“世界”更为相关，因此赋予了它更高的权重。
  在实际应用中，嵌入向量会大得多（通常是 512 或 1024 维度），投影矩阵（W_Q, W_K, W_V）会在训练过程中学习得到。上面的例子简化了过程，以展示 QKV 机制的基本运作方式。