Build a Large Language Model (From Scratch) 第三章 编码注意力机制

本章内容

  • 探索在神经网络中使用注意力机制的原因
  • 介绍一个基础的自注意力框架,并进一步探讨一种增强型自注意力机制
  • 实现一个因果注意力模块,使大语言模型能够一次生成一个词元
  • 使用 dropout 随机掩码部分注意力权重来降低过拟合风险
  • 将多个因果注意力模块堆叠成一个多头注意力模块

到目前为止,你已经知道如何将文本分割成单词和子词,并将这些词元编码成向量表示(嵌入),从而为大语言模型提供输入数据。

本章将探讨大语言模型架构中的一个核心部分------注意力(attention)机制,如图 3-1 所示。 我们将大篇幅单独探讨注意力机制,并专注于其工作原理。然后,我们将编写自注意力机制的其他部分代码,观察其运行过程,并构建一个生成文本的模型。

本章将实现 4 种注意力机制的变体,如图 3-2 所示。这些变体中的每一个都是在前一个的基础上逐步建立的,最终目的是实现一种紧凑、高效的多头注意力机制,并在第 4 章中将其集成到大语言模型的架构中。

3.1 长序列建模中的问题

深入探讨大语言模型核心的自注意力机制之前,让我们考虑一下在大语言模型出现之前的没有注意力机制的架构中所存在的问题。假设我们想要开发一个将文本从一种语言翻译成另一种语言的语言翻译模型。如图 3-3 所示,由于源语言和目标语言的语法结构不同,我们无法简单地逐个单词进行翻译。

为了处理这个问题,通常使用一个包含编码器和解码器两个子模块的深度神经网络。编码器首先读取和处理整个文本,解码器则负责生成翻译后的文本。

在 Transformer 出现之前,循环神经网络(recurrent neural network,RNN)是语言翻译中最流行的编码器-解码器架构。RNN 是一种将前一步骤的输出作为当前步骤的输入的神经网络,它非常适合处理像文本这样的序列数据。如果你不熟悉 RNN 也没关系,因为理解这次讨论并不需要了解 RNN 的具体工作原理。这里我们主要关注编码器-解码器架构的基本概念。

在编码器-解码器 RNN 中,输入文本被传递给编码器以逐步处理。编码器在每一步都会更新其隐藏状态(隐藏层的内部值),试图在最终的隐藏状态中捕捉输入句子的全部含义,如图 3-4 所示。然后,解码器使用这个最终的隐藏状态开始逐字生成翻译后的句子。解码器同样在每一步更新其隐藏状态,该状态应包含为下一单词预测所需的上下文信息。

这个图很多细节没有说清楚,箭头有点乱,让人迷惑

I need to refine the diagram, particularly the feedback loop, keep the existing explanation, add a section detailing how RNNs work internally, and format everything as a markdown blockquote note rather than using headers.

RNN 编码器-解码器 翻译笔记(德语 → 英语)

任务:把 "Kannst du mir helfen?" 翻译成 "Can you help me?"


编码器(读德语,压缩成一个记忆向量)

复制代码
输入词:   Kannst      du         mir        helfen
            │          │          │           │
            ▼          ▼          ▼           ▼
         ┌─────┐    ┌─────┐    ┌─────┐    ┌─────┐
 h₀ ────▶│ RNN │─h₁▶│ RNN │─h₂▶│ RNN │─h₃▶│ RNN │───▶ h₄
         └─────┘    └─────┘    └─────┘    └─────┘
                                               │
                                               ▼
                                    h₄ = 整句话的"总结记忆"
                                    (context vector 上下文向量)
  • 每个方块是同一个 RNN 单元,按时间步一次处理一个词。
  • 横向的 h₁ h₂ h₃ 就是隐藏状态,像接力棒一样往后传。
  • 处理完最后一个词得到 h₄,它浓缩了整句德语的信息,交给解码器。

解码器(生成英语,关键在"把上一个词喂回去")

复制代码
           <start>      "Can"       "you"       "help"
             │            │           │           │
             │       ┌────┘      ┌────┘      ┌────┘
             │       │ 喂回      │ 喂回      │ 喂回
             ▼       ▼           ▼           ▼
  h₄ ─────┌─────┐ ┌─────┐    ┌─────┐    ┌─────┐
   起点   │ RNN │─│ RNN │─s₂─│ RNN │─s₃─│ RNN │──▶ ...
          └─────┘ └─────┘    └─────┘    └─────┘
             │       │           │           │
             ▼       ▼           ▼           ▼
           "Can"   "you"       "help"       "me"

每一步解码器有两个输入:横向的隐藏状态 s(内部记忆)+ 上一步刚生成的词(往回喂)。

一步一步拆解:

  • 第 1 步:输入 = 编码器记忆 h₄ + 起始标记 <start> → 算出隐藏状态 s₁ → 输出 "Can"
  • 第 2 步:输入 = s₁ + 刚生成的 "Can" → 算出 s₂ → 输出 "you"
  • 第 3 步:输入 = s₂ + 刚生成的 "you" → 算出 s₃ → 输出 "help"
  • ...... 循环,直到输出结束标记 <end>

为什么要"把 Can 喂回去":下一个词是什么,取决于上一个词说了什么。已经说了 "Can",模型才知道后面大概率接主语 "you";说完 "you" 才知道该接动词 "help"。不告诉它上一个词,它就"断片"了。这种"自己的输出变成自己下一步输入"的做法叫自回归(autoregressive)。

两条信息流对比:

信息流 是什么 作用
隐藏状态 s(横向 →) 内部记忆 记住整体进度和上下文
上一个输出词(往回喂 ↑) 具体的词 告诉模型"我刚说完哪个词"

RNN 单元内部长什么样(把上面每个方块拆开看)

复制代码
         隐藏状态 hₜ₋₁ (上一时刻的记忆)
                 │
                 │      当前输入 xₜ (这一步的词)
                 │            │
                 ▼            ▼
             ┌───────────────────┐
             │   拼在一起后做:     │
             │   hₜ = tanh(       │
             │      W·[hₜ₋₁, xₜ]  │   ← W 是要学习的权重
             │      + b )         │
             └───────────────────┘
                 │
                 ├──────────────▶ hₜ  (传给下一个时刻,横向接力棒)
                 │
                 ▼
             输出 yₜ (可选,比如预测的词)
  • 输入两样东西:上一时刻的隐藏状态 hₜ₋₁ + 当前这一步的词向量 xₜ。
  • 把这两个拼接起来,乘上权重矩阵 W、加偏置 b,再过一个 tanh 激活函数(把数值压到 -1~1 之间)。
  • 得到新的隐藏状态 hₜ,它同时干两件事:往横向传给下一步、往上产生这一步的输出。
  • 关键:整个序列自始至终用的是同一组权重 W(参数共享),所以不管句子多长,RNN 的参数量不变。

用一句话串起来:RNN 就是一个"带记忆的复读机"------每读一个词,就把新词和旧记忆揉在一起更新出新记忆,再往下传。编码器靠它读完全句得到 h₄,解码器靠它一边记进度一边逐词吐译文。


这个结构的弱点:整句话全压在 h₄ 一个向量里,句子一长就记不全。这正是后来 Attention 机制要解决的问题------让解码器每生成一个词都能回头看编码器的所有隐藏状态(h₁~h₄),而不是只靠一个 h₄。

虽然我们不需要了解这些编码器-解码器 RNN 的内部工作原理,但关键点在于,编码器部分会将整个输入文本处理成一个隐藏状态(记忆单元)。然后解码器会使用这个隐藏状态来生成输出。你可以将这个隐藏状态视为一种嵌入向量(第 2 章中讨论过的一个概念)。

编码器-解码器 RNN 的一个主要限制是,在解码阶段,RNN 无法直接访问编码器中的早期隐藏状态。因此,它只能依赖当前的隐藏状态,这个状态包含了所有相关信息。这可能导致上下文丢失,特别是在复杂句子中,依赖关系可能跨越较长的距离。

幸运的是,构建大语言模型不需要深入了解 RNN。只需记住,编码器-解码器 RNN 存在的缺陷对注意力机制的设计起到了促进作用。

3.2 使用注意力机制捕捉数据依赖关系

尽管 RNN 在翻译短句时表现良好,但在处理较长文本时效果不佳,因为它无法直接访问输入中靠前的单词。这种方法的一个主要缺点是,RNN 在将信息传递给解码器之前,必须将整个编码后的输入存储到一个单独隐藏状态中(参见图 3-4)。

因此,研究人员在 2014 年为 RNN 开发了 Bahdanau 注意力机制(以该研究论文的第一作者命名,更多信息请参见附录 B),该机制对编码器-解码器 RNN 进行了修改,使得解码器在每个解码步骤中可以选择性地访问输入序列的不同部分,如图 3-5 所示。

有趣的是,仅仅 3 年后,研究人员发现 RNN 并不是构建自然语言处理深度神经网络的必需架构,并提出了最初的 Transformer 架构(在第 1 章中讨论过),其中包括受 Bahdanau 注意力机制启发的自注意力机制。

在计算序列表示时,自注意力机制允许输入序列中的每个位置关注同一序列中的所有位置。 自注意力机制是基于 Transformer 架构的当代大语言模型(如 GPT 系列)的关键组成部分。

本章的重点是编码并理解类 GPT 模型中使用的自注意力机制,如图 3-6 所示。在第 4 章中, 我们将对大语言模型的其余部分进行编码。

3.3 通过自注意力机制关注输入的不同部分

现在我们将深入探讨自注意力机制的内部工作原理,并学习如何从头开始编写它的代码。自 注意力机制是所有基于 Transformer 架构的大语言模型的基石。值得注意的是,这一主题可能需 要大量的关注与注意力(不含双关意图),但一旦掌握了其基本原理,你就能攻克本书以及大语 言模型实现过程中最具挑战性的部分之一。

自注意力机制中的"自"

在自注意力机制中,"自"指的是该机制通过关联单个输入序列中的不同位置来计算注意 力权重的能力 。它可以评估并学习输入本身各个部分之间的关系和依赖,比如句子中的单词或图像中的像素。

这与传统的注意力机制形成对比。传统的注意力机制关注的是两个不同序列元素之间的 关系,比如在序列到序列模型中,注意力可能在输入序列和输出序列之间,如图3-5中的示例所示。

考虑到自注意力机制可能会显得复杂,特别是对于第一次接触它的人,因此本章将从一个简化版本入手。接下来,我们会实现在大语言模型中使用的带可训练权重的自注意力机制。

3.3.1 没有可训练权重的简单自注意力机制

让我们首先实现一个不包含任何可训练权重的简化的自注意力机制变体,如图 3-7 所示。目标是在引入可训练权重之前,阐明自注意力中的一些关键概念。

图 3-7 显示了一个输入序列,记为 xxx,它由 TTT 个元素组成,分别表示为 x(1)x^{(1)}x(1) 到 x(T)x^{(T)}x(T)。这个序列通常代表文本(如一个句子),并且该文本已经被转换为词元嵌入。

例如,考虑输入文本 "Your journey starts with one step."。在这种情况下,文本序列中的每个元素(如 x(1)x^{(1)}x(1))都对应一个 ddd 维的嵌入向量,该向量代表了一个特定的词元,比如 "Your"。在图 3-7 中,这些输入向量被表示为三维嵌入。

在自注意力机制中,我们的目标是为输入序列中的每个元素 x(i)x^{(i)}x(i) 计算上下文向量 z(i)z^{(i)}z(i)。上下文向量(context vector)可以被理解为一种包含了序列中所有元素信息的嵌入向量。

为了说明这个概念,我们重点关注第二个输入元素 x(2)x^{(2)}x(2)(对应于词元 "journey")的嵌入向量及其对应的上下文向量 z(2)z^{(2)}z(2),如图 3-7 底部所示。这个增强的上下文向量 z(2)z^{(2)}z(2) 是一个嵌入,包含了关于 x(2)x^{(2)}x(2) 及其他所有输入元素(x(1)x^{(1)}x(1) 到 x(T)x^{(T)}x(T))的信息。

上下文向量在自注意力机制中起着关键作用。它们的目的是通过结合序列中其他所有元素的信息,为输入序列(如一个句子)中的每个元素创建丰富表示,如图 3-7 所示。这在大语言模型中至关重要,因为这些模型需要理解句子中单词之间的关系和相关性。稍后我们将添加可训练的权重,以帮助大语言模型学习如何构建这些上下文向量,使它们能用于生成下一个词元。但首先,我们将实现一个简化的自注意力机制,逐步计算这些权重和上下文向量。

考虑以下输入句子,该句已按照第 2 章中讨论的方式嵌入为三维向量。我们选择较小的嵌入维度进行说明,以确保内容能够在页面上显示而无须换行。

实现自注意力机制的第一步是计算中间值 www ,即所谓注意力分数,如图 3-8 所示。由于空间有限,图中显示的前几个输入张量的数值是截断后的版本,比如 0.87 被截断为 0.8。在这个截断版本中,单词"journey"和"starts"的嵌入可能会由于随机原因看起来相似。

图 3-8 展示了如何计算查询词元与每个输入词元之间的中间注意力分数。我们通过计算查询 2 词元 x(2)x^{(2)}x(2) 与其他所有输入词元的点积来确定这些分数:

python 复制代码
import torch

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)
)

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

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

print("attention score_2:", attn_scores_2)
复制代码
tensor([0.4300, 0.1500, 0.8900]) tensor([0.5500, 0.8700, 0.6600])
tensor([0.5500, 0.8700, 0.6600]) tensor([0.5500, 0.8700, 0.6600])
tensor([0.5700, 0.8500, 0.6400]) tensor([0.5500, 0.8700, 0.6600])
tensor([0.2200, 0.5800, 0.3300]) tensor([0.5500, 0.8700, 0.6600])
tensor([0.7700, 0.2500, 0.1000]) tensor([0.5500, 0.8700, 0.6600])
tensor([0.0500, 0.8000, 0.5500]) tensor([0.5500, 0.8700, 0.6600])
attention score_2: tensor([0.9544, 1.4950, 1.4754, 0.8434, 0.7070, 1.0865])

理解点积

点积本质上是将两个向量逐个元素相乘然后对乘积求和的简洁方法,可以像下面这样进 行演示:

python 复制代码
res = 0.
for idx, element in enumerate(inputs[0]):
	res += inputs[0][idx] * query[idx] 
print(res) 
print(torch.dot(inputs[0], query))

输出证实了元素乘法的总和与点积的结果相同:

复制代码
tensor(0.9544) 
tensor(0.9544)

点积不仅被视为一种将两个向量转化为标量值的数学工具, 而且也是度量相似度的一种方式,因为它可以量化两个向量之间的对齐程度:点积越大,向量之间的对齐程度或相似度就越高。在自注意机制中,点积决定了序列中每个元素对其他元素的关注程度:点积越大,两个元素之间的相似度和注意力分数就越高

在下一步中,如图 3-9 所示,我们将对先前计算的每个注意力分数进行归一化处理。归一化的主要目的是获得总和为 1 的注意力权重。这种归一化是一个惯例,有助于解释结果,并能维持大语言模型的训练稳定性。以下是一种实现这一归一化步骤的简单方法。

python 复制代码
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 函数进行归一化更为常见,而且是一种更可取的做法。这种方法更好地处理了极值, 并在训练期间提供了更有利的梯度特性。 以下是用于归一化注意力分数的 softmax 函数的基础实现。如以下输出所示,softmax 函数同样达到了目标,成功地对注意力权重进行了归一化,使它们的总和为 1:

python 复制代码
def softmax_naive(x):
    return torch.exp(x) / torch.exp(x).sum(dim=0)

attn_weights_2_naive = softmax_naive(attn_scores_2)

print("Attention weights:", attn_weights_2_naive)
print("Sum:", attn_weights_2_naive.sum())
复制代码
Attention weights: tensor([0.1385, 0.2379, 0.2333, 0.1240, 0.1082, 0.1581])
Sum: tensor(1.)

sum(dim=0) 的含义是:让行索引这个维度"消失"。要消掉行索引,就得把不同行号的元素加到一起。于是:

结果的第 0 个 = 行0列0 + 行1列0

结果的第 1 个 = 行0列1 + 行1列1

然后形状是 (6,),只有一个维度 dim=0。

x.sum(dim=0) = 把这 6 个数全加起来 = 一个标量

另外,softmax 函数可以保证注意力权重总是正值,这使得输出可以被解释为概率或相对重要性,其中权重越高表示重要程度越高。

