Deep Gradient Compression: Reducing the Communication Bandwidth for Distributed Training(DGC)论文解读

简介

Deep Gradient Compression: Reducing the Communication Bandwidth for Distributed Training (arxiv.org)

本文针对当前的分布式机器学习形式提出了以下待解决的困难：

1. bandwidth，梯度信息在节点中传输，而网络带宽不足以满足频繁的传输需求
2. redundant，绝大多数（99.9%）的梯度信息的传输是多余的

针对这两种困难形式，作者提出了Deep Gradient Compression (DGC)，来有效的减少梯度信息交换所需要的网络带宽，在DGC过程中，由于梯度压缩会影响梯度所包含的信息量，进而影响模型准确率，作者提出了4个方法，来减轻梯度压缩对准确率带来的影响.

1. momentum correction.
2. local gradient clipping.
3. momentum factor masking.
4. warm-up training.

前两个方法保证了模型的表现，后两个方法克服了staleness陈旧数据对梯度下降的影响

Deep Gradient Compression (DGC)

Gradient Sparsification (梯度稀疏化)

只发送重要的梯度信息来减少通信所需要的带宽，即只发送大于某个阈值的梯度。为了避免信息的损失造成模型效果减弱，小于阈值的梯度会随着后续训练继续累加，直至超过阈值。

立即发送大梯度，小梯度累计发送

算法流程如下：

1. 输入训练集dataset，每个minibatch的大小（每个节点的minibatch都相同），节点的个数，最优化方法SGD，初始化参数.
2. (Line 1)初始化节点k的梯度为0
3. (Line 2)对于每一轮循环，(Line 3)先将上一次更新所得的梯度值复制给这次即将需要更新的梯度值，(Line 4-7)，对于minibatch中的每个元素\(x\)，利用Line 6的公式计算每个元素的梯度的\(\frac{1}{Nb}\)倍加到总梯度中（Nb是每轮更新中所有N个节点的所有minibatch b中包含的所有数据点的个数）
4. (Line8 - 13)对于参数\(w\)的每个维度，根据本轮梯度中第\(j\)的维度的大小，选择其一定的比例作为阈值. 并将大于阈值的部分保存到\(\widetilde{G}_t^{k}[j]\)中，小于阈值的保存到\(G_t^{k}[j]\)中
5. (Line 14-15)最后，将大于阈值的部分通过编码后进行all-reduce操作保存到当前的全局梯度中，随后根据此全局梯度，使用SGD优化全局参数

如果不懂all-reduce是什么的话可以移步：

Broadcast,Scatter,Gather,Reduce,All-reduce分别是什么？_aiwanghuan5017的博客-CSDN博客

AllReduce_jony0917的博客-CSDN博客_allreduce

已经超过阈值的梯度会被全局模型用来更新，那么未超过阈值的梯度该如何累计呢？

我们定义\(F(w)\)为优化参数\(w\)的损失函数，通过SGD（随机梯度下降）对模型参数进行优化更新，公式如下：

\[F(w) = \frac{1}{\mathcal{X}}\sum_{x \in \mathcal{X}}f(x,w), \quad w_{t+1} = w_t - \eta\frac{1}{Nb}\sum_{k=1}^{N}\sum_{x \in \mathcal{B}_{k,t}} \nabla f(x,w_t) \tag{1} \]

\(F(w)\)为训练数据集\(\mathcal{X}\)中所有数据的损失的平均值，\(w_{t+1}\)通过计算每个节点\(k\)上的每个minibatch的所有数据的平均值进行更新。

对于权重\(w^{(i)}\)，在第T轮迭代后，有

\[w_{t+T}^{(i)} = w^{(i)}_{t} - \eta T \cdot \frac{1}{NbT}\sum_{k=1}^{N}( \sum_{\tau = 0}^{T-1} \sum_{x \in \mathcal{B}_{k,t}}\nabla ^{(i)} f(x,w_{t+\tau})) \tag{2} \]

从公式中可以看出局部梯度积累可以看作是将批次量的大小从Nb增加到NbT (对\(\tau\)的二次求和), 其中T是两次迭代之间的稀疏更新间隔的长度，即每进行T次迭代就发送一次\(w(i)\)的梯度。学习率缩放是处理大量小批量下降的常用技术, \(\eta T\)和\(NbT\)中的\(T\)将会自动抵消

Improving The Local Gradient Accumulation

如果不注意稀疏性，当稀疏度极高时，稀疏更新会极大地危害收敛

作者提出Momentum Correction与Local Gradient Clipping来解决这个问题

1. Momentum Correction

