Sharp Analysis of Stochastic Optimization under Global Kurdyka-Łojasiewicz Inequality
Anonymous Author(s)
AffiliationAddress email
Abstract
We study the complexity of finding the global solution to stochastic nonconvexoptimization when the objective function satisfies global Kurdyka-Łojasiewicz(KL) inequality and the queries from stochastic gradient oracles satisfy mild expected smoothness assumption.
We first introduce a general framework to analyze Stochastic Gradient Descent (SGD) and its associated nonlinear dynamics underthe setting.
As a byproduct of our analysis, we obtain a sample complexity of O ( ϵ − (4 − α ) /α ) for SGD when the objective satisfies the so called α -PŁ condition,where α is the degree of gradient domination.
Furthermore, we show that a modified SGD with variance reduction and restarting (PAGER) achieves an improvedsample complexity of O ( ϵ − 2 /α ) when the objective satisfies the average smoothness assumption.
This leads to the first optimal algorithm for the important case ofα = 1 which appears in applications such as policy optimization in reinforcementlearning.
Introduction
Nonconvex optimization problems are ubiquitous in machine learning domains such as training deepneural networks [21] or policy optimization in reinforcement learning [50].
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 nonconvex setting has become animportant avenue of research in recent years [20, 4, 28, 42, 22, 14, 57].
We are interested in solving the unconstrained stochastic , nonconvex optimization problem of theform
[Equation]
(1) 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]
For a general nonconvex differentiable objective f : R d → R , finding a global minimum of f isin general intractable [40, 52]. There are two common strategies to analyze optimization methodsfor nonconvex functions.
The first one is to scale down the requirements on the solution of interestfrom global optimality to some relaxed version, e.g., first-order stationary point.
However, suchsolutions do not exclude the possibility of approaching a suboptimal local minima or a saddle point.
Another approach is to study nonconvex problems with additional structural assumption in the hope ofconvergence to global solutions.
In this direction, several relaxations of convexity have been proposed and analyzed, for instance, star convexity, quasar-convexity, error bounds condition, restricted secantinequality, and quadratic growth [25, 23, 22].
Many of the aforementioned relaxations have limitedapplication 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 [8, 9].
Of particular interest is the family of functions thatsatisfy global KŁ inequality.
Specifically, we say that f ( · ) satisfies (global) KŁ inequality if thereexists some continuous function ϕ ( · ) such that ||∇ f ( x ) || ≥ ϕ ( f ( x ) − inf x f ( x )) for all x ∈ R d .
If this inequality is satisfied for ϕ ( t ) = √ 2 µ t 1 /α , then we say that (global) α -PŁ condition issatisfied for f ( · ) .
The special case of KŁ condition, 2 -PŁ, often referred as Polyak-Łojasiewiczor PŁ condition, was originally discovered independently in the seminal works of B. Polyak, T. Ležanski and S. Łojasiewicz [46, 31, 36, 37].
Notably, this class of problems has found manyinteresting emerging applications in machine learning, for instance, policy gradient (PG) methodsin reinforcement learning [39, 1, 54], generalized linear models [38], over-parameterized neuralnetworks [2, 55], linear quadratic regulator in optimal control [12, 18], and low-rank matrix recovery[7].
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[25, 22, 28, 51, 47].
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
Stochastic gradient descent. A plethora of existing works has studied the sample complexity of SGD and its variants for finding an ϵ -stationary point of general nonconvex function f , that is, apoint x ∈ R d for which E [ ||∇ f
(ˆ x ) || ] ≤ ϵ .
For instance, [20] showed that for a smooth objective(one with Lipschitz gradients) under bounded variance (BV) assumption, SGD with properly chosenstepsizes reaches an ϵ -stationary point with the sample complexity of O ( ϵ − 4 ) .
Recently, [28, 54]further extended the result under a much milder expected smoothness (ES) assumption on stochasticgradient.
While this sample complexity is known to be optimal for general nonconvex functions, anaive application of this result to the function value using α -PŁ condition would lead to a suboptimal O ( ϵ − 4 / α f ) sample complexity for finding an ϵ f -optimal solution, i.e., E [ f ( x ) − f ⋆ ] ≤ ϵ f .
Recently,[19] studied SGD and established convergence rates for α -PŁ functions under BV assumption.
Theirsample complexity result is O ( ϵ − (4 − α ) / α f ) in our notation.
Later [33] considered SGD scheme withrandom reshuffling under local and global KŁ conditions and provided convergence in the iterates forα ∈ (1 , 2] .
We note that our proof techniques are different from [19] and [33] and are not limitedmerely to the case of BV assumption.
In this work, we will answer the following open question: What is the exact performance limit of SGD under global KŁ condition and a more practical modelof stochastic oracle?
Variance reduction. There has been extensive research on development of algorithms whichimprove the dependence on n and/or ϵ for both problems (1) and (2) (over simple methods such as
SGD and Gradient Descent ( GD ) ).
One important family of techniques 1 is variance reduction , whichhas emerged from the seminal work of Le Roux et.
al [48].
The main idea of variance reductionis to make use of the stochastic gradients computed at previous iterations to construct a bettergradient estimator at a relatively small computational cost.
Various modifications, generalizations,and improvements of the variance reduction technique appeared in subsequent work, for instance,[48, 27, 15] to name a few.
Finite-sum case. A number of recently proposed algorithms such as SNVRG [56], SARAH [43],
STORM [14], SPIDER [16], and PAGE [34] achieve the sample complexity O (cid:0) n + √ nϵ 2 (cid:1) whenminimizing a general nonconvex function with finite sum structure (2).
This result is also known tobe optimal in this setting [34].
[26] studies SARAH in finite sum case under local KŁ assumption andproves convergence in the iterates.
The study in [26] 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.
[32] proposes an SVRG -based algorithm for KŁ functions and [53, 43, 44]study other variance reduction techniques, but they only analyze the special case α = 2 . Under
2-PŁ condition, these methods further improve to O (cid:0) ( n + κ √ n ) log( 1 ϵ f ) (cid:1) sample complexity 2 forfinding an ϵ f -optimal solution.
However, it is not clear if it is possible to provide any non-asymptoticguaranties for variance reduced methods under α -PŁ condition for any α ∈ [1 , 2) .
In our work, wewill answer the following open question: What is the extent of improvement any variance reductionscheme can provide under global α -PŁ condition for finite-sum objectives of the form (2) ?
Online/streaming case. While variance reduction methods have been initially designed for problems of the form (2), it was later discovered that they also improve over SGD when solving (1)[30, 35] 3 .
The analysis of these methods was obtained for general nonconvex functions (for minimizing the norm of the gradient, E [ ∥∇ f
(ˆ x ) ∥ ] ≤ ϵ ) and later extended to 2 -PŁ objectives for minimizingthe function value, E [ f ( x ) − f ⋆ ] ≤ ϵ f . For example, the methods in [56, 43, 16, 34] achieve O ( ϵ − 3 )complexity improving over O ( ϵ − 4 ) complexity of SGD for finding an ϵ -stationary point. Under the2 -PŁ condition, these results can be extended to global convergence with O ( ϵ − 1 f ) sample complexity[34]. However, in contrast to a general nonconvex case, variance reduction under 2 -PŁ assumptiondoes not show any improvement over SGD in terms of ϵ f .
We highlight that all existing analysisof variance reduction under α -PŁ condition is restricted only to a special case α = [-
-]2 .
We referthe reader to Appendix C, where we elaborate on the key difficulties in the analysis for the casesα ∈ [1 , 2) .
Since the direct analysis for α ∈ [1 , 2) is challenging, in order to obtain the globalconvergence in this setting, one could naively translate the complexity for finding a stationary pointof a general nonconvex function (which is O ( ϵ − 3 ) ) to convergence in a function value by using α -PŁcondition: √ 2 µ ( f (ˆ x ) − f ⋆ ) 1 / α ≤ ||∇ f
(ˆ x ) || .
This would result in O (cid:0) ϵ − 3 / α f (cid:1) sample complexity.
However, there are two serious issues with this approach.
First, this complexity does not provide anyimprovement over SGD in the most interesting practical case α = 1 and gives strictly worse resultfor all α > 1 .
Second, the guarantees for general nonconvex optimization hold on average, in thesense that the point ˆ x is sampled uniformly from all the iterates of the algorithm.
It would be moredesirable to instead derive last iterate convergence guarantees under KŁ ( α -PŁ) condition.
In thiswork, we will address the following open question: Is it possible to accelerate the O (cid:0) ϵ − (4 − α ) / α f (cid:1)sample complexity of SGD under global α -PŁ 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 affirmative.
More precisely, our contributions are as follows
• We provide a new framework for the analysis of the dynamics of SGD under global KŁ
condition (see Section 3).
It is based on the analysis of SGD dynamic which is governed by
a recursive inequality (see Equation (6)).
As a result of this analysis, we introduce a set of
conditions (see Theorem 1) for designing proper stepsizes to guarantee convergence.
• Using this framework, we provide sharp analysis of SGD under a general ES assumption
(Assumption 4) and demonstrate that the sample complexity O (cid:0) ϵ − (4 − α ) / α f (cid:1) is tight for the
dynamical system describing SGD .
• Next, we propose PAGER , a new variance reduction scheme with parameter restart.
A
carefully chosen sequence of parameters of PAGER allows the algorithm to adapt to the
nonconvex geometry of the problem and establish state-of-the-art convergence guarantees for minimizing α -PŁ functions.
In online setting (1), we obtain O (cid:16) ϵ − 2 / α f
(cid:17)
sample complexity
of PAGER , which beats O (cid:16) ϵ − (4 − α ) / α f (cid:17) complexity of SGD for the whole spectrum of parameters α ∈ [1 , 2) .
In particular, for the important special case of 1 -PŁ, this leads to the first optimal algorithm with O ( ϵ − 2 f ) sample complexity, which already matches with the
lower bound known for stochastic convex optimization [40].
• Furthermore, we obtain faster rates with PAGER in finite sum case (2), providing the first acceleration over GD and SGD under α -PŁ condition. Under additional assumptions such as smoothness of individual
functions
f ξ ( · ) or even milder condition such as average L - smoothness (Assumption 6). 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 in Table 2 in the Appendix. [Table]
Assumptions and Discussion
In this section, we introduce the assumptions we make throughout the paper.
Assumption 1.
The gradient of f ( · ) is Lipschitz continuous, that is, for all x and y ,∥∇ f ( x ) − ∇ f ( y )
∥
≤ L ∥ x − y ∥ , L > 0 is referred to as the Lipschitz constant.
Furthermore, we assume that the objective function f is lower bounded, i.e., f ∗ :=
inf x f ( x ) > −∞ ,and it satisfies the following inequality
Assumption 2 (global KŁ 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]
Assumption 3 ( α -PŁ or Polyak-Łojasiewicz) .
There exists α ∈ [1 , 2] and µ > 0 such that
[Equation]
We refer to α as the PŁ power and µ as the PŁ constant.
It is straightforward to see that the α -PŁ is a special case of KŁ with ϕ ( t ) = √ 2 µ t 1 /α .
Connections with other assumptions.
Another commonly adopted way to define the global KŁproperty is to assume that ρ ′ ( f ( x ) − f ⋆ ) · ∥∇ f ( x ) ∥ ≥ 1 for all x ∈ R d , where ρ ( t ) is called adisingularizing function and ρ ′ ( · ) denotes its derivative.
Moreover, disingularizing function satisfiesthe following conditions, it is continuous, concave, ρ (0) = 0 , and ρ ′ ( t ) > 0 [33].
If the above assumption holds for ρ ( t ) :
= 1 θ t θ and θ > 0 , then Assumption 3 is satisfied with PŁpower α = 1 1 − θ .
However, α -PŁ condition 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 ′ ( x ) | ≥ f ( x )for all x .
Thus f ( x ) satisfies Assumption 3 with α = 1 and µ = 1 / 2 .
Moreover, this function isconvex, but it does not satisfy inequality ρ ′
( f ( x ) − f ⋆ ) · ∥∇ f ( x ) ∥ ≥ 1 for any choice of ρ ( t ) .
Wealso provide several non-convex problems for which α -PŁ holds with α ∈ [1 , 2] in the Appendix A.
Other forms of ϕ ( t ) also appear in practice, e.g., squared cross entropy loss function satisfies the KLcondition with ϕ ( t ) =
min { t, √ t } , [49].
The intuition behind the special case α = 1 is that the function is allowed to be flat 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 ∇ f ( x ) at any given point x and its second moment satisfies
[Equation]
where A, B, C are non-negative constants.
h : R + → R + is a concave continuously differentiablefunction with h ′ ( t ) ≥ 0 , h (0)
= 0 .
The expectation is taken over random vector ξ ∼ D .
We call b kthe cost of such estimator.
This assumption encompasses previous assumptions in the literature.
For instance, it is straightforwardto see that an estimator satisfies the standard bounded variance assumption [20] when h ( t ) = 0 , B =1 , and b k = 1 in (5).
Gradient estimators with relaxed growth assumption [6, 10] are also specialcases of (5) for h ( t ) = {+
+} and b k = 1 .
A closely related assumption to the relaxed growth wasintroduced in [51] which holds when h ( t ) = 0 and C = 0 in (5).
Expected smoothness assumption[28, 23, 54] is the closest assumption to k -ES
and it holds when h ( t ) =
t and b k = 1 .
Notably, ES assumption is satisfied in practical scenarios such as mini-batching, importance sampling andcompressed communication [28].
More recently, it has been shown that PG method with softmaxpolicies and log barrier regularization can be modeled using ES assumption [54].
Note that due tothe first term in (5), i.e., 2 A · h ( f ( x ) − f ∗ ) , the second moment can be large when the objectivegap at x is large.
Such property is not captured by standard bounded variance (Assumption 5).
Theflexibility and advantages of introducing such additional term are elaborated in detail in the literature[23, 28, 24, 54].
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 withsuch { b k } k ≥ 0 sequences.
When a gradient estimator satisfies Assumption 4, unless b k is boundedfor all k , it essentially means that we have a mechanism to reduce its variance.
More precisely, thevariance decreases according to the sequence 1 /b k .
As we show in Section 4, if such estimator exists,it results in a better convergence rate compared to vanilla SGD and the improvement is captured bysequence b k .
On the other hand, usually the access to such estimator comes with a cost proportionalto b k , e.g., mini-batch setting described in Section 4.
Thus we refer to b k as the cost of the estimatorg k ( x, ξ ) .
Stochastic Gradient Method
Algorithm 1 summarizes the steps of a slightly modified SGD which we analyze in this work.
Wecall this algorithm 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 iterates areupdated via x t +1 = x t − ηg k ( x t , ξ t ) , where η is the step-size and { ξ t } t ≥ 0 are independent randomvectors.
The cost b k of the gradient estimator g k ( x, ξ ) remain the same within the inner loop of
Algorithm 1.
Let { x t } t ≥ 0 be the sequence of points generated by the inner loop of Algorithm 1, and Assumptions1, 2, 4 are satisfied. Then the dynamics of SGD in the inner-loop of Algorithm 1 is characterized by
Lemma 1.
Under Assumptions 1, 2, and 4 with constant cost, i.e., b :=
b k , we obtain
[Equation]
(6) 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 4 point of (6) when the inequalityis replaced by equality and for a fixed step-size, η k = η , which we denote by r ( η ) .
ii) Selecting thestep-sizes { η k } and sequence { b k } k ≥ 0 , such that the corresponding stationary points { r ( η k ) } k ≥ 0(defined below) converge to zero as k increases.
The stationary point of (6) after replacing inequality with equality must satisfy the equation:
[Equation]
Let us call this stationary point r ( η ) .
To complete the first step, we approximate r ( η ) by a polynomialfunction of η .
In other words, we find ν ∈ R + such that r ( η ) = Θ( η ν ) .
For the second step of our framework, we should design the stepsizes.
Next result introduces a setof conditions that allow us to design the stepsizes, which will guarantee convergence.
The detailedderivations are presented in the Appendix.
Theorem 1.
Suppose there exist ν ≥ 0 , { ω j } j ≥ 0 , and ζ ≥ 0 such that η k = Θ( k − ζ ) , r ( η k ) =
Θ( k − ζν ) , | 1 − ω k | < 1 , and
[Equation]
[Equation]
Then, δ k
= O
( k − ζν ) and the iteration complexity of Algorithm 1 with T =Ω(1 / min j ω j ) is O ( ϵ − 1 / ( ζν ) f ) .
As a consequence of Theorem 1, we present the iteration complexity of SGD for α -PŁ functions.
Corollary 1.
Consider a special case of Assumption 4 with h ( t ) =
t β and b k = k τ , where β ∈ (0 , 1]and τ ≥ 0 .
Suppose the objective function f satisfies 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)
If γ = 2 ( α = 2 and β = 1 ), we have
N = O ( ϵ − 11+ τ f ) , with η k = Θ( k − 1 ) .
ii)
If γ < 2 , we have
N =
O (cid:0) ϵ −
[Equation]
( τ +1) f
(cid:1) with η k = Θ( k − τ +1 2 − α/ 2 + τ ) if τ ≤ γ 4 − α − γ , and N =
O (cid:0) ϵ − 4 − α − γ
α f (cid:1) with η k = Θ( k − 2 − γ 4 − α − γ ) if τ >
γ
− α −
γ . To verify the above result empirically, we simulated δ t in (6) 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 figure, the above convergence ratescorrectly capture the behaviour of the dynamics in (6).
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 γ = 1 . 1 , α = 1 . 4 , and τ = 0 .
Based on Corollary 1, the convergence rate of Algorithm 1 for this setting is O ( ϵ − 0 .
98 f ) or equivalentlylog( δ k ) =
( − α ( τ +1) 4 − α ) log( k ) ≈ − 1 .
02 log( k ) which is shown by red dashed line.
Next, we discusshow the results in Corollary 1 generalizes the existing work in the literature.
Comparison to related works.
Authors in [28] studied the convergence of SGD for 2 -PŁ objectives,under a stronger assumption than Assumption 4.
More precisely, they assumed an estimator thatsatisfies Assumption 4 with τ = {+
+} and h ( t ) =
t and obtained the convergence rate of O ( ϵ − 1 f ) .
Thisis consistent with our rate presented in the Corollary 1.
It is worth noting that in this setting, O ( ϵ − 1 f )is optimal [28].
The authors in [54] studied the performance of SGD for 1 -PŁ objectives.
Assuming that the gradientestimator satisfies Assumption 4 with τ = 0 and h ( t )
=
t , they obtain O ( ϵ − 3 f ) sample complexity. This result can be recovered from Corollary 1 by setting τ = 0 , γ = 1 , and α = 1 .
Note that inthis case, the cost of each iteration is b k = 1 , which means that the iteration complexity coincideswith the sample complexity.
We note that our proof technique is different than in [54] and allows toconsider more general assumptions. Finally, [19] studied SGD with α -PŁ objectives for α ∈ [1 , 2]under bounded variance Assumption 5 with b k = 1 and obtained similar convergence rate to ours.
Werecover their result as a special case by setting A = 0 and τ = 0 in Corollary 1, if we set T = 1 wealso recover the same (up to a constant) step-sizes η k = Θ (cid:16) k − 2 4 − α (cid:17) .
However, we highlight that ourproof technique is different from [19], and more generic since it holds for a general Assumption 4.
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 .
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 .
The result of Corollary 1 suggests that by increasing the cost of the gradient estimator b k over theiterations, one can achieve a better iteration complexity of Algortihm 1.
In particular, it improveswith τ until it reaches the minimum N = O (cid:16) ϵ −
[Equation]
(cid:17) at τ = γ 4 − α − γ and does not change forlarger values of τ . However, we are merely interested in the iteration complexity in practice, sincethe computational cost at each iteration can be prohibitively large. A more adequate measure is thetotal computational cost (sample complexity) of the method.
It is interesting whether increasing b kover the iterations may also result in a better sample complexity for finding an ϵ f -optimal solution,than for the constant choice, e.g., b k = 1 . The following lemma shows the contrary. Proposition 1.
Let the assumptions of Corollary 1 hold, b k = Θ ( k τ ) , T = Θ (1) .
Then the expectedtotal computational cost (sample complexity) of Algorithm 1 is
[Equation]
The above result implies that increasing the cost of the gradient estimator with iterations does notimprove the total sample complexity of Algorithm 1. Therefore, one can simply select b k = 1 ( τ = 0 )and obtain O (cid:16) ϵ − 4 − αα f (cid:17) sample complexity.
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 (6) describing the progress of SGD .
More precisely, if there exists a function f and agradient estimator satisfying the assumptions in Corollary 1 such that its corresponding recursiveinequality (6) is an equality, then its convergence rate, presented in Corollary 1 cannot be improvedby any choices of stepsizes { η k } k ≥ 0 .
Next proposition summarizes our results about the tightness ofour convergence rates in Corollary 1.
Proposition 2.
Consider the following recursion
[Equation]
where a ≥ 0 , d > 0 , h ( t ) =
t β with β ∈ (0 , 1] , ϕ ( t ) = √ 2 µt 1 /α with α ∈ [1 , 2] , and b k = Θ(1) .
Then δ k = Ω( k − α 4 − α ) for any sequence of { η k } k ≥ 0 . Moreover, this rate is achieved by the choiceη k = Θ( k − 1 2 − α/ 2 ) . 1: Initialization: ¯ x 0 , ¯ g 0 , K, { Λ k
= η k , T k , p k , b k , b ′ k ) : k = 0 , ..., K − 1 } 2: for k = 0 , . . . , K − 1 do 3: ( x 0 , g 0 )
← (¯ x k , ¯ g k ) 4: ( η, p, b, 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 χ ∼ Bernoulli ( p ) 8: if χ = 1 then 9: g t +1 = 1 b (cid:80) bi =1 ∇ f ξ it +1 ( x t +1 ) 10: else 11: g t +1 =
g t + 1 b
cid:80) b ′ i =
1 ∇ f ξ it +1 ( x t +1 )
− 1 b ′ (cid:80)
b ′ i =
1 ∇ f ξ it +1 ( x t ) 12: (¯ x k +1 , ¯ g k +1 )
← ( x t +1 , g t +1 ) 13:
Return: ¯ x K
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 g k ( x t , ξ t ) : = 1 b k (cid:80) b
k i =1 ∇ f ξ it ( x t ) , where ξ t := (cid:16) ξ 1 t , . . .
, ξ b k t (cid:17) is arandom vector of independent entries, ξ t are independent for all iterations, {∇ f ξ it ( x t ) }
b k i =1 are queriesprovided by an oracle such that E [ ∇ f ξ it ( x t )] = ∇ f ( x t ) and E [ ||∇ f ξ it ( x t ) − ∇ f ( x t ) || 2 ] ≤ σ 2 forall t ≥ 0 .
The variance of this estimator diminishes linearly in the size of the mini-batch b k , i.e.,g k ( x t , ξ t ) satisfies
Assumption 5 ( k -BV, bounded variance) . Let Assumption 4 hold with A = 0 , B = 1 and C = σ 2 ,i.e., E [Equation]
Additionally, we assume that we have access to a gradient estimator g ′ k
( x, ξ ) , which satisfies thefollowing
Assumption 6 (Average L -smoothness (of order k )) .
Let g ′ k
( x, ξ ) :
= 1 b ′ k (cid:80) b ′ k
i =1 ∇ f ξ i ( x ) andg ′ k ( y, ξ ) :
= 1 b ′ k (cid:80) b ′ k
i =1 ∇ f ξ i
( y ) be unbiased mini-batch estimators of the gradient of f ( · ) at points xand y , respectively for shared stochasticity ξ i ∼ D for each i = 1 , . . .
, b ′ k and ξ = ( ξ 1 , . .
. , ξ b k ) .
Define (cid:101) ∆( x, y ) :
= g ′ k
( x, ξ ) − g ′ k
( y, ξ ) .
The average L -smoothness (of order k ) holds if thereexists L ≥ 0 such that E (cid:20)(cid:13)(cid:13)(cid:13)(cid:101) ∆( x, y ) − ∆( x, y ) (cid:13)(cid:13)(cid:13) 2 (cid:21) ≤ L 2 b ′ k
∥ x − y
∥for all x, y ∈ R d , where∆( x, y ) :
= ∇ f ( x ) − ∇ f ( y ) .
Remark 1.
The Assumption 6 holds in several standard settings.
For instance, if each ∇ f ξ i ( x )is Lipschitz with constant ¯ L (almost surely or on average), then Assumption 6 holds with L ≤ ¯ L .
Another example is when f ( · ) is of the form (2) and b ′ k = n , then L = 0 .
4.1 PAGER – a new variance reduction for α -PŁ objectives
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 sufficient 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 will be the main building block for our Algorithm 2. PAGE was introduced and analyzed in [34] and is known to be optimal for finding a first order stationary point.
Moreover, it is easy to implement and designed via a small modification to mini-batch SGD
[Equation]
where p is a small probability and mini-batch sizes satisfy b > b ′ .
However, while the method looks simple, the extension of its analysis to α -PŁ functions faces severaldifficulties.
Therefore, we introduce a new method, which we call PAGER (Algorithm 2) – a Probabilistic Average Gradient Estimator with parameter Restart .
It takes as input the sequenceof parameters { Λ k :=
( η k , T k , p k , b k , b ′ k ) : k = 0 , ..., K − 1 } , where T k is the length of stage k ,η k , p k , b k , b ′ 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 fixed set ofparameters 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 for each α insettings (1) and (2).
4.2 Online case
We present convergence guarantees for Algorithm 2 in the setting (1).
Theorem 2.
Let f ( · ) have the form (1) and satisfy Assumptions 1, 3 (with α ∈ [1 , 2) ), 5 and6, let the sequences 6 in Algorithm 2 be chosen as b
′ k = Θ (cid:0) 2 (2 − α ) k α (cid:1) , p k = Θ (cid:0) 2 − (2 − α )
k α (cid:1) , b k =Θ (cid:0) 2 2 kα (cid:1) , T k = Θ (cid:0) 2 (2 − α ) k α (cid:1) , η k = Θ (cid:0) 1 (cid:1) .
Then, for any ϵ f > 0 Algorithm 2 returns a pointx with E [ f ( x ) − f ⋆ ] ≤ ϵ f after N := {+(cid:80) K − 1 k =
+}
T k
= O (cid:0) ϵ − 2 − αα f (cid:1) iterations.
The expected totalcomputational cost (sample complexity) {+is O (cid:16) ϵ − 
+}
α f (cid:17) .
Improvement over SGD .
Theorem 2 implies that PAGER improves the sample complexity of SGDfrom O {+(cid:0) ϵ − 4 − αα f (cid:1) to O (cid:0) ϵ − 
 α+} f (cid:1) under α -PŁ condition for the whole spectrum of parameters α ∈ [1 , 2) .
