首页LLM

DeepNet

本篇文章作为Transformer 中的 Layer Normalization与梯度稳定性的后继文章，进一步进行Layer Normalization的相关学习

Warm-up 与初始化方式对于梯度下降稳定性的影响

在本节设置了三个实验组进行对比

  • 对照组1: Post-LN + 标准Xavier初始化 + No Warm-up
  • 对照组2: Post-LN + 标准Xavier初始化 + Warm-up
  • 对照组3: Post-LN + Xavier缩放初始化 + Warm-up

Xavier的缩放初始化是对于第 ll 层，l∈[1,N]l\in[1,N], 有随着深度递减的缩放因子kl=N+1−lk_l = N+1-l, 使得

Wol∼N(0,1kl2d′)W_o^l\sim \mathcal{N}(0,\frac{1}{k_l^2d'})

scaling_init

其采用的评估指标为整个模型各轮训练的输出更新的模, θi\theta_i 表示第ii轮训练后的参数, θ0\theta_0 表示初始化参数

∥ΔFi∥=∥F(x,θi)−F(x,θ0)∥\|\Delta F_i\| = \|F(x,\theta_i)-F(x,\theta_0)\|

在图4.a中，对照组1 的Model Update在初始步就出现了爆炸。这一结论和 Transformer 中的 Layer Normalization与梯度稳定性 中的结论是一致的。

通过图3的子图对比，作者将Model Update 而非梯度模作为评估梯度下降稳定性的评估指标，因为在对照组2、3的对比中，对照组3的梯度模更大，但是具有更好的收敛效果。

DeepNet 的结构

本文的主体创新点，即DeepNorm 架构。相对于Pre-LN,就是在主输入信号处增加了一个可学习参数 α>1\alpha>1, 对输入tensor进行了放大。

deepnet_architecture

Lemma 1.: 对于 X=(x1,⋯ ,xn)∈Rn×dX=(\mathbf{x}_1,\cdots,\mathbf{x}_n)\in \mathbb{R}^{n\times d}, Var(xi)=1,E(xi)=0\mathrm{Var}(\mathbf{x}_i) = 1, \mathbf{E}(\mathbf{x}_i) = 0, qi∈[0,1]q_i\in [0,1], 则有

Softmax(q1,⋯ ,qn)⋅X≍xi\mathrm{Softmax}(q_1,\cdots, q_n)\cdot X \asymp \mathbf{x}_i

这里的 ≍\asymp 表示 equal bound of magnitude，即只比较数量级上的上下界，而不要求两个向量在方向或逐元素取值上相同。
Proof:

Softmax(q1,⋯ ,qn)⋅X=(exp⁡qi∑exp⁡qj)i⋅(x1,⋯ ,xn)T=∑xiexp⁡qi∑exp⁡qj\begin{aligned} \mathrm{Softmax}(q_1,\cdots, q_n) \cdot X &= \left(\frac{\exp{q_i}}{\sum \exp{q_j}}\right)_{i}\cdot \left(\mathbf{x}_1,\cdots ,\mathbf{x}_n \right)^T\\ &=\sum \frac{\mathbf{x}_i\exp q_i}{\sum\exp q_j} \end{aligned}

由于

n≤∑exp⁡qj≤nen\leq\sum \exp q_j\leq ne

有

1∑exp⁡qj=Θ(1n)\frac{1}{\sum \exp q_j} = \Theta\left(\frac{1}{n}\right)

因此可以写成凸集的形式

Softmax(q1,⋯ ,qn)X=∑i=1nwixi,wi=exp⁡qi∑exp⁡qj=Θ(1n).\mathrm{Softmax}(q_1,\cdots,q_n)X = \sum_{i=1}^n w_i\mathbf{x}_i, \quad w_i=\frac{\exp{q_i}}{\sum \exp q_j}=\Theta\left(\frac1n\right).

因此，由 wi=Θ(1n)w_i=\Theta\left(\frac1n\right) 可得

∑i=1nwixi≍∑i=1n1nxi.\sum_{i=1}^n w_i\mathbf{x}_i \asymp \sum_{i=1}^n \frac1n\mathbf{x}_i .

又因为每个 xi\mathbf{x}_i 都满足 E(xi)=0\mathbf{E}(\mathbf{x}_i)=0 且 Var(xi)=1\mathrm{Var}(\mathbf{x}_i)=1，所以所有 xi\mathbf{x}_i 具有相同的 magnitude bound。于是 nn 个同阶向量经过 1n\frac1n 量级的权重加权求和后，其输出仍然具有与单个 xi\mathbf{x}_i 相同的 magnitude bound，即

∑i=1n1nxi≍xi.\sum_{i=1}^n \frac1n\mathbf{x}_i \asymp \mathbf{x}_i .

因此

Softmax(q1,⋯ ,qn)X=∑i=1nwixi≍xi.\mathrm{Softmax}(q_1,\cdots,q_n)X = \sum_{i=1}^n w_i\mathbf{x}_i \asymp \mathbf{x}_i .

Lemma 1对Attention输出头的阶进行了估计，保证了输入和输出的阶是相等的，Attention不会导致Model Update发生爆炸

Lemma 2 方差的阶的可加性

Var(X+Y)≍Var(X)+Var(Y)\mathrm{Var}(X+Y)\asymp \mathrm{Var}(X) + \mathrm{Var}(Y)

Proof:
展开得

Var(X+Y)=E[(X+Y)2]−[E(X+Y)]2=E(X2)−E2(X)+E(Y2)−E2(Y)+2E(XY)−E(X)E(Y)=Var(X)+Var(Y)+2cov(X,Y)\begin{aligned} \mathrm{Var}(X+Y) &= \mathbb{E}[(X+Y)^2] - [\mathbb{E}(X+Y)]^2\\ & = \mathbb{E}(X^2) - \mathbb{E}^2(X) +\mathbb{E}(Y^2) - \mathbb{E}^2(Y) +2\mathbb{E}(XY) - \mathbb{E}(X)\mathbb{E}(Y)\\ & = \mathrm{Var}(X) +\mathrm{Var}(Y)+2\mathrm{cov}(X,Y) \end{aligned}

Theorem 1 对于NN层DeepNet结构 F(x,θ)F(x,\theta) , 每一层的DeepNet可表示为 xl+1=f(xl,θl)=LN(αxl+Gθl(xl))x_{l+1} = f(x_l,\theta_l) = \mathrm{LN}(\alpha x_l+G_{\theta_l}(x_l)), θi∈{θ1,⋯ ,θ2N}\theta_i\in \left\{\theta_1,\cdots,\theta_{2N}\right\}

其中奇数参数是Selt-Attention的参数权重，偶数参数是FFN的参数权重。则有

∥ΔF∥≤∑i=1nvi2+wi2α∥θi∗−θi∥\|\Delta F\|\leq \sum_{i=1}^n \frac{\sqrt{v_i^2+w_i^2}}{\alpha}\|\theta^*_i-\theta_i\|

基于最简化的情况，Self-Attention Layer 与 FFN layer 都视为 Gθ(xl)≍vlwlxlG_\theta(x_l) \asymp v_lw_lx_l, 则每一个Layer的递推输出为

xl+1=f(xl,θl)=αxl+Gθl(xl)Var(αxl+Gθl(xl))≍(α+vlwl)xl(α2+vl2wl2)Var(xl)\begin{aligned} x_{l+1} = f(x_l,\theta_l) &= \frac{\alpha x_l+G_{\theta_l}(x_l)}{\sqrt{\mathrm{Var}(\alpha x_l+G_{\theta_l}(x_l))}}\\ & \asymp \frac{(\alpha+v_lw_l)x_l}{\sqrt{(\alpha^2+v_l^2w_l^2)\mathrm{Var}(x_l)}}\\ \end{aligned}

根据Lemma 1, Var(xl)≍Var(x0)=1\mathrm{Var}(x_l) \asymp \mathrm{Var}(x_0) = 1, 故

xl+1=f(xl,θl)≍α+vlwlα2+vl2wl2xlx_{l+1} = f(x_l,\theta_l) \asymp \frac{\alpha+v_lw_l}{\sqrt{\alpha^2+v_l^2w_l^2}}x_l

因此

∂fl∂xl≍α+vlwlα2+vl2wl2\frac{\partial f_l}{\partial x_l} \asymp \frac{\alpha+v_lw_l}{\sqrt{\alpha^2+v_l^2w_l^2}}

视 θl=(vl,wl)\theta_l = (v_l,w_l)

∂fl∂θl≍αxl(α−vlwl)(α2+vl2wl2)32(wl,vl)\frac{\partial f_l}{\partial \theta_l} \asymp \frac{\alpha x_l(\alpha -v_lw_l)}{(\alpha^2+v_l^2w_l^2)^{\frac{3}{2}}} \left(w_l,v_l\right)

因此可展开Layer Model Update

∥ΔF∥=∥f(x2N∗,θ2N∗)−f(x2N,θ2N)∥=∥∂f∂xl(x2N,θ2N)(x2N∗−x2N)+∂f∂θl(x2N,θ2N)(θ2N∗−θ2N)∥≤∥∂f∂xl(x2N,θ2N)∥∥x2N∗−x2N∥+∥∂f∂θl(x2N,θ2N)∥∥θ2N∗−θ2N∥≍α+v2Nw2Nα2+v2N2w2N2∥x2N∗−x2N∥+αx2N(α−v2Nw2N)(α2+v2N2w2N2)32∥(w2N,v2N)∥∥θ2N∗−θ2N∥∣≤∥x2N∗−x2N∥+αx2N(α−v2Nw2N)(α2+v2N2w2N2)32w2N2+v2N2∥θ2N∗−θ2N∥≍∥x2N∗−x2N∥+w2N2+v2N2α∥θ2N∗−θ2N∥ \begin{aligned} \|\Delta F\| &= \|f(x_{2N}^*,\theta_{2N}^*)-f(x_{2N},\theta_{2N})\|\\ & = \|\frac{\partial f}{\partial x_l}(x_{2N},\theta_{2N})(x^*_{2N}-x_{2N}) +\frac{\partial f}{\partial \theta_l}(x_{2N},\theta_{2N})(\theta_{2N}^*-\theta_{2N})\|\\ &\leq \|\frac{\partial f}{\partial x_l}(x_{2N},\theta_{2N}) \| \|x^*_{2N}-x_{2N}\| +\|\frac{\partial f}{\partial \theta_l}(x_{2N},\theta_{2N})\|\|\theta_{2N}^*-\theta_{2N}\|\\ &\asymp \frac{\alpha+v_{2N}w_{2N}}{\sqrt{\alpha^2+v_{2N}^2w_{2N}^2}}\|x^*_{2N}-x_{2N}\|+ \frac{\alpha x_{2N}(\alpha -v_{2N}w_{2N})}{(\alpha^2+v_{2N}^2w_{2N}^2)^{\frac{3}{2}}} \|\left(w_{2N},v_{2N}\right)\|\|\theta_{2N}^*-\theta_{2N}\||\\ &\leq \|x^*_{2N}-x_{2N}\|+\frac{\alpha x_{2N}(\alpha -v_{2N}w_{2N})}{(\alpha^2+v_{2N}^2w_{2N}^2)^{\frac{3}{2}}} \sqrt{w_{2N}^2+v_{2N}^2}\|\theta_{2N}^*-\theta_{2N}\|\\ &\asymp \|x^*_{2N}-x_{2N}\|+\frac{\sqrt{w_{2N}^2+v_{2N}^2}}{\alpha}\|\theta_{2N}^*-\theta_{2N}\|\ \end{aligned}

递推得

∥ΔF∥≤∑i=12Nwi2+vi2α∥θi∗−θi∥\|\Delta F\|\leq\sum_{i=1}^{2N}\frac{\sqrt{w_{i}^2+v_{i}^2}}{\alpha}\|\theta_{i}^*-\theta_{i}\|

由于每一个wi,viw_i,v_i 的初始化都与scaling factor β\beta 相关， 因此整个Model Update 的 Bound 能通过 α,β\alpha, \beta 控制，我们能通过设计这两个值的大小实现稳定训练Deeper的Network且保证收敛