别再死记硬背QKV了！用Python手写一个Self-Attention，带你从几何视角彻底搞懂

从几何视角手写Self-Attention：用Python和NumPy拆解Transformer核心机制

在深度学习领域，Transformer架构已经彻底改变了自然语言处理的游戏规则。而在这个架构中，Self-Attention机制无疑是最闪耀的明星。但令人沮丧的是，大多数教程都停留在抽象的数学公式层面，让学习者陷入矩阵乘法的迷雾中。本文将带你用Python和几何视角，亲手构建一个Self-Attention层，让那些神秘的Q、K、V矩阵变得触手可及。

1. 准备工作：理解向量的几何意义

在开始编写代码之前，我们需要建立一些直观的几何理解。想象你站在一个明亮的房间里，墙上投射着各种形状的影子——这就是向量投影的物理类比。

向量的内积在几何上表示什么？它实际上衡量的是两个向量的"对齐程度"：

import numpy as np # 定义两个二维向量 v1 = np.array([1, 0]) v2 = np.array([0.707, 0.707]) # 45度方向的单位向量 # 计算内积 dot_product = np.dot(v1, v2) print(f"内积值: {dot_product:.3f}") # 输出 ≈0.707

这个0.707的值正是cos(45°)，也就是v2在v1方向上的投影长度。当两个向量完全对齐时，内积最大；垂直时，内积为零。

提示：在自然语言处理中，词向量的内积可以理解为词语之间的语义相关性。相关性高的词语在向量空间中会更"靠近"。

2. 构建最简Self-Attention

让我们从最基础的版本开始，暂时不考虑Q、K、V矩阵。假设我们有一个包含三个词的句子，每个词用二维向量表示：

# 三个词的嵌入向量 X = np.array([ [1.0, 0.5], # 词1 [0.8, 1.2], # 词2 [0.3, 1.5] # 词3 ])

Self-Attention的第一步是计算注意力分数，也就是词与词之间的相关性：

# 计算注意力分数 (原始版本) attention_scores = X @ X.T # 矩阵乘以其转置 print("原始注意力分数:\n", attention_scores)

输出结果会是一个3x3的矩阵，其中每个元素表示两个词向量的内积。但这样直接使用内积有两个问题：

1. 数值大小依赖于向量维度
2. 未经归一化，不适合直接作为权重

因此我们需要引入缩放和softmax：

# 缩放点积注意力 d_k = X.shape[1] # 向量维度(这里是2) scaled_scores = attention_scores / np.sqrt(d_k) # 应用softmax归一化 def softmax(x): exp_x = np.exp(x - np.max(x, axis=-1, keepdims=True)) return exp_x / np.sum(exp_x, axis=-1, keepdims=True) attention_weights = softmax(scaled_scores) print("归一化后的注意力权重:\n", attention_weights)

最后，我们用这些权重对原始向量进行加权求和：

# 加权求和 output = attention_weights @ X print("自注意力输出:\n", output)

这个输出矩阵中的每一行都是原始词向量基于注意力权重的"新表示"——它已经融入了上下文信息。

3. 引入Q、K、V矩阵

现在我们来理解那些神秘的Q(Query)、K(Key)、V(Value)矩阵。它们实际上是原始输入的三个不同线性变换：

# 随机初始化Q、K、V的权重矩阵 np.random.seed(42) W_Q = np.random.randn(2, 2) W_K = np.random.randn(2, 2) W_V = np.random.randn(2, 2) # 计算Q、K、V Q = X @ W_Q K = X @ W_K V = X @ W_V print("Query矩阵:\n", Q) print("Key矩阵:\n", K) print("Value矩阵:\n", V)

为什么要引入这三个矩阵？几何上看，它们让模型能够学习：

• Query：当前词"想知道"什么
• Key：每个词"能提供"什么
• Value：实际传递的信息

完整的缩放点积注意力实现：

def self_attention(X, W_Q, W_K, W_V): Q = X @ W_Q K = X @ W_K V = X @ W_V d_k = Q.shape[-1] scores = Q @ K.T / np.sqrt(d_k) weights = softmax(scores) return weights @ V # 使用自注意力 output = self_attention(X, W_Q, W_K, W_V) print("完整自注意力输出:\n", output)