Momentum SGD经常被用来代替普通的SGD.但是算法1不直接适用于带有动量项的SGD,因为它忽略了稀疏更新间隔之间的折扣因子

Momentum ：动量。模拟的是物理中物体运动的惯性，更新时需要保留一部分当前的运动信息，利用当前的梯度微调最终的更新方向，可以在一定程度上增加稳定性，从而学习地更快，并且还有一定摆脱局部最优的能力。

那我们如何将算法1改进为Momentum SGD呢？我们先从普通的\(momentum-SGD\)来观察。N个节点的普通\(momentum-SGD\)算法如下：

\[u_t = mu_{t-1} + \sum_{k=1}^{N}(\nabla_{k,t}),\quad w_{t+1} = w_{t} - \eta u_{y}\tag{3} \]

其中，\(\nabla_{k,t} = \frac{1}{Nb}\sum_{x \in \mathcal{B}_{k,t}} \nabla f(x,w_t)\)，让我们考虑\(w^{(i)}\)在T轮更新后的值：

\[w_{t+T}^{(i)} = w^{(i)}_{t} - \eta\begin{bmatrix} \cdots + (\sum_{\tau=0}^{T-2}m^{\tau})\nabla_{k,t+1}^{(i)} + (\sum_{\tau=0}^{T-1}m^{\tau})\nabla_{k,t}^{(i)} \end{bmatrix} \tag{4} \]

如果上式直接应用到稀疏梯度场景下，即(Line 15 in Algorithm 1)，则更新就不等价与公式3，如下：

\[v_{k,t} = v_{k,t-1} + \nabla_{k,t},\quad u_t = mu_{t-1} + \sum_{k=1}^{N}sparse(v_{k,t}),\quad w_{t+1} = w_t - \eta u_t\tag{5} \]

\(v_{k,t}\)为节点k的local gradient accumulation，当\(v_{k,t}\)大于阈值时，会通过sparse函数并encode后，通过第二个公式项计算梯度，并由第三个公式更新。更新公式为：

\[w_{t+T}^{(i)} = w_{t}^{(i)} - \eta (\cdots + \nabla_{k,t+1}^{(i)} +\nabla_{k,t}^{(i)})\tag{6} \]

较于公式(4)的区别在于取消掉了累积折扣因素 \(\sum_{\tau}^{T-1} m^\tau\) 导致收敛性能的损失，如图a所示，C为公式6的结果，B为公式4的结果

不使用动量修正的局部梯度梯度积累会导致改变最后的优化方向，为了避免这种错误,我们需要动量修正上方的公式(5), 确保稀疏更新等价于公式(3)中的密集更新，结果如B所示，公式如下：

\[u_{k,t} = mu_{k,t-1} + \nabla_{k,t},\quad v_{k,t} = v_{k,t-1} + u_{k,t},\quad w_{t+1} = w_t - \eta \sum_{k=1}^{N}sparse(v_{k,t}) \tag{7} \]

2. Local Gradient Clipping

为了避免梯度爆炸问题，使用梯度裁剪 gradient clipping. 该方法会在梯度的L2范数之和超过某一阈值时对梯度进行重缩放, 这一步通常在所有节点的梯度聚合之后执行。我们用当前梯度除以\(N^{-\frac{1}{2}}\)进行缩放.

Overcoming The Staleness Effect

因为我们延迟了小梯度的更新，当这些更新发生时，可能经过了几百至数千轮的训练过程，那么当时的数据对于现在来说就不具有时效性了，会降低收敛速度，降低模型性能。本文使用Momentum Factor Masking与Warm-up training来解决。

1. Momentum Factor Masking

对公式7中的累计梯度\(v_{k,t}\)与momentum factor \(u_{k,t}\)也使用Mask掩码.

\[Mask \leftarrow |v_{k,t}| > thr,\quad v_{k,t} \leftarrow v_{k,t} \odot \neg Mask, \quad u_{k,t} \leftarrow u_{k,t} \odot \neg Mask \]

这个mask阻止stale梯度的动量，防止陈旧的动量携带权重(weight)在错误的方向

2. Warm-up training

在训练的早期阶段，网络正在迅速变化，梯度更加多样化，且梯度值非常大。稀疏梯度限制了模型的变化范围，从而延长了网络剧烈变化的周期。同时，前期剩余的大梯度在被选择为下一次更新之前被累积，因此可能会超过最新的梯度，从而误导优化方向。

在大型小批量训练中引入的热身训练方法是有帮助的。

在热身阶段，我们使用较低的学习率来减缓神经网络在训练开始时的变化速度，并使用较低的梯度稀疏性来减少被延迟的极端梯度的数量。