Training Transformers with 4-bit Integers

目录

• 概
• 符号说明
• 4-bit FQT
  • Learned Step Size Quantization
  • Hadamard Quantization
  • Bit Splitting and Leverage Score Sampling
• 代码

Xi H., Li C., Chen J. and Zhu J. Training transformers with 4-bit integers. NeurIPS, 2023.

概

本文针对 4-bit 中训练中一些特点 (针对 transformers) 提出了一系列解决方法.

符号说明

• 本文主要考虑 transformer 的 4-bit 量化训练, 其中主要涉及如下的矩阵乘法操作:
  $$\mathbf{Z} = \mathbf{XW}^T, \quad \mathbf{Z} \in \mathbb{R}^{N \times C}, \mathbf{X} \in \mathbb{R}^{N \times D}, \mathbf{W} \in \mathbb{R}^{C \times D}.$$

4-bit FQT

Learned Step Size Quantization

• 首先, 作者采用 learned step size quantizer (LSQ):

  $$\text{int}_{s_X} (\mathbf{X}) := \lfloor \text{clamp}(\mathbf{X} / s_X, -Q_N, Q_P) \rceil,$$

  其中 $s_X$ 是可学习的参数, 旨在保证将输入的范围放缩到 $[-Q_N, Q_P]$. 量化后的值属于: $\{-Q_N, -Q_N + 1, \ldots, Q_P\}$, 对于 4-bit 量化, $Q_N = Q_P = 7$.

• 反量化过程为:

  $$\text{float}(\text{int}_{s_X} (\mathbf{X})) = s_X \text{int}_{s_X}(\mathbf{X}) \approx \mathbf{X}.$$

• 给定 LSQ, 量化后的矩阵乘法可以用如下方式近似:

  $$\mathbf{Y} = \mathbf{XW}^T \approx s_X s_W \text{int}_{s_X} (\mathbf{X}) \text{int}_{s_W} (\mathbf{W})^T.$$

Hadamard Quantization

• 但是, 单纯使用 LSQ 会产生很大的误差, 主要原因是, 数据的分布如上图 (a) 所示极不均匀, 所以提出用 Hadamard matrix 来进行一个 smooth.

• Hadamard matrix 是一个 $2^k \times 2^k$ 大小的正交矩阵:

  $$\mathbf{H}_0 = [1], \quad \mathbf{H}_k = \frac{1}{\sqrt{2}} \left [ \begin{array}{cc} \mathbf{H}_{k-1} & \mathbf{H}_{k-1} \\ \mathbf{H}_{k-1} & -\mathbf{H}_{k-1} \end{array} \right ].$$

  容易证明 $\mathbf{H}_k^T = \mathbf{H}_k, \mathbf{H}_k \mathbf{H}_k = \mathbf{I}$. 如 (b) 所示, 明显分布更加均匀了.

• 而且, 这种方式下的量化矩阵乘法的计算也是方便的:

  $$\mathbf{Y} = \mathbf{XW}^T \approx s_X s_W \text{int}_{s_X} (\mathbf{XH}) \underbrace{\mathbf{HH}}_{=\mathbf{I}} \text{int}_{s_W} (\mathbf{HW}^T).$$

Bit Splitting and Leverage Score Sampling

• 在 Backpropagation 过程中, 作者发现, transformers (但是卷积网络不行) 的梯度呈现出明显的结构稀疏性, 经常一整行为 0. 作者选择抛弃对很小的梯度的量化, 用省下的空间去建模那些较大的值. 这些操作通过 bit splitting 和 leverage score sampling 实现. 感兴趣的请回看原文.

代码

[official-code]