ERAGED SOLUTIONS
Year: 2019

Abstract:
Although stochastic gradient descent (SGD) method and its variants (e.g., stochastic momentum methods, ADAGRAD) are algorithms of choice for solving nonconvex problems (especially deep learning), big gaps still remain between the theory and the practice with many questions unresolved. For example, there is still a lack of theories of convergence for SGD and its variants that use stagewise step size and return an averaged solution in practice. In addition, theoretical insights of why adaptive step size of ADAGRAD could improve non-adaptive step size of SGD is still missing for non-convex optimization. This paper aims to address these questions and fill the gap between theory and practice. We propose a universal stagewise optimization framework for a broad family of non-smooth nonconvex problems with the following key features: (i) at each stage any suitable stochastic convex optimization algorithms (e.g., SGD or ADAGRAD) that return an averaged solution can be employed for minimizing a regularized convex problem; (ii) the step size is decreased in a stagewise manner; (iii) an averaged solution is returned as the final solution. Our theoretical results of stagewise ADAGRAD exhibit its adaptive convergence, therefore shed insights on its faster convergence than stagewise SGD for problems with slowly growing cumulative stochastic gradients. To the best of our knowledge, this is the first work for addressing the unresolved issues of existing theories mentioned earlier. Besides theoretical contributions, our empirical studies show that our stagewise SGD and ADAGRAD improve the generalization performance of existing variants/implementations of SGD and ADAGRAD.

heading: 1 INTRODUCTION
text: Non-convex optimization has recently received increasing attention due to its popularity in emerging machine learning tasks, particularly for learning deep neural networks. One of the keys to the success of deep learning for big data problems is the employment of simple stochastic algorithms such as SGD or ADAGRAD (Krizhevsky et al., 2012; Dean et al., 2012). Analysis of these stochastic algorithms for non-convex optimization is an important and interesting research topic, which already attracts much attention from the community of theoreticians (Ghadimi & Lan, 2013; 2016; Yang et al., 2016; Davis & Drusvyatskiy, 2018b; Ward et al., 2018; Li & Orabona, 2018). However, one issue that has been largely ignored in existing theoretical analysis is that the employed algorithms in practice usually differ from their plain versions that are well understood in theory. Below, we will discuss several important heuristics used in practice for training deep neural networks and the gap between the practice and the theory to motivate this work.
First, a trick for setting the step size in training deep neural networks is to change it in a stagewise manner from a large value to a small value, i.e., a constant step size is used in a stage for a number of iterations and is decreased for the next stage (Krizhevsky et al., 2012; Ren et al., 2018). Although
∗Equal contribution. Correspondence to: tianbao-yang@uiowa.edu
this trick has been adopted by most open-sourced libraries, e.g., Caffe (Jia et al., 2014), TensorFlow (Abadi et al., 2015), Pytorch (Paszke et al., 2017), it still lacks theoretical analysis to date for non-convex optimization. A related question that stands out is how and when to decrease the step size. On standard benchmark datasets for academic use such as CIFAR-10, CIFAR-100 (Krizhevsky et al.), people could follow the setting reported in previous studies to get a good result, which however might not work well for new datasets. Hence, a better solution to setting the stagewise step size with insights from theory would be much preferred. However, in the existing literature of theory for non-convex optimization (Ghadimi & Lan, 2013; Davis & Drusvyatskiy, 2018b), only strategies based on an iteratively decreasing step size or a small constant step size have been well analyzed. For example, the existing theory usually suggests an iteratively decreasing step size proportional to 1/ √ t at the t-th iteration or a small constant step size, e.g., proportional to 2 with 1 for finding an -stationary solution whose gradient’s magnitude (in expectation) is small than .
Second, the averaging heuristic is usually used in practice, i.e., an averaged solution is returned for prediction (Bottou, 2010), which could yield improved stability and generalization (Hardt et al., 2016). However, existing theory for many stochastic non-convex optimization algorithms only provides guarantee on a uniformly sampled solution or a non-uniformly sampled solution with decreasing probabilities for latest solutions (Ghadimi & Lan, 2013; Yang et al., 2016; Davis & Drusvyatskiy, 2018b). In particular, if an iteratively decreasing step size proportional to 1/ √ t at the t-th iteration is employed, the convergence guarantee was provided for a random solution that is non-uniformly selected from all iterates with a sampling probability proportional to 1/ √ t for the t-th iterate. This means that the latest solution always has the smallest probability to be selected as the final solution, which contradicts to the common wisdom. If a small constant step size is used, then usually a uniformly sampled solution is provided with convergence guarantee. However, both options are rarely used in practice and cannot justify the heuristic that returns the last solution.
A third common approach in practice is to use adaptive coordinate-wise step size of ADAGRAD (Dean et al., 2012). Although adaptive step size has been well analyzed for convex problems (e.g., why and when it could yield faster convergence than SGD) (Duchi et al., 2011; Chen et al., 2018b), it still remains an mystery for non-convex optimization with missing insights from theory. Several recent studies have attempted to analyze ADAGRAD for non-convex problems (Ward et al., 2018; Li & Orabona, 2018; Chen et al., 2018a; Zou & Shen, 2018). Nonetheless, none of them are able to exhibit the adaptive convergence of ADAGRAD to data as in the convex case and its potential advantage over SGD for non-convex problems.
To overcome the shortcomings of existing theories for stochastic non-convex optimization, this paper analyzes new algorithms that employ some or all of these commonly used heuristics in a systematic framework, aiming to fill the gap between theory and practice. The main results and contributions are summarized below:
• We propose a universal stagewise optimization framework for solving a family of non-convex problems, i.e., weakly convex problems, which is broader than smooth non-convex problems and includes some non-smooth non-convex problems. At each stage, any suitable stochastic convex optimization algorithms (e.g., SGD, ADAGRAD) with a constant step size parameter can be employed for optimizing a regularized convex problem with a number of iterations, which usually return an averaged solution. The step size parameter is decreased in a stagewise manner following a polynomial decaying scheme.
• We analyze several variants of the proposed framework by employing different basic algorithms, including SGD, ADAGRAD, stochastic heavy-ball (SHB) method, and stochastic Nesterov’s accelerated gradient (SNAG) method. We prove the convergence of their stagewise versions for an averaged solution that is randomly selected from all stagewise averaged solutions.
• To justify a heuristic approach that returns the last averaged solution in stagewise learning, we present and analyze a non-uniform sampling strategy over stagewise averaged solutions with sampling probabilities increasing as the stage number.
• Regarding the convergence results, for stagewise SGD, SHB, SNAG, we establish the same order of iteration complexity for finding a nearly stationary point as the existing theories of their nonstagewise variants. For stagewise ADAGRAD, we establish an adaptive convergence for finding a nearly stationary point, which is provably better than (stagewise) SGD, SHB, and SNAG when the cumulative growth of stochastic gradient is slow.
• Besides theoretical contributions, we also empirically verify the effectiveness of the proposed stagewise algorithms. In particular, our empirical studies show that (i) the stagewise ADAGRAD dramatically improves the generalization performance of existing variants of ADAGRAD, (ii) stagewise SGD, SHB, SNAG also outperform their plain variants with an iteratively decreasing step size; (iii) the proposed stagewise algorithms achieve similar if not better generalization performance than their heuristic variants implemented in existing libraries on standard benchmark datasets.