请注意,这种简单的 softmax 实现(softmax_naive)在处理大输入值或小输入值时可能会遇到数值稳定性问题,比如溢出和下溢。因此,在实践中建议使用 softmax 的 PyTorch 实现,该实现经过了大量性能优化:

通过减去最大值就可以解决大输出和小输出比如10000和-10000这种问题。

python 复制代码
def softmax_stable(x):
   x_max = x.max(dim=0, keepdim=True).values
   x_shifted = x - x_max                    # 先平移,最大值变成0
   exp_x = torch.exp(x_shifted)             # 每个都 <= 1,不会溢出
   return exp_x / exp_x.sum(dim=0)
python 复制代码
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-10 所示,通过将嵌入的输入词元 x(i)x^{(i)}x(i) 与相应的注意力权重相乘,再将得到的向量求和来计算上下文向量 z(2)z^{(2)}z(2) 。因 2 此,上下文向量 z(2)z^{(2)}z(2)是所有输入向量的加权总和,通过将每个输入向量与其对应的注意力权重相乘而获得:

python 复制代码
query = inputs[1] # 2nd input token is the query

context_vec_2 = torch.zeros(query.shape)
for i,x_i in enumerate(inputs):
    context_vec_2 += attn_weights_2[i]*x_i

print(context_vec_2)
复制代码
tensor([0.4419, 0.6515, 0.5683])

接下来,我们将推广这个计算上下文向量的过程,以便同时计算所有上下文向量。

3.3.2 计算所有输入词元的注意力权重

到目前为止,我们已经计算了输入 2 的注意力权重和上下文向量,如图 3-11 中突出显示的那一行所示。接下来,我们将扩展这个计算过程,以计算所有输入的注意力权重和上下文向量。

图中每个元素 aija_{ij}aij 的意义是:

计算第 iii 个词的新表示时,从第 jjj 个词获取信息的权重。

  • :当前需要更新的词(Query)
  • :被关注、提供信息的词(Key/Value)
  • 每一行的权重之和为 1

例如,journey 对应的注意力权重为:

来源词 Your journey starts with one step
权重 0.13 0.23 0.23 0.12 0.10 0.15

这表示生成 journey 的上下文表示时:

zjourney=0.13xYour+0.23xjourney+0.23xstarts+0.12xwith+0.10xone+0.15xstepz_{\text{journey}} =0.13x_{\text{Your}} +0.23x_{\text{journey}} +0.23x_{\text{starts}} +0.12x_{\text{with}} +0.10x_{\text{one}} +0.15x_{\text{step}}zjourney=0.13xYour+0.23xjourney+0.23xstarts+0.12xwith+0.10xone+0.15xstep

这么做的意义是:让每个词的表示不再只包含它自身的信息,而是根据注意力权重,有选择地融合句子中其他词的信息

因此,journey 更新后的向量不再只是孤立的 journey,而是一个结合了整句话上下文的表示。

一句话概括:每个元素表示"当前词应该从另一个词那里获取多少信息";整个注意力矩阵用于为每个词生成包含上下文的新表示。

如图 3-12 所示,我们遵循与之前相同的 3 个步骤,唯一的区别是在代码中进行了一些修改,以计算所有上下文向量,而不仅仅是第二个上下文向量 z(2)z^{(2)}z(2)

python 复制代码
attn_scores = torch.empty(6, 6)

for i, x_i in enumerate(inputs):
    for j, x_j in enumerate(inputs):
        attn_scores[i, j] = torch.dot(x_i, x_j)

print(attn_scores)
复制代码
tensor([[0.9995, 0.9544, 0.9422, 0.4753, 0.4576, 0.6310],
        [0.9544, 1.4950, 1.4754, 0.8434, 0.7070, 1.0865],
        [0.9422, 1.4754, 1.4570, 0.8296, 0.7154, 1.0605],
        [0.4753, 0.8434, 0.8296, 0.4937, 0.3474, 0.6565],
        [0.4576, 0.7070, 0.7154, 0.3474, 0.6654, 0.2935],
        [0.6310, 1.0865, 1.0605, 0.6565, 0.2935, 0.9450]])

张量中的每个元素表示每对输入之间的注意力分数,如图 3-11 所示。请注意,图中的值已经归一化,这就是它们与之前张量中的未归一化注意力分数不同的原因。我们稍后会处理归一化的问题。

在计算前面的注意力分数张量时,我们使用了 Python 中的 for 循环。然而,for 循环通常较慢,因此可以使用矩阵乘法来得到相同的结果:

python 复制代码
attn_scores = inputs @ inputs.T
print(attn_scores)

优雅的方式

复制代码
inputs (6×3)              inputs.T (3×6)
[x1]  [0.43 0.15 0.89]    转置后每一列是一个词向量
[x2]  [0.55 0.87 0.66]    [0.43 0.55 0.57 0.22 0.77 0.05]
[x3]  [0.57 0.85 0.64]    [0.15 0.87 0.85 0.58 0.25 0.80]
[x4]  [0.22 0.58 0.33]    [0.89 0.66 0.64 0.33 0.10 0.55]
[x5]  [0.77 0.25 0.10]
[x6]  [0.05 0.80 0.55]

在图 3-12 的第(2)步中,我们现在对每一行进行归一化,以确保每一行中的值总和为 1:

python 复制代码
attn_weights = torch.softmax(attn_scores, dim=-1)
print(attn_weights)

这会返回如下注意力权重张量,其值与图 3-11 中所示的值一致:

复制代码
tensor([[0.2098, 0.2006, 0.1981, 0.1242, 0.1220, 0.1452],
        [0.1385, 0.2379, 0.2333, 0.1240, 0.1082, 0.1581],
        [0.1390, 0.2369, 0.2326, 0.1242, 0.1108, 0.1565],
        [0.1435, 0.2074, 0.2046, 0.1462, 0.1263, 0.1720],
        [0.1526, 0.1958, 0.1975, 0.1367, 0.1879, 0.1295],
        [0.1385, 0.2184, 0.2128, 0.1420, 0.0988, 0.1896]])

在使用 PyTorch 时,像 torch.softmax 这样的函数中的 dim 参数用于指定输入张量的计算维度。 将 dim 设置为-1 表示让 softmax 函数在 attn_scores 张量的最后一个维度上进行归一化。如果 attn_scores 是一个二维张量(比如形状为行, 列),那么它将对列进行归一化,使得每行的值(在列维度上的总和)为 1。

python 复制代码
row_2_sum = sum([0.1385, 0.2379, 0.2333, 0.1240, 0.1082, 0.1581])
print("Row 2 sum:", row_2_sum)

print("All row sums:", attn_weights.sum(dim=-1))
复制代码
Row 2 sum: 1.0
All row sums: tensor([1.0000, 1.0000, 1.0000, 1.0000, 1.0000, 1.0000])

PyTorch 中 dim 的含义

直觉上,我们可能会认为:

  • dim=0 代表"行",所以应该让每行之和为 1;
  • dim=1 代表"列",所以应该让每列之和为 1。

但在 PyTorch 中,dim 的含义并不是"对哪一个对象归一化",而是:

沿着哪个维度(坐标轴)执行计算。


二维矩阵中的维度

假设有一个形状为 (2, 3) 的矩阵:

python 复制代码
x = torch.tensor([
    [1.0, 2.0, 3.0],
    [4.0, 5.0, 6.0]
])

可以把它的两个维度理解为:

text 复制代码
                 dim=1(列轴)
                      →
dim=0(行轴)    [[1, 2, 3],
     ↓           [4, 5, 6]]
  • dim=0 是第 0 个轴,方向是从上到下;
  • dim=1 是第 1 个轴,方向是从左到右。

dim=0 表示沿着第 0 个轴,也就是从上到下进行计算。

此时参与同一组计算的元素是:

text 复制代码
1 和 4
2 和 5
3 和 6

例如:

python 复制代码
x.sum(dim=0)

计算过程是:

text 复制代码
[1 + 4, 2 + 5, 3 + 6]

结果为:

text 复制代码
tensor([5., 7., 9.])

因此:

python 复制代码
torch.softmax(x, dim=0)

会分别对以下三组数据执行 softmax:

text 复制代码
[1, 4]
[2, 5]
[3, 6]

最终效果是:

每一列的元素之和为 1。

可以这样验证:

python 复制代码
weights = torch.softmax(x, dim=0)
print(weights.sum(dim=0))

输出:

text 复制代码
tensor([1., 1., 1.])

dim=1 表示沿着第 1 个轴,也就是从左到右进行计算> 此时组计算的元素是:

text 复制代码
1、2、3
4、5、6

例如:

python 复制代码
x.sum(dim=1)

计算过程是:

text 复制代码
[1 + 2 + 3, 4 + 5 + 6]

结果为:

text 复制代码
tensor([6., 15.])

因此:

python 复制代码
torch.softmax(x, dim=1)

会分别对每一行执行 softmax:

text 复制代码
第一组:[1, 2, 3]
第二组:[4, 5, 6]

最终效果是:

每一行的元素之和为 1。

可以这样验证:

python 复制代码
weights = torch.softmax(x, dim=1)
print(weights.sum(dim=1))

输出:

text 复制代码
tensor([1., 1.])

为什么 PyTorch 要这样设计?

因为张量不一定只有"行"和"列"两个维度。

例如,一个三维张量的形状可能是:

text 复制代码
[batch, sequence, embedding]

对应的维度编号为:

text 复制代码
dim=0:batch 维度
dim=1:sequence 维度
dim=2:embedding 维度

对于更高维张量,使用"按行计算"或者"按列计算"就不够准确了,因为第三维及后续维度无法简单地称为行或列。

因此,PyTorch 使用了一个统一的规则:

dim 指定沿着哪个坐标轴执行计算。

这个规则可以适用于任意维度的张量。


从"维度被压缩"理解 dim

对于 summeanmax 等操作,还可以把 dim 理解成:

指定哪个维度参与运算并被压缩掉。

原矩阵的形状为:

text 复制代码
(2, 3)

执行:

python 复制代码
x.sum(dim=0)

结果形状为:

text 复制代码
(3,)

因为大小为 2 的第 0 维被求和压缩掉了。

执行:

python 复制代码
x.sum(dim=1)

结果形状为:

text 复制代码
(2,)

因为大小为 3 的第 1 维被求和压缩掉了。

需要注意的是,softmax 不会真正改变张量的形状,但 dim 的含义仍然相同:沿指定维度上的元素进行 softmax 计算。

总结

dim=0 不是"对每一行计算",而是"沿第 0 个轴上下计算",所以每列之和为 1;dim=1 是沿第 1 个轴左右计算,所以每行之和为 1。

在图 3-12 的第(3)步,也就是最后一步,我们用这些注意力权重通过矩阵乘法计算出所有上下文向量,生成的输出张量中,每一行包含一个三维的上下文向量

python 复制代码
all_context_vecs = attn_weights @ inputs
print(all_context_vecs)
复制代码
tensor([[0.4421, 0.5931, 0.5790],
        [0.4419, 0.6515, 0.5683],
        [0.4431, 0.6496, 0.5671],
        [0.4304, 0.6298, 0.5510],
        [0.4671, 0.5910, 0.5266],
        [0.4177, 0.6503, 0.5645]])

all_context_vecs 中的每一行,都是原来对应单词经过注意力计算后得到的上下文化表示。上下文向量不是简单覆盖原词向量,而是原词向量经过上下文信息增强后得到的新表示;真实 Transformer 还会通过残差连接保留原始信息,不过这个是后话了

至此,简单自注意力机制的代码讲解就结束了。接下来,我们将添加可训练的权重,使大语言模型能够从数据中学习,并提升其在特定任务上的性能。

3.4 实现带可训练权重的自注意力机制

接下来,我们将实现在原始 Transformer 架构、GPT 模型和大多数其他流行的大语言模型中使用的自注意机制。这种自注意力机制也被称为缩放点积注意力(scaled dot-product attention)。 图 3-13 展示了在实现整个大语言模型的过程中,自注意力机制是如何嵌入的。

如图 3-13 所示,带有可训练权重的自注意力机制是建立在先前概念之上的:我们希望将上下文向量计算为某个特定输入元素对于序列中所有输入向量的加权和。你会看到,带有可训练权重的自注意力机制与我们之前实现的基础自注意力机制只有些微的不同。

最显著的区别是这里引入了在模型训练期间更新的权重矩阵。这些可训练的权重矩阵至关重要,这样模型(特别是模型内部的注意力模块)才能学会产生"好的"上下文向量。(请注意, 我们将在第 5 章中训练大语言模型。)

3.4.1 节和 3.4.2 节将讨论这种自注意力机制。首先,我们会像以前一样逐步编写代码。其次,我们会把代码组织成一个紧凑的 Python 类,以便能够导入到大语言模型架构中。

3.4.1 逐步计算注意力权重

本节将通过引入 3 个可训练的权重矩阵 WqWqWq 、WkWkWk 和 WvWvWv ,一步一步地实现自注意力机制。这 3 个矩阵用于将嵌入的输入词元 x(i)x^{(i)}x(i) 分别映射为查询向量、键向量和值向量,如图 3-14 所示。

之前,当我们通过注意力权重计算上下文向量 z(2)z^{(2)}z(2) 时,将第二个输入元素 x(2)x^{(2)}x(2)定义为了查询。 然后,我们将这一方法推广到了计算所有上下文向量 z(1)⋅⋅⋅z(T)z^{(1)} ··· z^{(T)}z(1)⋅⋅⋅z(T) ,应用于 6 个词的输入句子"Your journey starts with one step."

对,你的理解基本正确。

之前的简化版自注意力中,输入只有词元向量 x(i)x^{(i)}x(i),没有通过可训练矩阵显式生成 QQQ、KKK、VVV。

但在计算过程中,(x) 实际上已经分别承担了 Query、Key 和 Value 的角色:

q(2)=x(2),k(i)=x(i),v(i)=x(i)q^{(2)}=x^{(2)},\qquad k^{(i)}=x^{(i)},\qquad v^{(i)}=x^{(i)}q(2)=x(2),k(i)=x(i),v(i)=x(i)

因此,查询词元 x(2)x^{(2)}x(2) 与其他词元之间的注意力分数为:

ω2i=x(2)⋅x(i)\omega_{2i}=x^{(2)}\cdot x^{(i)}ω2i=x(2)⋅x(i)

经过 Softmax 归一化后:

α2i=softmax⁡(ω2i)\alpha_{2i}=\operatorname{softmax}(\omega_{2i})α2i=softmax(ω2i)

最终得到上下文向量:

z(2)=∑iα2ix(i)z^{(2)}=\sum_i\alpha_{2i}x^{(i)}z(2)=∑iα2ix(i)

因此,简化版可以理解为:

Q=K=V=XQ=K=V=XQ=K=V=X

这里没有对它们进行独立的、可学习的线性变换。

在标准的自注意力机制中,会引入三个不同的可训练权重矩阵:

q(i)=x(i)WQq^{(i)}=x^{(i)}W_Qq(i)=x(i)WQ

k(i)=x(i)WKk^{(i)}=x^{(i)}W_Kk(i)=x(i)WK

v(i)=x(i)WVv^{(i)}=x^{(i)}W_Vv(i)=x(i)WV

此时,对于查询词元 x(2)x^{(2)}x(2),注意力分数变为:

ω2i=q(2)⋅k(i)\omega_{2i}=q^{(2)}\cdot k^{(i)}ω2i=q(2)⋅k(i)

注意力权重通过 Softmax 计算:

α2i=softmax⁡(q(2)⋅k(i)dk)\alpha_{2i} =\operatorname{softmax}\left( \frac{q^{(2)}\cdot k^{(i)}}{\sqrt{d_k}} \right)α2i=softmax(dk q(2)⋅k(i))

然后使用注意力权重对 Value 向量进行加权求和:

z(2)=∑iα2iv(i)z^{(2)}=\sum_i\alpha_{2i}v^{(i)}z(2)=∑iα2iv(i)

对所有输入词元进行计算时,可以写成矩阵形式:

Q=XWQ,K=XWK,V=XWVQ=XW_Q,\qquad K=XW_K,\qquad V=XW_VQ=XWQ,K=XWK,V=XWV

Z=softmax⁡(QK⊤dk)VZ=\operatorname{softmax}\left( \frac{QK^\top}{\sqrt{d_k}} \right)VZ=softmax(dk QK⊤)V

