现代大规模机器学习应用程序要求在分布式计算体系结构上实现随机优化算法，一个关键瓶颈是在不同worker之间交换信息（例如随机梯度）的通信开销。 在本文中，为了降低通信成本，我们提出了一种凸优化公式，以最小化随机梯度的编码长度。 关键思想是随机删除随机梯度向量的坐标，并适当地放大其余的坐标，以确保稀疏的梯度无偏。为了有效地解决最优稀疏性问题，提出了一种简单快速的近似解算法，并为稀疏性提供了理论上的保证。在实验阶段，我们使用了L2正则化逻辑回归，支持向量机和卷积神经网络的实验来验证稀疏化方法。

大规模机器学习需要我们将随机优化算法扩展到分布式计算架构或多核系统中，在同步随机梯度算法中，每个worker利用本地数据训练模型得到更新梯度后将梯度广播给其他worker进行聚合更新全局模型，这个步骤需要重复多次知道模型实现收敛。在这个过程中，worker之间的通信问题会严重影响优化算法的速度。
在本篇论文中，我们提出一个新的算法，该算法通过稀疏化随机梯度以减少通信成本，而迭代次数只增加了一点点，该算法的核心思想是我们会随即丢弃随机梯度中的某些梯度，而对于其余的梯度我们会进行适当地缩放以保证无偏性。该稀疏化方式会大幅减少随机梯度的编码长度，并只会给随机梯度增加少量的方差。
本文提出了一个凸公式来实现方差和稀疏性的最优权衡：给定任何固定方差预算，都可以得到每个样本坐标的最优概率来决定是否丢弃该坐标。为了在线性时间内求解该优化问题，提出了几种在稀疏性保证下寻找近似最优解的有效算法。

给定一个数据集 { x n } n = 1 N \{x_n\}_{n=1}^N {xn​}n=1N​和 N N N个损失函数 { f n } n = 1 N \{f_n\}_{n=1}^N {fn​}n=1N​，每个损失函数会将对应的数据样本映射到损失值，我们使用 w ∈ R d w\in \mathbb{R}^d w∈Rd来表示训练的模型的参数，并使用随机优化的算法来解决下述问题：
m i n w f ( w ) = 1 N f n ( w ) ,   w t + 1 = w t − η t g t ( w t ) min_w f(w)=\frac{1}{N}f_n(w),\,w_{t+1}=w_t-\eta_tg_t(w_t) minw​f(w)=N1​fn​(w),wt+1​=wt​−ηt​gt​(wt​)
其中的 t t t代表着迭代的次数，而 g t ( w t ) g_t(w_t) gt​(wt​)则是对于正式梯度 ∇ f ( w t ) \nabla f(w_t) ∇f(wt​)的无偏估计，我们可以选择使用SGD或者SVRG算法来计算 g t ( w t ) g_t(w_t) gt​(wt​)。
接下来我们考虑如何通过稀疏化梯度 g t ( w t ) g_t(w_t) gt​(wt​)，定义为 Q ( g ( w t ) ) Q(g(w_t)) Q(g(wt​))来减少机器学习过程中的通信费用，并且保证 Q ( g ( w t ) ) Q(g(w_t)) Q(g(wt​))是无偏的以及产生的方差是相对较小的。为了方便后续描述，我们使用 g g g来表示通过SGD或SVRG训练得到的梯度， g ∈ R d g\in \mathbb{R}^d g∈Rd，而 g i g_i gi​代表梯度中的第 i i i个元素。我们会通过算法获得一个概率向量 p ∈ R d p\in \mathbb{R}^d p∈Rd，其第 i i i个元素 p i p_i pi​代表第 i i i个梯度不进行丢弃的概率，并且定义一个二值向量 Z ∈ { 0 , 1 } d Z\in\{0,1\}^d Z∈{0,1}d， Z i Z_i Zi​为1的概率有 p i p_i pi​，为了保证稀疏化后的梯度的无偏性，对于不被丢弃的梯度要进行缩放，将其修改为 g i → g i / p i g_i \rightarrow g_i/p_i gi​→gi​/pi​，因此稀疏化算法具体如下所述：
Q ( g ) = [ Z 1 g 1 p 1 , Z 2 g 2 p 2 , . . . , Z d g d p d ] Q(g)=[Z_1\frac{g_1}{p_1},Z_2\frac{g_2}{p_2},...,Z_d\frac{g_d}{p_d}] Q(g)=[Z1​p1​g1​​,Z2​p2​g2​​,...,Zd​pd​gd​​]
因为 E [ Q ( g ) i ] = p i ∗ g i p i + ( 1 − p i ) ∗ 0 = g i \mathbb{E}[Q(g)_i]=p_i*\frac{g_i}{p_i}+(1-p_i)*0=g_i E[Q(g)i​]=pi​∗pi​gi​​+(1−pi​)∗0=gi​，因此该稀疏化能实现对原梯度的无偏估计。我们定义总共有M个worker在分布式地训练机器学习模型，其训练过程如算法1所示。

