LSQ+: Improving low-bit quantization through learnable offsets and better initialization

URL

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

TL;DR

• 典型量化算法（如：DoReFa、PACT、LSQ）对 feature 的量化通常使用无符号量化（负半轴量化到零点）：
  • 这种量化算法适用于使用 Relu 激活函数的网络。
  • 但对使用新式激活函数（如：swish, H-swish, Mish, Leaky-ReLU）的网络（如：EfficientNet、MixNet）会造成较大的量化误差。
  • LSQ+ 作为一种非对称量化算法可以学习 scale 和 offset，去适应需要负数激活函数的网络。
• 量化网络训练收敛较难，所以 LSQ+ 提出了一种比较高效的 scale 和 $\beta$ 初始化方法：MSE 最小误差初始化方法优于最大最小值。

Algorithm

1. LSQ 算法存在的问题

• LSQ 公式
  $\bar{x} = \lfloor clamp (\frac{x}{s}, n, p)\rceil$
  $\hat{x} = \bar{x} \times s$
• LSQ 对例如 Leaky-ReLU、Swish 这种存在负半轴的激活函数量化有两种方法：
  • 截断负半轴，即 $n=0,\ p=2^b-1$，显然降低模型表现能力
  • 正负半轴相同尺度量化，即 $n=-2^{b-1},\ p=2^{b-1}-1$，由于正半轴的信息量远高于负半轴，同尺度量化会增大正半轴的量化误差

2. LSQ+ 的解决方案

• LSQ+ 的改进主要包含两个方面：
  • 在 LSQ 设置了可学习的参数 scale 的基础上，在 activation 量化上 （weight 量化没有 offset）增加了另外一个可学习参数 offset
  • 对网络初始化的方法进行了改进

2.1 增加可学习 offset 参数

• LSQ+ 公式
  $\bar{x} = \lfloor clamp (\frac{x-\beta}{s}, n, p)\rceil$
  $\hat{x} = \bar{x} \times s + \beta$
• 对 s 的梯度计算:
  $\frac{\partial\hat{x}}{\partial{s}}$ = $\frac{\partial\bar{x}}{\partial{s}}s + \bar{x} \simeq \begin{cases} -\frac{x-\beta}{s}+\lfloor\frac{x-\beta}{s}\rceil &,if\ n < \frac{x-\beta}{s} < p\\ n&, if\ \frac{x-\beta}{s} < n\\ p&,if\ p<\frac{x-\beta}{s}\end{cases}$
• 对 $\beta$ 的梯度计算
  $\frac{\partial\hat{x}}{\partial{\beta}}=\frac{\partial\bar{x}}{\partial{\beta}}+1\simeq \begin{cases} 0&,if\ n<\frac{x-\beta}{s} <p\\ 1&,otherwise\end{cases}$
• 加入了 $\beta$ 参数后，对存在负半轴的激活函数进行非对称量化将几乎没有额外开销
  $\hat{w}\hat{x}=(\bar{w}\times s_w)(\bar{x}\times s_x + \beta)=\bar{w}\bar{x}s_ws_x+\begin{matrix}\underbrace{\beta s_w\bar{w}}\\ bias\end{matrix}$

2.2 更合理的 scale 和 $\beta$ 参数初始化方法

• 低比特量化神经网络训练的最终性能与参数初始化方法关联性较大，这在深度可分离卷积网络中更为明显（例如： MobileNet系列）

2.2.1 weight 量化的 scale 参数初始化方法

• LSQ 对于对称量化的 weight 的 scale 初始化公式是：$s_{init} = \frac{2<|w|>}{\sqrt{p}}$
• 作者认为这样的 scale 初始化方法会导致初始 scale 太大，所以 LSQ+ 对 weight 对称量化的 scale 初始化方法是：$s_{init}=\frac{max(|\mu-3\times\sigma|, |\mu+3\times\sigma|)}{2^{b-1}}$
  • 其中，$\mu,\ \sigma$ 分别表示本层权重的均值和标准差

由于 weight 是对称量化，所以不需要 $\beta$ 参数

2.2.2 feature 量化的 scale 和 $\beta$ 参数初始化方法

• 一个最理想的量化方式是 x 只被量化，没有被 clamp，因此根据 LSQ+ 的量化公式可知：
  • $\frac{x_{min}-\beta_{init}}{s_{init}}\rightarrow n,\ \frac{x_{max}-\beta_{init}}{s_{init}}\rightarrow p$
  • 化简后：$s_{init}=\frac{x_{max}-x_{min}}{p-n},\ \beta_{init} = x_{min} - n \times s_{init}$
  • 但是，这种完全不 clamp 的操作容易被离群点干扰，所以 $s_{init},\ \beta_{init}$ 还是使用量化前和反量化后的数据最小 MSE loss 来确定，即：$s_{init},\ \beta_{init}=argmin_{s,\beta} ||\hat{x}-x||^2_F$

Thought

• 量化算法中多一个可以被硬件无痛吸收的自由度自然是好事