同样,为了便于说明,这里我们只计算一个上下文向量 z(2)z^{(2)}z(2)。之后我们会修改这段代码来计算所有上下文向量。

首先,定义几个变量:

python 复制代码
print(inputs) # 防止遗忘本章节一开始定义的的inputs是什么
x_2 = inputs[1] # second input element
d_in = inputs.shape[1] # the input embedding size, d=3
d_out = 2 # the output embedding size, d=2
复制代码
tensor([[0.4300, 0.1500, 0.8900],
        [0.5500, 0.8700, 0.6600],
        [0.5700, 0.8500, 0.6400],
        [0.2200, 0.5800, 0.3300],
        [0.7700, 0.2500, 0.1000],
        [0.0500, 0.8000, 0.5500]])

请注意,在类 GPT 模型中,输入和输出的维度通常是相同的,但为了便于理解计算过程,这里我们使用不同的输入维度(d_in=3)和输出维度(d_out=2)。

GPT 中输入和输出维度通常相同,是因为 Transformer 使用了残差连接:

y=x+F(x)y = x + F(x)y=x+F(x)

要进行逐元素相加,xxx 和 F(x)F(x)F(x) 的维度必须一致。同时,保持统一的隐藏维度 dd_{\text{}}d 也便于堆叠多个 Transformer 层。

注意,这里指的是 Transformer 层的输入和输出维度相同,并不代表内部投影维度或最终词表输出维度也必须相同。

然后,初始化图 3-14 中的 3 个权重矩阵 WqWqWq 、W 和 WvWvWv 。设置 requires_grad=False 以减少输出中的其他项,但如果要在模型训练中使用这些权重矩 阵,就需要设置 requires_grad=True,以便在训练中更新这些矩阵。接下来,计算查询向量、键向量和值向量,因为我们通过 d_out 将对应的权重矩阵的列数设置为了 2,所以查询的输出结果是一个二维向量。

python 复制代码
torch.manual_seed(123)

W_query = torch.nn.Parameter(torch.rand(d_in, d_out), requires_grad=False)
W_key   = torch.nn.Parameter(torch.rand(d_in, d_out), requires_grad=False)
W_value = torch.nn.Parameter(torch.rand(d_in, d_out), requires_grad=False)

print(x_2)
query_2 = x_2 @ W_query # _2 because it's with respect to the 2nd input element
key_2 = x_2 @ W_key 
value_2 = x_2 @ W_value

print(query_2)
复制代码
tensor([0.5500, 0.8700, 0.6600])
tensor([0.4306, 1.4551])

权重参数与注意力权重

在权重矩阵 W 中,"权重"是"权重参数"的简称,表示在训练过程中优化的神经网络参 数。这与注意力权重是不同的。正如我们已经看到的,注意力权重决定了上下文向量对输入的不同部分的依赖程度(网络对输入的不同部分的关注程度)。

总之,权重参数是定义网络连接的基本学习系数,而注意力权重是动态且特定于上下文 的值。

虽然目前我们的目标只是计算一个上下文向量 z(2)z^{(2)}z(2) ,但仍然需要所有输入元素的键向量和值向量,因为它们参与了计算相对于查询 q(2)q^{(2)}q(2) 的注意力权重(参见图 3-14)。

可以通过矩阵乘法得到所有的键向量和值向量,从输出中可以看出,我们成功地将 6 个输入词元从三维空间映射到了二维嵌入空间。

python 复制代码
keys = inputs @ W_key 
values = inputs @ W_value

print("keys.shape:", keys.shape)
print("values.shape:", values.shape)
复制代码
keys.shape: torch.Size([6, 2])
values.shape: torch.Size([6, 2])

接下来是计算注意力分数,如图 3-15 所示。

首先,计算出注意力分数 ω22ω22ω22,未归一化的注意力评分结果如下所示

python 复制代码
keys_2 = keys[1] # Python starts index at 0
attn_score_22 = query_2.dot(keys_2)
print(attn_score_22)
复制代码
tensor(1.8524)

同样,可以通过矩阵乘法将这个计算推广到所有的注意力分数,如你所见,经过快速检查,输出中的第二个元素与之前计算的 attn_score_22 一致。

python 复制代码
attn_scores_2 = query_2 @ keys.T # All attention scores for given query
print(attn_scores_2)
复制代码
tensor([1.2705, 1.8524, 1.8111, 1.0795, 0.5577, 1.5440])

现在,我们想要将注意力分数转换为注意力权重,如图 3-16 所示。我们通过缩放注意力分数并应用 softmax 函数来计算注意力权重。不过,此时是通过将注意力分数除以键向量的嵌入维度的平方根来进行缩放(取平方根在数学上等同于以 0.5 为指数进行幂运算)。运行上述代码得到的注意力权重如下所示。

python 复制代码
d_k = keys.shape[1]
attn_weights_2 = torch.softmax(attn_scores_2 / d_k**0.5, dim=-1)
print(attn_weights_2)
复制代码
tensor([0.1500, 0.2264, 0.2199, 0.1311, 0.0906, 0.1820])

为什么要除根号dk

输入矩阵 X ∈ ℝ^(6×3),理解为 6 个词元、每个词元 3 个输入特征。权重矩阵 W_Q, W_K, W_V ∈ ℝ^(3×2),把每个词元从 3 维投影到 2 维,得到 Q = XWQXW_QXWQ、K = XWKXW_KXWK、V = XWVXW_VXWV,均为 ℝ^(6×2)。因此:6 是词元数量,3 是输入维度 d_in,2 是投影后的 Query/Key/Value 维度。此例中 dk = 2,因为每个 Key 有两个分量,每次点积累加两个乘积。

一行 Query 得到 6 个分数:q(2)=0.4306,1.4551q^{(2)} = 0.4306, 1.4551q(2)=0.4306,1.4551,与 KTK^TKT 相乘得到 6 个数。每个分数 s2j=q1(2)⋅k1(j)+q2(2)⋅k2(j)s_2j = q_1^{(2)}·k_1^{(j)} + q_2^{(2)}·k_2^{(j)}s2j=q1(2)⋅k1(j)+q2(2)⋅k2(j),内部有 2 项相加。输出为 [1.2705, 1.8524, 1.8111, 1.0795, 0.5577, 1.5440]。两个层次要分清:外层 6 个 Key → 6 个分数;每个分数内部 2 项相加 → dk = 2。

为什么不按 Key 数量缩放:即使序列从 6 个词元增加到 100 个,只要 Q、K 仍是二维,每个分数内部仍只加 2 项,缩放因子仍是 √2 而非 √100。序列长度决定产生多少个分数,dk 决定每个分数内部累加多少项,这是两件不同的事。

为什么点积典型大小是 √dk:设点积中每个乘积项 X_r = q_r·k_r 可正可负、典型波动约为 1。一项时 Var(S₁) = 1,Std = 1;两项独立相加 Var(S₂) = 2,但 Std = √2;64 项时 Var = 64,Std = 8。所以 dk 个正负随机项相加时 Var(S) ≈ dk,而典型波动(标准差)Std(S) ≈ √dk。关键区别:方差约为 dk,标准差约为 √dk。要缩放的是原始分数 S 而不是方差,所以除以它的标准差 √dk。

为什么不是除以 dk:正负随机数相加,典型大小不会线性增长为 dk。若每项只取 +1 或 −1,100 项相加通常不会得到 +100,因为正负互相抵消,典型结果更接近 √100 = 10。项数为 dk ⇒ 典型波动为 √dk。只有当所有项始终同号、强相关时和才可能按 dk 尺度增长,但缩放点积注意力采用的是近似独立、零均值的初始化分析。

从方差公式严格推:设 S = q·k = Σ q_r·k_r,令 X_r = q_r·k_r,假设各项近似独立且 Var(X_r) ≈ 1。独立变量相加方差相加,故 Var(S) = 1 + ... + 1(dk 项)= dk。缩放 S' = S/√dk,由 Var(aX) = a²·Var(X) 得 Var(S') = (1/√dk)²·dk = (1/dk)·dk = 1。根号就出现在这一步:因为方差会把缩放系数平方,所以原始数值只需除以 √dk,方差就除以 dk。

缩放点积注意力的原理

对嵌入维度进行归一化是为了避免梯度过小,从而提升训练性能。例如,在类 GPT 大语 言模型中,嵌入维度通常大于 1000,这可能导致点积非常大,从而在反向传播时由于 softmax 函数的作用导致梯度非常小。当点积增大时,softmax 函数会表现得更像阶跃函数,导致梯度接近零。这些小梯度可能会显著减慢学习速度或使训练停滞。

因此,通过嵌入维度的平方根进行缩放解释了为什么这种自注意力机制也被称为缩放点 积注意力机制。

现在,最后一步是计算上下文向量,如图 3-17 所示。

与计算上下文向量时对输入向量进行加权求和(参见 3.3 节)的方式类似,现在通过对值向量进行加权求和来计算上下文向量。在这里,注意力权重作为加权因子,用于权衡每个值向量的重要性。和之前一样,可以使用矩阵乘法一步获得输出结果:

python 复制代码
context_vec_2 = attn_weights_2 @ values
print(attn_weights_2)
print(values)
print(context_vec_2)
复制代码
tensor([0.1500, 0.2264, 0.2199, 0.1311, 0.0906, 0.1820])
tensor([[0.1855, 0.8812],
        [0.3951, 1.0037],
        [0.3879, 0.9831],
        [0.2393, 0.5493],
        [0.1492, 0.3346],
        [0.3221, 0.7863]])
tensor([0.3061, 0.8210])

到目前为止,我们只计算了一个上下文向量 z(2)z^{(2)}z(2)。在 3.4.2 节中,我们将扩展代码来计算输入序列 中 z(1)z^{(1)}z(1) 到 z(T)z^{(T)}z(T) 的所有上下文向量。

为什么要用查询、键和值

在注意力机制中,"键"(key)、"查询"(query)和"值"(value)这些术语借用自信息检 索和数据库领域,这些领域使用类似的概念来进行信息存储、搜索和检索。

查询类似于数据库中的搜索查询。它代表了模型当前关注或试图理解的项(比如句子中的一个单词或词元)。查询用于探测输入序列中的其他部分,以确定对它们的关注程度。

键类似于用于数据库索引和搜索的键。在注意力机制中,输入序列中的每个项(比如句 子中的每个单词)都有一个对应的键。这些键用于与查询进行匹配。

在这种背景下,值类似于数据库中键-值对中的值。它表示输入项的实际内容或表示。一 旦模型确定哪些键以及哪些输入部分与查询(当前关注的项)最相关,它就会检索相应的值。

在这里补充讲解"归一化"到底是什么

归一化的本质 :所有叫"归一化"的操作,本质上都在做同一件事------通过一个变换,让数据满足某个预先规定的性质,从而消除掉"不该影响结果的差异"。关键词是那个"规定的性质",它就是这个操作的约束 。不同的归一化,只是规定的性质不同。

为什么容易混乱 :中文的"归一化"其实对应了英文两个不同词根。一是 Normalization(normal = 规范、标准),让数据符合某种标准形态,如 BatchNorm、LayerNorm、数据标准化;二是 Normalize by the norm(norm = 范数,即向量长度),除以自身的长度,如 L2 归一化。它们叫同一个名字,但目标不是一回事------这就是混乱的根源。

第一类:消除"数值尺度/分布"的差异(Normalization)

要解决的问题:训练时各层激活值分布乱飘,数值忽大忽小,让训练不稳定。

规定的标准:让数值变成零均值、单位方差(长得像标准正态分布)。

公式:x^=x−μσ\hat{x} = \dfrac{x - \mu}{\sigma}x^=σx−μ。减去均值 μ\muμ 使均值为 0,除以标准差 σ\sigmaσ 使方差为 1。任何一堆数据经过它都被"归到"零均值单位方差这个统一标准。

BatchNorm 与 LayerNorm 的区别 :公式完全一样,唯一区别是 μ\muμ 和 σ\sigmaσ 拿哪些数算。

例:一个 batch 两个样本,样本1 = 1,2,3,样本2 = 4,5,6

BatchNorm:对每一列(同一特征、跨样本)算------1,4 一套、2,5 一套、3,6 一套。消除"某特征在不同样本间的尺度差异",但需要 batch 里样本多才算得准,batch 小就崩。

LayerNorm:对每一行(同一样本、跨特征)算------1,2,3 一套、4,5,6 一套。每个样本自己管自己,与 batch 大小无关。Transformer 用它,因为句子长短不一、batch 可能很小。

一句话:BatchNorm 跨样本,LayerNorm 跨特征,其余一模一样。

第二类:消除"长度"差异,只保留方向(L2 归一化)

要解决的问题:只关心向量指向哪,不关心多长(如比较人脸 embedding 像不像)。

规定的标准:让每个向量的模长变成 1,全部落到单位球面上。

公式:x^=x∥x∥2\hat{x} = \dfrac{x}{\|x\|_2}x^=∥x∥2x,其中 ∥x∥2=x12+x22+⋯\|x\|_2 = \sqrt{x_1^2 + x_2^2 + \cdots}∥x∥2=x12+x22+⋯ 。向量除以自身长度,结果长度必为 1(如 3,4 ÷ 5 = 0.6,0.8,长度正好 1)。方向不变,长度统一。这里根本不关心均值方差,标准就是"长度=1"。

第三类:变成"概率"(Softmax)

要解决的问题:把一堆分数变成"属于每一类的概率"。概率的标准是每个非负、加起来等于 1。

公式:softmax(x)i=exi∑jexj\text{softmax}(x)_i = \dfrac{e^{x_i}}{\sum_j e^{x_j}}softmax(x)i=∑jexjexi。先取指数保证每项为正(非负),再除以总和保证加起来为 1(归一)。

三类对比

想消除什么差异 规定的标准 靠什么公式
标准化类 (BN/LN) 数值的尺度和分布飘动 零均值、单位方差 减均值 ÷ 标准差
L2 归一化 向量长度差异 模长 = 1 ÷ 自身模长
Softmax 分数不是概率 非负且和为 1 指数 ÷ 求和

用三个问题串起来

  1. 希望满足什么约束(先问这个)→ 决定属于哪一类:稳定分布?单位长度?概率?
  2. 用什么公式 → 约束定了公式基本唯一(要零均值单位方差就只能减均值除标准差)。
  3. 沿哪个维度 → 只在第一类里纠结(BN 还是 LN),本质是"想消除的差异在哪个方向":跨样本用 BN,样本内部用 LN。

3.4.2 实现一个简化的自注意 Python 类

到目前为止,我们已经完成了多个步骤来计算自注意力的输出。这些步骤主要是为了演示清晰,以便逐步了解每个环节。在实际操作中,为了实现第 4 章中的大语言模型,最好将这些代码组织成一个 Python 类,如代码清单 3-1 所示。

在这段 PyTorch 代码中,SelfAttention_v1 是一个从 nn.Module 派生出来的类。nn.Module 是 PyTorch 模型的一个基本构建块,它为模型层的创建和管理提供了必要的功能。

__init__方法初始化了可训练的权重矩阵(W_query、W_key 和 W_value),这些矩阵用 于查询向量、键向量和值向量,每个矩阵将输入维度 d_in 转换为输出维度 d_out。

在前向传播过程中,我们通过使用 forward 方法将查询向量和键向量相乘来计算注意力分数(attn_scores),然后使用 softmax 对这些分数进行归一化。最后,我们通过使用这些归一化的注意力分数对值向量进行加权来创建上下文向量。

可以通过以下方式来使用这个类:

python 复制代码
import torch.nn as nn
import torch

class SelfAttention_v1(nn.Module):

    def __init__(self, d_in, d_out):
        super().__init__()
        self.W_query = nn.Parameter(torch.rand(d_in, d_out))
        self.W_key   = nn.Parameter(torch.rand(d_in, d_out))
        self.W_value = nn.Parameter(torch.rand(d_in, d_out))

    def forward(self, x):
        keys = x @ self.W_key
        queries = x @ self.W_query
        values = x @ self.W_value
        
        attn_scores = queries @ keys.T 
        attn_weights = torch.softmax(
            attn_scores / keys.shape[-1]**0.5, dim=-1
        )

        context_vec = attn_weights @ values
        return context_vec

torch.manual_seed(123)


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)
)
d_in = inputs.shape[1] # the input embedding size, d=3
d_out = 2 # the output embedding size, d=2
sa_v1 = SelfAttention_v1(d_in, d_out)
print(sa_v1(inputs))

