[人工智能]Paper Reading||Distance-aware Quantization

论文(ICCV 2021)：https://arxiv.org/abs/2108.06983
源码(Pytorch)：https://github.com/cvlab-yonsei/DAQ

1 Motivation

1. 使用 STE，虽然使得Round操作处处可导(梯度为1)，但会引入梯度不匹配、不合理的问题。
2. 软量化器(Soft Quantizer)方法虽然能有效改进梯度不匹配问题，但从软量化器直接“跳变”为硬量化器会带来 Quantizer Gap 问题。
3. 在软量化器中引入能够逐渐增大的 Temperature 参数，虽然能缓解 Quantizer Gap 问题，但又会导致梯度流变得不稳定。

?　对于 Motivation 2 ，作者认为，软量化器是有效的，但其仅在训练阶段使用，而推理阶段使用的还是硬量化器。那么如果直接从软量化器转化成硬量化器，就会导致待量化的数据映射得比较“剧烈”，跳变还是比较僵硬，即所谓的 Quantizer Gap 问题。
?　虽然 Motivation 3 中，有人通过引入 Temperature 参数，让软量化器逐渐过渡成硬量化器，但节奏把控不好又会导致梯度流变得不稳定。

2 Method

2.1 Overview

?　设全精度输入为 $\hat{x}$ (可以代表 Weight $\hat{w}$ 或 Activation $\hat{a}$ )，经过 $b$-bit 量化器 $Q$ 会得到量化后的值 Q ( x ^ ) ∈ { 0 , 1 , . . . , 2 b ? 1 } Q(\hat{x})\in \{0, 1, ..., 2^b-1 \} Q(x^)∈{0,1,...,2b?1} 。整个流程如 Fig 2 所示。

?　在此之前，先要对 $\hat{x}$ 进行 Clip 和 Normalize，将数据分布范围限制到量化区间内，如下公式所示：

?　其中 $u$ 为上边界， $l$ 为下边界， $x$ 就是那个标准化后的输入。然后再经过反量化，还原出 $w_q\in [?1,1]$ 和 $a_q\in [0,1]$ ：

?　设有张量形式的 ${\mathbf{w}}_q$ 和 ${\mathbf{a}}_q$，那么一个卷积操作就可以写成：

?　其中 $s$ 为一个用于调节卷积输出尺寸的可学习参数。

2.2 DASR

?　四舍五入操作即将一个全精度的输入仿射到最邻近的量化值。这个过程可以看成两步：第一步，计算全精度值与量化值的距离。第二步，根据 $argmin$ 操作将全精度值映射成最短距离的那个量化值，但整个过程不可导。如 Fig 1 所示。

?　针对这个问题，作者提出 DASR，用可微的赋值算子来近似离散四舍五入函数。与两布操作类似，DASR 经过公式 (1) 得到 $x$ ，然后根据有限个的量化值 q ∈ q , ? q = { 0 , ? 1 , ? . . . , ? 2 b ? 1 } q \in \bm{q},~\bm{q}=\{0,~1,~...,~2^b-1\} q∈q,?q={0,?1,?...,?2b?1} 来计算 Distance Scores。接着在利用 $softargmax$ 函数，将 $x$ 分配成与量化值非常接近的一个浮点值。具体描述如下：

• Distance socre。计算公式如下：

?　当 $x$ 越接近量化值 $q$ ，这个距离分数就会越大。需要注意的是，该公式的计算复杂度随着位宽的增加而增大，因为参与计算的 $q$ 的个数会增多。为了解决这个问题，作者为每个 $x$ 只计算两次，一次是和最邻近取 $floor$ 得到的量化值 $q_f$ ，一次是和最邻近取 $ceil$ 得到的量化值 $q_c$ 。 $q_c?q_f=1$ 。

• Soft assignment。作者使用 Kernel Soft Argmax 来近似 Discrete Argmax，使得梯度变得可导：

?　其中 $\beta$ 为 Temperature 参数，$m_x$ 函数输出为起到加权作用的距离置信度。上式就是在 $q_f$ 和 $q_c$ 之间权衡取值。$m_x$ 函数进一步定义如下：

?　其中 $k_x(q_i)$ 为以 $q_i$ 为中心的一维高斯核函数，$x$ 越接近 $q_i$ ，则 $k_x(q_i)$ 越大，反之亦然。综上， $?$ 函数的效果如 Fig 3 所示，避免了梯度爆炸问题。

2.3 Temperature

