别再死记硬背了！用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, heads8) # 自注意力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, heads8) 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, heads8) 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, maskNone): 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(), cmapviridis) 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)), dim3) 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_size512, heads8) # 编码器自注意力 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, i1) plt.plot(sample_weights[i].detach().numpy()) plt.title(fHead {i1}) plt.tight_layout() plt.show()通过实际运行这些代码你可以清晰地看到QKV矩阵在不同注意力机制中的生成过程和数据流动。这种动手实践的方式比单纯看理论图示更能加深对Transformer核心机制的理解。