归一化的两个 -1 的含义

这两个 -1 含义完全不同,只是写法上都用了 -1 而已。

第一个:keys.shape[-1]

shape 是一个元组,[-1] 表示取最后一个元素。keys 的形状是 (6, 2),所以 keys.shape[-1] 就是 2,也就是 d_out(每个 key 向量的维度)。

这里的作用是做缩放(scaled dot-product attention 里的 scaling)。注意力分数除以 √(维度),即 keys.shape[-1]**0.5 = √2。这样做是为了防止点积结果随维度增大而变得过大,导致 softmax 梯度变小、训练不稳定。用 [-1] 只是为了稳妥地拿到最后一维的大小,不管张量有几个维度都能正确取到向量维度。

第二个:dim=-1

这是 softmax 的参数,指定沿哪个维度做归一化。-1 表示最后一个维度。

attn_scores 形状是 (6, 6)attn_scores[i][j] 表示第 i 个 query 对第 j 个 key 的关注分数。dim=-1 表示对每一行做 softmax,让每行权重和为 1。这正是注意力想要的:某个 token 的新表示是所有词 value 的加权平均,权重必须和为 1。若错用 dim=0(对每列归一化),就变成「所有 query 对同一个 key 的权重和为 1」,语义就乱了。

为什么用 -1 而不是写死数字:

dim=-1 永远指向最后一维,对形状变化「免疫」;写死的正数索引一旦维度增加就会指错。以对 key 维归一化为例:

张量形状 含义 最后一维索引 dim=-1 写死 dim=1
(6,6) (query, key) 1
(32,6,6) (batch, query, key) 2 ❌ 指到 query
(32,8,6,6) (batch, heads, query, key) 3 ❌ 指到 heads

加了 batch 或多头后,dim=-1 一行代码都不用改,而写死的索引每次都得手动改、还容易出错。

位置 含义 值/作用
keys.shape[-1] 取形状元组的最后一个元素 得到向量维度 2,用于缩放
dim=-1 softmax 沿最后一个维度归一化 让每行权重和为 1

它们只是恰好都用了 -1 这个索引写法,实际语义没有关联。使用 -1 而不是写死具体数字,是一种更通用、更健壮的写法。

由于输入包含 6 个嵌入向量,因此我们会得到一个用于保存这 6 个上下文向量的矩阵:

复制代码
tensor([[0.2996, 0.8053],
        [0.3061, 0.8210],
        [0.3058, 0.8203],
        [0.2948, 0.7939],
        [0.2927, 0.7891],
        [0.2990, 0.8040]], grad_fn=<MmBackward0>)

通过快速检查,可以看到第 2 行(0.3061, 0.8210)与 3.4.1 节中的 context_vec_2 内容相符。

图 3-18 总结了我们刚刚实现的自注意力机制。

自注意力机制包含了可训练的权重矩阵 Wq 、Wk 和 Wv 。这些矩阵将输入数据转换为查询向量、键向量和值向量,这些组件在注意力机制中至关重要。随着模型在训练中接触更多数据,它会调整这些可训练的权重,后续章节会对此进行介绍。

可以通过使用 PyTorch 的 nn.Linear 层来进一步优化 SelfAttention_v1 的实现,当偏置 单元被禁用时,nn.Linear 层可以有效地执行矩阵乘法。相比手动实现 nn.Parameter(torch. rand(...)),使用 nn.Linear 的一个重要优势是它提供了优化的权重初始化方案,从而有助于模型训练的稳定性和有效性,如代码清单 3-2 所示。

python 复制代码
class SelfAttention_v2(nn.Module):

    def __init__(self, d_in, d_out, qkv_bias=False):
        super().__init__()
        self.W_query = nn.Linear(d_in, d_out, bias=qkv_bias)
        self.W_key   = nn.Linear(d_in, d_out, bias=qkv_bias)
        self.W_value = nn.Linear(d_in, d_out, bias=qkv_bias)

    def forward(self, x):
        keys = self.W_key(x)
        queries = self.W_query(x)
        values = self.W_value(x)
        
        attn_scores = queries @ keys.T
        attn_weights = torch.softmax(attn_scores / keys.shape[-1]**0.5, dim=-1)

        context_vec = attn_weights @ values
        return context_vec

torch.manual_seed(789)

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)
)
d_in = inputs.shape[1] # the input embedding size, d=3
d_out = 2 # the output embedding size, d=2

sa_v2 = SelfAttention_v2(d_in, d_out)
print(sa_v2(inputs))

v1 vs v2:两种写法的区别

两者做的事情完全一样:把输入 x(形状 (num_tokens, d_in))变成输出(形状 (num_tokens, d_out)),算出来的结果相同,只是写法和权重存储方向不同。

v1:手写矩阵

python 复制代码
self.W_key = nn.Parameter(torch.rand(d_in, d_out))  # 权重存成 (d_in, d_out)
keys = x @ self.W_key

维度对账:(num_tokens, d_in) @ (d_in, d_out) = (num_tokens, d_out)

权重方向本来就对,直接乘,不用转置。

v2:用 nn.Linear

python 复制代码
self.W_key = nn.Linear(d_in, d_out, bias=False)  # 内部权重存成 (d_out, d_in)
keys = self.W_key(x)   # 内部执行 x @ weight.T

虽然写的参数是 (d_in, d_out),但 PyTorch 内部把 weight 存成反方向 (d_out, d_in),所以计算时要 .T 转回来对齐维度。

为什么 PyTorch 要存成 (d_out, d_in)

  1. 语义清晰 ------ 每一行正好是一个输出神经元的全部权重,weight[i] 就是第 i 个神经元。
  2. 内存连续 ------ 矩阵按行连续存放,点积时读取每个神经元的权重最快(缓存友好)。
  3. 底层约定 ------ 和 BLAS/cuBLAS 等线性代数库的高效调用方式一致。

代价:数学公式要写成 x @ W.T,多一个转置对齐维度。

一句话总结

v1 权重存 (d_in, d_out),直接乘;v2 的 nn.Linear 为了"每行=一个神经元"的清晰语义和计算效率,把权重存成 (d_out, d_in),所以内部用 x @ W.T 补偿。结果一样,这也是搬权重时要加 .T 的原因。

请注意,SelfAttention_v1 和 SelfAttention_v2 因为使用了不同的初始权重矩阵而给出了不同的输出,这是由 nn.Linear 使用了一个更复杂的权值初始化方案所导致的。

复制代码
tensor([[-0.0739,  0.0713],
        [-0.0748,  0.0703],
        [-0.0749,  0.0702],
        [-0.0760,  0.0685],
        [-0.0763,  0.0679],
        [-0.0754,  0.0693]], grad_fn=<MmBackward0>)

练习 3.1比较 SelfAttention_v1 和 SelfAttention_v2

注意, SelfAttention_v2 中的 nn.Linear 使用了与 SelfAttention_v1 中的 nn.Parameter(torch.rand(d_in, d_out))不同的权重初始化方式,这导致两个机制的输出结果不同。为了确认 SelfAttention_v1 和 SelfAttention_v2 的其他方面是否相同,可以将 SelfAttention_v2 对象的权重矩阵转移到 SelfAttention_v1 对象中,这样两个对象就会产生相同的结果。

你的任务是将 SelfAttention_v2 实例中的权重正确地分配给 SelfAttention_v1 实例。为此,需要了解两个版本中的权重之间的关系。(提示:nn.Linear 以转置的形式存储权重矩阵。)完成分配后,应该能够观察到这两个实例产生了相同的输出。

接下来,我们将改进自注意力机制,重点是在机制中引入因果机制和多头机制。因果机制的作用是调整注意力机制,防止模型访问序列中未来的信息,这在语言建模等任务中尤为重要,因为每个词的预测只能依赖之前出现的词。

多头机制涉及将注意力机制分成多个"头"。每个头会学习数据的不同特征,使模型能够在 不同的位置同时关注来自不同表示子空间的信息。这能够提升模型在复杂任务中的性能。

从期望到方差,再到神经网络里的初始化与注意力缩放

一、离散情况下的"期望"(均值)

期望就是加权平均:每个可能的取值,乘以它出现的概率,再全部加起来。

掷一个骰子,取值 1~6,每个概率都是 1/6:

复制代码
E[X] = 1·(1/6) + 2·(1/6) + 3·(1/6) + 4·(1/6) + 5·(1/6) + 6·(1/6) = 3.5

每一项都是"取值 × 该取值的概率"。写成通式:

复制代码
E[X] = Σ (取值 × 概率) = Σ xᵢ · P(xᵢ)

那个 xᵢ 就是"取值本身",所以期望公式里一定会出现取值 x,它不是凭空来的。同理 E[X²] 是把"取值的平方"做加权平均:

复制代码
E[X²] = Σ xᵢ² · P(xᵢ)

二、过渡到连续情况:x 从哪来

连续分布(比如 U(0,1))取值有无穷多个,不能一个一个加,所以把求和 Σ 换成积分 ∫,把概率 P(xᵢ) 换成概率密度 f(x)dx

复制代码
离散:  E[X] = Σ  x  ·  P(x)
连续:  E[X] = ∫  x  ·  f(x) dx
              ↑        ↑
            取值     概率(密度)

那个 x 就是离散公式里的"取值",一直都在。f(x) 只是"概率密度"这一部分。对 U(0,1),因为 f(x)=1:

复制代码
E[X]  = ∫₀¹ x · f(x) dx  = ∫₀¹ x · 1 dx  = ∫₀¹ x dx
E[X²] = ∫₀¹ x² · f(x) dx = ∫₀¹ x² · 1 dx = ∫₀¹ x² dx

三、把积分算出来

规则:∫ xⁿ dx = x^(n+1)/(n+1)(幂次加 1,再除以新幂次),然后代入上下限(上限减下限)。

复制代码
E[X]  = ∫₀¹ x dx  = [x²/2] 从0到1 = 1²/2 − 0²/2 = 1/2
E[X²] = ∫₀¹ x² dx = [x³/3] 从0到1 = 1³/3 − 0³/3 = 1/3

四、方差公式

方差的定义是"取值离均值的平均偏离程度":

复制代码
Var(X) = E[(X − μ)²]      μ 是均值

展开括号:(X − μ)² = X² − 2μX + μ²。对整个式子取期望,利用期望的线性(可逐项拆开,常数可提到外面),且 E[X] = μ

总体均值(期望):μ = E[X],是随机变量分布的理论平均值,是一个固定的常数。这跟"期望"完全等价。

样本均值:x̄ = (1/n)Σxᵢ,是你从实际抽到的 n 个数据算出来的平均,是一个随样本变化的估计值。

在你贴的这段方差推导中,μ 指的是第一种------总体均值,也就是期望 EX

复制代码
E[(X−μ)²] = E[X²] − 2μ·E[X] + μ² = E[X²] − 2μ² + μ² = E[X²] − μ²

于是得到好用的公式:

复制代码
Var(X) = E[X²] − (E[X])²

直观理解:方差 = "平方的平均" 减去 "平均的平方"。


五、算 U(0,1) 的方差

已知 EX = 1/2(所以 (EX)² = 1/4),E = 1/3:

复制代码
Var(X) = 1/3 − 1/4 = 4/12 − 3/12 = 1/12 ≈ 0.083

一般地,对 U(a, b),方差 = (b−a)²/12,即"区间宽度的平方除以 12"。U(0,1) 宽度为 1,1²/12 = 1/12。这个"宽度²/12"是后面的关键工具。


六、为什么 nn.Linear 用 U(−k, k), k = 1/√d 而不是 torch.rand

一个线性层做加权求和(d = d_in 是输入维度):

复制代码
y = w₁x₁ + w₂x₂ + ... + w_d·x_d + b

我们希望网络刚初始化时每层输出和输入"规模差不多",不要爆炸也不要消失。

torch.rand 给的是 U(0,1),有两个真正的问题:

  • 问题 1(均值不为 0 → 系统性正向偏移):U(0,1) 均值是 1/2,权重全是正数。做加权和 z = Σ wᵢxᵢ 时,每一项期望都是正的,d 项累加,z 的期望随 d 线性增长,d 越大偏得越狠。注意"权重全正"和"均值为正"是同一件事,不是两个独立原因。

  • 问题 2(方差随 d 放大):假设 wᵢ、xᵢ 独立,和的方差是各项方差相加:

    Var(y) = Σ Var(wᵢ·xᵢ) = d · Var(w) · Var(x)

那个 d 会把输出方差成正比放大,层层叠加迅速爆炸。

附带说明------权重全正还有一个坏处:反向传播时同一神经元的这些权重梯度倾向于同号(一起变大或一起变小),导致参数更新只能"锯齿形"前进,收敛效率低(和 sigmoid 输出非零中心的问题同一类)。

需要澄清一个常见误解:这不是"对称性打不破"。对称性打不破(symmetry breaking failure)指的是同一层所有神经元初始值完全相同(比如全初始化为 0 或同一常数),导致每个神经元输出、梯度、更新方向都一样,永远保持相同,等价于只有一个神经元------那是常数初始化的问题。U(0,1) 是随机采样的,每个权重都是不同的随机数,对称性其实是被打破的。

对症下药:

  • 对策 1:用对称区间 U(−k, k),均值 (−k+k)/2 = 0,输出不偏移,同时权重有正有负,避免梯度同号的锯齿问题。
  • 对策 2:让 Var(w) 反比于 d 来抵消那个 d。要 Var(y) ≈ Var(x),就要 d · Var(w) = 1,即 Var(w) = 1/d

七、把 k 解出来

对 U(−k, k),宽度是 2k,套用"宽度²/12":

复制代码
Var(w) = (2k)²/12 = 4k²/12 = k²/3

令它等于 1/d:

复制代码
k²/3 = 1/d   ⟹   k² = 3/d   ⟹   k = √(3/d)

严格保方差应该是 k = √(3/d)(Xavier/Kaiming uniform 的形式)。而 PyTorch 默认用 k = 1/√d,差一个 √3 系数。原因:网络里还有非线性激活(ReLU 会砍掉一半方差,需要补偿系数),具体常数是根据激活函数和工程经验微调的。关键抓住的是那个 1/√d 的形状------权重典型大小要随输入维度的平方根反比缩小,常数系数是次要的。PyTorch 源码里默认用 fan_in(即 d_in)作为 d。


八、这和注意力除以 √d 是一回事吗?是的

核心思想完全一样:d 个东西加起来方差会变成 d 倍,都要想办法拉回 1,只是"动手时机"不同。

注意力打分 score = q·k = q₁k₁ + ... + q_d·k_d,若 qᵢ、kᵢ 均值 0、方差 1、独立,则每项方差约 1,d 项相加:

复制代码
Var(score) = d · 1 = d      ⟹    标准差 = √d

d 大时 score 里有些值非常大,softmax 变得极度尖锐(近似 one-hot),梯度趋近 0 学不动。所以除以 √d:

复制代码
score / √d   ⟹   Var = d/(√d)² = d/d = 1
起因 做法
Linear 初始化 输出 y 是 d 项加权和,方差 ∝ d 事前:把权重方差设成 1/d,一开始就压住
注意力缩放 score 是 d 项乘积和,方差 = d 事后:算完 score 再除以 √d

一个"提前把料放少点",一个"做好了再稀释",目标都是让方差 ≈ 1。


九、具体例子(d = 100,输入均值 0、方差 1、独立)

用 rand = U(0,1)(均值 1/2,方差 1/12):

复制代码
Var(y) = d · Var(w) · Var(x) = 100 · (1/12) · 1 = 100/12 ≈ 8.3

每过一层方差乘以 8.3,层层叠加:

复制代码
rand:1 → 8.3 → 69 → 578 → 4800 → 40000 ...   爆炸(5层后标准差≈200,很快溢出成inf)

用 U(−k, k), k = 1/√100 = 0.1(范围 −0.1, 0.1,均值 0):

复制代码
Var(w) = (2k)²/12 = 0.2²/12 = 1/300 ≈ 0.0033
Var(y) = 100 · (1/300) · 1 = 1/3 ≈ 0.33

k=1/√d:  1 → 0.33 → 0.11 → 0.037 ...   稳定(缓慢收缩,不爆炸不消失)

为什么是 ×1/3 而不是精确 ×1?理想保方差需要 k = √(3/d) ≈ 0.173,PyTorch 用 0.1 小了 √3 倍,所以略微收缩。但两者都是 1/√d 的形状,差的只是常数系数------这个系数按激活函数和经验微调,不影响主要结论。


