由于良好的可扩展性，随机梯度下降（SGD）的并行实现是最近研究的热点。实现并行化SGD的关键障碍就是节点间梯度更新时的高带宽开销。因此，研究者们提出了一些启发式的梯度压缩方法，使得节点间只传输压缩后的梯度。尽管这些启发式方法在实践中很有效，但它们有时并不会收敛。
本文提出了量化SGD（Quantization SGD，QSGD），它是一类具有收敛保证且在实践中性能良好的压缩模式。QSGD允许用户平滑得权衡通信带宽和收敛时间：节点可以在每轮迭代时调整发送的比特数，代价可能是更高的方差。这种权衡是固有的，因为将其提高到某个阈值会违反信息理论的下限。QSGD在异步情况下保证了凸与非凸目标函数的收敛性，并且可以使用随机方差削减技术扩展。
当在图像识别与语音识别任务上训练深度神经网络时，QSGD显著地降低了端到端的训练时间。
1. Introduction
目前，许多研究者正在研究如何扩展大规模SGD算法。SGD算法的定义如下。令\(f:\mathbb{R}^n\rightarrow\mathbb{R}\)是我们要最小化的目标函数。我们可以得到随机梯度\(\widetilde{g}\)，即\(\mathbb{E}[\widetilde{g}(x)]=\triangledown f(x)\)。通过多次迭代，SGD可以收敛到最小值。
其中\(\boldsymbol{x}_t\)是当前，\(\eta_t\)是变量的步长。给出由未知分布\(D\)生成的独立同分布的数据点\(X_1,X_2,\cdots,X_m\)，和一个衡量模型\(\theta\)在数据点\(X\)的损失函数\(l(X,\theta)\)。我们希望找到一个最优模型\(\theta^\ast\)，使得\(f(\theta)=\mathbb{E}_{X\sim D}[l(X,\theta)]\)。该框架可以解决神经网络训练等基础性任务。
本文主要关注并行SGD方法，特别地，我们考虑将大型数据集分割到\(K\)个处理器上，它们共同优化目标函数\(f\)。每个处理器维护一个参数向量\(\boldsymbol{x}_t\)的本地副本，在每次迭代中，该处理器会获得对应参数向量的梯度更新。随后每个处理器向其他处理器广播它们的梯度更新，并聚合其他处理器的梯度更新以计算得到下一次迭代的参数向量\(\boldsymbol{x}_{t+1}\)
在许多并行化的SGD算法的实现中，每次迭代时每个处理器都需要将所有的梯度更新发送给所有其他的处理器。如果梯度向量比较稠密，每次迭代时每个处理器需要发送或接收\(n\)个浮点数并维护参数向量\(\boldsymbol{x}\)。实践中发现，梯度传输所产生的通信开销是一个重要的性能瓶颈。一种通用的降低通信开销的方法就是在略微损失训练精度的情况下对梯度进行压缩。一个简单的实现是降低表示的精度，该方法在凸性和稀疏性假设下会收敛。更为激进的量化技术是1BitSGD，它将梯度的每个分量减少到1比特，通过对\(\widetilde{g}\)的坐标平均缩放，可以在本地累积误差。在某些条件下，实验观察到1BitSGD可以保证收敛。然而，目前尚不清楚1BitSGD是否在强假设下也是如此，并且不清楚是否可以实现更高的压缩率。
2. Preliminaries
令\(\mathcal{X} \subseteq \mathbb{R}^n\)为一个已知的凸集，\(f:\mathcal{X} \rightarrow \mathbb{R}\)是一个未知凸函数，它具有可微和平滑的性质。假设我们要重复访问\(f\)在（可能是随机的）输入\(x\)处的随机梯度，并期望在正确的方向上前进。
Definition 2.1固定\(f:\mathcal{X} \leftarrow \mathbb{R}\)，\(f\)的随机梯度是一个随机函数\(\widetilde{g}(x)\)且\(\mathbb{E}[\widetilde{g}(x)]=\triangledown f(x)\)。对于所有的\(x\in \mathcal{X}\)，若满足\(\mathbb{E}[\left\|\widetilde{g} \right\|_2^2]\leq B\)则说明随机梯度的二阶矩最多为\(B\)；若满足\(\mathbb{E}[\left\|\widetilde{g}(x)-\triangledown f(x) \right\|_2^2]\leq \sigma^2\)则说明方差最大为\(\sigma^2\)。
Theorem 2.1令\(\mathcal{X} \subseteq \mathbb{R}^n\)为一个已知的凸集，\(f:\mathcal{X} \leftarrow \mathbb{R}\)是一个未知凸函数，它是拉普拉斯平滑的。对于给定的\(x\in \mathcal{X}\)，令\(R^2 = \sup_{x\in \mathcal{X}}\left\|x-x_0\right\|^2\)，固定迭代次数\(T>0\)，且\(f\)的随机梯度的方差上界为\(\sigma^2\)，每次迭代的步长为\(\eta_t = \frac{1}{L+1/\gamma}\)，其中\(\gamma = \frac{R}{\sigma}\sqrt{\frac{2}{T}}\)，那么就有：
Minibatch SGD小批量SGD（Minibatch SGD）是传统SGD的一种变体。在小批量SGD中，更新公式为\(x_{t+1} = \prod_{\mathcal{X}}(x_t-\eta_t \widetilde{G}_t(x_t))\)，其中\(\widetilde{G}_t(x_t) = \frac{1}{m}\sum_{m}^{i=1}\widetilde{g}_{t,i}\)。显然，如果随机梯度\(\widetilde{g}_{t,i}\)的方差上界为\(\sigma^2\)，那么随机梯度\(\widetilde{G}_t\)的方差上界就是\(\sigma^2/m\)。由Theorem2.1可知，只要等式\(\ref{2}\)中的第一项占主导地位，小批量SGD的收敛次数是传统SGD的\(1/m\)。
Data-parallel SGD我们主要考虑多GPU环境下的数据并行SGD算法的通信开销。我们有\(K\)个同步的处理器\(p_1,p_2,\cdots,p_K\)，使用点对点方式进行通信。其中每个处理器维护一个\(n\)维向量\(\boldsymbol{x}\)。
如算法1所示，在每次同步迭代中，每个处理器聚合\(\boldsymbol{x}\)的值，然后得到\(\boldsymbol{x}\)不同部分的随机梯度更新，接着将这些更新发送给其他所有的处理器，最终在本地聚合并应用这些更新。重要的是，我们在算法第3行（广播之前）与第7行（接收之后）添加了编码与解码的步骤。在不同的SGD算法的变体中，只有编码与解码函数是不同的。注意这里解码函数无须覆盖原始梯度\(\widetilde{g}^l\)，因为我们通常使用近似值。
对于原始的并行SGD（即没有编码与解码函数）,在每个处理器中，如果\(\boldsymbol{x}_t\)是处理器在迭代\(t\)之前维护的\(\boldsymbol{x}\)值，那么在此迭代结束时\(\boldsymbol{x}\)的更新值是\(\boldsymbol{x}_{t+1}=\boldsymbol{x}_t -(\eta_t /K)\sum_{l=1}^K\widetilde{g}^l(\boldsymbol{x}_t)\)，其中每个\(\widetilde{g}^l\)都是随机梯度。特别地，此更新是大小为\(K\)的小批量更新。因此，综合上述讨论与Theorem2.1，我们有如下推论：
Corollary 2.2令\(\mathcal{X},f,L,x_0\text{和}R\)的定义与Theorem2.1中相同。固定\(\epsilon > 0\)。假设我们在\(K\)个处理器上运行并行SGD，每个处理器访问独立的随机梯度且二阶矩的界为\(B\)，步长为\(\eta_t=1/(L+\sqrt{K}/\gamma)\)，其中\(\gamma\)定义与Theorem2.1中相同。如果
大多数情况下，式\(\ref{3}\)中最大的第一项将主导必要的迭代次数。具体来说，迭代次数与二阶矩\(B\)线性相关。
3. Quantized Stochastic Gradient Descent(QSGD)
本节主要介绍QSGD，符号约定如下。\(\log\)是以2为底的对数函数，浮点数约定为32位。对于任意向量\(v\in \mathbb{R}^n\)，令\(\left\|v\right\|_0\)表示\(v\)的非零数。对于任意字符串\(\omega \in \{0,1\}^*\)，我们令\(|\omega|\)表示该字符串的长度。对于任意标量\(x\in \mathbb{R}\)，我们用\(\text{sng}(x)\in \{-1,+1\}\)表示它的符号，其中\(\text{sgn}(0)=1\)。
3.1. Generalized Stochastic Quantization and Coding
Stochastic Quantization首先考虑一个通用的、参数化的有损压缩模式。量化函数定义为\(Q_s(v)\)，其中\(s \geq 1\)是一个需要调整的超参数，代表量化的程度。直观地，我们定义\(s\)为0和1之间的均匀分布的级别，每个值以保持期望值不变的方式进行量化，并引入最小方差。如图1所示。
对于任意\(v\in \mathbb{R}^n\)且\(v \neq 0\)，\(Q_s(v)\)的定义为：
其中\(\xi(v,s)\)是按照下式定义的随机变量。令\(0\leq l < s\)是一个满足\(|v_i|/\|v\|_2 \in [l/s,(l+1)/s]\)条件的整数。也就是说，\([l/s,(l+1)/s]\)是\(|v_i|/\|v\|_2\)对应的量化区间。那么
这里，\(p(a,s)=as-l\)对于所有的\(a\in[0,1]\)都成立。若\(v=0\)，则定义\(Q_s(v)=0\)。\(\xi_i(v,s)\)分布具有最小方差且它的期望满足\(\mathbb{E}[\xi_i]=|v_i|/\|v\|_2\)。
Lemma 3.1对于所有的向量\(v\in\mathbb{R}^n\)，我们有（1）\(\mathbb{E}[Q_s(v)]=v\)（2）\(\mathbb{E}[\|Q_s(v)-v\|_2^2]\leq \min(n/s^2,\sqrt{n}/s)\)（3）\(\mathbb{E}[\|Q_s(v)\|_0]\leq s(s+\sqrt{n})\)
Efficient Coding of Gradients
对于任意向量\(v\)，\(Q_s(v)\)的输出是一个三元组\((\|v\|_2,\sigma,\zeta)\)，其中\(\sigma\)是\(v_i\)的符号，而\(\zeta\)是由整数值\(s·\xi_i(v,s)\)构成的向量。编码方案背后的关键思想是并非所有整数值\(s·\xi_i(v,s)\)都以相同的概率出现，具体来说，较大的整数出现的频率会较低。我们将通过专门的Elias整数编码来实现它。对于任何正整数k，其编码（表示为Elias(k)）从k的二进制表示开始，其前缀为该表示的长度，然后它递归编码这个前缀。对于任意正整数k，最终编码长度为\(|Elias(k)|=\log k+\log\log k + \cdots + 1 \leq (1+o(1))\log k +1\)，Elias编码与解码都可以很高效地实现。
对于给定三元组\((\|v\|_2,\sigma,\zeta)\)表示的梯度向量，使用\(s\)表示量化程度，我们的编码算法的输出是如下定义的字符串\(S\)。首先，使用32位来编码\(\|v\|_2\)。接着，它使用Elias递归编码\(\zeta\)的第一个非零项的位置。然后用一位来表示符号\(\sigma_i\)，再后面是\(Elias(s·\xi_i(v,s))\)。迭代地，它继续编码从当前位置到下一个非零项的距离，并以相同的方式编码\(\sigma_i\)和\(\zeta_i\)。解码过程很直观：我们首先读取前32位来构建\(\|v\|_2\)，然后迭代地使用Elias递归解码\(\zeta\)中非零项的位置、非零项的符号和值。
Theorem 3.2固定\(f:\mathbb{R}^n \rightarrow \mathbb{R}\)，然后取任意\(\boldsymbol{x} \in\mathbb{R}^n\)，固定量化级别\(s \geq 2\)。若\(\widetilde{g}(\boldsymbol{x})\)是\(f\)在\(\boldsymbol{x}\)处的随机梯度且二阶矩上限为\(B\)，则\(Q_s(\widetilde{g}(\boldsymbol{x}))\)是\(f\)在\(\boldsymbol{x}\)处的随机梯度且方差上界为\(\min\left(\frac{n}{s^2},\frac{\sqrt{n}}{s}\right)B\)。此外，因为编码函数的存在，所以\(Q_s(\widetilde{g}(\boldsymbol{x}))\)需要通信的比特数上界为：
3.2. QSGD Guarantees
将上面给出的通信和方差的界与在平滑凸函数上的SGD算法的收敛保证放在一起,产生以下结果：
Theorem 3.4(Smooth Convex QSGD)令\(\mathcal{X},f,L,x_0\text{和}R\)与Theorem2.1中定义相同。固定\(\epsilon > 0\)，假设我们在\(K\)个处理器上以\(s\)量化级别来运行并行QSGD且随机梯度的二阶矩的界为\(B\)，步长为\(\eta_t=1/(L+\sqrt{K}/\gamma)\)，其中\(\gamma\)在Theorem 2.1中定义；再令\(\sigma=B'\)，其中\(B'=\min\left(\frac{n}{s^2},\frac{\sqrt{n}}{s}\right)B\)。那么若\(T=O\left(R^2·\max\left(\frac{2B'}{K\epsilon^2},\frac{L}{\epsilon}\right)\right)\)，则\(\mathbb{E}\left[f\left(\frac{1}{T}\sum_{t=0}^T\boldsymbol{x}_t\right)\right] - \min_{\boldsymbol{x}\in\mathcal{X}}f(\boldsymbol{x})\leq\epsilon\)。此外，QSGD在每一轮需要\(\left(3+\left(\frac{3}{2}+o(1)\right)\log\left(\frac{2(s^2+n)}{s^2+\sqrt{n}}\right)\right)\)Bit来进行通信。当\(s=\sqrt{n}\)时，通信所需的比特数降至\(2.8n+32\)。
Theorem 3.5(QSGD for smooth non-convex optimization)令\(f:\mathbb{R}^n \rightarrow \mathbb{R}\)是一个拉普拉斯平滑（可能非凸）函数，令\(\boldsymbol{x}_1\)是任意的初始点。令\(T>0\)，\(s>0\)。然后在\(\{1,2,\cdots,N\}\)上有一个随机停止时间\(R\)，因此拥有常数步长\(\eta=O(1/L)\)且\(f\)的随机梯度的二阶矩的界为\(B\)的QSGD满足\(\frac{1}{L}\mathbb{E}[\|\triangledown f(\boldsymbol{x})\|^2_2]\leq O\left(\frac{\sqrt{L(f(\boldsymbol{x}_1)-f^*)}}{N}+\frac{\min(n/s^2,\sqrt{n}/s)B}{L}\right)\)。另外，通信开销与Theorem3.4相同
3.3. Quantized Variance-Reduced SGD
假设我们有\(K\)个处理器和一个大于0的参数\(m\)，每个处理器\(i\)都访问函数\(\{f_{im/K},\cdots,f_{(i+1)m/K-1}\}\)。目标是近似最小化\(f=\frac{1}{m}\sum_{i=1}^mf_i\)。对于处理器\(i\)，令\(h_i=\frac{1}{m}\sum_{j=im/K}^{(i+1)m/K-1}f_i\)是该处理器所知道的f的一部分，因此有\(f=\sum_{i=1}^K h_i\)。
那么应该如何在并行化SVRG算法上应用随机量化方法以减少通信开销呢？经过检查，我们注意到最终的更新将破坏标准SVRG。我们最终证明可以使用随机量化技术来量化SVRG更新，并获得相同的收敛界限。
Algorithm Description令\(\widetilde{Q}(\boldsymbol{v})=Q(\boldsymbol{v},\sqrt{n}\)，其中\(Q(\boldsymbol{v},s)\)在3.1节中定义。给定任意初始点x0，我们令\(\boldsymbol{y}^{(1)}=\boldsymbol{x}_0\)。在第p轮次开始时，每个处理器广播一个未量化的梯度\(\triangledown h_i(\boldsymbol{y}^{(p)})\)，并对其他处理器发来的梯度进行聚合\(\triangledown f(\boldsymbol{y}^{(p)})=\sum_{i=1}^m\triangledown h_i(\boldsymbol{y}^{(p)}\)。在每一轮次，对于每次迭代\(t=1,2,\cdots,T\)和每个处理器\(i=1,2,\cdots,K\)，我们令\(j_{i,t}^{(p)}\)是一个随机整数，那么在每次迭代中，处理器\(i\)广播更新向量\(\boldsymbol{u}_{t,i}^{(p)}=\widetilde{Q}\left(\triangledown f_{j_{i,t}^{(p)}}(\boldsymbol{x}_t^{(p)})-\triangledown f_{j_{i,t}^{(p)}}(\boldsymbol{y}_t^{(p)})+\triangledown f(\boldsymbol{y}_t^{(p)})\right)\)。然后，每个处理器计算所有的更新\(\boldsymbol{u}_t^{(p)}=\frac{1}{K}\sum_{i=1}^K \boldsymbol{u}_{t,i}\)，并令\(\boldsymbol{x}_{t+1}^{(p)}=\boldsymbol{x}_t^{(p)}-\eta \boldsymbol{u}_t^{(p)}\)。在第p轮次结束时，每个处理器会执行\(\boldsymbol{y}^{(p+1)}=\frac{1}{T}\sum_{t=1}^T\boldsymbol{x}_t^{(p)}\)。
Theorem 3.6令\(f(x)=\frac{1}{m}\sum_{i=1}^mf_i(\boldsymbol{x})\)，其中\(f\)是一个强凸函数，对于所有的\(i\)来说，\(f_i\)是一个强凸且拉普拉斯平滑的函数。令\(x^\ast\)是\(f\)在\(\mathbb{R}^n\)上的最小值点。若\(\eta=O(1/L)\)且\(T=O(L/l)\)，那么初始点为\(\boldsymbol{y}^{(1)}\)的QSVRG对任意\(p\geq1\)的轮次保证了\(\mathbb{E}[f(\boldsymbol{y}^{(p+1)})]-f(\boldsymbol{x}^{\ast})\leq0.9^p (f(\boldsymbol{y}^{(1)})-f(\boldsymbol{x}^{\ast}))\)。此外，每个轮次拥有\(T\)次迭代的QSVRG至多需要\((F+2.8n)(T+1)+Fn\)比特进行通信。
4. QSGD Variants
首先，我们可以通过量化到大小固定为d的buckets中来控制量化方差。如果我们把每个梯度都看成一维向量v，桶的定义是向量中连续的d个值，比如v[(i-1)d+1:i·d]。我们可以使用qsgd单独量化每个桶中的值。d=1的时候相当于没进行量化，也就是普通的sgd。d=n相当于前面小节提到的全量化。显然，使用桶量化后，由lemma3.1得到的收敛保证将取决于d而不是完整的维度n。我们可以通过桶来控制方差，代价就是需要额外存储缩放因子。举例来说，如果我们令桶的大小为512，量化4位，那么方差的上界就变为\(\sqrt{512} / 2^4=1.41\)。
其次，与理论不同的是，我们会按照向量的最大值而不是2范数进行缩放。直观地讲，通过最大值进行归一化可以保留更多的值，并且对于相同的迭代次数具有略高的精度。由于较低的比特宽度（例如，每个维度32Bit到2Bit），两种方法在baseline上降低了相同的带宽，但是通过最大值进行归一化理论上不提供稀疏性保证。在实践中，最大值归一化会产生非平凡的稀疏性。
5. Experiments
Setup我们在amazon ec2p2.16xlarge实例上做的实验，每个实例包括16块tesla k80显卡。实例间使用GPUdirect点对点通信，不支持NCCL。我们在CNTK上实现了QSGD，CNTK提供一个高效地基于MPI-Bcast的GPU到GPU的通信方式，并且实现了1BitSGD。代码开源在GitLab上。
我们在图像分类和语音识别任务上做了实验。其中图像分类使用ilsvrs2015（imagenet）,cifar10和mnist数据集，语音识别使用cmu an4数据集。对于图像分类任务，我们实现了alexnet，vgg，resnet和带有bn的inception-net。对于语音识别任务，我们训练了lstm。实验详情见表1。
Protocol在实验中我们使用的都是标准的网络，除非特别说明，否则实验中用到的网络模型的超参数都是CNTK2.0中默认最优的。随着GPU数量的增多，我们增大batchsize以平衡计算和通信开销。我们采用双缓冲技术在计算的同时进行通信和量化。量化技术在较低的学习率下效果很好，因此我们使用32Bit的学习率，并且减少桶大小以降低方差。我们并不量化较小的（元素数量小于10K）梯度矩阵，因为这样会得不偿失。然而，在实验中99%的梯度都被量化
Communication vs. Compution图2给出了不同网络模型在图像分类任务的结果。首先，基于通信与计算的比率，我们可以粗略地将神经网络模型分为通信密集型（AlexNet，VGG，LSTM）和计算密集型（Inception，ResNet）。对于这两种网络类型，随着GPU数量的增加，通信时间的占比会显着增加。对于32位版本的传统SGD，所有的网络模型都可以从降低通信中获益。举例来说，在16块GPU上训练的AlexNet，当batchsize为1024时，超过80%的训练时间都花在通信上。而对于在2块GPU上训练的LSTM模型，当batchsize为256时，通信占总训练时间的71%。
接下来，我们看一下QSGD的通信时间与总训练时间的关系（这里通信时间包括压缩和解压梯度的时间）。我们测试了使用2Bit量化与桶大小为128的QSGD以及使用4Bit量化、8Bit量化与桶大小为512的QSGD。这两种变体的实验结果非常相似，因为不同的桶大小意味着4Bit量化仅比2Bit量化发送的数据多77％（但是比32Bit少了约8倍）。本实验中的桶大小仅保证不错的收敛性，没有进行仔细地调参。对于在16块GPU上以batchsize为1024训练的AlexNet，4BitQSGD减少了4倍的通信时间，总训练时间减少了2.5倍。对于LSTM模型，通信量减少了6.8倍，总训练时间减少了2.7倍。
Accuracy我们在imagenet数据集上训练了alexnet和resent，在an4数据集上训练lstm，在cifar10上训练reset110，在mnist训练了一个2层的感知机。实验结果由图3给出，完整的数据在表1中。qsgd在使用8块gpu时性能最好，一般来说，4Bit或8Bit梯度量化足以达到甚至略微提高32Bit的传统SGD的精度，同时确保非平凡的加速比。注意到在训练深度网络时向梯度添加噪声的好处，而量化可以被视为一种零均值噪声，这恰好使得通信更有效。与此同时，我们注意到更高程度的量化可能会损害准确性。比如，使用桶大小为8192的4BitQSGD训练的AlexNet与使用32BitSDG相比，损失了0.57%的top5准确率和0.68%的top1准确率。
需要详细地研究的一个问题是哪些层对量化更敏感。如果过度地量化卷积层（例如2BitQSDG）会导致精度损失，但使用4BitQSGD或8BitQSGD几乎不损失精度。这表明，视觉任务的现代架构（如ResNet或Inception，几乎完全是卷积层）可能比LSTM等网络更少地受益于量化技术。
Additional Experiments论文的完整版本包含其他实验，包括与1BitSGD的完整比较。简言之，QSGD在本文考虑的网络和参数方面优于或等于1BitSGD的性能和最终精度。
6. Conclusions and Future Work
QSGD是一类允许用户在每轮的通信量和训练时间进行平滑权衡的算法。实验表明，QSGD在各种任务都表现良好。但是还有很多工作没有完成，最重要的是利用QSGD创造的稀疏性。MPI的当前实现不提供对稀疏类型的支持，但将来的工作可以对这部分进行探索。此外，还可以研究一下QSGD在大规模应用（如超级计算机）中的潜力。
来源：oschina
链接：https://my.oschina.net/u/4357381/blog/4294473