第03章:Scaled Dot-Product Attention——那个√d_k到底在防什么?

39 阅读7分钟

第03章:Scaled Dot-Product Attention——那个√d_k到底在防什么?

论文链接:Attention Is All You Need (Vaswani et al., NIPS 2017) 本章对应:Section 3.2.1, Footnote 4

本系列包含12章深度技术文章,涵盖Attention机制、架构设计等核心内容。 由于掘金平台限制,完整版(含架构图和数学公式)请访问: 📖 GitHub: <github.com/Yunzenn/blo… 或其他博客平台同名账号

核心困惑

为什么Attention公式里要除以dk?不除会怎样?

这个dk看起来不起眼,但它在防一个致命问题。如果你在面试时被问”Scaled Dot-Product Attention的Scaled是什么意思”,答不出来就直接挂了。

完整的Attention公式是:

Attention(Q,K,V)=softmax(QKTdk)V

为什么是dk而不是dk或者dk2?这背后有严格的数学推导。

前置知识补给站

1. 向量点积的几何意义

两个向量q和k的点积:

q⋅k=∑i=1dkqiki=‖q‖‖k‖cos⁡θ

几何意义:

  • 点积衡量两个向量的相似度
  • 点积越大,向量越相似(夹角越小)
  • 点积可以是负数(夹角>90°)

2. 随机变量的方差

对于随机变量X:

Var(X)=E[(X−E[X])2]=E[X2]−(E[X])2

方差的性质:

  • 如果和独立:

3. Softmax函数的饱和区

Softmax函数:

饱和区问题:

  • 当某个远大于其他时,,其他位置
  • 此时梯度()
  • 这叫”饱和”,会导致梯度消失

论文精读:为什么需要缩放?

原论文的解释

Section 3.2.1:

“We call our particular attention ‘Scaled Dot-Product Attention’. The input consists of queries and keys of dimension , and values of dimension . We compute the dot products of the query with all keys, divide each by , and apply a softmax function to obtain the weights on the values.”

Footnote 4(关键):

“We suspect that for large values of , the dot products grow large in magnitude, pushing the softmax function into regions where it has extremely small gradients. To counteract this effect, we scale the dot products by .”

翻译成人话:

  1. 当很大时,点积的值会很大
  2. 大的点积会让Softmax进入饱和区
  3. 饱和区的梯度极小,导致梯度消失
  4. 除以可以把点积的值控制在合理范围

但原论文没有证明”为什么是“。我们来严格推导。

第一性原理推导:为什么是√d_k?

推导1:随机向量点积的方差

假设(理想化数学模型):

  • 和是维随机向量
  • ,,其中和在初始化时独立
  • 每个分量的均值为0,方差为1
  • 的各分量不相关,的各分量不相关

注:实践中这些假设通过LayerNorm近似满足。

目标:计算点积的方差。

推导:

首先计算期望:

在初始化时,和独立初始化,因此对于固定的输入,和条件独立:

由于投影后的向量均值为0(假设),因此:

接下来计算方差:

对于的交叉项,由于和来自不同的投影矩阵:

如果的各分量不相关且均值为0,则()。同理。因此交叉项期望为0。

对于对角项:

因此:

结论:点积的方差是。

标准差:

归一化:如果我们除以:

这就是为什么是:它把点积的方差归一化到1,无论多大。

推导2:Softmax饱和区的数值演示

问题:为什么点积过大会导致Softmax饱和?

Softmax的梯度:

数值示例:

假设有3个位置,点积为。

情况1:点积适中(,已缩放)

梯度正常,可以学习。

情况2:点积过大(,未缩放)

梯度几乎为0,无法学习。

可视化对比:

d_k标准差3σ范围Softmax状态梯度
64(已缩放)1[-3, 3]✅ 正常~0.2
64(未缩放)8[-24, 24]❌ 饱和~0.0
512(未缩放)22.6[-68, 68]❌ 极饱和~0.0

推导3:为什么不是其他缩放因子?

候选方案:

  1. 除以:,方差太小
  2. 除以:,方差刚好
  3. 除以:,方差仍然随增长

只有能把方差归一化到1。

Scaled vs Unscaled的实验对比

原论文没有直接对比Scaled和Unscaled的实验(因为缩放的必要性可以通过数学推导证明)。但Table 3 row (B)的消融实验说明:在缩放了的情况下,的维度选择仍然重要。

原论文Table 3 row (B):

  • : PPL 5.16, BLEU 25.1
  • : PPL 5.01, BLEU 25.4
  • (base): PPL 4.92, BLEU 25.8

