Escaping Saddle Points with Compressed SGD

TL;DR

This paper extends the analysis of compressed stochastic gradient descent (SGD) to guarantee convergence to second-order stationary points, improving communication efficiency in distributed machine learning.

Contribution

It provides the first convergence guarantees of compressed SGD to second-order stationary points, including cases with non-Lipschitz gradients, and demonstrates improved communication efficiency.

Findings

Converges to $oldsymbol{ ext{ε}}$-SOSP with compressed SGD.

Achieves same iteration complexity as uncompressed SGD.

Reduces total communication by a factor of $ ilde ext{Θ}(rac{ ext{√d}}{ ext{ε}^{3/4}})$.

Abstract

Stochastic gradient descent (SGD) is a prevalent optimization technique for large-scale distributed machine learning. While SGD computation can be efficiently divided between multiple machines, communication typically becomes a bottleneck in the distributed setting. Gradient compression methods can be used to alleviate this problem, and a recent line of work shows that SGD augmented with gradient compression converges to an $\varepsilon$-first-order stationary point. In this paper we extend these results to convergence to an $\varepsilon$-second-order stationary point ($\varepsilon$-SOSP), which is to the best of our knowledge the first result of this type. In addition, we show that, when the stochastic gradient is not Lipschitz, compressed SGD with RandomK compressor converges to an $\varepsilon$-SOSP with the same number of iterations as uncompressed SGD [Jin et al.,2021] (JACM), while improving the total communication by a factor of $\tilde \Theta(\sqrt{d} \varepsilon^{-3/4})$, where $d$ is the dimension of the optimization problem. We present additional results for the cases when the compressor is arbitrary and when the stochastic gradient is Lipschitz.