0x00 Preface#
本文着眼于分布式机器学习中的收敛性分析，作为第四部分，主要介绍了分布式学习中的收敛性证明，参考主要还是。同时我们把假设拓展到分布式的 setting 下。
Gradient Unbiased Estimate Assumption#
$$
\mathbb{E}[g(w;\xi)] = \nabla f(w)
$$
Gradient Bounded Variance Assumption#
$$
\mathrm {Var}(\nabla f(w;\xi)) \leq \sigma^2
$$
由 $\frac{1}{m} \sum_{i=1}^{m} g\left(w, \xi_{i}\right)$ 为 $\nabla f(w)$ 的可得
$$
\mathrm {Var}\left(\frac{1}{m} \sum_{i=1}^{m} g\left(w, \xi_{i}\right)\right) = \mathbb{E}\left[\left(\frac{1}{m} \sum_{i=1}^{m} g\left(w, \xi_{i}\right)\right)^{2}\right] - \mathbb{E}\left[\frac{1}{m} \sum_{i=1}^{m} g\left(w, \xi_{i}\right)\right]^{2} \leq \frac{\sigma^{2}}{bm}
$$
即
$$
\mathbb{E}_{\xi}[||g(w; \xi)||^2] \leq ||\nabla f(w)||^{2} + \frac{\sigma^2}{bm}
$$
0x01 Distributed Synchronous Stochastic Gradient Descent#
先介绍一下几个分布式的优化方法，首先就是同步的分布式随机梯度下降，其实就是在随机梯度下降上的扩展。假设我们有 $m$ 个 worker node，在 Parameter Server Framework 就叫做 synchronous SGD。迭代过程如下：
$$
\begin{aligned}
w_{t+1} &=w_{t}-\frac{\eta}{m} \sum_{i=1}^{m} g(w_{t}, \xi_{i}) \\
&=w_{t}-\frac{\eta}{m} \sum_{i=1}^{m} \frac{1}{b} \sum_{j=1}^{b} \nabla f\left(w_{t}, \xi_{i, j}\right)
\end{aligned}
$$
对于同步的 SGD，我们有如下定理：
defination
假设目标函数 $f$ 是 $R^{d}$ 上的$\mu$-强凸函数，并且 $L$-光滑，同时满足梯度无偏估计和梯度方差有界假设，当步长 $\eta \leq \frac{1}{L}$ 且起始点为 $w_0$ 时，经过 $t$ 步 Distributed Synchronous mini-batch SGD 迭代，$\mathbb{E}[f(w_{t})]$ 被 bounded 如下：
$$
\mathbb{E}[f(w_{t})] - f(w^{*}) - \frac{\eta L \sigma^{2}}{2\mu bm} \leq (1- \eta \mu)^{t}\left (\mathbb{E}[f(w_{0})] - f(w^{*}) - \frac{\eta L \sigma^{2}}{2\mu bm} \right )
$$
其实根据上面假设也可以看出来，基本就是在 SGD 的结论上多加了一个 $m$，把 $b$ 替换成了 $bm$ 即可，这里就不进行证明了，读者可以自行推导。
除此之外，还有 K-Synchronous SGD, K-batch-Synchronous SGD，证明基本雷同，将 $m$ 替换成 $K$ 即可。
0x02 Distributed Asynchronous Stochastic Gradient Descent#
同步的情况下，往往根据木桶原理需要等待最慢的，如果有一个 worker node 出现了故障，那么整个系统就会被阻塞。虽然 $K-Sync$ 和 $K-batch-sync$ 能够在一定程度上缓解，但也存在浪费算力的问题（因为需要取消其他的计算任务、白白消耗了大量资源）。如果采用异步的方式，那么就能够避免这种情况，但是代价则是滞后的梯度（因为不取消其他计算任务意味着并不是从相同的参数起开始计算梯度的，the cost of gradient staleness）
这一部分则是本文的重点，先说假设依旧是 $L-smooth$ 和 $\mu$-strongly convex，同时满足梯度无偏估计和梯度方差有界假设。但是这里的梯度无偏估计和梯度方差有界假设和之前的定义略有不同
defination
Unbiased Stochastic Gradients: 随机梯度是真实梯度的无偏估计
$$
\mathbb{E}_{\xi_{j}|w_{k}}[g(w_{k};\xi_{j})] = \nabla f(w_{k}), \forall k \leq j
$$
相较于之前的定义 $\mathbb{E}_{\xi_{j}}[g(w;\xi_{j})] = \nabla f(w)$ 更加严格了，因为 $k<j$ 的时候并非是独立同分布的，只有当前的迭代轮次达到 $j$ 之后，或者说能够 access 到 $\xi_{j}$ 后才是独立同分布的。
defination
Bounded Gradient Variance: 梯度方差有界
$$
\mathbb{E}_{\xi_{j}|w_{k}}[||g(w_{k};\xi_{j})-\nabla f(w_{k})||^2] \leq \frac{\sigma^2}{b}, \forall k \leq j
$$
除此之外，我们还需要一个滞后性的假设，即
defination
Bounded Staleness Assumption: 我们假设 $\gamma \leq 1$
$$
\mathbb{E}[||\nabla f(w_{j})-\nabla f(w_{\tau(j)})||^2] \leq \gamma \mathbb{E}[||\nabla f(w_{j})||^2]
$$
其中 $\tau(j)$ 表示 worker 收到的模型版本的 index，$\tau(j)\leq j$，是个随机变量。同时，$\tau(j)$ 受一个参数 $p_{0}$ 影响，在迭代中接收新梯度的概率的下界。
defination
假设 $p_{0}^{(j)}$ 是 $\tau(j)=j$ 的条件概率，在给定过去所有延迟和参数的情况下，$p_{0}\leq p_{0}^{(j)}$，有
$$
\mathbb{E}[||\nabla f(w_{\tau(j)})||^2] \geq p_{0} \mathbb{E}[||\nabla f(w_{j})||^2]
$$
Proof:
$$
\begin{align}
\mathbb{E}[||\nabla f(w_{\tau(j)})||^2] &= p_{0}^{(j)} \mathbb{E}[||\nabla f(w_{\tau(j)})||^2|\tau(j)=j] + (1-p_{0}^{(j)}) \mathbb{E}[||\nabla f(w_{\tau(j)})||^2|\tau(j)\neq j] \\
&\geq p_{0} \mathbb{E}[||\nabla f(w_{j})||^2]
\end{align}
$$
如果每个 worker 计算梯度的时间服从指数分布，则 $\tau(j)$ 是服从几何分布的（geometrically distributed），并且 $p_{0}=\frac{1}{m}$。这是因为，当服从指数分布是，PS在一个时间窗口内收到的梯度个数是服从泊松分布的。
根据以上定义和假设，我们有
defination
假设目标函数 $f$ 是 $R^{d}$ 上的$\mu$-强凸函数，并且 $L$-光滑，同时满足梯度无偏估计、梯度方差有界和滞后性假设，当步长 $\eta \leq \frac{1}{2L}$ 且起始点为 $w_0$ 时，经过 $t$ 步 Distributed Asynchronous mini-batch SGD 迭代，$\mathbb{E}[f(w_{t})]$ 被 bounded 如下：
$$
\mathbb{E}[f(w_{t})] - f(w^{*}) - \frac{\eta L \sigma^{2}}{2\mu b{\gamma}'} \leq (1- \eta \mu{\gamma}')^{t}\left (\mathbb{E}[f(w_{0})] - f(w^{*}) - \frac{\eta L \sigma^{2}}{2\mu b{\gamma}'} \right )
$$
其中 ${\gamma}' = 1-\gamma+\frac{p_{0}}{2}$
Proof:
由 $L$-smooth 可得
$$
\begin{align}
f(w_{j+1}) &\leq f(w_{j}) + (w_{j+1}-w_{j})^{\top}\nabla f(w_{j}) + \frac{L}{2}||w_{j+1}-w_{j}||^2 \\
&= f(w_{j}) - \eta g(w_{\tau(j)})^{\top}\nabla f(w_{j}) + \frac{L\eta^2}{2}||g(w_{\tau(j)})||^2 \\
&= f(w_{j}) - \frac{\eta}{2} ||\nabla f(w_{j})||^2 - \frac{\eta}{2} ||g(w_{\tau(j)})||^2 + \frac{\eta}{2} ||\nabla f(w_{j})-g(w_{\tau(j)})||^2 + \frac{L\eta^2}{2}||g(w_{\tau(j)})||^2
\end{align}
$$
最后一行由 $2a^{\top}b = ||a||^2 + ||b||^2 - ||a-b||^2$ 得到。由于 $g(w_{\tau(j)})$ 是无偏估计，对两边求期望得
$$
\begin{align}
\mathbb{E}[f(w_{j+1})] - \mathbb{E}[f(w_{j})] &\leq - \frac{\eta}{2} \mathbb{E}[||\nabla f(w_{j})||^2] - \frac{\eta}{2} \mathbb{E}[||g(w_{\tau(j)})||^2] + \frac{\eta}{2} \mathbb{E}[||\nabla f(w_{j})-g(w_{\tau(j)})||^2] + \frac{L\eta^2}{2}\mathbb{E}[||g(w_{\tau(j)})||^2] \\
&= - \frac{\eta}{2} \mathbb{E}[||\nabla f(w_{j})||^2] - \frac{\eta}{2} \mathbb{E}[||g(w_{\tau(j)})||^2] + \frac{\eta}{2} \mathbb{E}[||\nabla f(w_{j})-\nabla f(w_{\tau(j)}) +\nabla f(w_{\tau(j)}) - g(w_{\tau(j)})||^2] + \frac{L\eta^2}{2}\mathbb{E}[||g(w_{\tau(j)})||^2] \\
&= - \frac{\eta}{2} \mathbb{E}[||\nabla f(w_{j})||^2] - \frac{\eta}{2} \mathbb{E}[||g(w_{\tau(j)})||^2] + \frac{L\eta^2}{2}\mathbb{E}[||g(w_{\tau(j)})||^2] + \frac{\eta}{2} \mathbb{E}[||\nabla f(w_{j})-\nabla f(w_{\tau(j)})||^2] + \frac{\eta}{2} \mathbb{E}[||\nabla f(w_{\tau(j)}) - g(w_{\tau(j)})||^2] \\
&= - \frac{\eta}{2} \mathbb{E}[||\nabla f(w_{j})||^2] - \frac{\eta}{2} \mathbb{E}[||g(w_{\tau(j)})||^2] + \frac{L\eta^2}{2}\mathbb{E}[||g(w_{\tau(j)})||^2] + \frac{\eta}{2} \mathbb{E}[||\nabla f(w_{j})-\nabla f(w_{\tau(j)})||^2] + \frac{\eta}{2} (\mathbb{E}[||\nabla f(w_{\tau(j)})||^2] - 2\mathbb{E}[\nabla f(w_{\tau(j)})^{\top}g(w_{\tau(j)})] + \mathbb{E}[||g(w_{\tau(j)})||^2]) \\
&= - \frac{\eta}{2} \mathbb{E}[||\nabla f(w_{j})||^2] - \frac{\eta}{2} \mathbb{E}[||g(w_{\tau(j)})||^2] + \frac{L\eta^2}{2}\mathbb{E}[||g(w_{\tau(j)})||^2] + \frac{\eta}{2} \mathbb{E}[||\nabla f(w_{j})-\nabla f(w_{\tau(j)})||^2] + \frac{\eta}{2} \mathbb{E}[||g(w_{\tau(j)})||^2] - \frac{\eta}{2} \mathbb{E}[||\nabla f(w_{\tau(j)})||^2] \\
&= - \frac{\eta}{2} \mathbb{E}[||\nabla f(w_{j})||^2] - \frac{\eta}{2} \mathbb{E}[||\nabla f(w_{\tau(j)})||^2] + \frac{L\eta^2}{2}\mathbb{E}[||g(w_{\tau(j)})||^2] + \frac{\eta}{2} \mathbb{E}[||\nabla f(w_{j})-\nabla f(w_{\tau(j)})||^2] \\
&\leq - \frac{\eta}{2} \mathbb{E}[||\nabla f(w_{j})||^2] - \frac{\eta}{2} \mathbb{E}[||\nabla f(w_{\tau(j)})||^2] + (\frac{L\eta^2\sigma^{2}}{2b} + \frac{L\eta^2}{2}\mathbb{E}[||\nabla f(w_{\tau(j)})||^2]) + \frac{\eta \gamma}{2} \mathbb{E}[||\nabla f(w_{j})||^2] \\
&= \frac{\eta}{2}(1-\gamma) \mathbb{E}[||\nabla f(w_{j})||^2] + \frac{L\eta^2\sigma^{2}}{2b} - \frac{\eta}{2}(1-\eta L) \mathbb{E}[||\nabla f(w_{\tau(j)})||^2] \\
&\leq \frac{\eta}{2}(1-\gamma) \mathbb{E}[||\nabla f(w_{j})||^2] + \frac{L\eta^2\sigma^{2}}{2b} - \frac{\eta}{4}p_{0}\mathbb{E}[||\nabla f(w_{j})||^2] \\
&= -\frac{\eta}{2}(1-\gamma+\frac{p_{0}}{2}) \mathbb{E}[||\nabla f(w_{j})||^2] + \frac{L\eta^2\sigma^{2}}{2b}
\end{align}
$$
后面，就是利用强凸性
$$
2\mu (f(w) - f(w^{*})) \leq ||\nabla f(w)||^2
$$
有
$$
\mathbb{E}[f(w_{j})] - f(w^{*}) - \frac{\eta L \sigma^{2}}{2b} \leq (1- \eta \mu(1-\gamma+\frac{p_{0}}{2}))^{t}\left (\mathbb{E}[f(w_{j})] - f(w^{*}) \right )
$$
令 ${\gamma}' = 1-\gamma+\frac{p_{0}}{2}$，再配平（？），累乘，即可得到结论。
$$
\mathbb{E}[f(w_{j})] - f(w^{*}) - \frac{\eta L \sigma^{2}}{2\mu b{\gamma}'} \leq (1- \eta \mu{\gamma}')^{t}\left (\mathbb{E}[f(w_{0})] - f(w^{*}) - \frac{\eta L \sigma^{2}}{2\mu b{\gamma}'} \right )
$$
0x03 K-Async SGD#
此时，
$$w_{j+1}=w_{j}-\frac{\eta}{K} \sum_{i=1}^{K} g(w_{\tau(k,j)};\xi_{k,j})$$
其中，$k=1,2,\cdots,K$ 表示在当前轮贡献的 worker 索引，$\xi_{k,j}$ 表示第 $k$ 个 worker 在第 $j$ 轮迭代中的样本，$\tau(k,j)$ 表示第 $k$ 个 worker 在第 $j$ 轮迭代中的模型版本的 index，$\tau(k,j)\leq j$，是个随机变量。同时，$g(w_{\tau(k,j)};\xi_{k,j}) = \frac{1}{b} \sum_{\xi \in \xi_{k,j}} \nabla f(w_{\tau(k,j)};\xi_{k,j})$。
此时结论也是略有变动，${\gamma}' = 1-\gamma+\frac{p_{0}}{2}$
$$
\mathbb{E}[f(w_{j})] - f(w^{*}) - \frac{\eta L \sigma^{2}}{2\mu b{\gamma}'K} \leq (1- \eta \mu{\gamma}')^{t}\left (\mathbb{E}[f(w_{0})] - f(w^{*}) - \frac{\eta L \sigma^{2}}{2\mu b{\gamma}'K} \right )
$$
Reference#
