Normalization methods

LayerNorm#

對 sample 裡面的所有 dim 進行 norm, 讓每個 embedding 的分佈接近：

\frac{x - \mathrm{E}[x]}{\sqrt{\mathrm{Var}[x] + \epsilon}} * \gamma + \beta

其中 $\epsilon$ 為固定常數(1e-5), $\gamma$ 以及 $\beta$ 兩個可學習參數避免 overfitting

Pre-norm, Post-norm#

在 FFN 之後的 Add&Norm Layer, 若先進行 add 才做 norm 則為 post-norm, 反之為 pre-norm

早期的 Transformer 多使用 post-norm, 近代的研究則比較多使用 pre-norm, pre-norm 的優勢在梯度傳播時較為穩定（因為在做完 norm 之後才做殘差連接), post-norm 則是模型的效果會更好

且 pre-norm 的模型在 fine-tune 時表現相較 post-norm 會更好

RMSNorm#

給定一個輸入向量 $a = (a_1, a_2, \dots, a_n)$ , 先算RMS：

$\text{RMS}(a) = \sqrt{\frac{1}{n}\sum_{i=1}^{n} a_i^2}$

然後除以這個 RMS, 再乘上 $g_i$ ：

$\bar a_i = \frac{a_i}{\text{RMS}(a)} \cdot g_i$

pRMSNorm（partial RMSNorm)#

只用前面 $p\%$ 的維度來估計 RMS (e.g. $p=6.25\%$ ), 進一步省計算, 效果仍相近

實作上還會在分母加一個很小的 $\epsilon$ 防止除零：

$\bar a_i = \frac{a_i}{\sqrt{\frac1n\sum_k a_k^2 + \epsilon}} \cdot g_i$

與 LayerNorm 的比較#

RMSNorm（ $\bar a_i = a_i / r$ , $r=\sqrt{\frac1n\sum_k a_k^2}$ ）的 Jacobian 是：

$\frac{\partial \bar a_i}{\partial a_j} = \frac{1}{r}\left(\delta_{ij} - \frac{1}{n}\bar a_i \bar a_j\right)$

LayerNorm (多減一個 $\mu$ ) 的 Jacobian 是：

$\frac{\partial \bar a_i}{\partial a_j} = \frac{1}{\sigma}\left(\delta_{ij} - \frac{1}{n} - \frac{1}{n}\bar a_i \bar a_j\right)$

只差一個 $-\frac{1}{n}$ 項 ( $\partial\mu/\partial a_j = 1/n$ )

把那個 $-\frac{1}{n}$ 項代進反向傳播 $\frac{\partial L}{\partial a_j}=\sum_i \frac{\partial \bar a_i}{\partial a_j}\frac{\partial L}{\partial \bar a_i}$ ，它對每個 $j$ 的貢獻是：

$-\frac{1}{n\sigma}\sum_i \frac{\partial L}{\partial \bar a_i}$

但其實這項是在做 re-centering, 沒辦法控制梯度的量級, 只能在常數範圍上縮小, RMSNorm 只留 re-scaling, 避免計算均值的額外成本