别再死记硬背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的矩阵其中每个元素表示两个词向量的内积。但这样直接使用内积有两个问题数值大小依赖于向量维度未经归一化不适合直接作为权重因此我们需要引入缩放和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, keepdimsTrue)) return exp_x / np.sum(exp_x, axis-1, keepdimsTrue) 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_width0.1, length_includes_headTrue) plt.text(vec[0], vec[1], f词{i1}) 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_width0.1, length_includes_headTrue, colorred) plt.text(vec[0], vec[1], f新词{i1}) 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_model4, num_heads2): 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_modeld_model, num_heads2) output_multi mha(X_multi) print(多头注意力输出形状:, output_multi.shape)多头注意力的优势在于允许模型在不同位置共同关注来自不同表示子空间的信息提供更丰富的上下文表示增强模型的表达能力6. 自注意力与卷积、循环网络的对比为了深入理解Self-Attention的价值让我们将其与传统方法对比特性自注意力RNNCNN长程依赖优秀直接连接困难逐步传递中等需要深层并行计算完全并行序列依赖部分并行计算复杂度O(n²)O(n)O(kn)位置信息需要位置编码自然包含通过卷积核隐含解释性注意力图可解释难以解释难以解释自注意力的几何解释特别有趣——它实际上是在学习一种动态的、内容相关的卷积核。与固定权重的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(scalescale, sizeshape) W_Q glorot_init((d_model, d_model))2. 注意力掩码在处理变长输入或生成任务时需要掩码来防止关注不该看的位置def get_attention_mask(seq_len): mask np.triu(np.ones((seq_len, seq_len)), k1) 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_model64, num_heads8): 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, maskNone): 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, maskNone): 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_modeld_model) output_demo attention_layer(X_demo) print(完整自注意力层输出形状:, output_demo.shape)这个实现包含了实际Transformer中Self-Attention的所有关键要素多头注意力机制缩放点积注意力可选的注意力掩码最终的线性变换9. 性能优化技巧当处理真实数据时我们需要考虑计算效率。以下是几个优化方向1. 矩阵乘法优化使用高效的BLAS库或者利用分块矩阵计算# 分块矩阵乘法示例 def block_matmul(A, B, block_size32): 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:iblock_size, j:jblock_size] \ A[i:iblock_size, k:kblock_size] B[k:kblock_size, j:jblock_size] return C2. 稀疏注意力对于长序列完全注意力计算代价高昂。可以使用局部注意力稀疏注意力模式近似注意力3. 内存优化注意力权重矩阵可能消耗大量内存。可以通过梯度检查点混合精度训练内存高效的注意力实现10. 扩展与应用Self-Attention的思想远不止于Transformer。以下是一些创新应用方向1. 计算机视觉Vision Transformer (ViT)自注意力增强的CNN2. 图神经网络图注意力网络 (GAT)用注意力替代消息传递3. 跨模态学习图像-文本联合注意力多模态融合4. 高效变体LinformerPerformerLongformer在实现这些变体时核心的几何直觉保持不变——通过向量间的相互作用内积来建模关系只是计算方式更加高效或 specialized。