#推荐

3.实现注意力机制

本章涵盖以下内容：

• 探讨在神经网络中使用注意力机制的原因
• 介绍一个基本的自注意力框架，并逐步深入到改进的自注意力机制
• 实现一个因果注意力模块，使 LLM 能够一次生成一个token
• 使用 dropout 随机掩盖部分注意力权重，以减少过拟合

• 3.1 长序列建模的问题
• 3.2 通过注意力机制捕捉数据依赖关系
• 3.3 通过自注意力机制关注输入的不同部分
  • 3.3.1 一种不含可训练权重的简化自注意力机制。
  • 3.3.2 为所有输入的 token 计算注意力权重
• 3.4 实现带有可训练权重的自注意力机制
  • 3.4.1 逐步计算注意力权重
  • 3.4.2 实现一个简洁的自注意力机制 Python 类
• 3.5 使用因果注意力机制来屏蔽后续词
  • 3.5.1 应用因果注意力掩码
  • 3.5.2 使用 dropout 遮掩额外的注意力权重
  • 3.5.3 实现一个简洁的因果注意力类
• 3.6 从单头注意力扩展到多头注意力
  • 3.6.1 堆叠多层单头注意力层
  • 3.6.2 通过权重分割实现多头注意力机制
• 3.7 本章摘要

3.1 长序列建模的问题

个人思考： 虽然本书没有涉及对RNN的过多讨论，但了解从RNN到注意力机制的技术变迁对于核心内容的理解至关重要。让我们通过一个具体的示例来理解这种技术变迁：

1. RNN的局限性
   假设我们有一个长句子：“The cat, who was sitting on the windowsill, jumped down because it saw a bird flying outside the window.”
   假设任务是预测句子最后的内容，即要理解“it”指的是“the cat”而不是“the windowsill”或其他内容。对于 RNN 来说，这个任务是有难度的，原因如下：
   • 长距离依赖问题：在 RNN 中，每个新输入的词会被依次传递到下一个时间步。随着句子长度增加，模型的隐状态会不断被更新，但早期信息（如“the cat”）会在层层传播中逐渐消失。因此，模型可能无法在“it”出现时有效地记住“the cat”是“it”的指代对象。
   • 梯度消失问题：RNN 在反向传播中的梯度会随着时间步的增加逐渐减小，这种“梯度消失”使得模型很难在长句中保持信息的准确传播，从而难以捕捉到长距离的语义关联。
2. 注意力机制的解决方法
   为了弥补 RNN 的这些不足，注意力机制被引入。它的关键思想是在处理每个词时，不仅依赖于最后的隐藏状态，而是允许模型直接关注序列中的所有词。这样，即使是较远的词也能在模型计算当前词的语义时直接参与。
   在上例中，注意力机制如何帮助模型理解“it”指代“the cat”呢？
   • 注意力机制的工作原理：当模型处理“it”时，注意力机制会将“it”与整个句子中的其他词进行相似度计算，判断“it”应该关注哪些词。
     • 由于“the cat”与“it”在语义上更相关，注意力机制会为“the cat”分配较高的权重，而其他词（如“windowsill”或“down”）则获得较低的权重。
   • 信息的直接引用：通过注意力机制，模型可以跳过中间步骤，直接将“it”与“the cat”关联，而不需要依赖所有的中间隐藏状态。
3. 示例中的注意力矩阵
   假设使用一个简单的注意力矩阵，模型在处理“it”时，给每个词的权重可能如下（至于如何计算这些权重值后文会详细介绍）：

   词	The	cat	who	was	sitting	…	it	saw	bird	flying	…	window
   权重	0.1	0.3	0.05	0.05	0.05	…	0.4	0.05	0.02	0.01	…	0.02

   在这个注意力矩阵中，可以看到it”对“the cat”有较高的关注权重（0.3），而对其他词的关注权重较低。这种直接的关注能力让模型能够高效捕捉长距离依赖关系，理解“it”与“the cat”的语义关联。

3.2 通过注意力机制捕捉数据依赖关系

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

“自我”在自注意力机制中的含义