In the case α = 1 , which holds in many interesting applications (see Appendix A for examples), thisleads to O (cid:0) ϵ − 2 f (cid:1) sample complexity compared to the best known O (cid:0) ϵ − 3 f (cid:1) for SGD .
Relation to convex optimization and last iterate convergence. As a consequence of our analysiswe obtain the optimal sample complexity for convex stochastic optimization under the additionalassumption that the iterates of the method remain bounded, i.e., ∥ x t − x ⋆ ∥ 2 ≤ D for all t ≥ 0 ,where x ⋆ ∈ arg min x f ( x ) .
For 1 -PŁ objectives, PAGER has O (cid:16) ϵ − 2 f (cid:17) sample complexity.
Sincethe iterates of the algorithm are bounded, convexity ⟨∇ f ( x ) , x − x ⋆ ⟩ ≥ f
( x ) − f ( x ⋆ ) implies 1 -PŁwith µ = 12 D .
This observation implies convergence of PAGER for convex objectives with O (cid:16) ϵ − 2 f (cid:17)sample complexity, which is known to be non-improvable for convex stochastic optimization [40]. Moreover, we highlight that this result holds for the last iterate of PAGER , while the standard analysisof first order methods for convex functions guarantees convergence for the average iterate [29].
Thelast iterate convergence for convex objectives was only recently established for SGD by following aninvolved analysis with a careful control of iterates via suffix-averaging scheme [19].
4.3 Finite sum case
Let f ( · ) have the finite sum form (2).
Then we obtain the following result.
Theorem 3.
Let f ( · ) have the form (2) and satisfy Assumptions 1, 3 (with α ∈ [1 , 2) ) and 6, letthe sequences be chosen as p k = 1 n +1 , b ′ k = 1 , b k = n , T k
= Θ (cid:0) 2 (2 − α )
k α (cid:1) , η k = Θ (cid:0) 1 (cid:1) .
Then,for any ϵ f > 0 , Algorithm 2 returns a point x with E [ f ( x ) − f ⋆ ] ≤ ϵ f after N := {+(cid:80) K − 1 k =
+}
T k =(cid:101) O (cid:0) n + √ nκϵ − 2 − αα f (cid:1) iterations.
The expected total computational cost (sample complexity) is (cid:101) O (cid:0) n +√ nκϵ − 2 − αα f (cid:1) .
Theorem 3 quantifies the improvement of PAGER over GD in the finite sum setting in terms of nand over SGD in terms of ϵ f , see Table 1 for comparison.
Recall that GD has sample complexity
O (cid:0) nκϵ − (2 − α ) α f (cid:1) .
When n is large, we get the improvement of order √ n .
Notice that in the limitα → 2 , it matches the best known result for 2 -PŁ objectives [34].
Checklist
For all authors...
(a) Do the main claims made in the abstract and introduction accurately reflect the paper’s contributions and scope?
[Yes] We accurately state the main contributions inintroduction.(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 tothem?
[Yes]2.
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 thetechnical proofs in Appendix.3.
If you ran experiments...
(a) Did you include the code, data, and instructions needed to reproduce the main experimental 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 theywere chosen)?
[N/A](c)
Did you report error bars (e.g., with respect to the random seed after running experiments 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 identifiable
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 amount spent on participant compensation?
[N/A]