A Near-Optimal Algorithm for Stochastic Bilevel Optimization via Double-Momentum
Year: 2021

Abstract:
This work proposes a new algorithm – the Single-timescale Double-momentum Stochastic Approximation (SUSTAIN) – for tackling stochastic unconstrained bilevel optimization problems. We focus on bilevel problems where the lower level subproblem is strongly-convex and the upper level objective function is smooth. Unlike prior works which rely on two-timescale or double loop techniques, we design a stochastic momentum-assisted gradient estimator for both the upper and lower level updates. The latter allows us to control the error in the stochastic gradient updates due to inaccurate solution to both subproblems. If the upper objective function is smooth but possibly non-convex, we show that SUSTAIN requires O(✏ 3/2) iterations (each using O(1) samples) to find an ✏-stationary solution. The ✏-stationary solution is defined as the point whose squared norm of the gradient of the outer function is less than or equal to ✏. The total number of stochastic gradient samples required for the upper and lower level objective functions match the best-known complexity for single-level stochastic gradient algorithms. We also analyze the case when the upper level objective function is strongly-convex.

heading: 1 Introduction
text: Many learning and inference problems take a “hierarchical" form, wherein the optimal solution of one problem affects the objective function of others [27]. Bilevel optimization is often used to model problems of this kind with two levels of hierarchy [27, 8], where the variables of an upper level problem depend on the optimizer of certain lower level problem. In this work, we consider unconstrained bilevel optimization problems of the form:
minx2Rdup `(x) = f(x, y ⇤(x)) := E⇠[f(x, y⇤(x); ⇠)]
s.t. y⇤(x) = argminy2Rdlo g(x, y) := E⇣ [g(x, y; ⇣)] ,
(1)
where f, g : Rdup ⇥ Rdlo ! R with x 2 Rdup and y 2 Rdlo ; f(x, y; ⇠) with ⇠ ⇠ ⇡f (resp. g(x, y; ⇣) with ⇣ ⇠ ⇡g) represents a stochastic sample of the upper level objective (resp. lower level objective). Note here that the upper level objective f depends on the minimizer of the lower level objective g, and we refer to `(x) as the outer function. Throughout this paper, g(x, y) is assumed to be strongly-convex in y, which implies that `(x) is smooth but possibly non-convex.
The applications of (1) include many machine learning problems that have a hierarchical structure. Examples are meta learning [13, 31], data hyper-cleaning [35], hyper-parameter optimization [12,
⇤Corresponding Author: Mingyi Hong.
35th Conference on Neural Information Processing Systems (NeurIPS 2021).
13, 29], and reinforcement learning [22], etc.. To better contextualize our study, below we describe examples on meta-learning problem and data hyper-cleaning problem:
Example 1: Meta learning. The meta learning problem aims to learn task specific parameters that generalize to a diverse set of tasks [30]. Suppose we have M tasks {Ti, i = 1, . . . ,M} and each task has a corresponding loss function L(x, yi; ⇠i) with ⇠i representing a data sample for task Ti, x 2 Rdup the model parameters shared among tasks, and yi 2 Rd i lo the task specific parameters. The goal of meta learning is then to solve the following problem:
minx2Rdup Lts(x, ȳ⇤(x)) := 1 M PM i=1 E⇠i⇠Di [L(x, y⇤i (x); ⇠i)]
s.t. ȳ⇤(x) 2 argmin ȳ2R PM i=1 d i lo Ltr(x, ȳ) := 1 M PM i=1 E⇣i⇠Si [L(x, yi; ⇣i)] +R(yi) , (2)
where ȳ = [yT1 , . . . , yTM ] T , R(·) is a strongly convex regularizer while Si and Di are the training and testing datasets for task Ti. Compared to the number of tasks, the dataset sizes are usually small for meta-learning problems, so the stochasticity in tackling (2) results from the fact that at each iteration we can only sample a subset m out of M tasks. Note that this problem is a special case of (1). ⇤ Example 2: Data hyper-cleaning. The data hyper-cleaning is a hyperparameter optimization problem that aims to train a classifier model with a dataset of randomly corrupted labels [35]. The optimization problem is formulated below:
minx2Rdup `(x) := P i2Dval L(a > i y ⇤(x), bi) (3)
s.t. y⇤(x) = argminy2Rdlo ckyk
2 + P
i2Dtr (xi)L(a > i y, bi) .
In this problem, we have dup = |Dtr| and dlo is the dimension of the classifier. Moreover, (ai, bi) is the ith data point; L(·) is the loss function, with y being the model parameter; xi is the parameter that determines the weight for the ith data sample, and : R! R+ is the weight function; c > 0 is a regularization parameter; Dval and Dtr are validation and training sets, respectively. Clearly, (3) is a special case of (1) where the lower level problem finds the classifier y⇤(x) with the training set Dtr, and the upper level problem finds the best weights x with respect to the validation set Dval. ⇤ A natural approach to tackling (1) is to apply alternating stochastic gradient (SG) updates. Let ,↵ > 0 be some step sizes, one performs the recursion
y + y r̂yg(x, y), x + x ↵r̂x ˆ̀(x; y) (4)
such that r̂yg(x, y), r̂x ˆ̀(x; y) are stochastic estimates of ryg(x, y), r`(x), respectively. Notice that (4) is significantly different from the standard alternating primal-dual gradient algorithm for saddle point problems. Particularly, the design of r̂x ˆ̀(x; y) is crucial to the SG scheme in (4). Observe that r`(x) can be computed using the implicit function theorem, and its evaluation requires f(·, ·) and y?(x), the minimizer of g(x, y) given x (cf. (5)). This gives rise to a unique challenge to bilevel optimization, where y?(x) can only be approximated by y obtained in the first relation of (4).
In light of the above observations, previous endeavors have considered two approaches to improve the estimate of y?(x) while r̂x ˆ̀(x; y) is used as a biased approximation of r`(x). The first approach is to apply the double-loop algorithms. For example, [14] proposed to repeat the y+ update for multiple times to obtain a better estimate of y?(x) before performing the x+ update, [19] proposed
to take a large batch size to estimate ryg(x, y). While simple to analyze, these algorithms may suffer from a poor sample complexity for the inner problem. The second approach is to apply single-loop algorithms where the y+-updates are performed simultaneously with the x+-updates. Instead, advanced techniques are utilized that allows y+ to accurately track y?(x). For example, [18] suggested to tune the step size schedule with ↵, [5, 17] proposed single-timescale algorithms with advanced variance reduction techniques. However, the latter two algorithms require Hessian projections onto a compact set along with Hessian matrices inversion which scales poorly with dimension (i.e., in O(d3lo)). We summarize and compare the complexity results of the state-of-the-art algorithms in Table 1.
A careful inspection on the above results reveals a gap in the iteration/sample complexity compared to single-level stochastic optimization. For instance, an optimal stochastic gradient algorithm finds an ✏-stationary solution [cf. Definition 1.1] to minx E⇠[`(x; ⇠)] in O(✏ 3/2) iterations [10, 7, 37, 42]. For bilevel optimization, the fastest rate available is only O(✏ 2) to the best of the authors’ knowledge. In comparison, the proposed algorithm achieves a rate of O(✏ 3/2). During the preparation of the current paper, a preprint [17] has appeared which extended [5], and achieves an improved rate of O(✏ 3/2). We remark that the latter work follows a different design philosophy from ours and maybe less efficient; see the detailed discussion at the end of Sec. 3.
Contributions. In this paper, we depart from the prior developments which focused on finding better inner solutions y⇤(x) to approximate r̂x ˆ̀(x; y) ⇡ r`(x). Our idea is to exploit the gradient estimates from prior iterations to improve the quality of the current gradient estimation. This leads to momentum-assisted stochastic gradient estimators for both ryg(x, y) and r`(x) using similar techniques in [7, 37] for single-level stochastic optimization. The resultant algorithm only requires O(1) samples at each update, and updates x and y using step sizes of the same order, hence the name single-timescale double-momentum stochastic approximation(SUSTAIN) algorithm. Additionally, it is worth noting that our algorithm has a O(d2lo) per iteration complexity, compared to the O(d 3 lo) complexity of STABLE [5] and SVRB [17]. That is, the SUSTAIN algorithm is both sample and computation efficient. Our specific contributions are: • We propose the SUSTAIN algorithm for bilevel problems which matches the best complexity
bounds as the optimal SGD algorithms for single-level stochastic optimization. That is, it requires O(✏ 3/2) [resp. O(✏ 1)] samples to find an ✏-stationary solution for non-convex (resp. stronglyconvex) bilevel problems; see Table 1. Furthermore, the algorithm utilizes a single-loop update with step sizes of the same order for both upper and lower level problems. Such complexity bounds match the optimal sample complexity of stochastic gradient algorithms for single-level problems.
• By developing the Lipschitz continuous property of the (biased) stochastic estimates of r`(x), we show that obtaining a good estimate of r`(x) does not require explicit (sampled) Hessian inversion. This key result ensures that our algorithm depends favorably on the problem dimension.
• Comparing with prior works such as TTSA [18], BSA [14], STABLE [5] and SVRB [17], our analysis reveals that improving the gradient estimation quality for bothryg(x, y) and r`(x) is the key to obtain a sample and computation efficient stochastic algorithm for bilevel optimization.
Related works. The study of the bilevel problem (1) can be traced to that of game theory [36] and was formally introduced in [2–4]. It is also related to the broader class of problems of Mathematical Programming with Equilibrium Constraints [26]. Related algorithms include approximate descent [9, 38], and penalty-based methods [40]; see [6] and [25] for a comprehensive survey.
In addition to the works cited in Table 1, recent works on bilevel optimization have focused on algorithms with provable convergence rates. In [34], the authors proposed BigSAM algorithm for solving simple bilevel problems (with a single variable) with convex lower level problem. Subsequently, the works [24, 23] utilized BigSAM and developed algorithms for a general bilevel problem for the cases when the solution of the lower level problem is not a singleton. Note that all the aforementioned works [34, 24, 23] assumed the upper level problem to be strongly-convex with convex lower level problem. In a separate line of work, backpropagation based algorithms have been proposed to approximately solve bilevel problems [12, 35, 16, 15]. However, the major focus of these works was to develop efficient gradient estimators rather than on developing efficient optimization algorithms.
Notation. For any x 2 Rd, we denote kxk as the standard Euclidean norm; as for X 2 Rn⇥d, kXk is induced by the Euclidean norm. For a multivariate function f(x, y), the notationrxf(x, y) [resp.ryf(x, y)] refers to the partial gradient taken with respect to (w.r.t.) x [resp. y]. For some µ >
0, a function f(x, y) is said to be µ-strongly-convex in x if f(x, y) µ2 kxk 2 is convex in x. For some L > 0, the map A : Rd ! Rm is said to be L-Lipschitz continuous if kA(x) A(y)k  Lkx yk for any x, y 2 Rd. A function f : Rd ! R is said to be L-smooth if its gradient is L-Lipschitz continuous. Uniform distribution over a discrete set {1, . . . , T} is represented by U{1, . . . , T}.
Finally, we state the following definitions for the optimality criteria of (1). Definition 1.1 (✏-Stationary Point). A point x is called ✏-stationary if kr`(x)k2  ✏. A stochastic algorithm is said to achieve an ✏-stationary point in t iterations if E[kr`(xt)k2]  ✏, where the expectation is over the stochasticity of the algorithm until time instant t. Definition 1.2 (✏-Optimal Point). A point x is called ✏-optimal if `(x) `⇤  ✏, where `⇤ := minx2Rdup `(x). A stochastic algorithm is said to achieve an ✏-optimal point in t iterations if E[`(xt) `⇤]  ✏, where the expectation is over the stochasticity of the algorithm until time instant t.

heading: 2 Preliminaries
text: We discuss the assumptions on (1) to specify the problem class of interest. We also preface the proposed algorithm by describing a practical procedure for estimating the stochastic gradients. Assumption 1 (Upper Level Function). f(x, y) satisfies the following conditions:
(i) rxf(x, y) and ryf(x, y) are Lipschitz continuous w.r.t. (x, y) 2 Rdup ⇥ Rdlo , and with constants Lfx 0 and Lfy 0, respectively.
(ii) For any (x, y) 2 Rdup ⇥ Rdlo , we have kryf(x, y)k  Cfy , for some Cfy 0. Assumption 2 (Lower level Function). g(x, y) satisfies the following conditions:
(i) For any x 2 Rdup and y 2 Rdlo , g(x, y) is twice continuously differentiable in (x, y). (ii) ryg(x, y) is Lipschitz continuous w.r.t. (x, y) 2 Rdup ⇥ Rdlo , and with constant Lg 0.
(iii) For any x 2 Rdup , g(x, ·) is µg-strongly-convex in y for some µg > 0. (iv) r2xyg(x, y) and r2yyg(x, y) are Lipschitz continuous w.r.t. (x, y) 2 Rdup ⇥ Rdlo , and with
constants Lgxy 0 and Lgyy 0, respectively. (v) For any (x, y) 2 Rdup ⇥ Rdlo , we have kr2xyg(x, y)k2  Cgxy for some Cgxy > 0.
Assumption 3 (Stochastic Functions). Assumptions 1 and 2 hold for f(x, y; ⇠) and g(x, y; ⇣), for all ⇠ 2 supp(⇡f ) and ⇣ 2 supp(⇡g) where supp(⇡) is the support of ⇡. Moreover, we assume the following variance bounds.
E ⇥ krxf(x, y) rxf(x, y; ⇠)k 2 ⇤  2 fx , Ekryf(x, y) ryf(x, y; ⇠)k 2  2 fy ,
Ekr2xyg(x, y) r2xyg(x, y; ⇠)k2  2gxy for some fx 0, fy 0 and gxy 0.
These assumptions are standard in the analysis of bilevel optimization [14]. For example, they are satisfied by a range of applications such as the meta learning problem (2), data hypercleaning problem (3) with linear classifier. Notice that under these assumptions, the gradient r`(·) is well-defined. By utilizing Assumption 2–(i) and (ii) along with the implicit function theorem [33], it is easy to show that for a given x̄ 2 Rdup , the following holds [14, Lemma 2.1]:
r`(x̄) = rxf(x̄, y ⇤(x̄)) r2xyg(x̄, y ⇤(x̄))[r2yyg(x̄, y ⇤(x̄))] 1ryf(x̄, y ⇤(x̄)). (5)
Obtaining y⇤(x) in closed-form is usually a challenging task, so it is natural to use the following gradient surrogate. At any (x̄, ȳ) 2 Rdup⇥dlo , define:
r̄f(x̄, ȳ) = rxf(x̄, ȳ) r 2 xyg(x̄, ȳ)[r 2 yyg(x̄, ȳ)] 1 ryf(x̄, ȳ). (6)
Evaluating (6) requires computing the exact gradients and Hessian inverse which can be non-trivial. Below, we describe a practical procedure from [14] to generate a biased estimate of r̄f(x̄, ȳ). Stochastic gradient estimator forr`(x). The estimator requires a parameter K 2 N and is based on a collection of K+3 independent samples ⇠̄ := {⇠, ⇣(0), ..., ⇣(K), k(K)}, where ⇠ ⇠ µ, ⇣(i) ⇠ ⇡g , i = 0, ...,K, and k(K) ⇠ U{0, ...,K 1}. We set
r̄f(x, y; ⇠̄) = rxf(x, y; ⇠) K Lg r 2 xyg(x, y; ⇣ (0))
k(K)Y
i=1
✓ I r 2 yyg(x, y; ⇣ (i))
Lg
◆ ryf(x, y; ⇠), (7)
where we have used the convention Qj
i=1 Ai = I if j = 0. It has been shown in [14, 18] that the bias with the gradient estimator (7) decays exponentially fast with K, as summarized below: Lemma 2.1. Under Assumptions 1, 2. For any K 1, the gradient estimator in (7) satisfies
kr̄f(x, y) E⇠̄[r̄f(x, y; ⇠̄)]k  CgxyCfy
µg
✓ 1 µg
Lg
◆K , 8 (x, y) 2 Rdup ⇥ Rdlo . (8)
The detailed statement of the above lemma is included in Appendix C. We remark that each computation of r̄f(x, y; ⇠̄) requires at most K Hessian-vector products, and later we will show that setting K = O(log(T )) is necessary for the proposed algorithm. Sincer2yyg(x, y; ⇣) is of size dlo⇥ dlo, the total complexity of this step is O(log(T )d2lo). On the contrary, STABLE [5] and SVRB [17] require O(d3lo) to estimate the Hessian inverse, which is more computationally expensive when dlo 1. Indeed, it has been explicitly mentioned in [5] that “our algorithm (STABLE) is preferable in the regime where the sampling is more costly than computation or the dimension d is relatively small”.
Notice that (7) is not the only option for estimating the gradient surrogate r̄f(x, y). For ease of presentation, below we abstract out the conditions on the stochastic estimates ofryg, r̄f required by our analysis as the following assumption: Assumption 4 (Stochastic Gradients). For any (x, y) 2 Rdup⇥Rdlo , there exists constants f , g 0 such that the estimatesryg(x, y; ⇣), r̄f(x, y; ⇠̄) satisfy:
(i) The gradient estimate of the upper level objective satisfies:
E⇠̄ ⇥ kr̄f(x, y; ⇠̄) r̄f(x, y) B(x, y)k2 ⇤  2 f , (9)
where B(x, y) = E⇠̄[r̄f(x, y; ⇠̄)] r̄f(x, y) is the bias in estimating r̄f(x, y).
(ii) The gradient estimate of the lower level objective satisfies
E⇣ ⇥ kryg(x, y; ⇣) ryg(x, y)k 2 ⇤  2 g . (10)
As observed from Lemma 2.1, the gradient estimator (7) satisfies Assumption 4(i).
Lastly, the approximate gradient defined in (6), the true gradient (5), as well as the optimal solution of the lower level problem are Lipschitz continuous, as proven below: Lemma 2.2. [14, Lemma 2.2] Under Assumptions 1, 2 and 3, we have
kr̄f(x, y) r`(x)k  Lky⇤(x) yk, ky⇤(x1) y⇤(x2)k  Lykx1 x2k,
kr`(x1) r`(x2)k  Lfkx1 x2k, (11)
for all x, x1, x2 2 Rdup and y 2 Rdlo . The above Lipschitz constants are defined as:
L = Lfx + LfyCgxy
µg + Cfy
✓ Lgxy
µg +
LgyyCgxy
µ2g
◆ , Lf = L+ LCgxy
µg , Ly =
Cgxy
µg . (12)
The first result in (11) reveals that r̄f(x, y) approximates r`(x) when y ⇡ y⇤(x). This suggests that a double-loop algorithm which solves the strongly-convex lower level problem to sufficient accuracy can be applied to tackle (1). Such approach has been pursued in [14, 19]. Next, we propose an algorithm which rely on single-loop updates with improved sample efficiency.
3 The proposed SUSTAIN algorithm
Equipped with a practical stochastic gradient estimator for r`(x) [cf. (7)], our next endeavor is to develop a single-loop algorithm to tackle (1) through drawing O(1) samples for upper and lower level problems at each iteration. Our main idea is to adopt the recursive momentum techniques developed in [7, 37]. Notice that these works utilize unbiased stochastic gradients evaluated at consecutive iterates to construct a variance reduced gradient estimate for single-level stochastic optimization.
In the context of bilevel stochastic optimization (1), a few key challenges are in order:
Algorithm 1 The Proposed SUSTAIN Algorithm
1: Input: Parameters: { t}T 1t=0 , {↵t} T 1 t=0 , {⌘ f t } T 1 t=0 , and {⌘ g t } T 1 t=0 with ⌘ f 0 = ⌘ g 0 = 1 2: Initialize: x0, y0; set x 1 = y 1 = hf 1 = h g 1 = 0 3: for t = 0 to T 1 do 4: (y-update) Compute the gradient estimator hgt by (13) and set yt+1 = yt th g t . 5: (x-update) Compute the gradient estimator hft by (14) and set xt+1 = xt ↵th f t . 6: end for 7: Return: xa(T ) where a(T ) ⇠ U{1, ..., T}.
• Recall from Lemma 2.1 that obtaining an unbiased estimator for the outer gradientr`(x) requires using K ! 1 samples in (7), this calls for the new techniques to control the bias arising from approximating r`(x).
• The gradient estimator (7) has a more complicated structure than a plain gradient estimator, as it involves up to three different stochastic vectors/matrices related to rxf(x, y), ryf(x, y), rxyg(x, y), and one stochastic inversion that is related to [ryyg(x, y)] 1. It is not clear which are the most important objects for which variance reduction shall be applied.
Our key innovation is to develop a useful estimate of r̄f(x, y) by using a novel double-momentum technique. First, we build a recursive momentum estimator for ryg(x, y), based upon which the variable y gets updated. Then, with such a "stabilized" inner iteration, we compute an estimate of r̄f(x, y) as given in (7), by using the four stochastic vectors/matrices mentioned above but without performing any variance reduction. Such a stochastic estimator will then be used to construct a recursive momentum estimator for r̄f(x, y). The intuition is that as long as y is accurate enough, then the stochastic terms in (7) are also accurate enough, so they can be used to construct the estimator for the outer gradient. Our approach only tracks two vector estimators, while still being able to leverage the low-complexity sample-based Hessian inversion as given in (7).
The SUSTAIN algorithm is summarized in Algorithm 1. Define ⌘gt 2 [0, 1], ⌘ f t 2 [0, 1]. For the lower level problem involving y, it utilizes the following momentum-assisted gradient estimator, h g t 2 Rdlo , defined recursively as
h g t = ⌘ g tryg(xt, yt; ⇣t) + (1 ⌘ g t ) h g t 1 +ryg(xt, yt; ⇣t) ryg(xt 1, yt 1; ⇣t) ; (13)
For the upper level problem involving x, we utilize a similar estimate, hft 2 Rdup , defined as
h f t = ⌘ f t r̄f(xt, yt; ⇠̄t) + (1 ⌘ f t ) h f t 1 + r̄f(xt, yt; ⇠̄t) r̄f(xt 1, yt 1; ⇠̄t) . (14)
The gradient estimators hgt and h f t are computed from the current and past gradient estimates ryg(xt, yt; ⇣t),ryg(xt 1, yt 1; ⇣t) and r̄f(xt, yt; ⇠̄t), r̄f(xt 1, yt 1; ⇠̄t). Note that the stochastic gradients at two consecutive iterates are computed using the same sample sets ⇣t for hgt and ⇠̄t for h f t .
Both x and y-update steps mark a major departure of the SUSTAIN algorithm from existing algorithms on bilevel optimization [14, 18, 19]. The latter works apply the direct gradient estimator r̄f(xt, yt+1; ⇠̄t) [cf. (7)] to serve as an estimate to r̄f(x, y) [and subsequently r`(x)]. To guarantee convergence, these works focused on improving the tracking performance of yt+1 ⇡ y?(xt) by employing double-loop updates, e.g., by repeatedly applying SG step multiple times for the inner problem; or a sophisticated two-timescale design for the step sizes, e.g., by setting t/↵t !1.
A recent preprint [17] suggested the SVRB algorithm which applies a similar recursive momentum technique as SUSTAIN. However, SVRB is different from SUSTAIN as the momentum estimator is applied exhaustively to all the individual random quantities involved in (7) and requires Hessian projection. As a result, the SVRB algorithm entails a high complexity in storage and computation as the latter has to store matrix variables of size dlo ⇥ dlo and computes a matrix inverse for each iteration. In comparison, the SUSTAIN algorithm only requires storing the gradient estimators hgt , h f t of size dlo, dup, respectively, and the computation complexity is only O(d2loK) for each iteration.

