Deep Gradient Compression: Reducing The Communication Bandwidth For Distributed Training

Deep Gradient Compression: Reducing The Communication Bandwidth For Distributed Training

Under Review

Intro

分布式的SGD optimization问题中，会发现computation时间是降低了，通过data parallellism和增加node，但是communication相应增加，尤其是gradient非常复杂的情况下。这篇paper提出来Deep Gradient Compression的方法来解决。和前面已有的方法比较，DGC的优势在于能够不损失精度。

Deep Gradient Compression

Gradient Sparsification

每次只传输重要的gradient，通过设置某一个threshold。然后为了避免损失信息，将剩余的gradient进行accumulate。最终，这个gradient会积累起来，一直到达到可以传输。这种处理方法等同于增加了batch－size，如公式2描述。

Improving the Local Gradient Accumulation

当gradient是sparse的情况下，convergence会变差。使用momentum correction和local gradient clipping来缓解这个问题。

Momentum Correction:

vanilla momentum如下：

$$ w{t+1} = w_t - \eta u{t+1}\ u_{t+1} = m u_t + \nabla f

$$

而因为local gradient会进行accumulation，所以$$\nabla f_t = \frac{1}{N_b} \sum_k \sum_x \nabla f(x, w_t)$$

在这里的设定下，第k个节点上的gradient会accumulate，那么momentum变成了：

$$ u{t+1}^k = m u_t^k + \nabla f_t^k\ v{t+1}^k = vt^k + u{t+1}^k\ w{t+1}^k = w_t^k - \eta \sum_k sparse(v{t+1}^k)

$$

类似的操作对于Nestrov momentum SGD。

Local Gradient Clipping: gradient clipping经常用于避免gradient爆炸问题。Pascanu提出了一个解决方案，是将gradient rescale，每当它们的L2 norm超过某一个阈值。

Overcoming The Staleness Effect

因为我们延迟了small gradient update，所以当这些gradient发生的时候，已经是过时或者叫做stale。在我们的实验中，大多数参数每个600到1000个iteration进行一次update。而这种staleness会使得convergence变慢，从而影响模型性能。为了缓解这个问题，使用momentum factor masking和warm-up training。

Momentum Factor Masking: 提出了momentum factor masking。对于gradient $$G_t$$和momentum factor $$U_t$$使用相同的mask（见appendix）

$$U_t[j] \leftarrow U_t[j] \odot \neg Mask$$

Warm-up Training: 先用比较小的learning rate训练几个epoch。原因是训练模型的早期，gradient是非常diverse和aggressive，而使用sparse gradient会限制住model改变的范围，从而使得drastic gradient的变化影响拉长。

Appendix

这里提到另外一篇解决staleness的paper： Asynchrony begets Momentum, with an Application to Deep Learning