深度神经网络已经成为解决计算机视觉、语音识别和自然语言处理等机器学习任务的最先进的技术。尽管如此,深度学习算法是计算密集型和存储密集型的,这使得它难以被部署到只有有限硬件资源的嵌入式系统上。
为了解决这个限制,可以使用深度压缩来显著地减少神经网络所需要的计算和存储需求。例如对于具有全连接层的卷积神经网络(如Alexnet和VGGnet),深度压缩可以将模型大小减少35到49倍。即使对于全卷积神经网络(如GoogleNet和SqueezeNet),深度压缩也可以将模型大小减少10倍。而且上述两种压缩情况都不会降低模型预测的精度。
当前的训练方法有不足之处
压缩模型而不丢失其精确度意味着在训练好的模型中有严重的冗余,这说明当前的训练方法有不足之处。为了解决这个问题,我和来自NVIDIA的JeffPool、百度的Sharan Narang和Facebook的Peter Vajda合作开发了“密集-稀疏-密集”(DSD)的训练方法。这是一种新的方法,它首先通过稀疏约束的优化方法将模型正则化,然后通过恢复和重新训练被剪枝的连接的权重来提高预测精度。在测试时,由DSD训练得到的最终模型仍然跟原始密集型模型具有相同的架构和维度,并且DSD训练不会增加任何推理开销。我们对主流的神经网络(如CNN / RNN / LSTM)架构用DSD训练方法进行了图像分类、图像描述和语音识别的实验,发现模型有显著的性能改进。
在本文中,我们会首先介绍深度压缩,然后介绍“密集-稀疏-密集”(DSD)训练方法。
深度压缩
深度压缩的第一步是“突触剪枝”。 人类大脑是有这一剪枝过程的。从婴儿时期到成年,人脑会有5成的突触会被修剪掉。
类似的规则是否适用于人工神经网络呢?答案是肯定的。在早期的工作中,网络剪枝已经被证明是一种减少网络复杂度和过度拟合的有效方法。这种方法也适用于现代神经网络。首先我们通过常规神经网络训练来学习网络连接权重。然后我们会剪枝权重值较小的连接:即删除网络中权重值低于某一阈值的所有连接。最后,我们重新训练网络,得到剩余稀疏连接的权重值。剪枝方法使AlexNet和VGG-16模型的参数数量分别减少了9倍和13倍。
图1. 剪枝一个神经网络。所有图片由Song Han 友情提供
深度压缩的下一步是权重共享。我们发现神经网络对低精度权重值具有非常高的容忍度:极度粗略的权重值并不会降低预测精度。如图2所示,蓝色权重值最初为2.09、2.12、1.92和1.87,然后让它们共享相同的权重值2,网络的预测精确度仍然不受影响。因此我们可以只存储非常少量的权重值,称之为“编码本”。并让许多其他权重共享这些相同的权重值,且只在码本中存储其索引即可。
索引可以用非常少的比特数来表示。例如在下图中存在四种颜色,因此仅需要两位来表示一个权重而不用原来的32位。另一方面,编码本占用的存储空间几乎可以忽略不计。我们的实验发现,就权衡压缩比和精度而言,这种权重共享技术是优于线性量化的方法的。
图2. 训练权重共享的神经网络
图3显示了使用深度压缩的总体结果。Lenet-300-100和Lenet-5是在MNIST数据集上评估的,而AlexNet、VGGNet、GoogleNet和SqueezeNet是在ImageNet数据集上评估的。压缩比从10倍到49倍不等。即使对于那些全卷积神经网络(如GoogleNet和SqueezeNet),深度压缩仍然可以将它们压缩一个数量级。我们重点看一下SqueezeNet,它比有相同预测精度的AlexNet少50倍的参数,但仍然还可以再压缩10倍使其模型大小只有470KB,这使它可以很容易地在片上SRAM里使用。而访问SRAM比DRAM更快更节能。
我们还尝试了其他压缩方法,例如基于低秩近似的方法,但是压缩比没有那么高。你可以在Deep Compression的论文中找到完整的讨论。
图3 深度压缩的实验结果
DSD 训练
深度神经网络可以被大量剪枝和压缩的事实意味着我们当前的训练方法具有一些局限性。它不能充分利用密集模型的全部容量来找到最佳局部最小值,而一个剪枝过的有着更少神经突触的稀疏模型也可以达到相同的精度。这带来了一个问题:我们是否可以通过恢复并重新学习这些权重来达到更好的精度吗?
让我们拿奥运会田径比赛的训练做个比喻。教练首先会让跑步运动员在高海拔的山地上训练,那里会有很多的限制:低氧、寒冷等。结果当跑步运动员再次返回平原地区时,他的速度就会有提高。对于神经网络来说也是相同的道理:给定严格约束的稀疏训练方法得到的网络模型有跟密集网络模型一样的性能。一旦你解除了这些约束,模型可以工作得更好。
理论上,以下是DSD训练能行之有效的因素:
1.避开鞍点:优化深度网络最大的困难之一是鞍点的扩散。 DSD训练方法通过剪枝和重新密集化的框架来避开鞍点。对收敛的模型做剪枝干扰了模型学习的模式并使得网络模型能够避开鞍点,这使模型有机会可以收敛到一个更好的局部或全局最小值。这个想法也跟“模拟退火”算法类似。虽然模拟退火算法在搜索图上随着概率降低而随机跳跃,但是DSD训练方法会确定性的偏离收敛值。这一收敛值是通过在第一次密集模型训练阶段中去除小权重值和增强稀疏支持得到的。
2.正规化和稀疏训练:稀疏训练步骤中的稀疏正规化将模型优化降维到较低维空间,在这个空间中的损失函数表面更平滑并且对噪声更鲁棒。很多数值实验证实了稀疏训练和最终DSD方法可以降低方差和减少误差。
3. 强大的重新初始化:权重初始化在深度学习中起着重要的作用。传统的训练方法只有一次初始化的机会,而DSD方法在训练过程中给模型优化第二次(或更多)的机会。它基于更鲁棒的稀疏训练结果来重新初始化。我们基于稀疏模型结果来重新构建密集网络,这可以理解为剪枝权重的零初始化。其它的初始化方法也值得尝试。
4. 打破对称性:隐藏单元的置换对称性会使权重对称,因此在训练中容易相互影响。在DSD方法中,权重剪枝打破了与权重相关的隐藏单元的对称性,而且在最终的密集模型中是不对称的。
我们在几个主流的CNN/RNN/LSTM模型上进行了图像分类、图像描述和语音识别数据集的实验,发现这种“密集-稀疏-密集”的训练流程能够显著地提高模型精度。我们的DSD训练方法采用了三个步骤:密集,稀疏,密集。图4展示了每个步骤。
图4 密集-稀疏-密集的训练流程
1.初始的密集训练:第一个“密集”步骤通过在密集网络上的常规网络训练来学习连接权重。然而跟传统训练不同的是,该“密集”步骤的目的不是学习连接权重的最终值,而是学习哪些连接是重要的。
2. 稀疏训练:“稀疏”步骤会剪枝掉权重值较低的连接并重新训练稀疏网络。我们对试验中的所有层都使用相同的稀疏度,因此会有一个单一的超参数:稀疏度(sparsity)。我们对每一层的参数进行排序,从网络中去掉最小的N* sparsity个稀疏参数,将密集网络转换为稀疏网络。我们发现稀疏比率为50%-70%的效果非常好。然后我们重新训练稀疏网络,这可以在稀疏约束下完全复原模型的精度。
3. 最终的密集训练:最后的“密集”步骤会恢复已剪枝的连接,使网络再次变的密集。这些之前剪枝的连接会初始化为零并重新训练。恢复被剪枝的连接增加了网络的维度,并且更多的参数更容易使网络的鞍点向下滑动以获得更好的局部最小值。
我们将DSD训练方法应用于各种类型的神经网络和不同领域的数据集。 我们发现DSD训练方法提高了所有这些神经网络的精确度。 神经网络选自CNN、RNN和LSTM;数据集是从图像分类、语音识别和图像描述领域中选择的, 结果如图5所示。DSD模型可以在DSD Model Zoo上下载。
图 5. DSD训练提高了预测精度
生成图像描述
我们把在图像描述任务中使用DSD训练方法的效果做了可视化(见图6)。我们把DSD训练方法应用于NeuralTalk中,NeuralTalk是一个用于生成图像自然语言描述的长短时记忆模型(LSTM)。基准模型不能很好的描述图片1、4和5。例如图片1,基准模型将女孩错误地描述为男孩,并且把女孩的头发错误地描述为岩墙。稀疏模型可以在图片中识别出有一个女孩,而DSD模型可以进一步识别出秋千。
在第二张图片中,DSD训练方法可以识别出运动员正在尝试投篮,而基准模型只能描述成运动员正在打球。值得注意的是稀疏模型有时比DSD 模型表现得更好。在最后一张图片中,稀疏模型能够正确地识别出泥潭,而DSD模型只能从背景中识别出森林。DSD训练方法的优良性能不仅限于这些例子,此论文的附录中提供了更多的由DSD训练方法生成的图像描述结果。
图6. DSD训练方法提高图像描述的性能的可视化展现
稀疏模型的优点
用于把深度神经网络压缩为较小模型的深度压缩和用于神经网络正则化的DSD训练方法都是利用模型稀疏性来实现更小的模型或者更高的预测精度的技术。除了模型大小和预测精度,我们还研究了可以利用稀疏性的其他两个方面:速度和能耗,这超出了本文的讨论范围。读者可以查看EIE论文作为进一步的参考。