在自注意力机制中，“self”指的是该机制通过关联同一输入序列中的不同位置来计算注意力权重的能力。它评估并学习输入内部各部分之间的关系和依赖性，例如句子中的单词或图像中的像素。这与传统注意力机制不同，传统机制关注的是两个不同序列间的关系，例如序列到序列模型中，注意力可能存在于输入序列和输出序列之间，这一点在图 3.5 中有示例说明。

由于自注意力机制对于初次接触的读者可能显得较为复杂，我们将在下一小节中首先介绍一个简化版的自注意力机制。随后，在第 3.4 节中，我们将实现带有可训练权重的自注意力机制，这种机制被用于大语言模型（LLM）中。

3.3.1 一种不含可训练权重的简化自注意力机制。

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]                                               #A
attn_scores_2 = torch.empty(inputs.shape[0])
for i, x_i in enumerate(inputs):
    attn_scores_2[i] = torch.dot(x_i, query)
print(attn_scores_2)


#A 第二个输入 token 用作查询向量

个人思考： 这里对于注意力得分的计算描述的比较笼统，仅仅说明了将当前的输入Token向量与其它Token的向量进行点积运算计算注意力得分，实际上，每个输入Token会先通过权重矩阵W分别计算出它的Q、K、V三个向量，这三个向量的定义如下：

• Q向量（查询向量）：查询向量代表了这个词在寻找相关信息时提出的问题
• K向量（键向量）：键向量代表了一个单词的特征，或者说是这个单词如何”展示”自己，以便其它单词可以与它进行匹配
• V向量（值向量）：值向量携带的是这个单词的具体信息，也就是当一个单词被”注意到”时，它提供给关注者的内容

更通俗的理解： 想象我们在图书馆寻找一本书（Q向量），我们知道要找的主题（Q向量），于是查询目录（K向量），目录告诉我哪本书涉及这个主题，最终我找到这本书并阅读内容（V向量），获取了我需要的信息。

具体生成Q、K、V向量的方式主要通过线性变换：

Q1 = W_Q * (E1 + Pos1)
K1 = W_K * (E1 + Pos1)
V1 = W_V * (E1 + Pos1)

依次类推，为所有token生成Q，K，V向量，其中W_Q，W_K和W_V是Transformer训练出的权重（每一层不同）

针对每一个目标token，Transformer会计算它的 Q向量 与其它所有的token的 K向量 的点积，以确定每个词对当前词的重要性（即注意力分数）

假如有句子：“The cat drank the milk because it was hungry”

例如对于词 cat 的 Q向量 Q_cat，模型会计算：

• score_cat_the = Q_cat · K_the — 与the的语义相关度
• score_cat_drank = Q_cat · K_drank — 与 drank 的语义相关度
• score_cat_it = Q_cat · K_it — 与 it 的语义相关度
• 依此类推，得到cat与句子中其它所有token的注意力分数 [score_cat_the、score_cat_drank、socre_cat_it、……]

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

理解点积

点积运算本质上是一种将两个向量按元素相乘后再求和的简单方式，我们可以如下演示：

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)

除了将点积运算视为结合两个向量并产生标量结果的数学工具之外，点积也是一种相似度的衡量方法，因为它量化了两个向量的对齐程度：较高的点积值表示向量之间有更高的对齐程度或相似度。在自注意力机制的背景下，点积决定了序列中元素之间的关注程度：点积值越高，两个元素之间的相似度和注意力得分就越高。

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)

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

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

个人思考： 这里稍微延伸探讨一下Softmax, 它是一种常用的激活函数，尤其在神经网络的分类任务中被广泛使用。它的作用是将一个任意的实数向量转换为一个概率分布，且所有元素的概率之和为 1。下面通过例子来说明 softmax 的原理、好处，以及它在神经网络中的使用原因。

1. Softmax 的原理
   Softmax 函数的公式如下：
   其中z_i是输入的每个分数（即未激活的原始值），e 是自然对数的底。这个公式的作用是将输入向量中的每个元素转换为一个概率值，且所有值的和为 1。
