【深度学习】实验 — 动手实现 GPT【二】：注意力机制、注意力掩码、多头注意力机制

注意力机制

简单示例：单个元素的情况

• 假设我们有以下输入句子，已按照第 3 章中的描述嵌入为 3 维向量（此处使用非常小的嵌入维度，仅用于说明，方便在页面上显示而不换行）：

inputs = torch.tensor(
  [[0.43, 0.15, 0.89], # Your     (x^1)
   [0.55, 0.87, 0.66], # journey  (x^2)
   [0.57, 0.85, 0.64], # starts   (x^3)
   [0.22, 0.58, 0.33], # with     (x^4)
   [0.77, 0.25, 0.10], # one      (x^5)
   [0.05, 0.80, 0.55]] # step     (x^6)
)

• （在本书中，我们遵循机器学习和深度学习的常见惯例，即训练样本表示为行，特征值表示为列；在上面的张量中，每一行表示一个词，每一列表示一个嵌入维度。）

• 本节的主要目的是演示如何使用第二个输入序列 $x^{(2)}$ 作为查询，计算上下文向量 $z^{(2)}$。

• 图示展示了该过程的初始步骤，其中通过点积操作计算 $x^{(2)}$ 与所有其他输入元素之间的注意力分数 ω。

• 我们使用输入序列中的元素 2，即 $x^{(2)}$，作为示例来计算上下文向量 $z^{(2)}$；在本节稍后，我们将推广此方法来计算所有的上下文向量。
• 第一步是通过计算查询 $x^{(2)}$ 与所有其他输入词元之间的点积，得到未归一化的注意力分数：

query = inputs[1]  # 2nd input token is the query

attn_scores_2 = torch.empty(inputs.shape[0])
for i, x_i in enumerate(inputs):
    attn_scores_2[i] = torch.dot(x_i, query) # dot product (transpose not necessary here since they are 1-dim vectors)

print(attn_scores_2)

输出

tensor([0.9544, 1.4950, 1.4754, 0.8434, 0.7070, 1.0865])

• 步骤 2： 将未归一化的注意力分数（“omegas”，$\omega$）归一化，使其总和为 1。
• 以下是一种简单的归一化方法，使未归一化的注意力分数总和为 1（这种方式是约定俗成的，有助于解释，并对训练稳定性非常重要）：

attn_weights_2_tmp = attn_scores_2 / attn_scores_2.sum()

print("Attention weights:", attn_weights_2_tmp)
print("Sum:", attn_weights_2_tmp.sum())

输出

Attention weights: tensor([0.1455, 0.2278, 0.2249, 0.1285, 0.1077, 0.1656])
Sum: tensor(1.0000)

• 然而，在实际操作中，通常推荐使用 softmax 函数进行归一化，因为它在处理极端值方面更有效，并且在训练过程中具有更理想的梯度特性。

attn_weights_2 = torch.softmax(attn_scores_2, dim=0)

print("Attention weights:", attn_weights_2)
print("Sum:", attn_weights_2.sum())

输出

Attention weights: tensor([0.1385, 0.2379, 0.2333, 0.1240, 0.1082, 0.1581])
Sum: tensor(1.)

• 步骤 3：通过将嵌入的输入词元 $x^{(i}$