heading: 2 RELATED WORK
text: SGD for unconstrained smooth non-convex problems was first analyzed by Ghadimi & Lan (2013), who established an O(1/ 4) iteration complexity for finding an -stationary point x in expectation satisfying E[‖∇f(x)‖] ≤ , where f(·) denotes the objective function. As mentioned earlier, the returned solution is either a uniformly sampled solution or a non-uniformly sampled one with sampling probabilities proportional to the decreasing step size. Similar results were established for the stochastic momentum variants of SGD (i.e., SHB, SNAG) by Yang et al. (2016); Ghadimi & Lan (2016). Recently, SGD was also analyzed for (constrained) weakly convex problems, whose objective function is non-convex and not necessarily smooth, by Davis & Drusvyatskiy (2018b). However, none of these studies provide results for algorithms that return an averaged solution, and these analyzed algorithms differ significantly from that used in practice for achiving the state-of-the-art results (Krizhevsky et al., 2012; Ren et al., 2018; Loshchilov & Hutter, 2017).
Although adaptive variants of SGD, e.g., ADAGRAD (Duchi et al., 2011), ADAM (Kingma & Ba, 2015; Reddi et al., 2018), were widely used for training deep neural networks, there are few studies on theoretical analysis of these algorithms for non-convex problems. Several recent studies attempted to analyze ADAGRAD for non-convex problems (Ward et al., 2018; Li & Orabona, 2018; Chen et al., 2018a; Zou & Shen, 2018). Although these studies have established an iteration complexity of O(1/ 4) for different variants of ADAGRAD for finding an -stationary solution of a stochastic non-convex optimization problem, none of them can exhibit the potential adaptive advantage of ADAGRAD over SGD as in the convex case. Besides that, these studies also suffer from the following shortcomings: (i) they all assume smoothness of the objective function, while we consider non-smooth and non-convex problems; (ii) their convergence is provided on a solution with minimum magnitude of gradient that is expensive to compute, though their results also imply a convergence on a random solution selected from all iterates with decreasing sampling probabilities. In contrast, these shortcomings do not exist in this paper. To the best of our knowledge, our result is the first one that explicitly shows that coordinate-wise adaptive step size could yield faster convergence than using non-adaptive step size for non-smooth non-convex problems, which is similar to that in the convex case and was observed in practice for deep learning (Dean et al., 2012).
The proposed stagewise algorithm is similar to several existing algorithms in design (Xu et al., 2017; Davis & Grimmer, 2017), which are originated from the proximal point algorithm (Rockafellar, 1976). I.e., at each stage a regularized convex subproblem is formed and then a stochastic algorithm is employed for optimizing the regularized subproblem inexactly with a number of iterations. Xu et al. (2017) used this idea for solving problems that satisfy a local error bound condition, aiming to achieve faster convergence than vanilla SGD. Davis & Grimmer (2017) are probably the first who analyzed this idea for solving non-smooth weakly convex problems. In these two papers SGD with decreasing step sizes for a strongly convex problem is employed at each stage for solving the regularized subproblem. Our stagewise algorithm is developed following the similar idea. The key differences from (Xu et al., 2017; Davis & Grimmer, 2017) are that (i) we focus on non-convex problems instead of convex problems considered in (Xu et al., 2017); (ii) we analyze a non-uniform sampling strategy with sampling probabilities increasing as the stage number to justify a heuristic approach that uses the last averaged solution for prediction, unlike the uniform sampling used in (Davis & Grimmer, 2017); (iii) we present a unified algorithmic framework and convergence analysis, which enable one to employ any suitable stochastic convex optimization algorithms at each stage. Importantly, it brings us several interesting variants including stagewise stochastic momentum methods and stagewise ADAGRAD. Finally, we note that a similar idea that adds a strongly convex term into the non-convex objective has been considered in several recent works (Allen-Zhu, 2017; Lan & Yang, 2018; Carmon et al., 2016). However, the key difference between these works and the present work is that they need to assume the objective function is smooth or has a smooth component. In addition, Allen-Zhu (2017); Lan & Yang (2018) consider finite-sum problems, and Carmon et al.
(2016) consider deterministic problems. In contrast, we consider more general stochastic problems without assuming the objective function is smooth.

heading: 3 PRELIMINARIES
text: The problem of interest in this paper is:
min x∈Ω φ(x) = Eξ[φ(x; ξ)], (1)
where Ω ⊆ Rd is a closed convex set, ξ ∈ U is a random variable, φ(x) and φ(x; ξ) are non-convex functions, with the basic assumptions on the problem given in Assumption 1.
To state the convergence property of an algorithm for solving the above problem. We need to introduce some definitions. These definitions can be also found in related literature, e.g., Davis & Grimmer (2017); Davis & Drusvyatskiy (2018b). In the sequel, we let ‖ · ‖ denote an Euclidean norm, [S] = {1, . . . , S} denote a set, and δΩ(·) denote the indicator function of the set Ω. Definition 1. (Fréchet subgradient) For a non-smooth and non-convex function f(·),
∂F f(x) = {v ∈ Rd|f(y) ≥ f(x) + v>(y − x) + o(‖y − x‖), ∀y ∈ Rd}
denotes the Fréchet subgradient of f .
Definition 2. (First-order stationarity) For problem (1), a point x ∈ Ω is a first-order stationary point if 0 ∈ ∂F (φ + δΩ)(x), where δΩ denotes the indicator function of Ω. Moreover, a point x is said to be -stationary if dist(0, ∂F (φ + δΩ)(x)) ≤ , where dist denotes the Euclidean distance from a point to a set.
Definition 3. (Moreau Envelope and Proximal Mapping) For any function f and λ > 0, the following function is called a Moreau envelope of f
fλ(x) = min z f(z) +
1
2λ ‖z− x‖2. (2)
Further, the optimal solution to the above problem denoted by
proxλf (x) = arg minz f(z) +
1
2λ ‖z− x‖2
is called a proximal mapping of f .
Definition 4. (Weakly convex) A function f is µ-weakly convex (µ > 0), if f(x) + µ2 ‖x‖ 2 is convex. It is known that if f(x) is µ-weakly convex and γ < µ−1, then its Moreau envelope fγ(x) is C1-smooth with the gradient given by (see e.g., Davis & Drusvyatskiy (2018b))
∇fγ(x) = γ−1(x− proxγf (x)).
The tool of Moreau envelope is introduced to measure the convergence for optimizing non-smooth and non-convex functions. A small norm of ∇fγ(x) has an interpretation that x is close to a point that is nearly stationary. In particular for any x ∈ Rd, let x̂ = proxγf (x), then we have (Davis & Drusvyatskiy, 2018b)
f(x̂) ≤ f(x), ‖x− x̂‖ = γ‖∇fγ(x)‖, dist(0, ∂f(x̂)) ≤ ‖∇fγ(x)‖, (3)
where the first inequality follows from the definition of proxλf (x) such that f(x̂) + 1 2λ‖x̂− x‖ 2 ≤ f(x) (Please see Appendix H for a detailed proof of (3)). This means that a point x satisfying ‖∇fγ(x)‖ ≤ is close to a point in distance of O( ) that is -stationary. It is notable that for a non-smooth function f(·), there could exist a sequence of solutions {xk} such that ∇fγ(xk) converges while dist(0, ∂f(xk)) may not converge (Drusvyatskiy & Paquette, 2018). A simple example is to consider minx∈R |x|. As long as x 6= 0, dist(0, ∂f(x)) = 1 6= 0 no matter how close is x to the stationary point 0. To handle such a challenging issue for non-smooth non-convex problems, we will follow existing works (Davis & Drusvyatskiy, 2018a; Drusvyatskiy & Paquette, 2018; Davis & Grimmer, 2017) to prove the near stationarity in terms of ∇fγ(x). In the case when f is smooth, ‖∇fγ(x)‖ is closely related to the magnitude of the gradient. In
particular, let us define the projected gradient below, which is used as a criterion for non-convex smooth optimization in the presence of constraints (Ghadimi & Lan, 2016; Allen-Zhu, 2017):
Gγ(x) = 1
γ (x− proxγδΩ(x− γ∇f(x))). (4)
Note that the projected gradient becomes the normal gradient when the constraint x ∈ Ω is absent. It was shown that when f(·) is smooth with L-Lipschitz continuous gradient (Drusvyatskiy & Lewis, 2016):
(1− Lγ)‖Gγ(x)‖ ≤ ‖∇fγ(x)‖ ≤ (1 + Lγ)‖Gγ(x)‖,∀x ∈ Ω. (5)
Thus, the near stationarity in terms of ∇fγ(x) implies the near stationarity in terms of Gγ(x) for a smooth function f(·) for a properly chosen γ > 0 (e.g., γ = 1/(2L)). In this work, we define φγ(x) as the Moreau envelope of φ(x) + δΩ(x) as in (2) with f(x) replaced by φ(x) + δΩ(x) and study the convergence in terms of of∇φγ(x). Now, we are ready to state the basic assumptions of the considered problem (1). Assumption 1. (i) There is a measurable mapping g : Ω×U → R such that Eξ[g(x; ξ)] ∈ ∂Fφ(x) for any x ∈ Ω; (ii) for any x ∈ Ω, E[‖g(x; ξ)‖2] ≤ G2; (iii) there exists ∆φ > 0 such that φ(x)−minz∈Ω φ(z) ≤ ∆φ for any x ∈ Ω; (iv) the objective function φ is µ-weakly convex; Remark: Assumption 1(i), (ii) assume a stochastic subgradient is available for the objective function and its Euclidean norm square is bounded in expectation, which are standard assumptions for nonsmooth optimization. Assumption 1(iii) assumes that the objective value with respect to the optimal value is bounded. Assumption 1(iv) assumes weak convexity of the objective function, which is weaker than assuming smoothness. Below, we present some examples of objective functions in machine learning that are weakly convex.
Ex. 1: Smooth Non-Convex Functions. If φ(·) is a L-smooth function (i.e., its gradient is LLipschitz continuous), then it is L-weakly convex. This will include the objective function for neural networks with a smooth activation function (e.g., the sigmoid function) and a smooth loss function (e.g., softmax loss).
Ex. 2: Convex and Smooth Composition. Consider φ(x; ξ) = h(c(x; ξ)) where h(·) : Rm → R is closed convex and M -Lipschitz continuous, and c(x; ξ) : Rd → Rm is nonlinear smooth mapping with L-Lipschitz continuous gradient. This class of functions has been considered in (Drusvyatskiy & Paquette, 2018) and it was proved that φ(x; ξ) is ML-weakly convex. An interesting example is phase retrieval, where φ(x; a, b) = |(x>a)2 − b|. Another example related to deep learning is that if c(x; ξ) denotes the mapping function of a neural network parameterized by x with a smooth activation function and h denotes a non-smooth Lipschitz continuous convex loss function (e.g., hinge loss), then the resulting loss h(c(x; ξ)) is a weakly convex function of x. More examples of this class can be found in (Davis & Drusvyatskiy, 2018a).

