Scalable Methods for 8-bit Training of Neural Networks

目录

• 概
• Range Batch Normalization
• 代码

Banner R., Hubara I., Hoffer E. and Soudry D. Scalable methods for 8-bit training of neural networks. NeurIPS, 2018.

概

本文针对 Batch Norm 模块在低精度 (8-bit) 的情况下进行一个合适的改进.

Range Batch Normalization

• 对于一个 $n \times d$ 的输入 $x = (x^{(1)}, x^{(2)}, \ldots, x^{(d)})$, Batch Norm 模块 normalize 每个维度:

  $$\tag{1} \hat{x}^d = \frac{ x^{(d)} - \mu^d }{ \sqrt{ \text{Var}[x^{(d)}] } },$$

  其中 $\mu^d$ 是 $x^{(d)}$ 上的平均, $\text{Var}[x^{(d)}] = \frac{1}{n} \| x^{(d)} - \mu^d \|_2^2$.

• 容易发现, $\sqrt{\text{Var}[x^{(d)}]}$ 这一项涉及平方和, 在低精度的情况下很容易导致数值不稳定. 因此, 本文希望提出一个 Range BN 模块来实现低精度下的一个鲁棒模拟:

  $$\tag{2} \hat{x}^d = \frac{ x^{(d)} - \mu^d }{ C(n) \cdot \text{range}(x^{(d)} - \mu^d) },$$

  其中 $\text{range}(x) = \max(x) - \min(x)$, $C(n) = 1 / \sqrt{2 \ln (n)}$.

• 我们来解释为什么这么做:

  1. 首先, 如果 $x$ 服从高斯分布, 我们有
     $$0.23 \sigma \cdot \sqrt{\ln(n)} \le \mathbb{E}[\max(x^{(d)} - \mu^d)] \le \sqrt{2} \sigma \cdot \sqrt{\ln(n)}.$$

  2. 根据, 如果 $x$ 关于 $0$ 对称, 我们有 $\mathbb{E}[\max(x)] = -\mathbb{E}[\min (x)]$, 我们有
     $$0.23 \sigma \cdot \sqrt{\ln(n)} \le -\mathbb{E}[\min(x^{(d)} - \mu^d)] \le \sqrt{2} \sigma \cdot \sqrt{\ln(n)}.$$

  3. 上面二式相加可得:
     $$0.325\sigma \le C(n) \cdot \text{range}(x^{(d)}- \mu^d) \le 2 \sigma.$$

• 因此, (2) 可以作为 (1) 的有效的稳定的模拟.

代码

[official-code]