Transformer自注意力机制中除以$\sqrt{d_k}$深度剖析

【 Transformer 系列，故事从$\sqrt{d_k}$说起】

LLM这么火，Transformer厥功甚伟，某天心血来潮~，再去看看！

它长这个样子： 深入浅出 Transformer

看完后，想起了老生常谈$\sqrt{d_k}$问题，必须一探究竟：Transformer 中缩放点积注意力机制探讨：除以根号 dk 理由及其影响

感觉不够清楚，还是再Review下考研概率论，有了：基于考研概率论知识解读 Transformer：为何自注意力机制要除以根号 dk，中间会涉及初始化、标准化、Sofrmax函数，于是继续

【初始化相关】：深度学习中的常见初始化方法：原理、应用与比较
【标准化系列】： 数据为什么要进行标准化：Z-标准化的神奇蜕变，带出了关联知识点: 深度 “炼丹” 术之 Batch Normalization 与 Z - 标准化：开启数据的神秘转换
【Softmax复习】：Softmax 层反向传播梯度计算实例解析，中间想到了经常配套使用的交叉熵，于是梳理了交叉熵的前世今生

KL 散度：多维度解读概率分布间的隐秘 “距离”
熵与交叉熵：从不确定性角度理解 KL 散度
机器学习、深度学习关于熵你所需要知道的一切

摘要

本文深入探讨了Transformer自注意力机制中除以$\sqrt{d_k}$这一关键操作的原因。通过详细的推导过程揭示$\sqrt{d_k}$的来源，并结合Softmax函数的特性，分析不除以$\sqrt{d_k}$以及除以结果偏离$\sqrt{d_k}$时对模型造成的后果及其内在原因，旨在为理解Transformer的工作原理提供全面且深入的视角。

引言

Transformer架构在自然语言处理及其他诸多领域取得了巨大成功，其自注意力机制是核心创新点之一。在自注意力机制的计算过程中，除以$\sqrt{d_k}$（其中$d_k$是键（Key）向量的维度）这一操作对模型的稳定性和性能起着至关重要的作用，本文结合考研中的概率知识对除以$\sqrt{d_k}$进行理解。

---

一、考研概率论内容复习

在Transformer自注意力机制中关于方差推导主要用到了以下考研概率知识：

1. 期望与方差的基本定义及性质
   • 期望：期望 $E(X)$ 表示随机变量 $X$ 取值的平均水平。
   • 方差：方差 $Var(X)=E[(X-E(X){)}^2]=E[X^2]-(E[X]{)}^2$，用于衡量随机变量取值的离散程度。
2. 独立随机变量的性质
   • 若 $X$ 和 $Y$ 是两个独立的随机变量，则 $E[XY]=E[X]E[Y]$ 。
   • 对于独立随机变量 $X$ 和 $Y$，$Var(XY)=E[X^2]E[Y^2]-(E[X]E[Y]{)}^2$ 。
   • 若 $X_1,X_2,\cdots ,X_n$ 相互独立，那么 $Var({\sum }_{i=1}^n{X}_i)={\sum }_{i=1}^{n}Var(X_i)$ 。
3. 随机变量的分布
   假设随机变量 $q_i$ 和 $k_i$ 独立同分布，简化了分析过程，使得可以基于相同的分布特性对所有变量进行统一处理。

二、选择除以$\sqrt{d_k}$的原因

1. 稳定方差

在自注意力机制中，计算注意力分数时涉及查询（Query）向量与键（Key）向量的点积。假设Query向量$Q$和Key向量$K$的元素是独立同分布的随机变量，均值为0，方差为1（实际模型参数随机初始化时接近此假设）。对于长度为$d_k$的向量$Q=(q_1,q_2,\cdots ,q_{d_k})$和$K=(k_1,k_2,\cdots ,k_{d_k})$，它们的点积$$S=\sum _{i=1}^{d_k}{q}_i{k}_i$$

