别再死记硬背Self-Attention公式了!用NumPy手搓一个,5分钟搞懂QKV和缩放点积

别再死记硬背Self-Attention公式了!用NumPy手搓一个,5分钟搞懂QKV和缩放点积 用NumPy手搓Self-Attention从矩阵操作到直觉理解在自然语言处理和计算机视觉领域Transformer架构已经成为革命性的模型而Self-Attention机制正是其核心所在。许多教程从数学公式出发堆叠各种矩阵运算符号让初学者望而生畏。本文将采用完全不同的路径——我们直接用NumPy实现一个完整的Self-Attention层通过观察每一步的矩阵变化来建立直观理解。1. 准备工作理解输入与基础概念假设我们有一个包含3个单词的句子每个单词用4维向量表示实际应用中维度通常为512或768。首先创建这个输入矩阵import numpy as np # 输入矩阵3个token每个token 4维嵌入 X np.array([ [1.0, 0.5, -1.2, 0.8], [0.6, -0.3, 1.5, -0.7], [-0.9, 1.1, 0.4, -1.0] ]) print(输入矩阵X:\n, X)Self-Attention的核心是Q(Query)、K(Key)、V(Value)三个矩阵。它们都是通过线性变换从输入X得到的# 随机初始化权重矩阵 W_Q np.random.randn(4, 3) # 将4维输入映射到3维Q空间 W_K np.random.randn(4, 3) # 将4维输入映射到3维K空间 W_V np.random.randn(4, 3) # 将4维输入映射到3维V空间 # 计算Q,K,V Q X W_Q K X W_K V X W_V print(\nQuery矩阵Q:\n, Q) print(\nKey矩阵K:\n, K) print(\nValue矩阵V:\n, V)注意实际应用中这些权重矩阵是通过训练学习得到的这里为了演示使用随机初始化2. 计算注意力分数点积与缩放注意力分数的计算分为两步点积和缩放。我们先看原始点积结果# 计算原始注意力分数 raw_attention Q K.T # Q和K的点积 print(\n原始注意力分数:\n, raw_attention)这时我们会发现一个问题——当嵌入维度较大时点积结果的值会变得非常大导致softmax后的梯度消失。解决方案是引入缩放因子# 计算缩放后的注意力分数 d_k K.shape[1] # 这里是3 scaled_attention raw_attention / np.sqrt(d_k) print(\n缩放后的注意力分数:\n, scaled_attention)缩放因子√dₖ的引入是Transformer论文中的关键设计之一。它确保了无论嵌入维度多大注意力分数的尺度都保持稳定。3. Softmax归一化与加权求和接下来我们对每一行应用softmax将注意力分数转化为概率分布# 对每一行应用softmax def softmax(x): exp_x np.exp(x - np.max(x, axis1, keepdimsTrue)) # 数值稳定 return exp_x / np.sum(exp_x, axis1, keepdimsTrue) attention_weights softmax(scaled_attention) print(\n注意力权重:\n, attention_weights)观察这个权重矩阵你会发现每行和为1softmax保证较大的值表示更强的关注对角线元素通常较大token关注自身最后用这些权重对Value矩阵进行加权求和# 计算加权求和 output attention_weights V print(\n输出矩阵:\n, output)这个输出矩阵就是Self-Attention的最终结果——每个位置都包含了全局信息的聚合表示。4. 完整实现与常见陷阱将上述步骤整合为一个完整的Self-Attention函数def self_attention(X, W_Q, W_K, W_V): Q X W_Q K X W_K V X W_V d_k K.shape[1] attention_scores (Q K.T) / np.sqrt(d_k) attention_weights softmax(attention_scores) return attention_weights V # 测试完整函数 output self_attention(X, W_Q, W_K, W_V) print(\n完整Self-Attention输出:\n, output)实现时常见的陷阱包括忘记缩放导致softmax输入值过大梯度消失softmax应用错误应该在行方向而非整个矩阵上应用维度不匹配确保Q、K、V的维度一致数值稳定性实现softmax时需要考虑数值溢出5. 可视化理解注意力机制为了更直观地理解我们可以可视化注意力权重import matplotlib.pyplot as plt # 示例可视化第一个token的注意力分布 plt.figure(figsize(8,4)) plt.bar(range(3), attention_weights[0]) plt.xlabel(Token位置) plt.ylabel(注意力权重) plt.title(第一个token对其他token的关注程度) plt.show()这种可视化揭示了Self-Attention的核心能力——动态地为每个位置计算与其他所有位置的相关性权重。与传统的RNN不同Self-Attention能够直接捕获长距离依赖关系不受序列长度的限制。6. 多头注意力与扩展思考实际Transformer中使用的是多头注意力(Multi-Head Attention)其本质是并行运行多个Self-Attention机制def multi_head_attention(X, num_heads2): heads [] for _ in range(num_heads): # 每个头有自己的权重矩阵 W_Q np.random.randn(4, 3) W_K np.random.randn(4, 3) W_V np.random.randn(4, 3) head self_attention(X, W_Q, W_K, W_V) heads.append(head) # 拼接所有头的输出 return np.concatenate(heads, axis1) # 测试多头注意力 multi_head_output multi_head_attention(X) print(\n多头注意力输出:\n, multi_head_output)多头机制允许模型同时关注来自不同表示子空间的信息类似于CNN中使用多个滤波器提取不同特征。