0x00 Preface

本篇把 Part 31 的 mini-batch SGD 分析拓展到分布式场景。默认 $f$ 为欧氏空间上的 $L$-光滑、$\mu$-强凸函数，最优解存在，$f^{\ast}=f(w^{\ast})$。$m$ 为 worker 数，$b$ 为每个 worker 的 batch 大小，$\sigma^2$ 为单样本梯度噪声的方差上界。

本文的分布式界以条件无偏、噪声方差界和相应独立性为前提。对于不同 worker 持有不同分布的数据、按完成速度选择 worker，或计算时间依赖样本的情形，必须重新验证这些条件。机器更多或更新更快，本身不能保证收敛界成立。

0x01 Distributed Synchronous SGD

所有 worker 使用同一个 $w_t$，各自计算

$$ g_{i,t}=\frac1b\sum_{r=1}^b\nabla\ell(w_t;\xi_{i,t,r}),\qquad \bar g_t=\frac1m\sum_{i=1}^mg_{i,t},\qquad w_{t+1}=w_t-\eta\bar g_t. $$

给定采样前历史 $\mathcal F_t$，假设 worker 的梯度对全局目标条件无偏、噪声相互独立，且每个 worker 的噪声方差不超过 $\sigma^2/b$。因此

$$ \mathbb E[\bar g_t\mid\mathcal F_t]=\nabla f(w_t),\qquad \mathbb E[\|\bar g_t-\nabla f(w_t)\|^2\mid\mathcal F_t]\le\frac{\sigma^2}{bm}. $$

defination

对于 $0<\eta\le1/L$，令 $\mathcal E_t=\mathbb E[f(w_t)]-f^{\ast}$，有
$$ \mathcal E_T\le C_m+(1-\eta\mu)^T(\mathcal E_0-C_m), \qquad C_m=\frac{\eta L\sigma^2}{2\mu bm}. $$

证明就是在 Part 3 的一步递推中使用 $\sigma^2/(bm)$，再展开几何级数。固定步长时，这是向误差上界平台收缩的保证，不是到零误差的线性收敛保证。

K-Sync and K-Batch-Sync

K-Sync 接收最快的 $K\le m$ 个 worker 的梯度，并取消其余任务；K-batch-sync 接收在同一参数上计算的 $K$ 个 batch，可以让快速 worker 连续贡献多个 batch。

若入选之后的 batch 仍条件无偏，且噪声条件独立或互不相关，则将方差替换为 $\sigma^2/(bK)$，误差上界平台为 $C_K=\eta L\sigma^2/(2\mu bK)$。时间与采样独立、每个 worker 都从同一全局分布采样，是支持这些条件的一种理想模型。

如果快慢与样本类别相关，选择最快的 K 个可能引入偏差。非 IID worker 的局部梯度也不一定是全局梯度的无偏估计。因此不能只把 $m$ 换成 $K$ 而跳过假设检查。

0x02 Distributed Asynchronous SGD

异步更新接受一个 worker 的梯度，不等待所有 worker：

$$ w_{j+1}=w_j-\eta v_j,\qquad v_j=g(w_{\tau_j};\xi_j),\qquad0\le\tau_j\le j. $$

$\tau_j$ 是该梯度计算时读取的参数版本，$d_j=j-\tau_j$ 是滞后步数。$\xi_j$ 表示一个大小为 $b$ 的 batch；模型版本与采样 batch 都可能随机。

Assumptions for the Proof

为使下面的证明逐步成立，明确采用以下条件。

  1. 条件噪声模型。 令 $\mathcal H_j$ 包含本次更新前的服务器参数历史、入选 worker、旧参数和延迟信息，但不包含待应用梯度的样本噪声。写 $a_j=\nabla f(w_{\tau_j})$，$v_j=a_j+e_j$，要求
$$ \mathbb E[e_j\mid\mathcal H_j]=0,\qquad \mathbb E[\|e_j\|^2\mid\mathcal H_j]\le\frac{\sigma^2}{b}. $$

