A White Paper on Neural Network Quantization

URL

• https://arxiv.org/pdf/2106.08295.pdf

TL;DR

• 本文用比较简洁的方式给出了神经网络的通用量化方法，是量化领域的必读论文。

Algorithm

1. 量化基础知识

1.1 硬件背景

• 一个 $y=Wx+b$ 实际上是由 乘法器 和 累加器 组合而成的，实际的计算过程如下：
  

卷积实际上也是通过 image to column 操作变成 $y=Wx+b$ 操作

• 常见的 int8 量化会将上述过程变成如下过程：
  

weight 和 input 都被量化为 int8 ，同时保留各自的量化 scale，乘法操作是整形乘法器（更快），累加器是 int32 类型，最后再量化为 int8 放到 OCM 上

1.2 均匀仿射量化

• 均匀仿射量化也被称为 非对称量化，由三个量化参数定义：
  • 比例因子 scale
  • 零点 zero_point
  • 比特宽度 bits
• 非对称量化：
  • for unsigned integers: $X_{int} = clamp(\lfloor\frac{X}{s}\rceil+z;0,2^b-1)$
  • for signed integers: $X_{int} = clamp(\lfloor\frac{X}{s}\rceil+z;-2^{b-1},2^{b-1}-1)$
  • 这里的 $\lfloor\rceil$ 表示 round 运算
• 对称量化是非对称量化的简化版本，是将零点 zero_point 固定为 0
• 对称量化：
  • for unsigned integers: $X_{int} = clamp(\lfloor\frac{X}{s}\rceil;0,2^b-1)$
  • for signed integers: $X_{int} = clamp(\lfloor\frac{X}{s}\rceil;-2^{b-1},2^{b-1}-1)$
• 对称量化和非对称量化的含义：
  
  
• 2 的指数幂量化：
  • 限制 $s=2^{-k}$
  • 优势：scale 过程变成了硬件移位，对硬件更友好。
  • 劣势：会使得 round 和 clip 误差的权衡变难。
• 量化颗粒度：
  • per-tensor: 硬件更友好，但限制了量化的自由度。
  • per-channel: 反之。

1.3 量化模拟

• 量化模拟是指在浮点计算设备上模拟定点计算设备的过程，通常用于训练。
  

左边是定点计算过程，右边是用浮点设备模型定点计算的过程

• 为了减少数据搬运和不必要的量化步骤，通常会做：
  • batch norm 折叠：batch norm 在推理时是静态的，因此可以和前面的 conv 等层合并。
  • 激活函数融合：在实际的硬件解决方案中，通常会在非线性操作（如 ReLU）之后直接进行量化，而不是先将激活写入内存然后再加载回计算核心。

1.4 实践考量

• 对称量化和非对称量化：
  • 对称量化：zero-point == 0
  • 非对称量化：zero-point != 0
• 为了方便计算，通常情况下，会将权重设置为对称量化（$z_w=0$），将特征设置为非对称量化（$z_x\ne 0$）
  • 原因分析：
    • $W=S_w(W_{int} - Z_w)$
    • $X=S_x(X_{int} - Z_x)$
    • $WX=S_wS_x(W_{int} - Z_w)(X_{int} - Z_x)\\=S_wS_xW_{int}X_{int}-S_wS_xZ_wX_{int}-S_wS_xZ_xW_{int}+S_wS_xZ_wZ_x$
    • 在推理阶段：$S_w,\ S_x,\ Z_w,\ Z_x,\ W_{int}$ 已知，因此：
      • 等式的第三项和第四项可提前算出，无需推理耗时。
      • 第一项和第二项由于关联动态输入 $X_{int}$，因此需要额外耗时；但是如果设置 $Z_w=0$，则第二项恒等于0，可节省计算量。

2. 训练后量化（PTQ，post-training quantization）

• 训练后量化是指用 float32 精度训练的模型直接转成量化模型，无需任何数据和训练。

2.1 量化范围的设置

• 最大最小值法（min-max）：$q_{min}=minV,\ \ q_{max}=maxV$，$V$ 是待量化 tensor
• 均方差法（MSE）：$\argmin_{q_{min},q_{max}}||V-\hat{V}(q_{min}, q_{max})||^2_F$
• 交叉熵法（cross entropy）：$\argmin_{q_{min},q_{max}}=H(softmax(V),softmax(\hat{V}(q_{min},q_{max})))$，其中 $H$ 表示 cross entropy function
• 批量归一化法（BN based）：$q_{min}=min(\beta-\alpha\gamma),\ \ q_{max}=max(\beta+\alpha\gamma)$，其中 $\beta,\ \gamma$ 分布表示 batch norm 学到的 per channel 的 shift 和 scale，$\alpha>0$ 是超参数
• 组合法（comparsion）：以上方法的自由组合
  
  

使用不同量化方法分别量化 weight 和 activation 后的精度

2.2 跨层均衡（Cross-Layer Equalization）

• 这是一种 通过修改模型权重 来改善神经网络量化性能的技术，CLE 的目的是减少网络中不同 channel 之间由于量化引起的性能不平衡，这种问题在 depth-wise conv layer 中尤其容易出现。
  

mobilenetv2 第一个 depth-wise conv 层的 per output channel weight range

• 想要实现跨层均衡的模型，需要激活函数满足交换律，即：$f(sx)=sf(x)$，常见的 ReLU 和 PReLU 都满足。
  
