Escaping Saddle Points with Compressed SGD

additional results for the cases when the compressor is arbitrary and when the stochastic gradient is Lipschitz . 1 Introduction . Stochastic Gradient Descent ( SGD ) and its variants are the main workhorses of modern machine learning . Distributed implementations of SGD on a cluster of machines with a central server and a large number of workers are frequently used in practice due to the massive size of the data . In distributed SGD each machine holds a copy of the model and the computation proceeds in rounds . In every round , each worker finds a stochastic gradient based on its batch of examples , the server averages these stochastic gradients to obtain the gradient of the entire batch , makes an SGD step , and broadcasts the updated model parameters to the workers . With a large number of workers , computation parallelizes efficiently while communication becomes the main bottleneck [ 12 , 38 ] , since each worker needs to send its gradients to the server and receive the updated model parameters . Common solutions for this problem include : local SGD and its variants , when each machine performs multiple local steps before communication [ 36 ] ; decentralized architectures which allow pairwise communication between the workers [ 30 ] and gradient compression , when a compressed version of the gradient is communicated instead of the full gradient [ 6 , 37 , 27 ] . In this work , we consider the latter approach , which we refer to as compressed SGD . Most machine learning models can be described by a d-dimensional vector of parameters x and the model quality can be estimated as a function f ( x ) . Hence optimization of the model parameters can be cast a minimization problem minx f ( x ) , where f : Rd → R is a continuous function , which can be optimized using continuous optimization techniques , such as SGD . Fast convergence of compressed SGD to a first-order stationary point ( FOSP , ‖∇f ( x ) ‖ < ε ) was shown recently for various gradient compression schemes [ 6 , 37 , 27 , 23 , 2 ] . However , even an exact FOSP can be either a local minimum , a saddle point or a local maximum . While local minima often correspond to good solutions in machine learning applications [ 21 , 39 , 7 ] , saddle points and local maxima are always suboptimal and 35th Conference on Neural Information Processing Systems ( NeurIPS 2021 ) . it is important for an optimization algorithm to avoid converging to them . In particular , [ 13 ] show that for neural networks many local minima are almost optimal , but the corresponding loss functions have a combinatorial explosion in the number of saddle points . Furthermore , [ 15 ] show that saddle points can significantly slow down SGD convergence and hence it is important to be able to escape from them efficiently . Since finding a local minimum is NP-hard in general [ 5 ] , a common relaxation of this requirement is to find an approximate second-order stationary point ( SOSP ) , i.e . a point with a small gradient norm ( ‖∇f ( x ) ‖ < ε ) and the smallest ( negative ) eigenvalue being small in absolute value ( λmin ( ∇2f ( x ) ) > −εH ) . When f has ρ-Lipschitz Hessian ( i.e . ‖∇2f ( x ) −∇2f ( y ) ‖ ≤ ρ‖x− y‖ for all x , y ) , a standard choice of εH is √ ρε [ 32 ] , and such approximate SOSP is commonly referred as an ε-SOSP . While second-order optimization methods allow one to escape saddle points , such methods are typically substantially more expensive computationally . A line of work originating with the breakthrough of [ 20 ] shows that first-order methods can escape saddle points when perturbations are added at certain iterations . In particular , a follow-up [ 25 ] show that SGD converges to an ε-SOSP in an almost optimal number of iterations . In this paper , we show that even compressed SGD can efficiently converge to an ε-SOSP . To the best of our knowledge , this is the first result showing convergence of compressed methods to a second-order stationary point . 1.1 Related Work . Escaping from saddle points While it is known that gradient descent with random initialization converges to a local minimum almost surely [ 28 ] , existence of saddle points may result in exponential number of steps with non-negligible probability [ 17 ] . Classical approaches for escaping from saddle points assume access to second-order information [ 32 , 14 ] . Although these algorithms find a secondorder stationary point in O ( ε−3/2 ) iterations , each iteration requires computation of the full Hessian matrix , which can be prohibitive for high-dimensional problems in practice . Some approaches relax this requirement , and instead of full Hessian matrix they only require access to a Hessian-vector product oracle [ 9 , 1 ] . While in certain settings , including training of neural networks , it ’ s possible to compute Hessian-vector products ( HVP ) efficiently [ 33 , 35 ] , such an oracle might not be available in general . Furthermore , in practice HVP-based approaches are significantly more complex compared to SGD ( especially if the workers aren ’ t communicating in every iteration in the distributed setting ) and require additional hyperparameter tuning . Moreover , HVP is typically used for approximate an eigenvector computation , which may increase the number of iterations by a logarithmic factor . Limitations of second-order methods motivate a long line of recent research on escaping from saddle points using first-order algorithms , starting from [ 20 ] . [ 24 ] show that perturbed gradient descent finds ε-SOSP in Õ ( ε−2 ) iterations . Later , this is improved by a series of accelerated algorithms [ 10 , 1 , 11 , 26 ] which achieves Õ ( ε−7/4 ) iteration complexity . There are also a number of algorithms designed for finite sum setting where f ( x ) = ∑n i=1 fi ( x ) [ 34 , 4 , 18 ] , or in case when only stochastic gradients are available [ 40 , 25 ] , including variance reduction techniques [ 3 , 18 ] . The sharpest rates in these settings have been obtained by [ 18 ] , [ 42 ] and [ 19 ] . Compressed SGD While gradient compression may require a complex communication protocol , from theoretical perspective this process is often treated as a black-box function : a ( possibly randomized ) function C is called a µ-compressor if E [ ‖x− C ( x ) ‖2 ] < ( 1− µ ) ‖x‖2 . In a simplified form , the update step in compressed SGD can be expressed as xt+1 ← xt − ηC ( ∇f ( x ) ) 1. Notable examples of compressors include the following : SIGN [ 6 ] sends the sign of each coordinate , QUANTIZATION [ 2 ] splits all possible values into buckets and communicates bucket indices , TOPK preserves only k largest ( by the absolute value ) coordinates while setting other coordinates to 0 , and RANDOMK preserves k random coordinates . Among them , it ’ s not clear how to implement SIGN and TOPK in distributed settings ; [ 23 ] suggest an efficient sketch-based approximate implementation of TOPK . Both RANDOMK and sketch-based TOPK are k/d-compressors requiring Õ ( k ) communication . 1The actual update equation is more complicated , see Algorithm 1 . While it is known that compressed SGD converges ( e.g . [ 27 ] and the works above ) , the convergence was shown only to a FOSP . The crucial idea to facilitate convergence is to use error-feedback [ 37 ] : the difference between the true gradient and its compression is propagated to the next iteration . 1.2 Our Contributions . Our main contribution is the analysis showing that perturbed compressed SGD with error-feedback can escape from saddle points efficiently . Moreover , we show faster convergence rate for a certain type of compressors and show that such compressors exist . Inspired by the ideas from [ 25 ] and [ 37 ] , we present an algorithm ( Algorithm 1 ) which uses perturbed compressed gradients with error-feedback and converges to an ε-second-order stationary point ( see Theorem 3.1 ) . Our main results shows that compressed SGD with RANDOMK compressor achieves substantial communication improvement:2 Theorem 1.1 ( Informal , Theorem 3.2 and Corollary 3.3 ) Assume that f has Lipschitz gradient and Lipschitz Hessian . Let α = 1 when the stochastic gradient is Lipschitz and α = d otherwise . Then SGD with RANDOMK compressor ( which selects k random coordinates ) with k = dε 3/4 √ α converges to an ε-SOSP after Õ ( α ε4 ) iterations , with Õ ( d √ α ε3+1/4 ) total communication per worker . Compared with the uncompressed case , the total communication improves by ε−1/4 when the stochastic gradient is Lipschitz and by √ dε−3/4 otherwise ( the sharpest results for SGD are by [ 19 ] and [ 25 ] respectively ) . In Theorem 1.1 , we heavily rely on the following property of RANDOMK : when its randomness ( i.e . sampled k coordinates ) is fixed , the compressor becomes a linear function . For other compressors , this property doesn ’ t necessarily hold ; in this case , we show convergence with a slower convergence rate : Theorem 1.2 ( Informal , Theorem 3.1 and Corollary 3.4 ) Assume that f has Lipschitz gradient and Lipschitz Hessian . Let α = 1 when the stochastic gradient is Lipschitz and α = d otherwise . Let C be a k/d-compressor requiring Õ ( k ) communication . Then SGD with compressor C with k = d √ dε 3/4 √ α converges to an ε-SOSP after Õ ( α ε4 ) iterations , with Õ ( d √ d √ α ε3+1/4 ) total communication per worker . Compared with the uncompressed case , the total communication improves by ε −1/4 √ d when the stochastic gradient is Lipschitz ( note that this is the only setting where the convergence improvement is conditional , requiring ε = o ( d−2 ) ) and by ε−3/4 otherwise . Table 1 in Section 3.1 shows communication improvements for various choices of compression parameters . We outline our main techniques and technical contributions in Section 3.2 and give the complete proof in Appendix A and B . 2 Preliminaries . Function Properties For a twice differentiable nonconvex function f : Rd → R , we consider the unconstrained minimization problem minx∈Rd f ( x ) . We use the following standard [ 25 , 19 , 41 , 3 , 43 ] assumptions about the objective function f : Assumption A f is fmax-bounded , L-smooth and has ρ-Lipschitz Hessian , i.e . for all x , y : |f ( x ) − f ( y ) | ≤ fmax , ‖∇f ( x ) −∇f ( y ) ‖ ≤ L‖x− y‖ , ‖∇2f ( x ) −∇2f ( y ) ‖ ≤ ρ‖x− y‖ Assumption B Access to an unbiased stochastic gradient oracle ∇F ( x , θ ) , whose randomness is controlled by a parameter θ ∼ D3 , with bounded variance , i.e . for all x : Eθ∼D [ ∇F ( x , θ ) ] = ∇f ( x ) , Eθ∼D [ ‖∇F ( x , θ ) −∇f ( x ) ‖2 ] ≤ σ2 As shown by the above works , smoothness allows one to achieve fast convergence for nonconvex optimization problems ( namely , to use the folklore descent lemma ) . Similarly , Lipschitz Hessian 2Õ hides polylogarithmic dependence dependence on d and ε and polynomial dependence on other parameters . 3E.g . θ is a minibatch selected at the current iteration allows one to show fast second-order convergence , since , within a certain radius , the function stays close to its quadratic approximation ( see e.g . [ 8 , Section 9.5.3 ] ) . As common in the literature , in our convergence rates we treat L , ρ and σ2 as constants . We also consider an additional optional assumption ( see [ 25 ] for a justification ) : Assumption C ( Lipschitz stochastic gradient , optional ) For any x , y , θ : ‖∇F ( x , θ ) −∇F ( y , θ ) ‖ ≤ ˜̀‖x− y‖ . Gradient Compression Our goal is to optimize f in a distributed setting [ 16 , 29 ] : given W workers , for each worker i we have a corresponding data distribution Di . Then the each worker has a corresponding function fi ( x ) = Eθi∼Di [ F ( x , θi ) ] and f = ∑W i=1 fi . In a typical distributed SGD setting , each worker computes a stochastic gradient ∇Fi ( x , θi ) and sends it to the coordinator machine . The coordinator machine computes the average of these gradients v = 1W ∑W i=1∇Fi ( x , θi ) and broadcasts it to the workers , which update the local parameters x← x− ηv ( η is the step size ) . With this approach , with increase of the number of machines , the computation can be perfectly parallelized . However , with each machine required to send its gradient , communication becomes the main bottleneck [ 12 , 38 ] . There exist various solutions to this problem ( see Section 1 ) , including gradient compression , when each machine sends an approximation of its gradient . Then coordinator averages these approximations and broadcasts the average to all machines ( possibly compressing it again , see discussion on TOPK and SIGN in Section 1.1 ) . Depending on the compression method , this protocol results in different gradient approximation and different communication per machine . There is a natural trade-off between approximation and communication , and it ’ s not clear whether having smaller per-iteration communication results in smaller total communication required for convergence . The approximation quality can be formalized using the following definition : Definition 2.1 ( [ 37 ] ) Function C ( x , θ̃ ) , whose randomness is controlled by a parameter θ̃ ∼ D̃4 , is a µ-compressor if Eθ̃∼D̃ [ ‖x− C ( x , θ̃ ) ‖2 ] < ( 1− µ ) ‖x‖2 . Section 1.1 provides examples of important compressors . In our analysis , we consider general and linear compressors separately , and in the latter case , we show an improved convergence rate . Definition 2.2 C is a linear compressor if C ( · , θ̃ ) is a linear function for any θ̃ . One example of a linear compressor is RANDOMK , which sets all but k coordinates of a vector to 0 ; it ’ s a k/d-compressor [ 37 ] and can be computed easily in the distributed setting . Stationary Points The optimization problem of finding a global minimum or even a local minimum is NP-hard for nonconvex objectives [ 31 , 5 ] . Instead , as is standard in the literature , we show convergence to an approximate FOSP or an approximate SOSP , see Section 1 . Definition 2.3 If f is differentiable then x is an ε-First-Order Stationary Point if ‖∇f ( x ) ‖ ≤ ε . An ε-FOSP can be a local maximum , a local minimum or a saddle point . While local minima typically correspond to good solutions , saddle points and local maxima are inherently suboptimal . Assuming non-degeneracy , saddle points and local maxima have escaping directions , corresponding to Hessian ’ s negative eigenvectors . Following [ 32 ] we refer to points with no escape directions ( up to a second-order approximation ) as approximate second-order stationary points : Definition 2.4 ( [ 32 ] ) If f is a twice differentiable ρ-Hessian Lipschitz function then x is an ε-SecondOrder Stationary Point if ‖∇f ( x ) ‖ ≤ ε and λmin ( ∇2f ( x ) ) ≥ − √ ρε , where λmin is the smallest eigenvalue5 . 4E.g . for RANDOMK , θ̃ is the set of indices of preserved coordinates . 5While one can consider two threshold parameters – εg for ∇f and εH for ∇2f – we follow convention of [ 32 ] which selects εH = − √ ρε , which , intuitively , balances first-order and second-order variability . Algorithm 1 : Compressed SGD parameters : η – step size , T – number of iterations , r2 – variance of the artificial noise , I – the number of iterations required for escaping , R – escaping radius input : objective f , compressor function C , starting point x0 output : ε-SOSP of f 1 e0 ← 0d 2 t′ ← 0 3 for t = 0 . . . T − 1 do 4 // Reset the error after I iterations or in case we moved far from the initial point 5 if ( t− t′ > I or ‖xt′ − ( xt − ηet ) ‖ > R then 6 if f ( xt ) < f ( xt′ ) then xt ← xt − ηet else xt ← xt′ ; 7 t′ ← t , et ← 0d 8 end 9 Sample ξt ∼ Nd ( 0d , r2 ) , θt ∼ D , θ̃t ∼ D̃ 10 gt ← C ( et +∇F ( xt , θt ) + ξt , θ̃t ) // Compressed gradient 11 xt+1 ← xt − ηgt // Compressed gradient descent step 12 et+1 ← et +∇F ( xt , θt ) + ξt − gt // Update the error 13 end 14 return xT An important property of points which are not ε-SOSP is that they are unstable : adding a small perturbation allows gradient descent to escape them [ 20 ] ( similar results were shown for e.g . stochastic [ 25 ] and accelerated [ 26 ] gradient descent ) . In this work we show that this property holds even for SGD with gradient compression .

Summary: This paper is concerned with finding local minima, or second order stationary points of a function $\mathbb{R}^d \to \mathbb{R}$. The key contribution of the paper is to combine the ideas of stochastic gradient descent (SGD) and compressed gradient descent (CGD). This leads to an algorithm which takes care of both the convergence rate and communication complexity, and is favorable for distributed computing. Theoretical results of the convergence and communication complexity are rigorously established, and some empirical studies are conducted.