Paper Reading||Differentiable Dynamic Quantization with Mixed Precision and Adaptive Resolution

论文(ICML 2021 Accepted)：https://arxiv.org/abs/2106.02295
暂无源码

1 Motivation

1. 不同的网络架构在不同的层中分配的内存和计算复杂度是不同的
   ，会导致量化处于次优状态，特别是当量化参数(如比特宽度、动态范围和步长)逐层冻结训练时(可以参考 PROFIT)。
2. (在没有 STE 之前，)基于梯度的低比特宽度模型训练是困难的，因为先前量化函数的梯度可能会消失，即量化后的 DNN 在反向传播时可能返回零梯度。
3. 量化往往会在轻量级架构上表现较差，如 MobileNet 系列。这主要因为这些模型已经是高效和紧凑的，没有太多的量化空间。

 　作者主要讨论了 Motivation 3。一般大家认为模型中的 Weight 统计分布都是钟形分布，这就不适用于 MobileNet 系列、EfficientNet 系列、ShuffleNet 系列，如 Fig 1 (b)(c)(d) 的分布图均是 MobileNet 系列中可能的 Weight 分布。这些不常规的分布尤其出现在深度可分离卷积层和组卷积层中。一个直观的感受是，这些不同的分布形状会偏好不同的量化参数，如位宽(Bit-width)、量化电平(Quantization Level)、量化间隔(Interval)、动态范围(Dynamic Range)等。

 　综上所述，作者提出了自己的方法，DDQ：可以针对不同的模型架构来可微地学习对应的量化参数，包括位宽、量化电平和动态范围。形象的解释还是如 Fig 1 所示。图中 (a) 列比较了均匀量化曲线和二的幂次量化曲线，意在说明，二的幂次量化往往能更有效地作用于钟形分布，以减小量化误差。而 (b)(c)(d) 列的量化曲线均来自于 DDQ，也就是说 DDQ 会根据数据的分布情况，找出 Dense Region，然后做出类似二的幂次量化的量化。Dense Region 可以有多个，如 (d) 中的双峰分布，就导致了曲线中有两个 Dense Region。

 　作者提到 DDQ 的三个好处：

• 其适用于高效且紧凑的模型框架，能根据不同的层的情况来定制不同的量化参数以达到更高的准确度。
• 其对硬件友好，很容易利用通用矩阵乘(GEMM,General Matrix Multiplication)来部署在 ARM 架构上。设 $b$ 为位宽，DDQ还能利用矩阵再参数化的方法将矩阵的储存空间从 $O(2^{2b})$ 压缩至 $O(log{2}^b)$ 。
• 广泛的实验表明 DDQ 在 ImageNet 和 CIFAR10 上比以往的量化算法表现要好，且第一次能在 MobileNet 上几乎不损失准确度地量化到 4-bit。

2 Method

2.1 Preliminary and Notations

 　在模型量化中，一个连续的值 $x\in \mathbb{R}$ 会被量化成离散的值 $x_q$ ，来自于集合 $\mathbf{q}=[q_1,q_2,...,q_n{]}^{\mathsf{T}}$ 。集合 $\mathbf{q}$ 内的元素就是量化电平，其中 $x_q\in \mathbf{q}$ ， $n=2^b$ ， $b$ 为位宽。目前的量化方法就是通过均匀量化或二的幂次量化来映射 $\mathbf{q}$ 。设量化有关参数集为 θ = { b , c } T \bm{\theta}=\{b,c\}^{\mathsf{T}} θ={b,c}T ，$c$ 为表征对称范围的截断阈值， $\bar{x}$ 为量化零点。

• 均匀量化：虽然均匀量化可能简单有效，但它假设权重分布是均匀的，这在许多 DNN 中是不可靠的。一个 b-bit 对称均匀量化器的量化电平 $q_u$ 为：

• 二的幂次量化：其可以产生一个 Dense Region 来集中表达单峰钟形分布。当然通过预定义组合方案也可以产生多个 Dense Region、表达多个峰。一个 b-bit 对称的二的幂次量化器的量化电平 $q_p$ 为：

 　以上这两个量化器均会固定码本 $\mathbf{q}$ 然后优化量化尺度(或称为量化步长) $d$ 和截断阈值 $c$ 。但它们有一个明显的缺点——即量化电平是预定义的，这就无法适应训练期间每个层的权重或激活的不同分布。

2.2 Dynamic Differentiable Quantization (DDQ)