heading: 4 STAGEWISE OPTIMIZATION: ALGORITHMS AND CONVERGENCE
text: In this section, we will present the proposed stagewise algorithms and their convergence results. We will first present a Meta algorithmic framework highlighting the key features of the proposed algorithms and then present several variants of the Meta algorithm.
The Meta algorithmic framework is described in Algorithm 1. There are several key features that differentiate Algorithm 1 from existing stochastic algorithms that come with theoretical guarantee. First, the algorithm is run with multiple stages. At each stage, a basic stochastic algorithm (SA) is called to optimize a regularized problem fs(x) inexactly that consists of the original objective function and a quadratic term, which is guaranteed to be convex due to the weak convexity of φ and γ < µ−1. The convexity of fs allows one to employ any suitable existing stochastic algorithms (cf. Theorem 1) that have convergence guarantee for convex problems. The returned solution from the (s − 1)-th stage is used as a reference point for constructing fs and as an initial solution for warm-start. It is notable that SA usually returns an averaged solution xs at each stage. Second, a decreasing sequence of step size parameters ηs is used. At each stage, the SA uses a constant step size parameter ηs and runs the updates for a number of Ts iterations. We do not initialize Ts as it might be adaptive to the data as in stagewise ADAGRAD. The setup of ηs and Ts will depend on the specific choice of SA, which will be exhibited later for different variants.
Algorithm 1 A Meta Stagewise Algorithm: Stagewise-SA 1: Initialize: a sequence of decreasing step size parameters {ηs}, x0 ∈ Ω, γ < µ−1 2: for s = 1, . . . , S do 3: Let fs(·) = φ(·) + 12γ ‖ · −xs−1‖ 2
4: xs = SA(fs,xs−1, ηs, Ts) xs is usually an averaged solution 5: end for
To illustrate that Algorithm 1 is a universal framework such that any suitable SA algorithm can be employed, we present the following result by assuming that SA has an appropriate convergence for a convex problem.
Theorem 1. Let f(·) be a convex function, x∗ = arg minx∈Ω f(x) and Θ denote some problem dependent constants. Suppose for x+ = SA(f,x0, η, T ), we have
E[f(x+)− f(x∗)] ≤ ε1(η, T,Θ)‖x0 − x∗‖22 + ε2(η, T,Θ)(f(x0)− f(x∗)) + ε3(η, T,Θ). (6)
Under Assumption 1(i), (iii) and (iv), by running Algorithm 1 with γ = 1/(2µ), and with ηs, Ts satisfying ε1(ηs, Ts,Θ) ≤ 1/(48γ), ε2(ηs, Ts,Θ) ≤ 1/2, ε3(ηs, Ts,Θ) ≤ c3/s for some c3 > 0, we have
E [ ‖∇φγ(xτ )‖2 ] ≤ 32∆φ(α+ 1)
γ(S + 1) +
48c3(α+ 1)
γ(S + 1) ,
where τ is randomly selected from {0, . . . , S} with probabilities pτ ∝ (τ + 1)α, α ≥ 1. Remark: It is notable that the convergence guarantee is provided on a stagewise average solution xτ . To justify a heuristic approach that returns the final average solution for prediction, we analyze a new sampling strategy that samples a solution among all stagewise average solutions with sampling probabilities increasing as the stage number increases. This sampling strategy is better than uniform sampling strategy or a strategy with decreasing sampling probabilities in the existing literature. The convergence upper bound in (6) of SA covers the results of a broad family of stochastic convex optimization algorithms. When ε2(ηs, Ts,Θ) = 0 (as in SGD), the upper bound can be improved by a constant factor. Moreover, we do not optimize the value of γ. Indeed, any γ < 1/µ will work, which only has an effect on constant factor in the convergence upper bound.
Next, we present several variants of the Meta algorithm by employing SGD, ADAGRAD, and stochastic momentum methods as the basic SA algorithm, to which we refer as stagewise SGD, stagewise ADAGRAD, and stagewise stochastic momentum methods, respectively. It is worth mentioning that one can follow similar analysis to analyze other stagewise algorithms by using their basic convergence for stochastic convex optimization, including RMSProp (Mukkamala & Hein, 2017), AMSGrad (Reddi et al., 2018), stochastic alternating direction methods of multipliers (ADMM) (Ouyang et al., 2013; Suzuki, 2013).

heading: 4.1 STAGEWISE SGD
text: In this subsection, we analyze the convergence of stagewise SGD, in which SGD shown in Algorithm 3 in the Appendix is employed in the Meta framework. Besides Assumption 1, we impose the following bounded domain assumption in this subsection.
Assumption 2. There exists D > 0 such that ‖x− y‖ ≤ D for any x,y ∈ Ω. It is worth mentioning that bounded domain assumption is imposed for simplicity, which is usually assumed in convex optimization. For machine learning problems, one usually imposes some bounded norm constraint to achieve a regularization. Nevertheless, the bounded domain assumption is not essential for the proposed algorithm. We present a more involved analysis in subsection 4.3 for unbounded domain Ω = Rd. The following is a standard basic convergence result of SGD (Zinkevich, 2003), which clearly satisfies the bound in (6).
Lemma 1. For Algorithm 3, assume that f(·) is convex and E‖gt‖2 ≤ G2,∀t, then for any x ∈ Ω
E[f(x̂T )− f(x)] ≤ ‖x− x0‖2 2η(T + 1) + ηG2 2
The following theorem exhibits the convergence of stagewise SGD.
Algorithm 2 ADAGRAD(f,x0, η, ∗) 1: Initialize: x1 = x0, g1:0 = [], H0 ∈ Rd×d 2: while T does not satisfy the condition in Theorem 3 do 3: Compute a stochastic subgradient gt for f(xt) 4: Update g1:t = [g1:t−1,g(xt)], st,i = ‖g1:t,i‖2 5: Set Ht = H0 + diag(st) and ψt(x) = 12 (x− x1)
>Ht(x− x1) 6: Let xt+1 = arg min
x∈Ω ηx>
( 1 t ∑t τ=1 gτ ) + 1tψt(x)
7: end while 8: Output: x̂T = ∑T t=1 xt/T
Theorem 2. Suppose Assumption 1 and 2 hold. By setting γ = 1/(2µ), ηs = η0/s, Ts = 12γs/η0 where η0 > 0 is a free parameter, then stagewise SGD (Algorithm 1 employing SGD) guarantees that
E [ ‖∇φγ(xτ )‖2 ] ≤ 16µ∆φ(α+ 1)
S + 1 +
24µη0Ĝ 2(α+ 1)
(S + 1) ,
where Ĝ2 = 2G2 + 2γ−2D2, and τ is similarly defined as in Theorem 1. Remark: To find a solution with E [ ‖∇φγ(xτ )‖2 ] ≤ 2, we can set S = O(1/ 2) and the total
iteration complexity ∑S s=1 Ts is in the order of O(1/
4). The above theorem is essentially a corollary of Theorem 1 by applying Lemma 1 to fs(·) at each stage. We present a complete proof in the appendix.

