当前位置：首页 > article >正文

Chapter 3 Coding Attention Mechanisms

article 2025/1/6 13:22:55

文章目录

Coding Attention Mechanisms
- The problem with modeling long sequences
- Capturing data dependencies with attentionmechanism
- Attending to different parts of the input with self-attention
- - A simple self-attention mechanism without trainableweights
  - Computing attention weights for all input tokens
- Implementing self-attention with trainable weights
- - Computing the attention weights step by step
  - Implementing a compact self-attention Python class
- Hiding future words with causal attention
- - Applying a causal attention mask
  - Implementing a compact causal attention class
- Extending single-head attention to multi-head attention
- - Stacking multiple single-head attention layers
  - Implementing multi-head attention with weight splits
- Summary

Coding Attention Mechanisms

本章节包含
1. 探讨在神经网络中使用注意力机制的原因
2. 介绍基本的自注意力框架，并逐步发展到增强的自注意力机制
3. 实现一个因果注意力模块，使大型语言模型能够一次生成一个token
4. 使用dropout随机掩盖选定的注意力权重以减少过拟合
5. 将多个因果注意力模块堆叠成多头注意力模块
在本章中，我们将了解LLM架构本身的一个组成部分，即注意力机制

在本章中将实现注意力机制的四种不同变体

上中展示的不同注意力变体是相互构建的，目标是在本章结束时实现一个紧凑且高效的多头注意力机制，然后我们可以将其插入到下一章我们将编写的LLM架构中。

The problem with modeling long sequences

因为源语言和目标语言之间的语法结构存在差异，例如德语翻译成英语，逐字翻译一段文本是不可行的。
为了解决无法逐字翻译文本的问题，通常使用具有两个子模块（所谓的编码器和解码器）的深度神经网络。

在 Transformer 模型出现之前，编码器-解码器 RNN 是机器翻译的流行选择。编码器将源语言中的标记序列作为输入，其中编码器的隐藏状态（中间神经网络层）对整个输入序列的压缩表示进行编码。然后，解码器使用其当前的隐藏状态开始逐个标记的翻译。

但是，编码器-解码器 RNN 的最大问题和限制是，RNN 无法在解码阶段直接从编码器访问早期隐藏状态。因此，它仅依赖于当前隐藏状态，其中封装了所有相关信息。这可能会导致上下文丢失，尤其是在依赖关系可能跨越很长距离的复杂句子中。

Capturing data dependencies with attentionmechanism

RNN 可以很好地翻译短句子，但不能很好地翻译较长的文本，因为它们无法直接访问输入中的先前单词。因此，研究人员在 2014 年为 RNN 开发了所谓的 Bahdanau 注意力机制，该机制修改了编码器-解码器 RNN，使得解码器可以在每次解码时选择性地访问输入序列的不同部分。
受 Bahdanau 注意力机制启发的自注意力机制在三年后(2017)被使用，并提出了原始 transformer 架构。Transformer中的自注意力机制是一种技术，旨在通过使序列中的每个位置能够与同一序列中的其他位置交互并确定其相关性，从而增强输入表示。

Attending to different parts of the input with self-attention

如下图，传统的注意力机制关注的是两个不同序列的元素之间的关系，例如在sequence-to-sequence模型中，注意力可能位于输入序列和输出序列之间

与此相反，在自注意力机制中，“自”指的是该机制能够通过关联单个输入序列中的不同位置来计算注意力权重。它评估并学习输入本身各个部分之间的关系和依赖性，例如句子中的单词或图像中的像素。

A simple self-attention mechanism without trainableweights

在本节中，我们实现了自注意力的简化变体，没有任何可训练的权重。假设我们有一个输入序列 $x^{(1)}$ 到 $x^{(T)}$ ,输入是一个文本（例如，句子“Your journey starts with one step”）,已经转换为 token embedding,例如， $x^{(1)}$ 是一个表示单词“Your”的d维向量，依此类推。

目标：在自注意力中，我们的目标是为输入序列中的每个元素 $x^{(i)}$ 计算上下文向量 $z^{(i)}$ （其中 $z$ 和 $x$ 具有相同的维度）。