Formulation for Arbitrary Quantizaion

 　DDQ 则会学习所有的量化参数，即还会学习量化电平。设有一个函数 $Q(x;\mathbf{\theta })$ 能将连续值 $x$ 转换成 $x_q$ ，$x_q\in \mathbf{q}$ ，即 $x_q=Q(x;\mathbf{\theta })$ ，其中 $\mathbf{q}$ 初始化为均匀量化的格式。设 $x_o^i\in {\mathbf{x}}_{\mathbf{o}}$ ， x o ∈ { 0 , 1 } n × 1 \bm{x_{o}} \in \{0, 1\}^{n \times 1} xo​∈{0,1}n×1 ，其中 ${\mathbf{x}}_{\mathbf{o}}$ 为二进制 One-hot 向量，用于选取量化电平。设 U ∈ { 0 , 1 } n × n \bm{U} \in \{0, 1\}^{n \times n} U∈{0,1}n×n 为二进制分块对角矩阵， $n$ 为量化电平个数，$Z_U$ 为一个用于均衡离散值的常数标准化因子。DDQ 就由下面的低精度矩阵向量乘积公式表示：

 　其中 $\mathbf{q}$ 会被 $\tilde{\mathbf{q}}$ 重新设置，如下：

 　其中 $R(\cdot )$ 为将 $\tilde{\mathbf{q}}$ 变为 $b_q$ 比特表征数的均匀量化函数( $b_q<b$)。于是公式 (3) 就有可学习参数 θ = { q ~ , U } \bm{\theta}=\{\tilde{\bm{q}},\bm{U}\} θ={q~​,U} ，接受 SGD 梯度更新，让 DDQ 能自动捕获全精度的权重分布，如 Fig (b) (c) (d) 所示的那样。直观地说，就是不同的 $x_o,\mathbf{U},\mathbf{q}$ 使得 DDQ 能实现不同的量化方法。为了方便理解，Fig 2 (a) 比较了 DDQ 和 其他 Rounding-based 方法的计算图。可以发现 DDQ 的所有量化电平都会被学习，而不是像现有技术那样仅学习量化尺度 $d$ 。

Discussions of Representation Capacity

 　DDQ能够表征成的量化方法非常广泛。当 公式 (1)中的 $\mathbf{q}={\mathbf{q}}_u$ ， $Z_U=1$ ， $\mathbf{U}=\mathbf{I}$ ，$\mathbf{I}$ 为单位矩阵时，公式 (3) 就可以表示成均匀量化。当 公式 (2) 的 $\mathbf{q}={\mathbf{q}}_p$ ，$Z_U=1$ ， $\mathbf{U}=\mathbf{I}$ 时，公式 (3) 就可以表示成二的幂次量化。当 $\mathbf{q}$ 被梯度更新时，DDQ 就可以表示具有不同密集区域的任意量化电平。
 　此外，当 $\mathbf{U}$ 为分块对角阵时，DDQ 就可以表示成混合精度量化。例如，如 Fig 2 (b) 所示，当 $\mathbf{q}$ 的长度为8(即位宽为 3 bit )，$Z_U=\frac{1}{2}$， $\mathbf{U}=\mathrm{D}\mathrm{i}\mathrm{a}\mathrm{g}(1_{2\times 2},\cdot \cdot \cdot ,1_{2\times 2})$ ， $\mathrm{D}\mathrm{i}\mathrm{a}\mathrm{g}(\cdot )$ 能够产生对角分块并填入相应值的矩阵(其余位置为0)，$1_{2\times 2}$ 表示一个 $2\times 2$ 的值为1的矩阵，于是在公式 (3) 中 $\mathbf{U}$ 能通过加权 $\mathbf{q}$ 中相邻的离散值来表征成一个 2-bit 量化器。再例如，当 $\mathbf{U}=\mathrm{D}\mathrm{i}\mathrm{a}\mathrm{g}(1_{4\times 4},\cdot \cdot \cdot ,1_{4\times 4})$ 且 $Z_U=\frac{1}{4}$ 时，公式 (3) 就会变成一个 1-bit 量化器。并且，当 ${\mathbf{x}}_{\mathbf{o}}$ 为软 One-hot 向量(即有多个元素是非零的)时，公式 (3) 就能表示成软量化(Soft Quantization)。

