论文(ICCV 2019):https://arxiv.org/abs/1908.05033
源码(Pytorch Unofficial):https://github.com/ricky40403/DSQ
DSQ 概要
1 Motivation
- 由于量化电平是有限的,当模型量化到极低比特(如 2-bit)时,目前的二值量化和均匀量化仍然会导致非常严重的性能下降。
- QAT 的反向传播很难获得准确的梯度,在文献中基本都用直通估计器(STE)来近似。但当数据被量化到极低比特时,梯度误差会被放大,导致优化明显不稳定。
- 量化本身不可避免地会导致原始数据与其量化值之间出现较大偏差,从而常常导致性能下降。量化通常由两个操作完成:剪裁(Clip)和取整(Round)。前者将数据限制在较小的范围内,而后者将原始值映射到最近的量化点。这两种操作都会导致量化损失。因此,为了缓解性能下降,找到合适的剪裁范围并在剪裁和舍入之间取得平衡也很重要。
2 Method
整个文中,作者在标准 1-bit 二值化和多 bit 均匀量化的语境下展开讨论。总的来说,作者的设计思路如 Fig 1 所示。
2.1 Preliminaries
对于 1-bit 二值量化,Activation 和 Weight 会限制为 -1 或 +1,具体如下公式:
而对于多比特均匀量化,设位宽为
b
b
b ,限制在
[
l
,
u
]
[l,~u]
[l, u] 范围内的浮点 Activation/Weight 为
x
x
x,那么一次完整的量化-反量化过程可以定义为:
其中范围
[
l
,
u
]
[l,~u]
[l, u] 被划分成
2
b
−
1
2^b-1
2b−1 个间隔
P
i
,
i
∈
(
0
,
1
,
.
.
.
,
2
b
−
1
)
\mathcal{P}_i,~i\in(0, 1,...,2^b-1)
Pi, i∈(0,1,...,2b−1) ,
Δ
=
u
−
l
2
b
−
1
\Delta=\frac{u-l}{2^b-1}
Δ=2b−1u−l 为间隔的宽度。
2.2 Quantization function
为了解决量化公式的微分问题,作者提出一个可微渐近函数来处理落在不同间隔 P i \mathcal{P}_i Pi 的 x x x 点:
其中缩放因子 s s s 保证邻近间隔的 t a n h tanh tanh 函数之间能平滑地连接到一起,具体如 Fig 2 (a) 所示。
由于
t
a
n
h
tanh
tanh 函数的高度对称性,
φ
\varphi
φ 是处处可微的。而系数
k
k
k 决定了该渐近函数的形状——换句话说,
k
k
k 越大,该渐近函数就长得越像均匀量化器的阶梯函数,即如 Fig 1 所示,渐近曲线越像均匀量化器的标准分段曲线。
进一步地,DSQ的均匀量化过程就可以写成:
上式用于模拟前向过程中真实量化的影响,其也在计算反向传播过程中的梯度时表现良好,也就是说,梯度可以更好地反映更新的正确方向。值得注意的是,如 Fig 2 (b) 所示,当只有一个间隔的时候,该公式也能起到很好拟合效果。
2.3 Evolution to the standard quantization
然而,DSQ 对标准量化的逼近程度在很大程度上会影响量化后的模型的行为。要求在训练过程中能够自适应地选择合适的 DSQ 估计参数,从而根据量化网络的优化目标保证受控逼近。Fig 3 (a) 展示了 DSQ 在一个间隔里的没有缩放到
[
−
1
,
+
1
]
[-1, +1]
[−1,+1] 范围内的映射曲线。
很容易证明当曲线的最大点到上界 +1 的距离足够小时,DSQ 函数可以完美地逼近标准量化器。 鉴于此,作者又提出了一个特征变量 α \alpha α:
于是缩放因子
s
s
s ,又可以重写为:
而因为
Δ
=
u
−
l
2
b
−
1
\Delta=\frac{u-l}{2^b-1}
Δ=2b−1u−l 和
φ
(
0.5
Δ
)
=
1
\varphi(0.5\Delta)=1
φ(0.5Δ)=1,所以系数
k
k
k 又可以重写为:
Fig 3(b) 阐述了
α
\alpha
α 对二值化模型的影响。其中曲线分别代表纯粹 DSQ、DSQ + 接着使用 sign 转换为标准量化、标准量化。可以看到,当
α
\alpha
α 很小时,DSQ 可以很好地逼近均匀量化的性能。适当的
α
\alpha
α 会促进 DSQ 训练高准确度的量化模型。
于是作者将
α
\alpha
α 自适应化,在网络中优化它,并在训练期间迫使 DSQ 演变为标准量化器。形式上,作者将其表述为关于各层输入
x
x
x 和输出
y
y
y 的网络损耗最小化问题:
其中对
α
\alpha
α 使用 L2 正则化。可以进一步计算出
α
\alpha
α 的梯度:
2.4 Balancing clipping error and rounding error
Clip 和 Round 操作都会导致量化误差。作者将 Clip 的上界和下界设为可学习参数,联合优化二者以追求剪裁误差和舍入误差之间的平衡。 有对应的梯度表示为:
观察 公式 (11) 和 (12) 可知,过大的离群点对
u
u
u 的更新起主要贡献作用,且被
u
u
u 截断;小的离群点对
l
l
l 的更新起主要贡献作用,且被
u
u
u 截断;而落在合理范围内的数据点同时影响
u
u
u 和
l
l
l 。当 Clip 带来的误差主导整个量化误差时,离群点产生的梯度会很大,从而成为权重更新的主要力量。否则,当舍入误差主导误差时,范围内的点在反向传播过程中会影响更多。
综上,可以发现 DSQ 有三个关键的可学习参数: α \alpha α, l l l 和 u u u。
2.5 Training and Deploying
Algorithm 1 中列出了卷积网络的详细微调过程。
为了部署在计算资源有限的设备上,作者还实现了低比特计算内核,以加速 ARM 架构上的推理。在卷积网络中,乘法和累加是通用矩阵乘法 (GEMM) 的核心运算,其可以通过 ARM 上的 MLA 指令高效地完成。 MLA 将 8 位寄存器中的两个数字相乘,并将结果累加到 8 位寄存器中。 若累加器会溢出,作者就会通过 SADD 指令将值传输到16位寄存器。图4显示了 GEMM 内核的完整数据流。
实际上比特数越低并不代表加速效果越好,且SADD 操作会比 MLA 需要额外的计算成本,但作者巧妙地降低了数据传输发生的概率,实现了有效的加速。 给定两个 b-bit 有符号数,最多可以调用 MLA 指令
2
7
−
1
(
−
2
b
−
1
)
2
\frac{2^7-1}{(-2^{b-1})^2}
(−2b−1)227−1 次,之后再通过 SADDW 转换至 16-bit 寄存器。Table 1 列出了调用两个指令的频率之比。
也就是说,使用较低的量化比特可以获得更好的加速性能。
3 Experiment
Fig 5 展示的是 ResNet-20 在 CIFAR-10 上使用 DSQ 前、后的可视化效果。
Fig 6 展示的是训练过程中的收敛效果。
Fig 7 是训练过程中形变参数
α
\alpha
α 的变化过程。
Table 2 为某一次训练完毕后各个层的形变参数
α
\alpha
α 取值。可以发现 Weight 对应的形变参数
α
\alpha
α 比 Activation 的较小,表明 Activation 对量化更敏感;不同层对量化有不同的敏感程度。
Table 3 是对 1-bit 量化下 DSQ 机制做的消融实验;Table 4 是对各个可学习参数做的消融实验;Table 5 表明 DSQ 在某些情况下是可以和其他方法组合使用的。
Table 6 和 Table 7 是和 SOTA 做的比较。
Table 8 展示的是作者利用自己设计的 GEMM 核的加速效果。Table 9 中作者将自己的 GEMM 核的加速效果与 NCNN 做对比。
个人思考
一个困惑的地方,没看明白:
个人觉得意思应该是每一次 DSQ 量化完后再过一遍
s
i
g
n
sign
sign 函数来,达到真正有限离散化的效果。既然这样,作者在训练过程中也用到了 STE,因为 sign() 函数大部分地方给梯度的影响为零,少部分地方甚至不可导。
还困惑的是,Piecewise Quantization 和 Binary Quantization 的区别是什么,为什么会在 Table 3 中作比较。个人初步的理解是,Piecewise Quantization 可以是非对称量化,比如量化电平为 -0.9 和 0.8,到原点距离不一致;而 Binary Quantization 只能是对称量化?