heading: 4.2 STAGEWISE ADAGRAD
text: One of the main contributions of the present work is to develop a variant of ADAGRAD with adaptive convergence to data for stochastic non-convex optimization. In this subsection, we analyze stagewise ADAGRAD and establish its adaptive complexity. In particular, we consider the Meta algorithm that employs ADAGRAD in Algorithm 2 to optimize each fs. The key difference of stagewise ADAGRAD from stagewise SGD is that the number of iterations Ts at each stage is adaptive to the history of learning. It is this adaptiveness that makes the proposed stagewise ADAGRAD achieve adaptive convergence. It is worth noting that such adaptive scheme has been also considered in (Chen et al., 2018b) for solving stochastic strongly convex problems. In contrast, we consider stochastic weakly convex problems. Similar to previous analysis of ADAGRAD (Duchi et al., 2011; Chen et al., 2018b), we assume ‖g(x; ξ)‖∞ ≤ G,∀x ∈ Ω in this subsection. Note that this is stronger than Assumption 1 (ii). We formally state this assumption required in this subsection below. Assumption 3. ‖g(x; ξ)‖∞ ≤ G for any x ∈ Ω. The convergence analysis of stagewise ADAGRAD is build on the following lemma, which is attributed to Chen et al. (2018b) with a proof sketch provided in the Appendix. Lemma 2. Let f(x) be a convex function, H0 = GI with G ≥ maxt ‖gt‖∞, and iteration number T satisfy T ≥ M max{G+maxi ‖g1:T,i‖2c , c ∑d i=1 ‖g1:T,i‖} for some c > 0. Algorithm 2 returns an averaged solution x̂T such that
E[f(x̂T )− f(x∗)] ≤ c
Mη ‖x0 − x∗‖2 +
η
Mc , (7)
where x∗ = arg minx∈Ω f(x), g1:t = (g(x1), . . . ,g(xt)) and g1:t,i denotes the i-th row of g1:t.
The convergence property of stagewise ADAGRAD is described by following theorem. Theorem 3. Suppose Assumption 1(i), (iii), (iv), Assumption 2 and Assumption 3 hold. By setting γ = 1/(2µ), ηs = η0/ √ s, Ts ≥ Ms max{(Ĝ + maxi ‖gs1:Ts,i‖)/(2c), c ∑d i=1 ‖gs1:Ts,i‖} where η0, c > 0 are free parameters, and Msηs ≥ 24γc, then stagewise ADAGRAD guarantees that
E[‖∇φγ(xτ )‖2] ≤ 16µ∆φ(α+ 1)
S + 1 +
4µ2η20(α+ 1)
c2(S + 1) ,
where Ĝ = G + γ−1D, gs1:t,i denotes the cumulative stochastic gradient of the i-th coordinate at the s-th stage, and τ is similarly defined as in Theorem 1.
Remark: Note that the free parameter c is introduced to balance the two terms in the lower bound of Ts due to that (Ĝ + maxi ‖gs1:Ts,i‖) ∝ Ĝ √ Ts and ∑d i=1 ‖gs1:Ts,i‖ ∝ dĜ √ Ts have different
orders when d is very large. One way to balance these two terms is to set c in the order of √
1/d, resulting O(d/(S + 1)) for the second term in the above convergence bound. Another way is to choose c in the s-th stage such that the two terms in the max of the lower bound of Ts match each other. One can derive a similar order of O(d/(S+ 1)) for the second term in the above convergence bound. It is obvious that the total number of iterations ∑S s=1 Ts is adaptive to the data. Next,
let us present more discussion on the iteration complexity. Note that Ms = O( √ s) by setting
c as a constant. By the boundness of stochastic gradient ‖gs1:Ts,i‖ ≤ O( √ Ts), therefore Ts in the order of O(s) will satisfy the condition in Theorem 3. Thus, in the worst case the iteration complexity for finding E[‖∇φγ(xτ )‖2] ≤ 2 is in the order of ∑S s=1O(s) ≤ O(1/ 4). To show the potential advantage of adaptive step size as in the convex case, let us consider a good case when the cumulative growth of stochastic gradient is slow, e.g., assuming ‖gs1:Ts,i‖ ≤ O(Ts α) with α < 1/2. Then Ts = O(s1/(2(1−α))) will work, and then the total number of iterations∑S s=1 Ts ≤ S1+1/(2(1−α)) ≤ O(1/ 2+1/(1−α)), which is better than O(1/ 4).

heading: 4.3 STAGEWISE STOCHASTIC MOMENTUM (SM) METHODS
text: Finally, we present stagewise stochastic momentum (SM) methods and their convergence results. In the literature, there are two popular variants of stochastic momentum methods, namely, stochastic heavy-ball method (SHB) and stochastic Nesterov’s accelerated gradient method (SNAG). Both methods have been used for training deep neural networks (Krizhevsky et al., 2012; Sutskever et al., 2013), and have been analyzed by (Yang et al., 2016) for non-convex optimization. To contrast with the results in (Yang et al., 2016), we will consider the same unified stochastic momentum methods that subsume SHB, SNAG and SGD as special cases when Ω = Rd. The updates are presented in Algorithm 4 in the Appendix. There are two additional parameters: β ∈ (0, 1) is the momentum parameter and ρ is a parameter that can vary between [0, 1/(1 − β)]. By changing the value of ρ, we can obtain the three variants, SHB (ρ = 0), SNAG (ρ = 1) and SGD (ρ = 1/(1− β)). Due to the limit of space, we only present the convergence of stagewise SM methods below. Theorem 4. Suppose Assumption 1 holds. By setting γ = 1/(2µ), ηs = (1− β)γ/(96s(ρβ + 1)), Ts ≥ 2304(ρβ + 1)s, with τ similarly defined as in Theorem 1 stagewise SM methods guarantee
E[‖∇φγ(xτ )‖2] ≤ 16µ∆φ(α+ 1) S + 1 + G2(β + 48(2ρβ + 1)(1− β))(α+ 1) 96(ρβ + 1)(1− β)(S + 1) .
Remark: The bound in the above theorem is in the same order as that in Theorem 2. The total iteration complexity for finding a solution xτ with E [ ‖∇φγ(xτ )‖2 ] ≤ 2 is O(1/ 4) similar to that achieved in (Yang et al., 2016).