十、三个要点总结

  1. 期望 = 取值 × 概率再累加,连续情况是 ∫ x·f(x)dx;方差 = E − (EX)²,均匀分布方差 = 宽度²/12。
  2. "d 个东西相加,方差变成 d 倍"是同一个数学事实:Linear 层事前把权重方差设成 1/d(即 k=1/√d)压住;注意力事后除以 √d 拉回 1。
  3. 几乎从不该用 torch.rand 初始化权重:它均值为正会导致系统性偏移、梯度同号走锯齿,且方差随 d 爆炸(注意这不是"对称性打不破",那是常数初始化的问题)。手动初始化用 torch.nn.init.kaiming_uniform_xavier_uniform_ 这类专为"保持方差"设计的函数。

3.5 利用因果注意力隐藏未来词汇

对于许多大语言模型任务,你希望自注意力机制在预测序列中的下一个词元时仅考虑当前位置之前的词元。因果注意力(也称为掩码注意力)是一种特殊的自注意力形式。它限制模型在处理任何给定词元时,只能基于序列中的先前和当前输入来计算注意力分数,而标准的自注意力机制可以一次性访问整个输入序列。

现在,我们将通过修改标准自注意力机制来创建因果注意力机制,这是在后续章节中开发大语言模型的关键步骤。要在类 GPT 模型中实现这一点,对于每个处理的词元,需要掩码当前词元之后的后续词元,如图 3-19 所示。我们会掩码对角线以上的注意力权重,并归一化未掩码的注意力权重,使得每一行的权重之和为 1。稍后,我们将通过代码来实现这一掩码和归一化过程。

3.5.1 因果注意力的掩码实现

接下来我们将在代码中实现因果注意力掩码。如图 3-20 所示,为了实现应用因果注意力掩码的步骤并得到掩码后的注意力权重,我们将使用前面介绍的注意力分数和权重来编写因果注意力机制代码。

在第(1)步中,按照之前的方法,通过 softmax 函数计算注意力权重:

python 复制代码
# Reuse the query and key weight matrices of the
# SelfAttention_v2 object from the previous section for convenience
queries = sa_v2.W_query(inputs)
keys = sa_v2.W_key(inputs) 
attn_scores = queries @ keys.T

attn_weights = torch.softmax(attn_scores / keys.shape[-1]**0.5, dim=-1)
print(attn_weights)
复制代码
tensor([[0.1921, 0.1646, 0.1652, 0.1550, 0.1721, 0.1510],
        [0.2041, 0.1659, 0.1662, 0.1496, 0.1665, 0.1477],
        [0.2036, 0.1659, 0.1662, 0.1498, 0.1664, 0.1480],
        [0.1869, 0.1667, 0.1668, 0.1571, 0.1661, 0.1564],
        [0.1830, 0.1669, 0.1670, 0.1588, 0.1658, 0.1585],
        [0.1935, 0.1663, 0.1666, 0.1542, 0.1666, 0.1529]],
       grad_fn=<SoftmaxBackward0>)

可以使用 PyTorch 的 tril 函数来实现第(2)步,该函数可以创建一个对角线以上元素为 0 的掩码:

python 复制代码
context_length = attn_scores.shape[0]
mask_simple = torch.tril(torch.ones(context_length, context_length))
print(mask_simple)
复制代码
tensor([[1., 0., 0., 0., 0., 0.],
        [1., 1., 0., 0., 0., 0.],
        [1., 1., 1., 0., 0., 0.],
        [1., 1., 1., 1., 0., 0.],
        [1., 1., 1., 1., 1., 0.],
        [1., 1., 1., 1., 1., 1.]])

现在,可以把这个掩码矩阵和注意力权重矩阵相乘,使对角线上方的值变为 0:

python 复制代码
masked_simple = attn_weights*mask_simple
print(masked_simple)
复制代码
tensor([[0.1921, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000],
        [0.2041, 0.1659, 0.0000, 0.0000, 0.0000, 0.0000],
        [0.2036, 0.1659, 0.1662, 0.0000, 0.0000, 0.0000],
        [0.1869, 0.1667, 0.1668, 0.1571, 0.0000, 0.0000],
        [0.1830, 0.1669, 0.1670, 0.1588, 0.1658, 0.0000],
        [0.1935, 0.1663, 0.1666, 0.1542, 0.1666, 0.1529]],
       grad_fn=<MulBackward0>)

第(3)步是重新归一化注意力权重,使每一行的总和再次为 1。可以通过将每行中的每个元素除以每行中的和来实现这一点:

python 复制代码
row_sums = masked_simple.sum(dim=-1, keepdim=True)
print(row_sums)
masked_simple_norm = masked_simple / row_sums
print(masked_simple_norm)

结果是一个注意力权重矩阵,其中对角线以上的注意力权重已被归 0,每一行之和为 1。

复制代码
tensor([[0.1921],
        [0.3700],
        [0.5357],
        [0.6775],
        [0.8415],
        [1.0000]], grad_fn=<SumBackward1>)
tensor([[1.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000],
        [0.5517, 0.4483, 0.0000, 0.0000, 0.0000, 0.0000],
        [0.3800, 0.3097, 0.3103, 0.0000, 0.0000, 0.0000],
        [0.2758, 0.2460, 0.2462, 0.2319, 0.0000, 0.0000],
        [0.2175, 0.1983, 0.1984, 0.1888, 0.1971, 0.0000],
        [0.1935, 0.1663, 0.1666, 0.1542, 0.1666, 0.1529]],
       grad_fn=<DivBackward0>)

信息泄露

当我们应用掩码并重新归一化注意力权重时,初看起来,未来的词元(打算掩码的)可能仍然会影响当前的词元,因为它们的值会参与 softmax 计算。然而,关键的见解是,在掩码后重新归一化时,我们实际上是在对一个较小的子集重新计算 softmax(因为被掩码的位置不参与 softmax 计算)。

softmax 函数在数学上的优雅之处在于,尽管最初所有位置都在分母中,但掩码和重新归一化之后,被掩码的位置的效果被消除------它们不会以任何实际的方式影响 softmax 分数。

简而言之,掩码和重新归一化之后,注意力权重的分布就像最初仅在未掩码的位置计算 一样。这保证了不会有来自未来或其他被掩码的词元的信息泄露。

信息泄露

产生的疑惑

第(1)步算 softmax 时,用的是完整的一整行分数:

python 复制代码
attn_weights = torch.softmax(attn_scores / keys.shape[-1]**0.5, dim=-1)

softmax 的公式是:

复制代码
softmax(x_i) = exp(x_i) / (exp(x_1) + exp(x_2) + ... + exp(x_n))
                ↑ 分母里包含了所有位置,包括"未来"的词

也就是说,分母里已经掺进了未来词元的信息。即使后面用掩码把未来位置置 0,可"未来词元已经参与过分母的计算了",那不就泄露了吗?这就是"初看起来可能仍会影响当前词元"的意思。

为什么实际上没有泄露

关键在第(3)步的重新归一化。以第 2 行为例:

掩码后的原始 softmax 结果(分母是全部 6 个位置算出来的):

复制代码
[0.2041, 0.1659, 0.0000, 0.0000, 0.0000, 0.0000]

重新归一化:每个数除以这一行的和 (0.2041 + 0.1659 = 0.3700):

复制代码
0.2041 / 0.3700 = 0.5517
0.1659 / 0.3700 = 0.4483
→ [0.5517, 0.4483, 0, 0, 0, 0]

如果一开始就只用前 2 个位置算 softmax(即理想中"根本不知道未来"的情况),结果会是:

复制代码
softmax([s1, s2]) = exp(s1)/(exp(s1)+exp(s2)),  exp(s2)/(exp(s1)+exp(s2))

而原始的 0.2041 和 0.1659 本身就是 exp(score) / (全部 6 项之和)。当我们把它们相除做归一化时:

复制代码
0.2041 / (0.2041+0.1659)
= [exp(s1)/Z] / ([exp(s1)+exp(s2)]/Z)   ← Z 是那个"包含未来"的大分母
= exp(s1) / (exp(s1)+exp(s2))           ← Z 被约掉了!

那个包含未来词元信息的大分母 Z,在除法里被完全约掉了。

结论

数学上,"先用全部位置算 softmax → 掩码 → 重新归一化" 的结果,和 "一开始就只用未掩码位置算 softmax" 的结果完全相同。所以虽然计算过程中未来词元的值曾经出现在分母里,但重新归一化把它的影响彻底消掉了,最终没有信息泄露。

补充:实际实现里通常不这么绕,而是直接把未来位置的分数设成 -∞,这样 exp(-∞)=0,未来词元从一开始就不进分母。这是更高效的做法,上面的"乘掩码再归一化"主要是为了建立直觉。

尽管可以在此时完成对因果注意力的实现,但我们仍然可以进行改进。让我们利用 softmax 函数的数学特性,以更少的步骤更高效地计算掩码后的注意力权重,如图 3-21 所示。

softmax 函数会将其输入转换为一个概率分布。当输入中出现负无穷大值(--∞)时,softmax 函数会将这些值视为零概率。(从数学角度来看,这是因为 e--∞ 无限接近于 0。)

可以通过创建一个对角线以上是 1 的掩码,并将这些 1 替换为负无穷大(-inf)值,来实 现这种更高效的掩码"方法":

python 复制代码
mask = torch.triu(torch.ones(context_length, context_length), diagonal=1)
masked = attn_scores.masked_fill(mask.bool(), -torch.inf)
print(masked)
复制代码
tensor([[0.2899,   -inf,   -inf,   -inf,   -inf,   -inf],
        [0.4656, 0.1723,   -inf,   -inf,   -inf,   -inf],
        [0.4594, 0.1703, 0.1731,   -inf,   -inf,   -inf],
        [0.2642, 0.1024, 0.1036, 0.0186,   -inf,   -inf],
        [0.2183, 0.0874, 0.0882, 0.0177, 0.0786,   -inf],
        [0.3408, 0.1270, 0.1290, 0.0198, 0.1290, 0.0078]],
       grad_fn=<MaskedFillBackward0>)

torch.triu(triangular upper,上三角):保留矩阵的上三角部分,其余位置置为 0。

torch.triu(input, diagonal=0)

  • input:输入张量(通常是二维矩阵)
  • diagonal:控制从哪条对角线开始保留,默认 0
    • diagonal=0:保留主对角线及其以上的所有元素
    • diagonal=1:主对角线也置为 0,只保留主对角线严格以上的部分
    • diagonal=-1:连主对角线下方一条也保留

对全 1 的 6×6 矩阵执行 torch.triu(torch.ones(6, 6), diagonal=1)

复制代码
tensor([[0., 1., 1., 1., 1., 1.],
        [0., 0., 1., 1., 1., 1.],
        [0., 0., 0., 1., 1., 1.],
        [0., 0., 0., 0., 1., 1.],
        [0., 0., 0., 0., 0., 1.],
        [0., 0., 0., 0., 0., 0.]])

1 出现的位置正好是「每个 token 不该看到的未来位置」。用 diagonal=1 是因为主对角线(自己看自己)允许保留,不能被掩掉。

masked_fill:根据布尔掩码,把张量中被标记为 True 的位置替换成指定值。

tensor.masked_fill(mask, value)

  • mask:布尔张量,形状要能和 tensor 广播对齐,值为 True 的位置会被替换
  • value:用来填充的标量值

mask.bool() 把 0/1 浮点矩阵转成布尔矩阵(0 → False,1 → True),然后 attn_scores.masked_fill(mask.bool(), -torch.inf) 在 mask 为 True(上三角、未来位置)的地方,把注意力分数替换成 -inf。

为什么用 -inf :后续要对每一行做 softmax,exp(-inf) = 0,所以未来位置的注意力权重精确变成 0,且不需要手动重新归一化,一步到位。这比「先置 0 再重新归一化」更高效。

注意:masked_fill 返回新张量;带下划线的 masked_fill_ 是原地(in-place)修改版本。

现在只需要对这些掩码结果应用 softmax 函数,就可以完成整个过程:

python 复制代码
attn_weights = torch.softmax(masked / keys.shape[-1]**0.5, dim=-1)
print(attn_weights)

从输出结果来看,每行中的值总和为 1,因此不需要再进行额外的归一化处理:

复制代码
tensor([[1.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000],
        [0.5517, 0.4483, 0.0000, 0.0000, 0.0000, 0.0000],
        [0.3800, 0.3097, 0.3103, 0.0000, 0.0000, 0.0000],
        [0.2758, 0.2460, 0.2462, 0.2319, 0.0000, 0.0000],
        [0.2175, 0.1983, 0.1984, 0.1888, 0.1971, 0.0000],
        [0.1935, 0.1663, 0.1666, 0.1542, 0.1666, 0.1529]],
       grad_fn=<SoftmaxBackward0>)

现在可以利用修改后的注意力权重通过 context_vec = attn_weights @ values 计算上下 文向量了,就像 3.4 节中所描述的那样。然而,在此之前,我们将先讨论一种对因果注意力机制的小调整,这在训练大语言模型时可以有效减少过拟合。

3.5.2 利用 dropout 掩码额外的注意力权重

dropout 是深度学习中的一种技术,通过在训练过程中随机忽略一些隐藏层单元来有效地"丢弃"它们。这种方法有助于减少模型对特定隐藏层单元的依赖,从而避免过拟合。需要强调的是,dropout 仅在训练期间使用,训练结束后会被取消。

在 Transformer 架构中,一些包括 GPT 在内的模型通常会在两个特定时间点使用注意力机制中的 dropout:一是计算注意力权重之后,二是将这些权重应用于值向量之后。如图 3-22 所示,我们将在计算注意力权重之后应用 dropout 掩码,因为这是实践中更常见的做法。

下面的代码示例中使用了50%的 dropout 率,这意味着掩码一半的注意力权重。(当我们在接下来的章节中训练 GPT 模型时,将使用较低的 dropout 率,比如 10%或 20%。)为了便于操作,我们首先将 PyTorch 的 dropout 实现应用于一个由 1 组成的 6×6 张量:

python 复制代码
torch.manual_seed(123)
dropout = torch.nn.Dropout(0.5) # dropout rate of 50%
example = torch.ones(6, 6) # create a matrix of ones

print(dropout(example))
复制代码
tensor([[2., 2., 2., 2., 2., 2.],
        [0., 2., 0., 0., 0., 0.],
        [0., 0., 2., 0., 2., 0.],
        [2., 2., 0., 0., 0., 2.],
        [2., 0., 0., 0., 0., 2.],
        [0., 2., 0., 0., 0., 0.]])

在对注意力权重矩阵应用 50%的 dropout 率时,矩阵中有一半的元素会随机被置为 0。为了补偿减少的活跃元素,矩阵中剩余元素的值会按 1/0.5 = 2 的比例进行放大。这种放大对于维持注意力权重的整体平衡非常重要,可以确保在训练和推理过程中,注意力机制的平均影响保持一致。

期望守恒:为什么要乘那个 1/(1-p)

不缩放会怎样?

假设一层输出是 [1, 1, 1, 1],丢弃率 p=0.5。如果只丢弃不缩放,随机丢一半后:

复制代码
[1, 0, 1, 0]   → 总和 = 2

原本总和是 4,信号缩水了一半。问题在于:训练时喂给下一层的信号量约是 2,但推理时 dropout 关闭,下一层收到的是完整的 4。同一层训练时看到"2 的量级"、推理时变成"4 的量级"------量级翻倍。下一层权重是按训练时的量级学的,推理时数值就全乱了。缩放就是为了让训练和推理的信号量级对得上。

缩放怎么补回来

丢弃后信号平均缩水到 1-p 倍,想补回原样就乘它的倒数 1/(1-p)。p=0.5 时就是乘 2:

复制代码
[1, 0, 1, 0] × 2 = [2, 0, 2, 0]   → 总和 = 4

总和回到 4,和推理时的完整信号量级一致。

"期望守恒"到底在说什么

关键是区分单次和平均两个视角。

单看某个元素的某一次:要么被丢变 0,要么被留变 2,没有一次是 1,跟原值对不上。

但换成平均视角------同一个元素反复前向很多次:

  • 50% 次数被丢 → 值是 0
  • 50% 次数被留 → 值是 2

平均起来:

复制代码
0.5 × 0 + 0.5 × 2 = 1

长期平均正好回到原值 1,这就是期望守恒(期望 = 长期平均)。

一般公式:

复制代码
E[输出] = (1-p) × [原值 × 1/(1-p)] + p × 0 = 原值

被留概率 (1-p) 乘放大后的值,加被丢概率 p 乘 0,那个 (1-p)1/(1-p) 正好抵消,结果就是原值。这个 1/(1-p) 就是精心设计来让两项抵消的。

现在,对注意力权重矩阵进行 dropout 操作:

python 复制代码
torch.manual_seed(123)
print(dropout(attn_weights))

在处理后的注意力权重矩阵中,部分元素已被置为 0,其余元素则被重新缩放:

复制代码
tensor([[2.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000],
        [0.0000, 0.8966, 0.0000, 0.0000, 0.0000, 0.0000],
        [0.0000, 0.0000, 0.6206, 0.0000, 0.0000, 0.0000],
        [0.5517, 0.4921, 0.0000, 0.0000, 0.0000, 0.0000],
        [0.4350, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000],
        [0.0000, 0.3327, 0.0000, 0.0000, 0.0000, 0.0000]],
       grad_fn=<MulBackward0>)

请注意,由于操作系统的差异,最终的 dropout 输出可能会有所不同。有关这种不一致性的详细信息,请查看 PyTorch 的 issue tracker。

理解了因果注意力和 dropout 掩码之后,现在我们可以开发一个简洁的 Python 类。这个类的目的是高效地应用这两种技术。

3.5.3 实现一个简化的因果注意力类

现在我们将把因果注意力和 dropout 修改应用到 3.4 节中开发的 SelfAttention Python 类 中。这个类将成为开发多头注意力的基础,而多头注意力是我们最终实现的注意力类。

但在开始之前,需要确保代码可以处理包含多个样本的批次,以便 CausalAttention 类能够支持第 2 章中实现的数据加载器产生的批量输出。

为简单起见,可以通过复制输入文本示例来模拟批量输入

python 复制代码
batch = torch.stack((inputs, inputs), dim=0)
print(inputs)
print(batch.shape) # 2 inputs with 6 tokens each, and each token has embedding dimension 3
print(batch)

这将生成一个三维张量,其中包含两个输入文本,每个文本有 6 个词元,每个词元是一个三维的 嵌入向量。

复制代码
tensor([[0.4300, 0.1500, 0.8900],
        [0.5500, 0.8700, 0.6600],
        [0.5700, 0.8500, 0.6400],
        [0.2200, 0.5800, 0.3300],
        [0.7700, 0.2500, 0.1000],
        [0.0500, 0.8000, 0.5500]])
torch.Size([2, 6, 3])
tensor([[[0.4300, 0.1500, 0.8900],
         [0.5500, 0.8700, 0.6600],
         [0.5700, 0.8500, 0.6400],
         [0.2200, 0.5800, 0.3300],
         [0.7700, 0.2500, 0.1000],
         [0.0500, 0.8000, 0.5500]],

        [[0.4300, 0.1500, 0.8900],
         [0.5500, 0.8700, 0.6600],
         [0.5700, 0.8500, 0.6400],
         [0.2200, 0.5800, 0.3300],
         [0.7700, 0.2500, 0.1000],
         [0.0500, 0.8000, 0.5500]]])

代码清单 3-3 中的 CausalAttention 类与先前实现的 SelfAttention 类类似,唯一的不 同是增加了 dropout 和因果掩码组件。

python 复制代码
class CausalAttention(nn.Module):
    def __init__(self, d_in, d_out, context_length,
                 dropout, qkv_bias=False):
        super().__init__()
        self.d_out = d_out
        self.W_query = nn.Linear(d_in, d_out, bias=qkv_bias)
        self.W_key   = nn.Linear(d_in, d_out, bias=qkv_bias)
        self.W_value = nn.Linear(d_in, d_out, bias=qkv_bias)
        self.dropout = nn.Dropout(dropout)  # 新增
        self.register_buffer('mask', torch.triu(torch.ones(context_length, context_length), diagonal=1))  # 新增
    def forward(self, x):
        b, num_tokens, d_in = x.shape  # 新增的批次维度 b
        # 对于 `num_tokens` 超过 `context_length` 的输入,下面的掩码创建过程会出错。
        # 实际上这并不是问题,因为 LLM(第 4-7 章)会确保输入在到达此 forward 方法之前
        # 不会超过 `context_length`。
        keys = self.W_key(x)
        queries = self.W_query(x)
        values = self.W_value(x)
        attn_scores = queries @ keys.transpose(1, 2)  # 修改后的 transpose
        attn_scores.masked_fill_(  # 新增,带 _ 的操作为原地操作
            self.mask.bool()[:num_tokens, :num_tokens], -torch.inf)  # `:num_tokens` 用于处理批次中 token 数量小于所支持的 context_size 的情况
        attn_weights = torch.softmax(
            attn_scores / keys.shape[-1]**0.5, dim=-1
        )
        attn_weights = self.dropout(attn_weights)  # 新增
        context_vec = attn_weights @ values
        return context_vec
torch.manual_seed(123)
context_length = batch.shape[1]
ca = CausalAttention(d_in, d_out, context_length, 0.0) # 这里官方代码给的dropout是0.0,所以看不到效果,换成0.5即可print(attn_weights)看到效果。
context_vecs = ca(batch)
print(context_vecs)
print("context_vecs.shape:", context_vecs.shape)

虽然此时所有新增的代码行都应该是熟悉的, 但我们在 __init__ 方法中增加了一个self. register_buffer()调用。虽然在 PyTorch 中使用 register_buffer 并非所有情况下都是必需的,但在这里具有一些优势。例如,当我们在大语言模型中使用 CausalAttention 类时,缓冲区会与模型一起自动移动到适当的设备(CPU 或 GPU),这在训练大语言模型时非常重要。这意味着我们无须手动确保这些张量与模型参数在同一设备上,从而避免了设备不匹配的错误。

register_buffer 是什么

register_buffernn.Module 提供的方法,用来把一个张量注册为模块的"缓冲区(buffer)"。缓冲区是模块状态的一部分,但不是可训练参数

代码中:

python 复制代码
self.register_buffer('mask', torch.triu(torch.ones(context_length, context_length), diagonal=1))

这把上三角掩码矩阵注册进了模块,之后可以像普通属性一样用 self.mask 使用它。

参数(parameter)与缓冲区(buffer)的区别

参数 (nn.Parameter) 缓冲区 (buffer)
参与梯度更新 是,优化器会更新 否,训练中保持不变
.to(device) 移动
出现在 state_dict 是(默认)
典型例子 权重、偏置 掩码、BatchNorm 的 running_mean/var

mask 是固定常量、不需要学习,所以不应是参数;但它又是运行时必需的状态,所以用缓冲区最合适。

为什么需要用 register_buffer

  1. 自动跟随设备移动:调用 model.to("cuda") 时,PyTorch 会自动把参数和缓冲区 一起搬到 GPU。若写成普通属性 self.mask = torch.triu(...),它会留在 CPU 上,当输入 x 在 GPU 时 masked_fill_ 会因设备不匹配报错。
  2. 随模型一起保存和加载:缓冲区默认出现在 state_dict() 中,torch.save / load_state_dict 会一并处理。(可从 context_length 推导的常量,也可用 persistent=False 排除保存。)
  3. 明确表达语义:清楚表明这是模型状态的一部分,但不是可训练的权重。
    transpose 方法

tensor.transpose(dim0, dim1) 的作用是交换指定的两个维度,其他维度保持不动。

python 复制代码
x = torch.randn(2, 3, 4)
x.transpose(1, 2).shape   # (2, 4, 3)  只交换了 dim1 和 dim2

它和 .T 的区别:.T 会反转所有 维度,而 transpose 只精确交换你指定的那两个。

为什么这里用 transpose(1, 2)

加了批次维度后,keys 是三维的:

复制代码
keys.shape == (b, num_tokens, d_out)

注意力打分需要 queries @ keys^T,但批次维度 dim0 要保持不动 ,只交换后面的 num_tokensd_out

python 复制代码
keys.transpose(1, 2)   # (b, d_out, num_tokens)  ✓
keys.T                 # (d_out, num_tokens, b)   ✗ 批次维度被打乱

这样相乘得到每个 batch 独立的注意力矩阵:

复制代码
(b, num_tokens, d_out) @ (b, d_out, num_tokens) = (b, num_tokens, num_tokens)

所以用 transpose(1, 2) 而不是 .T,是为了只转置最后两个维度、保留批次维度。
attn_scores.masked_fill_(self.mask.bool():num_tokens, :num_tokens, -torch.inf)为什么要有num_tokes

因为:

  • self.mask__init__ 里就被永久固定成 6x6
  • attn_scores 的大小由实际输入长度 决定,比如是 3x3
  • masked_fill_ 要求 mask 和被填充的张量形状能对上,3x3 对不上 6x6,于是崩溃

改成官方写法 self.mask.bool()[:num_tokens, :num_tokens] 后:

  • 第二次调用时 num_tokens=3,从 6x6 里切出左上角 3x3
  • 尺寸对上了,正常运行

为什么真实场景一定会遇到长度不一样的输入?

context_length 是模型支持的最大 长度(比如 GPT 的 1024),你只在建模型时定一次。但推理/训练时喂进来的句子长度是变化的------有时 10 个词,有时 50 个词,很少正好等于最大值。模型不可能每来一个新长度就重建一次 mask,所以做法是:一次性按最大长度建好 mask,每次调用时按实际长度切一小块用

所以attn_scores.masked_fill_( self.mask.bool(), -torch.inf) 就行了啊不是"错",而是"不通用"------它只在输入长度恰好等于 context_length 时才成立。切片是为了让同一个模型能安全处理任意 ≤ context_length 的输入。

下面画图说明。

原始的 6×6 mask(__init__ 里按 context_length=6 建好的,1=要屏蔽的未来位置):

复制代码
          列0  列1  列2  列3  列4  列5
        ┌────────────────────────────────┐
   行0  │  0    1    1  │  1    1    1     │
   行1  │  0    0    1  │  1    1    1     │
   行2  │  0    0    0  │  1    1    1     │
        ├───────────────┘                 │
   行3  │  0    0    0     0    1    1     │
   行4  │  0    0    0     0    0    1     │
   行5  │  0    0    0     0    0    0     │
        └────────────────────────────────┘

当输入只有 3 个词时,执行 self.mask[:3, :3],只取左上角虚线框住的部分:

复制代码
          列0  列1  列2
        ┌────────────────┐
   行0  │  0    1    1   │
   行1  │  0    0    1   │
   行2  │  0    0    0   │
        └────────────────┘

这个 3×3 的小 mask 正好能对上 3 个词产生的 3×3 注意力分数。

复制代码
tensor([[[-0.4519,  0.2216],
         [-0.5874,  0.0058],
         [-0.6300, -0.0632],
         [-0.5675, -0.0843],
         [-0.5526, -0.0981],
         [-0.5299, -0.1081]],

        [[-0.4519,  0.2216],
         [-0.5874,  0.0058],
         [-0.6300, -0.0632],
         [-0.5675, -0.0843],
         [-0.5526, -0.0981],
         [-0.5299, -0.1081]]], grad_fn=<UnsafeViewBackward0>)
context_vecs.shape: torch.Size([2, 6, 2])

图 3-23 总结了目前我们所取得的进展。我们集中讨论了神经网络中因果注意力的概念和 实现。接下来,我们将基于这一概念,开发一个并行实现了多个因果注意力机制的多头注意力 模块。

3.6 将单头注意力扩展到多头注意力

在本节中,我们将进行最后一步操作,即把先前实现的因果注意力类扩展到多个头上。这也被称为多头注意力。

"多头"这一术语指的是将注意力机制分成多个"头",每个"头"独立工作。在这种情况下, 单个因果注意力模块可以被看作单头注意力,因为它只有一组注意力权重按顺序处理输入。

我们将从因果注意力扩展到多头注意力。首先,我们将直观地通过堆叠多个 CausalAttention 模块来构建多头注意力模块。然后,我们将用一种更复杂但计算上更高效的方式来实现这个多头注意力模块。

3.6.1 叠加多个单头注意力层

在实际操作中,实现多头注意力需要构建多个自注意力机制的实例(参见图 3-18),每个实例都有其独立的权重,然后将这些输出进行合成。虽然这种方法的计算量可能会非常大,但它对诸如基于 Transformer 的大语言模型之类的模型的复杂模式识别是非常重要的。

图 3-24 展示了多头注意力模块的结构,它是由图 3-18 所示的多个单头注意力模块依次叠加在一起组成的。

正如前面提到的,多头注意力的主要思想是多次(并行)运行注意力机制,每次使用学到的不同的线性投影------这些投影是通过将输入数据(比如注意力机制中的查询向量、键向量和值向量)乘以权重矩阵得到的。在代码中,可以通过实现一个简单的 MultiHeadAttentionWrapper 类来达到这一目标,MultiHeadAttentionWrapper 类堆叠了多个之前实现的 CausalAttention模块实例,如代码清单 3-4 所示。

python 复制代码
class MultiHeadAttentionWrapper(nn.Module):

    def __init__(self, d_in, d_out, context_length, dropout, num_heads, qkv_bias=False):
        super().__init__()
        self.heads = nn.ModuleList(
            [CausalAttention(d_in, d_out, context_length, dropout, qkv_bias) 
             for _ in range(num_heads)]
        )

    def forward(self, x):
        return torch.cat([head(x) for head in self.heads], dim=-1)


torch.manual_seed(123)

context_length = batch.shape[1] # This is the number of tokens
d_in, d_out = 3, 2
mha = MultiHeadAttentionWrapper(
    d_in, d_out, context_length, 0.0, num_heads=2
)

context_vecs = mha(batch)

print(context_vecs)
print("context_vecs.shape:", context_vecs.shape)

如果采用这个具有两个注意力头( num_heads=2 )以及 CausalAttention 输出维度为 d_out=2 的 MultiHeadAttentionWrapper 类, 那么我们就会得到一个四维的上下文向量 (d_out*num_heads=4),如图 3-25 所示。

多头注意力包装类(MultiHeadAttentionWrapper)用一个 nn.ModuleList 装了多个独立的 CausalAttention 头,每个头单独学习不同的注意力模式。

nn.ModuleList 是专门用来存放子模块的列表,能让 PyTorch 正确地注册和管理这些模块的参数。

forward 中,[head(x) for head in self.heads] 让每个头分别处理输入 x。因为 CausalAttention 继承自 nn.Module,所以 head(x) 会通过 nn.Module__call__ 自动调用 forward 方法(不要直接写 head.forward(x),会跳过 PyTorch 的 hook 机制)。

每个头输出形状为 (batch, num_tokens, d_out)torch.cat(..., dim=-1) 把它们沿最后一个维度(特征维度)拼接起来。

torch.cat(tensors, dim) 用于拼接张量,除拼接维度外其他维度形状必须一致;它只在已有维度上扩展,不新增维度(新增维度要用 torch.stack)。

所以 2 个头、每个 d_out=2,拼接后最后一维变成 num_heads * d_out = 4,最终 context_vecs.shape(batch, num_tokens, 4)

复制代码
tensor([[[-0.4519,  0.2216,  0.4772,  0.1063],
         [-0.5874,  0.0058,  0.5891,  0.3257],
         [-0.6300, -0.0632,  0.6202,  0.3860],
         [-0.5675, -0.0843,  0.5478,  0.3589],
         [-0.5526, -0.0981,  0.5321,  0.3428],
         [-0.5299, -0.1081,  0.5077,  0.3493]],

        [[-0.4519,  0.2216,  0.4772,  0.1063],
         [-0.5874,  0.0058,  0.5891,  0.3257],
         [-0.6300, -0.0632,  0.6202,  0.3860],
         [-0.5675, -0.0843,  0.5478,  0.3589],
         [-0.5526, -0.0981,  0.5321,  0.3428],
         [-0.5299, -0.1081,  0.5077,  0.3493]]], grad_fn=<CatBackward0>)
context_vecs.shape: torch.Size([2, 6, 4])

结果中的 context_vecs 张量的第一维是 2,因为我们有两个输入文本(输入文本是重复的,所以这些上下文向量完全相同)。第二维表示每个输入中的 6 个词元。第三维表示每个词元的四维嵌入。

练习 3.2

返回二维嵌入向量

更改 MultiHeadAttentionWrapper(..., num_heads=2)调用的输入参数,使输出上下文向量是二维而不是四维,同时保持设置 num_heads=2。(提示:不需要修改类实现,只需要改变另一个输入参数。)