其中：
1. 上下文向量 $z^{(i)}$ 是输入序列 $x^{(1)}$ 到 $x^{(T)}$ 的加权和
2. 上下文向量可以解释为丰富的嵌入向量。
3. 以第二个输入 $x^{(2)}$ 为例（注意是第二个输入，所以权重 a 的下标第一个是 2 ）
  $z^{(2)} = a_{21}x^{(1)} + a_{22}x^{(2)} + \cdots + a_{2T}x^{(T)}$
  第二个上下文向量 $z^{(2)}$ 是所有输入 $x^{(1)}$ 到 $x^{(T)}$ 的加权和，权重 $a_{2i}$ , $1,2,\dots,T$ 与第二个输入元素 $x^{(2)}$ 相关,注意力权重是决定在计算 $z^{(2)}$ 时，每个输入元素对加权和的贡献程度的权重。
4. 简而言之，可以将 $z^{(2)}$ 视为 $x^{(2)}$ 的修改版本，它还结合了与当前任务相关的所有其他输入元素的信息。
在自注意力机制中，上下文向量起着至关重要的作用。它们的目的在于通过结合序列中所有其他元素的信息，为输入序列（如句子）中的每个元素创建丰富的表示。这对于大型语言模型（LLMs）至关重要，因为它们需要理解句子中单词之间的关系和相关性。稍后，我们将引入可训练的权重，帮助LLM学习构建这些上下文向量，使其对生成下一个标记具有相关性。

实现自注意力（简化版）的步骤
1. step 1:步骤1计算未归一化的注意力分数 $\omega$ ：
  $\omega_{21} = x^{(1)} q^{(2)\top}\\ \omega_{22} = x^{(2)} q^{(2)\top}\\ \omega_{23} = x^{(3)} q^{(2)\top}\\ ...\\ \omega_{2T} = x^{(T)} q^{(2)\top}\\$
  其中， $\omega$ 是希腊字母“omega”，用于表示未归一化的注意力分数。下标“21”表示使用输入序列元素2作为查询，与输入序列元素1进行比较。
  
  我们以输入序列元素2， $x^{(2)}$ ，为例来计算上下文向量 $z^{(2)}$ ；通过计算查询 x(2)x(2) 与所有其他输入标记的点积来计算未归一化的注意力分数：
```
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)
)
print(f"inputs.shape: {inputs.shape}")
query = inputs[1]  # 第二个 input x^2,journey

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)

"""输出"""
inputs.shape: torch.Size([6, 3])

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)
```
2. step 2步骤 2： 将未归一化的注意力分数（“omega”, $\omega$ ）进行归一化，使其总和为 1以下是一种简单的方法，将未归一化的注意力分数归一化，使其总和为 1（这是一种惯例，有助于解释，并且对训练稳定性很重要）：
```
# 简单的归一化
attn_weights_2_tmp = attn_scores_2 / attn_scores_2.sum()

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

"""输出"""
Attention weights: tensor([0.1455, 0.2278, 0.2249, 0.1285, 0.1077, 0.1656])
Sum: tensor(1.0000)
```
  然而，在实践中，使用 softmax 函数进行归一化更为常见且推荐，因为它在处理极端值方面表现更好，并且在训练过程中具有更理想的梯度特性。以下是一个简单的 softmax 函数实现，用于缩放，同时将向量元素归一化，使其总和为 1：
```
# 使用自定义的 sorftmax
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.)
```
  上述简单的实现可能会因为输入值过大或过小而遇到数值不稳定问题，导致溢出或下溢。因此，在实践中，建议使用 PyTorch 中的 softmax 实现，该实现已针对性能进行了高度优化：
```
# pytorch中的 softmax
attn_weights_2 = torch.softmax(attn_scores_2, dim=0)

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

"""输出"""
Attention weights: tensor([0.1385, 0.2379, 0.2333, 0.1240, 0.1082, 0.1581])
Sum: tensor(1.)
```
3. step3:步骤3通过将嵌入的输入标记 $x^{(i)}$ 与注意力权重相乘，并对结果向量求和，计算上下文向量 $z^{(2)}$
```
# 计算 z^2
query = inputs[1] # 第二个 input x^2,journey

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])
```
至此我们完成了计算输入 2 的注意力权重和上下文向量。

Computing attention weights for all input tokens

接下来，我们将计算所有注意力权重和上下文向量。

我们遵循与之前相同的三个步骤，不同之处在于我们对代码进行了一些修改以计算所有上下文向量，而不是仅计算第二个上下文向量 $z^{(2)}$

step 1: **Compute attention scores,**在步骤 1 中，我们添加一个额外的 for 循环来计算所有输入对的点积。

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

step2:Compute attention weights我们现在对每一行进行标准化，以便每一行中的值总和为 1：

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

setp3:Compute context vectors我们现在使用这些注意力权重通过矩阵乘法计算所有上下文向量：

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

可以根据结果看到之前计算的context_vec_2和第二行的一样

print(context_vec_2)
print(all_context_vecs[1]) # 下标1 为第二行

"""输出"""
tensor([0.4419, 0.6515, 0.5683])
tensor([0.4419, 0.6515, 0.5683])

一个简单的自注意力机制（）的代码演练到此结束。

Implementing self-attention with trainable weights

一个概念框架，展示了本节中开发的自注意力机制如何融入本书和本章的整体叙述和结构：

Computing the attention weights step by step

注意：Weight parameters vs attention weight

需要注意的是，在权重矩阵 $W$ 中，“权重”（weight）是 “权重参数”（weight parameters）的简称，指的是在训练过程中被优化的神经网络参数。这与注意力权重（attention weights）不同，不应混淆。正如我们在上一节中看到的，注意力权重决定了上下文向量依赖输入不同部分的程度，即网络在多大程度上关注输入的不同部分。

总结来说，权重参数是定义网络连接的基本学习系数，而注意力权重是动态的、依赖于上下文的值。
本节介绍的自注意力机制（即“缩放点积注意力”）是原始 Transformer 架构、GPT 模型及大多数大型语言模型的核心，通过引入可训练的权重矩阵，使模型能够学习生成更优的上下文向量。
step1: 逐步实现自注意力机制（self-attention mechanism），我们将首先引入三个训练权重矩阵 $W_q$ 、 $W_k$ 和 $W_v$ 。

这三个矩阵用于通过矩阵乘法将嵌入的输入标记 $x^{(i)}$ 投影为查询向量（query vector）、键向量（key vector）和值向量（value vector）：
- 查询向量： $q^{(i)} = W_q \,x^{(i)}$
- 键向量： $k^{(i)} = W_k \,x^{(i)}$
- 值向量： $v^{(i)} = W_v \,x^{(i)}$
输入 $x$ 和查询向量 $q$ 的嵌入维度可以是相同的，也可以是不同的，具体取决于模型的设计和实现细节。在 GPT 模型中，输入和输出的维度通常是相同的，但为了便于理解计算过程，我们在这里选择不同的输入（3）和输出维度（2）进行说明：
```
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)
)

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
```
初始化q,k,v权重矩阵,计算q、k、v向量
```
# 初始化q,k,v权重矩阵
torch.manual_seed(123)
# 出于说明目的，设置 require_grad=False 来减少输出中的混乱
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)
# 计算 q,k,v向量
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.4306, 1.4551])
```
通过矩阵乘法计算所有的
```
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])
```
我们成功地将 6 个输入标记从 3D 空间投影到了 2D 嵌入空间
setp2:计算注意力分数

注意力得分的计算是一种点积计算，类似于我们在第 3.3 节中使用的简化自注意力机制。这里的新内容是，我们不是直接计算输入元素之间的点积，而是通过各自的权重矩阵对输入进行变换，得到查询（query）和键（key）后再进行计算。

我们通过计算查询向量（query）与每个键向量（key）之间的点积，来计算未归一化的注意力得分：
```
keys_2 = keys[1] # Python starts index at 0
attn_score_22 = query_2.dot(keys_2)
print(attn_score_22)

"""输出"""
tensor(1.8524)
```
因为我们有 6 个 input,所以
```
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)

"""输出"""
tensor([0.1500, 0.2264, 0.2199, 0.1311, 0.0906, 0.1820])
```
setp3:计算上下文向量
```
context_vec_2 = attn_weights_2 @ values
print(context_vec_2)

"""输出"""
tensor([0.3061, 0.8210])
```
到目前为止，我们只计算了一个上下文向量 $z^{2}$ 。

Implementing a compact self-attention Python class

综上所述，我们可以按如下方式实现自注意力机制：

import torch.nn as nn

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

由于输入包含六个嵌入向量，因此会产生一个存储六个上下文向量的矩阵：

我们可以使用PyTorch的线性层（Linear layers）来简化上述实现，如果禁用偏置单元，线性层等价于矩阵乘法;使用nn.Linear而非手动定义nn.Parameter(torch.rand(...))的另一个巨大优势在于，nn.Linear具有优选的权重初始化方案，这有助于实现更稳定的模型训练。

