Universal  Stagewise Learning for Non-Convex Problems with  Convergence on  Averaged Solutions

1 INTRODUCTION . Non-convex optimization has recently received increasing attention due to its popularity in emerging machine learning tasks , particularly for learning deep neural networks . One of the keys to the success of deep learning for big data problems is the employment of simple stochastic algorithms such as SGD or ADAGRAD ( Krizhevsky et al. , 2012 ; Dean et al. , 2012 ) . Analysis of these stochastic algorithms for non-convex optimization is an important and interesting research topic , which already attracts much attention from the community of theoreticians ( Ghadimi & Lan , 2013 ; 2016 ; Yang et al. , 2016 ; Davis & Drusvyatskiy , 2018b ; Ward et al. , 2018 ; Li & Orabona , 2018 ) . However , one issue that has been largely ignored in existing theoretical analysis is that the employed algorithms in practice usually differ from their plain versions that are well understood in theory . Below , we will discuss several important heuristics used in practice for training deep neural networks and the gap between the practice and the theory to motivate this work . First , a trick for setting the step size in training deep neural networks is to change it in a stagewise manner from a large value to a small value , i.e. , a constant step size is used in a stage for a number of iterations and is decreased for the next stage ( Krizhevsky et al. , 2012 ; Ren et al. , 2018 ) . Although ∗Equal contribution . Correspondence to : tianbao-yang @ uiowa.edu this trick has been adopted by most open-sourced libraries , e.g. , Caffe ( Jia et al. , 2014 ) , TensorFlow ( Abadi et al. , 2015 ) , Pytorch ( Paszke et al. , 2017 ) , it still lacks theoretical analysis to date for non-convex optimization . A related question that stands out is how and when to decrease the step size . On standard benchmark datasets for academic use such as CIFAR-10 , CIFAR-100 ( Krizhevsky et al . ) , people could follow the setting reported in previous studies to get a good result , which however might not work well for new datasets . Hence , a better solution to setting the stagewise step size with insights from theory would be much preferred . However , in the existing literature of theory for non-convex optimization ( Ghadimi & Lan , 2013 ; Davis & Drusvyatskiy , 2018b ) , only strategies based on an iteratively decreasing step size or a small constant step size have been well analyzed . For example , the existing theory usually suggests an iteratively decreasing step size proportional to 1/ √ t at the t-th iteration or a small constant step size , e.g. , proportional to 2 with 1 for finding an -stationary solution whose gradient ’ s magnitude ( in expectation ) is small than . Second , the averaging heuristic is usually used in practice , i.e. , an averaged solution is returned for prediction ( Bottou , 2010 ) , which could yield improved stability and generalization ( Hardt et al. , 2016 ) . However , existing theory for many stochastic non-convex optimization algorithms only provides guarantee on a uniformly sampled solution or a non-uniformly sampled solution with decreasing probabilities for latest solutions ( Ghadimi & Lan , 2013 ; Yang et al. , 2016 ; Davis & Drusvyatskiy , 2018b ) . In particular , if an iteratively decreasing step size proportional to 1/ √ t at the t-th iteration is employed , the convergence guarantee was provided for a random solution that is non-uniformly selected from all iterates with a sampling probability proportional to 1/ √ t for the t-th iterate . This means that the latest solution always has the smallest probability to be selected as the final solution , which contradicts to the common wisdom . If a small constant step size is used , then usually a uniformly sampled solution is provided with convergence guarantee . However , both options are rarely used in practice and can not justify the heuristic that returns the last solution . A third common approach in practice is to use adaptive coordinate-wise step size of ADAGRAD ( Dean et al. , 2012 ) . Although adaptive step size has been well analyzed for convex problems ( e.g. , why and when it could yield faster convergence than SGD ) ( Duchi et al. , 2011 ; Chen et al. , 2018b ) , it still remains an mystery for non-convex optimization with missing insights from theory . Several recent studies have attempted to analyze ADAGRAD for non-convex problems ( Ward et al. , 2018 ; Li & Orabona , 2018 ; Chen et al. , 2018a ; Zou & Shen , 2018 ) . Nonetheless , none of them are able to exhibit the adaptive convergence of ADAGRAD to data as in the convex case and its potential advantage over SGD for non-convex problems . To overcome the shortcomings of existing theories for stochastic non-convex optimization , this paper analyzes new algorithms that employ some or all of these commonly used heuristics in a systematic framework , aiming to fill the gap between theory and practice . The main results and contributions are summarized below : • We propose a universal stagewise optimization framework for solving a family of non-convex problems , i.e. , weakly convex problems , which is broader than smooth non-convex problems and includes some non-smooth non-convex problems . At each stage , any suitable stochastic convex optimization algorithms ( e.g. , SGD , ADAGRAD ) with a constant step size parameter can be employed for optimizing a regularized convex problem with a number of iterations , which usually return an averaged solution . The step size parameter is decreased in a stagewise manner following a polynomial decaying scheme . • We analyze several variants of the proposed framework by employing different basic algorithms , including SGD , ADAGRAD , stochastic heavy-ball ( SHB ) method , and stochastic Nesterov ’ s accelerated gradient ( SNAG ) method . We prove the convergence of their stagewise versions for an averaged solution that is randomly selected from all stagewise averaged solutions . • To justify a heuristic approach that returns the last averaged solution in stagewise learning , we present and analyze a non-uniform sampling strategy over stagewise averaged solutions with sampling probabilities increasing as the stage number . • Regarding the convergence results , for stagewise SGD , SHB , SNAG , we establish the same order of iteration complexity for finding a nearly stationary point as the existing theories of their nonstagewise variants . For stagewise ADAGRAD , we establish an adaptive convergence for finding a nearly stationary point , which is provably better than ( stagewise ) SGD , SHB , and SNAG when the cumulative growth of stochastic gradient is slow . • Besides theoretical contributions , we also empirically verify the effectiveness of the proposed stagewise algorithms . In particular , our empirical studies show that ( i ) the stagewise ADAGRAD dramatically improves the generalization performance of existing variants of ADAGRAD , ( ii ) stagewise SGD , SHB , SNAG also outperform their plain variants with an iteratively decreasing step size ; ( iii ) the proposed stagewise algorithms achieve similar if not better generalization performance than their heuristic variants implemented in existing libraries on standard benchmark datasets . 2 RELATED WORK . SGD for unconstrained smooth non-convex problems was first analyzed by Ghadimi & Lan ( 2013 ) , who established an O ( 1/ 4 ) iteration complexity for finding an -stationary point x in expectation satisfying E [ ‖∇f ( x ) ‖ ] ≤ , where f ( · ) denotes the objective function . As mentioned earlier , the returned solution is either a uniformly sampled solution or a non-uniformly sampled one with sampling probabilities proportional to the decreasing step size . Similar results were established for the stochastic momentum variants of SGD ( i.e. , SHB , SNAG ) by Yang et al . ( 2016 ) ; Ghadimi & Lan ( 2016 ) . Recently , SGD was also analyzed for ( constrained ) weakly convex problems , whose objective function is non-convex and not necessarily smooth , by Davis & Drusvyatskiy ( 2018b ) . However , none of these studies provide results for algorithms that return an averaged solution , and these analyzed algorithms differ significantly from that used in practice for achiving the state-of-the-art results ( Krizhevsky et al. , 2012 ; Ren et al. , 2018 ; Loshchilov & Hutter , 2017 ) . Although adaptive variants of SGD , e.g. , ADAGRAD ( Duchi et al. , 2011 ) , ADAM ( Kingma & Ba , 2015 ; Reddi et al. , 2018 ) , were widely used for training deep neural networks , there are few studies on theoretical analysis of these algorithms for non-convex problems . Several recent studies attempted to analyze ADAGRAD for non-convex problems ( Ward et al. , 2018 ; Li & Orabona , 2018 ; Chen et al. , 2018a ; Zou & Shen , 2018 ) . Although these studies have established an iteration complexity of O ( 1/ 4 ) for different variants of ADAGRAD for finding an -stationary solution of a stochastic non-convex optimization problem , none of them can exhibit the potential adaptive advantage of ADAGRAD over SGD as in the convex case . Besides that , these studies also suffer from the following shortcomings : ( i ) they all assume smoothness of the objective function , while we consider non-smooth and non-convex problems ; ( ii ) their convergence is provided on a solution with minimum magnitude of gradient that is expensive to compute , though their results also imply a convergence on a random solution selected from all iterates with decreasing sampling probabilities . In contrast , these shortcomings do not exist in this paper . To the best of our knowledge , our result is the first one that explicitly shows that coordinate-wise adaptive step size could yield faster convergence than using non-adaptive step size for non-smooth non-convex problems , which is similar to that in the convex case and was observed in practice for deep learning ( Dean et al. , 2012 ) . The proposed stagewise algorithm is similar to several existing algorithms in design ( Xu et al. , 2017 ; Davis & Grimmer , 2017 ) , which are originated from the proximal point algorithm ( Rockafellar , 1976 ) . I.e. , at each stage a regularized convex subproblem is formed and then a stochastic algorithm is employed for optimizing the regularized subproblem inexactly with a number of iterations . Xu et al . ( 2017 ) used this idea for solving problems that satisfy a local error bound condition , aiming to achieve faster convergence than vanilla SGD . Davis & Grimmer ( 2017 ) are probably the first who analyzed this idea for solving non-smooth weakly convex problems . In these two papers SGD with decreasing step sizes for a strongly convex problem is employed at each stage for solving the regularized subproblem . Our stagewise algorithm is developed following the similar idea . The key differences from ( Xu et al. , 2017 ; Davis & Grimmer , 2017 ) are that ( i ) we focus on non-convex problems instead of convex problems considered in ( Xu et al. , 2017 ) ; ( ii ) we analyze a non-uniform sampling strategy with sampling probabilities increasing as the stage number to justify a heuristic approach that uses the last averaged solution for prediction , unlike the uniform sampling used in ( Davis & Grimmer , 2017 ) ; ( iii ) we present a unified algorithmic framework and convergence analysis , which enable one to employ any suitable stochastic convex optimization algorithms at each stage . Importantly , it brings us several interesting variants including stagewise stochastic momentum methods and stagewise ADAGRAD . Finally , we note that a similar idea that adds a strongly convex term into the non-convex objective has been considered in several recent works ( Allen-Zhu , 2017 ; Lan & Yang , 2018 ; Carmon et al. , 2016 ) . However , the key difference between these works and the present work is that they need to assume the objective function is smooth or has a smooth component . In addition , Allen-Zhu ( 2017 ) ; Lan & Yang ( 2018 ) consider finite-sum problems , and Carmon et al . ( 2016 ) consider deterministic problems . In contrast , we consider more general stochastic problems without assuming the objective function is smooth .

Summary: Non-convex optimization is a hot topic since many machine learning problems can be formulated as non-convex problems. In this paper, the authors propose a universal stage-wise algorithm for weakly convex optimization problems. The idea is to add a strongly convex regularizer centered at an iterate of previous stage to the objective function. This builds a convex function which can be optimized by any standard methods in the convex optimization setting. The authors developed convergence rates in expectation in terms of the gradient of envelope. Empirical results are also reported to show the effectiveness of the method.