Efficient Inference on Hardware

 　DDQ 是一个统一的量化器，既支持自适应梯度更新的 $\mathbf{q}$ 又支持预定义成均匀量化或二的幂次量化的超参数的 $\mathbf{q}$ 。其对资源有限的硬件很友好。当 DDQ 作为均匀量化，部署要硬件上时，会用正常的 Round 操作和 Scale 操作来完成量化，而没有 $\mathbf{U}$ 和矩阵-向量乘积。并且，当 DDQ 的 $\mathbf{q}$ 为自适应梯度更新时，还可以使用 通用矩阵乘(GEMM)。例如，设 $y$ 为神经元的 Activation， $x_q$ 为上一层的 Feature Map 离散值， $w$ 为待量化的连续权重值， $\mathbf{U}$ 为二进制矩阵， ${\mathbf{w}}_o$ 为 One-hot 向量，那么就有公式 $y=Q(w;\mathbf{\theta })x_q={\mathbf{q}}^{\mathsf{T}}\frac{\mathbf{U}}{Z_U}{\mathbf{w}}_o{x}_q$ 。于是就能利用 GEMM 来加速 $\frac{\mathbf{U}}{Z_U}{\mathbf{w}}_o{x}_q$ 的计算，然后再乘与 1-d 向量 $\mathbf{q}$。这个 $\mathbf{q}$ 是所有的卷积权重共享的，其可以在所要求的硬件环境中置为 Float32、Float16 或 INT8。

2.3 Matrix Reparameterization of U

 　在公式 (3) 中， $\mathbf{U}$ 是一个可学习矩阵变量，其比较棘手，表现在两方面：

• 为了保证 DDQ 为一个有效的量化器， $\mathbf{U}$ 必须为二进制的块对角结构，而直接使用 SGD 作用于整个矩阵则会破坏这个结构。
• 当位宽增大(如 $2^{2b}$)， $\mathbf{U}$ 的参数个数也会增多。
  为了解决以上问题，作者通过一系列的小矩阵来显示合成构造 $\mathbf{U}$ ，再用于前向传播。

A Kronecker Composition for Quantization

 　 $\mathbf{U}$ 的参数可以由 $2^{2b}$ 减少成 $\mathrm{l}\mathrm{o}\mathrm{g}{2}^b$ 。设一系列 $2\times 2$ 大小的矩阵构成集合 { U 1 , U 2 , ⋅ ⋅ ⋅ , U b } \{\bm{U}_{1}, \bm{U}_{2}, ···,\bm{U}_{b}\} {U1​,U2​,⋅⋅⋅,Ub​} ，那么 $\mathbf{U}$ 就可以写成 $\mathbf{U}={\mathbf{U}}_1\otimes {\mathbf{U}}_2\otimes \cdot \cdot \cdot \otimes {\mathbf{U}}_b$ ，其中 $\otimes$ 指的是克罗内克积(Kronecker Product)。每个 ${\mathbf{U}}_i(i=1,\cdot \cdot \cdot ,b)$ 要么为单位矩阵(定义为 $I$)，要么为全1矩阵(定义为 $1_{2\times 2}$ )，共同组成块对角矩阵 $\mathbf{U}$ 。
 　为了进一步压缩储存空间，作者为每个 ${\mathbf{U}}_i$ 定义了可学习的门值参数(Gate Value)，表示为 ${\mathbf{U}}_i=g_i\mathbf{I}+(1-g_i)1_{2\times 2}$ ，其中 $g_i=H(\widehat{g_i})$ ， $H(\cdot )$ 为海维赛德阶跃函数(输入大于等于0则输出为1，反之为输出为0)，也就是说 $g_i$ 为二进制值。形象展示如 Fig 2 (c) 所示。也就是说，可学习参数不是整个 $\mathbf{U}$，而是 { g ^ i } i = 1 b \{\hat{{g}}_{i} \}^{b}_{i=1} {g^​i​}i=1b​ ，就能显著地减少参数量，从 $2^b\times 2^b=2^{2b}$ 变为 $b$ 。
 　在克罗内克积作用下，DDQ 的量化参数变为 θ = { q , { g ^ i } i = 1 b } \bm{\theta}=\{ \bm{q}, \{\hat{{g}}_{i} \}^{b}_{i=1}\} θ={q,{g^​i​}i=1b​} ，可以根据 Per-channel 或 Per-tensor 力度来独立设置。

