最近在细读Tensor Programs的时候，发现一个有意思的细节，Yang 2022.在将mup应用到Transformer的时候，对Attention Logits使用$1/d_k$进行放缩，而不是传统上的$1/\sqrt{d_k}$。

我不禁回想起很多年前在知乎的一个问题：

transformer中的attention为什么scaled? - Nil-9的回答

当时笔者沿着Vaswani 2017.的一个脚注，展开写了这样一篇回答，其数学原理并不复杂，浓缩版本是这样的：

假设$\boldsymbol{q}$,$\boldsymbol{k}\in\mathbb{R}^{d_k}$的每个分量均是是独立采样的随机变量，均值和方差分别是0、1，$\forall i$满足分量$q_i,k_i$互相独立，则点积$\langle\boldsymbol{q},\boldsymbol{k}\rangle$的量级是$\Theta(\sqrt{d_k})$。为了防止点积过大导致softmax出现梯度消失问题，因此需要再点积基础上乘上$1/\sqrt{d_k}$，将点积稳定在$\Theta(1)$量级。

这个问题后来也成为了某些公司的所谓「机器学习八股文」之一。但这个放缩是不是严谨的呢？我们来仔细思考一下。

这里比较成问题的是「互相独立」的假设。我们先建立如下的基本的结论。

假设$X,Y$是两个随机变量，已知： $$ \mathbb{E}[X]= \mathbb{E}[Y] = 0\\ \text{Var}[X]=\text{Var}[Y] = 1 $$

那么乘积$XY$满足： $$ \mathbb{E}[XY]= \rho $$ 如果进一步地，$X,Y$均为正态分布随机变量，则有 $$ \text{Var}[XY]= 1 + \rho^2 $$ 其中$\rho$是$X$和$Y$的相关系数。

proof

按照相关系数定义： $$ \rho=\frac{\text{Cov}(X, Y)}{\sqrt{\text{Var}[X]\cdot \text{Var}[Y]}}\\ \text{Cov}(X, Y)=\mathbb{E}[XY]-\mathbb{E}[X]\mathbb{E}[Y] $$ 反过来可以得到 $$ \begin{align} \mathbb{E}[XY] &= \text{Cov}(X, Y)+\mathbb{E}[X]\mathbb{E}[Y]\\ &=\rho\cdot \sqrt{\text{Var}[X]\cdot \text{Var}[Y]}+\mathbb{E}[X]\mathbb{E}[Y]\\ &=\rho \end{align} $$

如果$X,Y$均为正态分布随机变量，按照方差的定义： $$ \begin{align} \text{Var}[XY]&=\mathbb{E}[(XY)^2]-(\mathbb{E}[XY])^2\\ &= \mathbb{E}[X^2Y^2]-\rho^2\tag{1}\\ \end{align} $$

对$\mathbb{E}[X^2Y^2]$应用Isserlis定理（显然$\mathbb{E}[X^2]=\mathbb{E}[Y^2]=1$）：

$$ \begin{align} \mathbb{E}[X^2Y^2] &= \mathbb{E}[X^2]\mathbb{E}[Y^2] + \mathbb{E}[XY]\mathbb{E}[XY] \cdot 2\\ &=1 + 2\rho^2 \end{align} $$

带入$(1)$，则可以得到 $$ \begin{align} \text{Var}[XY]&=1 + 2\rho^2-\rho^2\\ &=1+\rho^2 \end{align} $$

$\mathbb{E}[XY]$和$\text{Var}[XY]$为我们分析内积的量级随向量维度的变化关系提供了很好的工具。对于两个随机向量$\boldsymbol{x},\boldsymbol{y}\in\mathbb{R}^{n}$，而言，假设每个分量$x_i,y_i$都是独立地由某个联合正态分布中采样得到的。由于内积可以写成一系列乘积的和 $$ \langle\boldsymbol{x},\boldsymbol{y}\rangle = \sum_i x_i y_i $$

由大数定理我们有，当$n\to\infty$， $$ \frac{1}{n}\sum_{i}^n x_i y_i \overset{P}{\to} \mathbb{E}[XY] $$

由中心极限定理，有 $$ \frac{1}{\sqrt{n}}\sum_{i}^n \left(x_i y_i-\mathbb{E}[XY]\right) \overset{\mathcal{D}}{\to} \mathcal{N}\left(0, \text{Var}[XY]\right) $$

这里讨论两种典型的情况：

（1）当$x_i,y_i$完全独立时，对应$\rho=0$，$\mathbb{E}[XY]=0,\text{Var}[XY]=1$.

此时$\frac{1}{\sqrt{n}}\sum_i x_i y_i\overset{\mathcal{D}}{\to}\mathcal{N}(0, 1)$，内积的量级被中心极限定理主导，$\langle\boldsymbol{x},\boldsymbol{y}\rangle$具有$\Theta(\sqrt{n})$的典型值。

（2）当$x_i,y_i$具有一定相关性时，例如在极端情况下，二者完全正相关对应$\rho=1$，$\mathbb{E}[XY]=1,\text{Var}[XY]=2$.

此时我们主要关注大数定理下的极限$\frac{1}{n}\sum_{i}^n x_i y_i \overset{P}{\to} 1$，$\langle\boldsymbol{x},\boldsymbol{y}\rangle$具有$\Theta(n)$的典型值1。

在典型的Attention模型（e.g., Transformer）中，$\boldsymbol{q},\boldsymbol{k}$更接近于哪一种情况呢？

这里可以借用Dao 2022.的梯度推导来获取一个直观的印象（P19页）（请注意下面公式中$q_i$是指第$i$个query向量，而不是query的某个分量，其他符号类似）:

$$ dq_i = \sum_{j} dS_{ij} k_j = \sum_{j} P_{ij} (dP_{ij} - D_i) k_j = \sum_{j} \frac{e^{q_i^T k_j}}{L_i} (do_i^T v_j - D_i) k_j. $$ $$ dk_j = \sum_{i} dS_{ij} q_i = \sum_{i} P_{ij} (dP_{ij} - D_i) q_i = \sum_{i} \frac{e^{q_i^T k_j}}{L_i} (do_i^T v_j - D_i) q_i. $$

可以看到对于$q_i$而言，其梯度是$\{k_j\}$的线性组合，对$k_j$亦然。由于神经网络训练一般是由一系列梯度累积形成最终的参数&函数，因此我们可以很合理地假设$\boldsymbol{q},\boldsymbol{k}$具有一定相关性（但无法达到$\rho=1$）。

因此也就有了mup中对Attention Logits使用$1/d_k$的放缩系数的做法。

不过这里需要注意，上面的分析是在$n\to\infty$下推导出来的，在mup这样的无穷宽理论框架下使用这种假设是相对自然的。 但在典型的head维度设定下，使用$1/\sqrt{d_k}$自然也算不上不好（在初始化状态下，$\boldsymbol{q},\boldsymbol{k}$分量互相独立仍然是合理的假设）。

  1. 大数定理vs.中心极限定理的思维框架在TP中也被用于理解其他量纲。 ↩︎