?　Temperature 参数 $\beta$ 可以调节距离置信度函数 $m_x$ 的形状分布，即也影响了 Soft Assignment $?$ 。但是固定的 $\beta$ 是有缺点的：

1. $\beta$ 太小(如 $\beta =4$) 时，会造成 Quantizer Gap 问题，过渡点附近的值 Tuning 不充分。
2. $\beta$ 太大(如 $\beta =24$) 时，虽然能有效缓解 Quantizer Gap 问题，让过渡点附近的 Weight 积极更新，但会对量化值附近的 Weight 引入梯度消失问题。

?　综上，Temperature 参数 $\beta$ 需要进一步地设计：

• Temperature controller。作者根据 $x$ 到过渡点的距离将 $\beta$ 设计成自适应式，在 Quantizer Gap 和 梯度消失问题上做了很好的折衷，有效缓解了这两个问题。定义自适应的 Temperature 参数 ${\beta }^{?}$ ：

?　其中 $\gamma$ 为正值常量， $s_x(q_i)$ 为加权过的距离得分：

?　当 $x$ 非常接近于过渡点时， $s_x(q_f)$ 和 $s_x(q_f)$ 就几乎相等。具体来说，当 $x$ 越接近过渡点时，公式 (7) 的分母会越小，使得对应的 Temperature 参数 ${\beta }^{?}$ 增大，有效缓解 Quantizer Gap 问题。另一方面，当 $x$ 越远离过渡点时，公式 (7) 的分母会越大，使得对应的 Temperature 参数 ${\beta }^{?}$ 减小，从而有效缓解消失梯度问题。
?　总的来说，作者就是让过渡点附近的温度大些，让量化值附近的温度小些，温度会随着 $x$ 的变化而变化，也可以理解为每个 $x$ 都独立地有对应的量化分配方式(其实同一个 Quantization Bin 内的值均量化到上一个量化值或下一个量化值，仅梯度的分配有差别，后面作者会提到)。这本质上也是是适应地调整了梯度的大小。具体来说，对于 Fig 4 中的 $x_1$ 点来说，其温度为 4 ，那么其量化分配方式为 $?(x;\beta =4)$ 。对于 Fig 4 中的 $x_2$ 点来说，其温度为 6 ，那么其量化分配方式为 $?(x;\beta =6)$ 。

?　进一步地，公式 (5) 可以化简成：

?　其中 $\lambda =\frac{1}{e^{\gamma }+1}$ 。$q_t$ 为过渡点，$q_t=\frac{q_f+q_c}{2}$ 。可以注意到， Fig 4 中不同的 Input $x$，分配到的离散值都是一样的(如灰线所示)。而实际是梯度的分配有细微差别，如 Fig 5 (灰线)所示。但是，$?$ 公式分配到的离散值并不是量化电平对应的离散值，故而公式 (9) 的输出需要经过一次标准化缩放成量化电平。有缩放公式为：

?　而最终应有：

?　上式即为 DAQ 的真正输出，即 $Q(\hat{x})$ 。

3 Experiment

?　Table 2 、Table 3、Table 4 为 ImageNet上的比较。Table 5 为 CIFAR10 上的比较。

?　Table 6 为针对 Temperature 参数的消融实验，Table 7 评估的是 DAQ 的训练开销。Fig 6 为训练过程可视化。

4 Appendix

4.1 Derivation of Eq. (9)

?　作者给出了 $?$ 公式最终详细的推导过程：

?　而 ${\beta }^{?}$ 为：

?　${\beta }^{?}$ 可以展开为：

?　所以进一步地有

?　上式中 $\lambda =\frac{1}{exp(\gamma )+1}$ 。进一步地，可以整理为：

?　注意 $q_c?q_f=1$。故最终就有：

4.2 Overall Process of DAQ

4.3 Design of DAQ

?　关于距离得分函数和采样核函数，作者也尝试过其他的公式，但是效果都不如原文里的好，故最后选择原文的方法。

4.4 Experiments on the SR Task

?　作者还在超分辨率任务上做了实验，具体在这里就不展示了。感兴趣的旁友可以翻翻原文。

个人思考

• 有一个细节，作者说自己的方法能避免梯度爆炸，但是随着温度的上升，最后还是会长得像 Sigmoid 函数，那么梯度在过渡点附近不也形成梯度爆炸的情况了吗？所以实质性地说，梯度爆炸情况没有消失，在过渡点的局部值仍会出现梯度爆炸，只是其他大部分的值梯度合理。
• 作者的这个温度不是随时间变化，而是随要量化的原值而变化。和EWGS有异曲同工之妙。