别再死记硬背了！用Python+PyTorch手把手图解自注意力机制（附完整代码）

别再死记硬背了！用Python+PyTorch手把手图解自注意力机制（附完整代码）

理解自注意力机制最有效的方式不是背诵公式，而是亲手实现它。本文将带你用PyTorch从零构建一个可交互的自注意力模块，并通过动态可视化揭示其核心计算逻辑。无论你是准备面试的开发者，还是正在学习Transformer架构的研究者，这套代码实验都能让你真正掌握"注意力"的本质。

1. 环境准备与数据建模

我们先构建一个极简的文本处理场景：输入4个单词的嵌入向量，模拟Transformer中的单头自注意力计算。这里使用PyTorch的自动微分功能，避免手动计算矩阵导数。

import torch import torch.nn as nn import matplotlib.pyplot as plt from matplotlib.animation import FuncAnimation # 模拟输入：4个单词的嵌入向量（维度=64） tokens = ["deep", "learning", "is", "fun"] embed_dim = 64 x = torch.randn(4, embed_dim) # 形状：[序列长度, 嵌入维度]

定义可训练的权重矩阵（实际项目中这些参数会自动学习）：

class SelfAttention(nn.Module): def __init__(self, embed_dim): super().__init__() self.W_q = nn.Linear(embed_dim, embed_dim, bias=False) self.W_k = nn.Linear(embed_dim, embed_dim, bias=False) self.W_v = nn.Linear(embed_dim, embed_dim, bias=False) def forward(self, x): Q = self.W_q(x) # 查询向量 K = self.W_k(x) # 键向量 V = self.W_v(x) # 值向量 return Q, K, V

2. 动态计算注意力分数

自注意力的核心是计算单词间的关联程度。我们通过查询-键点积得到原始分数，然后用softmax归一化：

def compute_attention(Q, K): scores = torch.matmul(Q, K.transpose(0, 1)) # 点积运算 scores = scores / (embed_dim ** 0.5) # 缩放防止梯度消失 attn_weights = torch.softmax(scores, dim=-1) return attn_weights # 实例化并计算 attn_layer = SelfAttention(embed_dim) Q, K, V = attn_layer(x) attn_weights = compute_attention(Q, K)

用热力图实时显示注意力矩阵的变化：

fig, ax = plt.subplots() im = ax.imshow(attn_weights.detach().numpy(), cmap='viridis') def update(i): # 模拟训练过程中权重更新 with torch.no_grad(): attn_layer.W_q.weight += 0.01 * torch.randn_like(attn_layer.W_q.weight) Q, K, V = attn_layer(x) im.set_data(compute_attention(Q, K).detach().numpy()) return [im] ani = FuncAnimation(fig, update, frames=20, interval=500) plt.colorbar(im) plt.show()

这段代码会生成一个动态图，展示随着权重矩阵更新，各单词间注意力分布的变化过程。你会直观看到某些单词组合（如"deep"和"learning"）逐渐形成强关联。

3. 权重聚合与输出生成

获得注意力权重后，我们需要用它加权求和值向量：

def weighted_sum(attn_weights, V): return torch.matmul(attn_weights, V) # 形状：[序列长度, 嵌入维度] output = weighted_sum(attn_weights, V)

为了验证效果，可以对比输入输出向量的相似度：

cos = nn.CosineSimilarity(dim=1) print("输入输出相似度:", cos(x, output))

典型输出可能显示：

输入输出相似度: tensor([0.3124, 0.2897, 0.2568, 0.3012])

4. 扩展为多头注意力

单头注意力只能捕捉一种模式的关系。实际Transformer使用多头机制：

class MultiHeadAttention(nn.Module): def __init__(self, embed_dim, num_heads=8): super().__init__() self.head_dim = embed_dim // num_heads self.W_o = nn.Linear(embed_dim, embed_dim) # 输出投影 def split_heads(self, x): return x.view(x.size(0), -1, self.head_dim) def forward(self, x): Q, K, V = attn_layer(x) Q = self.split_heads(Q) # [序列长度, 头数, 头维度] K = self.split_heads(K) V = self.split_heads(V) # 各头独立计算 attn_outputs = [] for i in range(Q.size(1)): attn = compute_attention(Q[:,i], K[:,i]) attn_outputs.append(weighted_sum(attn, V[:,i])) # 拼接并投影 combined = torch.cat(attn_outputs, dim=1) return self.W_o(combined)

关键改进点：

1. 查询/键/值被分割到不同子空间
2. 每个头独立计算注意力
3. 最终结果通过线性层融合

5. 可视化技巧进阶

使用NetworkX库绘制动态注意力图：

import networkx as nx def draw_attention_graph(weights, tokens): G = nx.DiGraph() G.add_nodes_from(tokens) for i, src in enumerate(tokens): for j, dst in enumerate(tokens): G.add_edge(src, dst, weight=weights[i,j].item()) pos = nx.circular_layout(G) nx.draw(G, pos, with_labels=True, edge_color=[G[u][v]['weight'] for u,v in G.edges()], width=[2*G[u][v]['weight'] for u,v in G.edges()])

调用示例：

draw_attention_graph(attn_weights, tokens)

这会生成带权重的有向图，边的粗细和颜色深度反映注意力强度。通过对比不同层的注意力图，可以直观理解Transformer如何构建层级表征。