根据方差性质，$$Var(S)=\sum _{i=1}^{d_k}Var(q_i{k}_i)$$由于$q_i$和$k_i$相互独立，$$Var(q_i{k}_i)=E[q_i^2]E[k_i^2]-(E[q_i]E[q_i]{)}^2$$已知$E[q_i]=E[k_i]=0$，$Var(q_i)=Var(k_i)=1$，则$E[q_i^2]=E[k_i^2]=1$，所以$Var(q_i{k}_i)=1$，进而$$Var(S)=d_k$$

这表明点积$S$的方差随维度$d_k$增大而线性增大。为使点积结果方差稳定，除以$\sqrt{d_k}$，即新变量$S^{\prime }=\frac{S}{\sqrt{d_k}}$的方差$$Var(S^{\prime })=\frac{1}{d_k}Var(S)=1$$稳定的方差有助于后续Softmax函数输入处于合理范围，避免其输出概率分布过于极端，保证模型训练稳定性。

2. 平衡维度影响

除以$\sqrt{d_k}$可视为一种归一化操作，有助于平衡不同维度特征在注意力计算中的贡献。若不进行此缩放，高维度向量在点积运算中可能占据主导，使模型过度关注某些维度信息，忽略其他维度。通过除以$\sqrt{d_k}$，不同维度影响更均衡，模型能全面考虑序列元素关系。

三、$\sqrt{d_k}$的来源推导

如前文所述，基于随机向量点积方差分析。设$Q$和$K$向量元素为独立同分布随机变量，均值为0，方差为1。点积$$S=\sum _{i=1}^{d_k}{q}_i{k}_i$$，其方差$Var(S)=d_k$。

为将点积结果方差稳定在1附近，需对其归一化。自然的想法是除以与方差相关量，即除以$\sqrt{Var(S)}$，也就是$\sqrt{d_k}$。这样处理后，点积结果的方差得以稳定，为后续Softmax函数提供合适输入。

四、不恰当缩放的后果

1. 不除以$\sqrt{d_k}$

若不除以$\sqrt{d_k}$，随着$d_k$增大，点积结果方差增大。输入到Softmax函数的数值方差过大，会使Softmax输出概率分布极其不均匀，大部分概率集中在少数较大数值对应元素上。

在反向传播中，这种不均匀分布导致梯度问题。一方面，概率接近0的元素对应梯度极小，在反向传播中对梯度贡献可忽略不计，多个此类元素累积使得梯度消失，模型参数更新缓慢甚至停滞。另一方面，概率接近1的元素对应梯度可能极大，在反向传播中累积导致梯度爆炸，使模型参数更新幅度过大，无法收敛。

2. 除以结果大于$\sqrt{d_k}$

若除以大于$\sqrt{d_k}$的值，如$\alpha \sqrt{d_k}$（$\alpha >1$），会过度缩放点积结果。这使得Softmax函数输入数值范围变小，输出概率分布相对均匀，各元素概率差异不明显。

在这种情况下，模型难以区分不同元素重要性，无法有效捕捉序列中关键信息，导致模型性能下降。因为自注意力机制依赖概率分布差异来聚焦重要信息，过于均匀分布削弱了这种聚焦能力。

3. 除以结果小于$\sqrt{d_k}$

若除以小于$\sqrt{d_k}$的值，如$\beta \sqrt{d_k}$（$0<\beta <1$），点积结果缩放不足，方差仍较大。Softmax函数输入数值方差大，输出概率分布仍会过度集中，类似不除以$\sqrt{d_k}$的情况，导致梯度消失或爆炸问题，阻碍模型正常训练。

五、结论

在Transformer自注意力机制中，除以$\sqrt{d_k}$是经过精心设计的关键操作。它源于对随机向量点积方差的理论分析，旨在稳定方差、平衡维度影响。不恰当的缩放，无论是不除以$\sqrt{d_k}$，还是除以结果偏离$\sqrt{d_k}$，都会因Softmax函数特性引发梯度问题或信息捕捉能力下降，严重影响模型性能。理解这一操作的原理和影响，对于深入理解Transformer架构及优化相关模型具有重要意义。