# 使用线性层来简化上述实现（禁用bias）
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)
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>
      )

请注意，SelfAttention_v1 和 SelfAttention_v2 给出不同的输出，因为它们对权重矩阵使用不同的初始权重，因为 nn.Linear 使用更复杂的权重初始化方案。

Hiding future words with causal attention

在本节中，我们修改了标准的self-attention机制，以创建一个causal attention机制，这对于后续章节中开发LLM（大型语言模型）至关重要。

Causal attention，也称为masked attention，是一种特殊的self-attention形式。它限制模型在处理任何给定token时，只能考虑序列中当前及之前的输入，换句话说这确保了下一个单词的预测仅依赖于前面的单词。这与标准的self-attention机制形成对比，后者允许模型一次性访问整个输入序列。

因此，在计算attention scores时，causal attention机制确保模型仅考虑序列中当前token或之前出现的token。
为了在类似GPT的LLM中实现这一点，对于每个处理的token，我们会屏蔽掉输入文本中当前token之后的未来token，如下图所示。

在causal attention中，我们屏蔽掉对角线以上的attention weights，使得在给定输入时，LLM在利用attention weights计算上下文向量时无法访问未来的token。例如，对于第二行中的单词“journey”，我们只保留其之前（“Your”）和当前位置（“journey”）的attention weights。

Applying a causal attention mask

在causal attention中，获得masked attention weight矩阵的一种方法是：对attention scores应用softmax函数，将矩阵对角线上方的元素置为零，并对结果矩阵进行归一化处理。

为了说明和实现causal self-attention，我们将使用上一节中的attention scores和weights进行操作：

上一节中的attention weight

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

屏蔽未来attention weights的最简单方法是使用PyTorch的tril函数创建一个mask，其中主对角线及其下方的元素设为1，主对角线上方的元素设为0：

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

然后，我们可以将attention weights与这个mask相乘，以将对角线上方的attention scores置为零：

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

然而，如果像上面那样在softmax之后应用mask，它会破坏softmax创建的概率分布。Softmax确保所有输出值的总和为1。我们可以通过将每行中的每个元素除以每行中的总和来使其每行的总和再次为1。

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

虽然我们现在在技术上已经完成了causal attention机制的编码，但让我们简要看看一种更高效的方法来实现相同的效果。

因此，与其将对角线上方的attention weights置为零并重新归一化结果，我们可以在这些未归一化的attention scores进入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>)

、

将 softmax 函数应用于这些屏蔽结果

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

后续可以使用修改后的注意力权重来计算上下文向量（这里就略过了）

此外，我们还应用dropout来减少训练期间的过拟合。（需要强调的是，dropout 仅在训练期间使用，之后就被禁用。）
Dropout可以应用于多个位置，例如，
1. 在计算attention weights之后；
2. 或在将attention weights与value vectors相乘之后。
在这里，我们将在计算attention weights之后应用dropout（如下图注意力掩码（左上）， dropout 掩码（右上），因为这是更常见的做法。

此外，在这个特定示例中，我们使用50%的dropout率，这意味着随机屏蔽掉一半的attention weights。（当我们稍后训练GPT模型时，我们将使用更低的dropout率，例如0.1或0.2。）

如果我们应用0.5（50%）的dropout率，未被屏蔽的值将按1/0.5 = 2的比例进行缩放。
缩放比例通过公式1 / (1 - dropout_rate)计算得出。

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

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

"""输出"""
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问题跟踪器上关于这种不一致性的信息。

Implementing a compact causal attention class

复制输入文本，模拟多个batch

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

"""输出"""
torch.Size([2, 6, 3])

现在下面的 CausalAttention 类与我们之前实现的 SelfAttention 类类似，只是我们现在添加了 dropout 和 causal mask 组件

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) # New
        self.register_buffer('mask', torch.triu(torch.ones(context_length, context_length), diagonal=1)) # New

    def forward(self, x):
        b, num_tokens, d_in = x.shape # New batch dimension b
        keys = self.W_key(x)
        queries = self.W_query(x)
        values = self.W_value(x)

        attn_scores = queries @ keys.transpose(1, 2) # Changed transpose
        attn_scores.masked_fill_(  # New, _ ops are in-place
            self.mask.bool()[:num_tokens, :num_tokens], -torch.inf)  # `:num_tokens` to account for cases where the number of tokens in the batch is smaller than the supported context_size
        attn_weights = torch.softmax(
            attn_scores / keys.shape[-1]**0.5, dim=-1
        )
        attn_weights = self.dropout(attn_weights) # New

        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)