2.4 Discussions of Relationship between $\mathbf{U}$ and $\mathbf{g}$

 　设 $\mathbf{g}=[g_1,\cdot \cdot \cdot ,g_b{]}^{\mathsf{T}}$ ，在用 $\mathbf{g}$ 来表征 $\mathbf{U}$ 时，需要注意以下两点：

• 排序。如 Fig 2 (c) 所示， $\mathbf{g}$ 内元素的不同的排序会导致不同的$\mathbf{U}$。当 $\mathbf{g}$ 内元素按照降序排列时，得到的 $\mathbf{U}$ 就是个正常合法的量化器；反之，当 $\mathbf{g}$ 内元素没有降序排列时，得到的 $\mathbf{U}$ 就不是块对角矩阵了，这就不正常，不能称作量化器了。
• 门值加和。设门值总和为 $s={\sum }_{i=1}^b{g}_i$，$0\le s\le b$。 $\mathbf{U}$ 是有 $2^s$ 个对角块的块对角矩阵，那么${\mathbf{U}}^{\mathsf{T}}q$ 就有 $2^s$ 个不同的离散值，表征一个 s-bit 量化器。需要提的一点是，DDQ 允许在给定内存条件下约束每一力度(Per-tensor 或 Per-channel)的 $s$ 值，以便可以将最佳比特宽度分配给 DNN 的不同层或核。

3 Training DNN with DDQ

 　DDQ 取得的结果一般是混合精度的，以最大化限定储存空间下的性能。

3.1 DNN with Memory Constraint

 　假设一个 DNN 有 $L$ 层，那么每一层的前向传播计算公式如下：

 　其中 $\ast$ 指卷积计算，${\mathbf{y}}^l$ 和 ${\mathbf{y}}^{l-1}$ 分别指 $l$-th 层的输出和输入。 $F$ 为激活函数，$Q$ 为 DDQ 的量化操作。设 ${\mathbf{W}}^l\in {\mathbb{R}}^{C_{out}^l\times C_{in}^l\times K^l\times K^l}$ 为核权重，${\mathbf{b}}^l$ 为 Bias，$C_{out}$ 为输出通道数， $C_{in}$ 为输入通道数， $K$ 为卷积核尺寸。 { g i l } i = 1 b \{g_{i}^{l} \}^{b}_{i=1} {gil​}i=1b​ 是 $l$-th 层的门值集，而该层的位宽设为 $s^l={\sum }_{i=1}^b{g}_i^l$。设一个模型的所有 Weight 占用的储存空间为 $\zeta$ ，则有：

 　如果说给定最大限度的储存空间为 $\zeta (b^1,\cdot \cdot \cdot ,b^L)$ ，由帕累托最优解课近似损失函数为：

 　可见损失函数被重新加权过了。 $\alpha$ 是超参数，当满足储存限制时，置 $\alpha =0$ ；当惩罚机制起作用时，置 $\alpha <0$ 。

3.2 Updating Quantization Parameters

 　该章节详细描述了各个量化参数的更新方式。往后的公式里为了简便会忽略 $‘l’$ 上标。

Gradients with Respect to $\mathbf{q}$

 　为了更新 $\mathbf{q}$，作者又设了个可学习参数 $\tilde{\mathbf{q}}$，有 $\mathbf{q}=\frac{R(\tilde{\mathbf{q}})(x_{max}-x_{min})}{2^b-1}+x_{min}$ ，以保证每个量化电平均在 $[x_{min},x_{max}]$ 中，其中 $x_{max}$ 和 $x_{min}$ 为一个层的 Weight 的原精度最值。 $R(\cdot )$ 定义为均匀量化函数，能将 $\tilde{q}$ 转换为 $b_q$ 比特 $(b_q<b)$ 。对于 $\mathbf{q}$ 中的每一个 $q_k$，有如下梯度：

 　其中 $x_q^i=Q(x^i;\mathbf{\theta })$ ，是 DDQ 量化器的输出。 ${\mathcal{S}}_k$ 表示离散到对应量化级别 $q_k$ 的值的索引集。观察 公式 (8)，可以发现 $q_k$ 的梯度是 $\frac{\partial \mathcal{L}}{\partial x_q^i}$ ，即说明在 Dense Region 中的量化级别将具有更大的梯度。

 　至于门值 { g i } i = 1 b \{ g_{i}\}_{i=1}^{b} {gi​}i=1b​ 的梯度，会在 Appendix A 中展示(附录没找到)。

