Sharp Analysis of Stochastic Optimization under Global Kurdyka-Łojasiewicz Inequality
Abstract
We study the complexity of finding the global solution to stochastic nonconvex optimization when the objective function satisfies global Kurdyka-Łojasiewicz (KŁ) 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 under the 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 improved sample 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 reinforcement learning.
Introduction
Nonconvex optimization problems are ubiquitous in machine learning domains such as training deep neural networks [22] or policy optimization in reinforcement learning [52].
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 an important avenue of research in recent years [21, 4, 30, 44, 23, 15, 59].
We are interested in solving the unconstrained stochastic , nonconvex optimization problem of the form
[Equation]
[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 written as the average of n smooth functions:
[Equation]
For a general nonconvex differentiable objective f : R d → R , finding a global minimum of f is in general intractable [42, 54].
There are two common strategies to analyze optimization methods for nonconvex functions.
The first one is to scale down the requirements on the solution of interest from global optimality to some relaxed version, e.g., first-order stationary point.
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 structural assumption in the hope of convergence 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 secant inequality, and quadratic growth [26, 24, 23].
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 [9, 10].
Of particular interest is the family of functions that satisfy global KŁ inequality.
Specifically, we say that f ( · ) satisfies (global) KŁ inequality if there exists 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 is satisfied for f ( · ) .
The special case of KŁ condition, 2 -PŁ, often referred as Polyak-Łojasiewicz or PŁ condition, was originally discovered independently in the seminal works of B. Polyak, T.
Ležanski and S. Łojasiewicz [48, 33, 38, 39].
Notably, this class of problems has found many interesting emerging applications in machine learning, for instance, policy gradient (PG) methods in reinforcement learning [41, 1, 56], generalized linear models [40], over-parameterized neural networks [2, 57], linear quadratic regulator in optimal control [13, 19], and low-rank matrix recovery [7].
Despite increased popularity of KŁ and α -PŁ assumptions, the analysis of stochastic optimization under it remains limited and the majority of works focus on deterministic gradient methods.
Indeed until recently, only the special case of α -PŁ with α = 2 was mainly addressed in the literature [26, 23, 30, 53, 49].
In this paper, we study the sample complexities of stochastic optimization for the 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, a point x ∈ R d for which E [ ||∇ f
(ˆ x ) || ] ≤ ϵ .
For instance, [21] showed that for a smooth objective (one with Lipschitz gradients) under bounded variance (BV) assumption, SGD with properly chosen stepsizes reaches an ϵ -stationary point with the sample complexity of O ( ϵ − 4 ) .
Recently, [30, 56] further extended the result under a much milder expected smoothness (ES) assumption on stochastic gradient.
While this sample complexity is known to be optimal for general nonconvex functions, a naive 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, [20] studied SGD and established convergence rates for α -PŁ functions under BV assumption.
Their sample complexity result is O ( ϵ − (4 − α ) / α f ) in our notation.
Later [35] considered SGD scheme with random 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 [20] and [35] 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 under global KŁ condition and a more practical model of stochastic oracle?
Variance reduction.
There has been extensive research on development of algorithms which improve 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 3 is variance reduction , which has emerged from the seminal works of Blatt et.
al [8].
The main idea of variance reduction is to make use of the stochastic gradients computed at previous iterations to construct a better gradient estimator at a relatively small computational cost.
Various modifications, generalizations, and improvements of the variance reduction technique appeared in subsequent work, for instance, [50, 28, 16] to name a few.
Finite-sum case.
A number of recently proposed algorithms such as SNVRG [58], SARAH [45],
STORM [15], SPIDER [17], and PAGE [36] achieve the sample complexity O (cid:0) n + √ n ϵ 2 (cid:1) when minimizing a general nonconvex function with finite sum structure (2).
This result is also known to be optimal in this setting [36].
[27] studies SARAH in finite sum case under local KŁ assumption and proves convergence in the iterates.
The study in [27] is only asymptotic analysis and the {+dependence 
 Another independent direction is to make use of higher order information [43, 18, 3]. on+} the parameters κ and n , which are important in practice for quantifying the improvement over GD and SGD are ignored.
[34] proposes an SVRG -based algorithm for KŁ functions and [55, 45, 46] 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 4 for finding an ϵ f -optimal solution.
However, it is not clear if it is possible to provide any non-asymptotic guaranties for variance reduced methods under α -PŁ condition for any α ∈ [1 , 2) .
In our work, we will answer the following open question: What is the extent of improvement any variance reduction scheme 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) [32, 37] 5 .
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 minimizing the function value, E [ f ( x ) − f ⋆ ] ≤ ϵ f .
For example, the methods in [58, 45, 17, 36] achieve O ( ϵ − 3 ) complexity improving over O ( ϵ − 4 ) complexity of SGD for finding an ϵ -stationary point.
Under the 2 -PŁ condition, these results can be extended to global convergence with O ( ϵ − 1 f ) sample complexity [36].
However, in contrast to a general nonconvex case, variance reduction under 2 -PŁ assumption does not show any improvement over SGD in terms of ϵ f .
We highlight that all existing analysis of variance reduction under α -PŁ condition is restricted only to a special case α = 2 .
We refer the 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 global convergence in this setting, one could naively translate the complexity for finding a stationary point of 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 any improvement over SGD in the most interesting practical case α = 1 and gives strictly worse result for all α > 1 .
Second, the guarantees for general nonconvex optimization hold on average, in the sense that the point ˆ x is sampled uniformly from all the iterates of the algorithm.
It would be more desirable to instead derive last iterate convergence guarantees under KŁ ( α -PŁ) condition.
In this work, we will address the following open question: Is it possible to accelerate the O (cid:0) ϵ − (4 − α ) / [Equation] f
[Equation]
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Ł condition and answer all the above questions.
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 [42].
Furthermore, we obtain faster rates with PAGER in finite sum case (2), providing the first acceleration over GD and SGD under α -PŁ condition.
In Table
[Table]
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 the set X ⋆ of optimal points and provide a summary in Table 2 in the Appendix.
[Table]
Assumptions and Discussion
In this section [Table] 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 continuous function 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 a disingularizing function and ρ ′ ( · ) denotes its derivative.
Moreover, disingularizing function satisfies the following conditions, it is continuous, concave, ρ (0) = 0 , and ρ ′ ( t ) > 0 [35].
If the above assumption holds for ρ ( t ) :
= 1 θ t θ and θ > 0 , then Assumption 3 is satisfied with PŁ power α = 11 − θ .
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 is convex, but it does not satisfy inequality ρ ′
( f ( x ) − f ⋆ ) · ∥∇ f ( x ) ∥ ≥ 1 for any choice of ρ ( t ) .
We also 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 KL condition with ϕ ( t ) =
min { t, √ t } , [51].
The intuition behind the special case α = 1 is that the function is allowed to be flat near the set of optimal 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 differentiable function with h ′ ( t ) ≥ 0 , h (0) = 0 .
The expectation is taken over random vector ξ ∼ D .
We call b k the cost of such estimator.
This assumption encompasses previous assumptions in the literature.
For instance, it is straightforward to see that an estimator satisfies the standard bounded variance assumption [21] when h ( t ) = 0 , B = 1 , and b k = 1 in (5).
Gradient estimators with relaxed growth assumption [6, 11] are also special cases of (5) for h ( t ) =
and b k = 1 .
A closely related assumption to the relaxed growth was introduced in [53] which holds when h ( t ) = 0 and C = 0 in (5).
Expected smoothness assumption [30, 24, 56] 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 and compressed communication [30].
More recently, it has been shown that PG method with softmax policies and log barrier regularization can be modeled using ES assumption [56].
Note that due to the first term in (5), i.e., 2 A · h ( f ( x ) − f ∗ ) , the second moment can be large when the objective gap at x is large.
Such property is not captured by standard bounded variance (Assumption 5).
The flexibility and advantages of introducing such additional term are elaborated in detail in the literature [24, 30, 25, 56].
We highlight two special cases for the sequence b k : b k = Θ( k τ ) with τ ≥ 0 and b k = Θ( q k ) with q > 1 .
For example, Monte Carlo sampling and mini-batching allow us to design estimators with such { b k } k ≥ 0 sequences.
When a gradient estimator satisfies Assumption 4, unless b k is bounded for all k , it essentially means that we have a mechanism to reduce its variance.
More precisely, the variance 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 by sequence b k .
On the other hand, usually the access to such estimator comes with a cost proportional to b k , e.g., mini-batch setting described in Section 4.
Thus we refer to b k as the cost of the estimator g k ( x, ξ ) .
Stochastic Gradient Method
Algorithm 1 summarizes the steps of a slightly modified SGD which we analyze in this work.
We call this algorithm SGD with restarts.
This algorithm updates the point x for T number of iterations within an inner-loop.
Note that, in the inner-loop, the step-size remains unchanged and the iterates are updated via x t +1 = x
t − ηg k
( x t , ξ t ) , where η is the step-size and { ξ t } t ≥ 0 are independent random vectors.
The cost b k of the gradient estimator g k ( x, ξ ) remain the same within the inner loop of
Algorithm 1.
3.1 Dynamics of SGD
Let { x t } t ≥ 0 be the sequence of points generated by the inner loop of Algorithm 1, and Assumptions 1, 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]
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 7 point of (6) when the inequality is replaced by equality and for a fixed step-size, η k = η , which we denote by r ( η ) .
ii) Selecting the step-sizes { η k } and sequence { b k } k ≥ 0 , such that the corresponding stationary points { r ( η k ) } k ≥(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 polynomial function 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 set of conditions that allow us to design the stepsizes, which will guarantee convergence.
The detailed derivations 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]
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, for any ϵ 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
[Equation] η k [Equation]
Θ
[Equation]
k [Equation] [Equation] τ ≤
[Equation]
N = O (cid:0) ϵ [Equation]
[Equation] f (cid:1) with η k = Θ( k − 2 − [Equation] [Equation]
α
[Equation]
) if τ >
[Equation] To verify the above result empirically, we simulated δ t in (6) throughout all iterations of Algorithmsfor different sets of parameters and presented the results in Figure 1 along with their corresponding convergence rates given in Corollary 1.
As it is shown in these figure, the above convergence rates correctly capture the behaviour of the dynamics in (6).
As an example, in Figure 1(b), the red solid curve 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 equivalently log( δ k ) =
( − α [Equation] +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.
Comparison to related works.
Authors in [30] studied the convergence of SGD for 2 -PŁ objectives, under a stronger assumption than Assumption 4.
More precisely, they assumed an estimator that satisfies Assumption 4 with τ = [-
-]0 and h ( t ) =
t and obtained the convergence rate of O ( ϵ − 1 f ) .
This 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 . is consistent with our rate presented in the Corollary 1.
It is worth noting that in this setting, O ( ϵ − 1 f ) is optimal [30].
The authors in [56] studied the performance of SGD for 1 -PŁ objectives.
Assuming that the gradient estimator 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 in this case, the cost of each iteration is b k = 1 , which means that the iteration complexity coincides with the sample complexity.
We note that our proof technique is different than in [56] and allows to consider more general assumptions.
Finally, [20] studied SGD with α -PŁ objectives for α ∈ [1 , 2] under bounded variance Assumption 5 with b k = 1 and obtained similar convergence rate to ours.
We recover their result as a special case by setting A = 0 and τ = 0 in Corollary 1, if we set T = 1 we also recover the same (up to a constant) step-sizes η k = Θ (cid:16) k − 2 4 − α (cid:17) .
However, we highlight that our proof technique is different from [20], and more generic since it holds for a general Assumption 4.
3.2 Sample complexity of SGD
The result of Corollary 1 suggests that by increasing the cost of the gradient estimator b k over the iterations, one can achieve a better iteration complexity of Algortihm 1.
In particular, it improves with τ until it reaches the minimum N = O (cid:16) ϵ
[Equation]
(cid:17) at τ = γ [Equation] γ and does not change for larger values of τ .
However, we are merely interested in the iteration complexity in practice, since the computational cost at each iteration can be prohibitively large.
A more adequate measure is the total computational cost (sample complexity) of the method.
It is interesting whether increasing b k over 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 expected total computational cost (sample complexity) of Algorithm 1 is
[Equation]
The above result implies that increasing the cost of the gradient estimator with iterations does not improve 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 for the dynamic (6) describing the progress of SGD .
More precisely, if there exists a function f and a gradient estimator satisfying the assumptions in Corollary 1 such that its corresponding recursive inequality (6) is an equality, then its convergence rate, presented in Corollary 1 cannot be improved by any choices of stepsizes { η k } k ≥ 0 .
Next proposition summarizes our results about the tightness of our convergence rates in Corollary 1.
Proposition 2.
Consider the following recursion
[Equation]
for all k
[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 ) .
Faster Rates with Variance Reduction
Algorithm
) 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
To simplify the exposition of the results in this section, let us assume that g k ( x t , ξ t ) is constructed explicitly 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 a random vector of independent entries, ξ t are independent for all iterations, {∇ f ξ it ( x t ) }
b k i
=1 are queries provided by an oracle such that E [ ∇ f ξ it ( x t )]
= ∇ f ( x t ) and E
[ ||∇ f ξ it ( x t ) − ∇ f ( x t ) || 2 ] ≤ σ 2 for all 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 the following
Assumption 6 (Average L -smoothness (of order k )) .
Let g ′ k
( x, ξ ) :
= 1 b ′ k (cid:80) b ′ k
i =1 ∇ f ξ i ( x ) and g ′ 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 x and 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 there exists L ≥ {+
+} 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 [36] 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 several difficulties.
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 sequence of 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 of parameters in the outer loop k = 0 , . .
, K − 1 and applies PAGE estimator with a fixed set of parameters 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 α in settings (1) and (2).
4.2 Online case
We present convergence guarantees for Algorithm 2 in the setting (1) and defer its formal proof to
Appendix C.
Theorem 2.
Let f ( · ) have the form (1) and satisfy Assumptions 1, 3 (with α ∈ [1 , 2) ), 5 and 6, let the sequences 9 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 point x with E [ f ( x ) − f ⋆ ] ≤ ϵ f after N := (cid:80) K − 1 k =
T k
= O (cid:0) κϵ − 2 − αα f (cid:1) iterations, where κ = L / µ .
The expected total computational cost (sample complexity) is O (cid:16)(cid:16) σ 2 µ + κ 2 (cid:17) ϵ −
α f (cid:17) .
Improvement over SGD .
Theorem 2 implies that PAGER improves the sample complexity of SGD from 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), this leads 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 analysis we obtain the optimal sample complexity for convex stochastic optimization under the additional assumption 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 ) .
10 For 1 -PŁ objectives, PAGER has O (cid:16)
ϵ − 2 f (cid:17) sample complexity.
Since the 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 [42].
Moreover, we highlight that this result holds for the last iterate of PAGER , while the standard analysis of first order methods for convex functions guarantees convergence for the average iterate [31].
The last iterate convergence for convex objectives was only recently established for SGD by following an involved analysis with a careful control of iterates via suffix-averaging scheme [20].
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, let the sequences be chosen as p k =
[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 =
T k = (cid:101) O (cid:0) n + √ nκϵ − 2 − αα f (cid:1) iterations, where κ = L / µ
The expected total computational cost (sample complexity) is (cid:101) O (cid:0) n + √ nκϵ − 2 − αα f (cid:1) .
The proof is deferred to Appendix C. Theorem 3 quantifies the improvement of PAGER over GD in the finite sum setting in terms of n and 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 [36].
Conclusion
We analyzed the complexity of SGD when the objective satisfies global KŁ inequality and the queries from stochastic gradient oracle satisfy weak expected smoothness.
We introduced a general framework for this analysis which resulted in a sample complexity of O ( ϵ − (4 − α ) /α ) for SGD with objectives satisfying α -PŁ condition.
We also demonstrated the tightness of this rate under the specific choice of stepsizes.
Last but not least, we developed a modified SGD with variance reduction and restarting ( PAGER ), which improves the sample complexity of SGD for the whole spectrum of parameters α ∈ [1 , 2) and achieves the optimal rate for the important case of 1 -PŁ objectives.
Acknowledgements
We would like to thank Anas Barakat and Anastasia Kireeva for valuable discussions.
The work of
I. Fatkhullin was supported by ETH AI Center doctoral fellowship.