2. Softmax 的好处
   • 归一化输出为概率：Softmax 将输出转换为 0 到 1 之间的概率，且所有类别的概率之和为 1，方便解释结果。例如，在分类任务中，输出可以直接表示模型对各类别的信心。
   • 平滑和放大效果：Softmax 不仅能归一化，还具有平滑和放大效果。较大的输入值会被放大，较小的输入值会被抑制，从而增强模型对最优类别的区分。
   • 支持多分类问题：与 sigmoid 不同，Softmax 适用于多类别分类问题。它可以输出每个类别的概率，使得模型可以处理多分类任务。
3. 神经网络为什么喜欢使用 Softmax
   在神经网络中，特别是分类模型（如图像分类、文本分类）中，Softmax 层通常用作最后一层输出。原因包括：
   • 便于优化：在分类任务中，Softmax 输出的概率分布可与真实的标签概率进行比较，从而计算交叉熵损失。交叉熵损失的梯度较为稳定，便于模型的优化。
   • 概率解释：Softmax 输出可以解释为“模型对每个类别的信心”，使得输出直观可理解。
   • 与交叉熵的结合：Softmax 与交叉熵损失函数结合效果特别好，可以直接将模型预测的概率分布与真实标签比较，从而更快收敛，效果更好。

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 为所有输入的 token 计算注意力权重

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

attn_scores = inputs @ inputs.T
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]])

attn_weights = torch.softmax(attn_scores, dim=-1)
print(attn_weights)

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

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

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

 print("Previous 2nd context vector:", context_vec_2)

 Previous 2nd context vector: tensor([0.4419, 0.6515, 0.5683])

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

3.4.1 逐步计算注意力权重

x_2 = inputs[1]                                                   #A
d_in = inputs.shape[1]                                            #B
d_out = 2                                                         #C


#A 第二个输入元素
#B 输入维度， d_in=3
#C 输出维度， d_out=2

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)

query_2 = x_2 @ W_query
key_2 = x_2 @ W_key
value_2 = x_2 @ W_value
print(query_2)

权重参数 VS 注意力权重

请注意，在权重矩阵 W 中，术语“权重”是“权重参数”的缩写，指的是神经网络在训练过程中被优化的数值参数。这与注意力权重不同，注意力权重用于确定上下文向量对输入文本的不同部分的依赖程度，即神经网络对输入不同部分的关注程度。

总之，权重参数是神经网络的基本学习系数，用于定义网络层之间的连接关系，而注意力权重则是根据上下文动态生成的特定值，用于衡量不同词语或位置在当前上下文中的重要性。

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

keys_2 = keys[1]                                                  #A
attn_score_22 = query_2.dot(keys_2)
print(attn_score_22)

#A 请牢记在Python中索引从0开始

个人思考： 之前一直有一个疑惑，相同的两个词在不同句子中语义相关度可能完全不同，那么它们的注意力得分是如何做到在不同的上下文中分数不一样的。例如考虑以下两个句子：

• 句子1：”The cat drank the milk because it was hungry.”
• 句子2：”The cat drank the milk because it was sweet.”

很明显，在这两个句子中，it的指代不同，第一个句子中，it指代cat,而在第二个句子中，it指代milk。

根据以下注意力得分的公式（Q_cat和K_it分别为cat和it的查询向量和键向量）可知，句子1中score_cat_it是要大于句子2中的score_cat_it，因为句子1中，it和cat的相关度更高，但是从公式中如何推断出实现呢？

score_cat_it = Q_cat · K_it

我们继续将公式拆解：

Q_cat= W_q * (E_cat + Pos_cat)

K_it = W_k * (E_it + Pos_it)

其中 E_cat和E_it是这两个词的嵌入向量，表示该词的基本语义信息，在不同的上下文中是固定的，根据公式可知，要使最终算出的score_cat_it与上下文语义相关，最重要的是W_q 和**W_k**这两个权重参数应该能反映出不同上下文语义的相关性。在标准的自注意力机制中，W、K、V向量都是固定的，然而，由于 GPT 模型是由多层自注意力模块堆叠而成，每一层都会根据当前输入和上下文信息，动态调整查询、键和值向量的权重矩阵。因此，即使初始的词嵌入和权重矩阵是固定的，经过多层处理后，模型能够生成与当前上下文相关的 Q、K、V 向量权重矩阵，最终计算出的Q、K、V 向量也就能反映出上下文的语义了。GPT多层的实现的细节后文会详述。

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

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