heading: 3.1 Convergence analysis
text: In the following, we present the convergence analysis for the SUSTAIN algorithm when `(·) is a smooth function [cf. consequence of Assumptions 1, 2 and 3]. Before proceeding to the main results, we present a lemma about the Lipschitzness of the gradient estimate r̄f(x, y; ⇠̄) given in (7): Lemma 3.1. Under Assumptions 1, 2 and 3, we have for any (x1, y1), (x2, y2) 2 Rdup ⇥ Rdlo ,
E⇠̄kr̄f(x1, y1; ⇠̄) r̄f(x2, y2; ⇠̄)k  L2K kx1 x2k+ ky1 y2k 2 ,
where
LK =
s
2L2fx + 6C2gxyL 2 fy K
2µgLg µ2g +
6C2fyL 2 gxyK 2µgLg µ2g +
6C2gxyC 2 fy L2gyyK 3
(Lg µg)2(2µgLg µ2g) , (15)
and K is the number of samples required to construct the stochastic gradient estimate given in (7).
The detailed proof can be found in Appendix C. We remark that the above result is crucial for analyzing the error of the gradient estimate hft defined in (14). To see this, let us first define the errors of the gradient estimates for the outer and inner functions as follows
e f t := h f t r̄f(xt, yt) Bt, e g t := h g t r̄yg(xt, yt), (16)
where Bt := B(xt, yt) denotes the bias. Rewriting eft using (14) gives the following recursion:
e f t = (1 ⌘ f t )e f t 1 + (1 ⌘ f t ) r̄f(xt, yt; ⇠̄t) r̄f(xt 1, yt 1; ⇠̄t)
(r̄f(xt, yt) +Bt r̄f(xt 1, yt 1) Bt 1) + ⌘ft r̄f(xt, yt; ⇠̄t) r̄f(xt, yt) Bt .
Lemma 3.1 allows us to control the variance of the second term in the above relation as O(↵2t kh f t 1k 2 + 2t kh g t 1k
2). This subsequently leads to a reduced error magnitude for E[keft k2]. Similarly, we can show a reduced error magnitude for E[kegt k2] for the inner gradient estimate. The above discussion suggests that we can track the gradientr`(x) using only stochastic gradient estimates (7), without needing to track each component stochastic vectors/matrices. This allows us to avoid costly Hessian inversions. In contrast, [5, 17] track the individual stochastic vectors/matrices of (7), and then combine them together to yield an estimate ofr`(x). This approach is unable to utilize the cheap stochastic estimates of Hessian and have to invert it directly.
Turning back to the convergence analysis of the SUSTAIN algorithm, the main idea of our analysis is to demonstrate reduction of a properly constructed potential function across iterations. For smooth (possibly non-convex) objective function, this potential function consists of a linear combination of the norms of the error terms E[keft k2] and E[ke g t k
2] along with the outer objective function `(xt) and the inner optimality gap kyt y⇤(xt)k2. We obtain: Theorem 3.2. Under Assumptions 1–4. Fix T 1 as the maximum iteration number. Set the number of samples used for the gradient estimator in (7) as K = (Lg/µg) log CgxyCfyT/µg and
↵t = 1
(w + t)1/3 , t = c ↵t, ⌘
f t = c⌘f↵ 2 t , ⌘ g t = c⌘g↵ 2 t , (17)
where w, c , c⌘f , c⌘g are defined in (29) of appendix. The iterates generated by Algorithm 1 satisfy
Ekr`(xa(T ))k2 = O ✓ `(x0) `⇤
T 2/3 + ky0 y
⇤(x0)k2
T 2/3 + log(T ) 2f T 2/3 + log(T ) 2g T 2/3
◆ . (18)
Details of the constants in the theorem and its proof can be found in Appendix D. The above result shows that to reach an ✏-stationary point, the SUSTAIN algorithm requires eO(✏ 3/2) (omitting logarithmic factors) samples of stochastic gradients from both the upper and lower level functions.
This sample complexity matches the best complexity bounds for single-level stochastic optimization like SPIDER [10], STORM [7], SNVRG [42] and Hybrid SGD [37]. We claim that this is a nearoptimal sample complexity for bilevel stochastic optimization since for example, we have imposed additional smoothness conditions on the Hessian of the lower level problem. We will leave this as an open question to investigate the lower bound complexity for bilevel stochastic optimization.
Strongly-convex `(x). We also discuss the case when in addition to smoothness, `(·) is µf -stronglyconvex. Here, a stronger guarantee can be obtained:
Theorem 3.3. Under Assumptions 1–4, and suppose `(x) is µf -strongly-convex. Fix any T 1, set the number of samples for the gradient estimator (7) as K = (Lg/2µg) log C 2 gxyC 2 fy T/µ 2 g and
↵t ⌘ ↵ 
⇢ 1
µf + 1 ,
1
2µg ĉ ,
µg
ĉ L 2 g
, 1 8L2K + Lf , L 2 + 2L2y 4L2KL 2 g ĉ 2 , ⌘ f t ⌘ (µf + 1)↵, t ⌘ ĉ ↵,
where ⌘gt ⌘ 1, ĉ = 8L2y + 8L2 + 2µf/µg and LK is defined in (15). The iterates generated by Algorithm 1 satisfy for any t 1 that:
E[`(xt) `⇤]  (1 µf↵)t̄0 + 1
µf
n 2 T + ⇥ (2ĉ2 + 8ĉ 2 L 2 K) 2 g + 2(µf + 1) 2 2 f ⇤ ↵ o , (19)
where ̄0 := `(x0) `⇤ + 2f + ky0 y ⇤(x0)k2.
The detailed proof can be found in Appendix E. For large T , setting ↵ ⇣ 1/T shows that the bound in (19) decreases at the rate of O(1/T ).
Theorem 3.3 shows that to reach an ✏-optimal point, the SUSTAIN algorithm requires eO(✏ 1) stochastic gradient samples from the upper and lower level problems, also see the detailed calculations in Appendix E. This improves over TTSA [18] which requires eO(✏ 1.5) samples, and BSA [14] which requires eO(✏ 1), O(✏ 2) samples for the upper and lower level problems, respectively. Again, we achieve similar sample complexity as SGD applied on strongly-convex single-level optimization.
Interestingly, in Theorem 3.3, we have selected ⌘gt ⌘ 1 where the momentum term in the lower level gradient vanishes. In this way, the SUSTAIN algorithm is reduced into a single-momentum algorithm where the recursive momentum acceleration is only applied to the upper level gradient. Similarly, in Theorem 3.2, if SUSTAIN utilizes only the upper level momentum, i.e., ⌘gt ⌘ 1, then with appropriate choice of parameters, we get Ekr`(xa(T ))k2  O(1/ p T ) (please see [20] for further details). This implies that to achieve an ✏-stationary solution SUSTAIN with only upper level momentum requires O(✏ 2) stochastic samples for both the upper and the lower level functions. Note that this improves over TTSA [18] which utilizes a vanilla SGD update for both the upper and the lower level problems, i.e., ⌘ft ⌘ 1 and ⌘ g t ⌘ 1 and requires O(✏ 5/2) stochastic samples for both upper and lower level functions.

heading: 4 Numerical experiments
text: In this section, we evaluate the performance of the SUSTAIN algorithm on two popular machine learning tasks: hyperparameter optimization and meta learning.
Hyperparameter optimization. We consider the data hyper-cleaning task (3), and compare SUSTAIN with several algorithms such as stocBiO [19] for different batch size choices, and the HOAG algorithm in [29]. Note that in [19], the authors shown that stocBio exhibits better practical performance compared with other bilevel optimization algorithms.
We consider problem (3) with L(·) being the cross-entropy loss (i.e., a data cleaning problem for logistic regression); (x) := 11+exp( x) ; c = 0.001; see [35]. The problem is trained on the FashionMNIST dataset [41] with 50k, 10k, and 10k image samples allocated for training, validation and testing purposes, respectively. The step sizes for different algorithms are chosen according to their theoretically suggested values. Let the outer iteration be indexed by t, for SUSTAIN we choose ↵t = t = 0.1/(1 + t)1/3 and tune for c⌘f and c⌘g (see Theorem 3.2), for stocBiO and HOAG we select ↵t = d↵, t = d and tune for parameters d↵ and d↵ in the range [0, 1].
In Figure 1, we compare the performance of different algorithms when the dataset has a corruption probability of 0.3. As observed, SUSTAIN outperforms stocBiO and HOAG. We remark that HOAG is a deterministic algorithm and hence requires full batch gradient computations at each iteration. Similarly, stocBio relies on large batch gradients which results in relatively slow convergence. This fast convergence of SUSTAIN results form the single timescale update with reduced variance resulting from the double-momentum variance reduced updates.
Meta learning. We consider a few-shot meta learning problem [11, 30] (cf. (2)) and compare the performance of SUSTAIN to ITD-BiO [19] and ANIL [30]. The task of interest is 5-way 5-shot
learning and we conduct experiments on the miniImageNet dataset [39, 32] with 100 classes and 600 images per class. We apply learn2learn [1] (available: https://github.com/learnables/ learn2learn) to partition the 100 classes from miniImageNet into subsets of 64, 16 and 20 for meta training, meta validation and meta testing, respectively. Similar to [1, 19], we implement a 4-layer convolutional neural network (CNN) with ReLU activation for the learning task. At each iteration, we sample a batch of 32 tasks from a set of 20000 tasks allocated for training and 600 each for validation and testing. For each algorithm, we implement 10 inner and 1 outer update. The performance is averaged over 10 Monte Carlo runs.
For ANIL and ITD-BiO, we use the parameter selection suggested in [1, 19]. Specifically, for ANIL, we use inner-loop stepsize of 0.1 and the outer-loop (meta) stepsize as 0.002. For ITD-BiO, we choose the inner-loop stepsize as 0.05 and the outer-loop stepsize to be 0.005. For SUSTAIN, we choose the outer-loop stepsize ↵t as /(1 + t)1/3 and choose  2 [0.1, 1], we choose the momentum parameter ⌘t as c̄↵2t /2 and tune for c̄ 2 {2, 5, 10, 15, 20}, finally, we fix the inner stepsize as 0.05. For the outer loop update ANIL and ITD-BiO utilize SGD optimizer whereas SUSTAIN uses the hybrid gradient estimator.
From Figure 2 which compares the training and testing accuracy against the iteration number, we observe that SUSTAIN achieves a better performance compared to ANIL and ITD-BiO on the meta learning task. Also, notice that in the initial iterations SUSTAIN converges faster but then converges probably as a consequence of diminishing stepsizes (and momentum parameter). In contrast, ANIL and ITD-BiO slowly improve in performance and catch up with SUSTAIN’s performance. In the appendix, we show that the SUSTAIN algorithm requires less computation time to achieve better performance compared to the ANIL and ITD-BiO.
For further evaluation of the performance of SUSTAIN, we have included additional experiments on hyperparameter optimization and meta learning on different datasets in the supplementary material.
Conclusions and limitations
We have developed the SUSTAIN algorithm for unconstrained bilevel optimization with strongly convex lower level subproblems. The proposed algorithm executes on a single-timescale, without the need to use either two-timescale updates, large batch gradients, or double-loop algorithm. We showed that SUSTAIN is both sample and computation efficient, because it matches the best-known sample complexity guarantees for single-level problems with non-convex and strongly convex objective (smooth) functions, while matching the best-known per-iteration computational complexity for the same class of bi-level problems. In the future, we plan to rigorously show the sample complexity lower bounds for the considered class of bilevel problems. Further, we plan to develop sample and communication efficient algorithms for a more general class of bilevel problems, such as those with constraints in the lower level problems.

heading: Acknowledgement
text: We thank the anonymous reviewers for their valuable comments and suggestions. The work of Prashant Khanduri, Siliang Zeng, and Mingyi Hong was supported by the National Science Foundation (NSF) through grant CIF-1910385. The work of Mingyi Hong was also supported by an IBM Faculty Research award. The work of Hoi-To Wai was supported by CUHK Direct Grant #4055113. Zhaoran Wang acknowledges National Science Foundation (Awards 2048075, 2008827, 2015568, 1934931), Simons Institute (Theory of Reinforcement Learning), Amazon, J.P. Morgan, and Two Sigma for their support. Zhuoran Yang acknowledges Simons Institute (Theory of Reinforcement Learning).

Review summaries:
This paper tackles stochastic unconstrained bilevel optimization problems, where the lower level subproblem is strongly-convex and the upper level objective function is smooth. By applying momentum-assisted gradient estimator for both the upper and lower level updates, the paper proposes a new bilevel stochastic optimization method termed SUSTAIN. It has been established that the sample complexity of SUSTAIN matches the best-known complexity for single-level stochastic gradient algorithms. 

This paper investigates a bi-level stochastic optimization problem in which the upper level function depends on the minimizer of the lower level objective. A new algorithm based on stochastic momentum-assisted gradient estimator is introduced, by exploiting prior iterations gradient estimates to improve the quality of the current gradient estimation (by reducing its variance). The new algorithm matches the best complexity bounds as the optimal SGD algorithms for single-level optimization. 

The work developed a new algorithm for unconstrained bilevel optimization with strongly convex lower level subproblems. The proposed algorithm executes on a single-timescale, without the need to use either two-timescale updates, large batch gradients, or double-loop algorithm. Authors showed that SUSTAIN is both sample and computation efficient, because it matches the best-known sample complexity guarantees on single-level problems with non-convex and strongly convex objective functions, while matching the best-known per-iteration computational complexity for the same class of bi-level problems.