Stochastic Proximal Point Algorithm for Large-scale Nonconvex Optimization: Convergence, Implementation, and Application to Neural Networks

1 INTRODUCTION . Algorithm design for large-scale machine learning problems have been dominated by the stochastic ( sub ) gradient descent ( SGD ) and its variants ( Bottou et al. , 2018 ) . The main reasons are two-fold : on the one hand , the size of the data set may be so large that obtaining the full gradient information is too costly ; on the other hand , solving the formulated problem to very high accuracy is typically unnecessary in machine learning , since the ultimate goal of most tasks is not to fit the training data but to generalize well on unseen data . As a result , stochastic algorithms such as SGD has gained tremendous popularity recently . There has been many variations and extensions of the plain vanilla SGD algorithm to accelerate its convergence rate . One line of research focuses on reducing the variance of the stochastic gradient , resulting in famous algorithms such as SVRG ( Johnson and Zhang , 2013 ) and SAGA ( Defazio et al. , 2014 ) , which results in extra time/memory complexities of the algorithm ( significantly ) . More recently , adaptive learning schemes such as AdaGrad ( Duchi et al. , 2011 ) and Adam ( Diederik P. Kingma , 2014 ) have shown to be more effective in keeping the algorithm fully stochastic and light-weight . In terms of theory , there has also been surging amount of work quantifying the best possible rate using first-order information ( Lei et al. , 2017 ; Allen-Zhu , 2017 ; 2018a ; b ) , as well as its ability to obtain not only stationary points but also local optima ( Ge et al. , 2015 ; Jin et al. , 2017 ; Xu et al. , 2018 ; Allen-Zhu , 2018 ) . 1.1 STOCHASTIC PROXIMAL POINT ALGORITHM ( SPPA ) . In this work , we consider a different type of stochastic algorithm called the stochastic proximal point algorithm ( SPPA ) , also known as incremental proximal point method ( Bertsekas , 2011a ; b ) or stochastic proximal iterations ( Ryu and Boyd , 2014 ) . Consider the following optimization problem with the objective function in the form of a finite sum of component functions minimize ) ∈R3 1 = =∑ 8=1 ℓ8 ( ) ) = ! ( ) ) . ( 1 ) SPPA takes the following simple form : 1 : repeat 2 : randomly draw 8 from { 1 , . . . , = } 3 : ) C+1 ← arg min ) _Cℓ8 ( ) ) + ( 1/2 ) ‖ ) − ) C ‖2 = Prox_Cℓ8 ( ) C ) 4 : until convergence The update rule in line 3 is called the proximal operator of the function _Cℓ8 evaluated at ) C . This is the stochastic version of the proximal point algorithm , which dates back to Rockafellar ( 1976 ) . Admittedly , SPPA is not as universally applicable as SGD , due to the abstraction of the per-iteration update rule . It is also asking for more information from the problem than merely the first-order derivatives . However , with the help of more information inquired , there is also hope that it provides faster and more robust convergence guarantees . As we will see in numerical experiments , SPPA is able to achieve good optimization performance by taking fewer number of passes through the data set , although it takes a little more computations for each batch . We believe in many cases it is worth trading off more computations for fewer memory accesses . To the best of our knowledge , convergence analyses of SPPA has only been studied for convex problems ( Bertsekas , 2011a ; Ryu and Boyd , 2014 ; Bianchi , 2016 ) . Their study shows that SPPA converges somewhat similar to SGD for convex problems , but the updates are much more robust to instabilities in the problem . Most authors also accept the premise that the proximal operator is sometimes difficult to evaluate , and thus proposed variations to the plain vanilla version to handle more complicated problem structures ( Wang and Bertsekas , 2013 ; Duchi and Ruan , 2018 ; Asi and Duchi , 2019b ; Davis and Drusvyatskiy , 2019 ) . In terms of nonconvex optimization problems , there is very little work until very recently ( Davis and Drusvyatskiy , 2019 ; Asi and Duchi , 2019a ) . However , their convergence analysis is somewhat unconventional . Typically for a nonconvex problem , we would expect a theoretical claim that the iterates generated by SPPA converges ( in expectation ) to a stationary point . This is not easy , and the result given by ( Davis and Drusvyatskiy , 2019 ) and ( Asi and Duchi , 2019a ) defined an imaginary sequence ( that is not computed in practice ) { ) ̃ C } as ) ̃ C = arg min ) _C ! ( ) ) + ( 1/2 ) ‖ ) − ) C ‖2 , i.e. , the proximal operator of the full loss function from the algorithm sequence { ) C } . Their results show that this imaginary sequence { ) ̃ C } converges to a stationary point in expectation . 1.2 CONTRIBUTIONS . There are two main contributions we present in this paper ; efficient implementations of proximal operator update for SPPA with application on regression and classification problems , and convergence to a stationary point for SPPA algorithm for general non-convex problems . The cost of this abstract per-iteration update rule has been the burden for SPPA algorithm , even though it has been shown to converge faster and more stable than the celebrated SGD algorithm . In this paper , we show that it is actually not a burden when the per-iteration update is efficiently implemented . In the implementation section we will discuss the implementation of abstract per-iteration update rule for non-linear least squares ( NLS ) problem and other non-linear problems . We present two different implementations for these two different categories of problems . SPPA-Gauss Newton ( SPPA-GN ) and SPPA-Accelerated , respectively for regression with nonlinear least squares and classification problems . We apply SPPA to a large family of nonconvex optimization problems and show that the seemingly complicated proximal operator update can still be efficiently obtained . Both implementations give results that are comparable with the state of the art stochastic algorithms . On the other hand , there is large group of nonlinear problems that are not necessarily expressed as a NLS problem ( non-NLS ) . Many of the classification problems in deep learning applications as of today , use different type of loss functions than mean least squared error . Hence , we suggest an alternative to SPPA-GN where we implement the per-iteration update rule based on well-known stochastic optimization algorithms , preferably with the ones that take less number of iterations to converge like L-BFGS ( Liu and Nocedal , 1989 ) . As a second contribution , we show that SPPA , when applied to a nonconvex problem , converges to a stationary point in expectation , as informally stated as follows Theorem 1 . ( informal . ) For SPPA with fixed _C = _ , the expected gradient converges to a region where the norm is bounded with radius proportional to _ . For diminishing step sizes such that lim ) →∞ ) ∑ ) C=1 _ −1 C , the expected gradient converges to zero . Finally , we apply SPPA with the novel efficient updates to some classical regression and classification using 4 different data sets for neural network training . Detailed algorithmic descriptions are provided , and we show the outstanding performance of SPPA when effectively executed . As a sneak peek of the numerical performance of the proposed SPPA implementation , Figures 1 and 2 shows the cross entropy loss and prediction error on the test set over the progression of SPPA verses various baseline algorithms . SPPA indeed converges much faster than all SGD-based methods , achieving peak prediction accuracy after only going through the data set approximately 10 times . We should stress that the per-iteration complexity of SPPA is higher than other methods , since it tries to solve a small optimization problem rather than a simple gradient update . However , as we have argued before , in many cases it is worth trading off more computations for fewer memory accesses . 2 EFFICIENT IMPLEMENTATION . In this section we introduce several efficient methods to calculate the proximal operator update given in pseudo-code of SPPA section 1.1 line 3 . The two methods emerge from the question of ‘ How can we implement the proximal operator update efficiently ? ’ . The answer to this question depends on the type of the objective function . We considered the problems in two different categories , nonlinear least squares ( NLS ) and other generic nonlinear problems . 2.1 SPPA-GN FOR NONLINEAR LEAST SQUARES . A nonlinear least squares ( NLS ) problem takes the following form minimize ) ∈R3 1 = =∑ 8=1 1 2 ( i8 ( ) ) ) 2 , ( 2 ) where each i8 is a general nonlinear function with respect to ) . It is a classical nonlinear programming problem ( Bertsekas , 1999 ; Boyd and Vandenberghe , 2018 ) with many useful applications , including least squares neural networks ( Van Der Smagt , 1994 ) , where each i8 corresponds to the residual of fitting for the 8th data sample for regression . We will show that problems of this form can be efficiently executed by SPPA-GN , despite its seemingly complication . To apply SPPA to NLS , the main challenge is to efficiently evaluate the proximal operator ) C+1 ← arg min ) _C 2 ( i8 ( ) ) ) 2 + 1 2 ‖ ) − ) C ‖2 . ( 3 ) Notice that this function itself is a nonlinear least squares objective , although with only one component function together with the proximal term . The traditional wisdom to solve a NLS problem is to apply the Gauss-Newton ( GN ) algorithm : at each iteration , we first take a first-order approximation of the vector-valued function inside the Euclidean norm , and set the update as the solution of the approximated linear least squares problem . It is a well-known algorithm that can be found in many standard textbooks , e.g. , ( Bertsekas , 1999 ; Nocedal and Wright , 2006 ; Boyd and Vandenberghe , 2018 ) . To apply GN to ( 3 ) , we first take linear approximation of i8 at the current update ) as i8 ( ) ) ≈ i8 ( ) ) + ∇i8 ( ) ) > ( ) − ) ) , and set the solution of the following problem ) + as the next update minimize ) _C 2 ( i8 ( ) ) + ∇i8 ( ) ) > ( ) − ) ) ) 2 + 1 2 ‖ ) − ) C ‖2 . ( 4 ) Obviously , ( 4 ) has a closed form solution ) + = ) C − ( 1 _C O + g8g > 8 ) −1 g8 ( i8 ( ) ) − g > 8 ( ) − ) C ) ) , ( 5 ) where we denote g8 = ∇i8 ( ) ) to simplify notation . Notice that the matrix to be inverted in ( 5 ) has a simple “ identity plus rank-one ” structure , implying that it can be efficiently computed in linear time . Using the “ kernel trick ” ( Boyd and Vandenberghe , 2018 , pp.332 ) ( G > G + UO ) −1G > = G > ( GG > + UO ) −1 , ( 6 ) update ( 5 ) simplifies to ) + = ) C − i8 ( ) ) − ∇i8 ( ) ) > ( ) − ) C ) _−1C + ‖∇i8 ( ) ) ‖2 ∇i8 ( ) ) . ( 7 ) Algorithm 1 SPPA-GN 1 : initialize ) 0 , C ← 0 2 : repeat 3 : randomly draw 8 from { 1 , . . . , = } 4 : ) + ← ) C 5 : repeat 6 : ) ← ) + 7 : ) + ← ) C − i8 ( ) ) −∇i8 ( ) ) > ( ) − ) C ) _−1C +‖∇i8 ( ) ) ‖2 ∇i8 ( ) ) 8 : until convergence 9 : ) C+1 ← ) + , C ← C + 1 10 : until convergence As we can see , each GN update only takes O ( 3 ) flops , which is as cheap as that of a SGD step . To fully obtain the proximal operator ( 3 ) , one has to run GN for several iterations . However , thanks to the superlinear convergence rate of GN near its optimal ( Nocedal and Wright , 2006 ) , which is indeed the case if we initiate at ) C because of the proximal term , it typically takes no more than 5–10 GN updates . The detailed description of the proposed algorithm , which we term SPPA-GN , for solving general NLS problems is shown in Algorithm 1 .

Summary: This paper studies the stochastic proximal point algorithm (SPPA) for large-scale nonconvex optimization problems. The authors propose to use Gauss-Newton to perform the proximal update in nonlinear least squares and L-BFGS or accelerated gradient for generic problems. The authors derive the convergence of SPPA to a stationary point in expectation for nonconvex problems, and perform numerical experiments to showcase the effectiveness of the proposed method compared to SGD and its variants.