RMSNorm

计算均方根

$$RMS(x)=\sqrt{\frac{1}{n}\sum _{i=1}^n{x}_i^2+ϵ}$$

• $ϵ$ 一般取 1e-6，用于防止计算结果为 0

归一化

$$y_i=\frac{x_i}{RMS(x)}\cdot {\gamma }_i$$

• ${\gamma }_i$ 为可学习的参数

Backward

$$\frac{\partial L}{\partial x_i}=\frac{1}{RMS}\left(\frac{\partial L}{\partial y_i}\cdot {\gamma }_i-\frac{{\hat{x}}_i}{n}\sum _{j=1}^n\frac{\partial L}{\partial y_j}{y}_j\right)$$ $$dX=\frac{1}{RMS}(dZ-\hat{X}\cdot mean(dZ\odot \hat{X}))$$ $$d\gamma =\sum _{\text{all token}}\hat{x}\cdot dy$$

其中

$$\hat{X}=\frac{X}{RMS}$$ $$dz=dy\odot \gamma$$