Sharp Analysis of Stochastic Optimization under Global Kurdyka- Ł ojasiewicz Inequality
Anonymous Author(s)
AfﬁliationAddress email
Abstract
We study the complexity of ﬁnding the global solution to stochastic nonconvex optimization when the objective function satisﬁes global Kurdyka- Ł ojasiewicz (KL)inequality and the queries from stochastic gradient oracles satisfy mild expectedsmoothness assumption.
We ﬁrst introduce a general framework to analyze Stochastic Gradient Descent (SGD) and its associated nonlinear dynamics under the setting.
As a byproduct of our analysis, we obtain a sample complexity of O ( ✏  (4  ↵ ) / ↵ )for SGD when the objective satisﬁes the so called ↵ -P Ł condition, where ↵ is thedegree of gradient domination.
We also demonstrate that the complexity boundof SGD is tight under the speciﬁc choice of stepsizes and cannot be improvedwith mini-batching.
Furthermore, we show that a modiﬁed SGD with variancereduction and restarting (PAGER) achieves an improved sample complexity of O ( ✏  2 / ↵ ) when the objective satisﬁes the average smoothness assumption.
Thisleads to the ﬁrst optimal algorithm for the important case of ↵ = 1 which appearsin applications such as policy optimization in reinforcement learning.
Introduction
Non-convex optimization problems are ubiquitous in machine learning domains such as in trainingdeep neural networks [16] or policy optimization in reinforcement learning [42].
Stochastic Gradient
Descent (SGD) and its variants are driving the practical success of machine learning approaches.
Naturally, understanding the limits of performance of SGD in the non-convex setting has become animportant avenue of research in recent years [15, 3, 23, 36, 17, 10, 50].
We are interested in solving the unconstrained stochastic , nonconvex optimization problem of theform
[Equation]
where f ( · ) is smooth and possibly nonconvex, and ⇠ is a random vector drawn from a distribution D .
Moreover, we are interested in an important special case of (1), when the expectation can be writtenas the average of n smooth functions:
[Equation]
As for a general nonconvex differentiable objective f : R d ! R , ﬁnding a global minimum of f isin general intractable [35, 44], there are two common strategies to analyze optimization methodsfor nonconvex functions.
The ﬁrst one is to scale down the requirements on the solution of interestfrom global optimality to some relaxed version, e.g., ﬁrst-order stationary point E
[ kr f (ˆ x ) k ]  ✏ .
However, such solutions do not exclude the possibility of approaching a suboptimal local minima or a saddle point.
Another approach is to study nonconvex problems with additional structuralassumption in the hope of being able to establish convergence to global solutions.
In this direction,several relaxations of convexity have been proposed and analyzed, for instance, star convexity, quasarconvexity, error bounds condition, restricted secant inequality, and quadratic growth [20, 18, 17].
Many of the aforementioned relaxations have limited application in real-world problems.
Recently, there has been a surge of interest in the analysis of functions satisfying the so-called
Kurdyka- Ł ojasiewicz (K Ł ) inequality [6, 7, 7].
Of particular interest is the family of functions thatsatisfy global K Ł inequality.
Speciﬁcally, we say that f ( · ) satisﬁes global K Ł inequality if there existssome continuous function  ( · ) such that || r f ( x ) ||   ( f ( x )  inf x f ( x )) for all x 2 R d .
If thisinequality is satisﬁed for  ( t ) = p 2 µ t 1 / ↵ , then we say that ↵ -P Ł condition is satisﬁed for f ( · ) .
Thespecial case of global K Ł inequality, 2 -P Ł , often referred as Polyak- Ł ojasiewicz or PL condition, wasoriginally discovered independently in the seminal works of B. Polyak, T. Le ž anski and S. Lojasiewicz[39, 26, 31, 32].
Notably, this class of problems have found many interesting emerging applicationsin machine learning, for instance, policy gradient (PG) methods in reinforcement learning [34, 1, 46],generalized linear models [33], over-parameterized neural networks [2, 48], linear quadratic regulatorin optimal control [9, 13], and low-rank matrix recovery [5].
Despite increased popularity of K Ł and ↵ -P Ł assumptions, the analysis of stochastic optimizationunder it remains limited and the majority of works focus on deterministic gradient methods.
Indeeduntil recently, only the special case of ↵ -P Ł with ↵ = 2 was mainly addressed in the literature[20, 17, 23, 43, 40].
In this paper, we study the sample complexities of stochastic optimization forthe broader class of nonconvex functions with global K Ł property.
1.1 Related Works and Open Questions
SGD : A plethora of existing works has studied the sample complexity of SGD and its variants forﬁnding an ✏ -stationary point of general nonconvex function f , that is, a point x 2 R d for which
E [ || r f (ˆ x ) || ]  ✏ .
For instance, [15] showed that for a smooth objective (one with Lipschitz gradients) 55 under bounded variance (BV) assumption, SGD with properly chosen stepsizes reaches an ✏ -stationarypoint with the sample complexity of O ( ✏  4 ) .
Recently, [23, 46] further extended the result undera much milder expected smoothness (ES) assumption on stochastic gradient.
While this samplecomplexity is known to be optimal for general noncovex functions, a naive application of this resultto a function value using ↵ -P Ł condition would lead to a suboptimal O ( ✏  4 / ↵ f ) sample complexity forﬁnding an ✏ f -optimal solution.
Recently, [14] studied SGD and established convergence rates for↵ -P Ł functions assuming variances are bounded.
Their sample complexity rates equal O ( ✏  (4  ↵ ) / ↵ f )in our notation.
Later [28] considered SGD scheme with random reshufﬂing under local and global K Ł and provided convergence in the iterates for ↵ 2 (1 , 2] .
We note that our proof techniques aredifferent from [14] and [28] and are not limited merely to the case of BV assumption.
In this work,we will answer the following open question: What is the exact performance limit of SGD underglobal K Ł condition and a more practical model of stochastic oracle?
Variance reduction : There has been extensive research on development of algorithms which improvethe dependence on n and/or ✏ for both problems (1) and (2).
One important family of techniques isvariance reduction , which has emerged from the seminal work of Le Roux et.
al [41].
The main ideaof variance reduction is to make use of the stochastic gradients computed at previous iterations toconstruct a better gradient estimator at a relatively small computational cost.
Various modiﬁcations,generalizations, and improvements of the variance reduction technique appeared in subsequent work,for instance, [41, 22, 11] to name a few.
Finite-sum case : A number of recently proposed algorithms such as SNVRG, SARAH, STORM,
SPIDER, and PAGE [49, 37, 10, 12] achieve the optimal sample complexity O  n + p n ✏ to obtain E [ kr f (ˆ x ) k ]  ✏ for minimizing a general nonconvex function [29].
Under 2-P Ł condition, thesemethods further improve to O [Equation]  sample complexity 1 for ﬁnding a point x suchthat E [ f ( x )  f ? ]  ✏ f .
[21] studies SARAH in ﬁnite sum case under local K Ł assumption andproves convergence in the iterates.
The analysis in [21] is only asymptotic analysis and the dependenceon the parameters  and n , which are important in practice for quantifying the improvement over GDand SGD are ignored.
[27] proposes an SVRG-based algorithm for K Ł functions and [45, 37, 38] study other variance reduction techniques, but they only analyze the special case ↵ = 2 .
In ourwork, we will answer the following open question: What is the extent of improvement any variancereduction scheme can provide under global K Ł condition for ﬁnite-sum objectives of the form (2) ?
Online/streaming case: While variance reduction methods have been initially designed for problemsof the form (2), later it was shown that they also improve over SGD in solving (1) [25, 30].
For example, the methods in [49, 37, 12, 29] achieve O ( ✏  3 ) complexity under a mild additional assumptionssuch as the average L - smoothness (Assumption 7) improving over O ( ✏  4 ) complexity of SGD forﬁnding an ✏ -stationary point.
Note that variance reduction for ↵ -P Ł functions has only been addressedfor the special case ↵ = {+
+}
and it does not yield an improvement over SGD in terms of ✏ since SGD isalready optimal. For other cases, taking ↵ = 1 as an example, to obtain the global convergence of SGD, we could translate the result O ( ✏  3 ) for ﬁnding a stationary point to a function value by using↵ -P Ł condition, and obtain O
[Equation]
sample complexity.
There are two issues with this approach.
First, this complexity does not provide any improvement over SGD in the most interesting practicalcase ↵ = 1 .
Second, the guaranties for general nonconvex optimization hold on average, in the sensethat the point ˆ x is sampled uniformly from all the iterates of the algorithm.
It would be more desirableto instead derive last iterate convergence guaranties under the KL condition.
In this work, we willanswer the following open question: Is it possible to accelerate the O  ✏
[Equation] / ↵ [Equation]  sample complexityof SGD under global K Ł condition for stochastic objectives of the form (1) ?
1.2 Contributions
In this work, we provide an extensive analysis of stochastic optimization under global K Ł conditionand answer all the above questions in the afﬁrmative.
More precisely, our contributions are as follows
[Figure]
• We provide a new framework for the analysis of the dynamics of SGD under global K Ł
condition (see Section 3).
It is based on analysing the dynamic of SGD which is govern by a
recursive inequality (see Equation (7)).
As a result of this analysis, we introduce a set of
conditions (see Theorem 1) for designing proper setpsizes to guarantee convergence.
• Using this framework, we provide sharp analysis of convergence rate of SGD under a
general ES assumption (see Assumption 4) and demonstrate that the sample complexity O  [Equation]
[Equation]
/ ↵ f
is tight for the dynamical system describing SGD.
More precisely, if the recur- sive inequality describing the behaviour of SGD (see Equation (7)) becomes an equality, then its convergence rate cannot be better rate than O [Equation]
unless additional assumptions are applied Next, we propose a variance reduction scheme with parameter restart PAGER .
Carefullychosen sequence of parameters of PAGER allows the algorithm to adapt to the nonconvex
K Ł geometry of the problem and establish state-of-the-art convergence rates for minimizing↵ -P Ł functions.
In online setting (1), we establish O
[Equation]
sample complexity of PAGER ,which beats O [Equation] ⌘ complexity of SGD for the whole spectrum of parameters↵ 2 [1 , 2) .
In particular, for the important special case of 1 -P Ł , this leads to the ﬁrst optimalalgorithm with O ( ✏  2 ) sample complexity, which already matches with the lower
bound
known for stochastic convex optimization [35].• Furthermore, we obtain faster rates with PAGER in ﬁnite sum case (2), providing the ﬁrstacceleration over GD and SGD under ↵ -P Ł condition.
In Table 1, we summarize the sample complexity results for stochastic optimization under ↵ -P Ł and
BV assumptions.
We also establish sharp convergence results for convergence in the iterates to theset X ? of optimal points and provide a summary Table 2 in the Appendix.
Assumptions and discussion
In this section, we introduce the assumptions we make throughout the paper.
Assumption 1.
We assume that the gradient of f is Lipschitz continuous, that is, for all x and y ,
[Equation]
where L > 0 is called the Lipschitz constant. [Table]
Furthermore, we assume that the objective function f is lower bounded, i.e., f ⇤ :=
inf x f ( x ) > 1 ,and it satisﬁes the following inequality
Assumption 2 (global KL or global Kurdyka- Ł ojasiewicz) .
Let  : R + ! R + be a continuousfunction such that  (0) = {+
+}
and  2 ( · ) is convex.
The function f ( · ) is said to satisfy global KurdykaŁ ojasiewicz inequality if
[Equation]
all
[Equation]
Assumption 3 ( ↵ -P Ł or Polyak- Ł ojasiewicz) .
There exists ↵ 2 [1 , 2] and µ > 0 such that
[Equation]
We refer to ↵ as the PL power and µ as the PL constant.
It is straightforward to see that the ↵ -P Ł is a special case of KL with  ( t ) = p 2 µ t 1 / ↵ .
Connections with other assumptions.
An alternative commonly adopted way to deﬁne the global KL property is: ⇢ 0 ( f ( x )  f ? ) · kr f ( x ) k  1 for all x 2 R d , where ⇢ ( t ) is called a disingularizingfunction and ⇢ 0 ( · ) denotes its derivative.
Moreover, disingularizing function satisﬁes the followingconditions, it is continuous and concave, ⇢ (0) = 0 , and ⇢ 0 ( t ) > 0 [28].
If the above assumption holds for ⇢ ( t ) :
= 1 ✓ t ✓ and ✓ > 0 , then Assumption 3 holds with power↵ = 11  ✓ .
However, ↵ -P Ł is more general since it allows to consider the case ↵ = 1 .
For example,consider the function of one variable f ( x ) =
( e x + e  x ) / 2  1 , then | f 0 ( x ) |  f ( x ) for all x .
Thusf ( x ) satisﬁes Assumption 3 with ↵ = 1 and µ = 1 / 2 .
Moreover, this function is convex, but it doesnot satisfy inequality ⇢ 0 ( f ( x )  f ? )
· kr f ( x ) k  1 for any choice of ⇢ ( t ) .
We also provide severalnon-convex examples for which ↵ -P Ł holds with ↵ = 1 in the Appendix.
The intuition behind the special case ↵ = 1 is that the function is allowed to be ﬂat near the set ofoptimal points X ?
= arg min x f ( x ) .
Assumption 4 ( k -ES, Expected
Smoothness of order k ) .
The stochastic gradient estimator g k ( x, ⇠ )is an unbiased estimate of the gradient r f ( x ) at any given point x and its second moment satisﬁes
[Equation]
where A, B, C are non-negative constants.
h : R + ! R + is a concave continuously differentiablefunction with h 0 ( t )  0 , h (0)
= 0 , and b k is a non-decreasing sequence of k .
The expectation istaken over random variable ⇠ ⇠ D .
We call b k the cost of such estimator.
This assumption encompasses previous assumptions in the literature.
For instance, it is straightforwardto see that an estimator satisﬁes bounded variance assumption [15] when h ( t ) = 0 , B = 1 , and ⌧ = 0in (6).
Gradient estimators with relaxed growth assumption [4, 8] are also special cases of (6) forh ( t ) = [-
-]0 and b k = 1 .
A closely related assumption to the relaxed growth was introduced in [43] whichholds when h ( t ) = 0 and b k = 1 in (6).
Expected smoothness assumption [23, 18, 47] is the closestassumption to ours and it holds
when h ( t ) =
t and b k = 1 .
Notably, ES assumption is satisﬁed inpractical scenarios such as mini-batching, importance sampling and compressed communication [23].
More recently, it has been shown that PG method with softmax policies and log barrier regularizationcan be modeled using ES assumption [46].
We highlight two special cases for the sequence b k : b k = ⇥ ( k ⌧ ) with ⌧  0 and b k = ⇥ ( q k ) withq > 1 .
For example, Monte Carlo sampling and mini-batching allow us to design estimators with suchb k s.
When a gradient estimator satisﬁes Assumption 4, unless b k is bounded for all k , it essentiallymeans that we have a mechanism to reduce its variance.
More precisely, the variance decreasesaccording to the sequence 1 /b k .
As we show in Section 4, if such estimator exists, it results in abetter convergence rate compared to vanilla SGD and the improvement is captured by sequence b k .
Note that due to the ﬁrst term in (6), i.e., 2 A · h ( f ( x )  f ⇤ ) , the second moment can be larger whenthe objective gap at x is large as well.
Such property does not hold for estimators with boundedvariance assumption.
An advantage of existing this term is discussed by [19]. Namely, [19] showsthat when h ( t ) = t and A > 0 , then this model allows functions with up to quadratic growth. Hence,it allows us to consider many classic problems which fundamentally have quadratic growth (e.g., anyquadratically regularized problem such as training a support vector machine).
Here we further extendthis model to cover functions beyond the quadratic growth regime, by allowing h ( t ) to be a largerclass of non-decreasing concave functions. Assumption 5 ( k -BV, bounded variance ) . Let Assumption 4 hold with A = 0 , B = 1 and C =  2 . Then k -BV assumption holds, i.e., E h k g k ( x, ⇠ )  r f ( x ) ki   2 b k .
Stochastic Gradient Method
Algorithm 1 summarizes the steps of a slightly modiﬁed SGD which we analyze in this work.
Wecall this algorithm a SGD with restarts.
This algorithm updates the point x for T number of iterationswithin an inner-loop.
Note that, in the inner-loop, the step-size remains unchanged and the iteratesare updated via x t +1 = x t  ⌘ g k ( x t , ⇠ t ) , where ⌘ is the step-size and ⇠ t s are independent randomvariables.
The cost and the variance of the gradient estimator g k ( x, ⇠ ) remain the same within theinner loop of Algorithm 1.
Algorithm 1:
Initialization do 3: ⌘ ⌘ k 4: do 5: x x  ⌘ g k ( x, ⇠ t ) 6: return x
Let { x t } t  0 be the sequence of points generated by the inner loop of Algorithm 1. Given Assumptions 189 1, 2, 4, and the update rule in Algorithm 1, the dynamics of SGD in the inner-loop are as follows.
Lemma 1.
Under Assumptions 1, 2, and 4 with constant cost, i.e., b :=
b k , we obtain
[Equation]
where  t := E [ f ( x t )  f ? ] , a :=
LA , d := LC 2 , ⌘ := ⌘ k .
Understanding the dynamics of this recursion, allows us to establish the global convergence of SGD.
Our approach consists of two main steps: i)
Finding the stationary 2 point of (7) when the inequalityis replaced by equality and for a ﬁxed step-size, ⌘ k = ⌘ , which we denote by r ( ⌘ ) .
ii) Selecting thestep-sizes { ⌘ k } and sequence { b k } , such that the corresponding stationary points { r ( ⌘ k ) } (deﬁned 196 below) converge to zero as k increases.
The stationary point of (7) after replacing inequality with equality must satisfy the equation:
[Equation]
Let us call this stationary point r ( ⌘ ) .
To complete the ﬁrst step, we approximate r ( ⌘ ) by a polynomialfunction of ⌘ .
In other words, we ﬁnd ⌫ 2 R + such that r ( ⌘ ) = ⇥ ( ⌘ ⌫ ) .
For the second step of our framework, we should design the stepsizes.
Next result introduces a set ofconditions that allow us to design the stepsizes and obtain its corresponding convergence rate.
Allproofs are presented in the Appendix.
Theorem 1.
Suppose there exist ⌫ , { ! j } j  0 , and ⇣  0 such that ⌘ k = ⇥ ( k  ⇣ ) , r ( ⌘ k ) = ⇥
( k  ⇣⌫ ) ,| 1  ! k | < 1 , and
[Equation]
Then,  k
= O
( k  ⇣⌫ ) and the iteration complexity of Algorithm 1 is O ( ✏  1 / ( ⇣⌫ ) f ) .
As an application of Theorem 1, we present the convergence rate of SGD for ↵ -P Ł functions.
Corollary 1.
Consider a special case of Assumption 4 with h ( t ) =
t  and b k = k ⌧ , where  2 (0 , 1]and ⌧  0 .
Suppose the objective function f satisﬁes Assumptions 1 and 3.
Let  :=
↵ .
Then, forany ✏ f > 0 , Algorithm 1 returns a point x with E [ f ( x )  f ? ]  ✏ f after N := K · T iterations.i) When
= 2 ( ↵ = 2 and  = 1 ), we have
N
[Equation]
O (
[Equation]
[Equation]
⌧ f ) , with ⌘ k = ⇥ k  1 ) .
ii)
When  < 2 , the convergence rate is
N =
O  ✏  4
[Equation]
[Equation] +1) f
with ⌘ k = ⇥ ( k  ⌧ +1 2  ↵ + ⌧ ) if ⌧ 
[Equation]   , and N = O
[Equation] with ⌘ k = ⇥ ( k  [Equation]
[Equation]
[Equation]
To verify the above result empirically, we simulated  t in (7) throughout all iterations of Algorithms 1for different sets of parameters and presented the results in Figure 1 along with their correspondingconvergence rates given in Corollary 1.
As it is shown in these ﬁgure, the above convergence ratescorrectly capture the behaviour of the dynamics in (7).
As an example, in Figure 1(b), the red solidcurve shows the rate of  k , i.e., log(  k ) as a function of log( k ) for when  = 1 . 1 , ↵ = 1 . 4 , and⌧ = 0 .
Based on Corollary 1, the convergence rate of Algorithm 1 for this setting is O ( ✏  0 .
98 f ) orequivalently log(  k ) =
(  ↵ ( ⌧ +1) 4  ↵ ) log( k ) ⇡  1 .
02 log( k ) which is shown by red dashed line.
Next,we discuss how the results in Corollary 1 generalizes the existing work in the literature.
Remark 1.
Authors in [23] studied the convergence of SGD for 2 -P Ł objectives, under a strongerassumption than Assumption (6) .
More precisely, they assumed an estimator that satisﬁes Assumption4 with ⌧ = [-
-]0 and h ( t ) =
t and obtained the convergence rate of O ( ✏  1 f ) .
This is consistent with ourrate presented in the Corollary 1.
It is noteworthy in this setting, O ( ✏  1 f ) is optimal [23].
The authors in [47] studied the performance of SGD for 1 -P Ł objectives.
Having an estimatorsatisfying Assumption 4 with ⌧ = 0 and h ( t )
=
t , they showed the convergence rate of O ( ✏  3 f ) whichcan be obtained from Corollary 1 by setting ⌧ = 0 ,  = 1 , and ↵ = 1 .[14]
studied SGD with ↵ -P Ł objectives for ↵ 2 [1 , 2] with an estimator that satisﬁes Assumption 4with A = 0 , B = 1 , and b k = 1 and obtained similar convergence rate to ours. More precisely, theirresults is a special case of Corollary 1, when  = 0 and ⌧ = 0 . 3.2 Mini-batching and SGD
In what follows, we prove that mini-batching technique cannot improve the rates that SGD achievesin Corollary 1. Consider a mini-batch setting in which the unbiased gradient estimator is given by g k ( x k , ⇠ k ) : =
1 b k P b k i =1 r f ⇠ ik ( x k ) , where { r f ⇠ i ( x k ) } i  1 are queries provided by an oracle such that E [ r f ⇠ i ( x k )] = r f ( x k ) and E [ || r f ⇠ i ( x k )  r f ( x k ) || 2 ]   2 for all k .
The variance of this estimator diminishes linearly in the size of the mini-batch b k , i.e., when E [ || g k ( x k , ⇠ k )  r f ( x k ) || 2 ]  237  2 /b k .
7.5 log(k) - - - - - - - l og ( k ) =1.1, =1.2 -0.38=1.1, =1.4 -0.
=1.1, =1.7 -0.74=1.1, = =2, =2 - (a) ⌧ = 0 .
6.1 6.2 6.3 6.4 log(k) - - - - - - l og ( k ) =1.1, =1.4 -1.02=1.1, =1.7 -1.
=1.1, =2 =2, =2 -1.
(b) ⌧ = 0 . 9 .
6.1 6.2 6.3 6.4 log(k) - - - - - - l og ( k ) =1.1, =1.4 -0.93=1.1, =1.7 -1.
=1.1, =2 -2.
(c) ⌧ = 2 .
Lemma 2.
Consider the setting of Corollary 1.
The optimal adaptive mini-batch size for Algorithm 1is b k = ⌦ ( k
[Equation]
Given the above result, Algorithm 1 requires total of
P
Nk =1 k  4  ↵   = ⌦ ( N 4  ↵ 4  ↵   ) number ofqueries after N iterations.
This means that Algorithm 1 with optimal mini-batch sizes ﬁnds a global✏ f -stationary point with the total complexity of the order
O  ✏  4  ↵↵ f  .
On the other hand, as it isshown in Corollary 1, Algorithm 1 with an estimator that satisﬁes Assumption 4 with b k = 1 (withoutmini-batching) can achieve the same sample complexity. In other words, mini-batching cannotimprove the complexity of Algorithm 1.
3.3 Tightness of rates in Corollary
In this section, we show that when ⌧ = 0 , the convergence rates presented in Corollary 1 are tight forthe dynamic in (7).
More precisely, if there exists a function f and a gradient estimator satisfying theassumptions in Corollary 1 such that its corresponding recursive inequality (7) is an equality, then itsconvergence rate, presented in Corollary 1 is optimal, i.e., it cannot be improved by any choices ofstepsizes. To do so, for any given 0 < "  1 , we introduce a sequence r 0 , " , r 1 , " , r 2 , " , ... such thatthere exists K 0 > 0 and for all k  K 0 , r k, "   k . Moreover, we show that to ensure r k, "  ✏ , kmust be of the order ⌦ ( ✏  4  ↵  " ↵ + " ) .
Next proposition summarizes our results about the tightness of ourconvergence rates in Corollary 1.
Proposition 1.
Consider the following recursion
[Equation]
where a, d > 0 , h ( t ) =
t  with  2 (0 , 1] ,  ( t ) =
p 2 µt 1 / ↵ with ↵ 2 [1 , 2] , and b k
= O (1) .
Then, the optimal convergence rate of  k is O ( k  ↵ 4  ↵ ) .
Moreover, this rate can be achieved by⌘ k = ⇥ ( k  1 2  ↵ / 2 ) .
Faster Rates with Variance Reduction
To simplify the exposition of the results in this section, let us assume that g k ( x t , ⇠ t ) is constructedexplicitly via mini-batching for which Assumption 5 holds.
Additionally, we assume that we haveaccess to g 0 k
( x, ⇠ ) , which satisﬁes Assumption 7.
Assumption 6 (Average L -smoothness (of order k )) .
Let g 0 k
( x, ⇠ ) :
= 1 b 0 k P b 0 k i =
1 r f ⇠ i ( x ) andg 0 k ( y, ⇠ ) : =
1 b 0 k
P b 0 k i =
1 r f ⇠ i ( y ) be unbiased mini-batch estimators of the gradient of f ( · ) at pointsx and y , respectively for shared stochasticity ⇠ i ⇠ D for each i = 1 , . . .
, b k and ⇠ = ( ⇠ 1 , . .
, ⇠ b k ) .
Deﬁne e  ( x, y ) :
= g 0 k
( x, ⇠ )  g 0 k
( y, ⇠ ) .
The average L -smoothness (of order k ) holds if there exists L  0 such that
[Equation]
all [Equation] where  ( x, y ) :
= r f ( x )  r f ( y ) .
Remark 2.
The Assumption 7 holds in several standard settings.
For instance, if each r f ⇠ i ( x )is Lipschitz with constant ¯ L (almost surely or on average), then Assumption 7 holds with L  ¯ L .
Another example is when f ( · ) is of the form (2) and b k = n , then L = 0 .
We remark from the analysis of Algorithm 1 in Section 3 that merely playing with choice of ⌘ k and b k (chosen as polynomial functions of k ) is not sufﬁcient to improve the convergence, hence, we need to construct more sophisticated gradient estimator and reduce the variance using control variate.
Now, we highlight the main algorithmic ingredients of our construction.
First, let us describe the variance reduced estimator named PAGE , which was originally proposed and shown to be optimal for ﬁnding ﬁrst order stationary point [29].
Moreover, PAGE is easy to implement and designed via a small modiﬁcation to mini-batch SGD.
The update of PAGE estimator writes as
[Equation]
where p is a small probability and mini-batch sizes satisfy b > b 0 .
Now we are ready to introduce PAGER (Algorithm 2) – a Probabilistic Average Gradient Estimatorwith parameter Restart, which is formally described in Algorithm 2.
It takes as input the sequenceof parameters { ⇤ k :=
( ⌘ k , T k , p k , b k , b 0 k ) : k = 0 , ..., K  1 } , where T k is the length of stage k , 276 ⌘ k , p k , b k , b 0 k step-size, probability, and batch-sizes at stage k .
PAGER updates this sequence ofparameters in the outer loop k = 0 , . .
, K  1 and applies PAGE estimator with a ﬁxed set ofparameters sequences in the inner loop t = 0 , . .
, T k  1 .
We will select { ⇤ k } k  0 depending on the P Ł power ↵ to capture the dependence on the geometry of the problem and establish fast rates foreach ↵ in settings (1) and (2). 1: Initialization } 2: for
1 do 3:
4: ( ⌘ , p, , b ) ( ⌘ k , p k , b k , b k ) 5: for t = 0 , . . , T k  1 do 6: x t +1 = x
t  ⌘ g t 7:
Sample  ⇠ Be ( p ) 8: if  = 1 then 9:
Query g k ( x t +1 , ⇠ t +1 ) with batch-size b k 10:
Set g t +1 =
g k ( x t +1 , ⇠ t +1 ) 11: else 12:
Query g 0 k ( x t +1 , ⇠ t +1 ) and g 0 k ( x t , ⇠ t +1 )
each with batch-size b 0 k 13:
Set g t +1 =
g t + g 0 k
( x t +1 , ⇠ t +1 )
g 0 k
( x t , ⇠ t +1 ) 14: (¯ x t +1 , ¯ g t +1 )
( x t +1 , g t +1 ) 15: Return: ¯ x T
4.1 Online case
We present convergence guaranties for Algorithm 2 in the setting (1).
Theorem 2.
Let f ( · ) has the form (1) and satisfy Assumptions 1, 3 (with ↵ 2 [1 , 2) ) and 7, letthe sequences in Algorithm 2 be chosen as b 0 k = ⇥
2 (2  ↵ ) k ↵  , p k = ⇥  2  (2  ↵ )
k ↵  , b k = ⇥  2 2 k ↵  , T k = ⇥  2 (2  ↵ ) k ↵  , ⌘ k = ⇥  2  (2  ↵ ) k ↵  .
Then, for any ✏ f > 0 Algorithm 2 returns a point x with E [ f ( x )  f ? ]  ✏ f after N :=[-(cid:80) K − 1 k =
-]
P
[Equation] = O [Equation] f  iterations.
The expected totalcomputational cost (sample complexity) [-is O (cid:16) ϵ − 
-]iscost :=
[Equation]
⌘ . Remark 3.
For brevity, in Theorem 6 we deﬁne the input sequences up to constants hidden in ⇥ ( · )notation. In fact, our analysis allows to specify these constants and we present detailed derivationsin Appendix. Theorem 2 implies that PAGER improves the convergence of [-SGDfrom O (cid:0) ϵ − 4 − αα f (cid:1) to O (cid:0) ϵ − 
 α f (cid:1) under α -PŁ condition for-]SGD by the order of O  ✏  2  ↵↵ f  under↵ -P Ł condition.
In the interesting case of ↵ = 1 , this leads to O  ✏  2 f  sample complexity comparedto the best known O  ✏  3 f  for SGD , see discussion in Section 3.
Additionally, as a consequence of our analysis we obtain optimal last iterate convergence for convexstochastic optimization under the additional assumption that the iterates of PAGER remain bounded ,i.e., k x t  x ? k 2  D for all t  0 , where x ? 2 arg min x f ( x ) .
3 Consider the interesting special caseof 1 -P Ł . In this setting PAGER has O  ✏  sample complexity for the last iterate of the algorithm.
Since the iterates of the algorithm are bounded, convexity hr f ( x ) , x  x ?
i  f ( x )  f ( x ? ) implies 299 1 -P Ł with µ = D .
This observation implies the last iterate convergence of PAGER for convexoptimization under additional average L -smoothness assumption. The standard analysis of SGD forconvex functions provides convergence for the average iterate [24].
The last iterate convergence forconvex functions was recently established for SGD in [14] by following an involved analysis with acareful control of iterates via sufﬁx-averaging scheme.
4.2 Finite sum case
Let f ( · ) have the ﬁnite sum form (2).
Then we obtain the following result.
Theorem 3.
Let f ( · ) has the form (2) and satisfy Assumptions 1, 3 (with ↵ 2 [1 , 2) ) and 7, letthe sequences in Algorithm 2 be chosen as p k = [Equation] +1 ,
[Equation] [Equation]
min [Equation]  .
Then, for any ✏ f > 0 , Algorithm 2 returns a point x with E [ f ( x )  f ? ] ✏ f after N := [-(cid:80) K − 1 k =
-][Equation]
iterations.
The expected total computational cost(sample complexity) is [Equation]
Theorem 3 quantiﬁes the improvement of PAGER over GD in the ﬁnite sum setting in terms of n .
Recall that GD has sample complexity O
[Equation]
, when n is large, we get the improvement oforder p n .
Notice that in the limit ↵ ! 2 , this improvement matches the best known improvementfor 2 -P Ł objectives [29].  Conclusion
We analyzed the complexity of SGD when the objective satisﬁes global KL inequality and thequeries from stochastic gradient oracle satisfy weak expected smoothness.
We introduced a generalframework for this analysis which resulted in a sample complexity of O ( ✏  (4  ↵ ) / ↵ ) for SGD withobjectives satisfying ↵ -P Ł condition.
We also demonstrated the tightness of this rate under thespeciﬁc choice of stepsizes.
Last but not least, we developed a modiﬁed SGD with variance
Checklist
For all authors...
(a) Do the main claims made in the abstract and introduction accurately reﬂect the pa- per’s contributions and scope?
[Yes] We accurately state the main contributions in introduction.
(b) Did you describe the limitations of your work?
[Yes]
(c) Did you discuss any potential negative societal impacts of your work?
[N/A]
(d) Have you read the ethics review guidelines and ensured that your paper conforms to them?
[Yes]
If you are including theoretical results...
(a) Did you state the full set of assumptions of all theoretical results?
[Yes]
(b) Did you include complete proofs of all theoretical results?
[Yes] We provide all the technical proofs in Appendix.3.
If you ran experiments...
(a) Did you include the code, data, and instructions needed to reproduce the main experi- mental results (either in the supplemental material or as a URL)?
[N/A]
(b) Did you specify all the training details (e.g., data splits, hyperparameters, how they were chosen)?
[N/A]
(c) Did you report error bars (e g., with respect to the random seed after running experi- ments multiple times)?
[N/A]
(d) Did you include the total amount of compute and the type of resources used (e g., type
of GPUs, internal cluster, or cloud provider)?
[N/A]
If you are using existing assets (e g., code, data, models) or curating/releasing new assets...
(a) If your work uses existing assets, did you cite the creators?
[N/A]
(b) Did you mention the license of the assets?
[N/A]
(c) Did you include any new assets either in the supplemental material or as a URL?
[N/A]
(d) Did you discuss whether and how consent was obtained from people whose data you’re
using/curating?
[N/A]
(e) Did you discuss whether the data you are using/curating contains personally identiﬁable
information or offensive content?
[N/A]
If you used crowdsourcing or conducted research with human subjects...
(a) Did you include the full text of instructions given to participants and screenshots, ifapplicable?
[N/A]
(b) Did you describe any potential participant risks, with links to Institutional Review
Board (IRB) approvals, if applicable?
[N/A]
(c) Did you include the estimated hourly wage paid to participants and the total amountspent on participant compensation?
[N/A]