| 《从零开始构建大型语言模型》一书的补充代码，作者：Sebastian Raschka 代码仓库：https://github.com/rasbt/LLMs-from-scratch |  |
多头注意力机制与数据加载
说明：本文档是对多头注意力机制（Multi-head Attention）的详细介绍，包含了数据加载管道和两种不同的实现方式。多头注意力是Transformer架构的核心组件，也是现代大型语言模型的基础。
# NBVAL_IGNORE_OUTPUT
from importlib.metadata import version
print("torch version:", version("torch"))
torch version: 2.2.2
完整的章节代码位于 ch03.ipynb。
本笔记本包含了主要内容：多头注意力机制的实现（以及来自第2章的数据加载管道）
第2章的数据加载器
数据加载器的作用：在训练大型语言模型时，我们需要将文本数据转换为模型可以处理的数字格式。数据加载器负责将原始文本分词、编码，并组织成批次供模型训练使用。
import tiktoken # OpenAI的分词器库
import torch
import torch.nn as nn
from torch.utils.data import Dataset, DataLoader
class GPTDatasetV1(Dataset):
"""
GPT数据集类，用于处理文本数据并创建训练样本
主要功能：
1. 将文本分词并转换为token ID
2. 使用滑动窗口创建重叠的序列
3. 为每个输入序列创建对应的目标序列（用于下一个token预测）
"""
def __init__(self, txt, tokenizer, max_length, stride):
self.input_ids = []
self.target_ids = []
# 对整个文本进行分词
# allowed_special参数允许特殊token如<|endoftext|>
token_ids = tokenizer.encode(txt, allowed_special={'<|endoftext|>'})
# 使用滑动窗口将文本分割成重叠的max_length长度序列
# stride控制窗口移动的步长，stride < max_length会产生重叠
for i in range(0, len(token_ids) - max_length, stride):
input_chunk = token_ids[i:i + max_length]
target_chunk = token_ids[i + 1: i + max_length + 1] # 目标序列向右偏移一位
self.input_ids.append(torch.tensor(input_chunk))
self.target_ids.append(torch.tensor(target_chunk))
def __len__(self):
return len(self.input_ids)
def __getitem__(self, idx):
return self.input_ids[idx], self.target_ids[idx]
def create_dataloader(txt, batch_size=4, max_length=256, stride=128, shuffle=True):
"""
创建数据加载器的便捷函数
参数说明：
- txt: 原始文本数据
- batch_size: 批次大小，影响训练效率和内存使用
- max_length: 每个序列的最大长度（上下文窗口大小）
- stride: 滑动窗口的步长
- shuffle: 是否随机打乱数据顺序
"""
# 初始化GPT-2分词器
tokenizer = tiktoken.get_encoding("gpt2")
# 创建数据集
dataset = GPTDatasetV1(txt, tokenizer, max_length, stride)
# 创建数据加载器
dataloader = DataLoader(dataset, batch_size=batch_size, shuffle=shuffle)
return dataloader
# 读取示例文本文件
with open("small-text-sample.txt", "r", encoding="utf-8") as f:
raw_text = f.read()
tokenizer = tiktoken.get_encoding("gpt2")
encoded_text = tokenizer.encode(raw_text)
# 模型配置参数
vocab_size = 50257 # GPT-2词汇表大小
output_dim = 256 # 嵌入维度
max_len = 1024 # 最大序列长度
context_length = max_len
# 创建嵌入层
# token_embedding_layer: 将token ID转换为向量表示
# pos_embedding_layer: 为每个位置添加位置编码
token_embedding_layer = nn.Embedding(vocab_size, output_dim)
pos_embedding_layer = torch.nn.Embedding(context_length, output_dim)
max_length = 4
dataloader = create_dataloader(raw_text, batch_size=8, max_length=max_length, stride=max_length)
# 处理一个批次的数据
for batch in dataloader:
x, y = batch # x是输入序列，y是目标序列
# 获取token嵌入
token_embeddings = token_embedding_layer(x)
# 获取位置嵌入
pos_embeddings = pos_embedding_layer(torch.arange(max_length))
# 将token嵌入和位置嵌入相加，得到最终的输入嵌入
# 这是Transformer架构的标准做法
input_embeddings = token_embeddings + pos_embeddings
break # 只处理第一个批次作为示例
print(input_embeddings.shape)
torch.Size([8, 4, 256])
输出解释：形状 [8, 4, 256] 表示：
- 8：批次大小（batch_size）
- 4：序列长度（max_length）
- 256：嵌入维度（output_dim）
第3章的多头注意力机制
多头注意力机制简介：多头注意力是Transformer架构的核心创新。它允许模型同时关注输入序列中的不同位置和不同类型的信息。通过使用多个"注意力头"，模型可以学习到更丰富的表示。
变体A：简单实现
设计思路：这种实现方式直观易懂，每个注意力头都是独立的CausalSelfAttention模块，最后将所有头的输出拼接起来。
class CausalSelfAttention(nn.Module):
"""
因果自注意力机制
"因果"意味着模型只能看到当前位置之前的信息，不能看到未来的信息。
这对于语言建模任务至关重要，确保模型在预测下一个token时不会"作弊"。
"""
def __init__(self, d_in, d_out, context_length, dropout, qkv_bias=False):
super().__init__()
self.d_out = d_out
# 创建查询(Query)、键(Key)、值(Value)的线性变换层
# 这是注意力机制的三个核心组件
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) # 防止过拟合
# 注册因果掩码：上三角矩阵，用于屏蔽未来信息
# diagonal=1表示主对角线上方的元素为1
self.register_buffer('mask', torch.triu(torch.ones(context_length, context_length), diagonal=1))
def forward(self, x):
b, n_tokens, d_in = x.shape # 批次大小、序列长度、输入维度
# 计算查询、键、值
keys = self.W_key(x)
queries = self.W_query(x)
values = self.W_value(x)
# 计算注意力分数：Q * K^T
attn_scores = queries @ keys.transpose(1, 2)
# 应用因果掩码，将未来位置的分数设为负无穷
# 这样在softmax后这些位置的权重就会接近0
attn_scores.masked_fill_(
self.mask.bool()[:n_tokens, :n_tokens], -torch.inf)
# 缩放点积注意力：除以sqrt(d_k)防止梯度消失
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
class MultiHeadAttentionWrapper(nn.Module):
"""
多头注意力包装器 - 简单实现版本
这个实现创建多个独立的注意力头，然后将它们的输出拼接起来。
虽然直观，但在实际应用中效率不如变体B。
"""
def __init__(self, d_in, d_out, context_length, dropout, num_heads, qkv_bias=False):
super().__init__()
# 创建多个注意力头
self.heads = nn.ModuleList(
[CausalSelfAttention(d_in, d_out, context_length, dropout, qkv_bias)
for _ in range(num_heads)]
)
# 输出投影层，用于整合多个头的信息
self.out_proj = nn.Linear(d_out*num_heads, d_out*num_heads)
def forward(self, x):
# 将所有头的输出在最后一个维度上拼接
context_vec = torch.cat([head(x) for head in self.heads], dim=-1)
return self.out_proj(context_vec)
# 测试简单实现
torch.manual_seed(123) # 设置随机种子确保结果可重现
context_length = max_length
d_in = output_dim
num_heads = 2 # 使用2个注意力头
d_out = d_in // num_heads # 每个头的输出维度
mha = MultiHeadAttentionWrapper(d_in, d_out, context_length, 0.0, num_heads)
batch = input_embeddings
context_vecs = mha(batch)
print("context_vecs.shape:", context_vecs.shape)
context_vecs.shape: torch.Size([8, 4, 256])
变体B：高效实现
优化思路：这种实现方式更加高效，它将所有注意力头的计算合并到一起，通过张量重塑和转置操作来实现并行计算，这是实际应用中的标准做法。
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 # 每个头的维度
# 注意：这里的线性层输出维度是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) # 最终的输出投影层
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
# 一次性计算所有头的Q、K、V
keys = self.W_key(x) # 形状: (b, num_tokens, d_out)
queries = self.W_query(x)
values = self.W_value(x)
# 关键步骤：重塑张量以分离不同的注意力头
# 从 (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)
# 应用softmax和dropout
attn_weights = torch.softmax(attn_scores / keys.shape[-1]**0.5, dim=-1)
attn_weights = self.dropout(attn_weights)
# 计算上下文向量并重新排列维度
# 形状: (b, num_heads, num_tokens, head_dim) -> (b, num_tokens, num_heads, head_dim)
context_vec = (attn_weights @ values).transpose(1, 2)
# 合并所有头的输出
# 从 (b, num_tokens, num_heads, head_dim) 重塑为 (b, num_tokens, d_out)
# 其中 d_out = num_heads * 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) # 使用相同的随机种子
context_length = max_length
d_in = output_dim
d_out = d_in
mha = MultiHeadAttention(d_in, d_out, context_length, 0.0, num_heads=2)
batch = input_embeddings
context_vecs = mha(batch)
print("context_vecs.shape:", context_vecs.shape)
context_vecs.shape: torch.Size([8, 4, 256])
实现对比总结：
变体A（简单实现）：
- ✅ 概念清晰，易于理解
- ✅ 每个注意力头完全独立
- ❌ 计算效率较低
- ❌ 内存使用较多
变体B（高效实现）：
- ✅ 计算效率高，适合大规模应用
- ✅ 内存使用优化
- ✅ 支持并行计算
- ❌ 实现复杂，需要理解张量操作
在实际应用中，变体B是标准选择，因为它在保持相同功能的同时提供了更好的性能。
多头注意力的核心优势：
- 多样性：不同的头可以关注不同类型的模式和关系
- 并行性：多个头可以并行计算，提高效率
- 表达能力：增强了模型捕获复杂依赖关系的能力
- 稳定性：多个头的平均效果通常比单头更稳定
扫一扫
1127
被折叠的 条评论
为什么被折叠?
到【灌水乐园】发言
