Root Mean Square Layer Normalization

目录

• 概
• RMSNorm
• 代码

Zhang B. and Sennrich R. Root mean square layer normalization. NIPS, 2019.

概

RMSNorm 节省时间.

RMSNorm

• 假设输入为 \(\mathbf{x} \in \mathbb{R}^m\), 然后

  \[\mathbf{a} = \mathbf{W} \mathbf{x} \in \mathbb{R}^{n}, \\ \mathbf{y} = f(\text{Norm}(\mathbf{a}) + \mathbf{b}) \in \mathbb{R}^{n}. \]

  其中 \(f(\cdot)\) 是 element-wise 的激活函数.

• LayerNorm 采取的是如下的方式 (注意, 下面的 \(/\) 是 element-wise 的):

  \[\text{LayerNorm}(\mathbf{a}) = \frac{\mathbf{a} - \bm{\mu}}{\bm{\sigma}} \odot \mathbf{g}, \]

  其中

  \[\bm{\mu} = \text{mean}(\mathbf{a}), \\ \bm{\sigma} = \sqrt{\text{mean}((\mathbf{a} - \bm{\mu})^2)}. \]

• RMSNorm 采用的是如下的方式:

  \[\text{RMSNorm}(\mathbf{a}) = \frac{\mathbf{a}}{\text{RMS}(\mathbf{a})} \odot \mathbf{g}, \]

  其中

  \[\text{RMS}(\mathbf{a}) = \sqrt{\text{mean}(\mathbf{a}^2)}. \]

• 由于不用计算均值, RMSNorm 所需的计算时间会少一点, 但是效果是差不多的:
  

• 此外, RMSNorm 保留了一些重要的不变性:

代码

[official]