第七篇：自注意力（二）—— 缩放、Softmax 与信息融合

第七篇自注意力二—— 缩放、Softmax 与信息融合系列文章第一篇预训练模型——站在巨人的肩膀上第二篇分词——文字如何变成数字第三篇向量与矩阵——理解一切的基石第四篇Embedding 层——从 ID 到向量的第一步第五篇位置编码——给 Transformer 装上顺序感第六篇自注意力一—— Q、K、V 的由来 本文目录开篇一、为什么需要缩放二、Softmax把分数变成概率三、加权求和融合信息完整的注意力公式代码实现完整的单头自注意力关键结果上下文感知参考资源本篇小结参考答案开篇前六篇我们完成了从文本到注意力分数的全部准备文本 → 分词 → Embedding → 位置编码 → Q、K、V → Q × Kᵀ 注意力分数 我喜欢人工智能 → [101, 2769, 4263, ...] → [0.021, -0.015, ...] → 位置向量 → × W_Q/W_K/W_V → Q, K, V → Q × Kᵀ [9, 9] 分数矩阵 ← 上一篇到这里但我们拿到的原始分数有两个问题数值范围不稳定点积结果随向量维度增大而发散不是概率无法直接作为权重来加权求和这篇我们将完成自注意力的后半程缩放 → Softmax → 加权求和。三个步骤一步比一步精妙。到这篇文章结束你将完整理解自注意力机制的每一个数学细节并能在代码中实现它。一、为什么需要缩放问题点积随维度变大而发散先看一个简单例子2 维向量 q [0.5, 0.5], k [0.5, 0.5] q·k 0.25 0.25 0.5 4 维向量 q [0.5, 0.5, 0.5, 0.5], k [0.5, 0.5, 0.5, 0.5] q·k 0.25×4 1.0 64 维向量BERT 每个注意力头的维度 q [0.5, 0.5, ..., 0.5], k [0.5, 0.5, ..., 0.5] q·k 0.25×64 16.0维度越大点积越大。快速回顾方差和标准差方差衡量一组数据有多分散。算法每个数减去平均值平方再求平均。数据[1, 2, 3, 4, 5]平均值 3 方差 ((1-3)² (2-3)² (3-3)² (4-3)² (5-3)²) / 5 (4 1 0 1 4) / 5 2 标准差 √方差 √2 ≈ 1.41方差大→ 数据很散比如 [1, 100]方差很大方差小→ 数据很集中比如 [49, 50, 51]方差很小标准差→ 方差的平方根和原始数据同单位更直观虽然实际向量不是全 0.5但统计上如果 q 和 k 是独立同分布随机变量均值为 0方差为 1则点积的均值为 0方差为 d_kd_k 就是向量的维度BERT 每个注意力头是 64。维度 d_k 越大分数的分布越散。这就是所谓的方差膨胀问题——注意力分数的方差正比于向量维度 d_k。为什么分数太散是问题——梯度消失来比较有缩放和没缩放的效果假设一组分数[10, 11, 9]d_k64 时可能的量级 Softmax (无缩放) e¹⁰ ≈ 22026, e¹¹ ≈ 59874, e⁹ ≈ 8103 总和 90003 P [0.245, 0.665, 0.090] ← 基本正常 假设分数[40, 44, 36]d_k1024 时可能的量级 Softmax (无缩放) e⁴⁰ ≈ 2.35e17, e⁴⁴ ≈ 1.29e19, e³⁶ ≈ 4.31e15 总和 1.31e19 P [0.018, 0.982, 0.0003] ← 最大值几乎占据全部概率当 softmax 输出极度接近 one-hot 时最大值位置以外的梯度 ≈ 0被压制的 token 对应的梯度在反向传播中几乎为零模型学不动了这个问题的数学根源是指数函数 exp(x) 对大的输入非常敏感——exp(44) / exp(36) ≈ e⁸ ≈ 2980小的差距被指数函数急剧放大。解法除以 √d_kBERT base hidden_size 768 num_heads 12 d_k hidden_size ÷ num_heads 768 ÷ 12 64 √d_k √64 8 缩放公式scaled_score raw_score ÷ √d_k除以 √d_k 后分数的方差稳定在 1 左右softmax 的梯度始终处于甜区。缩放要解决什么问题不缩放时分数会很大 q·k 64个数相加每个约±1 → 总和可能达到 ±20 甚至更大 大分数对 softmax 的影响 输入 [20, -10, 5] softmax → [0.9999, 0.0000, 0.0001] ← 极端分布接近 one-hot 问题这种极端分布下梯度几乎为 0 假设正确答案是第1个 梯度 [0.9999-1, 0.0000-0, 0.0001-0] ≈ [-0.0001, 0, 0] → 模型几乎学不到东西缩放后的效果除以 √64 8 后分数变小了 原分数 [20, -10, 5] → 缩放后 [2.5, -1.25, 0.625] 缩放后 softmax 的输出 [0.76, 0.02, 0.22] ← 有区分度但不极端 梯度正常了 梯度 [0.76-1, 0.02-0, 0.22-0] [-0.24, 0.02, 0.22] → 模型能正常学习总结缩放让 softmax 的输入保持在合理范围避免输出极端化从而保证梯度能正常传播。二、Softmax把分数变成概率Softmax 公式与直觉exp(xᵢ) softmax(xᵢ) ──────────────── Σⱼ exp(xⱼ)两个关键操作操作作用直观理解exp指数函数放大分数差距好的更好、差的更差归一化除以总和保证权重之和为 1形成竞争性概率分布为什么不用简单归一化除以最大值简单归一化允许负值且不是概率分布。Softmax 保证所有输出为正且和为 1——这是概率的基本要求。完整演算以喜为例以喜这行为例展示完整的缩放 Softmax 过程第1步原始分数来自 Q×Kᵀ [15.3, 22.7, 18.2, 2.1, 1.8, 3.2, 2.9, 5.6, 4.1] 对应 我 喜 欢 人 工 智 能 CLS SEP 第2步缩放除以 √d_k 8 [1.9125, 2.8375, 2.275, 0.2625, 0.225, 0.4, 0.3625, 0.7, 0.5125] 第3步对每个数求 exp e^1.91256.77, e^2.837517.07, e^2.2759.73, e^0.26251.30, e^0.2251.25, e^0.41.49, e^0.36251.44, e^0.72.01, e^0.51251.67 第4步分母 所有 exp 之和 42.73 第5步每个 exp ÷ 分母得到概率权重 6.77/42.730.16 ← 我16% 的注意力 17.07/42.730.40 ← 喜40%自己最高 9.73/42.730.23 ← 欢23%组成喜欢 1.30/42.730.03 1.25/42.730.03 1.49/42.730.03 1.44/42.730.03 2.01/42.730.05 1.67/42.730.04 ──── 合计 1.00 ✅Softmax 效果深度分析缩放前分数15.3 22.7 18.2 2.1 1.8 3.2 2.9 5.6 4.1 缩放后分数1.91 2.84 2.28 0.26 0.22 0.40 0.36 0.70 0.51 Softmax概率0.16 0.40 0.23 0.03 0.03 0.03 0.03 0.05 0.04 效果对比 - 最高分(22.7) 在 softmax 后占比 40% ← 强者恒强 - 次高分(18.2) 在 softmax 后占比 23% ← 第二强 - 低分(1.8) 在 softmax 后占比 3% ← 弱者愈弱 - 分数差别被非线性放大了缩放对 softmax 的影响实验importnumpyasnpdefsoftmax(x):e_xnp.exp(x-np.max(x))# 数值稳定版returne_x/e_x.sum()scoresnp.array([15.3,22.7,18.2,2.1,1.8])# 不缩放print(不缩放:,softmax(scores))# [0.001, 0.997, 0.002, ≈0, ≈0]# 缩放÷8print(缩放:,softmax(scores/8))# [0.187, 0.473, 0.270, 0.036, 0.035]不缩放的 softmax 几乎退化为 argmaxhard attention梯度消失。缩放后的 softmax 保留了更丰富的权重分布soft attention梯度良好。三、加权求和融合信息计算过程有了注意力权重概率最后一步就是用这些权重去加权求和所有字的 Value 向量喜的新向量 0.16 × V(我) 0.40 × V(喜) 0.23 × V(欢) 0.03 × V(人) 0.03 × V(工) 0.03 × V(智) 0.03 × V(能) 0.05 × V(CLS) 0.04 × V(SEP)展开前几步只展示前 3 维真实是 768 维V₁(我) [ 0.31, -0.12, 0.45, ...] × 0.16 → [ 0.050, -0.019, 0.072, ...] V₂(喜) [ 0.42, 0.15, -0.08, ...] × 0.40 → [ 0.168, 0.060, -0.032, ...] V₃(欢) [ 0.05, 0.28, -0.18, ...] × 0.23 → [ 0.012, 0.064, -0.041, ...] ... 全部9个加起来 ... 结果喜的新向量 [0.215, 0.142, -0.018, ...] (768维)代码实现# 假设已经算好了 attention_weights [1, 9, 9] 和 V [1, 9, 768]attn_outputtorch.matmul(attention_weights,V)print(attn_output.shape)# torch.Size([1, 9, 768])# 第2个字喜的新向量print(attn_output[0,1,:5])融合后的效果加权求和前喜的原始向量 [0.862, 0.525, 0.043, ...] ← 只有喜自己的信息 加权求和后喜的新向量 [0.215, 0.142, -0.018, ...] ← 已融合了我和欢的信息这个新向量的关键属性信息来源权重效果欢的信息23%喜欢整体含义被编码 ✅我的信息16%谁喜欢也被编码了 ✅“人、工、智、能”各 ~3%本字不太需要宾语信息 ✅喜自身40%保留自身语义 ✅一个类比集体智慧自注意力就像一个团队讨论 每个成员字先写下自己的想法V 再告诉别人我擅长什么K 同时提出我需要什么Q。 然后大家投票决定谁的发言对我最有价值 票数 softmax(Q×K/√d) 最后每个人综合所有人的意见更新自己的认识 新想法 加权求和(所有人的意见) 这个过程在每一层重复一次。完整的注意力公式三步合在一起就是论文里的那个经典公式Attention(Q, K, V) softmax(Q × Kᵀ / √d_k) × V 拆开看 ① Q × Kᵀ → 原始注意力分数 [9, 9] ② ÷ √d_k → 缩放方差稳定到1 ③ softmax → 变成概率分布每行和为1 ④ × V → 加权求和输出 [9, 768]代码实现完整的单头自注意力importtorchimporttorch.nnasnnimportmathclassSingleHeadAttention(nn.Module):完整的单头自注意力def__init__(self,hidden_size768):super().__init__()self.W_Qnn.Linear(hidden_size,hidden_size)self.W_Knn.Linear(hidden_size,hidden_size)self.W_Vnn.Linear(hidden_size,hidden_size)self.d_khidden_size# 每个头的维度defforward(self,x):Qself.W_Q(x)# [B, L, 768]Kself.W_K(x)# [B, L, 768]Vself.W_V(x)# [B, L, 768]# ① Q × Kᵀ → 注意力分数scorestorch.matmul(Q,K.transpose(-2,-1))# [B, L, L]# ② 缩放scoresscores/math.sqrt(self.d_k)# ③ Softmax → 权重attention_weightstorch.softmax(scores,dim-1)# [B, L, L]# ④ 加权求和 Voutputtorch.matmul(attention_weights,V)# [B, L, 768]returnoutput,attention_weights# 模拟 我喜欢人工智能 (9个token)xtorch.randn(1,9,768)attnSingleHeadAttention()output,weightsattn(x)print(f输出形状:{output.shape})# [1, 9, 768]print(f权重形状:{weights.shape})# [1, 9, 9]print(f权重每行之和:{weights[0,2].sum()})# ≈ 1.0输出输出形状: torch.Size([1, 9, 768]) 权重形状: torch.Size([1, 9, 9]) 权重每行之和: tensor(1.0000)关键结果上下文感知对比一下加权求和前后的变化喜 的原始向量Embedding 位置编码后 [0.862, 0.525, 0.043, ...] ← 只有喜自己的信息 喜 经过自注意力后的新向量 [0.215, 0.142, -0.018, ...] ← 融合了我(16%)和欢(23%)的信息这个新向量已经不再是孤立的喜——它包含了我主语和欢搭配词的语义信息。这就是上下文感知的含义同一个字在不同句子里经过自注意力后会得到不同的向量。参考资源[1] Vaswani et al., “Attention Is All You Need”, 2017. NeurIPS 2017 — 缩放点积注意力的原始论文[2] Jay Alammar, “The Illustrated Transformer”, 2018. Blog — Self-Attention 计算流程的可视化讲解[3] Stack Overflow, “Why do we scale by sqrt(d_k) in attention?”, 2020. Link — 关于缩放因子的讨论本篇小结本篇完成了自注意力的后半程缩放 → Softmax → 加权求和。步骤做了什么为什么缩放÷√d_k把分数除以 √648防止方差膨胀导致梯度消失Softmax把分数变成概率和为1形成竞争性权重分布加权求和×V按权重融合所有字的信息输出包含上下文的新向量完整公式Attention(Q, K, V) softmax(QKᵀ/√d_k) × V 输入 X → Q,K,V → Q×Kᵀ → ÷√d_k → softmax → ×V → 输出 上一篇 本篇三步 上下文感知的向量下一篇我们深入解读注意力矩阵的语言模式解释为什么 Q、K、V 必须是三个独立矩阵并用完整代码实现自注意力亲眼见证上下文感知如何诞生。思考题题1为什么 Softmax 前要除以 √d_k题2Softmax 的两个关键操作分别是什么题3加权求和后喜的新向量里包含了谁的信息答案见文末 → 参考答案参考答案题1为什么 Softmax 前要除以 √d_k因为点积的方差正比于维度 d_k。不缩放的话分数范围随维度增大而膨胀Softmax 输出会退化为接近 one-hot 的分布导致梯度消失模型无法学习。题2Softmax 的两个关键操作分别是什么exp指数函数放大分数差距让好的更好、差的更差归一化除以总和保证所有权重为正且和为 1形成概率分布。题3加权求和后喜的新向量里包含了谁的信息主要包含我16%和欢23%的信息——因为喜欢是一个词我是主语。人、工、智、能各只占约3%因为和喜语义关系较远。