$w_j$ 和 $w_{\tau_j}$ 对 $\mathcal H_j$ 可测。这个条件比只给定旧参数时无偏更明确；若延迟或入选过程泄露了样本信息，需要另行证明它成立。

  1. 相对梯度滞后界。 存在 $0\le\gamma\le1$，使每轮都有
$$ \mathbb E\|\nabla f(w_j)-a_j\|^2 \le\gamma\mathbb E\|\nabla f(w_j)\|^2. $$

这不是 $d_j\le d_{\max}$ 的延迟步数界，也不由光滑性自动推出。特别是在当前梯度很小时，这个条件可能很强。

  1. 新鲜梯度概率下界。 令 $\mathcal P_j$ 为本次入选延迟揭示前的历史，其中 $w_j$ 可测。假设存在 $p_0\in[0,1]$，使
$$ \Pr(\tau_j=j\mid\mathcal P_j)\ge p_0. $$

于是

$$ \mathbb E\|a_j\|^2 \ge\mathbb E\left[\mathbf1_{\{\tau_j=j\}}\|\nabla f(w_j)\|^2\right] \ge p_0\mathbb E\|\nabla f(w_j)\|^2. $$

这一步先在给定历史下使用概率下界，再取期望，避免把随机条件概率直接移到期望之外。

Convergence Bound

defination

令 $\gamma'=1-\gamma+p_0/2>0$。在以上条件与 $0<\eta\le1/(2L)$ 下，
$$ \mathcal E_T\le C_{\rm async}+(1-\eta\mu\gamma')^T(\mathcal E_0-C_{\rm async}), \qquad C_{\rm async}=\frac{\eta L\sigma^2}{2\mu b\gamma'}. $$

证明：在 $\mathcal H_j$ 下取条件期望，噪声交叉项消失。再取总期望，光滑性给出

$$ \mathbb E[f(w_{j+1})]-\mathbb E[f(w_j)] \le-\eta\mathbb E\langle\nabla f(w_j),a_j\rangle +\frac{L\eta^2}{2}\mathbb E\|a_j\|^2 +\frac{L\eta^2\sigma^2}{2b}. $$

使用 $2\langle u,a\rangle=\Vert u\Vert^2+\Vert a\Vert^2-\Vert u-a\Vert^2$，得到

$$ \begin{aligned} \mathbb E[f(w_{j+1})]-\mathbb E[f(w_j)] &\le-\frac\eta2(1-\gamma)\mathbb E\|\nabla f(w_j)\|^2 -\frac\eta2(1-L\eta)\mathbb E\|a_j\|^2 +\frac{L\eta^2\sigma^2}{2b}\\ &\le-\frac\eta2\left(1-\gamma+\frac{p_0}{2}\right) \mathbb E\|\nabla f(w_j)\|^2+\frac{L\eta^2\sigma^2}{2b}. \end{aligned} $$

第二步用了 $1-L\eta\ge1/2$ 和新鲜概率下界。再用 PL 不等式：