If we want to have an output dimension of 2, as earlier in single-head attention, we can have to change the projection dimension d_out to 1:

python 复制代码
torch.manual_seed(123)

d_out = 1
mha = MultiHeadAttentionWrapper(d_in, d_out, context_length, 0.0, num_heads=2)

context_vecs = mha(batch)

print(context_vecs)
print("context_vecs.shape:", context_vecs.shape)
复制代码
tensor([[[-0.5740,  0.2216],
         [-0.7320,  0.0155],
         [-0.7774, -0.0546],
         [-0.6979, -0.0817],
         [-0.6538, -0.0957],
         [-0.6424, -0.1065]],

        [[-0.5740,  0.2216],
         [-0.7320,  0.0155],
         [-0.7774, -0.0546],
         [-0.6979, -0.0817],
         [-0.6538, -0.0957],
         [-0.6424, -0.1065]]], grad_fn=<CatBackward0>)
context_vecs.shape: torch.Size([2, 6, 2])

到目前为止,我们已经实现了一个将多个单头注意力模块结合起来的 MultiHeadAttentionWrapper。不过,这些模块在 forward 方法中是通过head(x) for head in self.heads 依次处理的。我们可以通过并行处理所有的注意力头来改进这个实现。一种方法是在使用矩阵乘法的同时计算所有注意力头的输出。

3.6.2 通过权重划分实现多头注意力

到目前为止,我们已经创建了一个 MultiHeadAttentionWrapper,通过叠加多个单头注 意力模块来实现多头注意力。这是通过实例化并组合多个 CausalAttention 对象来完成的。

与其维护两个单独的类 MultiHeadAttentionWrapper 和 CausalAttention,不如将这两个概念合并成一个 MultiHeadAttention 类。此外,除了将 MultiHeadAttentionWrapper 与 CausalAttention 代码合并,我们还会进行一些其他调整,以更高效地实现多头注意力。

在 MultiHeadAttentionWrapper 中,多头机制通过创建 CausalAttention 对象的列表 (self.heads)来实现,每个对象代表一个独立的注意力头。CausalAttention 类单独执行注意力机制,每个头的结果会被拼接。相比之下,下面的 MultiHeadAttention 类会将多头功能整合到一个类内。它通过重新调整投影后的查询张量、键张量和值张量的形状,将输入分为多个头,然后在计算注意力后合并这些头的结果。

在进一步讨论之前,先来看一下 MultiHeadAttention 类,如代码清单 3-5 所示。

python 复制代码
class MultiHeadAttention(nn.Module):
    def __init__(self, d_in, d_out, context_length, dropout, num_heads, qkv_bias=False):
        super().__init__()
        assert (d_out % num_heads == 0), \
            "d_out 必须能被 num_heads 整除"

        self.d_out = d_out
        self.num_heads = num_heads
        self.head_dim = d_out // num_heads  # 减小投影维度以匹配期望的输出维度

        self.W_query = nn.Linear(d_in, d_out, bias=qkv_bias)
        self.W_key = nn.Linear(d_in, d_out, bias=qkv_bias)
        self.W_value = nn.Linear(d_in, d_out, bias=qkv_bias)
        self.out_proj = nn.Linear(d_out, d_out)  # 用于合并各注意力头输出的线性层
        self.dropout = nn.Dropout(dropout)
        self.register_buffer(
            "mask",
            torch.triu(torch.ones(context_length, context_length),
                       diagonal=1)
        )

    def forward(self, x):
        b, num_tokens, d_in = x.shape
        # 与 `CausalAttention` 一样,当输入的 `num_tokens` 超过 `context_length` 时,
        # 下面的掩码创建会产生错误。
        # 实际使用中这不是问题,因为 LLM(第 4-7 章)会确保输入
        # 在到达此 forward 方法之前不会超过 `context_length`。

        keys = self.W_key(x)  # 形状: (b, num_tokens, d_out)
        queries = self.W_query(x)
        values = self.W_value(x)

        # 我们通过添加一个 `num_heads` 维度来隐式地拆分矩阵
        # 展开最后一维: (b, num_tokens, d_out) -> (b, num_tokens, num_heads, head_dim)
        keys = keys.view(b, num_tokens, self.num_heads, self.head_dim)
        values = values.view(b, num_tokens, self.num_heads, self.head_dim)
        queries = queries.view(b, num_tokens, self.num_heads, self.head_dim)

        # 转置: (b, num_tokens, num_heads, head_dim) -> (b, num_heads, num_tokens, head_dim)
        keys = keys.transpose(1, 2)
        queries = queries.transpose(1, 2)
        values = values.transpose(1, 2)

        # 使用因果掩码计算缩放点积注意力(即自注意力)
        attn_scores = queries @ keys.transpose(2, 3)  # 对每个头做点积

        # 将原始掩码截断到 token 数量并转换为布尔类型
        mask_bool = self.mask.bool()[:num_tokens, :num_tokens]

        # 使用掩码填充注意力分数
        attn_scores.masked_fill_(mask_bool, -torch.inf)
        attn_weights = torch.softmax(attn_scores / keys.shape[-1]**0.5, dim=-1)
        attn_weights = self.dropout(attn_weights)

        # 形状: (b, num_tokens, num_heads, head_dim)
        context_vec = (attn_weights @ values).transpose(1, 2)
        # 合并各头,其中 self.d_out = self.num_heads * self.head_dim
        context_vec = context_vec.contiguous().view(b, num_tokens, self.d_out)
        context_vec = self.out_proj(context_vec)  # 可选的投影

        return context_vec


torch.manual_seed(123)

batch_size, context_length, d_in = batch.shape
d_out = 2
mha = MultiHeadAttention(d_in, d_out, context_length, 0.0, num_heads=2)

context_vecs = mha(batch)

print(context_vecs)
print("context_vecs.shape:", context_vecs.shape)

尽管 MultiHeadAttention 类中的张量重塑(.view)和转置(.transpose)在数学上看起来很复杂,但 MultiHeadAttention 类实现的概念仍与之前的 MultiHeadAttentionWrapper 类相同。

关于.view的方法

一个张量 = 数据本体 + 元信息。无论 numpy 的 ndarray 还是 PyTorch 的 tensor,它在内存里其实是两部分:数据本体是一条连续的一维数字直线;元信息(metadata)是一小组描述"该如何解读这条直线"的说明,本身很小,就是几个数字。

关键的元信息有三个:shape(形状,每个维度有多少个元素)、stride(步长,在每个维度上走一步需要在数据直线上跳过几个元素)、以及一个指向数据本体起始位置的指针(offset)。数据本体可能很大,但元信息永远只是那几个小数字。很多操作(view、transpose、切片)根本不碰数据本体,只改元信息,这就是它们快的原因。

以形状 (2, 3) 的张量为例:

复制代码
逻辑视图(你看到的):
  [[10, 11, 12],
   [13, 14, 15]]

数据本体(内存直线,永不改变):
  索引:  0    1    2    3    4    5
  值:    10   11   12   13   14   15

shape  = (2, 3)
stride = (3, 1)

stride (3, 1) 的意思是:在第 0 维(行)走一步跳 3 个位置,在第 1 维(列)走一步跳 1 个位置。验证取元素 [1, 2](应是 15):直线位置 = 1×3 + 2×1 = 5,直线5 = 15 ✓。所以"如何解读直线"完全由 shape + stride 决定。

view 做的事:只改元信息,不碰数据本体。把上面的 (2, 3) view 成 (3, 2)

复制代码
数据本体: 10 11 12 13 14 15   ← 一个字节都没变

变前:  shape=(2, 3),  stride=(3, 1)
变后:  shape=(3, 2),  stride=(2, 1)

读法(新 stride):
  行0: [10, 11]
  行1: [12, 13]
  行2: [14, 15]

view 的本质就是:换一套 shape 和 stride,让同一条数据直线呈现出新形状,数据本体原地不动。

在多头注意力代码里 keys.view(b, num_tokens, num_heads, head_dim) 是"拆分维度"。只看一个 token 的 8 个数(d_out=8num_heads=2, head_dim=4):

复制代码
数据本体:
  索引:  0    1    2    3    4    5    6    7
  值:    q0   q1   q2   q3   q4   q5   q6   q7

view 成 (2, 4),stride = (4, 1):
  头0(从索引0起取4个): [q0, q1, q2, q3]
  头1(从索引4起取4个): [q4, q5, q6, q7]

数据本体还是那 8 个数、还是那个顺序,只是新的 (shape, stride) 把它解读成了两个头。这就是"隐式拆分"------只动元信息。


为什么 view 有时会报错?根源在 stride。

"标准行优先 stride"就是一个固定的两步计算公式,从最后一维开始往前算:

第 1 步:最后一维的 stride 永远设为 1。因为最后一维(列)是挨着放的------内存直线里 10、11、12 就是紧挨着的,走一步跳 1 个,天经地义。

第 2 步:往前一维的 stride = 后一维的 stride × 后一维的长度。

复制代码
shape = (2, 3)
第 1 维(列)stride = 1
第 0 维(行)stride = 1 × 3 = 3   (要跨到下一行,得先跳过当前行的 3 个元素)
→ stride = (3, 1)

为什么这个公式算出来的 stride 就等于"连续"?"连续"的定义是:按内存直线从头读到尾(0→1→2→3→4→5),正好等于按逻辑顺序从头读到尾(先读完第 0 行,再读完第 1 行)。用 stride=(3,1) 挨个验证逻辑顺序落在直线的哪个位置:

复制代码
逻辑元素      计算位置 = 行×3 + 列×1     落在直线索引
[0,0]=10      0×3 + 0×1 = 0                0
[0,1]=11      0×3 + 1×1 = 1                1
[0,2]=12      0×3 + 2×1 = 2                2
[1,0]=13      1×3 + 0×1 = 3                3
[1,1]=14      1×3 + 1×1 = 4                4
[1,2]=15      1×3 + 2×1 = 5                5

最右列 0,1,2,3,4,5------正好是内存直线从头到尾、一个不跳、一个不重。这就是"连续":逻辑顺序和内存顺序完全吻合。只要 stride 是用标准公式算出来的,就必然吻合,所以把"stride 等于标准行优先公式的结果"直接当作"连续"的判定标准。


transpose 交换两个维度,同时把两个 stride 也交换,但数据本体不动:

复制代码
transpose 前:  shape=(2,3)  stride=(3,1)
transpose 后:  shape=(3,2)  stride=(1,3)   ← stride 也交换了
数据本体:      10 11 12 13 14 15  (没动)

转置后 shape=(3,2) 的标准 stride 本应是:最后一维 stride=1,前一维 stride=1×2=2,即 (2, 1)。但实际 stride 是 (1, 3),不等于标准值,说明它非连续。用实际 stride (1, 3) 验证逻辑顺序落在直线哪里:

复制代码
转置后逻辑视图:
  [[10, 13],
   [11, 14],
   [12, 15]]

逻辑元素      计算位置 = 行×1 + 列×3     落在直线索引
[0,0]=10      0×1 + 0×3 = 0                0
[0,1]=13      0×1 + 1×3 = 3                3   ← 跳到 3
[1,0]=11      1×1 + 0×3 = 1                1   ← 又跳回 1
[1,1]=14      1×1 + 1×3 = 4                4
[2,0]=12      2×1 + 0×3 = 2                2   ← 又跳回 2
[2,1]=15      2×1 + 1×3 = 5                5

最右列 0,3,1,4,2,5------跳着走,不是顺序的。逻辑顺序和内存顺序对不上,这就是非连续。


为什么非连续时 view 会报错?view 的工作方式是:你给它一个新 shape,它就用标准公式算出一套新 stride,然后套到现有数据本体上去读。但标准公式算出的 stride,只有当数据本体本身"按逻辑顺序连续排列"时才读得对。转置后的张量,数据本体顺序(10,11,12,13,14,15)和逻辑顺序(10,13,11,14,12,15)已经不一致,这时再套按逻辑顺序设计的标准 stride,读出来就是错乱的。PyTorch 不愿给出错误结果,于是直接报错,逼你先 .contiguous()

.contiguous() 会新开一块内存,把数据本体按当前逻辑顺序真正重新写一遍:

复制代码
contiguous 后的数据本体:
  索引:  0    1    2    3    4    5
  值:    10   13   11   14   12   15
  stride = (2, 1)  ← 又变回标准行优先了

现在数据顺序和逻辑顺序重新一致,stride 恢复标准形式,view 就能安全工作。所以代码里 context_vec.contiguous().view(b, num_tokens, self.d_out):前面 transpose 把它弄成非连续,先 contiguous 复原成连续,再 view。


和 numpy 的区别:两者都是"数据本体 + 元信息(shape, stride)"的结构,区别只在遇到"没法只靠改元信息完成变形"时的态度。numpy 的 reshape 会悄悄复制一份数据、重排好返回;PyTorch 的 view 承诺绝不复制,做不到就报错;PyTorch 的 reshape 行为像 numpy,能不复制就不复制、不行就复制。所以 PyTorch 里 view 是"保证零复制"的严格版,reshape 是"帮我搞定就行"的智能版。

在整体层面上,在之前的 MultiHeadAttentionWrapper 中,我们堆叠了多个单头注意力 层,并将其合并成一个多头注意力层。MultiHeadAttention 类采用了一种综合的方法。它从一个多头注意力层开始,然后在内部将这个层分割成单独的注意力头,如图 3-26 所示。

在 PyTorch 中,通过使用.view 方法进行张量重塑以及使用.transpose 方法进行张量转置, 我们实现了对查询张量、键张量和值张量的分割。输入首先经过线性层进行变换(针对查询矩阵、键矩阵和值矩阵),然后被重塑为多个头。

关键操作是将 d_out 维度分割为 num_heads 和 head_dim,其中 head_dim = d_out / num_heads。 此分割通过.view 方法来实现:维度为(b, num_tokens, d_out)的张量被重塑后的维度为(b, num_tokens, num_heads, head_dim)。 然后转置张量,使 num_heads 维度置于 num_tokens 维度之前,从而形成一个(b, num_heads, num_tokens, head_dim)的形状。 这种转置对于正确对齐不同头的查询矩阵、键矩阵和值矩阵, 以及有效地执行批处理矩阵乘法至关重要。

多头注意力完整图解

设定:

  • batch = 2(两个样本,同一句话复制两份)
  • num_tokens = 6(6 个词:Your, journey, starts, with, one, step)
  • token_dims = d_in = 3(每个词输入是 3 维)
  • d_out = 2(输出维度)
  • num_heads = 2(2 个头)
  • head_dim = d_out / num_heads = 1(每个头只有 1 维)

下面详细图解只画 batch 里第 1 个样本,第 2 个样本形状完全一样、并行处理。


步骤 0:输入

复制代码
inputs  →  (6, 3)          一个样本

            d_in = 3
            ┌────────────────────┐
  Your    → │ 0.43  0.15  0.89   │
  journey → │ 0.55  0.87  0.66   │
  starts  → │ 0.57  0.85  0.64   │
  with    → │ 0.22  0.58  0.33   │
  one     → │ 0.77  0.25  0.10   │
  step    → │ 0.05  0.80  0.55   │
            └────────────────────┘

堆叠成 batch  →  (2, 6, 3) = (batch, num_tokens, d_in)

步骤 1:QKV 投影(每个投影矩阵 W 形状 (d_in, d_out) = (3, 2))

复制代码
x (2, 6, 3)  @  W_query (3, 2)  =  queries (2, 6, 2)

对 K、V 同理:
  queries (2, 6, 2)
  keys    (2, 6, 2)
  values  (2, 6, 2)

每个词从 3 维被投影成 2 维(d_out)。单样本:

queries (6, 2)
            d_out = 2
            ┌──────────┐
  Your    → │ q00  q01 │
  journey → │ q10  q11 │
  starts  → │ q20  q21 │
  with    → │ q30  q31 │
  one     → │ q40  q41 │
  step    → │ q50  q51 │
            └──────────┘

步骤 2:view 拆分多头(把最后一维 2 = num_heads(2) × head_dim(1) 拆开,内存不变,只换读法)

复制代码
queries (2, 6, 2)  →  view  →  (2, 6, 2, 1)
                = (batch, num_tokens, num_heads, head_dim)

单样本 (6, 2, 1):
          head0    head1
         ┌─────┐  ┌─────┐
  Your → │ q00 │  │ q01 │
  jour   │ q10 │  │ q11 │
  star   │ q20 │  │ q21 │
  with   │ q30 │  │ q31 │
  one    │ q40 │  │ q41 │
  step   │ q50 │  │ q51 │
         └─────┘  └─────┘