4. 可视化理解注意力机制

为了真正理解Self-Attention，让我们可视化这个过程。我们将使用matplotlib绘制向量和它们的相互作用。

首先，绘制原始词向量：

import matplotlib.pyplot as plt plt.figure(figsize=(12, 5)) # 绘制原始向量 plt.subplot(1, 2, 1) for i, vec in enumerate(X): plt.arrow(0, 0, vec[0], vec[1], head_width=0.1, length_includes_head=True) plt.text(vec[0], vec[1], f'词{i+1}') plt.xlim(-0.5, 2) plt.ylim(-0.5, 2) plt.title('原始词向量') plt.grid()

然后绘制经过自注意力后的新向量：

# 绘制自注意力后的向量 plt.subplot(1, 2, 2) for i, vec in enumerate(output): plt.arrow(0, 0, vec[0], vec[1], head_width=0.1, length_includes_head=True, color='red') plt.text(vec[0], vec[1], f'新词{i+1}') plt.xlim(-0.5, 2) plt.ylim(-0.5, 2) plt.title('自注意力后的词向量') plt.grid() plt.tight_layout() plt.show()

通过对比左右两图，你会发现新向量已经不再是孤立的表示——它们包含了其他词向量的信息。例如，如果两个原始向量在空间中靠近（语义相关），它们在新表示中会相互"吸引"。

5. 多头注意力机制

真正的Transformer使用多头注意力，这相当于让模型从多个不同角度理解词语关系。实现多头注意力只需要将上述过程并行多次：

class MultiHeadAttention: def __init__(self, d_model=4, num_heads=2): self.d_model = d_model self.num_heads = num_heads assert d_model % num_heads == 0 self.depth = d_model // num_heads # 初始化所有权重矩阵 self.W_Q = np.random.randn(d_model, d_model) self.W_K = np.random.randn(d_model, d_model) self.W_V = np.random.randn(d_model, d_model) self.W_O = np.random.randn(d_model, d_model) def split_heads(self, x): # 将输入分割为多头 batch_size = x.shape[0] return x.reshape(batch_size, -1, self.num_heads, self.depth).transpose(0, 2, 1, 3) def __call__(self, X): Q = X @ self.W_Q K = X @ self.W_K V = X @ self.W_V # 分割多头 Q = self.split_heads(Q) K = self.split_heads(K) V = self.split_heads(V) # 计算缩放点积注意力 d_k = Q.shape[-1] scores = Q @ K.transpose(0, 1, 3, 2) / np.sqrt(d_k) weights = softmax(scores) output = weights @ V # 合并多头 output = output.transpose(0, 2, 1, 3).reshape(X.shape[0], -1, self.d_model) # 最终线性变换 return output @ self.W_O # 使用多头注意力 d_model = 4 X_multi = np.random.randn(3, d_model) # 3个词，每个词4维 mha = MultiHeadAttention(d_model=d_model, num_heads=2) output_multi = mha(X_multi) print("多头注意力输出形状:", output_multi.shape)

多头注意力的优势在于：

1. 允许模型在不同位置共同关注来自不同表示子空间的信息
2. 提供更丰富的上下文表示
3. 增强模型的表达能力

6. 自注意力与卷积、循环网络的对比

为了深入理解Self-Attention的价值，让我们将其与传统方法对比：

自注意力的几何解释特别有趣——它实际上是在学习一种动态的、内容相关的"卷积核"。与固定权重的CNN不同，Self-Attention的"感受野"会根据输入内容自适应调整。

7. 实际应用中的技巧与陷阱

在实现Self-Attention时，有几个关键点需要注意：

1. 初始化策略

Q、K、V矩阵的初始化对训练稳定性至关重要。通常使用较小的随机值：

# Xavier/Glorot初始化 def glorot_init(shape): fan_in, fan_out = shape scale = np.sqrt(2.0 / (fan_in + fan_out)) return np.random.normal(scale=scale, size=shape) W_Q = glorot_init((d_model, d_model))

