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: In this paper the authors study stochastic proximal point algorithm for nonconvex optimization, where the model is iteratively updated by solving a proximal optimization problem based on a randomly selected loss function. The authors develop efficient implementation for solving the proximal optimization problem: first for nonlinear least squares and then for general losses. Then the authors study the convergence rates for the developed algorithm. Upper bounds on the expected average squared gradients are developed for both constant step sizes and diminishing step sizes. Experimental results are also reported to support the algorithm in practical implementations.