Ultra-Low Precision 4-bit Training of Deep Neural Networks

目录

• 概
• 主要内容
  • Radix-4 FP4 format
  • GradScale
  • Two-Phase Rounding (TPR)

Sun X., Wang N., Chen C., Ni J., Agrawal A., Cui X., Venkataramani S. and Maghraoui K. E. and Srinivasan V. Ultra-low precision 4-bit training of deep neural networks. NeurIPS, 2020.

概

本文提出了一种特殊的 FP4 格式以及 Grad Scale 机制来实现 4-bit 量化.

主要内容

Radix-4 FP4 format

• 4-bit 的训练是很难的 (实际上 8-bit 的目前也才刚刚有了成色), 作者认为一大难点是这种情况下所能表示的数值范围过于有限了. 所以, 作者不采用常用的 [sign, exponent, mantissa] = [1, 2, 1] 的格式, 而是使用 [sign, exponent, mantissa] = [1, 3, 0] 的格式. radix-4 FP4 可以表示 (如果我没理解错):

  $$4^{-3}, 4^{-2}, \cdots, 4^2, 4^3, \\ -4^{-3}, -4^{-2}, \cdots, -4^2, -4^3.$$

• 然后 rounding 规则为:

  $$round(x) = \left \{ \begin{array}{ll} 4^{n-1} & x \le 4^n / 1.6, \\ 4^{n} & x > 4^n / 1.6. \\ \end{array} \right.$$

  注意到, $(4^n + 4^{n-1}) / 2 = 4^n / 1.6$.

GradScale

• FP16 的训练, 通常需要先将 loss scale 到最大表示范围处, 以尽可能用尽所有的资源, 然后梯度再 1 / scale 回去.

• FP4 仅通过一次 scale 是远远不够的. 所以作者给每一层都设置了可学习的 scale 参数 $scale_{i, L}$. 它的更新方式如下:

• 即如果梯度的最大值已经落在 $(Max / 2, Max]$ 这个合理区间内, 则不需要调整. 如果梯度的最大值落在 $[0, Max / 2]$ 这个 underflow 区间内, 则需要增大 $scale_{i, L}$ (更新方式是 $\times 2$). 如果梯度的最大值落在 $(Max, +\infty)$, 此时 overflow, 需要减小 $scale_{i, L}$ (更新方式是 $/ 2$).

Two-Phase Rounding (TPR)

• 作者还提出了两种 rounding 方式, 分别用于梯度回传中 $dL / dx, dL/dW$. 说实话, 我没看明白为什么这种做法能够保留更多信息?