Noise injection and clamping estimation for neural network quantization
文章目录
文章链接
Introduction
主要工作:提出了NICE(noise injection and clamping estimation)量化方式,主要分为以下两个步骤:
1:Noise injection during training that emulates the quantization noise introduced at inference time
2:Statistics-based initialization of parameter and activation clamping, for faster model convergence. In
addition, activation clamp is learned during train time
这种方法的好处在于不需要在训练的时候,修改网络,并且它将网络中的所有参数量化为fixed point(integer)。
Related work
Expressiveness based methods :如使用更大的网络或者加一个线性的尺度因子来弥补精度的损失。
Keeping full-precision copy of quantized weights:使用两份权值,一份是量化后的,一份是full precision的。前向时使用quantized weight,更新时使用full precision weight,如使用STE(straight-through estimator)来估算梯度,使用随机(stochastic)或者确定(deterministic)的方式量化参数。
Distillation:使用distillation。用student network来学习teacher network的输出,还挺有用的。
Model parametrization 、Optimization techniques 、Generalization bounds
Method
文章并不量化网络的第一层和最后一层,因为这两层对网络的性能影响很大。
这种方法的优势在于它在反向更新之后能迅速影响前向,相对于那些直接量化权重的策略而言,它们在反向时几乎不会更新来量化的权重。
为了实现一个类似drop-out的效果,作者使用了Bernoulli分布的一个mask M M M,来量化部分权值并向其他部分添加噪声。经验取 M ∼ B e r ( 0.05 ) M \sim Ber(0.05) M∼Ber(0.05)。并且,在前向时,使用的量化权重为: w ^ = ( 1 − M ) Q Δ ( w ) + M ( w − e ) \hat{w}=(1-M)\mathcal{Q}_{\Delta}(w)+M(w-e) w^=(1−M)QΔ(w)+M(w−e),其中 e ∼ U n i ( − Δ / 2 , Δ / 2 ) e \sim Uni(-\Delta/2, \Delta/2) e∼Uni(−Δ/2,Δ/2), Δ \Delta Δ表示quantization bin的大小。
Gradual Quantization
文章是先量化一部分参数,然后让剩下的参数adapt to the changes。对于梯度的量化,首先将网络的层分为 N N N个相同大小的block { B 1 , … , B N } \{B_1,\ldots,B_N\} {B1,…,BN},然后在第 i i i 阶段(at the i-th stage),将噪声从 B i B_i Bi块注入,它前面的blocks { B 1 , … , B i − 1 } \{B_1,\ldots,B_{i-1}\} {B1,…,Bi−1}都量化掉,后面的blocks B i + 1 , … , B N {B_{i+1},\ldots,B_{N}} Bi+1,…,BN都保持full precision,进行一次梯度更新,当完成 N N N个过程后,量化所有的层并使用STE方式训练。
这样训练就可以使得后面的层可以根据前面量化后的层的输入值的变化进行相应的调整,网络不会一下子变化特别大。之后的使用STE训练是为了使网络收敛,对于一个用来量化的已经训好的网络,作者发现最好的block size为带activation的一个single layer,但是使用上面的方式并没有使效果变得更好。
Clamping and Quantization
为了量化网络的权重,将权值限定在
[
−
c
w
,
c
w
]
[-c_w,c_w]
[−cw,cw]:
w
c
=
C
l
a
m
p
(
w
,
−
c
w
,
c
w
)
=
m
a
x
(
−
c
w
,
m
i
n
(
x
,
c
w
)
)
(1)
w_c=Clamp(w,-c_w,c_w)=max(-c_w,min(x, c_w)) \tag{1}
wc=Clamp(w,−cw,cw)=max(−cw,min(x,cw))(1)式中,对每一层定义一个
c
w
c_w
cw,并将其初始化为
c
w
=
mean
(
w
)
+
β
×
std
(
w
)
c_w=\operatorname{mean}(w)+\beta\times \operatorname{std}(w)
cw=mean(w)+β×std(w),其中
w
w
w是该层的权重,
β
\beta
β是一个超参数,对于一个给定的
c
w
c_w
cw,量化限定后的权重到
B
w
B_w
Bw bits:
w
^
=
[
w
c
2
B
w
−
1
−
1
c
w
]
c
w
2
B
w
−
1
−
1
(2)
\hat{w}=\left[w_{c} \frac{2^{B_{w}-1}-1}{c_{w}}\right] \frac{c_{w}}{2^{B_{w}-1}-1} \tag{2}
w^=[wccw2Bw−1−1]2Bw−1−1cw(2)式中,
[
⋅
]
\left[\cdot\right]
[⋅]表示rounding操作。同样的,激活值的量化也是一同样的方式。如ReLU被替换成clamped ReLU
a
c
=
Clamp
(
a
,
0
,
c
a
)
(3)
a_c=\operatorname{Clamp}(a,0,c_a) \tag{3}
ac=Clamp(a,0,ca)(3)其中,
a
a
a表示某一层的线性部分的输出,
c
a
c_a
ca是限定的范围,
a
c
a_c
ac表示在量化之前截断的一个非负的值,
c
a
c_a
ca是一个集合,每一层都有一个
c
a
c_a
ca,并通过反向传播学习更新。同样的,量化
a
c
a_c
ac到
B
a
B_a
Ba bits的方式也是
a
^
=
[
a
c
2
B
a
−
1
−
1
c
a
]
c
a
2
B
a
−
1
(4)
\hat{a}=\left[a_{c} \frac{2^{B_{a}-1}-1}{c_{a}}\right] \frac{c_{a}}{2^{B_{a}}-1} \tag{4}
a^=[acca2Ba−1−1]2Ba−1ca(4)由于
Round
\operatorname{Round}
Round函数是不可微的,因此使用STE进行反向传播,对于
c
a
c_a
ca的更新,可以计算
a
^
\hat{a}
a^对
c
a
c_a
ca的偏导为:
∂
a
^
∂
a
c
=
{
1
,
a
c
∈
[
0
,
c
a
]
0
,
otherwise
(5)
\frac{\partial \hat{a}}{\partial a_{c}}= \left\{ \begin{array}{ll} {1,} & {a_{c} \in\left[0, c_{a}\right]} \\ {0,} & {\text { otherwise }} \end{array} \right. \tag{5}
∂ac∂a^={1,0,ac∈[0,ca] otherwise (5)对于每一层的biases的量化就更为复杂,因此他们的尺度取决于activations和weights尺度。对于每一层,将截取偏差的值并初始化为:
c
b
=
(
c
a
2
B
a
−
1
⏟
Activation scale
⋅
c
w
2
B
w
−
1
−
1
⏟
Weight scale
)
⋅
(
2
B
b
−
1
−
1
⏟
Maximal bias value
)
(6)
c_{b}=(\underbrace{\frac{c_{a}}{2^{B_{a}}-1}}_{\text {Activation scale }} \cdot \underbrace{\frac{c_{w}}{2^{B_{w}-1}-1}}_{\text {Weight scale }}) \cdot(\underbrace{2^{B_{b}-1}-1}_{\text {Maximal bias value }}) \tag{6}
cb=(Activation scale
2Ba−1ca⋅Weight scale
2Bw−1−1cw)⋅(Maximal bias value
2Bb−1−1)(6)式中,
B
b
B_b
Bb表示bias的bitwidth,bias的截断和量化方式和weights相同。
Experiment Results
使用预训练好的FP-32模型。然后用NICE进行量化,用PyTorch实现,在ImageNet和CIFAR-10上测试分类准确率,在the MSR joint denoising and demosaicing dataset上测试回归。
ImageNet
量化Resnet-18/34/50,使用NICE来fine-tune网络,训了120个epoch,学习率为
1
0
−
4
10^{-4}
10−4,momentum为0.9,weight decay设为
4
×
1
0
−
5
4\times 10^{-5}
4×10−5,结果如下:
Ablation study
为了探究NICE每一部分的重要性,作者在ImageNet上使用了Resnet-18。从结果可以看出,对于high bitwidth,即5.5这种,noise addition和gradual training比clamp learning贡献更大。对于low bitwidth,结果则相反。
文章还对结果做了进一步的阐述,并且说明了在硬件上的加速,不是关注的重点,就不阐述了。