heading: 5 EXPERIMENTS
text: In this section, we present some empirical results to verify the effectiveness of the proposed stagewise algorithms. We use two benchmark datasets, namely CIFAR-10 and CIFAR-100 (Krizhevsky et al.) for our experiments. We implement the proposed stagewise algorithms in TensorFlow. We compare different algorithms for learning ResNet-20 (He et al., 2016) with batch normalization (Ioffe & Szegedy, 2015) adopted after each convolution and before ReLu activation. Baselines. We compare the proposed stagewise algorithms with their variants implemented in TensorFlow. It is notable that ADAGRAD has a step size (aka learning rate) parameter 1, which is a constant in theory (Li & Orabona, 2018; Chen et al., 2018a; Zou & Shen, 2018). However, in the deep learning community a heuristic fixed frequency decay scheme for the step size parameter is commonly adopted (Ren et al., 2018; Wilson et al., 2017). We thus compare two implementations of ADAGRAD - one with a constant learning rate parameter and another one with a fixed frequency decay scheme, which are referred to as ADAGRAD (theory) and ADAGRAD (heuristic). For each baseline algorithms of SGD, SHB, SNAG, we also implement two versions - a theory version with iteratively decreasing size η0/ √ t suggested by previous theories and a heuristic approach with fixed frequency decay scheme used in practice, using (theory) and (heuristic) to indicate them. The fixed 1note it is not equivalent to the step size in SGD.
frequency decay scheme used in the heuristic variants is similar as that in (He et al., 2016), i.e., the step size parameter is decreased by 10 at 40k, 60k iterations. We also compare stagwise ADAGRAD with AMSGrad (Reddi et al., 2018) - a corrected version of Adam. Parameters. The stagewise step size ηs = η0/ √ s is used in stagwise ADAGRAD, the number of iterations Ts in stagwise ADAGRAD is set according to Theorem 3 with some simplifications for dealing with unknown Ĝ, in particular we set Ts to the smallest value larger than T0 √ smaxi ‖gs1:Ts,i‖ ∑ i ‖gs1:Ts,i‖. For stagewise SGD, SHB, SNAG, the stagewise step size and
iteration number is set to ηs = η0/s and Ts = T0s, respectively. The involved parameters in the compared algorithms are tuned for the best performance, including the initial step size parameter η0 of all algorithms, the value of T0 and γ for our stagewise algorithms. Results. We consider two settings - with/without an `2 norm regularization on weights. For comparison, we evaluate the training error and testing error of obtained solutions in the process of training. For our stagewise algorithms, the evaluation is done based on the current averaged solution, and for other baselines the evaluation is done based on the current solution. Due to the limit of space, we only show the results of testing error on CIFAR-10 and CIFAR-100 for the setting without regularization in Figure 1. All results are included in the Appendix. From all results, we have several observations. (i) The proposed stagewise algorithms perform much better in terms of testing error than the existing theoretical versions reported in the literature (marked with theory in the legend). This indicates the proposed stagewise step size scheme is better than iteratively decreasing step size scheme. (ii) The proposed stagewise algorithms achieve similar and sometimes even better testing error than the heuristic approaches with a fixed frequency decay scheme used in practice. However, the heuristic approaches usually give smaller training error. This seems to indicate that the proposed algorithms are less vulnerable to the overfitting. In another word, the proposed algorithms have smaller generalization error, i.e., the difference between the testing error and the training error.

heading: 6 CONCLUSION & FUTURE WORK
text: In this paper, we have proposed a universal stagewise learning framework for solving stochastic nonconvex optimization problems, which employs well-known tricks in practice that have not been well analyzed theoretically. We provided theoretical convergence results for the proposed algorithms for non-smooth non-convex optimization problems. We also established an adaptive convergence of a stochastic algorithm using data adaptive coordinate-wise step size of ADAGRAD, and exhibited its faster convergence than non-adaptive stepsize when the growth of cumulative stochastic gradients is slow similar to that in the convex case. For future work, one may consider developing more variants of the proposed meta algorithm, e.g., stagewise AMSGrad, stagewise RMSProp, etc. We will also consider the empirical studies on the large-scale ImageNet data set.

heading: A MORE EXPERIMENTAL RESULTS
text: In this section, we present more experimental results. Comparison of training and testing error in the two settings (w/o regularization) on the two data sets are plotted in Figure 2, 3, 4, 5. We also report the final testing error (after running 80k iterations) of different algorithms in the two settings on the two datasets in Table 1. For parameter tuning, the initial step sizes of all algorithms are tuned in {0.1, 0.3, 0.5, 0.7, 0.9}. The value of γ of stagewise algorithms is tuned in {1, 10, 100, 500, 1000, 1500, 2000, 3000}. The initial value T0 for stagewise SGD, SHB, SNAG is tuned in {10, 100, 1k, 5k, 6k, 7k, 10k, 20k}, and that for stagewise ADAGRAD is tuned in {1, 10, 15, 20, 25, 50, 100}.

heading: B PROOF OF THEOREM 1
text: Proof. Below, we use Es to denote expectation over randomness in the s-th stage given all history before s-th stage. Define
zs = arg min x∈Ω fs(x) = proxγ(φ+δΩ)(xs−1) (8)
Then∇φγ(xs−1) = γ−1(xs−1− zs). By applying the convergence bound of SA to fs(x), we have
Es[fs(xs)− fs(zs)] ≤ ε1(ηs, Ts,Θ)‖xs−1 − zs‖22 + ε2(ηs, Ts,Θ)(fs(xs−1)− fs(zs)) + ε3(ηs, Ts,Θ)︸ ︷︷ ︸ Es .
It then follows that
Es [ φ(xs) + 1
2γ ‖xs − xs−1‖2
] ≤ fs(zs) + Es ≤ fs(xs−1) + Es
≤ φ(xs−1) + Es (9)
On the other hand, we have that
‖xs − xs−1‖2 =‖xs − zs + zs − xs−1‖2
=‖xs − zs‖2 + ‖zs − xs−1‖2 + 2〈xs − zs, zs − xs−1〉 ≥(1− α−1s )‖xs − zs‖2 + (1− αs)‖xs−1 − zs‖2
where the inequality follows from the Young’s inequality with 0 < αs < 1. Combining the above inequality with (9) we have that
Es
[ (1− αs)
2γ ‖xs−1 − zs‖2
] ≤ Es [ φ(xs−1)− φ(xs) +
(α−1s − 1) 2γ
‖xs − zs‖2 + Es ]
≤ Es [ φ(xs−1)− φ(xs) +
(α−1s − 1) γ(γ−1 − µ)
(fs(xs)− fs(zs)) + Es ]
≤ Es [ φ(xs−1)− φ(xs) +
α−1s − γµ (1− γµ)
Es ] ≤ Es [ φ(xs−1)− φ(xs) ] + Es [ α−1s − γµ (1− γµ) {ε1(ηs, Ts,Θ)‖xs−1 − zs‖2 + ε2(ηs, Ts,Θ)(fs(xs−1)− fs(zs)) + ε3(ηs, Ts,Θ)} ] ,
(10)14
where the second inequality uses the strong convexity of fs(x), whose strong convexity parameter is γ−1−µ. Next, we bound fs(xs−1)− fs(zs) given that xs−1 is fixed. According to the definition
of fs(·), we have
fs(xs−1)− fs(zs) = φ(xs−1)− φ(zs)− 1
2γ ‖zs − xs−1‖2
= φ(xs−1)− φ(xs) + φ(xs)− φ(zs)− 1
2γ ‖zs − xs−1‖2 = [φ(xs−1)− φ(xs)] + [ fs(xs)− fs(zs) + 1
2γ ‖zs − xs−1‖2 −
1
2γ ‖xs − xs−1‖2
] − 1
2γ ‖zs − xs−1‖2
≤ [φ(xs−1)− φ(xs)] + [fs(xs)− fs(zs)].
Taking expectation over randomness in the s-th stage on both sides, we have
fs(xs−1)− fs(zs) ≤ Es[φ(xs−1)− φ(xs)] + Es[fs(xs)− fs(zs)] ≤ E[φ(xs−1)− φ(xs)] + ε1(ηs, Ts,Θ)‖xs−1 − zs‖22 + ε2(ηs, Ts,Θ)(fs(xs−1)− fs(zs)) + ε3(ηs, Ts,Θ).
Thus,
(1− ε2(ηs, Ts,Θ))(fs(xs−1)− fs(zs)) ≤ E[φ(xs−1)− φ(xs)] + ε1(ηs, Ts,Θ)‖xs−1 − zs‖22 + ε3(ηs, Ts,Θ).
Assuming that ε2(ηs, Ts,Θ) ≤ 1/2, we have
ε2(ηs, Ts,Θ)(fs(xs−1)− fs(zs)) ≤ Es[φ(xs−1)− φ(xs)] + ε1(ηs, Ts,Θ)‖xs−1 − zs‖22 + ε3(ηs, Ts,Θ).
Plugging this upper bound into (10), we have
Es
[ (1− αs)
2γ ‖xs−1 − zs‖2
] ≤ Es [ φ(xs−1)− φ(xs) ] + Es [ α−1s − γµ (1− γµ) {2ε1(ηs, Ts,Θ)‖xs−1 − zs‖2 + φ(xs−1)− φ(xs) + 2ε3(ηs, Ts,Θ)} ] (11)
By setting αs = 1/2, γ = 1/(2µ) and assuming ε1(ηs, Ts,Θ) ≤ 1/(48γ), we have
Es
[ 1
8γ ‖xs−1 − zs‖2
] ≤ 4Es [ φ(xs−1)− φ(xs) ] + 6ε3(ηs, Ts,Θ)}
Define ws = sα. Multiplying both sides by ws, we have that wsγEs[‖∇φγ(xs−1)‖2] ≤ Es [ 32ws∆s + 48ε3(ηs, Ts,Θ)ws ] By summing over s = 1, . . . , S + 1, we have
S+1∑ s=1 wsE[‖∇φγ(xs−1)‖2] ≤ E [ 32 γ S+1∑ s=1 ws∆s + 48 γ S+1∑ s=1 wsε3(ηs, Ts,Θ) ] Taking the expectation w.r.t. τ ∈ {0, . . . , S}, we have that
E[‖∇φγ(xτ )‖2]] ≤ E [ 32 ∑S+1 s=1 ws∆s
γ ∑S+1 s=1 ws
+ 48 ∑S+1 s=1 wsε3(ηs, Ts,Θ))
γ ∑S+1 s=1 ws ] For the first term on the R.H.S, we have that S+1∑ s=1 ws∆s = S+1∑ s=1 ws(φ(xs−1)− φ(xs)) = S+1∑ s=1 (ws−1φ(xs−1)− wsφ(xs)) + S+1∑ s=1 (ws − ws−1)φ(xs−1)
= w0φ(x0)− wS+1φ(xS+1) + S+1∑ s=1 (ws − ws−1)φ(xs−1)
= S+1∑ s=1 (ws − ws−1)(φ(xs−1)− φ(xS+1)) ≤ ∆φ S+1∑ s=1 (ws − ws−1) = ∆φwS+1
Then,
E[‖∇φγ(xτ )‖2] ≤ 32∆φwS+1 γ ∑S+1 s=1 ws
+ 48 ∑S+1 s=1 wsε3(ηs, Ts,Θ)
γ ∑S+1 s=1 ws
The standard calculus tells that S∑ s=1 sα ≥ ∫ S 0 xαdx = 1 α+ 1 Sα+1
S∑ s=1 sα−1 ≤ SSα−1 = Sα,∀α ≥ 1, S∑ s=1 sα−1 ≤ ∫ S 0 xα−1dx = Sα α ,∀0 < α < 1
Combining these facts and the assumption ε3(ηs, Ts,Θ) ≤ c/s, we have that
E[‖∇φγ(xτ )‖2] ≤  32∆φ(α+1) γ(S+1) + 48c(α+1) γ(S+1) α ≥ 1
32∆φ(α+1) γ(S+1) + 48c(α+1) γ(S+1)α 0 < α < 1
In order to have E[‖∇φγ(xτ )‖2] ≤ 2, we can set S = O(1/ 2). The total number of iterations is S∑ s=1 Ts ≤ S∑ s=1 12γs ≤ 6γS(S + 1) = O(1/ 4)

heading: C PROOF OF THEOREM 2
text: Proof. The proof is almost a duplicate to that of Theorem 1. Define ws = sα. We apply Lemma 1 to each call of SGD in stagewise SGD,
E[fs(xs)− fs(zs)] ≤ ‖zs − xs−1‖2 2ηsTs + ηsĜ 2
2︸ ︷︷ ︸ Es ,
where Ĝ2 is the upper bound of E[‖g(x; ξ) + γ−1(x − xs−1)‖2], which exists and can be set to 2G2 + 2γ−2D2 due to the Assumption 1-(ii) and the bounded assumption of the domain. Then following the same analysis as that in the proof of Theorem 1, we have
Es
[ (1− αs)
2γ ‖xs−1 − zs‖2
] ≤Es [ φ(xs−1)− φ(xs) +
(α−1s − 1) 2γ
‖xs − zs‖2 + Es ]
≤E [ φ(xs−1)− φ(xs) +
(α−1s − 1) γ(γ−1 − µ)
(fs(xs)− fs(zs)) + Es ]
≤E [ φ(xs−1)− φ(xs) +
α−1s − γµ (1− γµ)
Es ]
(12)
Combining the above inequalities, we have that( (1− αs)γ −
γ2(α−1s − µγ) (1− µγ)ηsTs
) Es[‖∇φγ(xs−1)‖2] ≤ Es [ 2∆s + (α−1s − µγ)ηsĜ2
(1− µγ) ] Multiplying both sides by ws, we have that
ws ( (1− αs)γ −
γ2(α−1s − µγ) (1− µγ)ηsTs
) Es[‖∇φγ(xs−1)‖2] ≤ Es [ 2ws∆s + (α−1s − µγ)wsηsĜ2
(1− µγ) ] By setting αs = 1/2 and γ = 1/(2µ), Tsηs ≥ 12γ, we have
1 4 wsγEs[‖∇φγ(xs−1)‖2] ≤ Es[2ws∆s + 3wsηsĜ2]
By summing over s = 1, . . . , S + 1, we have
S+1∑ s=1 wsE[‖∇φγ(xs−1)‖2] ≤ E [ 16µ S+1∑ s=1 ws∆s + 24µ S+1∑ s=1 wsηsĜ 2 ] Taking the expectation w.r.t. τ ∈ {0, . . . , S}, we have that
E[‖∇φγ(xτ )‖2]] ≤ E [ 16µ ∑S+1 s=1 ws∆s∑S+1 s=1 ws + 24µ ∑S+1 s=1 wsηsĜ 2∑S+1 s=1 ws ] By similar analysis, we have that
E[‖∇φγ(xτ )‖2] ≤  16µ∆φ(α+1) S+1 + 24η0µĜ 2(α+1) S+1 α ≥ 1
16µ∆φ(α+1) S+1 +
24η0µĜ 2(α+1)
(S+1)α 0 < α < 1
In order to have E[‖∇φγ(xτ )‖2] ≤ 2, we can set S = O(1/ 2). The total number of iterations is S∑ s=1 Ts ≤ S∑ s=1 12γs ≤ 6γS(S + 1) = O(1/ 4)

heading: D PROOF OF THEOREM 4
text: We need the following lemma for the convergence bound of stochastic momentum methods for a strongly convex problem, whose proof is postponed to Section F. Lemma 3. For Algorithm 4, assume f(x) = φ(x) + 12γ ‖x− x0‖
2 is a λ-strongly convex function, gt = g(xt; ξ) + 1 γ (xt − x0) where g(x; ξ) ∈ ∂Fφ(xt) such that E[‖g(x; ξ)‖
2] ≤ G2, and η ≤ (1− β)γ2λ/(8ρβ + 4), then we have that
E[f(x̂T )− f(x∗)] ≤ (1− β)‖x0 − x∗‖2
2η(T + 1) + β(f(x0)− f(x∗)) (1− β)(T + 1) + 2ηG2(2ρβ + 1) 1− β + 4ρβ + 4 (1− β) η γ2 ‖x0 − x∗‖2 (13)
where x̂T = ∑T t=0 xt/(1 + T ) and x∗ ∈ arg minx∈Rd f(x).
Remark: It is notable that in the above result, we do not use the bounded domain assumption since we consider Ω = Rd for the unified momentum methods in this subsection. The key to get rid of bounded domain assumption is by exploring the strong convexity of f(x) = φ(x) + 12γ ‖x− x0‖ 2.
Proof. of Theorem 4 According to the definition of zs in (8) and Lemma 3, we have that
Es [ φ(xs) + 1
2γ ‖xs − xs−1‖2 ] ≤ fs(zs) +
β(fs(xs−1)− fs(zs)) (1− β)(Ts + 1) + (1− β)‖xs−1 − zs‖2 2ηs(Ts + 1) + 2ηsG 2(2ρβ + 1) 1− β + 1
24γ ‖xs−1 − zs‖2︸ ︷︷ ︸
Es
≤ φ(xs−1) + Es,
where the last inequality uses the value of ηs = (1−β)γ/(96s(ρβ+ 1)) ≤ (1−β)γ/(96(ρβ+ 1)), which also satisfies the condition in Lemma 3 by noting that λ = γ−1 − µ = 1/(2γ). Similar to the proof of Theorem 1, we have
(1− αs) 2γ ‖xs−1 − zs‖2 ≤Es[φ(xs−1)− φ(xs)] + α−1s − γµ (1− γµ) Es (14)
Plugging the expression of Es and rearranging above inequality, we have that( (1− αs)γ −
γ2(α−1s − µγ)(1− β) (1− µγ)ηs(Ts + 1) − α −1 s − γµ (1− γµ) γ 24
) ‖∇φγ(xs−1)‖2
≤2Es[∆s] + 2(α−1s − µγ)
(1− µγ)
[ β(fs(xs−1)− fs(zs))
(1− β)(Ts + 1) +
2ηsĜ 2(2ρβ + 1)
1− β ] The definition of fs gives that
fs(xs−1)− fs(zs) = φ(xs−1)− φ(zs)− 1
2γ ‖zs − xs−1‖2
On the other hand, the µ-weakly convexity of φ gives that
φ(zs) ≥ φ(xs−1) + 〈g(xs−1), zs − xs−1〉 − µ
2 ‖zs − xs−1‖2,
where g(xs−1) ∈ ∂Fφ(xs−1). Combing these two inequalities we have that
fs(xs−1)− fs(zs) ≤〈g(xs−1),xs−1 − zs〉 − µ
2 ‖zs − xs−1‖2
≤G 2 2µ + µ− µ 2 ‖zs − xs−1‖2 = G2 2µ
where the second inequality follows from Jensen’s inequality for ‖ · ‖ and Young’s inequality. Combining above inequalities and multiplying both side by ws, we have that
ws ( (1− αs)γ −
γ2(α−1s − µγ)(1− β) (1− µγ)ηs(Ts + 1) − α −1 s − γµ (1− γµ) γ 24
) ‖∇φγ(xs−1)‖2
≤2wsEs[∆s] + 2ws(α
−1 s − µγ)
(1− µγ)
[ βG2
2µ(1− β)(Ts + 1) +
2ηsG 2(2ρβ + 1)
1− β
] (15)
By setting αs = 1/2, ηs(Ts + 1) ≥ 24(1− β)γ, we have that
wsγ
4 ‖∇φγ(xs−1)‖2 ≤ 2wsEs[∆s] +
wsηsβG 2 4(1− β)2 + 12wsηsG 2(2ρβ + 1)
1− β Summing over s = 1, . . . , S + 1 and rearranging, we have
S+1∑ s=1 ws‖∇φγ(xs−1)‖2 = E [ S+1∑ s=1 8 γ ws∆s + wsηsG 2(β + 48(2ρβ + 1)(1− β)) γ(1− β)2 ] Following similar analysis as in the proof of Theorem 2, we can finish the proof.

heading: E PROOF OF THEOREM 3
text: Proof. Applying Lemma 2 with Ts ≥ Ms max{ Ĝ+maxi ‖gs1:Ts,i‖ 2c , c ∑d i=1 ‖gs1:Ts,i‖} Ms > 0, and the fact that φ(xs−1) ≥ φ(zs) + 12γ ‖xs−1 − zs‖ 2 in sth stage, we have that
Es [ φ(xs) + 1
2γs ‖xs − xs−1‖2
] ≤ fs(zs) + c
Msηs ‖xs−1 − zs‖2 + ηs Msc︸ ︷︷ ︸
Es
≤ φ(xs) + Es According to (14), we have that
(1− αs) 2γ Es[‖xs−1 − zs‖2] ≤φ(xs−1)− φ(xs) + (α−1s − 1) 2γ ‖xs − zs‖2 + Es
≤φ(xs−1)− φ(xs) + α−1s − γµ (1− γµ)
( c
Msηs ‖xs−1 − zs‖2 + ηs Msc
)
Rearranging above inequality then multiplying both side by ws, we have that
ws ( (1− αs)γ−
2γ2c(α−1s − µγ) (1− µγ)Msηs
) ‖∇φγ(xs−1)‖2
≤2wsEs[∆s] + 2wsηs(α
−1 s − µγ)
cMs(1− µγ)
By using Msηs ≥ 24γc and summing over s = 1, . . . , S + 1, we have that S+1∑ s=1 ws‖∇φγ(xs−1)‖2 ≤ E [ S+1∑ s=1 8ws∆s γ + wsη 2 s c2γ2 ] By the definition of τ in the theorem, taking expectation of ‖∇φγ(xτ )‖2 w.r.t. τ ∈ {0, . . . , S} we have that
E[‖∇φγ(xτ )‖2] =E [ 8
γ S+1∑ s=1 ws∆s∑S+1 i=1 wi ] + η20 c2γ2 S+1∑ s=1 sα−1∑S+1 i=1 wi
≤8∆φ(α+ 1) γ(S + 1) + η20(α+ 1) c2γ2(S + 1)αI(α<1) ,
where I(α < 1) is 1 if α < 1 and 0 otherwise.

heading: F PROOF OF LEMMA 3
text: Proof. Following the analysis in Yang et al. (2016), we directly have the following inequality,
E[‖xk+1 + pk+1 − x∗‖2] =
= E[‖xk + pk − x∗‖2]− 2η
1− β E[(xk − x∗)>∂f(xk)]−
2ηβ
(1− β)2 E[(xk − xk−1)>∂f(xk)]
− 2ρη 2β
(1− β)2 E[g>k−1∂f(xk)] +
( η
1− β
)2 E[|gk‖2]
We also note that
f(xk)− f(x∗) ≤ (xk − x∗)>∂f(xk)− λ
2 ‖xk − x∗‖2
f(xk)− f(xk−1) ≤ (xk − xk−1)>∂f(xk)− λ
2 ‖xk − xk−1‖2
− E[g>k−1∂f(xk)] ≤ E[‖gk−1‖2 + ‖∂f(xk)‖2] 2 ≤ 1 γ2 ‖xk−1 − x0‖2 + 1 γ2 ‖xk − x0‖2 + 2G2
Ek[‖gk‖2] ≤ 2
γ2 ‖xk − x0‖2 + 2G2
where the first two inequalities are due to the strong convexity of f(·) and the last three inequalities are due to the boundness assumption. Thus
E[‖xk+1 + pk+1 − x‖2] ≤ E[‖xk + pk − x‖2]− 2η
1− β E[(f(xk)− f(x))]
− 2ηβ (1− β)2
E[(f(xk)− f(xk−1))] + ( η
1− β
)2 (2ρβ + 1)4G2
− λη 1− β ‖xk − x∗‖2 − ληβ (1− β)2 ‖xk − xk−1‖2
+ 2ρβ (1− β)2 η2 γ2 ‖xk−1 − x0‖2 + 2ρβ + 2 (1− β)2 η2 γ2 ‖xk − x0‖2
By summarizing the above inequality over k = 0, . . . , T , we have
2η
1− β E [ T∑ k=0 (f(xk)− f(x∗)) ] ≤ E[‖x0 − x∗‖2] + 2ηβ (1− β)2 E[f(x0)− f(x∗)]
+
( η
1− β
)2 (2ρβ + 1)4G2(T + 1)
− ηλ 1− β T∑ k=0 ‖xk − x∗‖2 + 4ρβ (1− β)2 η2 γ2 T∑ k=0 ‖xk−1 − x∗‖2 + 4ρβ + 4 (1− β)2 η2 γ2 T∑ k=0 ‖xk − x∗‖2
+ 4ρβ + 4 (1− β)2 η2 γ2 (T + 1)‖x0 − x∗‖2
When η ≤ (1− β)γ2λ/(8ρβ + 4), we have
E [ (f(x̂T )− f(x∗)) ] ≤ (1− β)‖x0 − x∗‖ 2
2η(T + 1) +
β (1− β) f(x0)− f(x∗) T + 1 + η 1− β (2ρβ + 1)2G2
+ 4ρβ + 4 (1− β) η γ2 ‖x0 − x∗‖2

heading: G PROOF OF LEMMA 2
text: The proof is almost a duplicate of the proof of Proposition 1 in Chen et al. (2018b). For completeness, we present a proof here.
Proof. Let ψ0(x) = 0 and ‖x‖H = √ x>Hx. First, we can see that ψt+1(x) ≥ ψt(x) for any t ≥ 0. Define ζt = ∑t τ=1 gt and ∆τ = (∂F (xt)− gt)>(xt − x). Let ψ∗t be defined by
ψ∗t (g) = sup x∈Ω g>x− 1 η ψt(x)
Taking the summation of objective gap in all iterations, we have
T∑ t=1 (f(xt)− f(x)) ≤ T∑ t=1 ∂f(xt) >(xt − x) = T∑ t=1 g>t (xt − x) + T∑ t=1 ∆t
= T∑ t=1 g>t xt − T∑ t=1 g>t x− 1 η ψT (x) + 1 η ψT (x) + T∑ t=1 ∆t
≤ 1 η ψT (x) + T∑ t=1 g>t xt + T∑ t=1 ∆t + sup x∈Ω
{ −
T∑ t=1 g>t x− 1 η ψT (x)
}
= 1
η ψT (x) + T∑ t=1 g>t xt + ψ ∗ T (−ζT ) + T∑ t=1 ∆t
Note that
ψ∗T (−ζT ) = − T∑ t=1 g>t xT+1 − 1 η ψT (xT+1) ≤ − T∑ t=1 g>t xT+1 − 1 η ψT−1(xT+1)
≤ sup x∈Ω −ζ>T x−
1 η ψT−1(x) = ψ ∗ T−1(−ζT )
≤ ψ∗T−1(−ζT−1)− g>T∇ψ∗T−1(−ζT−1) + η
2 ‖gT ‖2ψ∗T−1
where the last inequality uses the fact that ψt(x) is 1-strongly convex w.r.t ‖ · ‖ψt = ‖ · ‖Ht and consequentially ψ∗t (x) is η-smooth wr.t. ‖ · ‖ψ∗t = ‖ · ‖H−1t . Thus, we have
T∑ t=1 g>t xt + ψ ∗ T (−ζT ) ≤ T∑ t=1 g>t xt + ψ ∗ T−1(−ζT−1)− g>T∇ψ∗T−1(−ζT−1) + η 2 ‖gT ‖2ψ∗T−1
= T−1∑ t=1 g>t xt + ψ ∗ T−1(−ζT−1) + η 2 ‖gT ‖2ψ∗T−1
By repeating this process, we have T∑ t=1 g>t xt + ψ ∗ T (−ζT ) ≤ ψ∗0(−ζ0) + η 2 T∑ t=1 ‖gt‖2ψ∗t−1 = η 2 T∑ t=1 ‖gt‖2ψ∗t−1 Then T∑ t=1 (f(xt)− f(x)) ≤ 1 η ψT (x) + η 2 T∑ t=1 ‖gt‖2ψ∗t−1 + T∑ t=1 ∆t (16) Following the analysis in Duchi et al. (2011), we have T∑ t=1 ‖gt‖2ψ∗t−1 ≤ 2 d∑ i=1 ‖g1:T,i‖2 Thus T∑ t=1 (f(xt)− f(x)) ≤ G‖x− x1‖22 2η + (x− x1)>diag(sT )(x− x1) 2η + η d∑ i=1 ‖g1:T,i‖2 + T∑ t=1 ∆t
≤ G+ maxi ‖g1:T,i‖2 2η ‖x− x1‖22 + η d∑ i=1 ‖g1:T,i‖2 + T∑ t=1 ∆t
Now by the value of T ≥M max{ (G+maxi ‖g1:T,i)‖2c , c ∑d i=1 ‖g1:T,i‖}, we have
(G+ maxi ‖g1:T,i‖2) 2ηT ≤ c ηM ,
η ∑d i=1 ‖g1:T,i‖2
T ≤ η Mc
Dividing by T on both sides and setting x = x∗, following the inequality (3) and the convexity of f(x) we have
f(x̂)− f∗ ≤ c
Mη ‖x0 − x∗‖2 +
η
Mc +
1
T T∑ t=1 ∆t
Let {Ft} be the filtration associated with Algorithm 1 in the paper. Noticing that T is a random variable with respect to {Ft}, we cannot get rid of the last term directly. Define the Sequence {Xt}t∈N+ as
Xt = 1
t t∑ i=1 ∆i = 1 t t∑ i=1 〈gi − E[gi],xi − x∗〉 (17)
where E[gi] ∈ ∂f(xi). Since E [gt+1 − E[gt+1]] = 0 and xt+1 = arg min x∈Ω
ηx> (
1 t ∑t τ=1 gτ ) +
1 tψt(x), which is measurable with respect to g1, . . . ,gt and x1, . . . ,xt, it is easy to see {∆t}t∈N is a martingale difference sequence with respect to {Ft}, e.g. E[∆t|Ft−1] = 0. On the other hand, since ‖gt‖2 is upper bounded by G, following the statement of T in the theorem, T ≤ N = M2 max{ (G+1) 2
4 , d 2G2} < ∞ always holds. Then following Lemma 1 in (Chen et al., 2018b) we
have that E[XT ] = 0.
Now taking the expectation we have that
E[f(x̂)− f∗] ≤ c
Mη ‖x0 − x∗‖2 +
ηc
M Then we finish the proof.

heading: H PROOF OF INEQUALITY (3)
text: Let us recall the definition of x̂ = proxγf (x).
proxγf (x) = arg minz f(z) + 1 2γ ‖z− x‖2 (18)
First, note that when f(z) is µ-weakly convex and γ < µ−1, the above problem is strongly convex and x̂ = proxγf (x) is well-defined and unique.
By the optimality condition of x̂, it is clear that f(x̂) + 12γ ‖x̂− x‖ 2 ≤ f(x), which proves the first inequality f(x̂) ≤ f(x). Note that when x = x∗ ∈ arg minz f(z), we can prove that x̂ = x = x∗. This is because that by definition of x̂, it is unique and satisfies 0 ∈ ∂f(x̂)+ 1γ (x̂−x). When x = x∗, we have 0 ∈ ∂f(x) and x∗ clearly satisifes the first-order condition 0 ∈ ∂f(x∗) + 1γ (x∗ − x). By the uniqueness of x̂, it follows that x̂ = x∗.
The first-order condition 0 ∈ ∂f(x̂) + 1γ (x̂− x) also gives that dist(0, ∂f(x̂)) ≤ 1 γ ‖x̂− x‖. Also, we have∇fγ(x) = 1γ (x̂−x) (Rockafellar, 1970), which implies the second and the third inequality in (3).

Review summaries:
Non-convex optimization is a hot topic since many machine learning problems can be formulated as non-convex problems. In this paper, the authors propose a universal stage-wise algorithm for weakly convex optimization problems. The idea is to add a strongly convex regularizer centered at an iterate of previous stage to the objective function. This builds a convex function which can be optimized by any standard methods in the convex optimization setting. The authors developed convergence rates in expectation in terms of the gradient of envelope. Empirical results are also reported to show the effectiveness of the method.