尽管上述稀疏化算法能减少通信消耗，但会增加梯度的方差，这会影响模型收敛的速度，我们需要权衡方差和稀疏性，稀疏化后的梯度的方差为
E ∑ i = 1 d [ Q ( g ) i 2 ] = ∑ i = 1 d [ g i 2 p i 2 ∗ p i ] = ∑ i = 1 d g i 2 p i \mathbb{E}\sum_{i=1}^d[Q(g)_i^2]=\sum_{i=1}^d[\frac{g_i^2}{p_i^2}*p_i]=\sum_{i=1}^d\frac{g_i^2}{p_i} Ei=1∑d​[Q(g)i2​]=i=1∑d​[pi2​gi2​​∗pi​]=i=1∑d​pi​gi2​​
此外，该稀疏化算法能实现的稀疏性的期望为 E [ ∣ ∣ Q ( g ) ∣ ∣ 0 ] = ∑ i = 1 d p i \mathbb{E}[||Q(g)||_0]=\sum_{i=1}^dp_i E[∣∣Q(g)∣∣0​]=∑i=1d​pi​，因此我们可以将平衡稀疏性和方差的问题公式化为
m i n p   ∑ i = 1 d p i s . t . ∑ i = 1 d g i 2 p i ≤ ( 1 + ϵ ) ∑ i = 1 d g i 2 min_p \, \sum_{i=1}^dp_i \qquad s.t.\sum_{i=1}^d\frac{g_i^2}{p_i}\le (1+\epsilon)\sum_{i=1}^dg_i^2 minp​i=1∑d​pi​s.t.i=1∑d​pi​gi2​​≤(1+ϵ)i=1∑d​gi2​
其中 0 < p i ≤ 1 0<p_i\le 1 0<pi​≤1，并且参数 ϵ \epsilon ϵ用来控制随机梯度的方差的增加。
Proposition   1. \textbf{Proposition 1.} Proposition 1.上述优化问题的解为一个概率向量，其中每个元素 p i = m i n ( λ ∣ g i ∣ , 1 ) p_i=min(\lambda|g_i|,1) pi​=min(λ∣gi​∣,1)，其中的 λ > 0 \lambda>0 λ>0，由 g g g和 ϵ \epsilon ϵ决定。
证明：通过引入拉格朗日乘子 λ \lambda λ和 μ i \mu_i μi​，我们可以得到
min ⁡ p max ⁡ λ max ⁡ μ L ( p i , λ , μ i ) = ∑ i = 1 d p i + λ 2 ( ∑ i = 1 d g i 2 p i − ( 1 + ϵ ) ∑ i = 1 d g i 2 ) + ∑ i = 1 d μ i ( p i − 1 ) \mathop {\min }\limits_p \mathop{\max }\limits_{\lambda} \mathop{\max}\mu L(p_i,\lambda,\mu_i)=\sum_{i=1}^dp_i+\lambda^2(\sum_{i=1}^d\frac{g_i^2}{p_i}-(1+\epsilon)\sum_{i=1}^dg_i^2)+\sum_{i=1}^d\mu_i(p_i-1) pmin​λmax​maxμL(pi​,λ,μi​)=i=1∑d​pi​+λ2(i=1∑d​pi​gi2​​−(1+ϵ)i=1∑d​gi2​)+i=1∑d​μi​(pi​−1)
根据KKT条件，我们可以得到 ∂ L ∂ p i = 0 \frac{\partial L}{\partial p_i}=0 ∂pi​∂L​=0，因此，有
1 − λ 2 g i 2 p i 2 + μ i = 0 1-\lambda^2\frac{g_i^2}{p_i^2}+\mu_i=0 1−λ2pi2​gi2​​+μi​=0
因为 μ i ( p i − 1 ) = 0 \mu_i(p_i-1)=0 μi​(pi​−1)=0，所以此时 p i = 1 p_i=1 pi​=1或者 μ i = 0 \mu_i=0 μi​=0，当 μ i = 0 \mu_i=0 μi​=0时可以得到 p i = λ ∣ g i ∣ p_i=\lambda|g_i| pi​=λ∣gi​∣。
根据上述描述，我们定义一个集合 S S S，若 p j = 1 p_j=1 pj​=1则 j ∈ S j\in S j∈S，因为 p i p_i pi​随着 g i g_i gi​的增大而增大，所以 S S S中的梯度下标必定是绝对值最大的那部分梯度的下标，我们对梯度排序后得到 g ( 1 ) , g ( 2 ) , . . . , g ( d ) g_{(1)},g_{(2)},...,g_{(d)} g(1)​,g(2)​,...,g(d)​，定义 ∣ S ∣ = k |S|=k ∣S∣=k，则代表着排序后的前 k k k个梯度必不会被丢弃。