context_vecs = ca(batch)

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

"""输出"""
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])

最后一定要注意：dropout仅在训练期间应用，而在推理期间不使用。

Extending single-head attention to multi-head attention

Stacking multiple single-head attention layers

以下是之前实现的自注意力机制（self-attention）的总结（为简化起见，未展示因果掩码和 dropout 掩码）：

下图是称为单头注意力（single-head attention）

我们只需堆叠多个单头注意力模块即可获得多头注意力模块（multi-head attention）

多头注意力的主要思想是通过不同的、可学习的线性投影多次（并行）运行注意力机制。这使得模型能够同时从不同位置的表示子空间中关注信息。

在代码中，我们可以通过实现一个简单的 MultiHeadAttentionWrapper 类来实现这一点，该类将我们之前实现的 CausalAttention 模块的多个实例堆叠起来：

使用 MultiHeadAttentionWrapper 时，我们指定了注意力头的数量（num_heads）。如果我们将 num_heads=2，如上面第二张图所示，我们会得到一个包含两组上下文向量矩阵的张量。在每一组上下文向量矩阵中，行表示与 token 对应的上下文向量，列表示通过 d_out=4 指定的嵌入维度。我们沿着列维度将这些上下文向量矩阵拼接起来。由于我们有 2 个注意力头，且每个头的嵌入维度为 2，因此最终的嵌入维度为 2 × 2 = 4。

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

Implementing multi-head attention with weight splits

虽然上述实现是多头注意力机制的一种直观且功能完整的实现方式（通过封装之前实现的单头注意力CausalAttention），但我们也可以编写一个独立的类 MultiHeadAttention 来实现相同的功能。

在这个独立的 MultiHeadAttention 类中，我们不会将单个注意力头拼接起来。

相反，我们会创建单一的 W_query、W_key 和 W_value 权重矩阵，然后将它们拆分为每个注意力头对应的独立矩阵：