Gradient Correction for ${\mathbf{x}}_q$

 　为了减少量化误差 $||{\mathbf{x}}_q-\mathbf{x}|{|}_2^2$，作者提出了一个“梯度修正项”来正则化 ${\mathbf{x}}_q$ 的梯度：

 　上式的第一个公式使用到了 STE。如此一来，修正过的梯度就能通过反向传播到公式 (8) 里的 $\mathbf{q}$ 和 { g i } i = 1 b \{ g_{i}\}_{i=1}^{b} {gi​}i=1b​ 上，且不会影响 $\mathbf{x}$ 的梯度。该梯度修正项是有效的，可以被理解为对 $||{\mathbf{x}}_q-\mathbf{x}|{|}_2^2$ 做了 ${\ell }_2$ 正则化。但是这不等价于直接在损失函数上施加 ${\ell }_2$ 正则化，否则由于 STE 作用，其会无效。

Implementation Details

 　DDQ 可以很容易地在 Pytorch、TensorFlow 等框架上实现。需要注意的是，以下梯度用到了 STE 思想：

• $\frac{\partial x_q}{\partial x}=1$，${\hat{q}}_{min}\le x\le {\hat{q}}_{max}$，${\hat{q}}_{min}$ 和 ${\hat{q}}_{max}$ 是 $\hat{\mathbf{q}}={\mathbf{U}}^{\mathsf{T}}\mathbf{q}$ 。
• $\frac{\partial g_k}{\partial {\hat{g}}_k}=1$， $|{\hat{g}}_k|\le 1$ 。

 　更多细节展示在 Algorithm 1 中。

4 Experiments

 　如果没有特殊说明，作者的所有实验都是在 $b_q=8$ 下进行的。

4.1 Evaluation on ImageNet

 　Table 1 比较了 DDQ 和现有模型。

 　Fig 3 展示了DDQ训练的MobileNetv2和ResNet18的每层的 Weight 的位宽的收敛结果(初始化为 8-bit)。作者发现：

• 这两种网络都倾向于在较浅的层中应用更多的比特宽度，较浅的层具有较少的参数，因此较少受到内存约束的影响。这允许学习更好的功能表示，从而减少性能下降。
• 深度卷积比常规卷积具有更大的比特宽度。不规则权重分布的深度卷积是导致MobileNet量化困难的主要原因。通过混合精度训练，DDQ为深度卷积分配了更多的比特宽度，以缓解这一困难。

 　Fig 4 更细致地展现了 ResNet18 下 DDQ 通过减小逐层位宽度，为每个层分配适当精度的过程(作者说这是 4-bit 的结果，我觉得有点不太理解，每一层的位宽都不一样呀，怎么能统称为 4-bit 呢)。

4.2 Ablation Study

Ablation Study I: mixed versus fixed precision

 　还是参考 Table 1，同时作者也做了有关实验，如 Table 2 所示。

Ablation Study II: adaptive resolution

 　接着作者固定 $\mathbf{U}$ (即位宽固定了)，只更新 $\mathbf{q}$，实验结果展示于 Table 3 中。

Ablation Study III: gradient correction

 　Fig 5 (a) 展示的是量化误差收敛曲线。Fig 5 (b) 展示的是量化电平的可视化结果。当然作者的实验也部分展示在 Table 3 中。

4.3 Evaluation on Mobile Devices

 　作者在 Table 4 中进一步评估了移动平台上的DDQ对准确性和延迟的权衡表现。具体细节展示在 Appendix C 中。

个人思考

• 个人感觉作者写的不是很好，比如一张图里的信息量太大了，居然需要三、四个小节的上下文串联才能看懂在干嘛，比如 Fig 2 (c) 中，非块对角矩阵居然是错误示例，但是一开始我一直以为是对的，直到我看到了 2.3 节。
• 首先比较好奇的是 $x_o^i$ 是怎么来的，但作者没有先阐述，导致接下来的阅读中一直在想这个事情。
• 梯度修正项好像有问题。个人感觉吧，或许是因为我笨没理解到位。
• 为什么是对硬件友好呢？作者的部署应该是利用了 FBGEMM，利用通用矩阵乘？部署在手机上？
• 作者对所有的层的 Weight 都进行量化了，但是作者自己跑了其他方法的实验，其他方法的实验也是Weight全都量化吗？
• 附录没找着在哪里获取，有机会看见了再整理了。
• 公式 (4) 的 $R(\cdot )$ 函数公式是什么，这我有点不太明白。