MLHPC 2016 | Communication Quantization for Data-parallel Training of Deep Neural Networks

本文主要研究HPC上进行数据并行训练的可行性。作者首先在HPC上实现了两种通信量化算法（1 Bit SGD以及阈值量化），然后提出了自适应量化算法以解决它们的缺点。此外，发挥出量化算法的性能，作者还自己实现了一个Allreduce算法。

1 Bit SGD可以实现良好的重构和较低的误差，但与阈值量化相比，它的计算开销更大，并且压缩率不能达到32倍以上。阈值量化速度很快，但是不同的模型需要设置不同的阈值，而且选择好的阈值也很困难，并且使用阈值\(\tau\)作为重建值是次优的。如果阈值设置的比较小，那么由于误差补偿的存在，可能会导致传输大量的数据。

自适应量化使用固定比例\(\pi\)来表示每次迭代时要发送的梯度更新比例。第一步，自适应量化要确定满足当前迭代所需比例的正阈值\(\tau^+\)和负阈值\(\tau^-\)。假设梯度向量中有\(k\)个非负值，我们只需要发送其中最大的\(\frac{k}{\pi}\)个值，因此，正阈值\(\tau^+\)就是梯度向量中第\(\frac{k}{\pi}\)个值。一般而言，我们可以通过快速选择等高效算法，在\(O(N)\)时间内找到第\(\frac{k}{\pi}\)个元素。负阈值的确定与正阈值类似，只不过是选择梯度向量中最小的\(\frac{k}{\pi}\)个值。注意，正阈值和负阈值都由同一个比例\(\pi\)确定的。同样，自适应量化需要使用误差补偿技术来防止模型不收敛。在接收端的重建阶段，自适应量化分别对大于正阈值\(\tau^+\)和小于负阈值\(\tau^-\)的元素求一个平均值，用这两个均值分别作为重建向量中的元素。

作者发现原始的MPI_Allreduce接口在传输压缩数据时表现不好，而且使用用户自定义操作时MPI_Allreduce会退化成recursive-doubling实现，因此他使用原始的MPI_Send和MPI_Recv方法实现了一个类似于Ring Allreduce的聚合通信算法。

下面是本文的一些实验结果。第一幅图展示了模型中第三个全连接层每次迭代时所发送的数据量。无量化和1 Bit量化的结果符合预期：要么发送完整的梯度矩阵，要么将数据量减少约32倍。自适应量化在\(\pi = 64\)的情况下，我们期望得到与1 Bit量化相似的数据压缩比。但是由于它使用采样来近似阈值以排除更新，因此它不能完全实现32倍的压缩比。但是，只要采样是一个足够好的近似值，那么数据量就不会过多或过少。阈值量化即使在\(\tau = 0.001\)的情况下也发送很少的数据，并且发送的数量非常不稳定。这就导致其测试精度明显低于其他方法，因为每个模型基本上仅从其本地的数据中进行学习。

下面这幅图展示了不同梯度量化方法的通信时间（同样是模型中第三个全连接层）。可以看到，阈值量化的通信时间是最短的，因为它只需要传输很少的数据。第二快的是传统的MPI_Allreduce，这是因为量化引入了额外的计算开销，当某一层本身的计算量较少时，这种额外的计算开销就会被放大。自适应量化紧随其后，而且相对于传统的MPI_Allreduce更加平稳。最慢的是1 Bit量化，因为它有一个额外的AdaGrad操作。