因为 head_dim=1,每个头就是原来 2 维向量里的一个数字:第 0 列归 head0,第 1 列归 head1。

步骤 3:transpose(1,2) 把 head 提前

python 复制代码
queries = queries.transpose(1, 2)
# (2, 6, 2, 1) → (2, 2, 6, 1)
# (batch, num_tokens, num_heads, head_dim) → (batch, num_heads, num_tokens, head_dim)
复制代码
单样本"按头分组",每个头拿到完整的 (6, 1) 矩阵:

head0 (6, 1):          head1 (6, 1):
       ┌─────┐              ┌─────┐
  Your │ q00 │         Your │ q01 │
  jour │ q10 │         jour │ q11 │
  star │ q20 │         star │ q21 │
  with │ q30 │         with │ q31 │
  one  │ q40 │         one  │ q41 │
  step │ q50 │         step │ q51 │
       └─────┘              └─────┘

关键:@ 只对最后两维做矩阵乘,前面维度当批次并行。把 num_heads 提到前面,两个头就能独立并行计算。


步骤 4:计算注意力分数 Q @ Kᵀ(keys 转置最后两维:(2, 2, 6, 1)(2, 2, 1, 6))

复制代码
每个头内部:
    Q          Kᵀ                attn_scores
 (6, 1)  @   (1, 6)      =         (6, 6)

                每个词 q 和每个词 k 的点积
                ┌─────────────────────────┐
                │ s00 s01 s02 s03 s04 s05 │ ← Your 对全部 6 个词
                │ s10 s11 s12 s13 s14 s15 │ ← journey ...
                │ ...                     │
                │ s50 s51 s52 s53 s54 s55 │ ← step ...
                └─────────────────────────┘

两个头同时算:attn_scores → (2, 2, 6, 6) = (batch, num_heads, num_tokens, num_tokens)

步骤 5:缩放 + 因果掩码 + softmax + dropout

python 复制代码
attn_scores / keys.shape[-1]**0.5   # keys.shape[-1] = head_dim = 1,本例除以 √1=1
复制代码
mask (6, 6):                      应用后(上三角置 -inf):
┌────────────────┐          ┌──────────────────────────────┐
│ 0 1 1 1 1 1    │          │ s00 -inf -inf -inf -inf -inf │
│ 0 0 1 1 1 1    │          │ s10  s11 -inf -inf -inf -inf │
│ 0 0 0 1 1 1    │ ─fill──→ │ s20  s21  s22 -inf -inf -inf │
│ 0 0 0 0 1 1    │          │ s30  s31  s32  s33 -inf -inf │
│ 0 0 0 0 0 1    │          │ s40  s41  s42  s43  s44 -inf │
│ 0 0 0 0 0 0    │          │ s50  s51  s52  s53  s54  s55 │
└────────────────┘          └──────────────────────────────┘

含义:每个词只能看到自己和它前面的词。
mask 是 2 维,分数是 4 维 (2,2,6,6),靠广播自动应用到所有 batch、所有 head。
softmax(最后一维归一化,-inf 变 0)再 dropout,形状不变:attn_weights → (2, 2, 6, 6)

步骤 6:注意力权重 @ V

python 复制代码
context_vec = attn_weights @ values
复制代码
每个头内部:
 attn_weights      V            output
   (6, 6)   @   (6, 1)   =     (6, 1)

context_vec → (2, 2, 6, 1) = (batch, num_heads, num_tokens, head_dim)

单样本:
head0 (6, 1):        head1 (6, 1):
       ┌─────┐              ┌─────┐
  Your │ c00 │         Your │ c01 │
  jour │ c10 │         jour │ c11 │
  ...  │ ... │         ...  │ ... │
  step │ c50 │         step │ c51 │
       └─────┘              └─────┘

步骤 7:transpose 还原 + view 合并

python 复制代码
context_vec = context_vec.transpose(1, 2)
# (2, 2, 6, 1) → (2, 6, 2, 1)
# (batch, num_heads, num_tokens, head_dim) → (batch, num_tokens, num_heads, head_dim)
复制代码
单样本:      head0  head1
  Your    → │ c00 │ c01 │
  journey → │ c10 │ c11 │
  starts  → │ c20 │ c21 │
  with    → │ c30 │ c31 │
  one     → │ c40 │ c41 │
  step    → │ c50 │ c51 │
python 复制代码
context_vec = context_vec.contiguous().view(2, 6, 2)
# → (2, 6, 2) = (batch, num_tokens, d_out)
复制代码
单样本 (6, 2):
            d_out = 2
            ┌──────────┐
  Your    → │ c00  c01 │
  journey → │ c10  c11 │
  starts  → │ c20  c21 │
  with    → │ c30  c31 │
  one     → │ c40  c41 │
  step    → │ c50  c51 │
            └──────────┘
           head0 head1

contiguous() 原因:transpose 只改了读取顺序(stride),内存没真正重排;view 要求内存连续,所以先 contiguous() 真正重排一遍。


步骤 8:out_proj 输出投影

python 复制代码
context_vec = self.out_proj(context_vec)
# (2, 6, 2) @ out_proj(2, 2) → (2, 6, 2),形状不变

为什么必须要有 out_proj?以 "Your" 为例:

复制代码
  Your → │ c00 │ c01 │
          head0  head1

问题:这只是把两个头的输出横向拼在一起,它们之间毫无信息交流。c00 完全由 head0 独立算出,c01 完全由 head1 独立算出,两者像两条平行线,互不相干。

打个比方:head0 和 head1 是两个专家,一个盯"语法关系",一个盯"语义关系"。它们各写了一份报告,现在只是把两份报告钉在一起,还没人把它们综合成结论。out_proj 就是那个"综合者",一个 (d_out, d_out) = (2, 2) 的全连接层。它输出的每一维都是所有头的加权组合:

复制代码
out_proj 对 "Your" 的计算:
  out[0] = w00·c00 + w01·c01 + b0
  out[1] = w10·c00 + w11·c01 + b1
           ↑head0  ↑head1
  两个头的信息被混合进了每一维

   拼接后(头之间隔离)         out_proj 后(头之间融合)
  ┌─────┬─────┐            ┌─────────────────┐
  │head0│head1│  ──proj──→ │ 每维都融合了      │
  │ c00 │ c01 │            │ 两个头的信息      │
  └─────┴─────┘            └─────────────────┘
   各干各的,无交流          交叉融合,信息互通

out_proj 的三个作用:

  1. 信息融合(最核心):让不同头学到的不同"视角"真正交互组合,而不是简单拼接。
  2. 可学习的重组:权重是训练出来的,模型能自己学会"如何搭配各头输出",哪个头重要就给更大权重。
  3. 维度衔接:保证输出是干净的 d_out,顺畅接入下一层(残差连接、前馈网络等)。

一句话:没有 out_proj,多头只是几个各干各的独立注意力被硬拼在一起;有了 out_proj,它们才真正"合作"成一个整体。


全流程形状速查表(书中例子):

步骤 操作 形状
0 输入 batch (2, 6, 3)
1 QKV 投影 (2, 6, 2)
2 view 拆头 (2, 6, 2, 1)
3 transpose(1,2) (2, 2, 6, 1)
4 Q @ Kᵀ (2, 2, 6, 6)
5 缩放/掩码/softmax/dropout (2, 2, 6, 6)
6 @ V (2, 2, 6, 1)
7 transpose 还原 + view 合并 (2, 6, 2)
8 out_proj (2, 6, 2)

主线:输入 (2, 6, 3) → 投影 (2, 6, 2) → 拆头 → 并行算注意力 → 合并 → out_proj 融合 → 输出 (2, 6, 2) = (batch, num_tokens, d_out)

为了说明这个批处理矩阵乘法,假设有下面所列的张量的例子,现在,在原始的张量和转置后的张量之间执行批处理矩阵乘法,其中我们转置了最后两个维度, 即 num_tokens 和 head_dim:

python 复制代码
# (b, num_heads, num_tokens, head_dim) = (1, 2, 3, 4)
a = torch.tensor([[[[0.2745, 0.6584, 0.2775, 0.8573],
                    [0.8993, 0.0390, 0.9268, 0.7388],
                    [0.7179, 0.7058, 0.9156, 0.4340]],

                   [[0.0772, 0.3565, 0.1479, 0.5331],
                    [0.4066, 0.2318, 0.4545, 0.9737],
                    [0.4606, 0.5159, 0.4220, 0.5786]]]])

print(a @ a.transpose(2, 3))
复制代码
tensor([[[[1.3208, 1.1631, 1.2879],
          [1.1631, 2.2150, 1.8424],
          [1.2879, 1.8424, 2.0402]],

         [[0.4391, 0.7003, 0.5903],
          [0.7003, 1.3737, 1.0620],
          [0.5903, 1.0620, 0.9912]]]])

在这种情况下,PyTorch 的矩阵乘法实现处理了四维输入张量,使得矩阵乘法在最后两个维度(num_tokens 和 head_dim)之间进行,并对每个头重复这一操作。 例如,上述方法提供了一种更简化的方式来单独计算每个头的矩阵乘法:

用实验验证一个关键事实: PyTorch 的批量矩阵乘法 a @ a.transpose(2, 3),和"手动把每个头单独拎出来做矩阵乘法",结果完全一样。

python 复制代码
first_head  = a[0, 0, :, :]      # 取出 head0,形状 (3, 4)
first_res   = first_head @ first_head.T   # (3,4)@(4,3) → (3,3)
print("First head:\n", first_res)

second_head = a[0, 1, :, :]      # 取出 head1,形状 (3, 4)
second_res  = second_head @ second_head.T # (3,3)
print("\nSecond head:\n", second_res)

最终结果与刚刚使用批处理矩阵乘法 print(a @ a.transpose(2, 3))得到的结果完全相同:

复制代码
First head:
 tensor([[1.3208, 1.1631, 1.2879],
        [1.1631, 2.2150, 1.8424],
        [1.2879, 1.8424, 2.0402]])

Second head:
 tensor([[0.4391, 0.7003, 0.5903],
        [0.7003, 1.3737, 1.0620],
        [0.5903, 1.0620, 0.9912]])

在 MultiHeadAttention 中,计算完注意力权重和上下文向量后,将所有头的上下文向量转置为(b, num_tokens, num_heads, head_dim)的形状。这些向量接着会被重塑(展平)为(b, num_tokens, d_out)的形状,从而有效地整合所有头的输出。

此外,我们在 MultiHeadAttention 中添加了一个输出投影层(self.out_proj),这是在合并多个头之后的步骤,而 CausalAttention 类中并不存在这个层。这个输出投影层并不是 必需的(更多细节参见附录 B),但它在许多大语言模型架构中被广泛使用,这就是我们在这里添加它以保持完整性的原因。

尽管 MultiHeadAttention 类因额外的张量重塑和转置显得比 MultiHeadAttentionWrapper 更复杂,但它的效率更高。原因是我们只需进行一次矩阵乘法来计算键矩阵,例如, keys = self.W_key(x)(查询矩阵和值矩阵也是如此)。在 MultiHeadAttentionWrapper 中,我们需要对每个注意力头重复进行这种矩阵乘法,而矩阵乘法是计算资源消耗较大的操作之一。

MultiHeadAttention 类的用法与我们之前实现的 SelfAttention 类和 CausalAttention 类类似:

【之前在写mha结构的地方也写了这个】

python 复制代码
torch.manual_seed(123)

batch_size, context_length, d_in = batch.shape
d_out = 2
mha = MultiHeadAttention(d_in, d_out, context_length, 0.0, num_heads=2)

context_vecs = mha(batch)

print(context_vecs)
print("context_vecs.shape:", context_vecs.shape)

结果显示,d_out 参数直接影响输出维度:

复制代码
tensor([[[0.3190, 0.4858],
         [0.2943, 0.3897],
         [0.2856, 0.3593],
         [0.2693, 0.3873],
         [0.2639, 0.3928],
         [0.2575, 0.4028]],

        [[0.3190, 0.4858],
         [0.2943, 0.3897],
         [0.2856, 0.3593],
         [0.2693, 0.3873],
         [0.2639, 0.3928],
         [0.2575, 0.4028]]], grad_fn=<ViewBackward0>)
context_vecs.shape: torch.Size([2, 6, 2])

现在,我们已经实现了将在实现和训练大语言模型时使用的 MultiHeadAttention 类。需要注意的是,尽管代码功能齐全,但为了保持输出的可读性,我们使用了相对较小的嵌入维度和注意力头数量。

相比之下,最小的 GPT-2 模型(参数量为 1.17 亿)有 12 个注意力头,上下文向量嵌入维度为 768,而最大的 GPT-2 模型(参数量为 15 亿)有 25 个注意力头,上下文向量嵌入维度为 1600。 请注意,在 GPT 模型中,词元输入和上下文嵌入的嵌入维度是相同的(d_in = d_out)。

练习 3.3 初始化 GPT-2 大小的注意力模块

使用 MultiHeadAttention 类初始化一个多头注意力模块,该模块应具有与最小的 GPT-2 模型相同数量的注意力头(12 个)。同时,确保使用与 GPT-2 相似的输入和输出嵌入维 度(768)。请注意,最小的 GPT-2 模型支持 1024 个词元的上下文长度。

python 复制代码
context_length = 1024
d_in, d_out = 768, 768
num_heads = 12

mha = MultiHeadAttention(d_in, d_out, context_length, 0.0, num_heads)

可选的参数量如下所示:

python 复制代码
def count_parameters(model):
    return sum(p.numel() for p in model.parameters() if p.requires_grad)

count_parameters(mha)
复制代码
2360064  # (2.36 M)

GPT-2 模型总共有 1.17 亿(117M)个参数,但正如我们所见,其大部分参数并不在多头注意力模块本身之中。

3.7 小结

  • 注意力机制可以将输入元素转换为增强的上下文向量表示,这些表示涵盖了关于所有输入的信息。
  • 自注意力机制通过对输入进行加权求和来计算上下文向量表示。
  • 在简化的注意力机制中,注意力权重是通过点积计算得出的。
  • 点积是两个向量的元素逐个相乘并将这些乘积相加的一种简洁计算方法。
  • 尽管矩阵乘法不是必需的,但它可以通过替代嵌套的 for 循环使计算更高效、更紧凑。
  • 在用于大语言模型的自注意力机制(也被称为"缩放点积注意力")中,我们引入了可训练的权重矩阵来计算输入的中间变换:查询矩阵、值矩阵和键矩阵。
  • 在处理从左到右读取和生成文本的大语言模型时,我们会添加一个因果注意力掩码,以防止模型访问未来的词元。
  • 除了使用因果注意力掩码将注意力权重置 0,还可以添加 dropout 掩码来减少大语言模型中的过拟合。
  • 基于 Transformer 的大语言模型中的注意力模块涉及多个因果注意力实例,这被称为"多头注意力"。
  • 可以通过堆叠多个因果注意力模块实例来创建多头注意力模块。
  • 创建多头注意力模块的一种更高效的方法是使用批量矩阵乘法。
相关推荐
奔跑中的小象1 小时前
统信UOS + 天数AI卡部署SGLang服务手册
人工智能·uos·sglang·天数智芯
DevSecOps选型指南1 小时前
中国版Mythos,为何是悬镜安全灵脉CodeAI?
人工智能·安全
Shockang2 小时前
LangGraph 状态机实战
人工智能
刹那芳华19922 小时前
循环神经网络的从零开始实现(RNN)
人工智能·rnn·深度学习
阿童木写作2 小时前
跨境图片翻译工具多合一,批量图片视频字幕翻译加智能抠图
人工智能·python·音视频·语音识别
科技之门2 小时前
AI3D从建模到贴图、绑骨和动画的完整流程怎么做?V2Fun完整工作流指南
人工智能·3d·贴图
宇宙第一小趴菜2 小时前
二、机器学习的应用领域和发展史
人工智能·机器学习
前端开发江鸟2 小时前
我写过 MCP Server,却一直以为 MCP 只有 Tool
人工智能
天国梦2 小时前
自习室智能化升级避坑指南:天学网AI智习室方案实测与选型建议
大数据·人工智能
阿里云云原生3 小时前
可用性从 99.9% 跃升至 99.995%:畅捷通如何用 AI 重塑运维底座?
运维·网络·人工智能