用Python从零实现Self-Attention：5行核心代码透视Transformer灵魂
当你第一次看到Transformer架构中那个神秘的"Scaled Dot-Product Attention"公式时，是否曾被各种矩阵运算绕得头晕目眩？让我们换种方式——用NumPy亲手构建一个自注意力机制，你会惊讶地发现，那些看似复杂的Q、K、V矩阵交互，本质上只是一系列精心设计的向量游戏。
1. 环境准备与数据建模
在开始编写注意力机制之前，我们需要先搭建一个极简的实验环境。打开你的Python环境（推荐Jupyter Notebook），导入以下基础工具包：
import numpy as np from scipy.special import softmax为了保持实验的直观性，我们构造一个包含3个单词的微型句子："Thinking Machines"。这里的每个单词将被表示为512维的嵌入向量——这是Transformer架构的典型配置：
# 构造3个单词的词嵌入矩阵 (3x512) word_embeddings = np.random.randn(3, 512) # 随机初始化代替真实嵌入 # 单词索引映射 word_index = {"Thinking": 0, "Machines": 1, "<END>": 2}提示：实际应用中，词嵌入通常通过训练获得，这里使用随机初始化仅用于演示目的
2. 权重矩阵：QKV的诞生之谜
自注意力机制的核心在于学会动态生成三种不同的向量表示：查询(Query)、键(Key)和值(Value)。这些向量都源自同一输入，但通过不同的权重矩阵进行投影：
# 定义可训练的权重矩阵 (512x64) W_Q = np.random.randn(512, 64) # 查询权重 W_K = np.random.randn(512, 64) # 键权重 W_V = np.random.randn(512, 64) # 值权重 # 生成Q、K、V矩阵 Q = word_embeddings @ W_Q # (3x512) @ (512x64) → 3x64 K = word_embeddings @ W_K # 同上 V = word_embeddings @ W_V # 同上为什么选择64维？这是原始Transformer论文中的设计选择——将512维的嵌入压缩到64维的注意力空间，既保留了足够的信息量，又控制了计算复杂度。
3. 注意力分数的计算艺术
现在进入最精妙的部分：通过查询和键的交互计算注意力分数。这个阶段决定了每个单词应该"关注"句子中的哪些部分：
# 计算缩放点积注意力分数 attention_scores = Q @ K.T / np.sqrt(64) # (3x64) @ (64x3) → 3x3 # 应用Softmax归一化 attention_weights = softmax(attention_scores, axis=1)让我们分解这个3x3的注意力矩阵：
- 行代表查询(query)单词
- 列代表键(key)单词
- 每个元素表示两个单词之间的关联强度
例如，attention_weights[0,1]表示"Thinking"对"Machines"的关注程度。通过打印这个矩阵，你会看到类似这样的模式：
[[0.9, 0.1, 0.0], [0.3, 0.6, 0.1], [0.1, 0.2, 0.7]]这显示第一个单词("Thinking")主要关注自身(0.9)，而对其他单词关注较少——这是自注意力在学习单词间关系时的典型行为。
4. 注意力输出：上下文感知的表示
最后一步是将注意力权重应用于值矩阵，生成每个单词的上下文相关表示：
# 计算加权值向量和 Z = attention_weights @ V # (3x3) @ (3x64) → 3x64这个输出矩阵Z的每一行都是原始单词嵌入的"精炼版"，其中包含了句子级别的上下文信息。例如：
- "Thinking"的新表示融合了它与其他单词的关系
- ""标记的表示现在包含了句子结构的线索
5. 可视化与扩展实践
为了更直观地理解，我们可以用热力图展示注意力权重：
import matplotlib.pyplot as plt plt.imshow(attention_weights, cmap='viridis') plt.colorbar() plt.xticks(range(3), ["Thinking", "Machines", "<END>"]) plt.yticks(range(3), ["Thinking", "Machines", "<END>"]) plt.title("Attention Weights Heatmap") plt.show()当你运行这段代码时，会看到一个色彩鲜明的矩阵，其中明亮的对角线通常表示单词对自身的强烈关注——这是自注意力机制的典型特征。
如果想进一步探索，可以尝试以下扩展实验：
- 修改输入句子的长度，观察注意力模式的变化
- 调整嵌入维度(如改为256维)，比较计算效率
- 添加位置编码，解决纯自注意力缺乏位置信息的问题
6. 自注意力与RNN/CNN的思维差异
与传统序列模型相比，自注意力机制展现出独特的优势：
| 特性 | RNN | CNN | Self-Attention |
|---|---|---|---|
| 并行计算 | 不支持 | 支持 | 支持 |
| 长距离依赖 | 困难 | 需要多层 | 直接建模 |
| 计算复杂度 | O(n) | O(kn) | O(n²) |
| 可解释性 | 低 | 中等 | 高(通过注意力权重) |
这种全局视野使Transformer能够同时处理序列中的所有关系，而不像RNN那样受制于顺序处理。在实际项目中，这种特性对于理解长文档中的指代关系特别有效。
7. 工程实践中的常见陷阱
虽然我们的示例代码简洁明了，但在真实场景中实现自注意力时，有几个关键点需要注意：
数值稳定性：当嵌入维度较大时，点积结果可能爆炸性增长，导致Softmax饱和
- 解决方案：缩放因子(1/√dₖ)必不可少
内存消耗：序列长度的平方复杂度限制了长文本处理
- 优化策略：采用分块计算或稀疏注意力模式
训练动态：随机初始化的QKV矩阵可能导致训练初期不稳定
- 技巧：使用Xavier初始化或残差连接缓解
# 改进的初始化示例 W_Q = np.random.randn(512, 64) / np.sqrt(512) # Xavier初始化8. 从单头到多头注意力的自然演进
我们实现的单头注意力已经能捕捉单词间的基本关系，但真正的Transformer使用多头注意力机制——本质上并行运行多个独立的注意力"子空间"，每个头学习不同的关系模式：
# 简化的多头注意力实现 num_heads = 8 head_dim = 64 // num_heads multi_head_Z = [] for _ in range(num_heads): # 每个头有独立的QKV权重 W_Q = np.random.randn(512, head_dim) / np.sqrt(512) W_K = np.random.randn(512, head_dim) / np.sqrt(512) W_V = np.random.randn(512, head_dim) / np.sqrt(512) # 计算单头注意力 Q = word_embeddings @ W_Q K = word_embeddings @ W_K V = word_embeddings @ W_V # 缩放点积注意力 attention_scores = Q @ K.T / np.sqrt(head_dim) attention_weights = softmax(attention_scores, axis=1) Z = attention_weights @ V multi_head_Z.append(Z) # 拼接所有头的输出 multi_head_Z = np.concatenate(multi_head_Z, axis=1) # (3x64)这种设计让模型能够同时关注不同位置的不同关系模式——就像人类阅读时会同时注意语法结构、语义关联和指代关系等多个方面。
