具有更低通信复杂度的方差缩减局部 SGD

为加速机器学习模型训练,分布式随机梯度下降(SGD)及其变体已被广泛采用,其利用多个工作节点并行以加快训练。其中,局部 SGD 因其较低通信成本而备受关注。然而,当工作节点上数据分布非独立同分布时,局部 SGD 需要 $O(T^{\frac{3}{4}} N^{\frac{3}{4}})$ 次通信来维持其\emph{线性迭代加速}性质,其中 $T$ 为总迭代次数,$N$ 为工作节点数。本文提出方差缩减局部 SGD (VRL-SGD) 以进一步降低通信复杂度。得益于消除了对工作节点间梯度方差的依赖,我们从理论上证明 VRL-SGD 即使在各工作节点访问非独立同分布数据集时,也能以更低通信复杂度 $O(T^{\frac{1}{2}} N^{\frac{3}{2}})$ 实现\emph{线性迭代加速}。我们在三项机器学习任务上开展实验,实验结果表明当工作节点间数据高度异构时,VRL-SGD 表现明显优于局部 SGD。