Abstract
我们提出了基于位置的缩放梯度(Position-based Scaled Gradient,简称PSG),该方法根据权重向量的位置对梯度进行缩放,使其更易于压缩。首先,我们从理论上证明了将PSG应用于标准梯度下降(Gradient Descent,简称GD),即PSGD,相当于在变形权重空间中的GD,这个变形权重空间是通过一个适当设计的可逆函数扭曲原始权重空间得到的。其次,我们从实验证明了PSG作为权重向量的正则化器,对于量化和剪枝等模型压缩领域是有利的。PSG减小了全精度模型和其压缩版本之间的权重分布差距。这使得模型可以根据资源的可用性灵活部署,既可以以非压缩模式运行,也可以以压缩模式运行。在CIFAR-10/100和ImageNet数据集上的实验结果显示,所提出的PSG在剪枝和量化领域均表现出了有效性,即使在极低比特情况下也是如此。代码已在Github上发布。
1 Introduction
许多正则化策略已经被提出,用于在神经网络中引入先验知识 [20, 38, 19, 23]。受到这些为特定目的引入先验的正则化方法的启发,本文提出了一种新颖的正则化方法,即为模型压缩问题非均匀缩放梯度的基于位置的缩放梯度(Position-based Scaled Gradient,简称PSG)。这种缩放梯度的缩放比例取决于权重的位置,限制了权重到一组易于压缩的网格点。我们用提出的PSG替换随机梯度下降法(SGD)中的传统梯度,并称之为PSG下降法(PSGD)。我们展示了在原始权重空间中应用PSGD相当于在变形空间中通过标准SGD优化权重,该变形空间通过一个可逆函数将原始权重空间扭曲而成。设计这个可逆扭曲函数,使得通过缩放梯度将原始空间的权重强制合并到期望的目标位置。
我们并不是第一个缩放梯度元素的方法。缩放梯度方法也被称为可变度量方法(variable metric method) [9],它通过将梯度向量与一个正定矩阵相乘来缩放梯度。它包括广泛的方法,如牛顿法、拟牛顿法和自然梯度法 [11, 34, 4]。通常,这些方法依赖于海森矩阵估计或费舍尔信息矩阵进行缩放。然而,我们的方法与它们不同,因为我们的缩放不依赖于损失函数,而只依赖于权重的当前位置。
我们将提出的PSG方法应用于模型压缩问题,如量化和剪枝。近年来,在受限的边缘设备(如智能手机和物联网设备)上部署深度神经网络(DNN)已成为一个非常重要的问题。基于这些原因,模型权重的减少位宽(量化)和删除不重要的模型权重(剪枝)被研究并广泛用于应用。大部分文献中的量化,称为量化感知训练(QAT)方法,通过在低精度域上微调预训练模型来实现,而不考虑全精度域,使用整个训练数据集。此外,这种场景在实际应用中是有限制的,因为需要额外的训练。在额外的训练阶段,需要完整的数据集和高计算资源,这禁止了DNN在需要快速部署的客户的边缘设备上容易和快速地部署。
为了解决这个问题,许多工作集中于后训练量化(PTQ)方法,这些方法不需要全规模训练 [25, 32, 2, 41]。例如,[32] 从一个预训练模型开始,只对权重进行微小的修改,通过均衡通道间的比例和校正偏差。然而,预训练模型的分布与量化模型的分布之间的内在差异过大,使得上面提到的方法无法抵消基本的分布差异。如图1所示,由于两种分布的差异,分类误差和量化误差(表示为均方误差)随着使用的更低比特数的增加而增加。因此,当进行逐层量化时,现有的后训练方法在量化低于6比特时会遭遇显著的准确度下降。
与此同时,量化研究领域最近出现了另一种从初始训练阶段进行量化的方法 [1]。我们的方法遵循从头开始训练的方案,如标准SGD,但我们获得了一个可以轻松量化为低精度模型的全精度模型,而不需要额外的后处理。本质上,我们的主要目标是训练一个在资源有限时可以轻松压缩的压缩友好型模型,无需重新训练、微调甚至访问数据。为此,我们通过按原始权重与其量化版本之间的误差成比例缩放梯度,限制原始权重合并到一组量化网格点(见附录A和图1(b))。对于剪枝,权重被正则化以合并到零。更多细节将在第三节中描述。
我们的贡献可以总结如下:
- 我们提出了一种新的正则化方法用于模型压缩,通过引入基于位置的缩放梯度(PSG),可以将其视为可变度量方法的一种变体。
- 我们从理论上证明,PSG下降(PSGD)等同于在变形权重空间中应用标准梯度下降。这使得权重在压缩和未压缩权重空间中都能收敛到性能良好的局部最小值(见附录A和图1)。
- 我们在量化和剪枝中应用PSG,并在CIFAR和ImageNet数据集上验证了PSG的有效性。我们还表明,PSGD对于极低比特量化非常有效。此外,当与现有的后训练量化方法结合使用时,PSGD预训练模型的性能优于经过SGD预训练的模型。
2 Related work
量化感知训练(QAT)方法在低精度领域表现越来越强,甚至可以达到2到3比特宽度的情况下取得良好的效果 [12, 14, 3, 22]。另一方面,后训练量化的目的是在不进行额外训练或不使用训练数据的情况下进行权重和激活的量化。最近文献中的大多数工作都是从用标准训练方案训练的预训练网络开始的 [41, 32, 2]。许多关于通道级量化方法的工作展示了显著的性能提升,即使在4比特的情况下 [2, 8]。然而,更适合硬件的分层量化方法(在每层存储量化参数而不是每通道)在低比特数情况下仍然存在性能问题 [32, 25, 41]。[32] 通过通道偏置校正和范围均衡,在8比特时达到了接近全精度的准确度,而[41] 通过分裂含有异常值的通道来减少截断误差。但是,它们在低于6比特时都遭遇了严重的精度下降。我们的方法改进了但不仅限于统一层级量化。与我们同时进行的[33]和[31]提出直接使用校准数据集最小化量化误差,以在低于6比特的情况下达到更高性能。我们在第5节中展示了使用PSGD预训练模型优于使用SGD预训练模型的结果。
与此同时,量化的另一条工作线集中在通过从初始训练阶段正则化权重分布以提高量化鲁棒性。[29]集中在最小化Lipshitz常数来正则化梯度,以增强对抗攻击的鲁棒性。类似地,[1]提出对梯度范数施加新的正则化项,以在不同比特宽度下提高量化鲁棒性。这使得能够"随需"量化到各种比特宽度。我们的方法没有显式的正则化项,而是通过缩放梯度来隐式正则化全精度域的权重,使其易于量化。此外,由于不需要梯度范数正则化,因此我们的方法不会引入显著的训练开销,而[1]需要双重反向传播,增加了训练的复杂性。其他一些与量化相关的工作旨在量化梯度向量以提高训练效率 [10],提出更具代表性的编码格式 [37],或学习最佳的混合精度比特宽度 [39]。
剪枝 在模型压缩中,另一个相关的研究领域是剪枝,即修剪不重要的单元,如权重、滤波器或整个块 [21, 28]。最近的工作集中在包括训练阶段修剪过程的剪枝方法 [35, 42, 30, 27]。其中,大量工作利用诱导稀疏性的正则化。[30]提出通过在个别权重上使用L0范数正则化器进行训练,以训练稀疏网络,使用预期的L0目标来放松原本不可微的正则化项。同时,其他工作集中在使用显著性标准。[27]利用掩码梯度作为重要性的代理,在一次操作中修剪网络。与[27]和[30]类似,我们的方法在训练过程中不需要启发式的剪枝计划,也不需要剪枝后的额外微调。在我们的方法中,剪枝被视为量化的一个子类,因为PSG可以通过将目标值设置为零而不是量化网格点,实现稀疏训练。
3 Proposed method
在本节中,我们描述了提出的基于位置的缩放梯度下降法(PSGD)。在PSGD中,一个缩放函数正则化原始权重,使其合并到表现良好的目标点中,这些目标点在未压缩和压缩域中效果都很好。这相当于在变形权重空间中通过SGD进行优化。通过一个特殊设计的可逆函数扭曲原始权重空间,使变形空间中的损失函数收敛到不同的局部极小点,比原始权重空间中驱动的解更加有利于压缩。我们首先证明在原始空间中使用PSGD进行优化等同于在变形空间中使用梯度下降进行优化。然后,我们演示了如何使用PSGD将权重约束到一组目标点。最后,我们解释为什么尽管有强正则化效应,这种方法在原始未压缩域中能够产生可与普通SGD相媲美的性能。
3.1 在变形空间中的优化
定理:设
F
:
X
→
Y
=
R
n
F : X \rightarrow Y = \mathbb{R}^n
F:X→Y=Rn是一个扭曲原始权重空间
X
X
X到
Y
Y
Y的任意可逆多变量函数,并且考虑损失函数
L
:
X
→
R
L : X \rightarrow \mathbb{R}
L:X→R和等效损失函数
L
′
=
L
∘
F
−
1
:
Y
→
R
L' = L \circ F^{-1} : Y \rightarrow \mathbb{R}
L′=L∘F−1:Y→R。那么,在变形空间
Y
Y
Y中的梯度下降法(GD)等价于在原始空间
X
X
X中应用缩放梯度下降法,满足
GD
(
y
,
∂
L
′
∂
y
)
≡
GD
(
x
,
(
J
F
∣
x
)
−
2
∂
L
∂
x
)
\text{GD}\left(y, \frac{\partial L'}{\partial y}\right) \equiv \text{GD}\left(x, (J_F \big|_x)^{-2} \frac{\partial L}{\partial x}\right)
GD(y,∂y∂L′)≡GD(x,(JF
x)−2∂x∂L)
其中
y
=
F
(
x
)
y = F(x)
y=F(x),且
∂
b
∂
a
\frac{\partial b}{\partial a}
∂a∂b和
J
∂
b
/
∂
a
J_{\partial b / \partial a}
J∂b/∂a分别表示函数
b
b
b相对于变量
a
a
a的梯度和雅可比矩阵。
证明:考虑时间点
t
t
t时在
X
X
X中的点
x
t
x_t
xt及其在
Y
Y
Y中的变形版本
y
t
y_t
yt。为了找到
L
(
y
)
L(y)
L(y)的局部最小值,可在变形空间中应用标准梯度下降法,步骤如下:
y
t
+
1
=
y
t
−
η
∂
L
′
∂
y
(
y
t
)
y_{t+1} = y_t - \eta \frac{\partial L'}{\partial y}(y_t)
yt+1=yt−η∂y∂L′(yt)
这里,
∂
L
′
∂
y
(
y
t
)
=
∂
L
∂
y
(
y
t
)
\frac{\partial L'}{\partial y}(y_t) = \frac{\partial L}{\partial y}(y_t)
∂y∂L′(yt)=∂y∂L(yt)是梯度,
η
\eta
η是学习率。应用反函数
F
−
1
F^{-1}
F−1到
y
t
+
1
y_{t+1}
yt+1,我们得到更新点
x
t
+
1
x_{t+1}
xt+1:
x
t
+
1
=
F
−
1
(
y
t
+
1
)
=
F
−
1
(
y
t
−
η
∂
L
∂
y
(
y
t
)
)
=
F
−
1
(
y
t
)
−
η
J
x
∣
y
(
y
t
)
∂
L
∂
y
(
y
t
)
x_{t+1} = F^{-1}(y_{t+1}) = F^{-1}(y_t - \eta \frac{\partial L}{\partial y}(y_t)) = F^{-1}(y_t) - \eta J_{x \big| y}(y_t) \frac{\partial L}{\partial y}(y_t)
xt+1=F−1(yt+1)=F−1(yt−η∂y∂L(yt))=F−1(yt)−ηJx
y(yt)∂y∂L(yt)
这里最后的等式来自于
y
t
y_t
yt附近的一阶泰勒展开,并且
J
x
∣
y
=
∂
F
−
1
∂
y
∈
R
n
×
n
J_{x \big| y}=\frac{\partial F^{-1}}{\partial y}\in \mathbb{R}^{n \times n}
Jx
y=∂y∂F−1∈Rn×n是
x
=
F
−
1
(
y
)
x = F^{-1}(y)
x=F−1(y)相对于
y
y
y的雅可比矩阵。根据链式法则,
∂
L
∂
y
=
∂
L
∂
x
⋅
J
x
∣
y
=
J
x
∣
y
⋅
∂
L
∂
x
\frac{\partial L}{\partial y} = \frac{\partial L}{\partial x} \cdot J_{x \big| y} = J_{x \big| y} \cdot \frac{\partial L}{\partial x}
∂y∂L=∂x∂L⋅Jx
y=Jx
y⋅∂x∂L
因为
J
x
∣
y
=
(
J
y
∣
x
)
−
1
=
(
J
F
∣
x
)
−
1
J_{x \big| y} = (J_{y \big| x})^{-1} = (J_{F \big| x})^{-1}
Jx
y=(Jy
x)−1=(JF
x)−1,我们可以将等式(3)重新表示为:
x
t
+
1
=
x
t
−
η
(
J
F
∣
x
(
x
t
)
)
−
2
∂
L
∂
x
(
x
t
)
x_{t+1} = x_t - \eta (J_{F \big| x}(x_t))^{-2} \frac{\partial L}{\partial x}(x_t)
xt+1=xt−η(JF
x(xt))−2∂x∂L(xt)
现在,等式(2)和等式(4)是等价的,并且等式(1)得以证明。换句话说,原始空间
X
X
X中的缩放梯度下降(PSGD),其缩放由矩阵
(
J
F
∣
x
)
−
2
(J_{F \big| x})^{-2}
(JF
x)−2决定,是等同于在变形空间
Y
Y
Y中的梯度下降。
3.2 基于位置的缩放梯度
在此部分中,我们介绍用于缩放梯度的可逆函数 F ( x ) F(x) F(x)的一个设计示例。这个可逆函数应使原始权重向量 x x x合并到一组期望的目标点 x ∗ x^* x∗上。这些期望的目标权重可以在优化过程中充当先验,以约束原始权重合并到特定位置。如何设置目标点的细节将在下一个小节中讨论。
权重依赖的梯度缩放的要点是简单的。对于一个给定的权重向量,如果特定的权重元素距离所需的目标点较远,则应用较高的缩放值以便更快地逃离该位置。另一方面,如果距离较小,则应用较低的缩放值以防止权重向量偏离该位置。从现在开始,我们重点介绍量化问题的缩放函数设计。剪枝的过程是类似的,我们会省略细节。
缩放函数:我们对每个坐标 x i x_i xi独立使用相同的扭曲函数 f f f,即 y = F ( x ) = [ f ( x 1 ) , f ( x 2 ) , … , f ( x n ) ] T y = F(x) = [f(x_1), f(x_2), \ldots, f(x_n)]^T y=F(x)=[f(x1),f(x2),…,f(xn)]T。因此,雅可比矩阵变得对角化, J F ∣ x = diag ( f ′ ( x 1 ) , f ′ ( x 2 ) , … , f ′ ( x n ) ) J_F \big|_x = \text{diag}(f'(x_1), f'(x_2), \ldots, f'(x_n)) JF x=diag(f′(x1),f′(x2),…,f′(xn)),我们的方法属于对角缩放梯度方法。
考虑以下扭曲函数:
f
(
x
)
=
2
sign
(
x
−
x
∗
)
(
∣
x
−
x
∗
∣
+
ϵ
)
+
c
(
x
)
f(x) = 2\text{sign}(x - x^*)(|x - x^*| + \epsilon) + c(x)
f(x)=2sign(x−x∗)(∣x−x∗∣+ϵ)+c(x)
其中目标
x
∗
x^*
x∗确定为最接近的网格点,
sign
(
x
)
\text{sign}(x)
sign(x)是符号函数,
c
(
x
)
c(x)
c(x)是一个依赖于特定网格点
x
∗
x^*
x∗的常数,使函数连续,
ϵ
\epsilon
ϵ是一个任意小的常数以避免无限梯度。然后,来自公式(4),逐元素缩放函数变为
s
(
x
)
=
1
[
f
′
(
x
)
]
2
s(x) = \frac{1}{[f'(x)]^2}
s(x)=[f′(x)]21,因此
s
(
x
)
=
∣
x
−
x
∗
∣
+
ϵ
1
.
s(x) = \frac{|x - x^*| + \epsilon}{1}.
s(x)=1∣x−x∗∣+ϵ.
使用逐元素缩放函数公式
s
(
x
)
s(x)
s(x),PSG下降(PSGD)的逐元素权重更新规则变为:
x
i
t
+
1
=
x
i
t
−
η
s
(
x
i
)
∂
L
∂
x
i
∣
x
t
x_i^{t+1} = x_i^t - \eta s(x_i) \frac{\partial L}{\partial x_i} \bigg|_{x_t}
xit+1=xit−ηs(xi)∂xi∂L
xt
其中,
η
\eta
η是学习率。
3.3 目标点
量化:本文中,我们使用统一对称量化方法 [25] 和为了硬件友好性的逐层量化方案。考虑模型权重的浮点范围 [min_x, max_x]。权重
x
x
x被量化为对于n比特精度的整数范围 [
−
2
n
−
1
+
1
,
2
n
−
1
−
1
-2^{n-1}+1, 2^{n-1}-1
−2n−1+1,2n−1−1]。权重的量化-反量化定义为步长(
Δ
\Delta
Δ)和截断值。整体量化过程如下:
x
Q
=
Clip
(
⌊
x
Δ
⌋
,
−
2
n
−
1
+
1
,
2
n
−
1
−
1
)
x_Q = \text{Clip}\left( \left\lfloor \frac{x}{\Delta} \right\rfloor, -2^{n-1}+1, 2^{n-1}-1 \right)
xQ=Clip(⌊Δx⌋,−2n−1+1,2n−1−1)
其中
⌊
⋅
⌋
\left\lfloor \cdot \right\rfloor
⌊⋅⌋是四舍五入到最接近的整数运算,
Clip
(
x
,
a
,
b
)
\text{Clip}(x, a, b)
Clip(x,a,b)为:
Clip
(
x
,
a
,
b
)
=
{
b
,
if
x
>
b
a
,
if
x
<
a
x
,
otherwise
\text{Clip}(x, a, b) = \begin{cases} b, & \text{if } x > b \\ a, & \text{if } x < a \\ x, & \text{otherwise} \end{cases}
Clip(x,a,b)=⎩
⎨
⎧b,a,x,if x>bif x<aotherwise
我们可以通过反量化过程得到量化权重
x
=
Δ
x
Q
x = \Delta x_Q
x=ΔxQ,并将其用作量化的目标位置。
剪枝:对于基于幅度的剪枝方法,靠近零的权重将被移除。因此,我们选择零作为目标值(即 x ∗ = 0 x^* = 0 x∗=0)。
3.4 深度网络中的PSGD
许多文献中关于使用随机梯度下降(SGD)优化深度神经网络(DNNs)都报告了多次实验在性能上的一致性,尽管DNNs有许多局部极小值(例如,参见 [6] 的第2部分)。[7] 分析了DNNs的损失表面,表明大规模网络有许多在测试集上性能相似的局部极小值,并且随机损失函数的最低关键值位于全局最小值以下的特定带内。从这一点出发,我们非正式解释深度网络中的PSGD是如何工作的。
如图2所示,我们假设在原始权重空间 X X X中存在许多性能相似的局部极小值(A和B),其中只有一部分(A)接近于某个目标点(0),在压缩域中也表现出高性能。如图2左所示,假设B的收敛区域比A的宽得多,这意味着从随机初始化中输出解B的可能性比A大得多。通过如上所述设计的扭曲函数 F F F(公式5),将原始空间 X X X扭曲到 Y Y Y,使得靠近目标点的区域被扩展,而远离目标的区域被收缩。如果我们在这个变形空间中应用梯度下降,损失函数将更有可能收敛到A。相应地,PSGD在原始空间中更可能输出A而不是B,这有利于压缩。注意 F F F将原始权重空间转换为变形空间 Y Y Y,而不是压缩域。
4 Experiments
本部分中,我们通过实验展示PSGD的有效性
为了验证我们的PSGD方法,我们首先通过将目标点设置为0进行稀疏训练的实验,然后进一步将我们的方法扩展到使用CIFAR [26]和ImageNet ILSVRC 2015 [36]数据集进行量化。我们首先通过与L0正则化 [30] 和SNIP [27] 的对比,展示了基于幅度剪枝的稀疏训练的有效性。[30]对非零模型参数进行惩罚,并共享在训练过程中正则化模型的方案。与我们的方法相似,[27]是一种单次剪枝方法,不需要修剪计划或额外的微调。
量化实验
对于量化,我们将我们的方法与(1)在初始训练阶段采用正则化的方法进行比较 [1, 16, 29]。我们选择了梯度L1范数正则化 [1] 方法和Lipschitz正则化方法 [29, 16] 作为基线,因为它们提出了与我们相似的在训练阶段使用的新正则化技术。请注意,[16] 在权重梯度上添加L2惩罚项,而不同于[1]的L1惩罚。我们还与(2)现有的从预训练模型开始的最先进的逐层后训练量化方法进行比较 [32, 41],以展示在低比特(4位)情况下的改进。有关比较方法的详细信息,请参见第2节。为了验证我们方法的有效性,我们还训练了极低比特(2位、3位)权重的模型。最后,我们展示了在各种网络架构上的实验结果,并将PSG应用于Adam优化器 [24],这些内容详见附录D。
实施细节
我们在所有实验中使用了Pytorch框架。对于表1中的剪枝实验,我们使用了CIFAR-100上的ResNet-32 [18],并遵循了[40]的训练超参数设置。我们使用了[30]发布的官方实现,并为Pytorch框架重新实现了[27]。对于表2和表3的量化实验,我们使用了ResNet-18,并遵循了[1]对CIFAR-10和ImageNet的设置。对于[41],实验中使用了发布的官方实现。所有其他数据均来自原始论文或重新实现。为了公平比较,所有量化实验跟随逐层统一对称量化 [25],并在量化激活时,使用批量归一化参数截断激活范围,如[32]所述,和[1]相同。对于ImageNet实验,PSGD从最后15个训练轮次开始应用,而对于CIFAR实验,PSGD从第一次学习率衰减轮次开始应用。我们为极低比特实验(表4)额外使用了30轮次。此外,我们为每个比特宽度和稀疏度调整了超参数 s s s。我们的搜索标准是确保未压缩模型的性能不下降,类似于[1]。更多细节见附录C。
4.1 剪枝
作为初步实验,我们首先展示了用零作为单一目标点来实现基于PSG的优化的可能性。然后,我们按照[17]的做法,在不同的稀疏性比例下应用基于幅度的剪枝。实验的目的是验证权重是否集中于零,权重在训练完成后并在未进行微调的情况下被修剪一次,并对模型进行评估,与[30]和我们的方法进行比较。对于[27],其在初始化时通过单次剪枝修剪权重,结果显示在单次剪枝中的表现。表1表明我们的方法在各种高稀疏比例下都优于这两种方法。虽然所有三种方法在低稀疏性(<10%)上能够保持准确率,[30]在20%稀疏性时出现了一些准确率下降,并在高稀疏性下严重受损。这与[13]的结果一致,该方法无法在不显著损害模型质量的情况下生成稀疏残差模型。与[27]相比,我们的方法即使在高稀疏性下也能保持更高的准确率,显示出在单次剪枝中的优势,无需剪枝计划或额外训练。图3显示了SGD和PSGD训练的模型中的权重分布。
4.2 量化
在量化领域中,我们首先在动态比特宽度问题上比较了PSGD与正则化方法的表现,即单个模型在不同比特宽度下进行评估。然后,我们与现有的最先进的逐层对称后训练方法进行比较,以验证在低比特宽度下由于权重分布差异所导致的准确率下降问题(见图1)。
正则化方法 表2显示了CIFAR-10和ImageNet数据集上正则化方法的结果。CIFAR-10实验中,我们将激活比特宽度固定为4位,然后将权重比特宽度从8位变化到4位。对于ImageNet实验,表2的设置是权重和激活都使用相同的比特宽度,遵循[1]。CIFAR-10实验中,所有方法似乎都能在4位量化时保持量化模型的性能。无论目标比特宽度是多少,PSGD都优于所有其他正则化方法。另一方面,ImageNet实验通常在6位之前显示出合理结果,但在4位时准确率急剧下降。PSGD针对8位和6位在所有比特宽度上都有边际提升,但也在4位时经历了急剧的准确率下降。相比之下,梯度L1(λ=0.05)和PSGD@W4即使在4位时也能保持量化模型的性能。与第二好的方法梯度L1(λ=0.05)相比,PSGD在所有比特宽度上表现更好。在全精度(FP)、8位、6位和4位时,与[1]相比,我们的方法的性能差距约为4.2%、3.9%、1.5%和8.1%。从表2来看,尽管某些情况下量化噪声可能略微降低准确率,但使用更多比特通常会导致更高的准确率。相比其他正则化方法,PSGD通过将全精度权重的分布约束为量化权重的分布,能在所有比特上保持合理性能。这种量化友好性是通过适当设计的缩放函数实现的。此外,与[1]不同,PSGD不需要额外的计算双重反向传播的开销。
后训练方法 表3显示了最先进的后训练方法OCS在4位时准确率急剧下降。对于OCS,按照原始论文选择了最佳的剪裁方法用于权重和激活。DFQ也有类似的趋势,在6位以下显示急剧的准确率下降,见DFQ原始论文 [32] 图1。这是由于全精度和量化权重分布之间的根本差异,如第1节和图1所述。另一方面,使用PSGD训练的模型因缩放函数具有相似的全精度和量化权重分布,从而导致较低的量化误差。在没有任何后训练和权重剪裁来处理异常值的情况下,我们的方法在4位时比OCS高出了大约19%。将PTQ方法应用于我们预训练的PSG模型将在第5节中展示。
极低比特量化
如图1所示,SGD在极低比特(如3位和2位)时表现出急剧的准确率下降。为了确认PSGD能够处理极低比特问题,我们进行了实验,PSGD目标是3位和2位,除了第一个和最后一个层被量化为8位。表4显示了应用PSGD的结果。尽管由于强约束,全精度准确率有所下降,但PSGD能够保持合理的准确率。这表明PSGD在极低比特后训练量化中的潜力,作为一种关键解决方案。
5 Discussion
在此部分中,我们首先通过一个玩具例子聚焦于PSG找到的局部最小值,以获得更深的理解。在这个玩具例子中,我们在MNIST数据集上用SGD和PSGD进行2位量化训练,使用的网络是具有两个隐藏层(50个和20个神经元)的全连接网络。在随后的部分中,我们将澄清PSGD和QAT的不同目的并分析为什么PSGD在低精度下无法达到接近全精度的性能。最后,我们通过将PSG应用于一种现有的PTQ方法来展示PSG的潜在应用。
量化和稀疏模型
通常,SGD会产生一种钟形分布的权重,这种分布并不适合低位量化 [41]。另一方面,PSGD总是提供多峰分布,且峰值位于量化值。在这个例子中,我们使用了三个目标点(2位量),因此权重合并为如图4a所示的三个模式。大部分权重接近于零。同样,我们注意到表2中显示的ResNet-18@W4的稀疏度在低精度时为72.4%。这是因为对称量化也包含零作为目标点。PSGD在W2A32时的准确率几乎与全精度相同(约96%)。然而,SGD在W2A32时的准确率约为9%,尽管全精度准确率为97%。这一趋势也在图1b中显示,表明PSGD减少了量化误差。
PSGD解决方案的曲率
在第3.4节和图2中,我们声称PSG找到的最小值处于尖锐的谷底,更易于压缩,但不易找到。由于Hessian特征向量方向上的曲率由相应的特征值决定 [15],我们通过评估特征值的大小来比较SGD和PSGD所得解的曲率,类似于[5]。SGD提供的最小值谷底相对宽阔,因为它有许多近似为零的特征值,这种趋势在[5]中也观察到。然而,PSGD训练的权重有更多的大正特征值,这意味着相比于SGD,解位于相对尖锐的谷底。具体来说,PSGD中的大特征值(> 10^3)的数量是SGD的9倍。从这个玩具例子中,我们确认PSG有助于找到更易于压缩的最小值(图4a),且位于难以通过简单SGD到达的尖锐谷底(图4b)。
量化感知训练(QAT)与PSGD
传统的QAT方法 [12, 14] 从最初用SGD训练的预训练模型开始,仅考虑低精度权重进一步更新权重。相反,正则化方法如我们的工作和[1]从头开始训练全精度权重,类似于SGD。在我们的工作中,PSGD的唯一目的是找到一组适合量化的全精度权重,这样可以在低精度(LP)下进行多种应用而无需进一步操作。因此,与SGD类似,正则化方法从初始训练阶段开始,而QAT方法从SGD和PSGD等初始训练阶段后的预训练模型开始。QAT方法的目的仅专注于低精度权重。通常使用粗略梯度通过前馈低精度权重来更新权重,而不是使用真实的全精度权重,并借助直通估计器(STE) [25] 。此外,量化方案被修改为包含依赖于低精度权重和激活的可训练参数。因此,QAT不能保持全精度的性能,因为它只关注低精度(如4位)下的性能。
具有PSGD训练模型的后训练
我们的模型在全精度下获得了与SGD相近的性能,即使在简单量化的低精度下也能获得合理的性能。因此,PSGD训练的模型可以作为QAT或PTQ方法的预训练模型。我们针对表3中的PSGD训练模型使用了一种现有的PTQ方法LAPQ [33]进行了额外实验,使用了官方代码。在W4A4下获得了66.5%的准确率,比仅使用PSGD和仅使用LAPQ分别高出3.1%和6.2个百分点。这表明PTQ方法可以从使用我们的预训练模型中受益。
6 Conclusion
在这项工作中,我们引入了基于位置缩放的梯度(PSG),该方法根据当前权重与相应目标点之间的距离来缩放梯度。我们证明了随机PSG下降(PSGD)等同于在扭曲空间中应用SGD。基于DNN在测试集上的性能相似的许多局部最小值这一假设,PSGD能够找到其他优化器难以达到的,易于压缩的最小值。因为PSGD减少了量化误差,弥合了压缩权重和未压缩权重分布之间的差距,所以PSGD可以成为低比特后训练量化的关键解决方案。由于目标点作为先验约束,将原始权重合并到特定位置,PSGD也可以通过简单地将目标点设为零来用于稀疏训练。在我们的实验中,我们通过在CIFAR-10/100和ImageNet等各种图像分类数据集上展示PSGD的有效性来验证PSGD在剪枝和量化领域的应用。我们还通过实验证明,PSGD找到的最小值位于比SGD更尖锐的谷底。我们相信PSGD将有助于进一步研究模型量化和剪枝。
7
PSG是一种根据权重向量的位置不同而对每个梯度分量进行不同缩放的基本方法。这项技术可以替代任何应用中需要对参数空间中特定位置进行区别处理的传统梯度。正如本文所示,最容易想到的应用是量化和剪枝,这些领域对特定形式的权重有明确的偏好。这些模型压缩技术是任何深度学习算法快速和轻量级部署的核心,因此,PSG可以对相关行业产生巨大影响。作为另一个潜在的相关研究主题,PSG有可能在优化领域中得到利用,例如整数量化和组合优化,作为在离散空间中优化目标函数的连续替代目标的工具。