缩放点积注意力机制的原理

对嵌入维度大小进行归一化的原因是为了避免出现小梯度，从而提高训练性能。例如，当嵌入维度增大时（在 GPT 类大型语言模型中通常超过一千），较大的点积在反向传播中应用 softmax 函数后，可能会导致非常小的梯度。随着点积的增大，softmax 函数的行为会更加类似于阶跃函数，导致梯度接近于零。这些小梯度可能会显著减慢学习速度，甚至导致训练停滞。

通过嵌入维度的平方根进行缩放，正是自注意力机制被称为‘缩放点积注意力’的原因。

个人思考： 这里再稍微解释一下上述关于缩放点积注意力的机制。在自注意力机制中，查询向量（Query）与键向量（Key）之间的点积用于计算注意力权重。然而，当嵌入维度（embedding dimension）较大时，点积的结果可能会非常大。那么大的点积对接下来的计算有哪些具体影响呢？

• Softmax函数的特性：在计算注意力权重时，点积结果会通过Softmax函数转换为概率分布。而Softmax函数对输入值的差异非常敏感，当输入值较大时，Softmax的输出会趋近于0或1，表现得类似于阶跃函数（step function）。
• 梯度消失问题：当Softmax的输出接近0或1时，其梯度会非常小，接近于零（可以通过3.3.1小节中提到的Softmax公式推断）。这意味着在反向传播过程中，梯度更新幅度会很小，导致模型学习速度减慢，甚至训练停滞。

为了解决上述问题，在计算点积后，将结果除以嵌入维度的平方根（即 ），其中 d_k 是键向量的维度。这样可以将点积结果缩放到适当的范围，避免Softmax函数进入梯度平缓区，从而保持梯度的有效性，促进模型的正常训练。

context_vec_2 = attn_weights_2 @ values
print(context_vec_2)

为什么使用Q、K和V向量？

在注意力机制的上下文中，“键”（key）、“查询”（query）和“值”（value）这些术语来源于信息检索和数据库领域，在这些领域中也使用类似的概念来存储、搜索和检索信息

查询（query）类似于数据库中的搜索查询。它代表模型当前关注或试图理解的项（如句子中的某个词或 token）。通过查询，模型可以探查输入序列中的其他部分，以确定对它们应关注的程度。

键（key）类似于数据库中用于索引和查找的键。在注意力机制中，输入序列的每个元素（例如句子中的每个单词）都对应一个关联的‘键’。这些‘键’用于与‘查询’进行匹配。

值（value）类似于数据库中的键值对中的“值”。它表示输入项的实际内容或表示。当模型确定哪些键（即输入中的哪些部分）与查询（当前的关注项）最相关时，就会检索出对应的值。

3.4.2 实现一个简洁的自注意力机制 Python 类

# Listing 3.1 A compact self-attention class
import torch.nn as nn
class SelfAttention_v1(nn.Module):
    def __init__(self, d_in, d_out):
        super().__init__()
        self.d_out = d_out
        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 # omega
        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)
sa_v1 = SelfAttention_v1(d_in, d_out)
print(sa_v1(inputs))

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

#  Listing 3.2 A self-attention class using PyTorch's Linear layers
class SelfAttention_v2(nn.Module):
    def __init__(self, d_in, d_out, 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)

    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)
sa_v2 = SelfAttention_v2(d_in, d_out)
print(sa_v2(inputs))

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 存储的是转置形式的权重矩阵。）分配完成后，你应该能观察到两个实例生成相同的输出。

3.5 使用因果注意力机制来屏蔽后续词

3.5.1 应用因果注意力掩码

queries = sa_v2.W_query(inputs)                                   #A
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)

#A 为了方便起见，我们复用上一节中 SelfAttention_v2 对象的query和key权重矩阵。

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

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

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

row_sums = masked_simple.sum(dim=1, keepdim=True)
masked_simple_norm = masked_simple / row_sums
print(masked_simple_norm)

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

信息泄露

当我们应用掩码并重新归一化注意力权重时，乍一看似乎未来的 token（即我们打算掩盖的部分）仍可能影响当前 token，因为它们的值仍然参与了 softmax 计算。然而，关键在于，当我们在掩码之后重新归一化注意力权重时，本质上是在一个更小的子集上重新计算 softmax（因为被掩盖的位置不会贡献到 softmax 的计算值中）。

softmax 算法的优雅之处在于，尽管最初所有位置都包含在分母中，但经过掩码处理和重新归一化后，被掩盖的位置的影响被抵消了——它们在任何实质性意义上都不会影响 softmax 得分。

简而言之，在应用掩码和重新归一化之后，注意力权重的分布就像一开始只在未被掩码的位置上计算的一样。这确保了不会有来自未来（或其他掩码位置）的信息泄露，从而实现了我们的预期。

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

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

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

3.5.2 使用 dropout 遮掩额外的注意力权重

torch.manual_seed(123)
dropout = torch.nn.Dropout(0.5)                                   #A
example = torch.ones(6, 6)                                        #B
print(dropout(example))


#A 我们使用的dropout率为0.5
#B 创建一个由1组成的矩阵

tensor([[2., 2., 0., 2., 2., 0.],
        [0., 0., 0., 2., 0., 2.],
        [2., 2., 2., 2., 0., 2.],
        [0., 2., 2., 0., 0., 2.],
        [0., 2., 0., 2., 0., 2.],
        [0., 2., 2., 2., 2., 0.]])

个人思考： 读到这一段时，我有些不解，Dropout相当于丢弃一定比例的注意力权重，这表明对输入中的某些token关注度降为0了（完全不关注），这样的处理方式难道对最终的预测效果没有影响么？另外如何理解Dropout之后的缩放操作是为了保持注意力在不同阶段的平衡？

经过查阅额外的资料及深度思考，我觉得可以从以下几个方面理解上述的疑问：

1. Dropout 的目的：提高模型的泛化能力
   dropout 的设计初衷是提高模型的泛化能力。通过随机丢弃一部分神经元或注意力权重，dropout 迫使模型在每次训练时学习略有不同的表示方式，而不是依赖某一特定的注意力模式。这种随机化的训练方式可以帮助模型在面对新数据时更具鲁棒性，减少过拟合的风险。
2. 注意力机制的冗余性
   在 Transformer 的注意力机制中，模型通常会对多个 token 进行注意力计算，实际上会有一些冗余信息。也就是说，不同 token 之间的信息通常会有部分重叠，并且模型能够从多个来源获取类似的信息。在这种情况下，dropout 随机丢弃一部分注意力权重并不会完全破坏模型的性能，因为模型可以依赖于其他未被丢弃的注意力路径来获取所需信息。
3. 缩放操作的作用
   在应用 dropout 时，一部分注意力权重被随机置零（假设 dropout 率为 p）。剩余的权重会被放大，其放大倍数为 。放大后的权重记为 z′：
   此时，未被置零的注意力权重 z′\mathbf{z}’z′ 将作为 Softmax 的输入。因此，dropout 后的缩放对 Softmax 有两个主要影响：
   • 增大未遮盖值的相对差异：放大剩余权重后，它们的数值相对于被置零的权重增大，从而拉大了非零元素之间的相对差异。这使得在 Softmax 计算中（通过前文提过的Softmax公式推导，输入值的差异越大，输出分布就会越尖锐；而输入值差异越小，输出分布就会越平滑），剩下的值之间的对比更明显。
   • 影响 Softmax 输出的分布形态：当未被置零的权重值被放大后，它们在 Softmax 输出中会更具代表性，注意力分布会更集中（即更尖锐），让模型更关注特定的 token。
4. 缩放后的 Softmax 输入导致注意力分布更倾向于少数的高权重 token，使得模型在当前步骤更关注这些 token 的信息。这对模型的影响包括：
   • 增强模型的选择性关注：在训练中，模型会在每个步骤中随机选择不同的 token 进行更高的关注，这使模型在学习时不会依赖特定 token 的注意力。
   • 确保总注意力强度保持一致：即便经过 dropout 丢弃了一部分权重，缩放保证了剩余权重在 Softmax 后的分布与未应用 dropout 时类似。
5. 训练过程中多次迭代弥补信息丢失
   在训练过程中，每个 batch 中的 dropout 掩码都是随机生成的。也就是说，在每次训练时被丢弃的注意力权重是随机的，并不会始终忽略相同的 token。这种随机性确保了在训练过程中，模型会在多个迭代中多次关注到每个 token。因此，即便某个 token 在当前的训练步中被忽略，在未来的训练步骤中它仍然会被关注到，从而在整体上避免了信息丢失的问题。

torch.manual_seed(123)
print(dropout(attn_weights))

tensor([[2.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000],
        [0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000],
        [0.7599, 0.6194, 0.6206, 0.0000, 0.0000, 0.0000],
        [0.0000, 0.4921, 0.4925, 0.0000, 0.0000, 0.0000],
        [0.0000, 0.3966, 0.0000, 0.3775, 0.0000, 0.0000],
        [0.0000, 0.3327, 0.3331, 0.3084, 0.3331, 0.0000]],
       grad_fn=<MulBackward0>

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

batch = torch.stack((inputs, inputs), dim=0)
print(batch.shape)                                              #A

#A 2个输入，每个输入有6个token，每个token的嵌入维度为3

# Listing 3.3 A compact causal attention class
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)                        #A
        self.register_buffer(
           'mask',
           torch.triu(torch.ones(context_length, context_length),
           diagonal=1)
        )                                                         #B

    def forward(self, x):
        b, num_tokens, d_in = x.shape                             #C
        keys = self.W_key(x)
        queries = self.W_query(x)
        values = self.W_value(x)
        attn_scores = queries @ keys.transpose(1, 2)              #C
        attn_scores.masked_fill_(                                 #D
            self.mask.bool()[:num_tokens, :num_tokens], -torch.inf)
        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


#A 与之前的 SelfAttention_v1 类相比，我们添加了一个 dropout 层
#B register_buffer 调用也是新添加的内容（后续内容会提供更多相关信息）
#C 我们交换第 1 和第 2 个维度，同时保持批次维度在第1个位置（索引0）
#D 在 PyTorch 中，带有下划线后缀的操作会在原有内存空间执行，直接修改变量本身，从而避免不必要的内存拷贝

torch.manual_seed(123)
context_length = batch.shape[1]
ca = CausalAttention(d_in, d_out, context_length, 0.0)
context_vecs = ca(batch)
print("context_vecs.shape:", context_vecs.shape)

context_vecs.shape: torch.Size([2, 6, 2])

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

3.6.1 堆叠多层单头注意力层

# Listing 3.4 A wrapper class to implement multi-head attention
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)

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

练习 3.2：返回二维嵌入向量

更改 MultiHeadAttentionWrapper(..., num_heads=2) 调用中的输入参数，使输出的上下文向量为 2 维而不是 4 维，同时保持 num_heads=2 的设置。提示：无需修改类的实现，只需更改其中一个输入参数即可。

3.6.2 通过权重分割实现多头注意力机制