解读:

  • 越大,效果越好(在缩放了的情况下)
  • 这说明更大的能提供更丰富的表示能力
  • 但如果不缩放,大会导致Softmax饱和,效果反而变差

如果不缩放会怎样:

  • 时,点积方差是64,标准差是8
  • Softmax输入的范围是(3倍标准差)
  • 这会导致严重的饱和

缩放后:

  • 点积方差归一化到1,标准差是1
  • Softmax输入的范围是
  • 梯度正常,可以训练

Dot-Product Attention vs Additive Attention

原论文提到了两种Attention机制:

1. Dot-Product Attention(原论文使用)

优点:

  • 计算高效:矩阵乘法,可以用高度优化的BLAS库
  • 并行性好:所有位置的score可以同时计算

缺点:

  • 需要缩放(否则方差随增长)

2. Additive Attention(Bahdanau et al., 2015)

优点:

  • 不需要缩放(tanh自带归一化)
  • 理论上表达能力更强(有非线性变换)

缺点:

  • 计算慢:需要两次矩阵乘法 + tanh
  • 参数多:需要

原论文Section 3.2.1:

“While for small values of the two mechanisms perform similarly, additive attention outperforms dot product attention without scaling for larger values of . We suspect that for large values of , the dot products grow large in magnitude, pushing the softmax function into regions where it has extremely small gradients. To counteract this effect, we scale the dot products by .”

结论:

  • 小时:Dot-Product和Additive差不多
  • 大时:Dot-Product需要缩放,否则不如Additive
  • 缩放后:Dot-Product更快,效果相当

完整的Scaled Dot-Product Attention流程

添加图片注释,不超过 140 字(可选)

关键步骤:

  1. 计算相似度:,得到的score矩阵
  2. 缩放:除以,把方差归一化到1
  3. Mask(可选):Decoder需要mask掉未来位置。注:Mask在Scale之后、Softmax之前。Mask的值通常是,Scale后仍然是极小值,因此先Scale后Mask是工程习惯
  4. 归一化:Softmax,把score转为概率分布
  5. 加权求和:用概率加权V,得到输出

2026年的批判性视角

1. 缩放因子的理论假设

原论文的推导假设独立同分布,均值0,方差1。

实际情况:

  • 经过LayerNorm后,这个假设基本成立
  • 但在训练初期,分布可能不稳定
  • 这可能是为什么需要warmup的原因之一

2. 其他缩放方案

后续研究提出了其他缩放方案:

T5的缩放(Raffel et al., 2020):

ALiBi的缩放(Press et al., 2022):

不需要除以,而是用位置偏置来控制范围。

3. Softmax的替代方案

Softmax不是唯一的归一化方案:

ReLU Attention(Shen et al., 2021):

不需要缩放,因为ReLU自带截断。

Linear Attention(Katharopoulos et al., 2020):

用特征映射替代Softmax,复杂度降到。

4. 原论文没有讨论的问题

  • 为什么不用BatchNorm:BatchNorm也能归一化方差,为什么用除以?

  • 答:BatchNorm需要统计整个batch的均值和方差,而除以是确定性的,推理时更稳定

  • 缩放因子是否可学习:能否把改成可学习的参数?

  • 答:可以,但实验表明固定的已经足够好

面试追问清单

⭐ 基础必会

  1. 为什么Attention要除以?
  • 提示:从点积的方差推导
  1. 如果不除以会怎样?
  • 提示:Softmax饱和、梯度消失
  1. 为什么是而不是?
  • 提示:方差的性质,

⭐⭐ 进阶思考

  1. 证明:如果独立同分布,均值0,方差1,则。
  • 提示:展开,利用独立性消去交叉项
  1. Dot-Product Attention和Additive Attention有什么区别?为什么Transformer选择Dot-Product?
  • 提示:计算效率、并行性
  1. 如果很小(如8),还需要缩放吗?
  • 提示:原论文说”for small values of the two mechanisms perform similarly”

⭐⭐⭐ 专家领域

  1. 原论文假设均值为0。但实际训练中,如何保证这个假设成立?
  • 提示:LayerNorm的作用
  1. 能否设计一个不需要缩放的Attention机制?
  • 提示:Additive Attention、ReLU Attention、Linear Attention
  1. 如果把缩放因子从改成可学习的参数,会有什么影响?
  • 提示:增加了灵活性,但可能过拟合;实验表明固定的已经足够

下一章预告:第04章将深入拆解Multi-Head Attention,回答”八个头,八个视角,还是八份低秩分解?”

论文原文传送门: