学习 Transformer 的初始化、参数化与标准化

尝试学习苏神的文章：https://kexue.fm/archives/8620 并做一些记录：

采样分布🤔

模型的初始化是随机采样的，一般情况下我们都是从指定均值 $\mu$ 和方差 $\sigma^2$ 的随机分布中进行采样来初始化，其中常用的随机分布有三个：正态分布，均匀分布，截尾正态分布（Truncated Normal）。

其中正态分布通常记为 $\mathcal{N}(\mu, \sigma^2)$；区间 $[a,b]$ 上的均匀分布一般记为 $U[a,b]$，其均值为 $\frac{a+b}{2}$，方差为 $\frac{(b-a)^2}{12}$，所以如果指定 $u$ 和 $\sigma^2$ 的话，对应的均匀分布为 $U[\mu-\sqrt{3}\sigma,\mu+\sqrt{3}\sigma]$。

一般来说正态分布的采样结果更多样化一些，但理论上他是无界的，如果采样到绝对值过大的结果可能不利于优化；而均匀分布是有界的，但是采样结果更单一。结合两者优点即可得到 截尾正态分布，他从 $\mathcal{N}(\mu, \sigma^2)$ 中采样，并确保采样结果在 $[a,b]$ 中（采样结果在区间中则保留，否则重复采样至符合区间要求）。

苏神提到 Tensorflow 中的 tf.random.truncated_normal，该函数采样结果实际均值为 $u$，而实际方差为 $\gamma\sigma^2$，其中 $\gamma=0.7737413…$（具体的推算过程我不会，我是菜鸡😭），但是我使用 Pytorch 的相似函数 torch.nn.init.trunc_normal_ 进行了试验（该函数的截断区间参数 $a$ 和 $b$ 默认是 2，但是可以自己调整），使用默认参数值时结果也是一样的：

import torch

w = torch. empty(300, 500)
torch.nn.init.trunc_normal_(w, mean=0, std=1)
# 计算均值
mean = torch.mean(w)
print(f"Empirical mean: {mean:.6f}")
# 计算方差
variance = torch.var(w)
print(f"Empirical variance: {variance:.6f}")

1 2	Empirical mean: -0.000635 Empirical variance: 0.778056

所以也可以使用苏神提到的同样的方法，若想得到方差为 $\sigma^2$ 的采样结果，传入函数的标准差为 $\frac{\sigma}{\sqrt{\gamma}} = \sigma \times 1.1368472…$

稳定二阶矩

一般的教程中，推导初始化方法的思想是尽量让输入输出具有同样的均值和方差，通常会假设输入是 $u$ 为 1 ，$\sigma^2$ 为 1 的随机向量，然后试图让输出的 $u$ 为 0，$\sigma^2$ 为 1。但是苏神说认为其没有必要，而且对于某些非负的激活函数来说（比如 ReLU），也做不到均值为 0。事实上，我们只需要一个衡量某个指标是否“消失”或者“爆炸”的指标，0均值、1方差是 非必要 的，这里我们用二阶（原点）矩来代替，它可以看成是 $L2$ 模长的变体，跟方差的作用类似，都可以用来衡量指标是否“消失”或者“爆炸”，但它相对来说更普适和简单。

二阶矩 通常指的是随机变量关于原点的二阶统计量，即 $E[X^2]$，其计算方式为：
$E[X^2] = Var(X) + (E[X])^2$
其中 $Var(X)$ 为随机变量 $X$ 的方差，当随机变量的均值 $E[X]$ 为0，此时二阶矩 $E[X^2]$ 也就等于方差 $Var(X)$；相比于方差，二阶矩不需要计算均值，因此在某些情况下更简洁，尤其是在均值不为零的情况下仍能有效衡量信号的幅度。

$L2$ 范数是向量空间中的一种范数，定义为向量元素的平方和的平方根，即：
$||x||_2 = \sqrt{\sum_{i}x_i^2}$

我们先看无激活函数的全连接层（输入数为 $m$，输出数为 $n$）：

$y_j = b_j + \sum_{i}x_iw_{i,j}$

简单起见，通常对偏置 $bj$ 使用全零初始化，并将 $w{i,j}$ 的均值 $E[w_{i,j}]$ 也设为 0，这有助于简化下面的结果，计算二阶矩：

$\begin{align} E[y_j^2] &= E\left[\left(\sum_ix_iw_{i,j}\right)^2\right] = E\left[\left(\sum_{i_1}x_{i_1}w_{i_1,j}\right)\left(\sum_{i_2}x_{i_2}w_{i_2,j}\right)\right] \\ &= E\left[\sum_{i_1,i_2}(x_{i_1}x_{i_2})(w_{i_1,j}w_{i_2,j})\right] = \sum_{i_1,i_2}E[x_{i_1}x_{i_2}]E[w_{i_1,j}w_{i_2,j}] \end{align}$

注意 $w{i_1,j},w{i2,j}$ 是独立同分布的，所以当 $i_1 \neq i_2$ 时 $E[w{i1,j}w{i2,j}] = E[w{i1,j}]E[w{i_2,j}] = 0$，因此只需要考虑 $i_1 = i_2 = i$ 的情况，假设输入的二阶矩为 1，那么：

$E[y_j^2] = \sum_iE[x_i^2]E[w_{i,j}^2] = mE[w_{i,j}^2]$

所以要使得 $E[yj^2]$ 为 1，那么 $E[w{i,j}^2] = 1/m$，综合均值为 0 的假设，我们得到 $w_{i,j}$ 初始化策略为：

从均值为 0、方差为 $1/m$ 的随机分布中独立重复采样，这就是 Xavier 初始化，该过程我们并没有对输入的均值做任何假设，因此它哪怕全是非负的也没问题。