从零训练GPT,《Let's build GPT》-Andrej Karpathy

Notes

偶然发现一个数字花园，感觉作者挺实在的。Maeiee 的 Digital Garden

• 《Let’s build GPT：from scratch, in code, spelled out.》— Andrej KarpathyMark
• 《Let’s build GPT》 和 3B1B 的 《Visualizing Attention》

单头注意力

现在，我们了解了多头注意力机制，但是对于单头注意力还是一头雾水，要知道，这是 Transformer 模型最精华的部分。

从 Transformer 模型提出的论文《Attention is All You Need》的标题即可看出，Attention 注意力机制的关键性。

注意力机制的灵感来自于人类的视觉注意力机制。当我们观察一幅图像或一个场景时，我们的视觉系统并不是平等地处理每一个细节，而是会有选择性地关注某些部分。例如，当我们看一张人脸的照片时，我们通常会更多地关注眼睛、鼻子、嘴巴等面部特征，而相对忽略其他次要的细节。这种选择性关注的机制使我们能够高效地提取关键信息，而不被大量无关的细节所淹没。

在自然语言处理任务中，注意力机制的作用与之类似。当我们处理一个句子或一段文本时，其中某些词或短语通常比其他部分更重要，包含了更多的信息。注意力机制允许模型去学习如何区分重要的信息和次要的信息，并根据这些重要的信息来做出判断或预测。

具体来说，注意力机制的核心思想可以用”查询-键-值（Query-Key-Value）“的框架来描述:

• 查询（Query）：我们想要关注的内容。在自然语言处理任务中，查询通常是我们当前正在处理的词或句子。
• 键（Key）：我们用来判断其他信息是否重要的参考。键可以是句子中的其他词，或者是来自其他来源的信息。
• 值（Value）:我们想要提取的信息。值通常与键是一一对应的。

注意力机制的过程可以概括为：对于每个查询，我们用它去和所有的键进行比较，计算出每个键与查询的相关性或重要性。然后，我们用这些重要性作为权重，对相应的值进行加权求和，得到最终的注意力结果。这个结果就是模型认为对于当前的查询最重要的信息。

通过这种方式，注意力机制使模型能够动态地调整对不同部分信息的关注，从而更好地理解和处理复杂的语言数据。

在 Transformer 模型中，注意力机制被用于捕捉句子内部和句子之间的依赖关系。例如,当模型处理一个代词时（如”it”）,注意力机制允许模型去关注前面提到的相关名词,从而正确地理解代词的指代对象。

---

单头注意力的数学推导

注意力机制的核心思想是，当查询（query）某个值（value）时，通过与一系列键（key）进行相似度计算来得到每个值的重要性权重，然后用这些权重对相应的值进行加权求和，得到最终的注意力结果。这个过程可以用数学公式表示为:

其中， 表示查询（query）， 表示键（key）， 表示值（value）， 是键向量的维度。

在单头注意力中，Q、K、V 实际上都是矩阵，而不是向量。它们的形状都是（B, T, head_size），其中 B 是 batch size,T 是序列长度，head_size 是每个注意力头的维度。

这里的 head_size 与嵌入维度 n_embd 有关。具体来说，在多头注意力中，我们将嵌入维度 n_embd 均分给每个注意力头。所以，每个头的维度 head_size = n_embd // num_heads。

Q、K、V 矩阵是从输入序列的词嵌入矩阵 X 计算得到的。设输入序列的词嵌入矩阵为 X,其形状为 (B, T, n_embd),即对于每个 batch 中的每个序列，我们都有 T 个 n_embd 维的词嵌入向量。

在自注意力机制中，查询、键、值都来自同一个输入序列。具体来说，对于输入序列的每个位置，们使用三个学习到的矩阵 、、 将其映射为查询向量、键向量和值向量:

其中，表示输入序列的嵌入表示。

换算成更加熟悉的代码表示：

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

其中，W_Q、W_K、W_V 的形状都是 (n_embd, head_size)。这个映射过程可以理解为，我们为每个注意力头学习一组特定的查询、键、值表示。

所以，Q、K、V 矩阵可以看作是从原始词嵌入空间映射到注意力头特定的查询、键、值空间。这种映射允许每个注意力头关注输入序列的不同方面。

在实际实现中，我们通常会一次性计算整个序列的注意力。这时，、、 都是矩阵，注意力计算可以用矩阵乘法高效地实现:

其中，、、 的形状分别为、、，注意力输出的形状为。

对于每个查询向量（Q 矩阵中的每一行），我们计算它与所有键向量（K 矩阵中的所有行）的相似度,得到一个注意力分数向量。这个过程可以用矩阵乘法高效地实现，换成代码表示:

scores = Q @ K.transpose(-1, -2)

这里的 scores 矩阵形状为 (B, T, T)，表示对于每个 batch 中的每个查询位置，我们都计算了它与所有键位置的相似度。

然后，我们对这些分数应用 softmax，得到注意力权重矩阵:

weights = softmax(scores, dim=-1)

最后，我们用这些权重对值矩阵 V 进行加权求和:

output = weights @ V

这里的 output 形状为 (B, T, head_size)，表示对于每个 batch 中的每个位置，我们都计算了一个 head_size 维的注意力输出向量。

计算注意力分数的过程可以解释为，对于每个查询向量，我们计算它与所有键向量的点积相似度，然后除以 进行缩放(这是为了防止点积结果过大)。接着，我们对这些相似度进行 softmax 归一化，得到注意力权重。最后，我们用这些权重对值向量进行加权求和，得到注意力输出。