【韩松】Deep Gradient Comression

《Deep Gradient Compression》

作者韩松，清华电子系本科，Stanford PhD，深鉴科技联合创始人。主要的研究方向是，神经网络模型压缩以及硬件架构加速。

 

论文链接：https://arxiv.org/pdf/1712.01887.pdf

本篇文章发表在 ICLR 2018，韩松之前有一系列模型压缩的研究，比如其博士毕业论文《EFFICIENT METHODS AND HARDWARE FOR DEEP LEARNING》，以及获得 ICLR 2016最佳论文的《Deep Compression：Compressing deep neural networks with pruning trained quantization and Huffman coding》。

 

算法与模型压缩这个方向，前期的主要工作都是针对于模型的Inference阶段，现在的情况是，Inference compression已经做得很不错了。所以自然得考虑Training阶段的模型压缩。

 

本篇文章即主要针对模型的Training过程进行压缩。因为Training的领域这么多年都没有什么大的变化，基本上都是围绕着GD（Gradient Descent）以及其变体（比如SGD、momentum、Adam）做文章。所以本文针对distributed SGD进行压缩，主要思想是，小于一定阈值的gradient对于系统的每一次迭代影响是非常小的，可以将其积累到一定的阈值，而后一并传递。但是这样的流程，会产生一定的时延效应影响进度，所以文章又采取了一些方法来克服这种方式的时延带来的影响。

 

下面是文章的主体梳理：

 

 

【Abstract】

多节点分布式训练会遇到带宽限制的问题，尤其是在移动设备上的训练。本文发现99.9%的gradint传递都是冗余的，所以提出一种Deep Gradient Compression（DGC）的方法来减少通信带宽communication bandwidth。

 

为了防止压缩导致精度损失，DGC采用四种方法：①动量修正（momentum correction）、②局部梯度修剪（local gradient clipping）、③动量因子掩藏（momentum factor masking）、④预训练（warm-up training）。

 

DGC可以达到270x到600x的无损压缩率，梯度大小，ResNet50从97M降到0.35M。可以在1G以太网中完成大规模分布式训练。

 

 

一、介绍

讲了一些分布式训练的工作，同步SGD是主流的方法。但是问题在于网络通信带宽成为进一步扩大训练规模的瓶颈，在移动设备上尤为严重（更低的网络带宽、断续的网络连接、更昂贵的移动数据协议）。

 

引出DGC是用来解决训练高通信带宽的问题，其实贡献主要的压缩率的方法（积累梯度传播）是另两篇文献提出的，DGC方法整合了这些方法，同时解决了这些方法（延时效应）的精度损失问题。

 

实验测试，测试了三种应用，用CNN图像分类（CIFAR10和ImageNet），RNN语言模型（Penn Treebank），语音识别（Librispeech Corpus）。结果显示最高可以无损压缩600倍。

 

 

二、相关工作

相关工作有，异步SGD。主要有梯度量化、梯度稀疏两大类压缩方法。

a)      梯度量化

二值SGD（1 bit），语音应用10X加速。QSGD。TernGrad三值SGD（对应于三值网络）。

b)     梯度稀疏

Aji&Heafield（2017）提出Gradient Dropping（文中多次提到这份工作），

 

 

三、DGC

a)      梯度稀疏 gradient sparsification

想法是通过只传送重要的gradient来降低通信带宽，可以简要概括成 稀疏化更新。 重要的梯度，简单认为是大于一定阈值的梯度。为了防止信息损失，其余“不重要”的梯度，被局部累加起来。当其累加到阈值的时候再进行传送。算法流程如下表1：

表1、简述流程，

l  输入数据集X、minibatch尺寸b、节点（分布式worker）数N、优化函数SGD，初始化权重w。

l  对于节点k上的gradient的集合 Gk，随时间t迭代Gkt。

l  在一个minibatch中，采样b个样本x，每个样本迭代，计算梯度，并累加到Gkt。

l  对于要所有送给w[j]的Gkt[j]，取最大绝对值的s%（预设值）作为其阈值thr（感觉原文这里写错了，应该是max| Gkt[j] |）。

l  每个w[j]对应的Gkt [j]，其绝对值若大于阈值thr，则置Mask=1，并将值Gkt[j]传递出去，表示为~Gkt=Gkt * Mask；

l  否则置Mask=0，保留在分布终端本地，不进行传递，即。即Mask代表是否进行传递（或根据mask翻译为是否掩藏）。

l  最后，得到一个稀疏化的sparse（~Gkt），将N个节点（分布式终端）的gradient进行汇总求和，得到第t次迭代权重wt更新所需要的Gt。

l  通过SGD的算法，得到第t+1次的权值wt+1。

 

换一种视角：局部梯度累加，等同于，随时间推移增加batch size。理论推导如下

F(w)是全体数据集的loss function，据此优化。同步分布式SGD更新公式为（with N training nodes）：

