论文(ICCV 2021):https://arxiv.org/abs/2108.06983
源码(Pytorch):https://github.com/cvlab-yonsei/DAQ
DAQ 概要
1 Motivation
- 使用 STE,虽然使得Round操作处处可导(梯度为1),但会引入梯度不匹配、不合理的问题。
- 软量化器(Soft Quantizer)方法虽然能有效改进梯度不匹配问题,但从软量化器直接“跳变”为硬量化器会带来 Quantizer Gap 问题。
- 在软量化器中引入能够逐渐增大的 Temperature 参数,虽然能缓解 Quantizer Gap 问题,但又会导致梯度流变得不稳定。
对于 Motivation 2 ,作者认为,软量化器是有效的,但其仅在训练阶段使用,而推理阶段使用的还是硬量化器。那么如果直接从软量化器转化成硬量化器,就会导致待量化的数据映射得比较“剧烈”,跳变还是比较僵硬,即所谓的 Quantizer Gap 问题。
虽然 Motivation 3 中,有人通过引入 Temperature 参数,让软量化器逐渐过渡成硬量化器,但节奏把控不好又会导致梯度流变得不稳定。
2 Method
2.1 Overview
设全精度输入为 x ^ \hat{x} x^ (可以代表 Weight w ^ \hat{w} w^ 或 Activation a ^ \hat{a} a^ ),经过 b b b-bit 量化器 Q Q 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 所示。
在此之前,先要对 x ^ \hat{x} x^ 进行 Clip 和 Normalize,将数据分布范围限制到量化区间内,如下公式所示:
其中
u
u
u 为上边界,
l
l
l 为下边界,
x
x
x 就是那个标准化后的输入。然后再经过反量化,还原出
w
q
∈
[
−
1
,
1
]
w_q \in [-1, 1]
wq∈[−1,1] 和
a
q
∈
[
0
,
1
]
a_q \in [0, 1]
aq∈[0,1] :
设有张量形式的
w
q
\bm{w}_q
wq 和
a
q
\bm{a}_q
aq,那么一个卷积操作就可以写成:
其中
s
s
s 为一个用于调节卷积输出尺寸的可学习参数。
2.2 DASR
四舍五入操作即将一个全精度的输入仿射到最邻近的量化值。这个过程可以看成两步:第一步,计算全精度值与量化值的距离。第二步,根据
a
r
g
m
i
n
argmin
argmin 操作将全精度值映射成最短距离的那个量化值,但整个过程不可导。如 Fig 1 所示。
针对这个问题,作者提出 DASR,用可微的赋值算子来近似离散四舍五入函数。与两布操作类似,DASR 经过公式 (1) 得到 x x 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。接着在利用 s o f t a r g m a x soft~argmax soft argmax 函数,将 x x x 分配成与量化值非常接近的一个浮点值。具体描述如下:
- Distance socre。计算公式如下:
当
x
x
x 越接近量化值
q
q
q ,这个距离分数就会越大。需要注意的是,该公式的计算复杂度随着位宽的增加而增大,因为参与计算的
q
q
q 的个数会增多。为了解决这个问题,作者为每个
x
x
x 只计算两次,一次是和最邻近取
f
l
o
o
r
floor
floor 得到的量化值
q
f
q_{f}
qf ,一次是和最邻近取
c
e
i
l
ceil
ceil 得到的量化值
q
c
q_{c}
qc 。
q
c
−
q
f
=
1
q_{c}-q_{f}=1
qc−qf=1 。
- Soft assignment。作者使用 Kernel Soft Argmax 来近似 Discrete Argmax,使得梯度变得可导:
其中
β
\beta
β 为 Temperature 参数,
m
x
m_x
mx 函数输出为起到加权作用的距离置信度。上式就是在
q
f
q_{f}
qf 和
q
c
q_{c}
qc 之间权衡取值。
m
x
m_x
mx 函数进一步定义如下:
其中
k
x
(
q
i
)
k_x(q_i)
kx(qi) 为以
q
i
q_i
qi 为中心的一维高斯核函数,
x
x
x 越接近
q
i
q_i
qi ,则
k
x
(
q
i
)
k_x(q_i)
kx(qi) 越大,反之亦然。综上,
ϕ
\phi
ϕ 函数的效果如 Fig 3 所示,避免了梯度爆炸问题。
2.3 Temperature
Temperature 参数 β \beta β 可以调节距离置信度函数 m x m_x mx 的形状分布,即也影响了 Soft Assignment ϕ \phi ϕ 。但是固定的 β \beta β 是有缺点的:
- β \beta β 太小(如 β = 4 \beta=4 β=4) 时,会造成 Quantizer Gap 问题,过渡点附近的值 Tuning 不充分。
- β \beta β 太大(如 β = 24 \beta=24 β=24) 时,虽然能有效缓解 Quantizer Gap 问题,让过渡点附近的 Weight 积极更新,但会对量化值附近的 Weight 引入梯度消失问题。
综上,Temperature 参数 β \beta β 需要进一步地设计:
- Temperature controller。作者根据 x x x 到过渡点的距离将 β \beta β 设计成自适应式,在 Quantizer Gap 和 梯度消失问题上做了很好的折衷,有效缓解了这两个问题。定义自适应的 Temperature 参数 β ∗ \beta^* β∗ :
其中
γ
\gamma
γ 为正值常量,
s
x
(
q
i
)
s_x(q_i)
sx(qi) 为加权过的距离得分:
当
x
x
x 非常接近于过渡点时,
s
x
(
q
f
)
s_x(q_f)
sx(qf) 和
s
x
(
q
f
)
s_x(q_f)
sx(qf) 就几乎相等。具体来说,当
x
x
x 越接近过渡点时,公式 (7) 的分母会越小,使得对应的 Temperature 参数
β
∗
\beta^*
β∗ 增大,有效缓解 Quantizer Gap 问题。另一方面,当
x
x
x 越远离过渡点时,公式 (7) 的分母会越大,使得对应的 Temperature 参数
β
∗
\beta^*
β∗ 减小,从而有效缓解消失梯度问题。
总的来说,作者就是让过渡点附近的温度大些,让量化值附近的温度小些,温度会随着
x
x
x 的变化而变化,也可以理解为每个
x
x
x 都独立地有对应的量化分配方式(其实同一个 Quantization Bin 内的值均量化到上一个量化值或下一个量化值,仅梯度的分配有差别,后面作者会提到)。这本质上也是是适应地调整了梯度的大小。具体来说,对于 Fig 4 中的
x
1
x_1
x1 点来说,其温度为 4 ,那么其量化分配方式为
ϕ
(
x
;
β
=
4
)
\phi(x;\beta=4)
ϕ(x;β=4) 。对于 Fig 4 中的
x
2
x_2
x2 点来说,其温度为 6 ,那么其量化分配方式为
ϕ
(
x
;
β
=
6
)
\phi(x;\beta=6)
ϕ(x;β=6) 。
进一步地,公式 (5) 可以化简成:
其中
λ
=
1
e
γ
+
1
\lambda=\frac{1}{e^\gamma+1}
λ=eγ+11 。
q
t
q_t
qt 为过渡点,
q
t
=
q
f
+
q
c
2
q_t=\frac{q_f+q_c}{2}
qt=2qf+qc 。可以注意到, Fig 4 中不同的 Input
x
x
x,分配到的离散值都是一样的(如灰线所示)。而实际是梯度的分配有细微差别,如 Fig 5 (灰线)所示。但是,
ϕ
\phi
ϕ 公式分配到的离散值并不是量化电平对应的离散值,故而公式 (9) 的输出需要经过一次标准化缩放成量化电平。有缩放公式为:
而最终应有:
上式即为 DAQ 的真正输出,即
Q
(
x
^
)
Q(\hat{x})
Q(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)
作者给出了
ϕ
\phi
ϕ 公式最终详细的推导过程:
而
β
∗
\beta^*
β∗ 为:
β
∗
\beta^*
β∗ 可以展开为:
所以进一步地有
上式中
λ
=
1
e
x
p
(
γ
)
+
1
\lambda=\frac{1}{exp(\gamma)+1}
λ=exp(γ)+11 。进一步地,可以整理为:
注意
q
c
−
q
f
=
1
q_c-q_f=1
qc−qf=1。故最终就有:
4.2 Overall Process of DAQ
4.3 Design of DAQ
关于距离得分函数和采样核函数,作者也尝试过其他的公式,但是效果都不如原文里的好,故最后选择原文的方法。
4.4 Experiments on the SR Task
作者还在超分辨率任务上做了实验,具体在这里就不展示了。感兴趣的旁友可以翻翻原文。
个人思考
- 有一个细节,作者说自己的方法能避免梯度爆炸,但是随着温度的上升,最后还是会长得像 Sigmoid 函数,那么梯度在过渡点附近不也形成梯度爆炸的情况了吗?所以实质性地说,梯度爆炸情况没有消失,在过渡点的局部值仍会出现梯度爆炸,只是其他大部分的值梯度合理。
- 作者的这个温度不是随时间变化,而是随要量化的原值而变化。和EWGS有异曲同工之妙。