在本节中我们将介绍如何来求解概率向量 p p p，因为 λ > 0 \lambda>0 λ>0，所以可以得到
∑ i = 1 d g i 2 p i − ( 1 + ϵ ) ∑ i = 1 d g i 2 = ∑ i = 1 k g ( i ) 2 + ∑ i = k + 1 d ∣ g ( i ) ∣ λ − ( 1 + ϵ ) ∑ i = 1 d g i 2 = 0 \sum_{i=1}^d\frac{g_i^2}{p_i}-(1+\epsilon)\sum_{i=1}^dg_i^2=\sum_{i=1}^kg_{(i)}^2+\sum_{i=k+1}^d\frac{|g_{(i)}|}{\lambda}-(1+\epsilon)\sum_{i=1}^dg_i^2=0 i=1∑d​pi​gi2​​−(1+ϵ)i=1∑d​gi2​=i=1∑k​g(i)2​+i=k+1∑d​λ∣g(i)​∣​−(1+ϵ)i=1∑d​gi2​=0
可以得到
λ = ( ϵ ∑ i = 1 d g i 2 + ∑ i = k + 1 d g ( i ) 2 ) − 1 ∑ i = k + 1 d ∣ g ( i ) ∣ \lambda=(\epsilon\sum_{i=1}^dg_i^2+\sum_{i=k+1}^dg_{(i)}^2)^{-1}\sum_{i=k+1}^d|g_{(i)}| λ=(ϵi=1∑d​gi2​+i=k+1∑d​g(i)2​)−1i=k+1∑d​∣g(i)​∣
并且，因为 λ ∣ g ( k + 1 ) ∣ < 1 \lambda|g_{(k+1)}|<1 λ∣g(k+1)​∣<1，所以可以得到：
∣ g ( k + 1 ) ∣ ( ∑ i = k + 1 d ∣ g ( i ) ∣ ) < ϵ ∑ i = 1 d g i 2 + ∑ i = k + 1 d g ( i ) 2 |g_{(k+1)}|(\sum_{i=k+1}^d|g_{(i)}|)<\epsilon\sum_{i=1}^dg_i^2+\sum_{i=k+1}^dg_{(i)}^2 ∣g(k+1)​∣(i=k+1∑d​∣g(i)​∣)<ϵi=1∑d​gi2​+i=k+1∑d​g(i)2​
我们应该寻找最小的 k k k值满足上述过程，因此可以得到算法

在算法2中，我们需要对梯度进行排序，这样计算代价比较大，因此我们设计一个贪心算法来近似求解这个过程。我们先预定义一个稀疏参数 κ ∈ ( 0 , 1 ) \kappa\in(0,1) κ∈(0,1)，并限制稀疏密度 ∑ i = 1 d p i / d ≈ κ \sum_{i=1}^dp_i/d\approx \kappa ∑i=1d​pi​/d≈κ，我们首先初始化 p i ~ = κ d ∣ g i ∣ / ∑ i = 1 d ∣ g i ∣ \widetilde{p_i}=\kappa d|g_i|/\sum_{i=1}^d|g_i| pi​ ​=κd∣gi​∣/∑i=1d​∣gi​∣，之后我们通过截断操作令 p i = min ⁡ ( p i ~ , 1 ) p_i=\mathop{\min}(\widetilde{p_i},1) pi​=min(pi​ ​,1)。使得系数密度小于 κ \kappa κ，之后我们固定概率为1的那部分概率，对于其余的概率迭代进行上述过程，具体如下所述，其中的 κ d − d + ∣ L ∣ \kappa d-d+|\mathcal{L}| κd−d+∣L∣代表去除概率为1那部分后其余部分的稀疏期望， ∑ i ∈ L p i j \sum_{i\in \mathcal{L}}p_i^j ∑i∈L​pij​则是代表那部分非1概率此时的稀疏期望，对这部分进行缩放。