• CLE 原理：
  • $y=f(W_2(W_1x+b_1)+b_2)\\=f(SW_2(S^{-1}W_1x+S^{-1}b_1)+b_2)\\=f(\hat{W_2}(\hat{W_1}x+\hat{b_1})+b_2)$
  • 其中：
    • $\hat{W_1}=S^{-1}W_1$
    • $\hat{b_1}=S^{-1}b_1$
    • $\hat{W_2}=SW_2$
    • $S_i=\frac{\sqrt {r_i^1r_i^2}}{r_i^2}$，其中 $r_i^j$ 表示 j tensor 的 i channel
• abosrbing high bias 是一种 解决模型中过大 bias 的技术，原理是：
  • $y=W_2h+b_2\\=W_2(f(W_1x+b_1))+b_2\\=W_2(f(W_1x+b_1)+c-c)+b_2\\=W_2(f(W_1x+\hat{b_1})+c)+b_2\\=W_2(f(W_1x+\hat{b_1}))+\hat{b_2}\\=W_2\hat{h}+\hat{b_2}$
  • 其中：
    • $\hat{b_2}=b_2+W_2c$
    • $\hat{h}=h-c$
    • $\hat{b_1}=b_1-c$
    • $c_i=max(0, min_x(W_{1i}x+b_{1i}))$

2.3 偏移校正（Bias correction）

• 偏移校正：通过修改 bias，使得对于相同输入作用在量化前后的权重的输出期望相同
• 用公式解释：$\mathbb E[Wx]\ne \mathbb E[\hat Wx]$
• 这种校正方式可以在不借助 calibration 数据集的情况下，单靠计算得到 $\Delta Bias$，比如通过 BN 层的 runing mean 得到 $\mathbb E[x]$，$\Delta bias=\mathbb E[\Delta W]\cdot \mathbb E[x]$

2.4 自适应舍入（AdaRound）

• AdaRound 目的是为了解决量化中极为关键的痛点：“四舍五入（Round-to-Nearest, RTN）并不一定是取整最优解”
• AdaRound 的核心思想：
  • 将 “取整” 这个离散的、不可导的动作，转成一个针对特定数据分布的、连续可导的优化问题。
  • 不再进行简单的四舍五入，而是对每一个权重引入一个可学习的决策变量，让模型结合少量的校准集，自己决定每个权重是应该“向下取整（Floor）”还是“向上取整（Ceil）”。
• 公式表示：

  $$w_q=\lfloor\frac{w}{s}\rfloor + h(v)$$

  • $\lfloor\cdot\rfloor$ 是向下取整
  • $V$ 是与权重维度完全相同的连续可学习参数
  • $h(V)$ 是一个连续可导的函数 $h(V) = \text{clip}(\sigma(V) (\zeta - \gamma) + \gamma, 0, 1)$
    • $\sigma(V)$：标准的 Sigmoid 函数
    • $\gamma$ 和 $\zeta$：控制拉伸范围的超参数，$\gamma = -0.1$ 和 $\zeta = 1.1$
    • $\text{clip}(x, 0, 1)$：截断函数，将最终结果强行限制在闭区间 $[0, 1]$ 内
• 损失函数设计

  $$\mathcal{L} = \sum || W x - \tilde{W}_q(V) x ||_F^2 + \lambda f_{reg}(V)$$

  • $W x$ 是原始浮点层的输出
  • $\tilde{W}_q(V) x$ 是带有软取整参数层的输出
  • $f_{reg}(V)$ 是一个正则化项（退火惩罚项）。随着优化的进行，这个正则化项会越来越强，逼迫 $h(V)$ 的值向 0 或 1 两端极化

    $$f_{reg}(V) = \sum_{i} \left( 1 - |2h(V_i) - 1|^\beta \right)$$

  • 在训练过程中，$\beta$ 会逐渐降低，从 20 逐渐降低到 2，$\beta$ 越大，曲线顶端越平
    

2.5 标准的 PTQ 流程

• 跨层均衡
• 添加量化器
• 权重范围设定
• 自适应舍入
• 偏移校正
• 激活范围设定

2.6 PTQ debug 流程

• FP32 健全性检查
• 权重或激活量化
• 修复权重量化
• 修复激活量化
• 逐层分析
• 可视化层数据分布
• 修复各个量化器

3. 训练感知量化（QAT，Quantization-aware training）

3.1 模拟量化反向传播

• 红色线表示直通估计器（STE, straight-through estimator），说人话就是 “由于四舍五入运算不可微，那就让梯度跳过这一层继续向前”

$$\frac{\partial\hat x_i}{\partial x_i}=\begin{cases} 1 & \text{if } q_{\min} \leq \mathbf{x}_{i} \leq q_{\max}, \\ 0 & \text{otherwise.} \end{cases}$$

• 被饱和的部分没有梯度（包括 weight 和 activation）

3.2 BN 折叠

• QAT 在训练过程中，如果 conv 可以做 per-channel 量化，那 conv 的 activation 在送入 BN 之前可以不需要做量化，BN 之后做 activation 的量化
• 部署时，BN 的 per-channel 缩放系数可以和 conv 的 per-channel 缩放系数合并

3.3 QAT 初始化

• 更好的初始化可以带来更好的 QAT 结果，但增益通常很小，并且随着训练持续时间的延长而消失。

3.4 标准 QAT 流程

• 跨层均衡
• 添加量化层
• 范围估计
• 量化可学习参数

Thoughts

• 非常全面的讲解了 PTQ 和 QAT 的底层原理，在 CNN 网络上非常适合
• 对于模型架构的持续迭代，以及的 GPU 对低比特浮点数的支持（包括分组量化浮点数类型），给模型量化带来了更多新的课题