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: # === Update ===  I have read the authors reviews as well as the response by the authors.  I'd like to thank the authors for engaging in a discussion around Lemma 1. It is unfortunate that the result does not appear to hold even in a restricted setting.  I will maintain my score. I wish the authors the best in their future work.  # ======  This submission investigates the convergence rate of the stochastic proximal point algorithm (SPPA) for convex, finite-sum functions. The authors derive new convergence rates for SPPA under the assumption that each function in the finite sum is convex (or strongly convex). Key to these derivations is a lemma characterizing the probabilities of sampling a sub-function form a time-reversible Markov chain.  In addition to these results, the authors derive efficient implementations of SPPA for a variety of loss functions, generally by reducing the problem to one-dimensional root-finding. The paper concludes with experiments comparing SPPA to SGD, Adam, and other baseline stochastic optimization algorithms for (i) convex optimization and (ii) training CNNs and ResNets on MNIST and CIFAR-10, respectively. The experiments show that SPPA converge faster than stochastic gradient methods and have a similar computational cost.