A Novel Convergence Analysis for the Stochastic Proximal Point Algorithm

1 INTRODUCTION . It has been widely accepted that when training large-scale machine learning models , the training algorithm should act in a sample-by-sample manner in order to reduce computational and memory overhead—the size of the data set may be too large that calculating the full gradient information is too costly . Moreover , most machine learning problems does not have to be solved with very high accuracy , since the ultimate goal is not to fit the training data but rather to generalize the algorithm well such that the performance is decent on unseen data . Most existing stochastic algorithms are based upon the stochastic gradient descent ( SGD ) framework ( Bottou et al. , 2018 ) . SGD is extremely easy to implement and provides asymptotic convergence , although the convergence rate is generally slow ( and subject to careful choice of step sizes ) . Various approaches have been proposed to accelerate the plain vanilla SGD . Reducing the variance of the stochastic gradient and introducing adaptive learning schemes are two main lines of research . SVRG ( Johnson & Zhang , 2013 ) and SAGA ( Defazio et al. , 2014 ) ( and their follow-up works such as ( Defazio , 2016 ) ) focus on reducing the variance of the stochastic gradient descent , at the cost of increased time or memory complexities ( to the order of the entire data set ) . On the other hand , AdaGrad ( Duchi et al. , 2011 ) and Adam ( Diederik P. Kingma , 2014 ) introduce adaptive learning schemes and effectively keep the algorithm fully stochastic and light-weight . Besides these practical improvements , theoretical progress has been made by on quantifying the best possible convergence rate using first-order information ( Lei et al. , 2017 ; Allen-Zhu , 2017 ; 2018a ; b ) . 1.1 STOCHASTIC PROXIMAL POINT ALGORITHM ( SPPA ) . In this paper , we explore a different type of stochastic algorithm called the stochastic proximal point algorithm ( SPPA ) , also known as stochastic proximal iterations ( Ryu & Boyd , 2014 ) or incremental proximal point method ( Bertsekas , 2011a ; b ) . Consider the following optimization problem with the objective function in the form of a finite sum of component functions minimize w∈R3 1 = =∑ 8=1 58 ( w ) = ( w ) . ( 1 ) SPPA takes the following simple form : Algorithm 1 Stochastic proximal point algorithm ( SPPA ) 1 : repeat 2 : randomly draw 8C uniformly from { 1 , . . . , = } 3 : wC+1 ← argminw _C 58C ( w ) + ( 1/2 ) ‖w − wC ‖2 = Prox_C 58C ( wC ) 4 : until convergence Line 3 of Algorithm 1 calculates the proximal operator of the function _C 58 evaluated at wC , denoted as Prox_C 58 ( wC ) . This is the stochastic version of the proximal point algorithm , which dates back to Rockafellar ( 1976 ) . SPPA has an abstract per-iteration update rule and it acquires more information from the problem than solely the first order derivatives , which makes it not as universally applicable as SGD . Yet , thanks to the more information inquired , it is possible to obtain faster and more robust convergence guarantees . While some ‘ accelerated ’ versions of SGD demand additional time or space overhead to go over the entire data set , SPPA does not have any overhead and it is suitable to be used in a completely online setting , it performs well even the data samples are not revisited again . To the best of our knowledge , people started studying the convergence behavior of SPPA only recently ( Bertsekas , 2011a ; Ryu & Boyd , 2014 ; Bianchi , 2016 ; Pătraşcu , 2020 ; Toulis et al. , 2021 ) . Somewhat surprisingly , the convergence analysis of SPPA draws little resemblance to its deterministic counterpart , the proximal point algorithm . This is unlike the case for SGD , of which the convergence analysis follows almost line-by-line to that of the subgradient method . Moreover , existing analysis of SPPA shows no improvement in terms of convergence rate , which seems counter-intuitive due to the nature of the updates . 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 & Bertsekas , 2013 ; Duchi & Ruan , 2018 ; Asi & Duchi , 2019 ; Davis & Drusvyatskiy , 2019 ) . 1.2 CONTRIBUTIONS . The main contribution of this paper is to provide a completely novel convergence analysis of SPPA for general convex problems . This contribution , together with the efficient implementation strategies discussed in the appendix , results in great practical results in almost all of the classical empirical risk minimization ( ERM ) problems in large-scale machine learning . While there exists some convergence results of SPPA for convex problems , the novel analysis provided in this work requires minimal assumptions ( nothing but the convexity of the loss functions ) . The new convergence analysis also shows great resemblance to its deterministic counterpart , which is not the case from other works . In the appendix we will discuss how to efficiently compute the abstract per-iteration update rule for SPPA with complexity that is comparable to SGD-type methods . An interesting observation is that in a lot of ERM examples the resulting update has a similar form like SGD , but with a smartly chosen step size derived from the proximal update . For a class of smooth risks such as the logistic loss , we appeal to reformulate the optimization problem such that we get close to closed form solution via bisection . We also briefly discuss how to modify the algorithm to have the ability to handle nonsmooth regularization terms such as the ! 1 norm , using the stochastic Douglas-Rachford splitting approach . Finally , we apply SPPA with efficient implementations to a large variety of classification and regression problems . Numerical experiments are conducted not only for linear classification and regression ( in the appendix ) models , but also nonconvex deep neural network problems . Although the convergence analysis provided in this paper does not cover nonconvex cases , empirical results suggest that it is still worth treating SPPA as an effective alternative for deep learning . 2 CONVERGENCE ANALYSIS . In this section we provide convergence analysis of SPPA for general convex loss functions equation 1 to a global minimum in expectation . In recent years there have been some work tackling the same problem , e.g. , Bertsekas ( 2011a ) ; Pătraşcu ( 2020 ) ; Toulis et al . ( 2021 ) . In this paper , however , we provide new convergence analysis that is much easier to understand while requiring nearly no assumptions other than convexity . For this reason we believe the theoretical contribution is significant enough as it broadens the applicability of SPPA . There is a well-known resemblance between proximal methods and gradient descent , assuming the loss function is differentiable : while a full gradient descent step takes the form wC+1 = wC − _C∇ ( wC ) , the definition of a full proximal step guarantees that _C∇ ( wC+1 ) = wC+1 − wC , meaning that wC+1 = wC − _C∇ ( wC+1 ) . Therefore , one might expect that a well-established convergence analysis of SGD for nonconvex problems , for example ( Bottou et al. , 2018 , §4.3 ) , can be seamlessly applied to SPPA . However , when applied in a stochastic fashion , the situation is a little more complicated . Consider E [ ∇ 58 ( wC ) | wC ] , where the expectation is taken over the sampling procedure conditioned on wC , for SGD we typically require ∇ 58 ( wC ) to be an unbiased estimator of ∇ ( wC ) , which is easy to satisfy if 8C is uniformly sampled from { 1 , . . . , = } given wC . For SPPA then one needs to consider E [ ∇ 58 ( wC+1 ) | wC ] , again over the sampling procedure conditioned on wC . This is in fact difficult to quantify because the update wC+1 depends on the sample that is drawn from the data set . The equation E [ ∇ 58 ( wC+1 ) | wC ] = ∇ ( wC+1 ) does not make sense because wC+1 on the right-hand-side is still random conditioned on wC . It is for this reason that existing analysis of SPPA differs drastically from its deterministic counterpart PPA ( Bertsekas , 2011a ; Ryu & Boyd , 2014 ; Bianchi , 2016 ) . What we can show , however , is that the distribution of 8C is still uniform conditioned on wC+1 instead of wC , as formalized in the following proposition . Lemma 1 . At every iteration of SPPA ( Algorithm 1 ) , we have the conditional probability ? ( 8C |wC+1 ) = 1/= . The proof of Lemma 1 is relegated to the supplementary material . What Lemma 1 suggests is that given the current iterate wC+1 , without knowing its predecessors wC and beyond , every component function 58 is equally likely to be picked as the proximal update that leads to wC+1 . In other words , conditioned on the current iterate without knowing the past , we do not gain additional information about which component function 58 is more likely to have been selected . As it turns out , Lemma 1 significantly simplifies the following convergence analysis while showing resemblance to that of deterministic PPA . What is more , the resulting analysis requires no assumptions other than convexity : typical assumptions such as Lipschitz smoothness or bounded variance are not required . Remark . Lemma 1 holds for C = 0 if the initialization w0 is random , and the distribution from which it is drawn has the sample space the same as the domain of the objective function equation 1 . In practice w0 is usually drawn from a distribution that is independent from the cost function , say N ( 0 , O ) . In this case , for a given value of w1 , there is no chance of eliminating possible w0 ’ s because any corresponding w0 is in the sample space of the initialization distribution . 2.1 GENERIC CONVEX CASE WITHOUT STRONG CONVEXITY . Under Lemma 1 , we have the following proposition , which serves as the stepping stone for our main convergence results . Proposition 1 . Suppose each 51 , . . . , 5= is convex , then at iteration C of SPPA ( Algorithm 1 ) we have 2_C ( ( wC+1 ) − ( w★ ) ) ≤ E [ ‖wC − w★‖2 | wC+1 ] − ‖wC+1 − w★‖2 , ( 2 ) where w★ denotes an optimal solution of equation 1 . Proof . We start by the equation ‖wC − w★‖2 = ‖wC − wC+1 + wC+1 − w★‖2 = ‖wC+1 − w★‖2 + 2 ( wC − wC+1 ) > ( wC+1 − w★ ) + ‖wC − wC+1‖2 . ( 3 ) According to the definition wC+1 = argminw _C 58C ( w ) + ( 1/2 ) ‖w − wC ‖2 , we know that wC − wC+1 is a subgradient of _C 58C at wC+1 . Therefore _C 58C ( wC+1 ) + ( wC − wC+1 ) > ( w − wC+1 ) ≤ _8 58C ( w ) at any w. Let w = w★ and substitute it in equation 3 , we have ‖wC − w★‖2 ≥ ‖wC+1 − w★‖2 + 2_C ( 58C ( wC+1 ) − 58C ( w★ ) ) + ‖wC − wC+1‖2 ≥ ‖wC+1 − w★‖2 + 2_C ( 58C ( wC+1 ) − 58C ( w★ ) ) . Finally , taking condition expectations on both sides given wC+1 ( but not wC nor 8C ) E [ ‖wC − w★‖2 | wC+1 ] ≥ ‖wC+1 − w★‖2 + 2_CE [ ( 58 ( wC+1 ) − 58 ( w★ ) | wC+1 ] . According to Lemma 1 , the conditional distribution over 8C is uniform , thus E [ 58C ( wC+1 ) − 58C ( w★ ) | wC+1 ] = ( wC+1 ) − ( w★ ) , and we obtain equation 2 . With the help of Proposition 1 , the rest of the results follows straight-forwardly . Theorem 1 . If all 51 , . . . , 5= are convex , and an initialization w0 is picked from a distribution with the sample space the same as the domain of , then the sequence { wC } generated by SPPA ( Algorithm 1 ) satisfies lim ) →∞ inf C≤ ) E [ ( wC ) ] − ( w★ ) ≤ 1 2 ∑∞ C=1 _C E‖w0 − w★‖2 , ( 4 ) where w★ denotes an optimal solution of equation 1 . The term E‖w0 − w★‖2 is a constant for any approach of initializing w0 . Proof . Taking total expectation over equation 2 , we have 2_C ( E [ ( wC+1 ) ] − ( w★ ) ) ≤ E [ ‖wC − w★‖2 ] − E [ ‖wC+1 − w★‖2 ] . Sum over all C = 1 , 2 , . . . , ) , we have 2 ) ∑ C=1 _C ( E [ ( wC+1 ) ] − ( w★ ) ) ≤ E‖w0 − w★‖2 − E‖w ) − w★‖2 ≤ E‖w0 − w★‖2 . On the left-hand-side we have ( by definition ) inf E [ ( wC ) ] ≤ E [ ( wg ) ] for any g. Applying that , dividing both sides by 2 ∑ C _C , and letting ) →∞ , we get equation 4 . Remark . The proofs here take conditional expectations backwards , i.e. , conditioned on wC+1 and average over the previous iteration . This may be a little counter-intuitive . However , there is nothing wrong mathematically—expectation is a linear operator so E [ + ] = E [ ] + E [ ] under all circumstances ; if an inequality 5 ( G ) ≤ 6 ( G ) holds for all values of G in its sample space , then E [ 5 ( G ) ] ≤ E [ 6 ( G ) ] for any distribution over G since it is just a nonnegative sum on both sides . Theorem 1 states a generic results on the convergence of SPPA . The left-hand-side of equation 4 is , by definition , nonnegative ; the right-hand-side , however , goes to zero if the infinite sum ∑ C _C →∞ . This implies that the infimum if E [ ( wC ) ] goes to zero for appropriately chosen step sizes . Somewhat surprisingly , the only assumption we made was that the functions are convex . The celebrated SGD , on the other hand , requires at least two more assumptions : Lipschitz smoothness and that the stochastic gradients have bounded variance . What is more , the flexible choice of _C means that we can make the convergence rate arbitrarily fast , by allowing _C to be increasing ( rather than decreasing in most gradient-based methods ) . Of course this may lead to large variance of the sequence , which we may want to avoid in practice . Here we provide the convergence rate for two commonly used step size rules . Corollary 1 . If all 51 , . . . , 5= are convex , then the sequence { wC } generated by SPPA ( Algorithm 1 ) with a constant step size _C = _ satisfies inf C≤ ) E [ ( wC ) ] − ( w★ ) ≤ 1 2_ ) ‖w0 − w★‖2 , ( 5 ) where w★ denotes an optimal solution of equation 1 . The proof is straight-forward by substituting _C with _ in equation 4 . This shows that a constant step size ( regardless of its value ) gives a O ( 1/ ) ) sublinear convergence rate . Corollary 2 . If all 51 , . . . , 5= are convex , then the sequence { wC } generated by SPPA ( Algorithm 1 ) with an increasing step size rule _C = C_ satisfies inf C≤ ) E [ ( wC ) ] − ( w★ ) ≤ 1 _ ) ( ) + 1 ) ‖w0 − w★‖ 2 , ( 6 ) where w★ denotes an optimal solution of equation 1 . The proof comes from the arithmetic series 1+2+· · ·+ ) = ) ( ) +1 ) /2 . This shows that the increasing step size rule _C = C_ gives a O ( 1/ ) 2 ) sublinear convergence rate , the same as Nesterov ’ s optimal gradient algorithm ( 2013 ) . While the expected convergence looks excellent , we notice that in practice the performance also depends on the variance of the sequence generated by SPPA , which could be large if the expected convergence rate is too fast . We can even have linear convergence rate by letting _C increase exponentially , but the variance would be so large that it makes little practical sense . As we will see in the experiment section , a constant step size rule still gives the best performance in most cases .

Summary: This paper provides a novel elegant proof of convergence for the Stochastic Proximal Point Algorithm, under the assumption of convexity. It proceeds by surprisingly conditioning on future iterates, reversing usual convergence proofs. Experiments show that the SPPA is competitive, including in nonconvex settings such as training neural networks for image classification.