我们接下来需要将稀疏化后的结果 Q ( g ) Q(g) Q(g)进行编码传输，我们先定义每个浮点数需要使用b比特来进行表示，我们使用两个向量， Q A ( g ) Q_A(g) QA​(g)存放着概率为1的下标， Q B ( g ) Q_B(g) QB​(g)存放着其他的下标。 Q A ( g ) Q_A(g) QA​(g)需要使用 l o g   d log\, d logd来表示索引以及 b b b来表示该索引的浮点数。而 Q B ( g ) Q_B(g) QB​(g)中的 Q i ( g ) = g i / p i = g i / ( λ ∣ g i ∣ ) = s i g n ( g i ) / λ Q_i(g)=g_i/p_i=g_i/(\lambda|g_i|)=sign(g_i)/\lambda Qi​(g)=gi​/pi​=gi​/(λ∣gi​∣)=sign(gi​)/λ，因此对于这部分数据我们只需要传输一个浮点数 1 / λ 1/\lambda 1/λ，并且使用 l o g   d log\, d logd表示索引及1bit表示该索引的符号。

在本节中我们将分布稀疏性的期望值 E [ ∣ ∣ Q ( g ) ∣ ∣ 0 ] = ∑ i = 0 d p i \mathbb{E}[||Q(g)||_0]=\sum_{i=0}^dp_i E[∣∣Q(g)∣∣0​]=∑i=0d​pi​，在实际情况中，当梯度向量中的值分布很不均衡时，稀疏算法能实现很好的效果。我们首先定义一下符号
Definition   2. \textbf{Definition 2.} Definition 2.一个梯度向量 g ∈ R d g\in \mathbb{R}^d g∈Rd如果存在一个子集 S S S，使得| S S S|= s s s，且 ∣ ∣ g S c ∣ ∣ 1 ≤ ρ ∣ ∣ g S ∣ ∣ 1 ||g_{S^c}||_1\le \rho||g_S||_1 ∣∣gSc​∣∣1​≤ρ∣∣gS​∣∣1​，其中 S c S^c Sc为集合 S S S的补集，则称该向量是( ρ , s \rho, s ρ,s)稀疏。
( ρ , s \rho, s ρ,s)稀疏用来衡量一个向量信息集中在一个大小为 s s s的子集的程度，我们可以用(1- ρ \rho ρ) s s s用来建立稀疏期望的界限，并且(1- ρ \rho ρ) s ≤ d s\le d s≤d，若向量值的分布高度不均衡，则(1- ρ \rho ρ) s ≪ d s\ll d s≪d。
Lemma   3. \textbf{Lemma 3.} Lemma 3.如果梯度向量 g g g是( ρ , s \rho, s ρ,s)稀疏，则我们可以将之前的参数 ϵ \epsilon ϵ定义为 ρ \rho ρ，并且稀疏期望的上界为 E [ ∣ ∣ Q ( g ) ∣ ∣ 0 ] ≤ ( 1 + ρ ) s \mathbb{E}[||Q(g)||_0]\le (1+\rho)s E[∣∣Q(g)∣∣0​]≤(1+ρ)s。
证明：令 ϵ = ρ \epsilon=\rho ϵ=ρ以及 S k = S S_k=S Sk​=S
E [ ∣ ∣ Q ( g ) ∣ ∣ 0 ] = ∑ i = 1 d p i = ∑ i ∈ S k p i + ∑ i ∉ S k p i = s + ∑ i ∉ S k ∣ g i ∣ ( ∑ j = k + 1 d ∣ g ( j ) ∣ ) ϵ ∑ j = 1 k g ( i ) 2 + ( 1 + ϵ ) ∑ j = k + 1 d g ( j ) 2 ( p i = λ ∣ g i ∣ ) = s + ∣ ∣ g S k c ∣ ∣ 1 2 ρ ∣ ∣ g S k ∣ ∣ 2 2 + ( 1 + ρ ) ∣ ∣ g S k c ∣ ∣ 2 2 = s + ρ 2 ∣ ∣ g S k ∣ ∣ 1 2 ρ ∣ ∣ g S k ∣ ∣ 2 2 + ( 1 + ρ ) ∣ ∣ g S k c ∣ ∣ 2 2 ( ∣ ∣ g S c ∣ ∣ 1 ≤ ρ ∣ ∣ g S ∣ ∣ 1 ) = s + ρ 2 s ∣ ∣ g S k ∣ ∣ 2 2 ρ ∣ ∣ g S k ∣ ∣ 2 2 + ( 1 + ρ ) ∣ ∣ g S k c ∣ ∣ 2 2 ( ( a + b ) 2 ≤ 2 ( a 2 + b 2 ) ) ≤ ( 1 + ρ ) s \begin{aligned} \mathbb{E}[||Q(g)||_0] &= \sum_{i=1}^dp_i\\ &= \sum_{i\in S_k}p_i+\sum_{i \notin S_k}p_i \\ &=s+\sum_{i \notin S_k}\frac{|g_i|(\sum_{j=k+1}^d|g_{(j)}|)}{\epsilon\sum_{j=1}^kg_{(i)}^2+(1+\epsilon)\sum_{j=k+1}^dg_{(j)}^2}\quad(p_i=\lambda|g_i|)\\ &= s+\frac{||g_{S_k^c}||_1^2}{\rho||g_{S_k}||_2^2+(1+\rho)||g_{S_k^c}||_2^2} \\ &=s+\frac{\rho^2||g_{S_k}||_1^2}{\rho||g_{S_k}||_2^2+(1+\rho)||g_{S_k^c}||_2^2}\quad(||g_{S^c}||_1\le \rho||g_S||_1)\\ &= s+\frac{\rho^2s||g_{S_k}||_2^2}{\rho||g_{S_k}||_2^2+(1+\rho)||g_{S_k^c}||_2^2}\quad((a+b)^2\le2(a^2+b^2))\\ &\le (1+\rho)s \end{aligned} E[∣∣Q(g)∣∣0​]​=i=1∑d​pi​=i∈Sk​∑​pi​+i∈/​Sk​∑​pi​=s+i∈/​Sk​∑​ϵ∑j=1k​g(i)2​+(1+ϵ)∑j=k+1d​g(j)2​∣gi​∣(∑j=k+1d​∣g(j)​∣)​(pi​=λ∣gi​∣)=s+ρ∣∣gSk​​∣∣22​+(1+ρ)∣∣gSkc​​∣∣22​∣∣gSkc​​∣∣12​​=s+ρ∣∣gSk​​∣∣22​+(1+ρ)∣∣gSkc​​∣∣22​ρ2∣∣gSk​​∣∣12​​(∣∣gSc​∣∣1​≤ρ∣∣gS​∣∣1​)=s+ρ∣∣gSk​​∣∣22​+(1+ρ)∣∣gSkc​​∣∣22​ρ2s∣∣gSk​​∣∣22​​((a+b)2≤2(a2+b2))≤(1+ρ)s​
Remark   1. \textbf{Remark 1.} Remark 1.Lemma 3表明稀疏化后梯度方差的增加受到因子 ( 1 + ρ ) (1+\rho) (1+ρ)的影响，稀疏期望值低于 ( 1 + ρ ) s (1+\rho)s (1+ρ)s，为了实现稀疏化后模型的精度不受影响，则需要增加的通信次数的因子为 ( 1 + ρ ) (1+\rho) (1+ρ)，因此所通信的浮点数的总数减少了 ( 1 + ρ ) 2 s / d (1+\rho)^2s/d (1+ρ)2s/d。
Theorem   4. \textbf{Theorem 4.} Theorem 4.假设梯度 g g g是( ρ , s \rho, s ρ,s)稀疏，且每个浮点数需要使用 b b b比特表示，则压缩后需要上传的数据量上限为 s ( b + l o g   d ) + ρ s l o g   d + b s(b+log\,d)+\rho slog\,d+b s(b+logd)+ρslogd+b。

在本篇论文中，针对分布式机器学习的通信问题，提出了一个稀疏化的算法，通过该算法来减轻通信压力，我们对这个稀疏化算法进行了详细的介绍，并列举了相关算法来求解需要的概率向量。之后，我们分析了使用该算法后造成的梯度方差，为了缓解这个方差我们需要增加一定的迭代次数，分析迭代次数增加后实现的通信效率。