ICML 2018 | Error Compensated Quantized SGD and its Applications to Large-Scale Distirbuted Optimization

如何降低大规模分布式训练中的通信开销是一个老生常谈的话题。本文提出了一种基于误差补偿的量化SGD(Error Compensated Quantized SGD, 下文简称ECQ-SGD)算法,能够大幅降低训练过程中的通信开销,提升训练效率。ECQSGD通过量化本地梯度和累积量化误差来降低通信开销,加速收敛过程。

1. Preliminaries

首先,考虑如下无约束优化:

其中$\mathbf{w} \in \mathbb{R}^{d}$,$f : \mathbb{R}^{d} \rightarrow \mathbb{R}$是一个可微的凸函数。假设我们使用数据并行的方式进行分布式训练,那么完整的数据集$\mathcal{D}$会被均匀地分发到$P$个节点上,第$p$个节点的数据子集就是$\mathcal{D}_p$。形式上,我们需要优化的问题如下:

图1描述了整个分布式训练的过程。注意到本文使用的Peer-to-Peer架构,而不是传统的参数服务器架构。首先,每个节点使用相同的随机数种子来初始化本地模型。在第t次迭代,每个节点随机选择一个mini-batch的训练样本计算本地梯度,然后将本地梯度广播到其他的所有节点,最终计算出全局梯度并更新模型参数。

2. Error Compensated Quantized SGD

接下来介绍ECQ-SGD算法。ECQ-SGD的主要思路是通过量化后的本地梯度更新模型参数。在每次迭代中,当前的本地梯度都会由之前所有迭代的累积量化误差进行补偿,然后输入到随机量化函数中进行压缩。

令$Q : \mathbb{R}^{d} \rightarrow \mathcal{C}^{d}$是一个无偏的随机量化函数。每次迭代中,每个节点在广播本地梯度之前会对梯度进行量化操作:

这里$\mathbf{g}_{p}^{(t)}$是第p个节点在第t次迭代时的本地梯度,$\tilde{\mathbf{g}}_{p}^{(t)}$就是量化后的本地梯度。当某个节点接收到所有其他节点发送来的梯度后,该节点就会按照下式计算全局梯度并更新本地模型的参数。

ECQ-SGD的核心思想就是在量化本地梯度时,当前梯度以及之前所有迭代的累积量化误差都会被考虑进去。令$\mathbf{h}_{p}^{(t)}$表示第$p$个节点在第$t$次迭代时的累积量化误差,那么

其中$\beta$是时间衰减因子($0 \leq \beta \leq 1$)。注意到这里$\mathbf{h}_{p}^{(t)}$可以增量式计算:

显然,$\mathbf{h}_{p}^{(0)}=0$。此时,量化的本地梯度可以根据下面的式子计算,$\alpha$是一个大于等于0的补偿因子:

关于随机量化函数,本文采用了与QSGD类似的量化函数:

上式中$\left| \mathbf{g} \right|$是缩放因子,可以是$l_{2}$范数或$l_{\infty}$范数。$\xi(\cdot)$是一个将标量映射到$\left\{0, \frac{1}{s}, \ldots, 1\right\}$的随机函数,当$\left|g_{i}\right| /\left|\mathrm{g}\right|$的值落入区间$\left[\frac{l}{s}, \frac{l+1}{s}\right)$时,该函数的取值如下式所示。式中的$s$定义了量化级别,$s$越大,量化的粒度越细,通信开销越大。

对本地梯度进行量化后,我们只需要$r=\left[\log _{2}(2 s+1)\right]$比特来对量化梯度的每个分量$\tilde{g}_{i}$进行编码,以及一个浮点数来表示标量$|\mathbf{g}|$。因此总的通信开销是$32+d r (r \ll 32)$比特,远小于使用32位浮点数时的通信开销($32d$比特)。此外,使用Huffman编码等技术,可以进一步减少通信开销。完整的算法由算法1描述。

