Self-attention 计算 QK^T 为什么要缩放？
在使用Transformer架构中的自注意力机制时，对Query（Q）和Key（K）的点积结果进行scaling（缩放）是一种常见的技术。这种操作的目的是为了提高模型的训练稳定性和性能。以下是进行scaling的几个关键原因：
- 控制方差：当输入向量的维度很大时，向量之间的点积通常会产生很大的值。在Transformer模型中，Query和Key通常是高维空间中的向量，它们的点积结果（即注意力分数）在没有缩放的情况下可能会非常大。
- 稳定梯度：大的注意力分数在通过Softmax函数后，会导致输出分布的概率分布非常极端（即一个值非常接近1，其他都接近0），这可以导致梯度很小，从而在反向传播中产生梯度消失的问题。缩放点积可以帮助维持梯度在一个更稳定的范围内，有助于避免这一问题。
那为什么梯度会很小呢？这需要进一步分析 softmax 了。
在自注意力机制中，大的注意力分数通过Softmax函数处理后可能导致输出概率分布极端化的现象，这背后的数学原理和它如何影响梯度，是理解Transformer模型训练中遇到的一些挑战的关键。下面我们来详细解释这一过程。
Softmax 函数的基本工作原理
Softmax 函数是一种在机器学习分类任务中常用的激活函数，它可以将一个含任意实数的K维向量“压缩”成另一个K维实向量，其中每一个元素的范围都是0到1之间，并且所有元素的和为1。这使得Softmax 函数的输出可以被解释为一个概率分布。函数的定义如下：
这里， $( z_i )$ 是输入向量中的第i个元素，分母是对整个输入向量的元素进行指数函数计算后求和的结果。
大的注意力分数如何导致极端概率分布
当Softmax 输入向量中的某些元素（注意力分数）非常大时，它们的指数 $e^{z_i}$ 会相对其他较小元素的指数显著增大。这种情况下，最大的数值会在Softmax的输出中占据主导地位，其对应的输出概率接近1，而其他较小的数值对应的输出概率则接近0。
梯度的影响
Softmax函数的这种特性在反向传播过程中会导致梯度问题。在机器学习中，我们通过反向传播算法调整模型的参数，这一过程依赖于损失函数对模型参数的梯度。Softmax函数的梯度由下面的公式给出：
其中， $( \delta_{ij} )$ 是克罗内克Delta函数，当 $i = j$ 时值为1，否则为0。
在极端概率分布的情况下，如果某一输出概率 $\text{Softmax}(z_i)$ 接近1而其他接近0，那么这个输出概率对于大部分输入 $z_j$ 的梯度 $\text{Softmax}(z_i)(1 - \text{Softmax}(z_j))$ 将会非常小，因为 $\text{Softmax}(z_j)$ 接近0。这意味着模型在这种状态下几乎无法从数据中学习，因为梯度消失使得参数更新非常缓慢。
这就是为什么在设计Transformer模型时，通过对Query和Key的点积结果进行适当的缩放（如除以 $\sqrt{d_k}$ ，其中 $d_k$ 是Key向量的维度），可以帮助避免Softmax输出过于极端，保持梯度在一个合理的范围内，从而改善模型的训练效率和稳定性。