class MultiHeadAttention(nn.Module):
    def __init__(self, d_in, d_out, context_length, dropout, num_heads, qkv_bias=False):
        """  
        Parameters:
        -------------------------------------------------------------------------------------------------------------------
        d_in:                type: int            desc:  输入特征的维度               
        d_out:               type: int            desc:  输出特征的维度       
        context_length:      type: int            desc:  上下文长度，即输入序列的长度  
        dropout:             type: float          desc:  Dropout概率，用于防止过拟合。             
        num_heads:           type: int            desc:  注意力头的数量            
        qkv_bias:            type: boolean        desc:  是否在查询（Query）、键（Key）、值（Value）的线性变换中使用偏置项，默认为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 # Reduce the projection dim to match desired output dim

        # 分别定义query\key\value的线性变换层，将输入维度d_in 映射到输出维度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.out_proj = nn.Linear(d_out, d_out)  # Linear layer to combine head outputs
        # 定义一个Dropout层，用于在训练过程中随机丢弃一部分神经元，以防止过拟合。
        self.dropout = nn.Dropout(dropout)
        # 注册因果掩码
        # self.register_buffer将一个张量注册为模型的缓冲区，这样它不会被当作模型参数进行优化，但会随着模型一起保存和加载。
        self.register_buffer(
            "mask",
            torch.triu(torch.ones(context_length, context_length),diagonal=1) # 生成一个上三角矩阵，对角线以上的元素为1，其余为0。这个矩阵用作因果掩码，确保在计算注意力时，每个位置只能关注到它之前的位置。
        )

    def forward(self, x):
        b, num_tokens, d_in = x.shape # b批次大小，num_tokens序列长度，d_in输入特征维度
        
        # 计算QKV
        keys = self.W_key(x) # Shape: (b, num_tokens, d_out)
        queries = self.W_query(x)
        values = self.W_value(x)
        
        # 重塑张量形状
        # We implicitly split the matrix by adding a `num_heads` dimension
        # Unroll last dim: (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)
        
        # 转置向量
        # Transpose: (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)

        # 计算注意力分数
        # 。keys.transpose(2, 3) 将键的形状从 (b, num_heads, num_tokens, head_dim) 转置为 (b, num_heads, head_dim, num_tokens)，以便进行矩阵乘法
        # Compute scaled dot-product attention (aka self-attention) with a causal mask
        attn_scores = queries @ keys.transpose(2, 3)  # Dot product for each head

        # 应用因果掩码
        # Original mask truncated to the number of tokens and converted to boolean
        mask_bool = self.mask.bool()[:num_tokens, :num_tokens] # 将因果掩码转换为布尔类型，并截取到当前序列长度。

        # Use the mask to fill attention scores
        attn_scores.masked_fill_(mask_bool, -torch.inf) # 使用掩码将注意力分数中不应该关注的位置设置为负无穷，这样在后续的Softmax操作中，这些位置的权重会趋近于0。
        
        # 计算注意力权重
        attn_weights = torch.softmax(attn_scores / keys.shape[-1]**0.5, dim=-1) # 除以 keys.shape[-1]**0.5 是为了缩放注意力分数，防止梯度消失或爆炸
        attn_weights = self.dropout(attn_weights) #  对注意力权重应用Dropout

        # 计算上下文向量
        # Shape: (b, num_tokens, num_heads, head_dim)
        context_vec = (attn_weights @ values).transpose(1, 2) 
        
        # 合并注意力头
        # Combine heads, where self.d_out = self.num_heads * self.head_dim
        context_vec = context_vec.contiguous().view(b, num_tokens, self.d_out) # 将多个注意力头的输出合并成一个张量，形状为 (b, num_tokens, d_out)。
        # 应用输出投影
        context_vec = self.out_proj(context_vec) # optional projection 对合并后的上下文向量进行线性变换，得到最终的输出

        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.\n',context_vecs)
print("context_vecs.shape:\n", context_vecs.shape)

"""输出"""
context_vecs.
 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实际本质上是 MultiHeadAttentionWrapper 的重写版本，但效率更高。另外需要注意的是，我们在上面的 MultiHeadAttention 类中添加了一个线性投影层（self.out_proj）。这只是一个不改变维度的线性变换。在 LLM 实现中使用这样的投影层是一种标准惯例，但它并不是严格必需的（最近的研究表明，移除它不会影响模型性能）。

如果你对上述实现的一个简洁且高效的版本感兴趣，也可以考虑使用 PyTorch 中的 torch.nn.MultiheadAttention 类。

如图 3.26 所示，查询（query）、键（key）和值（value）张量的分割是通过使用 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)。这种转置对于正确对齐不同注意力头的查询、键和值，并高效执行批量矩阵乘法至关重要。

为了说明这种批量矩阵乘法，假设我们有以下示例张量：

# (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 中的矩阵乘法实现会处理 4 维输入张量，使得矩阵乘法在最后两个维度（num_tokens 和 head_dim）之间进行，然后为每个注意力头重复执行。

例如，以下代码成为了一种更简洁的方式，可以分别为每个头计算矩阵乘法：

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

结果与我们之前使用批量矩阵乘法 print(a @ a.transpose(2, 3)) 获得的结果完全相同。

尽管 MultiHeadAttention 类由于额外的张量重塑和转置操作看起来比 MultiHeadAttentionWrapper 更复杂，但它效率更高。原因在于，我们只需要一次矩阵乘法来计算键（例如 keys = self.W_key(x)），对于查询和值也是如此。而在 MultiHeadAttentionWrapper 中，我们需要为每个注意力头重复这一矩阵乘法操作，这是计算中最耗时的步骤之一。

Summary

Attention mechanisms transform input elements into enhanced context vector representations that incorporate information about all inputs.
A self-attention mechanism computes the context vector representation as a weighted sum over the inputs.
In a simplified attention mechanism, the attention weights are computed via dot products.
A dot product is just a concise way of multiplying two vectors element-wise and then summing the products.
Matrix multiplications, while not strictly required, help us to implement computations more efficiently and compactly by replacing nested for-loops.
In self-attention mechanisms that are used in LLMs, also called scaled-dot product attention, we include trainable weight matrices to compute intermediate transformations of the inputs: queries, values, and keys.
When working with LLMs that read and generate text from left to right,we add a causal attention mask to prevent the LLM from accessing future tokens.
Next to causal attention masks to zero out attention weights, we can also add a dropout mask to reduce overfitting in LLMs.
The attention modules in transformer-based LLMs involve multiple instances of causal attention, which is called multi-head attention.
We can create a multi-head attention module by stacking multiple instances of causal attention modules.
A more efficient way of creating multi-head attention modules involves batched matrix multiplications.