1. WAGE背景介绍
WAGE(Training and Inference with Integers in Deep Neural Networks)是ICLR2018的Best Paper,这篇论文提出了一种量化神经网络的框架,可以将神经网络中的参数(Weight)、激活值(Activation)、梯度(Gradient)以及误差(Error)进行可选比特长度的整型量化。在本文中使用 W i W^i Wi、 a i a^i ai、 g i g^i gi、 e i e^i ei表示第 i i i层的参数、激活值、梯度以及误差,用下标 ⋅ q \cdot _q ⋅q来表示已进行量化的值。在WAGE框架中还需要定义好对4种值的量化函数 Q W ( ⋅ ) Q_W(\cdot) QW(⋅)、 Q a ( ⋅ ) Q_a(\cdot) Qa(⋅)、 Q g ( ⋅ ) Q_g(\cdot) Qg(⋅)和 Q e ( ⋅ ) Q_e(\cdot) Qe(⋅),用 [ ⋅ ] [\cdot] [⋅]表示量化之后的整型值在计算机中存储所需要的比特长度。
WAGE的前向传播(推理)过程如下式(1.1):
a
q
i
=
Q
a
(
Q
W
(
W
i
)
⋅
a
q
i
−
1
)
(1.1)
a^{i}_{q} = Q_a(Q_W(W_i)\cdot a^{i-1}_{q}) \tag{1.1}
aqi=Qa(QW(Wi)⋅aqi−1)(1.1)
WAGE的反向传播(训练)过程如下式(1.2):
e
q
i
=
Q
e
(
∂
L
∂
a
i
)
,
g
q
i
=
Q
g
(
∂
L
∂
W
i
)
(1.2)
e^{i}_{q} = Q_e(\frac{\partial \mathcal{L}}{\partial a^i}), \quad g^{i}_{q} = Q_g(\frac{\partial \mathcal{L}}{\partial W^i}) \tag{1.2}
eqi=Qe(∂ai∂L),gqi=Qg(∂Wi∂L)(1.2)
2. 移位线性映射和随机舍入
在WAGE中使用线性映射将浮点数映射到
k
k
k比特有符号整型值,那么整型值之间的最小不可分距离
σ
\sigma
σ可以从式(2.1)中计算出:
σ
(
k
)
=
2
1
−
k
,
k
∈
N
+
(2.1)
\sigma(k) = 2 ^ {1-k}, k \in \mathbb{N}_+ \tag{2.1}
σ(k)=21−k,k∈N+(2.1)
那么WAGE中的将浮点数
x
x
x量化为
k
k
k比特的有符号整型值的截断量化函数被定义为:
Q
(
x
,
k
)
=
C
l
i
p
{
σ
(
k
)
⋅
r
o
u
n
d
(
[
x
σ
(
k
)
]
)
,
−
1
+
σ
(
k
)
,
1
−
σ
(
k
)
}
(2.2)
Q(x, k) = Clip \{ \sigma(k) \cdot round([\frac{x}{\sigma(k)}]), -1+\sigma(k), 1-\sigma(k) \} \tag{2.2}
Q(x,k)=Clip{σ(k)⋅round([σ(k)x]),−1+σ(k),1−σ(k)}(2.2)
C
l
i
p
Clip
Clip函数将整型值限制在区间
[
−
1
+
σ
(
k
)
,
1
−
σ
(
k
)
]
[-1+\sigma(k), 1-\sigma(k)]
[−1+σ(k),1−σ(k)]中,去掉
−
1
-1
−1以保持对称性。在对浮点值进行量化之前会进行一个缩放(移位)运算,否则经过式(2.2)的处理,浮点值大部分会被饱和截断。缩放函数在本文被定义为
S
h
i
f
t
Shift
Shift操作:
S
h
i
f
t
(
x
)
=
2
r
o
u
n
d
(
l
o
g
2
x
)
(2.3)
Shift(x) = 2 ^ {round(log_2x)} \tag{2.3}
Shift(x)=2round(log2x)(2.3)
最终WAGE会对反向传播中的梯度值进行随机舍入运算,以确保梯度能够有效地传播来进行训练。
Q
g
(
⋅
)
Q_g(\cdot)
Qg(⋅)算子会通过一个16比特的随机数产生器来约束量化后的
k
g
k_g
kg比特整型值的随机性,在接下来的章节中会详细介绍四种量化算子的实现。
3. 参数初始化
Batch Normalization能够预防神经网络训练中的梯度爆炸和梯度消失问题,并缓解网络对参数初始化的需求,但是在整型运算中实现Batch Normalization十分复杂。由于种种原因,论文作者改进了何恺明提出的参数初始化方法:
W
∼
U
(
−
L
,
+
L
)
,
其
中
L
=
m
a
x
(
6
/
n
i
n
,
L
m
i
n
)
,
L
m
i
n
=
β
σ
(3.1)
W \sim U (-L, +L), \quad 其中 \quad L = max(\sqrt{6/n_{in}}, L_{min}), L_{min} = \beta \sigma \tag{3.1}
W∼U(−L,+L),其中L=max(6/nin,Lmin),Lmin=βσ(3.1)
n
i
n
n_{in}
nin是输入层的神经元数量,
L
m
i
n
L_{min}
Lmin约束均匀分布能够覆盖的最小区间(至少要大于
σ
\sigma
σ,因此
β
>
1
\beta > 1
β>1)。
4. 量化函数
4.1 参数量化函数 Q W ( ⋅ ) Q_W(\cdot) QW(⋅)
参数量化遵循式(2.2),并使用第3节中的初始化方法:
W
q
=
Q
W
(
W
)
=
Q
(
W
,
k
W
)
(4.1)
W_q = Q_W(W) = Q(W, k_W) \tag{4.1}
Wq=QW(W)=Q(W,kW)(4.1)
经过量化之后参数的方差被缩放了,这有可能导致训练过程中的梯度爆炸。为了缓解被放大的影响,论文作者提出在层与层之间增加一个缩放参数
α
\alpha
α:
α
=
m
a
x
{
S
h
i
f
t
(
L
m
i
n
/
L
)
,
1
}
(4.2)
\alpha = max \{ Shift(L_{min} / L), 1 \} \tag{4.2}
α=max{Shift(Lmin/L),1}(4.2)
4.2 激活值量化函数 Q W ( ⋅ ) Q_W(\cdot) QW(⋅)
作者在论文中不建议在卷积层之后使用Average Pooling,因为均值计算对精度要求更高(相比于Max Pooling)。论文中假设输出的Batch均值为0,Batch Normalization退化为一个缩放运算,在WAGE中可训练的缩放算子被式(4.2)中的
α
\alpha
α替代了,对激活值的量化方法:
a
q
=
Q
a
(
a
)
=
Q
(
a
/
α
,
k
a
)
(4.3)
a_q = Q_a(a) = Q(a / \alpha, k_a) \tag{4.3}
aq=Qa(a)=Q(a/α,ka)(4.3)
4.3 误差量化函数 Q e ( ⋅ ) Q_e(\cdot) Qe(⋅)
相较于激活值,误差的范围更大。论文中认为在训练时梯度的方向指导了整个网络的收敛,而梯度的大小作用不大。那么这样的话,可以直接使用整型值来表达梯度的方向来进行网络的训练。通过除以一个移位因子,WAGE将误差缩放到范围
[
−
2
,
+
2
]
[-\sqrt{2}, +\sqrt{2}]
[−2,+2]之内,量化函数如下:
e
q
=
Q
e
(
e
)
=
Q
(
e
/
S
h
i
f
t
(
m
a
x
(
∣
e
∣
)
)
,
k
e
)
(4.4)
e_q = Q_e(e) = Q(e / Shift(max(|e|)), k_e) \tag{4.4}
eq=Qe(e)=Q(e/Shift(max(∣e∣)),ke)(4.4)
4.4 梯度量化函数 Q g ( ⋅ ) Q_g(\cdot) Qg(⋅)
g = ∂ L ∂ W = ∂ L ∂ a ⋅ ∂ a ∂ W = e ⋅ ∂ a ∂ W (4.5) g = \frac{\partial \mathcal{L}}{\partial W} = \frac{\partial \mathcal{L}}{\partial a} \cdot \frac{\partial a}{\partial W} = e \cdot \frac{\partial a}{\partial W} \tag{4.5} g=∂W∂L=∂a∂L⋅∂W∂a=e⋅∂W∂a(4.5)
首先按照梯度方向决定训练的思路,WAGE对梯度进行了缩放,其中
η
\eta
η(
η
=
2
j
,
j
∈
N
\eta = 2^j, j \in \mathbb{N}
η=2j,j∈N)为学习率:
g
s
=
η
⋅
g
/
S
h
i
f
t
(
m
a
x
{
∣
g
∣
}
)
(4.6)
g_s = \eta \cdot g / Shift(max\{ |g| \}) \tag{4.6}
gs=η⋅g/Shift(max{∣g∣})(4.6)
g
s
g_s
gs表征了训练的最小步长,以及参数更新的方向。如果参数占用了
k
g
k_g
kg比特的空间,那么最小步长所对应的整型数和浮点数分别是
±
1
\pm1
±1和
±
σ
(
k
g
)
\pm \sigma(k_g)
±σ(kg)。在训练的开始阶段,设置学习率
η
≥
2
\eta \ge 2
η≥2使
g
s
>
1
g_s \gt 1
gs>1来加速训练,而在训练的后半段随着学习率的衰减,梯度也会变小。在参数的更新过程中,WAGE将
g
s
g_s
gs分解为整型部分和小数部分,然后使用16比特随机数生成器来约束浮点数随机映射到整型数:
Δ
W
=
Q
g
(
g
)
=
σ
(
k
g
)
⋅
s
g
n
(
g
s
)
⋅
{
⌊
∣
g
s
∣
⌋
+
B
e
r
n
o
u
l
i
(
∣
g
s
∣
−
⌊
∣
g
s
∣
⌋
)
}
(4.7)
\Delta W = Q_g(g) = \sigma(k_g) \cdot sgn(g_s) \cdot \{ \lfloor |g_s| \rfloor + Bernouli( |g_s| - \lfloor |g_s| \rfloor) \} \tag{4.7}
ΔW=Qg(g)=σ(kg)⋅sgn(gs)⋅{⌊∣gs∣⌋+Bernouli(∣gs∣−⌊∣gs∣⌋)}(4.7)
在更新参数之后仍然会有参数值超越范围
[
−
1
+
σ
,
1
−
σ
]
[-1+\sigma, 1-\sigma]
[−1+σ,1−σ]的情况存在,因此要使用
C
l
i
p
Clip
Clip函数来使参数值饱和,经过第
t
t
t轮的训练,可以得到:
W
t
+
1
=
C
l
i
p
{
W
t
−
Δ
W
t
,
−
1
+
σ
(
k
g
)
,
1
−
σ
(
k
g
)
}
(4.8)
W_{t+1} = Clip \{ W_t - \Delta W_t , -1 + \sigma(k_g) , 1 - \sigma(k_g) \} \tag{4.8}
Wt+1=Clip{Wt−ΔWt,−1+σ(kg),1−σ(kg)}(4.8)