CVPR 2022 | 升级传统模型量化STE反向求导?CMU，脸书，港科提出新思路

模型量化是一种非常实用的模型压缩方法，经常被用于减小神经网络模型大小和加速神经网络的运算，从而使得原本耗费计算量的大模型，可以运行在资源有限的移动端上。

近些年来，模型量化算法层出不穷，但主要都是研究量化算法的前向传播过程，而很少研究反向求导过程的优化。反向求导算法大多还是基于2013年Bengio 等人提出的straight-through estimator（STE）。STE隐含了一个约束条件，即量化输入的阈值区间步长和输出的步长必须等长，这就大大限制了量化器的设计空间。

另一方面，量化算法常常受限于拟合能力，均匀量化的拟合能力有限，而拟合能力更强的非均匀量化由于输出不均匀， 从而难以被硬件加速。因此，本文希望能结合均匀量化和非均匀量化的优点，通过学习输入区间的大小而固定输出区间长度，从而即达到输出是均匀的，可以像均匀量化一样进行快速的矩阵乘（matrix multiplication）运算，同时输入区间灵活可学，实现更高的网络拟合能力。

N2UQ的设计思路是把量化看作多段的阶梯，每一段都可以视作是一个二分类问题， 对应的导数可以从该段二分类函数的概率模型中推到而得，如下图所示，具体推导过程可以看原论文，非常简单易懂。

最终推导而得的量化函数不仅具有传统STE的传递导数到输入变量的作用，并且能根据输入的是数值分布和网络输出的需要， 自动学习每个区间的大小。而这个难以求导数的区间阈值的导数也可以由大一统的G-STE导数拟合方法得到。

推导得到的正向函数中，输入x^r被可学习的参数a划分到了多个量化区间内，输出均匀的量化结果：

根据G-STE，可以求得对应的反向拟合函数：

从这个反向拟合函数可以直接求得关于输入x^r 和参数a的导数。

可以看出，当所有区间都等长的时候，G-STE就退化成STE; 而当需要区间不等长时，G-STE可以根据输入输出分布自动学习最佳的区间长度。

作者将Nonuniform-to-Uniform Quantization (N2UQ)用于ResNet18, 34, 50 以及MobileNet V2 网络结构上， 均取得远超state-of-the-art 的结果。

ResNet：

可以看出随着量化位宽的增加，N2UQ的优势越发明显。

MobileNet：

论文代码已经开源，感兴趣的读者可以试用。