[Paper Review] Sparsified SGD with Memory

Sebastian U. Stich, Jean-Baptiste Cordonnier+1 more|arXiv (Cornell University)|Sep 20, 2018

Stochastic Gradient Optimization Techniques233 citations

TL;DR

The paper analyzes SGD with k-sparsification and an error-compensation memory, proving it converges at the same rate as vanilla SGD while drastically reducing communication.

ABSTRACT

Huge scale machine learning problems are nowadays tackled by distributed optimization algorithms, i.e. algorithms that leverage the compute power of many devices for training. The communication overhead is a key bottleneck that hinders perfect scalability. Various recent works proposed to use quantization or sparsification techniques to reduce the amount of data that needs to be communicated, for instance by only sending the most significant entries of the stochastic gradient (top-k sparsification). Whilst such schemes showed very promising performance in practice, they have eluded theoretical analysis so far. In this work we analyze Stochastic Gradient Descent (SGD) with k-sparsification or compression (for instance top-k or random-k) and show that this scheme converges at the same rate as vanilla SGD when equipped with error compensation (keeping track of accumulated errors in memory). That is, communication can be reduced by a factor of the dimension of the problem (sometimes even more) whilst still converging at the same rate. We present numerical experiments to illustrate the theoretical findings and the better scalability for distributed applications.

Motivation & Objective

• Motivate reducing communication in distributed SGD for large-scale learning.
• Introduce a memory-based error-compensation mechanism with sparsified updates.
• Prove convergence of SGD with memory under k-contraction operators.
• Show that top-k and random-k sparsification can match vanilla SGD rates.
• Demonstrate practical scalability through experiments on multi-core setups.

Proposed method

• Define k-contraction compression operators that reduce communicated updates.
• Introduce memory m_t to accumulate suppressed gradient information and inject it in subsequent iterations.
• Formulate Mem-SGD with updates g_t = comp_k(m_t + eta_t ∇f_i_t(x_t)) and memory update m_{t+1} = m_t + eta_t ∇f_i_t(x_t) - g_t.
• Prove convergence bounds for the averaged iterate x̄_T with specific step sizes eta_t = 8 / (mu(a+t)).
• Show that for T = Ω(d/k * sqrt(κ)) the method achieves the same rate as vanilla SGD, where κ = L/μ.
• Optionally extend to a parallel Mem-SGD variant with shared memory for multi-core setups.

Experimental results

Research questions

• RQ1 Does SGD with k-sparsification and memory retain the convergence rate of vanilla SGD?
• RQ2 How does memory compensation affect variance and convergence under top-k and rand-k operators?
• RQ3 What are the theoretical and practical impacts of the compression level k on convergence and communication?
• RQ4 Can Mem-SGD scale in parallel/shared-memory settings without sacrificing rate?

Key findings

• Mem-SGD converges at the same rate as vanilla SGD under 0<k≤d with appropriate stepsizes.
• Convergence rate expression includes terms that scale with d^2/k^2 and d^3/k^3, but with proper averaging the rate matches SGD.
• Top-k sparsification generally outperforms rand-k in sequential experiments, improving communication efficiency significantly.
• Parallel Mem-SGD achieves linear speedups in multi-core experiments, maintaining convergence while reducing communication.
• Empirical results show Mem-SGD requires orders of magnitude fewer bits transmitted than QSGD for similar accuracy.
• For dense data, top-1 sparsification reduces communication by ~10^3x compared to SGD; for sparse data, top-10 yields ~10x reduction.