文章的下一部分主要介绍了一些理论分析,包括量化误差的方差上界以及二次规划的收敛性证明,类似于QSGD中的证明,这里不再赘述。接下来,我们看一下ECQ-SGD在不同任务上的实验对比结果。

3. Experiments

3.1. Linear Model

第一个实验是拟合一个线性模型,用到的数据集是Syn-256、Syn512和Syn-1024。假设我们要拟合的模型是$y_{i}=\mathbf{w}^{\star T}\mathbf{x}_{i}+\epsilon_{i}$,其中$\mathbf{w}^\star\in\mathbb{R}^{d}$是我们要逼近的最优参数。实验中QSGD和ECQ-SGD都使用$l_2$范数,量化因子$s$设为4,对于ECQ-SGD,设置$\alpha = 0.2$,$\beta=0.9$。

图2中,第一行是损失函数的值,第二行是参数与最优解的距离。对于这三个数据集,ECQ-SGD损失函数的收敛更接近于标准SGD,而且收敛的还比QSGD快。另一方面,ECQ-SGD与最优解的距离明显小于QSGD,这表明量化误差对误差界限的贡献得到了很好的抑制。下面比较的是QSGD和ECQ-SGD在大数据集Syn-20k上的运行时间。图3显示了1000次迭代后各方法的时间开销和测试损失。尽管ECQ-SGD需要额外的编码与解码时间,但是在误差允许的范围内,总体的运行时间还是小于传统SGD。

本文还在YearPredictionMSDgisette数据集上使用不同的梯度量化方法训练了线性回归和logistic回归模型,比较它们的损失函数和通信开销。

根据图4,除1Bit-SGD以外,所有的量化方法都以相近的速度收敛到同一水平,这可能是因为YearPredictionMSD数据集的特征比较少,所以不同的量化方法没有很大的性能差距。

对于拥有5000个特征的gisette数据集,TernGrad和ECQ-SGD仍然能像传统的SGD一样收敛,但是1Bit-SGD和QSGD则有严重的性能下降。另一方面,ECQ-SGD拥有比TernGrad更高的压缩率。

3.2. Convolutional Neural Networks

第二个实验是在CIFAR-10数据集上使用不同的量化方法训练了一个ResNet-20模型,损失函数与通信开销如图5所示。对于所有的量化方法,batch size均为128,学习率从0.1开始,每40k和60k次迭代衰减为原来的1/10,整个训练过程共200epoch,约78k次迭代。

图5第一列比较了不同方法的损失值和通信开销。可以发现,所有的方法都以相似的速率收敛,但是ECQ-SGD可以降低80倍的通信开销。图5的后面两列详细比较了QSGD和ECQ-SGD,第二列使用$l_2$范数,第三列使用$l_\infin$范数。可以观察到,ECQ-SGD在收敛速度和分类精度方面始终优于QSGD,而在相同的超参数设置下,这两种方法的通信成本相似。

3.3. Performance Model

下面通过性能模型来分析ECQ-SGD的可扩展性。实验环境为:Intel Xeon E5-2680 CPU,Nvidia Tesla P40 GPU(8 units per node),Mellanox ConnectX-3 Pro网卡(40 Gb/s)。

图6是在ILSVRC-12数据集上训练ResNet-50模型时的吞吐量。随着GPU数量的增多,ECQ-SGD与标准SGD吞吐量的差距越来越大,使用512块GPU时,ECQ-SGD的吞吐量是标准SGD的1.435倍。

3.4. Parameter Study

最后一部分介绍的是超参数的学习,也就是如何设置ECQ-SGD中两个重要参数$\alpha$和$\beta$。和实验2一样,在CIFAR-10上面训练一个ResNet-20模型,固定$\alpha$为0.01,可以发现当$\beta$设置成1时错误率最低。当$\beta$固定为1时,$\alpha$的最优值在0.01到0.1之间。如果$\alpha$过大,那么模型最终可能不会收敛。