INTERSPEECH 2014 | 1-Bit Stochastic Gradient Descent and its Application to Data-Parallel Distributed Training of Speech DNNs

这篇文章之前也读过，不过读的不太仔细，论文中的一些细节并没有注意到。最近为了写开题报告，又把这篇论文细读了一遍。据笔者了解，这篇论文应该是梯度量化领域的开山之作，首次使用了梯度量化技术来降低分布式神经网络训练的通信开销。除此之外，这篇文章还提出了误差补偿机制，这种机制可以缓解梯度量化的负面影响，降低信息丢失所带来的模型精度损失。

对于数据并行式训练来说，最佳节点数量\(\hat{K}\)能够使得节点中计算和数据通信完全重叠，即同时使通信和计算资源饱和：

\[T_{calc}(\hat{K}) = T_{comm}(\hat{K}) \]

\(T_{calc}\)和\(T_{comm}\)分别是节点上每个minibatch所需要的计算和通信时间。如果我们将计算和通信开销细分，就能够求出最优节点数量：

\[\hat{K} = \frac{\frac{N}{2} \cdot T_{calc}^{frm}+C\cdot T_{post}^{calc}}{\frac{1}{Z}\cdot T^{float}_{comm}-T_{upd}^{calc}} \]

符号	含义	与参数量\(M\)的关系
\(T_{calc}^{frm}\)	处理数据的时间，这部分是可以并行的	\(\propto\frac{M}{\text{FLOPS}}\)
\(T_{calc}^{post}\)	对梯度后处理的时间，如momentum + AdaGrad操作	\(\propto\frac{M}{\text{RAM bindwidth}}\)
\(T_{comm}^{float}\)	传输梯度（由单精度浮点数表示）的时间	\(\propto\frac{M}{\text{Network bindwidth}}\)
\(T_{calc}^{upd}\)	模型更新时间，相对于\(K\)来说是固定的	\(\propto\frac{M}{\text{RAM bindwidth}}\)

为了实现更高的并行度，作者提出了双缓冲（double buffering）的概念，即在每个节点上把一个minibatch分成两部分，在交换其中一部分的梯度时进行另一部分的计算。然而，双缓冲机制会引入额外的更新开销，当通信开销小于更新开销（即\(T_{comm}^{float} \lt T_{calc}^{upd}\)）时，上述公式就不再成立，双缓冲机制也就失去了作用。

本文的核心思想就是将32位单精度浮点数压缩成1位，从而降低训练过程的数据交换开销。为了进一步降低通信开销，作者设定在节点间只进行梯度的交换，不进行模型参数的交换。为了降低量化误差带来的负面影响，作者使用了误差补偿技术：每次量化时，把上一次迭代的量化误差加到本次迭代的梯度上，然后再进行量化，接着求出本次量化操作的误差。这种误差补偿机制可以确保所有的梯度都会再一定程度上对模型更新产生作用，只不过这种作用分散在不同的迭代中——类似于一种延迟更新的形式。作者指出，使用误差补偿后，就可以在几乎不损失模型精度的情况下将梯度由32位量化成1位。

\[\begin{aligned} \tilde{g}_t &= Q(g_t + \text{Err}_{t-1})\\ \text{Err}_{t} &= g_t - Q^{-1}(\tilde{g}_t) \end{aligned} \]

在具体实现上，一种比较简单（且有效）的方法是将大于0的梯度值编码位1，小于等于0的梯度值编码为0。在反量化（解码）时，将1解码为+1，将0解码为-1，再进行聚合与更新操作。

本文作者并没有使用参数服务器架构，而是使用了一种类似于AllReduce的梯度聚合方法。每个节点负责聚合\(1/K\)的模型参数子集。也就是说，假设模型参数量为\(M\)，那么每个节点都从其他所有节点接收到\(M\times \frac{K-1}{K^2}\)个值。这些值被反量化后，在节点上进行聚合以及后处理（AdaGrad、momentum），然后进行第二次量化，再把量化后的梯度传给其他所有的节点并进行更新操作。

根据公式可以看出，有3种方法可以提高并行性：（1）增加batchsize，（2）增大压缩比率Z，（3）降低更新开销\(T_{calc}^{upd}\)。前面已经介绍如何对梯度进行压缩，这里主要介绍方法1和方法3。
对于方法1，一味地增加batchsize并不可取，因为这可能导致模型不收敛。为了解决较大batchsize下模型不收敛的问题，作者使用了较小的学习率和AdaGrad优化算法。但是一般来说，当batchsize增大时，学习率应该按照一定的比例进行放大[1, 2]，可能是因为模型不同的原因，作者使用较小的学习率可以获得不错的效果。设置全局学习率之后，AdaGrad每次通过全局学习率逐参数地除以历史梯度平方和的平方根，使得每个参数的学习率不同。AdaGrad的作用是在参数空间较为平缓的方向学习率会较大（因为平缓，所以累加梯度平方和较小，对应学习下降的幅度较小）；参数空间较为陡峭的方向学习率会较小，从而加快训练速度。作者指出，我们可以在三个不同的时间点应用AdaGrad：量化之前在每个节点上本地使用（可能会有助于量化，但也有可能在节点之间引入不一致的风险）；聚合后的数据交换期间（有可能引起量化干扰）；动量平滑之后（节省内存和更新开销，但由于平滑了峰值会降低了AdaGrad的效果）。通过实验验证，作者选择在梯度聚合后应用AdaGrad算法，（一个可能的）原因是量化操作将异常值分散在多个minibatch种，这样它们对标准差估计的影响较小。

对于方法3，作者使用了模型并行来降低更新开销。不过在作者的实验中，模型并行的效果貌似并不是很好，所以这里就不再讨论了。

参考资料
[1] Goyal P, Dollár P, Girshick R, et al. Accurate, large minibatch sgd: Training imagenet in 1 hour[J]. arXiv preprint arXiv:1706.02677, 2017.
[2] Krizhevsky A. One weird trick for parallelizing convolutional neural networks[J]. arXiv preprint arXiv:1404.5997, 2014.