# Listing 3.5 An efficient multi-head attention class
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 must be divisible by num_heads"
        self.d_out = d_out
        self.num_heads = num_heads
        self.head_dim = d_out // num_heads                        #A
        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)                   #B
        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
        keys = self.W_key(x)                                      #C
        queries = self.W_query(x)                                 #C
        values = self.W_value(x)                                  #C

        keys = keys.view(b, num_tokens, self.num_heads, self.head_dim)       #D
        values = values.view(b, num_tokens, self.num_heads, self.head_dim)   #D
        queries = queries.view(b, num_tokens, self.num_heads, self.head_dim) #D

        keys = keys.transpose(1, 2)                               #E
        queries = queries.transpose(1, 2)                         #E
        values = values.transpose(1, 2)                           #E

        attn_scores = queries @ keys.transpose(2, 3)              #F
        mask_bool = self.mask.bool()[:num_tokens, :num_tokens]    #G

        attn_scores.masked_fill_(mask_bool, -torch.inf)           #H

        attn_weights = torch.softmax(
            attn_scores / keys.shape[-1]**0.5, dim=-1)
        attn_weights = self.dropout(attn_weights)

        context_vec = (attn_weights @ values).transpose(1, 2)     #I

        context_vec = context_vec.contiguous().view(b, num_tokens, self.d_out)  #J
        context_vec = self.out_proj(context_vec)                  #K
        return context_vec


#A 将投影维度缩小，以匹配期望的输出维度
#B 使用线性层组合头部输出
#C 张量形状：(b, num_tokens, d_out)
#D 我们通过添加 num_heads 维度来隐式地拆分矩阵。然后展开最后一个维度，使其形状从 (b, num_tokens, d_out) 转换为 (b, num_tokens, num_heads, head_dim)
#E 将张量的形状从 (b, num_tokens, num_heads, head_dim) 转置为 (b, num_heads, num_tokens, head_dim)
#F 对每个注意力头进行点积运算
#G 掩码被截断到 token 的数量
#H 使用掩码填充注意力分数
#I 张量形状：（b, num_tokens, n_heads, head_dim）
#J 将多个注意力头的输出结果合并，其中输出维度 self.d_out 等于注意力头数 self.num_heads 与每个头的维度 self.head_dim 的乘积
#K 添加一个可选的线性投影层

a = torch.tensor([[[[0.2745, 0.6584, 0.2775, 0.8573],             #A
                    [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]]]])

#A 该张量的形状为 (b, num_heads, num_tokens, head_dim) = (1, 2, 3, 4)

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

first_head = a[0, 0, :, :]
first_res = first_head @ first_head.T
print("First head:\n", first_res)
second_head = a[0, 1, :, :]
second_res = second_head @ second_head.T
print("\nSecond head:\n", second_res)

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

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)

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

练习 3.3：初始化 GPT-2 规模的注意力模块

使用 MultiHeadAttention 类初始化一个多头注意力模块，该模块的注意力头数量与最小的 GPT-2 模型相同（12 个注意力头）。同时确保输入和输出的嵌入大小与 GPT-2 相似（768 维）。请注意，最小的 GPT-2 模型支持的上下文长度为 1024 个 tokens。

3.7 本章摘要

• 注意力机制将输入元素转换为增强的上下文向量表示，其中包含了所有输入的信息。
• 自注意力机制通过对输入的加权求和来计算上下文向量表示。
• 在简化的注意力机制中，注意力权重是通过点积计算的。
• 点积仅仅是对两个向量逐元素相乘后求和的一种简洁方式。
• 矩阵乘法虽然并不是绝对必要的，但通过替换嵌套的 for 循环，它帮助我们更高效和简洁地计算。
• 在 LLM 中使用的自注意力机制，也称为缩放点积注意力，我们引入可训练的权重矩阵，以计算输入的中间转换：查询、值和键。
• 在使用从左到右读取和生成文本的 LLM 时，我们添加一个因果注意力掩码，以防止模型访问未来的 token。
• 除了通过因果注意力掩码将注意力权重置为零之外，我们还可以添加 dropout 掩码，以减少 LLM 中的过拟合现象。
• 基于Transformer的 LLM 中的注意力模块包含多个因果注意力实例，这被称为多头注意力。
• 我们可以通过堆叠多个因果注意力模块的实例来创建一个多头注意力模块。
• 创建多头注意力模块的一种更高效的方法是采用批量矩阵乘法。

个人思考： 毫无疑问，本章的注意力机制是整本书中最重要的内容，也是最难的内容 ，这里强烈建议读者能多读几遍，按照文中的示例代码完整地实现一遍，对于不理解的地方多去查阅相关资料及深入思考，力求真正理解和掌握每个细节点。