$$ \mathcal E_{j+1}\le(1-\eta\mu\gamma')\mathcal E_j+\frac{L\eta^2\sigma^2}{2b}. $$

因为 $\mu\le L$、$\gamma’\le3/2$，步长限制保证收缩系数位于 $[0,1)$。展开递推即得定理。当 $\gamma’=0$ 时不能除以它，也不能从此推导得到该几何收缩界。

Interpreting Staleness

在同质 worker、独立无偏移指数服务时间、服务器处理和模型读取瞬时完成的理想模型中，下一次完成者在 $m$ 个 worker 中等概率出现。对于单梯度异步更新，新鲜梯度概率可取 $p_0=1/m$。

稳定运行后的滞后步数 $d_j$ 有几何尾部；绝对版本索引 $\tau_j$ 不是几何分布。启动阶段受 $0\le d_j\le j$ 限制，需要单独考虑。对于异构时间、偏移指数或一般服务分布，不直接使用这个概率结论。

0x03 K-Async SGD

K-Async 每次更新汇总 K 个不同 worker 的梯度；本轮完成的 worker 等待该次更新后读取新参数，尚未完成的 worker 保留旧计算。更新为

$$ w_{j+1}=w_j-\frac\eta K\sum_{i=1}^K v_{i,j},\qquad v_{i,j}=\frac1b\sum_{r=1}^b\nabla\ell(w_{\tau_{i,j}};\xi_{i,j,r}). $$

给定包含所有入选旧参数和延迟的 $\mathcal H_j$，写 $v_{i,j}=a_{i,j}+e_{i,j}$，要求噪声条件无偏、条件互不相关，且各自方差不超过 $\sigma^2/b$。因此平均噪声方差不超过 $\sigma^2/(bK)$。

另外要求平均相对滞后界和每个入选梯度的新鲜概率下界：

$$ \frac1K\sum_{i=1}^K\mathbb E\|\nabla f(w_j)-a_{i,j}\|^2 \le\gamma\mathbb E\|\nabla f(w_j)\|^2, $$ $$ \Pr(\tau_{i,j}=j\mid\mathcal P_j)\ge p_0,\qquad i=1,\ldots,K. $$

记 $\bar a_j=K^{-1}\sum_i a_{i,j}$。由 Jensen，$\Vert\bar a_j\Vert^2\le K^{-1}\sum_i\Vert a_{i,j}\Vert^2$。对每个 $a_{i,j}$ 使用内积恒等式后取平均，得到

$$ \begin{aligned} \mathbb E[f(w_{j+1})]-\mathbb E[f(w_j)] &\le-\frac\eta2(1-\gamma)\mathbb E\|\nabla f(w_j)\|^2 -\frac\eta{2K}(1-L\eta)\sum_i\mathbb E\|a_{i,j}\|^2 +\frac{L\eta^2\sigma^2}{2bK}\\ &\le-\frac\eta2\gamma'\mathbb E\|\nabla f(w_j)\|^2 +\frac{L\eta^2\sigma^2}{2bK}. \end{aligned} $$

因此在同样的步长限制和 $\gamma’>0$ 条件下，

$$ \mathcal E_T\le C_{K\text{-async}}+(1-\eta\mu\gamma')^T (\mathcal E_0-C_{K\text{-async}}),\qquad C_{K\text{-async}}=\frac{\eta L\sigma^2}{2\mu bK\gamma'}. $$

这里需要分别验证每个 K 对应的 $\gamma$ 和 $p_0$，不能把它们默认为与 K 无关，也不能把单梯度情形的 $p_0=1/m$ 无条件照搬。K-batch-async 允许一个 worker 连续贡献多个 batch；它有不同的运行时间模型，噪声与滞后条件也需要重新检查。

0x04 What the Analysis Does Not Guarantee

同步模型中的 $1/m$ 和异步模型中的 $1/K$ 方差缩减都依赖噪声结构。异步的相对梯度滞后界是附加条件；只给出延迟步数或更快的更新速度不能替代它。这些强凸结论也不直接覆盖一般非凸神经网络和非 IID 联邦学习。

如何把迭代界与随机耗时结合，见 分布式机器学习中的效率分析2。比较时必须同时考虑 batch 大小、误差上界平台、步长和每轮耗时。

Reference

本文使用单样本方差一致有界的简化模型，并把证明所需的条件噪声假设显式写出；原论文允许更一般的梯度依赖方差模型，不能省略其额外步长条件后直接照搬。

Sources: 3 4 5.

  1. Part 3 ↩︎

  2. 分布式机器学习中的效率分析 ↩︎

  3. Dutta, S., Joshi, G., Ghosh, S., Dube, P. and Nagpurkar, P., Slow and Stale Gradients Can Win the Race: Error-Runtime Trade-offs in Distributed SGD, AISTATS 2018, especially Theorem 3 and Supplement, Section 8. ↩︎

  4. Dutta, S., Joshi, G., Ghosh, S., Dube, P. and Nagpurkar, P., Slow and Stale Gradients Can Win the Race: Error-Runtime Trade-offs in Distributed SGD, AISTATS 2018, especially Theorem 3 and Supplement, Section 8. ↩︎

  5. Joshi, G., Optimization Algorithms for Distributed Machine Learning, Springer. ↩︎