别再死记硬背了！用PyTorch代码逐行拆解Transformer中的QKV矩阵计算
用PyTorch代码逐行拆解Transformer中的QKV矩阵计算
在自然语言处理领域，Transformer架构已经成为事实上的标准。但很多开发者发现，仅通过理论图示理解其核心的注意力机制仍然存在困难。本文将带你用PyTorch代码从零开始实现QKV矩阵的计算过程，通过实际运行和调试来直观感受信息流动。
1. 准备工作与环境搭建
首先确保你的开发环境已经安装了最新版本的PyTorch。如果你使用Colab，可以直接运行以下代码安装：
!pip install torch torchvision接下来导入必要的库：
import torch import torch.nn as nn import torch.nn.functional as F import numpy as np import matplotlib.pyplot as plt为了更清晰地观察矩阵变化，我们定义一个辅助函数来打印张量信息：
def print_tensor_info(name, tensor): print(f"{name}: shape={tensor.shape}, dtype={tensor.dtype}") print(tensor)2. 基础QKV计算实现
让我们从最基本的单头注意力开始，理解QKV矩阵的生成过程。
2.1 定义线性变换层
在Transformer中，QKV矩阵是通过对输入进行线性变换得到的：
class SelfAttention(nn.Module): def __init__(self, embed_size, heads): super(SelfAttention, self).__init__() self.embed_size = embed_size self.heads = heads self.head_dim = embed_size // heads assert ( self.head_dim * heads == embed_size ), "Embedding size needs to be divisible by heads" self.values = nn.Linear(embed_size, embed_size) self.keys = nn.Linear(embed_size, embed_size) self.queries = nn.Linear(embed_size, embed_size) self.fc_out = nn.Linear(embed_size, embed_size)2.2 生成QKV矩阵
现在我们实现前向传播过程，观察QKV矩阵的实际计算：
def forward(self, values, keys, query, mask): N = query.shape[0] # 批大小 value_len, key_len, query_len = values.shape[1], keys.shape[1], query.shape[1] # 线性变换得到QKV values = self.values(values) # (N, value_len, embed_size) keys = self.keys(keys) # (N, key_len, embed_size) queries = self.queries(query) # (N, query_len, embed_size) # 打印变换后的矩阵形状 print_tensor_info("Values after linear", values) print_tensor_info("Keys after linear", keys) print_tensor_info("Queries after linear", queries) # 分割多头 values = values.reshape(N, value_len, self.heads, self.head_dim) keys = keys.reshape(N, key_len, self.heads, self.head_dim) queries = queries.reshape(N, query_len, self.heads, self.head_dim) # 更多调试信息...3. 三种注意力机制的QKV实现差异
Transformer中有三种不同的注意力机制，它们的QKV来源各不相同。让我们分别实现并观察差异。
3.1 编码器自注意力
在编码器自注意力中，QKV都来自同一个输入：
# 模拟编码器输入 batch_size = 2 seq_length = 5 embed_size = 512 dummy_input = torch.randn(batch_size, seq_length, embed_size) # 初始化注意力层 encoder_attention = SelfAttention(embed_size, heads=8) # 自注意力：QKV都来自同一输入 Q = encoder_attention.queries(dummy_input) K = encoder_attention.keys(dummy_input) V = encoder_attention.values(dummy_input) print("Encoder Self-Attention:") print_tensor_info("Q", Q) print_tensor_info("K", K) print_tensor_info("V", V)3.2 解码器自注意力
解码器自注意力需要添加掩码，防止看到未来信息：
# 模拟解码器输入 decoder_input = torch.randn(batch_size, seq_length, embed_size) # 生成掩码 mask = torch.tril(torch.ones(seq_length, seq_length)).expand( batch_size, 1, seq_length, seq_length ) decoder_attention = SelfAttention(embed_size, heads=8) Q = decoder_attention.queries(decoder_input) K = decoder_attention.keys(decoder_input) V = decoder_attention.values(decoder_input) print("\nDecoder Masked Self-Attention:") print_tensor_info("Mask", mask) print_tensor_info("Q", Q) print_tensor_info("K", K) print_tensor_info("V", V)3.3 编码器-解码器注意力
这是跨注意力机制，Q来自解码器，KV来自编码器：
# 模拟编码器输出 encoder_output = torch.randn(batch_size, seq_length, embed_size) cross_attention = SelfAttention(embed_size, heads=8) Q = cross_attention.queries(decoder_input) # Q来自解码器 K = cross_attention.keys(encoder_output) # K来自编码器 V = cross_attention.values(encoder_output) # V来自编码器 print("\nEncoder-Decoder Attention:") print_tensor_info("Q (from decoder)", Q) print_tensor_info("K (from encoder)", K) print_tensor_info("V (from encoder)", V)4. 注意力计算与可视化
理解了QKV的来源后，让我们实现完整的注意力计算过程。
4.1 计算注意力分数
def scaled_dot_product_attention(Q, K, V, mask=None): d_k = Q.size(-1) attention_scores = torch.matmul(Q, K.transpose(-2, -1)) / torch.sqrt(torch.tensor(d_k)) if mask is not None: attention_scores = attention_scores.masked_fill(mask == 0, float("-1e20")) attention_weights = F.softmax(attention_scores, dim=-1) output = torch.matmul(attention_weights, V) return output, attention_weights4.2 可视化注意力权重
让我们可视化三种不同注意力机制的权重分布：
def plot_attention(attention_weights, title): plt.figure(figsize=(10, 5)) plt.imshow(attention_weights[0, 0].detach().numpy(), cmap='viridis') plt.colorbar() plt.title(title) plt.xlabel("Key Positions") plt.ylabel("Query Positions") plt.show() # 编码器自注意力 encoder_output, encoder_weights = scaled_dot_product_attention(Q, K, V) plot_attention(encoder_weights, "Encoder Self-Attention Weights") # 解码器自注意力（带掩码） decoder_output, decoder_weights = scaled_dot_product_attention(Q, K, V, mask) plot_attention(decoder_weights, "Decoder Masked Self-Attention Weights") # 编码器-解码器注意力 cross_output, cross_weights = scaled_dot_product_attention(Q, K, V) plot_attention(cross_weights, "Encoder-Decoder Attention Weights")5. 多头注意力实现
最后，我们实现完整的多头注意力机制，观察QKV在多个头中的不同表现。
5.1 多头注意力前向传播
def forward(self, values, keys, query, mask): N = query.shape[0] value_len, key_len, query_len = values.shape[1], keys.shape[1], query.shape[1] values = self.values(values) # (N, value_len, embed_size) keys = self.keys(keys) # (N, key_len, embed_size) queries = self.queries(query) # (N, query_len, embed_size) # 分割多头 values = values.reshape(N, value_len, self.heads, self.head_dim) keys = keys.reshape(N, key_len, self.heads, self.head_dim) queries = queries.reshape(N, query_len, self.heads, self.head_dim) # 计算注意力 energy = torch.einsum("nqhd,nkhd->nhqk", [queries, keys]) if mask is not None: energy = energy.masked_fill(mask == 0, float("-1e20")) attention = torch.softmax(energy / (self.embed_size ** (1/2)), dim=3) out = torch.einsum("nhql,nlhd->nqhd", [attention, values]).reshape( N, query_len, self.heads * self.head_dim ) out = self.fc_out(out) return out5.2 观察不同头的注意力模式
# 初始化多头注意力 multihead_attn = SelfAttention(embed_size=512, heads=8) # 编码器自注意力 output = multihead_attn(dummy_input, dummy_input, dummy_input, None) # 提取第一个样本的第一个token在各头的注意力权重 sample_weights = attention[0, :, 0, :] # (heads, key_len) # 绘制各头的注意力模式 plt.figure(figsize=(12, 6)) for i in range(8): plt.subplot(2, 4, i+1) plt.plot(sample_weights[i].detach().numpy()) plt.title(f"Head {i+1}") plt.tight_layout() plt.show()通过实际运行这些代码，你可以清晰地看到QKV矩阵在不同注意力机制中的生成过程和数据流动。这种动手实践的方式比单纯看理论图示更能加深对Transformer核心机制的理解。