其中X是训练集，w是权值，f是每一个样本的loss，N是training nodes，Bkt是N个minibatch的样本集（1）。

 

 

对权重的第i个值，在T次迭代后，有

即随时间迭代，可以视为一次的batch size从Nb增大为NbT。其中T称为稀疏更新间隔，即两次梯度传播的迭代间隔为T。

（但其实感觉这里有点牵强，后面也没有用到这个想法，因为后面的f仍然与wt~wt+T-1都有关）

 

又使用了Learning rate scaling（Goyal et al. 2017）的方法，即抛弃上面公式nT和NbT里的T。

 

b)      改进局部梯度累积 local gradient accumulation

但是简单采取上述方法，会严重影响收敛性。在Cifar10下导致1.0%精度损失。

这里采用动量修正（momentum correction）、局部梯度修剪（local gradient clipping）两种方法来解决高稀疏度下精度损失问题。

 

动量修正（momentum correction）

即将动量SGD的方法，应用到该稀疏化的策略中去。（因为动量SGD的算法并不能直接应用DGC）

 

 

传统的分布式动量SGD训练：

其中m是动量因子，N是节点数，

考虑第i个权重w(i)，T次迭代后，权值更新公式为：

              朴素动量的DGC

 

直接应用动量SGD到稀疏化DGC，有更新公式：

第一项为 将v视为gradient进行累加； 第二项为 应用DGC，大于阈值的累加值被传递给主服务器进行动量的更新，第三项为根据动量的权值更新。

 

 

则T次迭代后的权值更新公式：(原文此处得到的公式，有一定的近似不完全精确)

这个方法中对比传统动量SGD，略去的衰减因子，会导致收敛性能（精度）的损失。（个人理解：这其中较大的动量会不断累加，导致其影响能力超过gradient本身，从而导致信息误导）所以就需要，动量修正的DGC方法。

动量修正的DGC

思想为：

第一项的速度u视为“gradient”（其实是gradient的累加项），第二项可以视为针对“gradient”u应用传统SGD，第三项针对“gradient”u 应用DGC。

             

 

该算法的流程图如下：

（而后文中没给出，T次迭代后的表达式，应该是和上一种方法的相同。）

 

局部梯度修剪（local gradient clipping）

为了解决训练过程中的梯度爆炸（gradient exploding）问题。（注：其实都是一些老问题的老解决办法，放在这个新的DGC更新策略下，需要有所变动）

 

 

Pascanu et al. 2013 提出的方法，当gradient的L2-正则的求和达到一定阈值，就进行gradient rescale。是在所有gradient 聚合了以后执行。因为在DGC方法中，每个节点独立地进行gradient累加，所以在节点中（locally）执行gradient clipping，在将Gkt累加到Gkt-1之前执行。（在动量方法中就是Vt-1、Ut-1）。阈值设为N^(-1/2)。

c)      解决 迟滞效应（时延效应）

由于小的梯度用来积累，所以权重会存在一定的迟滞效应。实验中，当稀疏度为99.9%时，权重大约每600~1000次迭代会更新一次。迟滞会减慢收敛损失精度。采用两种方式来解决这个问题。

 

动量因子掩藏（Momentum Factor Masking）

 

Mitliagkas et al. (2016) 讨论了在异步训练中迟滞的影响。本文简单的用Mask因子来掩藏较小gradient的动量因子Ut。

这种方法可以停止延迟梯度产生的动量。（就好比当gradient没有传播时，由于其具有速度Ut，其积累会向着 旧的速度方向前进，而这个方向往往是偏斜的）

 

预训练（Warm-up Training）

在训练的初期，gradient往往是剧烈变动的（即梯度下降的方向变化很大），此时迟滞效应的影响会很大。Goyal et al.（2017）提出大minibatch训练的预训练方法。

 

在预热阶段，使用更小（原文：更不具侵略性的）的learning rate，较小的稀疏度，来减少被延迟的极端梯度数量。在前几个epochs，线性增加learning rate，e指数增加稀疏度（一开始很小）。

 

d)      所用技术总结概括，见下表1：

四、 实验结果

在Image Classification中，基于CIFAR10测了ResNet-110，基于ImageNet测了AlexNet和ResNet-50。

 

在Language Modeling中，使用Penn Treebank corpus数据集。测了一个两层的LSTM，每层1500个neuron。

 

在Speech Recognition中，使用了AN4和Librispeech数据集。测了一个5层LSTM，每层800个neuron。

 

结果如下，总结为同样高稀疏度下，可以不降精度（可能因为稀疏化，降低了一定过拟合，所以精度还有很小的提高）

 

 

（总结：总体来说，DGC方法的亮点：高稀疏度，采取一些方法解决精度损失，主要解决 降低 分布式训练中 通信带宽问题，即减少了节点间需要传输的梯度数量。）