2. 注意力掩码

在处理变长输入或生成任务时，需要掩码来防止关注不该看的位置：

def get_attention_mask(seq_len): mask = np.triu(np.ones((seq_len, seq_len)), k=1) return mask == 0 # 下三角为True，上三角为False mask = get_attention_mask(3) print("注意力掩码:\n", mask.astype(int))

3. 梯度问题

softmax可能导致梯度消失，特别是在长序列中。解决方案包括：

• 适当的缩放因子
• 注意力头的多样性
• 残差连接

注意：在实际Transformer实现中，通常会加入层归一化和残差连接，这对训练深度模型至关重要。

8. 从零实现完整Self-Attention层

结合以上所有内容，我们实现一个完整的Self-Attention层：

class SelfAttentionLayer: def __init__(self, d_model=64, num_heads=8): self.d_model = d_model self.num_heads = num_heads self.depth = d_model // num_heads # 初始化权重 self.W_Q = glorot_init((d_model, d_model)) self.W_K = glorot_init((d_model, d_model)) self.W_V = glorot_init((d_model, d_model)) self.W_O = glorot_init((d_model, d_model)) def scaled_dot_product_attention(self, Q, K, V, mask=None): d_k = Q.shape[-1] scores = Q @ K.transpose(-1, -2) / np.sqrt(d_k) if mask is not None: scores = np.where(mask, scores, -1e9) weights = softmax(scores) return weights @ V def split_heads(self, x): batch_size = x.shape[0] return x.reshape(batch_size, -1, self.num_heads, self.depth).transpose(0, 2, 1, 3) def combine_heads(self, x): batch_size = x.shape[0] return x.transpose(0, 2, 1, 3).reshape(batch_size, -1, self.d_model) def __call__(self, X, mask=None): batch_size = X.shape[0] Q = X @ self.W_Q K = X @ self.W_K V = X @ self.W_V Q = self.split_heads(Q) K = self.split_heads(K) V = self.split_heads(V) attention = self.scaled_dot_product_attention(Q, K, V, mask) output = self.combine_heads(attention) return output @ self.W_O # 使用示例 d_model = 64 batch_size = 2 seq_len = 10 X_demo = np.random.randn(batch_size, seq_len, d_model) attention_layer = SelfAttentionLayer(d_model=d_model) output_demo = attention_layer(X_demo) print("完整自注意力层输出形状:", output_demo.shape)

这个实现包含了实际Transformer中Self-Attention的所有关键要素：

1. 多头注意力机制
2. 缩放点积注意力
3. 可选的注意力掩码
4. 最终的线性变换

9. 性能优化技巧

当处理真实数据时，我们需要考虑计算效率。以下是几个优化方向：

1. 矩阵乘法优化

使用高效的BLAS库，或者利用分块矩阵计算：

# 分块矩阵乘法示例 def block_matmul(A, B, block_size=32): m, n = A.shape n, p = B.shape C = np.zeros((m, p)) for i in range(0, m, block_size): for j in range(0, p, block_size): for k in range(0, n, block_size): C[i:i+block_size, j:j+block_size] += \ A[i:i+block_size, k:k+block_size] @ B[k:k+block_size, j:j+block_size] return C

2. 稀疏注意力

对于长序列，完全注意力计算代价高昂。可以使用：

• 局部注意力
• 稀疏注意力模式
• 近似注意力

3. 内存优化

注意力权重矩阵可能消耗大量内存。可以通过：

• 梯度检查点
• 混合精度训练
• 内存高效的注意力实现

10. 扩展与应用

Self-Attention的思想远不止于Transformer。以下是一些创新应用方向：

1. 计算机视觉

• Vision Transformer (ViT)
• 自注意力增强的CNN

2. 图神经网络

• 图注意力网络 (GAT)
• 用注意力替代消息传递

3. 跨模态学习

• 图像-文本联合注意力
• 多模态融合

4. 高效变体

• Linformer
• Performer
• Longformer

在实现这些变体时，核心的几何直觉保持不变——通过向量间的相互作用（内积）来建模关系，只是计算方式更加高效或 specialized。