（QAT）DOREFA-NET：低比特在线量化训练方法

Paper name

DOREFA-NET: TRAINING LOW BITWIDTH CONVOLUTIONAL NEURAL NETWORKS WITH LOW BITWIDTH GRADIENTS

Paper Reading Note

URL: https://arxiv.org/pdf/1606.06160.pdf

TL;DR

• 提供了低比特权重、特征、梯度的在线量化训练方案
• 量化位宽的重要性排序为：梯度>特征>权重，所以实验中的配比建议为 1w4f6g（1bit权重，4bit特征，6bit梯度），配比的关系对应于 DoReFa 的论文名
• 首个成功在低于 8bit 梯度下得到较高精度的在线量化训练方案
• SVHN 和 ImageNet 数据集上基于 AlexNet 的实验结果证明精度可以与浮点模型相比

---

Introduction

• 之前的量化方法对于梯度的量化位宽没有低于 8bit 的（最少是 10bit，BNN 和 XNOR-Net 都是全精度梯度训练，导致大部分训练时间消耗在反传环节），这里提出一种能在低于 8bit 位宽梯度下的训练方案
• 这种梯度低比特量化的方案后面可能能支持在 FPGA、ASIC 上的低功耗训练
• 量化位宽的重要性排序为：梯度>特征>权重，配比的关系对应于 DoReFa 的论文名（1-bit weights, 1-bit activations and 2-bit gradients can lead to 93% accuracy on SVHN dataset）

---

Dataset/Algorithm/Model/Experiment Detail

实现方式

1. STRAIGHT-THROUGH ESTIMATOR（STE，直通滤波器）

• STE是对不可微分算子在计算反向传播梯度时的近似方法，比如量化操作不可微（量化函数一般是阶梯函数，处处的梯度为0），STE就假装量化没有发生，将梯度直接前传
• 比如对于伯努利采样函数， c是目标函数，伯努利分布不可微， $\frac{\partial (c)}{\partial (p)}$ 没有很好的定义，因此不能使用链式规则直接从正向过程构造反向过程，因为 $p$ 和 $q$ 是期望相等的，所以使用能够被正常定义的 $\frac{\partial (c)}{\partial (q)}$ 来近似 $\frac{\partial (c)}{\partial (p)}$
  
• 又比如对于如下量化函数，将 [0, 1] 之间的全精度数量化到 [0,1] 之间的 k-bit 量化数，注意前传的量化函数不唯一，可以是不同的量化形式
  
• STE 的问题：

1. loss curve会波动
2. 由于量化函数实际上和y=x这条直线并不一致，而我们用后者的梯度来代替前者，这就会导致梯度的不匹配

2. 权重量化

• 对于 1bit 权重情况下量化方法如下，其中的scaling factor： $E(|r_i|)$ 代表该层的权重绝对值的整体期望 ，这个 scaling factor 的作用是增加权重的取值范围；与 XNOR-Net 量化方法的区别是 XNOR-Net 中的 $E(|r_i|)$ 代表的是该层权重的每个 channel 的绝对值的期望，但是这里使用逐 channel 的 scaling factor 会导致计算在反向传播过程中，当计算梯度和权重之间的卷积时无法利用位卷积核
  

• 对于k-bit（k>1）的权重量化
  

  • 其中 $quantiz{e}_k$ 为量化函数，将 [0, 1] 之间的全精度浮点数聚类到 [0, 1] 区间的定点数，即将 [0, 1] 的浮点数聚类到 [0, 1] 中的 $2^k$ 个等距分割点上
    

• 权重量化源码

    def uniform_quantize(k):
      class qfn(torch.autograd.Function):
    
        @staticmethod
        def forward(ctx, input):
          if k == 32:
            out = input
          elif k == 1:
            out = torch.sign(input)
          else:
            n = float(2 ** k - 1)
            out = torch.round(input * n) / n
          return out
    
        @staticmethod
        def backward(ctx, grad_output):
          # STE (do nothing in backward)
          grad_input = grad_output.clone()
          return grad_input
    
      return qfn().apply
    
    
    class weight_quantize_fn(nn.Module):
      def __init__(self, w_bit):
        super(weight_quantize_fn, self).__init__()
        assert w_bit <= 8 or w_bit == 32
        self.w_bit = w_bit
        self.uniform_q = uniform_quantize(k=w_bit)
    
      def forward(self, x):
        if self.w_bit == 32:
          weight_q = x
        elif self.w_bit == 1:
          E = torch.mean(torch.abs(x)).detach()
          weight_q = self.uniform_q(x / E) * E
        else:
          weight = torch.tanh(x)
          max_w = torch.max(torch.abs(weight)).detach()
          weight = weight / 2 / max_w + 0.5
          weight_q = max_w * (2 * self.uniform_q(weight) - 1)
        return weight_q
  

3. 特征量化

• 特征是简单的直接量化处理
  

4. 梯度量化

• 梯度量化方法

  
  梯度对于量化更为敏感，所以前人的低比特梯度量化一般不 work，为了进一步补偿梯度量化引入的潜在偏差，这里加上随机均匀噪声从而增加鲁棒性，噪声的量级和量化误差类似：
  

• 一般板端部署模型时不需要在板端进行训练，所以一般都是使用全精度的梯度在服务器端进行模型训练，然后在板端部署，避免梯度量化产生的掉点

5. 整体算法流程

• 需要注意的是这些低比特的层之间会有高比特精度的 affine 层连接，从而节省了 conv 层等高功耗层的计算量

6. 不同层的量化位宽处理

• 第一层因为量化掉点明显并且第一层的 channel 数一般较少，所以一般第一层权重不量化，仅做特征量化
• 最后一层考虑到输出 channel 或者 fc 节点数较少，不量化的影响较小，为了保证精度一般最后一层的权重和特征都不做量化

实验结果

• 可以看出在梯度的位宽有保证的情况下相比对浮点模型的掉点较少
  
  

Thoughts

• 低比特量化的经典作品，梯度量化对于板端训练的意义重大