介绍
把先前实现的因果注意力类扩展到多个头上。这也被称为 多头注意力
多头
是将注意力机制分成 多个“头”，每个“头”独立工作。在这种情况下，单个因果注意力模块可以被看作单头注意力，因为它只有一组注意力权重按顺序处理输入。
叠加多个单头注意力层
实现多头注意力需要构建多个自注意力机制的实例如下图，每个实例都有其独立的权重，然后将这些输出进行合成
下图展示了多头注意力模块的结构，它是由上图所示的多个单头注意力模块依次叠加在一起组成的，即：两个堆叠在一起的单头注意力模块。
因此，不能使用单一的矩阵 W_v 来计算值矩阵，而是在一个 有两个头的多头注意模块中，有两个值权重矩阵： W_{v1} 和 W_{v2} 。
同样适用于其他的权重矩阵，比如 W_q 和 W_k 。得到了两组上下文向量 Z_1 和 Z_2 ，然后 合并成一个单一的上下文向量矩阵 Z
多头注意力的主要思想是 多次（并行）运行注意力机制，每次使用学到的不同的线性投影——这些投影是通过将输入数据（比如注意力机制中的查询向量、键向量和值向量）乘以权重矩阵得到的。
实现
准备代码
import torch
import torch.nn as nn
class CausalAttention(nn.Module):
def __init__(self, d_in, d_out, context_length, dropout, qkv_bias=False):
"""
初始化
:param d_in: 输入嵌入维度
:param d_out: 输出嵌入维度
:param context_length: 设置掩码矩阵
:param dropout: dropout忽略概率
:param qkv_bias:
"""
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)
# 增加 Dropout 层
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)
queries = self.W_query(x)
values = self.W_value(x)
# 将维度1和2转置，将批维度保持在第一个位置（0）
attn_scores = queries @ keys.transpose(1, 2)
# 直接修改传入的矩阵，而不是返回副本，减少内存
attn_scores.masked_fill_(
self.mask.bool()[:num_tokens, :num_tokens], -torch.inf)
attn_weights = torch.softmax(
attn_scores / keys.shape[-1]**0.5, dim=-1
)
# 对注意力权重矩阵进行dropout操作
attn_weights = self.dropout(attn_weights)
# 计算上下文向量
context_vec = attn_weights @ values
return context_vec
多头注意力封装类
在代码中，可以通过实现一个简单的 MultiHeadAttentionWrapper 类来达到这一目标，MultiHeadAttentionWrapper 类堆叠了多个之前实现的CausalAttention 模块实例
class MultiHeadAttentionWrapper(nn.Module):
def __init__(self, d_in, d_out, context_length,
dropout, num_heads, qkv_bias=False):
"""
初始化
:param d_in: 输入嵌入维度
:param d_out: 输出嵌入维度
:param context_length: 设置掩码矩阵
:param dropout: dropout忽略概率
:param num_heads: 多头注意力数量
:param qkv_bias:
"""
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)
使用
如下图，使用 MultiHeadAttentionWrapper，指定了 注意力头的数量 num_heads=2 ，就会得到一个具有 两组上下文向量矩阵的张量。
在每个上下文向量矩阵中：
- 行表示对应于词元的上下文向量
- 列则对应于通过
d_out=4指定的嵌入维度，沿着 列维度 连接 这些上下文向量矩阵
由于有 两个注意力头，并且 嵌入维度为2，因此最终的嵌入维度是 2×2=4
准备代码
# 输入嵌入维度d_in=3
d_in = inputs.shape[1]
# 输出嵌入维度d_out=2
d_out = 2
# 模拟批量输入，因为真正执行时，是通过数据加载器，批量读取的
batch = torch.stack((inputs, inputs), dim=0)
print("批量数据：\n", batch)
print("批量数据的shape：", batch.shape)
使用
# 计算上下文向量
print("\n\b计算上下文向量--------")
# 为了复现，设置随机种子
torch.manual_seed(123)
# 词元的数量，用于设置掩码矩阵
context_length = batch.shape[1]
print("context_length:", context_length)
# 传入参数：实例化
mha = MultiHeadAttentionWrapper(
d_in, d_out, context_length, 0.0, num_heads=2
)
# 计算上下文向量
context_vecs = mha(batch)
print("上下文向量：\n", context_vecs)
print("上下文向量.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>)
上下文向量.shape: torch.Size([2, 6, 4])
结果中的 context_vecs 张量维度 [2, 6, 4]的解释：
第一维：因为 有两个输入文本（输入文本是重复的，所以这些上下文向量完全相同）
第二维：表示每个输入中的
6个词元第三维：表示每个词元的四维嵌入