[Paper Review] Stochastic Gradient Push for Distributed Deep Learning

Mahmoud Assran, Nicolas Loizou+2 more|arXiv (Cornell University)|Nov 27, 2018

Stochastic Gradient Optimization TechniquesComputer Science137 citations

TL;DR

The paper introduces Stochastic Gradient Push (SGP), which combines PushSum gossip with stochastic gradient updates for distributed deep learning, proving convergence to a stationary point and node consensus, with empirical validation on large-scale vision and translation tasks.

ABSTRACT

Distributed data-parallel algorithms aim to accelerate the training of deep neural networks by parallelizing the computation of large mini-batch gradient updates across multiple nodes. Approaches that synchronize nodes using exact distributed averaging (e.g., via AllReduce) are sensitive to stragglers and communication delays. The PushSum gossip algorithm is robust to these issues, but only performs approximate distributed averaging. This paper studies Stochastic Gradient Push (SGP), which combines PushSum with stochastic gradient updates. We prove that SGP converges to a stationary point of smooth, non-convex objectives at the same sub-linear rate as SGD, and that all nodes achieve consensus. We empirically validate the performance of SGP on image classification (ResNet-50, ImageNet) and machine translation (Transformer, WMT'16 En-De) workloads. Our code will be made publicly available.

Motivation & Objective

• Address the inefficiencies of exact distributed averaging in synchronous methods under stragglers and delays.
• Develop a distributed optimization algorithm that achieves consensus while performing stochastic gradient updates.
• Provide theoretical guarantees for convergence to stationary points in non-convex settings and sub-linear convergence rates.
• Demonstrate practical effectiveness on large-scale computer vision and natural language processing workloads.

Proposed method

• Integrate PushSum-based gossip with stochastic gradient updates to form the Stochastic Gradient Push (SGP) algorithm.
• Model the network with augmented delay-aware mixing matrices to account for communication delays and virtual nodes.
• Provide a global viewpoint of the SGP update using augmented matrices and de-biased parameter tracking across delayed nodes.
• Prove convergence to a stationary point of smooth, non-convex objectives and consensus among all nodes with sub-linear rates.
• Validate empirically on ResNet-50 with ImageNet and Transformer with WMT’16 En-De workloads.

Experimental results

Research questions

• RQ1 Can Stochastic Gradient Push achieve consensus across distributed nodes while performing stochastic gradient updates in the presence of delays?
• RQ2 What convergence guarantees does SGP offer for non-convex objectives compared to standard SGD or fully synchronous methods?
• RQ3 Is the communication-delay-robustness of PushSum preserved when combined with stochastic gradients in a distributed deep learning setting?
• RQ4 Do practical benchmarks (vision and translation tasks) show competitive performance relative to existing asynchronous or synchronous distributed methods?

Key findings

• SGP converges to a stationary point of smooth, non-convex objectives at the same sub-linear rate as SGD.
• All nodes achieve consensus under the proposed delay-aware augmented framework.
• Empirical validation demonstrates effectiveness on image classification (ResNet-50/ImageNet) and machine translation (Transformer/WMT'16 En-De) workloads.