Contextual Stochastic Bilevel Optimization
Authors: Yifan Hu
Jie Wang
Yao Xie
Andreas Krause
Daniel Kuhn
Venue: NeurIPS 2023 poster (2023)
Keywords: stochastic optimization
bilevel optimization
contextual stochastic optimization
Multilevel Monte Carlo
TL;DR: We study stochastic bilevel optimization when the lower-level problem minimizes a conditional expectation based on some context information.

Abstract: We introduce contextual stochastic bilevel optimization (CSBO) -- a stochastic bilevel optimization framework with the lower-level problem minimizing an expectation conditioned on some contextual information and the upper-level decision variable. This framework extends classical stochastic bilevel optimization when the lower-level decision maker responds optimally not only to the decision of the upper-level decision maker but also to some side information and when there are multiple or even infinite many followers. It captures important applications such as meta-learning, personalized federated learning, end-to-end learning, and Wasserstein distributionally robust optimization with side information (WDRO-SI). Due to the presence of contextual information, existing single-loop methods for classical stochastic bilevel optimization are unable to converge. To overcome this challenge, we introduce an efficient double-loop gradient method based on the Multilevel Monte-Carlo (MLMC) technique and establish its sample and computational complexities. When specialized to stochastic nonconvex optimization, our method matches  existing lower bounds. For meta-learning, the complexity of our method does not depend on the number of tasks. Numerical experiments further validate our theoretical results.

Contextual Stochastic Bilevel Optimization
Yifan Hu" EPFL ETH Zürich Switzerland
Jie Wang Gatech United States
Yao Xie Gatech
United States
Andreas Krause
Daniel Kuhn
ETH Zürich Switzerland
EPFL Switzerland
Abstract
We introduce contextual stochastic bilevel optimization (CSBO) stochastic bilevel optimization framework with the lower-level problem minimizing an ex- pectation conditioned variable. This framework extends classical stochastic bilevel optimization when the lower-level decision maker responds optimally not only the decision of the upper-level decision maker but also to some side information and when there are multiple or infinite many followers. captures important applications such as meta-learning, personalized federated learning, end-to-end learning, and Wasserstein distributionally robust optimization with side information (WDRO-SI) Due to the presence of contextual information, existing single-loop methods for classical stochastic bilevel optimization are unable converge. To overcome this challenge, we introduce an efficient double-loop gradient method based on the Multilevel Monte-Carle (MLMC) technique and establish its sample and computa- tional complexities. When specialized to stochastic nonconvex optimization, our method matches existing lower bounds. For meta-learning, the complexity our method does not depend the number of tasks. Numerical experiments further validate our theoretical results.
contextual information and the upper-level decision
Introduction
A contextual stochastic bilevel optimization (CSBO) problem differs from classical stochastic bilevel optimization problem only in that its lower-level problem conditioned on a given context g.
min wERd>
F(x) R.g I [f v (x; [))
(upper level)
(I)
where
y" (xx;;) ;= argmin Bn [g(2, [)] V and x. (lower level)
Here { ~ Pg and " IP,je are random vectors, with P nn denoting the conditional distribution of y for a given The dimensions of the upper-level decision variable and the lower-level decision variable are dy and dy, respectively. The functions f and g continuously differentiable in (x, y) for any given sample pair ()) 17). The function f(x yi 77, {) can be nonconvex in x, but the function g(I, y: {) must be strongly convex in y for any given 17 and &. Thus, (x; the unique minimizer of the strongly convex lower-level problem for any given æ and ç. Note that, on its own, the lower-level problem be viewed as contextual stochastic optimization problem [Bertsimas and Kallus 2020) parametrized æ. We assume that the joint distribution of and 77 is unknown. However, assume that have access to any number of independent and
"oorespondence: yifan.hu@cpfi.c.
37th Conference on Neural Information Processing Systems (NcurIPS 2023).

---

identically distributed (i.i.d.) samples from Pg, and for any given realization of &. we can generate any number of i.idd samples from the conditional distribution P pnts- The bilevel structure generally makes the objective function F(x) nonconvex in the decision variable æ, except for few special cases. Thus we aim to develop efficient gradient- based algorithms for finding an e-stationary point of the nonconvex objective function F', i.e., point 5 satisfying the inequality E|) e².
CSBO generalizes the widely studied class of stochastic bilevel optimization (SBO) problems [Ghadimi and Wang 2018 whose lower-level problem minimizes an unconditional expectation.
min #
BE y (x)
where
(x) == argmin E, (9 g)
Indeed, (2) is a special case of CSBO if the upper- and lower-level objective functions are stochasti- cally independent. Another special case of CSBO is the conditional stochastic optimization (CSO) problem [Hu et al. 2020a He and Kasiviswanathan 2023 Goda and Kitade 2023) representable as
min Eg-ig (1,, Eap-Piic [b(a;; n. ():
ZER4e
Indeed, 3 is special case of CSBO if we set g( yi {) ||y h(2; 17 0)1 in which case the lower-level problem admits the unique closed-form solution y (x, {) = Er, 1,);
Applications. Despite the wide applicability of SBO to various machine learning and game theory paradigms, SBO cannot capture two important cases. The first case involves the lower-level decision maker responding optimally not only to the upper-level decision x but also to some side information & like weather, spatial, and temporal information. The second case involves multiple lower-level decision makers, especially when the total number is large. CSBO well captures these two settings and encompasses various important machine learning paradigms special cases, including meta- learning (Rajeswaran et al. 2019), personalized federated learning (Shamsian et al.; 2021 Xing
ct al., 2022, Wang et al.] 2023a]. hierarchical representation learning Yao et al. 2019 end-to -to-
end learning DDntt et al., 2017 Sadana et al.] [023] Rychener et al., [202]] Grigas et al. 2021]
Sinkhorn distributionally robust optimization (DRO) Wang et al.) [[02b]] Wasserstein DRO with
side information (Yang et al. 2022), information retrieval [Qiu et al. 20022, contrastive learning Quu
al. 2023 and instrumental variable regression [Muandet et al., 2020). Below provide a detailed
discussion of meta-learning, personalized federated learning, and end-to-end learning.
Meta-Learning and Personalized Federated Learning, Both applications can be viewed asspecial cases of CSBO. For meta-learning with M tasks or personalized federated learning with M users, the goal is to find common regularization center 0 shared by all tasks or all users. (upper level)
min EpepE [L (x),, D
where y7 (x:) = argmin, Expraio (16, Deroin) A 1196 x112]. E [M], (lower level)
2
(4)
Here, 1 is the empirical uniform distribution on [M]. The upper-level problem minimizes the generalization loss for all tasks/all users by tuning the joint regularization center æ. and the lower- level problem finds an optimal regularization parameter I; close to x for each individual task or user. Note that M may be as large as O(1²)) in meta-learming and as large as O(10)) in personalized federated learning. Thus, it is crucial to design methods with complexity bounds independent of M.
End-to-End Learning. Traditionally, applications from inventory control to online advertising involve a two-step approach: first estimating a demand function or the thickuhh rate, and then making decisions based on this estimation. End-to-end learning streamlines this into single-step method, allowing the optimization to account for estimation errors, thereby enabling more informed decisions. This can be framed as special case of CSBO, where the upper-level problem seeks the best estimator, while the lower-level problem makes optimal decisions based on the upper-level estimator and the contextual information Ç. For example, in online advertising, 2 represents the click-through rate estimator, and y" (æ; E) denotes the optimal advertisement display for a customer characterized by the feature vector &. For comprehensive review, the recent survey paper (Sadana et al. 2023).
Challenges. Given the wide applicability of CSBO, it is expedient to look for efficient solution algorithms. Unfortunately, when extended to CSBO, existing algorithms for SBO or CSO either

---

suffer from sub-optima convergence rates or are entirely unable to handle the contextual informa- tion. Indeed, a major challenge of CSBO to estimate (x; E) for (typically) uncountably many realizations of Ç. In the following, we explain in more detail why existing methods fail.
If the lower-level problem is strongly convex, then SBO can be addressed with numerous efficient single-loop algorithms [Guo and Yang 2021 Guo et al. 202 Chen et al.) 2022a 202 Hong
al. 2023 Yang al. 2021]. Indeed, as the unique minimizer (²)x) of the lower-level problem
in 2 depends only on the upper-level decision variable x, these algorithms can sequentially update
the upper- and lower-level decision variables x and y in a single loop while ensuring that the sequence {y'} approximates {v'(x") (x")). Specifically, these approaches leverage the approximation
V (xl - (0("") Vy (x') (x x²),
which is accurate if x is updated using small stepsizes. However, these algorithms generically fail to converge on CSBO problems because the minimizer y (x:;) of the lower-level problem in additionally depends on the context E. i.e., each realization of & corresponds to lower-level constraint. Consequently, there can be infinitely many lower-level constraints. It is unclear how samples from 1l corresponding to fixed context & can be reused to estimate the minimizer y" yx;;; {') corresponding to different context ç'. Since gradient-basee methods sample ç independently in each iteration f, no single sequence {y") h can approximate the function {y" (x', 3.. Guo and Yang [2021] and Hu et al. [2023] amalyze special case of the CSBO problem ([), in which is supported on M points as shown in However, the sample complexity of their algorithm grows linearly
with M. In contrast, we develop methods for general CSBO problems and show that their sample complexities are independent of the support of Ç.
SBO problems can also be addressed with double-loop stochustic grudient descent (DL-SGD), which solve the lower-level problem to approximate optimality before updating the upper-level decision variable [Ji et al. 2021 Ghadimi and Wang 2018). We will show that these DL-SGD algorithms can be extended to CSBO problems and will analyze their sample complexity as well as their computational complexity. Unfortunately, it turns out that, when applied to CSBO problems, DL-SGD incurs high per-iteration sampling and computational costs to obtain a low-bias gradient estimator for F. More precisely, solving the contextual lower-level problem to c-optimality for fixed & requires 3()=2) samples from F and gradient estimators for the function g, which leads to (c6-) total sample and computational complexity to obtain e-stationary point of F.
Methodology. Given these observations indicating that existing methods can fail or be sub-optimal for solving the CSBO problem, we next discuss the motivation for our algorithm design. Our goal is to build gradient estimators that share the same small bias as DL-SGD but require much fewer samples and incur much lower computational cost at the expense of a slightly increased variance.
To obtain estimators with low bias, variance, and a low sampling and computational cost, we propose here a multilevel Monte Carlo (MLMC) approach [Giles] 2015 Hu et al. 2021 Asi et al.) 2021], which is reminiscent of the control variate technique, and combine it with inverse propensity weighting GGlynn and Quinn 2010). We refer to the proposed method as random truncated MLMC (RT-MLMC) and demonstrate that the RT-MLMC estimator for F requires only ((1) samples
from P. nlc- This is a significant improvement vis-à-vis DL-SGD, which requires ((e-2) samples. Consequently, the sample complexity as well as the gradient complexity over g (i.e., the number g-gradient evaluations) of RT-MLMC for finding an e-stationary point of F is given by ((e-))
While the idea of using MLMC in stochastic optimization is not new (Hu et al. 2021 Asi et al.. 2021), the construction of MLMC gradient estimators for CSBO and the analysis of the variance of the RT-MLMC gradient estimators are novel contributions of this work.
1.1 Our Contributions
We introduce CSBO as a unifying framework for broad range of machine learning tasks and optimization problems. We propose two methods, DL-SGD and RT-MLMC, and analyze their complexities; see forr a summary. When specialized to SBO and CSO problems, RT-MLMC displays the same performance as the state-of-the-art algorithms for SBO [Chen et al. 2021 and CSO [Hu et al., 2021], respectively. When specialized to classical stochastic nonconvex optimization, RT-MLMC matches the lower bounds by Arjevani et al. 2023
3

---

Table 1: Complexity of the RI-MLMC and DL-SGD algorithms for finding an e-stationary point of F. The sample complexity refers to the total number of samples from Ieg as well as yye
Nonconvex CSBO
Sample Complexity
Gradient Complexity of and f Per-iteration Memory Cost
RT-MLMC
Ö(c") ") O(
O(d2 + dg)
DL-SGD
Ö(c-) | O(E-t)
O(d2 + dg)
For meta-learning with M tasks, the complexity bounds of RT-MLMC are constant in M. Thus, RT-MLMC outperforms the methods by [uuo and Yangl [2021] and Hu et al. [2023] when M is large. For Wasserstein DRO with side information Yang et al. 2022]. existing methods only cater for affine and parameeeet decision rules. In contrast, RT-MLM allows for neural network approximations. We also present the first sample and gradient complexity bounds for WDRO-SI.
For meta-learning and Wasserstein DRO with side information, our experiments show that the RT-MLMC gradient estimator can be computed an order of magnitude faster than the DL-SGD gradient estimator, especially when the contextual lower-level problem is solved to higher accuracy.
Rd and y E Rd,, we use Vıl, Vi and V2 th to denote the gradients of th with respect to (z, y). x and y. respectively. Similarly, use V²gh. Vi de and V22 to denote Hessians of i with respect to (z, y), and y, respectively. In addition, Vi2 stands for the (dæ x ) with entries 02 vh. A function 4p Rd -> TR
Preliminaries For any function & Rd x R4y with arguments
L-Lipschitz continuous if (<p(x)) - s<x')l S L||ar - x'Il for all x x' E I&d, and it is S-Lipschitz smooth if it is continuously differentiable and satisfies [[V)(()) - V<(2')l < S|I= x'Il for all xz x' E Rd. In addition, is called p-strongly convex it is continuously differentiable and if p(x) - (xx') Vi((z') (x x²) 2 +| 2²1 for all x x' Bd The identity matrix denoted by T. Finally, we use (-) as variant of the classical (.) symbol that hides logarithmic factors.
2 Algorithms for Contextual Stochastic Bilevel Optimization
Throughout the paper, make the following assumptions. Similar assumptions appear in the SBO
literature [Ghadimi and Wang 2018 Guo and Yang 2021 Chen et al. 2022a 2021 Hong et al. 2023
Assumption 2.1. The CSBO problem T satisfies the following regularity conditions:
(i) f is continuously differentiable in x and for any fixed n and &, and g is twice continuously differentiable in x and y for any fixed y and &.
(ii) g pg- pg-strongly convex in y for any fixed x, nn and E.
(iil) f. g. vf. Vg and V'g are Ly,, Lg.0. Lsn. Lg,1 and Lipsccitt continuous in (I, y) for any fixed n and respectively.
(iv) If (T,E) P(n.s) then V VI V,), is an unbiased estimator for (B (x yi O with variance of uniformly across all x and y. Also, i P 1l then Vg(a, y: n,E) {) is an unbiased estimator for VE,,E,, [l( y: n, E)] with variance o² and V2g(I, y: n, {) is an un- biased estimator for V'E [0(2, y: () with variance o3 uniformly across all x, y and g.
Assumption 2.1 ensures that problem 1 is well-defined. In particular, by slightly adapting the proofs of [Ghadimi and Wang 2018 Lemma 2.2] and [Hong et al., 2023 Lemma 2}. it allows us to show that F is -ipshihch continuous as well as Lipscchci smooth for some Sr > 0. Assumptions 2.1 (i-iii) also imply that the gradients of f and with respect (z,, y) can be interchanged with the expectations with respect to (1, {) Iy,, and 1) F Hence, Assumptions 2.1 (i-ii) readily imply the unbiasedness of the gradient estimators imposed in Assumption 2.1 (iv). In fact, only the uniform variance bounds do not already follow from Assumptions (i
In order to design SGD-type algorithms for problem T we first construct gradient estimators for F.
To this end, observe that the Jacobian V Viy" (æ;; {) Rdd.dd, exists and is Lipschitz continuous in x for any fixed & thanks to [Chen et al. 2021 Lemma 2]. By the chain rule, we therefore have
V F(x) ((.5))
Vi " (x: n. Vi E) V2l(z, of (n),))

---

Algorithm 1 EpochSGID(K, (.
Input: of epochs K, sample ç, upper-level decision initial iterate vL-
1: for = 1 K do
2: for == 0 to 2* do
3: Sample T/É from P nls and update vftt1 v/2 Bo2 1-,00~~Vg((1 2 T/.
4: end for
5: Update u = 2 2==0 vk.
6: end for
Output: vp. yk. and v
By following a similar procedure as in [Ghadimi and Wang 2018), we can derive an explicit formula for Viy" (ar; E) (for details we refer to Appendix B) and substitute it into the above equation to obtain
V (z) = (n.))) Vif (a; ()
[Enr~~aii Viggl, (): Max ht
where A(z, E) (ByyoPPis V22 y: n. Thus, the main challenges of constructing gradient estimator are to compute and store the minimizer y (x. well as the inverse expected Hessian matrix A (ar, y; E) for all (potentially uncountably many) realizations of Ç. Computing these two objects exactly would be too expensive. In the remainder of this section, thus derive estimators for y" (x; E) and (xx, y: E), and we combine these two estimators to construct an estimator for V Fxx).
Estimating y" (x; <). We estimate y (z; {) using the gradient-based method EpochSGD by Hazan and Kale| 2014), Asi et al. [202]], which involves K epochs of increasing lengths. Each epoch k ==1,, I, K starts from the average of the iterates computed in epoch k 1 and then applies 2 stochastic gradient steps to the lower-level problem with stepsize 2 - (see Algorithm 1). In the following
the output +1 of Algorithm with imputs K, x, & and yo an estimator for the minimizer (x; of the lower-level problem. We use EpochSGD for the following two reasons. First, EpochSGD attains the optimal convergence rate for strongly convex stochastic optimization in the gradient oracle model [Hazan and Kale, 2014). In addition, it is widely used in practical machine learning training procedures. Note that y (x;{) could also be estimated via classical SGD. Even though this would lead to similar complexity results, the analysis would become more cumbersome.
Estimating (ax y: 5). Following (Ghadimi and Wang 2018), one can estimate the inverse of an expected random matrix A with 0 A < using Neumann series argument. Specifically, we have
00 n
n
[EA-P, Al
Σ (I-EAPP A)
Σ
II EA.-P, An)
Σ
II
E. EA,- AP, (1 An).
21=0==1
The truncated series on the right hand side provides good approximation if N > 1. Assump- tion 2.1 (iii) implies that 0 < V229(2, yi Tu. {) < 2 Hence, the above formula can be applied
to 2bn,1 V229(1, 8i Th, E), which gives rise to an attractive estimator for A(z, y: E) of the form
if N 0,
A(z.v:)) =
2L
(5)
N
(1-21))
1 V229 x 7) 2T.211
if N 2 1.
I=
Here, N is random integer drawn uniformly from {0, NN 1} that is independent of the i..dd samples ..,,.., 7 from Py following properties. Its bias decays exponentaaly with N. its variance grows quadratically with N, and its sampling cost grows linearly with N. Below we call the approximation number.
[Cen et al. [2022b] showed that the estimator (5) displays the
Estimating Fz()) via DL-SGD. For any given K and N, we construct the DL-SGD estimator for the gradient of F by using the following procedure: (i) generate a sample from Pç., (ii) generate i..d. samples 7/ and r7" from the conditional distribution Pyls (iii) run EpochSGD as described in
Algorithm with an arbitrary initial iterate of to obtain y and (iv) construct Ä(v, v/++1: {) as in (5) Using these ingredients, we can now construct the DL-SGD gradient estimator as
g (x) = Vif(a, ii Vizg(2, vk V2) vk (6)
5

---

Algorithm 2 SGD Framework
Input: # of iterations T, stepsizes {ao)=11 initial iterate Ø1.
I: fort= = 1 to T do
2: Construct a gradient estimator v((2)) and update 2(+1 = 21 - (2((x11)
3: end for
Output: Σy uniformly sampled from "rr,....r.)..
In Lemma 2 below, we will analyze the bias and variance as well as the sampling and computational costs of the DL-SGD gradient estimator. We will see that small bias | (xx) (=))))l S
can be ensured by setting K (loo((=))) in which case EpochSGD computes { ( stochastic gradients of g. From now we refer to Algorithm 2 with v(x) = o (x) as the DL-SGD algorithm.
2.1 RT-MLMC Gradient Estimator
The bottleneck of evaluating the DL-SGD gradient estimators is the computation of vk+ The computational costs can be reduced, however, by exploiting the telescoping sum property
6K (x) = 7 +
(x) - -x)()]
(x) - (xx) = -x)(x) BRWP,
(xt) -8 - (1))
(x) +
Pk
Pk
P
k=
K replacing K, and where Igg is truncated geometric distribution with P3 (k- = k) p& x 2 for every ls = L...., K. This observation prompts us to construct the RT-MLMC gradient estimator as
where G is defined in with k l,..., K
(x)) 5' + ks
(x) 55 (x)).
The RT-MLMC gradient estimator has three key properties:
It is an unbiased estimator for the DL-SGD gradient estimator, i.e., BE [T(x)] = 7K (x))
Evaluating ((z) requires computing uh+1(2: (x, {) with probability pk. which decays exponentially K Olog(c-)), and thus pk (cc). Hence, most of the time, EpochSGD only needs to run over k < K epochs. As a result, the average sampling and computational costs are markedly smaller for RT-MLMC than for DL-SGD.
with k. To ensure small bias, we need
Since g*+1 (x) and 6 (x) differ only in v241 and ve. both of which are generated by EpochSGD and are thus highly correlated, g++1(()) (x)-6" (z:) has small variance thanks to control variate ef- fect [Nelson 1990]. Hence, the variance of RT-MLMC is well-controlled, as shown in Lemma
-0t
In Lemma 2 below, we will analyze the bias and variance as well as the sampling and computational costs of the RT-MLMC gradient estimator. We will see that it requires only O(1) samples to ensure that the bias drops to O(e). This is stark contrast to the DL-SGD estimator, which needs ((--)) samples. The lower sample complexity and the corresponding lower computational cost come at the expense of an increased variance of the order (log(( ()). The construction of the RT-MLMC
gradient estimator is detailed in Algorithm From now on, we refer to Algorithm 2 with v(z) ((x)
as the RT-MLMC algorithm.
2.2 Memory and Arithmetic Operational Costs
The per-iteration memory and arithmetic operational cost of DL-SGD as well as RT-MLMC is dominated by the cost of computing the matrix-vector product
ax, y: ==
(8)
By 5 AZz y: {) is a product of Ñ matrices of the form T- 1/1 Veng(z, ys Th {). and the n-th matrix coincides with the gradient of (y (y-117 12 nn {)) with respect to y. We can thus compute & recursively as follows. We first set v = V2f yi {). Next, we update by

---

Algorithm 3 RT-MLMC Gradient Estimator for Conditional Bilevel Optimization
Input: Iterate 2; largest epoch number K, initialization vL. approximation number N
1: Sample ç from 1e and sample N uniformly from NO...., - 1}.
2: Sample & from the truncated geometric distribution I
3: Run EpochSGD(k, x, &, v/) and obtain 1/2 o and v/-
4: Construct + (I)), T (ar), and 6¹ (x) according (6) and compute
((x) = हl (x) + P (x) o (z))
Output: ((x).
{)) To with respect to y. which is computed via automatic differentiation. This yields = (T- (1-1/(2T.,3) V22g y: n <)++)) V2f(( y; n/" E)- By using a backward recursion with respect to n, we can continue to multiply v from the left with the other matrices in the expansion of Ázz, y: E). This procedure is highly efficient because the memory and arithmetic operational cost of computing the product of Hessian matrix with constant vector via automatic differentiation is bounded up to universal constant- by the cost of computing the gradient of the same function [Rajeswaran et al. 2019]. See Algorithm in Appendix [ for details. The expected arithmetic operational costs of Algorithm is ((Nd) and the memory cost is (d).
setting to the gradient of (y (1-11/(2L)) 1V y:
3 Complexity Bounds
In this section we derive the sample and gradient complexities of the proposed algorithms. We first analyze the error of the general SGD framework detailed in Algorithm
Lemma 1 (Error Analysis of Algorithm 2}. If Algorithum 2 is used to minimize an Ly Lu-Licichitz continuous and Sippsschtttz smooth fucntion (xx) and if S 1/(2Sx), then we have
BIIV - 2 + + L ((((((,))) - Ve(a,)] + SpaBBlo((z)) - V+((+)1
where A₁ = (11) - ming +b (x).
Lemma CeemmaTsightt sightly generalizes Rakhlin et al. 2012 Ghadimi et al. 2016 Bottou al. 2018 We defer the proof to Appendix A Thus, to prove convergence to stationary point, we need to characterize the bias, variance, and computational costs of the DL-SGD and the RT-MLMC gradient estimators.
Lemma 2 (Bias, Variance, Sampling Cost and Computational Cost), We have the following results.
The biases of the DL-SGD and RT-MLMC estimators match and satisfy
||E* (x) F = [z) - VF(x) & Pg 21 (2g + + -K/2),
and the corresponding variances satisfy Var(i* (x)) O(N2) and Var(f(x)) O(KN4).
The numbers of samples and iterations needed by EpochSGD to build DI-SGD estimator are bounded by N + 2K+1 and 2 11, respectively. The expected numbers of samples and iterations needed for an RT-MLMC estimator are bounded by N + 3K and 3K, respectively.
Lemma 2 implics that setting N = ((log(E))) and K = (log(E 1)) reduces the bias to Oc)). In this case the RT-MLMC estimators have higher variances than the DL-SGD estimators, but their variances are still of the order (loggE On the other hand, using RT-MLMC estimators reduces the per-iteration sampling and computational costs from O(2*) O 2) to O(K) O(log(
Note that Hu et al. [202T] characterize the properties of general MLMC estimators and derive their complexity bounds. However, the proposed RT-MLMC estimators for CSBO problems are the first of their kind. In addition, as we need to estimate the Hessian inverse A(x, ""; E): E), our analysis more involved. In contrast to Asi et al. [222], who use MLMC techniques for estimating projections and proximal points, we use MLMC techniques for estimating gradients in bilevel optimization. The following main theorem summarizes complexity bounds.
7

---

R-6
122
K-11
Logas:: Vilow
2.0
1.3
1,3
1,0
1.4 1.5
1.4
5.5
104
4:33
5.0
Spodh
Figure 1:
Test error of DL-SGD (left figure), RT-MLMC (right figure), and MAMI. against upper-level iterations meta-learning K represents how accurately we solve the lower-level problem.
Theorem I (Complexity Bounds). If Assumplion 2.7 holds, then Algorithm 2 based the RT- MLMC or the DL-SGD estimator outputs an e-stationary point of F provided that N (log(( ')), O(e²) and T O When using the RT-MLMC estimator, the sample complexities of and 1, as well as the gradient complexities of g and f are Ö(e=t). When using the DL-SGD estimator, the sample complexity of nj and the gradient complexity of g are Ö(e=)) while and the sample complexity of and the gradient complexity of f are Ō(e-t).
K O(log( (),
Remark. Theorem 1 asserts that the sample complexity of and the gradient complexity of g much smaller for RT-MLMC than for DL-SGD, while the gradient complexities of f are comparable. When specialized to SBO or CSO problems, the complexity bounds of RT-MLMC match those of the state-of-the-art algorithms ALEST for SBO problems [Chen et al. 202T] and MLMC-based methods for CSO problems (Hu et al. 202T]. When restricted to classical stochastic nonconvex optimization, the complexity bounds RT RT-MLMC match the existing lower bounds [Arjevani et al. 2003). These observations further highlight the effectiveness of RT-MLMC across various settings.
4 Applications and Numerical Experiments
4.1 Meta-Learning
Optimization--aased meta-learning [Finn et al. 2017 Rajeswaran al. 2019] aims to find a common shared regularization parameter for multiple similar yet different machine learning tasks in order to avoid overfitting when training each task separately on their datasets that each only processes few data points. Recall Equation (4), the objective function of the optimization-basee meta-
learning Rajeswaran et al. 2019].
min Eii EE [e( D]
(9)
A
where v (x) = argmin,, Eperainap, (1), Dfrasn) + 11y: =113]. [M] and x.
where μ is the distribution over all M tasks, P; is the distribution of data from the task i, Derain and Dreat are the training and testing dataset of the task i, x is the shared parameter of all tasks and og (z) is the best parameter for regularized objective task i, 1, is a loss function that measures the average loss the dataset of the i-th task, and À is the regularization parameter to ensure the optimal solution obtained from the lower-level problem not too far from the shared parameter obtained from the upper-level problem. Note that such problem also occurs in personalized federated learning with each lower level being one user.
Note that the task distribution is usually replaced by averaging existing works [Guo and Yang 2021 Rajeswaran et al. 2019] only demonstrate a convergence rate that scales linearly with the number of tasks M. In contrast, the sample complexity of our proposed method does not depend on the number of tasks M, enabling substantially faster computation for larger M. The seminal work, Model- aggostic Meta-learning (MAML) [Finn et al.) 2017]. is an approximation of Problem (9 via replacing yf (x) with one-step gradient update, L..., (x(x)) :
all M tasks. In such cases,
-1-(, Derain).
We study the case where the loss function 4xx,, D), is a multi-class logistic loss using a linear classifier. The experiment is examined on tinylmageNet [Mnmoustafal 2017 by pre-processing

---

23
RUMMEME with Kinawwld RITM MC Kijecaall pelvren"
2.1
RT Me +id K
2311
LaggAin Los Vihee
At this paint the outer DEVIEr equais 1o4
1G5
0.25
2011
1.0
(213
5.0
300
Upora
Cuupataatio Tire (Sooada))
Figure 2: Left: Performance of our proposed RT-MLMC algorithm and BSVRB" against upper-level iterations on meta-learning. Right: Performance of algorithms against the total computational time on meta-learning.
using the pre-trained ResNet- 18 network (He et al. 2016 to extract linear features. Since the network has learned rich set of hierarchical features from the ImageNet dataset (Deng et al.) 20099. it typically extracts useful features for other image datasets. Note that each task consists of labels of similar characteristics.
Figure presents the average of logistic loss evaluated on the test dataset against the number of iterations, with each iteration representing upper-level update. From the plot, see that both DL-SGD and RT-MLMC methods tend to have better generalization performance when using a larger number of levels K. As shown in Table 2 RT-MLMC is about times faster to compute the upper-level gradient estimator than DL-SGD when K is large.
Table 2:
The computation time of DL-SGD/RT- MLMC gradient estimators on meta- leaanne...ing.
DL-SGD
RT-MLMC
K
Mean
Variance
Mean
Variance
2.65c-02
6.34e-03
2.73e-02
1.46e-02
7.23e-02
7.77e-03
3,41e-02
1.85e-02
10
2.48e-01
2.75e-02
4.93e-02
4.06e-02
12
9.38c-01
3.71c-02
1.08e-01
5,44c-02
In contrast, the MAML baseline does not have superior performance since the one-step gradient update does not solve the lower-level problem to approximate global optimality. In Appendix [ we provide numerical results for a modified MAML algorithm with multiple gradient updates, which achieves better performance compared to MAML but is still worse than proposed method.
After the initial submission of the paper, con-current work Hu et al. [0033] proposed two types of algorithms (BSVRB" and BSVRB*3) that apply to the meta-learning formulation (9). Their proposed algorithm BSVRB" is computationally expansive it requires the exact computation of the inverse of Hessian matrix (which of size 5120 x 5120 in this example) with respect to A in each iteration of the upper-level update. In the following, we compare the performance of our algorithm with their proposed Hessian-free algorithm BSVRB"2 in Figure 2
In the left plot of Figure we examine the performance of RT-MLMC method and BSVRB by rumning the same number of total epochs. It shows that RT-MLMC method has much better performance in terms of test error. In the right plot of Figure 2 we examine the performance of RT-MLMC method and BSVRB by running the same amount of computational time. Although the per-upper-level-iteee -ev-ettrtttiinoiooonono computational costs of BSVRB is small, takes a much longer time for B to achieve similar test error as RT-MLMC.
4.2 Wasserstein DRO with Side Information
The WDRO-SI Yang et al. 2022 studies robust stochastic optimization with side information Berrt simas and Kallus 2020]. Let denote the side information and 17 denote the randomness dependent on E. The WDRO-SI seeks to find parameterized mapping fæ; {) from the side information & to a decision that minimizes the expected loss w.r.t. 7) under the worst-case distributional shifts over (g. 1). Rigorously, with a penalty on the distributional robust constraint, WDRO-S admits the form
min max (BrEm)
AC(P,
(10)

---

WORD-6i
22
Testin Valie
Teaiing Loos Value
1.4
3.30
1.0 205
Spock
0.30
200
1000
420
Epech
600
800
1000
Sample Size
Figure 3: Test error on WDRO-SI against number upper-level updates for DL-SGD (left) and RT-MLMC (middle). Figure (Right) compares WDRO-SI with ERM and Wasserstein DRO that not incorporate side information. m mcans the number of samples of Z generated from Pzix for each realization of X.
where ((w;;) 17) the loss function dependent on the decision w and the random variable n. TPO is the nominal distribution of (S, n) that usually takes the form of an empirical distribution, and C(::)) a casual transport distance between distributions (Yang et al. 2022 Definition 1] variant of the Wasserstein distance that better captures the information from &. For distributionally robust feature-based newsvendor problems (Zhang et al., 2023 the covariate & can temporal, spatial, or weather information, is the random demand, f (z: denotes the ordering quantity for given and fa: f): characterizes the loss the ordering quantity fæ; {) does not match the demand
Incorporating the cost function of the casual transport distance used in (Ynn et al. 2022 and utilizing the dual form, the WDRO-SI problem (10) can be reformulated as special case of CSBO:
min Eg [1(f(x; v (x: (z: A (x; <|
(upper level)
x
where y (xxx;) argming Ba ((f(((:)) All E Ay and
(lower level)
(11)
The original work (Yang et al., 2022 only allows affine function f or non parameeretti approximation,
while our approach allows using neural network approximation such that f fx: {) is neural network parameterized by Using Theorem [ we obtain the first sample and gradient complexities for WDRO-SI. the distributionally robust feature-based newsvendor problems, we compare performance of DL-SGD and RT-MIMM. We compare with ERM and WDRO, which do not incorporate side information.
Fig. 3 (1eft) and (middle) present the results of test loss versus the number of upper-level iterations for DL-SGD and RT-MLMC, respectively. From the plot, using a larger number epochs K for the lower-level problem generally admits lower testing loss values, i.e., better generalization performance.
Fig. (right) highlights the importance of in- corporating side information as the performance of WDRO-SI outperforms the other two base- lines. In addition, more observations of for a given side information can enhance the per- formance. Table 3 reports the computational time for DL-SGD and RT-MLMC gradient esti- mators, and RT-MLMC is significantly faster since properly balances the bias-variance- computation trade-off for the gradient simula-
Table 3: Computation time of DL-SGD/RT- MLMC gradient estimators for WDRO-SI.
DL-SGD
RT-MLMC
K
Mean
Variance
Mean
Variance
1.27e-02
2.67e-03
5.04c-03
7.26e-04
5.25e-02
2.58e-03
1.25e-02
8.26e-03
1.68e-01
2.74e-03
2.02e-02
9,39e-03
4.63e-01
2.08e-03
3.41e-02
1.68e-02
tion.
Conclusion
We introduced the class of contextual stochastic bilevel optimization problems, which involve contextual stochastic optimization problem at the lower level. In addition, we designed efficient gradient-based solution schemes and analyzed their sample and gradient complexities. Numerical results two complementary applications showcase the expressiveness of the proposed problem class as well as the efficiency of the proposed algorithms. Future research should address generalized CSBO problems with constraints at the lower level, which will require alternative gradient estimators.
10

---

Acknowledgments and Disclosure of Funding
This research was supported by the Swiss National Science Foundation under the NCCR Automation, grant agreement 51NF40 180545, an NSF CAREER CCF-1650913, NSF DMS-2134033 CMMI- 2015787, CMMI-2112533, DMS- 1938106,
and the Coca-Cola Foundation.
References
Yossi Arjevani, Yair Carmon, John Duchi, Dylan J Foster, Nathan Srebro, and Blake Woodworth. Lower bounds for non-convex stochastic optimization. Mathematical Programming, 199(1-2): 165-2144 2023.
Hilal Asi, Yair Carmon, Arun Jambulapati, Yujia Jin, and Aaron Sidford. Stochastic bias-reduced gradient methods. In Advances in Neural Information Processing Systems, 2021.
Dimitris Bertsimas and Nathan Kallus. From predictive to prescriptive analytics. Management Science, 66(3):1025-1044 2020.
Léon Bottou, Frank E Curtis, and Jorge Nocedal. Optimization methods for large-scale machine learning. SIAM Review, 60(2):223-3) 2018.
Tianyi Chen, Yuejiao Sun, and Wotao Yin. Closing the gap: Tighter analysis of alternating stochastic gradient methods for bilevel problems. In Advances in Neural Information Processing Systems, 2021.
Tianyi Chen, Yuejiao Sun, Quan Xiao, and Wotao Yin. A single-timescale method for stochastic bilevel optimization. In International Conference on Artificial Intelligence and Statistics, 2022a.
Xin Chen, Niao He, Yifan Hu, and Zikun Ye. Efficient algorithms for minimizing compositions of convex functions and random functions and its applications in network revenue management. arXiv preprint arXiv:2205.0177: 2022b.
Jia Deng, Wei Dong, Richard Socher, Li-Jia Li, Kai Li, and Li Fei-Fei. Imagenet: A large-scale hierarchical image database. In JEEE Conference on Computer Vision and Puttern Recognition, 2009.
Priya Donti, Brandon Amos, and J Zico Kolter. Task-based end-to-end model learning in stochastic optimization. In Advances in Nneural Information Processing Systems, 2017.
Chelsea Finn, Pieter Abbeel, and Sergey Levine. Model-agnostic meta-learning for fast adaptation of deep networks. International Conference on Machine Learning. 2017.
Saeed Ghadimi and Mengdi Wang. Approximation methods for bilevel programming. arXiv preprint arXiv: 100::/800.2244 2018
Saeed Ghadimi, Guanghui Lan, and Hongchao Zhang. Mini- stochastic approximation methods for nonconvex stochastic composite optimization. Mathematical Programming, 155(1-2):267-305 2016.
Michael B Giles. Multilevel Monte Carlo methods. Acta Numerica, 24:259-328, 2015.
Adam N Glynn and Kevin M Quinn. An introduction to the augmented inverse propensity weighted estimator. Political Analysis, 18(1):36-5) 2010.
Takashi Goda and Wataru Kitade. Constructing unbiased gradient estimators with finite variance for conditional stochastic optimization. Mathematics and Computers in Simulation, 204:743-763, 2023.
Paul Grigas, Meng Qi, and Zuo-Jun Max Shen. Integrated conditional estimation-oppimization. arXiv preprint ww..v:22100125:5 2021.
Zhishuai Guo and Tianbao Yang. Randomized stochastic variance-reduced methods for stochastic bilevel optimization. arXiv preprint mxXiv22005.02266 2021.

---

Zhishuai Guo, Yi Xu, Wotao Yin, Rong Jin, and Tianbao Yang. On stochastic moving- estimators for non-convex optimization. arXiv preprint arXiv:2104.1484 2021.
Elad Hazan and Satyen Kale. Beyond the regret minimization barrier: Optimal algorithms for stochastic strongly-convex optimization. Journal of Machine Learning Research, 15(1):2489- 2512, 2014.
Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun. Deep residual learning for image recognition. In IEEE Conference Computer Vision and Pattern Recognition, 2016.
Lic He and Shiva Prasad Kasiviswanathan. Debiasing conditional stochastic optimization. arXiv preprint arXiv:2304. 10012, 2023.
Mingyi Hong, Hoi-To Wai, Zhaoran Wang, and Zhuoran Yang. A two-timescale stochastic algorithm framework for bilevel optimization: Complexity analysis and application to actor-critic. SIAM Journal on Optimization, 33(():147-180 2023.
Quanqi Hu, Zi-Hao Qiu, Zhishuai Guo, Lijun Zhang, and Tianbao Yang. Blockwise stochastic variance-reduced methods with parallel speedup for multi-block bilevel optimization. In Interna- fional Conference on Machine Learning, 2023.
Yifan Hu, Xin Chen, and Niao He. Sample complexity of sample average approximation for conditional stochastic optimization. SIAM Journal on Optimization, 3013):2103-2133, 2020a.
Yifan Hu. Siqi Zhang, Xin Chen, and Niao He. Biased stochastic first-order methods for conditional stochastic optimization and applications in meta learning. In Advances in Neural Information Processing Systems, 2020b.
Yifan Hu, Xin Chen, and Niao He. On the bias-variance-con tradeoff of stochastic optimization. In Advances in Neural Information Processing Systems, 2021.
Kaiyi Ji, Junjie Yang, and Yingbin Liang. Bilevel optimization: Convergence analysis and enhanced design. International Conference on Machine Leurning, 2021.
Mohammed Ali Mnmoustafa. Tiny imagenet, 2017. URL httph:ttp/// beattps://kaggle...e comppppttic/mmmmmmttiiionn//
tiny-imagenet
Krikamol Muandet, Arash Mchrjou, Si Kai Lee, and Anant Raj. Dual instrumental variable regression. In Advances in Neural Information Processing Systems, 2020.
Barry Nelson. Control variate remedies. Operations Research, 38(6):974 1990.
Zi-Hao Qiu, Quanqi Hu, Yongjian Zhong, Lijun Zhang, and Tianbao Yang. Large-scale stochastic optimization of NDCG surrogates for deep learning with provable convergence. In International Conference on Machine Learning, 2022.
Zi-Hao Qiu, Quanqi Hu, Zhuoning Yuan, Denny Zhou, Lijun Zhang, and Tianbao Yang. Not all semanties created equal: Contrastive sule-super learning with automatic temperature individualization. In International Conference on Machine Learning, 2023.
Aravind Rajeswaran, Chelsea Finn, Sham M Kakade, and Sergey Levine. Meta-learning with implicit gradients. In Advances in Neural Information Processing Systems, 2019.
Alexander Rakhlin, Ohad Shamir, and Karthik Sridharan. Making gradient descent optimal for strongly convex stochastic optimization. In International Conference on Machine Learning, 2012.
Yves Rychener, Daniel Kuhn, and Tobias Sutter. End-to-end learning for stochastic optimization: A Bayesian perspective. In International Conference on Machine Learning, 2023.
Utsav Sadana, Abhilash Chenreddy, Erick Delage, Alexandre Forel, Emma Frejinger, and Thibaut Vidal. A survey of contextual optimization methods for decision making under uncertainty. arXiv preprint arXiv:2306. 10374, 2023.
Aviv Shamsian, Aviv Navon, Ethan Fetaya, and Gal Chechik. Personalized federated learning using hypemetworks. In International Conference on Machine Learning, 2021.
12

---

Aman Sinha, Hongseok Namkoong, and John Duchi. Certifiable distributional robustness with principled adversarial training. In International Conference on Learning Representations, 2018.
Bokun Wang, Zhuoning Yuan, Yiming Ying, and Tianbao Yang. Memory-based optimization
methods for model-agnostic meta-learning and personalized federated learning. Journal of Machine Learning Research, 24:1-46, 2023a.
Jie Wang, Rui Gao, and Yao Xie. Sinkhorn distributionally robust optimization. arXiv preprint 1xXv:22009.11966 2023b.
Pengwei Xing. Songtao Lu, Lingfei Wu, and Han Yu. Big-fed: Bilevel optimization enhanced graph-aided federated learning. IEEE Transactions on Big Data, 2022.
Jincheng Yang, Luhao Zhang, Ningyuan Chen, Rui Gao, and Ming Hu. Decision-making with side information: A causal transport robust approach. Optimization Online, 2022.
Junjie Yang, Kaiyi Ji, and Yingbin Liang. Provably faster algorithms for bilevel optimization. In Advances in Neural Information Processing Systems, 2021.
Huaxiu Yao, Ying Wei, Junzhou Huang, and Zhenhui Li. Hierarchically structured meta-learning. In International Conference on Machine Learning, 2019.
Luhao Zhang, Jincheng Yang, and Rui Gao. Optimal robust policy for feature-based newsvendor. Management Science, 2023.
13

---

Appendices
A Proofs of Technical Results
Proof of L Since the function & Sy-Lipschitz smooth, we have
Bab( - ((r)))
<EVp(r)) (xre+) - 14) ..)+ + S
Bll - xx
2
(VVV(s)) v(ze)) + S 2
(2:)(2))11
S a EV (2)) (V -(zz)
- (V)())))1 -V
+ S
Spa, (xxs) Vo
(a.- -
- =(x))) + a Vo(a;)11
(a,
v(xx) [x]]]
+ Spaa 2 v
21]] + E vo(22)]1
S - |-
v(())) x/ =x1]|] a BB)((()))))))))))))))))))))))B (
where the first inequality uses Lipschitz smoothness of dh, the first equality uses the updates of the SGD algorithm, the second inequality uses the Kauchy-Schward inequality, the third equality uses the conditional expectation and the tower property, the third inequality uses the fact that as 1/(25,), and the fourth inequality uses the Cauchy-Setee Schwarz inequality.
Rearranging terms and setting a, we have
[//V(>=))] 2(Z B+((x(1))
a
+ 2E V (l - v(x)) æx]ll + 28 - Vu((a))1
Averaging from t = 1 to we have
E (+()))) 2(b(x,) min, (x))
af
2
2
+
E V 45 (z) || v(x))
Blle(x1) (x(=)|
T
2(v(11) min (x))
T ( + S.oBBl((z)) - vo(2)))?
<
+
aT
T
where the inequality holds as sb is Ly Lipscccsc continuous and thus [v(x))] & Ly for all ar.
To demonstrate the bias, variance as well as sampling and computational costs for building DL-SGD and RT-MLMC gradient estimators, we first show the iterate convergence of EpochSGD on the lower-level minimization problem with the side information. The analysis follows similarly as (Hazan)
and Kale 2014 Asi et al. 2021).
Lemma 3 (Error of EpochSGD). For given x and &, the iterates sh+1 of Algorithm with Bo = (4pg)= satisfies
b 3" (z: (1 < 22
important to note that the initial stepsize, denoted as Bo, doesn't 't necessarily have to be equal to (4m,)_1 Indeed, equivalent results can be achieved with constant 3n > 0. The choice of this
14

---

specific Sn value is primarily to streamline the analysis. In reality, there are numerous instances where the value of Hy is unknown beforehand. Under such circumstances, one practical approach could be to set 8o to a standard value, such as 0.4.
Proof of Lemma Denote Gix, E) := (z, yt 7), 7) For the case of notation, throughout the proof, E denotes taking full expectation conditioned on a given a and ç. Utilizing the update of y in the k-th epoch of EpochSGD algorithm, it holds for any y and any j 0,..., 2 that
B Elygi k +1 yll
E 3x 22 (z, u TT, 6) - yl2
= =Bll vol yll² + B3 V2g(E, vi: The 0)1 23, 29 vii (i TT... v)
<B v y|l² + 3 233 (EG((z ui: G( y: E)),
where the inequality utilizes the convexity of g and G in y. Rearranging terms, the above inequality yields
E 1
y/1
El
y13
B
EG G(
23
Summing up from = to j = 2* and dividing 2 on both sides, we obtain the relation
EG(z, i ()
2
E Blyy 4
y/1²
3x
(12)
EG(x, 5 G(.,,,,)
23,22
2
where the inequality uses the convexity of G in and Jensen's inequality. Since G(x, E) is -tronnly convex in y for any given and &, holds that
G(x,y(:)) G(a,y (z; E):
s V2G((x, (f: (w' (I; - v1) Pg |1y" (x; {) 2/112
S max { V2G(x, vP: Tty - vi) Ug 1 vill
VV2G(1, v E)1
- g.
2pg
2/1g
Next, we use induction on & to show
L2
BG v2: (m:))) <
2k
The base step for & = follows from the inequality established above. As for the induction step, suppose that
12.0
RG(x, (f G(a,y (x:
2th pg
15

---

holds for some & 2 n. Plugging 5=; into 12 we thus find
Ellv y² (a; 5)
BxLg.
BG (((1)) ((:l,,"(<x;<))))))))
23 EG(x, vh: E)
G(a, 3 (2; E):E)
BoLgg
2x+1
EG(x, v 5) G( y (x: E)
0g
2
2g
3
+
Mg502
2k+1
T22
1.2
222
Mg2h+
L
Mg 2,24
where the second inequality uses the strong convexity of in y and the fact that y (x; {) minimizes G, the first equality uses our assumption that Br 8//2* the third inequality uses the induction condition, and the second equality inequality uses the definition of fo = (4pt)- (MMg)-1. It concludes the induction. Therefore, we have
L2 11 122K+I*
EG(r, 2l
(x; S
By the My strong convexity of G(z. v. and the fact that y² (x: is the minimizer, we thus have
2 (GG(,,n44:8) GGx, (x;
I
Bllyk+1 - y" 2)12 S
n,22K
Ag
Proof of Lemma 2
We first demonstrate the properties of the RT-MLMC gradient estimator and then show that of the DL-SGD gradient estimator. To facilitate the analysis, we define
V(x)= = Pie ViS(E, (1;
(Ex Vizg(z, y" (z;
y
():
(
18 P [Az y (a
RT-MLMC gradient estimator
By the triangle inequality, we have
[ED(z) Fx))l [[ V(z) + [V(x) VF(=)].
From Chen et 2022b Lemma 2.2], we know for given x, y, and & that
Iw -P A(x, () A(z, v: 4)1 < (1
Hg N
|E
Mg 2Lg,1
Utilizing Lipschitz continuity of Vg and f. we know that
(x) [F(x)] < L..0 (1
Mgg
Mg
2L
On the other hand, by the definition of ((x), we have
[x(x) V(x)
-EV ff, ff BV
+ uk Va
2): (x 1a
16

---

By the Lipschitz continuity of f and Lemma 3
we have
EV [V, M (
5 EI,,,11lbk4+ - (x;
T/J.11.0 1.000 2g2K/2
Utilizing the Lipschitz continuity of VJ, Vg, and V in y. we have
[aT(()) V(x)II
Lgn
L5o0gg,N
5T,,1Bluk++ V (2;; +
Bluu+1 - (1; + L
BB Blybk+ (x;
Lg
N
1
+ I1 Lg N'=1
N'ElIKK v- (x::
N
L
L5.0
N2 Lg.0
+
Lgg. + L.g.oboi Lg,11 2g 222k/2
111 2222
Lgg11
-0 N
2K/2
Next, show the sampling and computational costs. To build up the RT-MLMC gradient estimator ((x), we need one sample of ç, and the number of samples of v) from pale is
N
K
2
1) T) 2-K-I
1-2-K-
< N + 3K.
N=1
On average, the iteration needed for EpochSGD is
2
(2 k+1 1) 1-2-K-1 < 3K.
Next, we demonstrate the variance of 0(z). Denote
Hx(1) Vif(z, Vif(a,ukin).. n
Hx(2) Viso(a, uk: n/ <) [(x uk: 6)
(1) Vif y (x:
H² Vi20(2, v/ (): 4) M(x, (z: (x: (
Thus one may rewrite (xx) and V(x) as
(xx) Hx+1 (1) - HA (1) (2) +2)+ Hx())
+ H1 H2)
Pk
(xx) = B|H (2)].
It holds that
E
(x)) E]
<B)
[2=)
<2 (1) 2E (1) - -(x)
<4B
(H+++ Hx( + +E
(2)+1(2) H2 (2)) - - -(())
Pk
Pk
where the first inequality holds by the definition of variance, the second inequality holds by the Sauchy-Schware inequality, the third inequality uses the
inequality and the definition
of (x)) It remains to analyze R
(H)+111) HA (1))
E pa (H)+1(2) Hg (2) and
Pa
E (1) (2) 22)
17

---

For the first term, we have
1
(H&+ 1 (1) H[](
pk
Pk 18 10 H1 (1)
K
Pk -1 T TE F1 - v212
k=1
K
Pk. 21,,1(BBlb (x(x (x: - -2213)
k=1
K
61.
Pk: 1,
1/226
k=1
6.2.g0 g.0
<KL 2
where the first inequality uses the Lipschitz continuity of Vf, the second inequality uses the Cauchy-Sehwar inequality, the third inequality uses Lemma and the last inequality uses the definition of pk-
For the second term, we may conduct a similar analysis.
1 (11x+1
Hu (2))
E
Pk
Pk H1111111 - H2((2))1
L2 13,0-3,222
Pk 6 1,2 N2 +
T. ) - -2211
K
12(
L2
L2
6 (L,,,N22
I2 12,012,222
+ 17,0.44)
Pk
1222
+
122k
12g1
k=1
36KL2
1.0 9.20 22
12
1
where the first inequality follows utilizing Lipschitz continuity of f. VS, Vg. V2g in y. the second inequality uses Lemma
3 and the last inequality uses the definition of pk-
For the third term, we have
E (1) (2) V(())
< H1 - EH 2E) (1 -2)
Notice that
E -
V EV ((r
+ +EIEV
- (x
12,2
<o² +
2/12
18

---

where the first equality holds Elja + bll = 1 + + + 2IEa b. and last inequality holds by Lemma 3 and the Lipschitz continuity of f.. On the other hand, we have
E H1 (2) (2)1
=E - EH [2 - (2)2
(2))) EH (2)12
5T, 12.00
Σ
N
+ +6(17,112
L 2,0
+ I. 1.7,0N
L2
N gL1
D2
2
N=1
N
19 1
9.1 1,00 gT.g.1
+
6
[.. 12,,,
12 No
2
2/12
where the first equality uses Elaa + =B|la|l + B||bI| by Lemma 3 and the Lipschitz continuity.
2Ea b, and last inequality holds
As result, we conclude that the variance satisfies that
(() - Ef (<x) - V(=) O(KN.)
DL-SGD gradient estimators the analysis of RT-MIMC.
Note that Rū(x) E (x)). Thus the bias follows directly from
Next, consider the per- per-iteration sampling costs and the average number of iterations for the EpochSGD. DL-SGD runs EpochSGD as a subroutine for 2 - number of iterations. Thus the sampling cost on average is N 2 1.
Consider the variance of the DL-SGD method. Note that
x (x) = (k+1(1) (2)
Following a similar decomposition as we did for the third term in bounding the variance of RT-MLMC estimators, we have
E 1 (x) E (x)]1
< T (x) (x)
k
(=))1
-
<E (k+1(1) 2E (2) (2)22
<2E| EEIK==+
2E (E(//)/)2EEEEK(((((() (1)| + 2E (2)
+ (H
(2H (2]1
17,12 13.0
N²
L
12
<207
T2? 1.
12 12 N3, N
14,
2 22
1..0 2L
12
2k+1
9.1
=O(N2 + N
where the first inequality uses the definition of variance, the second inequality uses the Cauchy- Schwarz inequality, and the third and the last equality uses Lipschitz continuity of f, f and Vg and follows similar argument as in bounding the third term for the variance of RT-MLMC estimators. Note that to control the bias, we both N and K to be of order (log(E-1)). Thus N2 is the dominating term in the variance of DL-SGD gradient estimator.
Next, we demonstrate the proof of Theorem
Proof of Theorem 7 We first demonstrate the analysis for RT-MLMC.
19

---

RT-MLMC gradient method:
combining Lemmas
and
we know that
E
((BF(()) ming F(I))
2
(2)
B
((xx)) 2x]]
aT
T
p0pa
(l) [=(2)11
2(EF((n)) min, F(x))
2 T
B
( - ((x)) æyll
aT
T
T
T
45pa
45pa
+
(l)(( v(xx)11
E
VF(x))]
T
T
t=1
t=1
1
Pg
N
KN
1
Fg
SF
+ SF
1g2 My2K/2 + aSp
+
Fg
2Lg
"2
11
2 2
To ensure that Tr is an e-stationarity point, it suffices to let the right hand side of the inequality above to (c.). Correspondingly, we set N O(log((-1)) K = O(log(<-))). a ((T-1/2),
and T = () As result, the total sampling and the gradient complexity over g are of order 3KT = ()-)). Since at each upper iteration, we only compute one gradient of f, the gradient complexity over f is T O)e
Next, we demonstrate the analysis for DL-SGD.
DL-SGD gradient method:
combining Lemmas 1 and 2
we have
E V F 1
2(EF(()) min, F(x))
2
VF (1) - JK (axe) (x()))) [x:]]
aT
T
p5pa
T
+
El ( VF(x))1
T
1=1
2(EF(s,) min, F(x))
2
+
( 8K (xx) xt]l
aT
T
t=1
4S
48FM
+
El (1x)
V(x:)11
Σ
[x)) VF(x:))1
T
T
1=1
-0
1 aT
SF
1
Pg
SF
2g22/2 aSpN2 + a
N²
1
Fig
2N
kg
2L..1
11
2Lg
To ensure that Tr is an c-stationarity point, it suffices let N ((og((c-1) 1)) K Ollo(( ?)),
= (T-1/22), and = ((e-)). As a result, the sample complexity and the gradient complexity over is of order (22 Ō At each upper iteration, we compute one gradient of f, and thus the gradient complexity over f is 7 O(6-1).
B Computing V (1y((x;)
To derive an explicit formula for V (x: we use the first-order optimality condition of the unconstrained lower-level problem. Indeed, as g(x, yi n, E) is strongly convex in y, for given & and g. y (z; E) is the unique solution of the equation
Bpate V29(2, (0: (x:E): 0.
20

---

Taking gradients with respect to on both sides and using the chain rule, we obtain
Er, V2g(( y (m; + 222 (z; () (x; = 0.
Since g(x, ys n. My -trsooyyyly convex y for any given x, 7, and &, the expected Hessian matrix Epat (2, 1:5) E - is invertible, and thus we find
VIv"(z; - (BP Vig( v² (x; 8)) (EP,Ie V229(1, (2: (xx ():.))))))
where MI Vi2g(r, ya () is in fact the transpose of Ex 221( yy ():
C Efficient Implementation for Hessian Vector Products
Our goal to compute a( y:
efficiently.
Algorithm 4 Hessian Vector Implementation for Computing 0(2, y:
Input: Iteration points (I, samples & and n/, {mm) In 7f" Pnis. N
1: Compute To (2f( ,E)
2: for = Ñ do
3: Get Hessian-vector product U, = (20 yi na, {) via automatic differentiation
4: Update rn Tn-1 Lg. Un
5: end for
6: Update r'ya Lg.1 Lo
.N.
7: Get Hessian-vector product c = Vi2((, v. ,)r's via automatic differentiation on
Output:
Remark: Step 4 and can be implemented at cost (dd) by as we compute V2g(2, yi The rn first and then differentiate over y to avoid calculating Hessian matrices directly.
D Implementation Details
In this section, we provide all relevant implementation details. We fine-tune the stepsize for all approaches using the following strategy: we pick fixed breakpoint denoted as to and adjust the stepsize accordingly. Specifically, for the t-th outer iteration when t S to, the stepsize is set as 1/V7, while for iterations beyond lo, the stepsize set as 1/L. We choose the breakpoint to such that training loss remains relatively stable the number of outer iterations approaches lo-
D.1 Meta-Learning
Let D (a, denote the training testing data, where E Idd represents the feature vector and b E [C] represents the label with C' categories. In this section, the loss (2 D) is defined as multi-class logistic loss, given by:
C(x; = -b= a log (1 eM "),
where the parameter D stands for the linear classifier, b 10 stands for the correspond- ing one-hot vector of the label The experiment utilizes the tiny LinyImaeNNet dataset, which consists of 100,000 images belonging to 200 classes. After data pre-processing, each image has dimension of 512. For dataset splitting, we divide the dataset such that 90% the images from each class are assigned to the training set, while the remaining 10% belong to the testing set. The meta-leariing task comprises tasks, with each task involving the classification 10 classes. Additionally, we set the hyper-parameter X 2.
Finally, we provide additional experiments meta-learning Figure In the left plot, we examine the performance of MAML by varying the stepsize of lnner-level gradient update from
21

---

{5e-3, 1e-2, 1e-1,ll1222 2e-11. From the plot we can tell that when using small stepsize, MAML tends have similar performance whereas the performance of MAML tend degrade when using large stepsize. In the right plot, we examine the performance of m-step MAMI. against upper-level iterations. The m-step MAMI, is an approximation of problem via replacing Ve(ar) with the m-step gradient update Fs. (z), which is defined recursively:
(1,(()) z, (1,(x)) Fi ~ ((05x-1(x)) 1 + 2 - x!|?
with stepsize and Di ~ Pi for ls = 1,... m. Here we take the number of gradient updates at inner level m from f1,4, 8, 12]. From the plot, we realize that multi-step MAML tends to have better performance, but still cannot outperform the RT-MLMC method.
2.2
2.2
XAMM.
XCAME
Jotep
MAMME (step
XEAME
XIMML pale
2.0
MAMML palcy die 2613
2.0
IEE-MELNCC with Kenase $2
Logatie Loss Value
RE-SEMSC
Logatic Lots Value
1.8
1.8
1.6
1.6
1.4
1.4
163
663
3e3
153
663
363
Epodi
Epoch
Figure 4: Left: Performance of one-step MAML against upper-level iterations on meta-learning. Right: Performance of m-step MAML against upper-level iterations on meta-learning.
D.2 Wasserstein DRO with Side Information
In this subsection, take the loss function
1(w, 17) h(w - 1)+ + b(1 w)=
where h 0 a constant representing the per-unit holding cost and b is constant representing the per-unit backlog cost. Since Assumption 2.] does not hold for the objective function due to its non-smooth structure, approximate the loss function with the smoothed version
h lu n) 3 log(1 el(
b
3 log(1 eb(n-w)),
where we specify the hyper-parameter 3 5 to balance the trade-off between loss function approxi- mation and smoothness. The synthetic dataset in this part generated the following procedure: the covariate & is sampled from the .00-dimensional uniform distribution supported on [-15, 15] 100 The demand 1 depends on & in a nonlinear way:
= (x: +e, e~ C~,,0, 1).
fNN( {) 10 Sigmoid(Ws, ReLU ReLU + b1) b2) + by):
where the neural network parameter v (W,, by,........ W3, by). In particular, the ground-truth neural network parameter is generated using the uniform initialization procedure by calling torch.mn init. uniform_ in pytorch. We quantify the performance of given e using the testing loss -(c,n))>-- = (x]f(x;( 7)], where the expectation is estimated using sample average approximation
based on 2 10 sample points. When creating training dataset, we generate M = 50 samples of
denoted as [EIME and for each æ;, generate E {10, 30, 50, 100) samples of from the conditional distribution P When generating the left two plots of Fig. 3 we specify m 100.
When solving the WDRO baseline, consider the formulation
min max (l E): m)) AM((PFF) }
(13)
22

---

with -(..) being the standard Wasserstein distance using the same cost function the casual transport distance. We apply the SGD algorithm developed in (Sinha et al. 2018 to solve the WDRO formulation. See Algorithm 5 for detailed implementation. The hyper-parameter À for WDRO-SI or WDRO formulation has been fine-tuned via grid search from the set (1, 10, 50, 100, 150} for optimal performance.
Algorithm 5 Gradient Descent Ascent Heuristic for Solving (13)
Input: of outer iterations and & of inner iterations Tin, stepsizes {ax)(=1 and (3,)2n initial iterate T1-
1: for = 1 to T do
2: Sample (5', rf²') from P and initialize = 5t.
3: for s 1 to Tin do
4: Generate 9s V D (f(; (1) AIEE 5913
5: Update 5. E + B,Adam(E', 9o)
6: end for
7: Compute G. = Valla( +)))))1)))
8: Update It+1 = It - a,G,
9: end for
Output: 2T
23


Reviews:
[{"review_id": "sd6qqxhJhi", "reviewer": "Reviewer_BuTe", "summary": "The paper investigates a generalization of the stochastic bilevel optimization model, in which the lower and upper optimization levels share a random variable. The authors design two gradient-based approaches named RT-MLMC and DL-SGD for solving this problem and analyze their performance. The two methods differ in the way they estimate the gradient of the objective: DL-SGD is rather straightforwardly based on existing results, while RT-MLMC improves upon the performance of DL-SGD by a carefully tailored sampling method. Finally, numerical examples are presented to demonstrate the generality and performance of the proposed model and methods.\n", "questions": "1. Is the dependence by $f$ on $\\eta$ in the upper level of (1) necessary? I believe that removing it does not hurt the generality of the model and may avoid some confusion (see typos below).\n\n2. In Lemma 1, $\\alpha$ should be defined.\n\n3. Typos:\n* In (5), $\\eta'$ is used as the argument for the $\\nabla_1 f$ and $\\nabla^2_{12}g$ terms while $\\eta''$ is used in the $\\nabla_2 f$ term. This is in contrast with the second formula for $\\nabla F$ on Page 4, where an independent RV is used in $\\nabla^2_{12}g$.\n* A similar typo appears on the bottom of Page 15, in the proof of Lemma 2, where $\\eta$ in the term $\\nabla^2_{12}g$ should be changed to $\\eta'$ (as in the definition of $V(x)$ above).\n* A similar typo appears on Page 16, where $H_K(1)$ and the last term in $H_K(2)$ should share the same variable $\\eta$ (compare the formula for $V(x)$ on page 16 versus the definition of $V(x)$ on Page 15).\n\n", "limitations": "N/A", "rating": 7, "confidence": 3, "soundness": 3, "presentation": 4, "contribution": 4, "strengths": "* The paper introduces a general model that is applicable in a wide range of situations. Although the model is more general than previous approaches, the proposed approximations appear to be as efficient as the methods developed for the special cases.\n\n* The paper is clear, focused, and well-written.\n\n* I have verified most of the math; it is easy to follow and, aside from a few typos, appears to be correct.\n", "weaknesses": "There are several typos in critical parts of the proofs (see \"Questions\" below). I believe that these typos do not affect the correctness of the proofs, however, this somewhat lowers my confidence in the results.\n", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your comments. Below, we address your comments.\n\n1. [Dependence on $f$ and $\\eta$] Indeed, the dependence between $f$ and $\\eta$ is not necessary. In many cases, we can remove $\\eta$ in $f$ and remove the expectation over $\\eta$ in the upper level. We are writing in the current way to be as most general as possible. Thanks for pointing it out!\n2. [Define $\\alpha$ in Lemma 1] We have fixed that.\n3. [Typos] Thanks for checking our paper carefully. Indeed, to align with the second $\\nabla F$ shown on page 4, one should use $\\eta^\\prime$ in $\\nabla_{12} g$ and use $\\eta^{\\prime\\prime}$ in $\\nabla_1 f$ and $\\nabla_2 f$ in Equation (5). We have made the modifications according to your suggestion. Note that the current Equation (5) is also valid because $\\nabla_{12} g$ uses a different sample compared to $\\nabla_2 f$. Since Equation (5) is of the form $\\nabla_1 f(\\eta^\\prime) - \\nabla_{12} g(\\eta^\\prime) \\Lambda \\nabla_2 f(\\eta^{\\prime\\prime})$ (where we omit other dependence), after taking the full expectation, it still aligns with the second $\\nabla F$ shown on Page 4. The key part is that we should not use the same sample of $\\eta$ for $\\nabla_{12} g$ and $\\nabla_2 f$, which will lead to correlation issues. The other two places are typos. Thanks for pointing it out.\n\nWe hope that the clarification could increase your confidence in our work. "}, {"review_id": "GGJLkGy5OX", "reviewer": "Reviewer_Z7Go", "summary": "**Summary:** The paper introduces a novel algorithm for solving contextual stochastic bilevel optimization (CSBO) problems. The authors develop DL-SGD and RT-MLMC gradient estimators for the problem along with a SGD-based algorithm for solving the CSBO problem. The authors analyze the properties of the gradient estimators and provide finite-time convergence guarantees for the proposed method. The theoretical results are corroborated by experiments on MAML, DRO with side information, and IV regression problems. ", "questions": "Please see the limitations above. ", "limitations": "n/a", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "**Strengths:** \n\n-\tThe problem formulation considered in the paper is challenging and is of significant interest to the ML community. \n-\tThe paper is well written, and the ideas are presented clearly with discussions. \n-\tThe authors provide theoretical guarantees for the proposed approaches with an analysis of the proposed gradient estimators. \n-\tThe authors have conducted experiments on multiple tasks including MAML, DRO with side information, and IV regression problem to evaluate the performance of the proposed framework. \n", "weaknesses": "**Weaknesses:**\n\n1. A major confusion I have is about the unbiasedness of the gradient estimator stated after line 127. Specifically, the expressions $\\nabla_1 y^\\ast$ and $\\nabla_2 f$ both depend on the random variable $\\xi$. This implies that the two expressions are dependent on each other which further means that the expectation given in the equation (after line 127) will not be the same as $\\nabla F$. Can the authors clarify why the given expression will be true? The same discussion holds for the approximate gradient expressions derived in eq (5) and the rest of the paper. This is a major issue since the proofs and the results are based on the independence of $\\nabla_1 y^\\ast$ and $\\nabla_2 f$. \\\n          If this issue is resolved, I am willing to raise my score. \n\n2. The authors should clarify the intuition behind using EpochSGD to estimate $y^\\ast$. In the discussion, the authors simply explain the algorithm without explaining the intuition. \n3. The gradient inverse estimator utilized from Ghadhimi’s work is clear, however, why the iterative estimator RT-MLMC works is difficult to understand. Can the authors please explain the working of the proposed estimator, i.e., why it works?\n4. The experiments on MAML and DRO do not compare the proposed approach with baseline algorithms in the area.\n\n    - Numerous bilevel algorithms solve MAML, the authors should compare the performance of the proposed approach against at least a few of them. \n                              \n    - Similarly, the authors need to show the performance of the proposed approach for solving the DRO problem against popular baselines. \n\n5. It would be easier if the authors keep the notations consistent in the experiment and the theory section of the paper.  \n\n6. To motivate the problem better the authors should include some examples in the introduction section. This way it will be easier for the reader to appreciate the considered formulation. \n\n----\nUpdated the score after the rebuttal. \n", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thanks for your comments. Below, we address your comments in detail.\n\n1. [Form of gradient estimator $\\nabla F(x)$] Indeed $y^*(x;\\xi)$ and $\\nabla_2 f(x,y^*(x;\\xi);\\eta,\\xi)$ both depend on $\\xi$. However, this does not prevent us from obtaining the expression for the gradient. Note that:\n$$F(x)=\\mathbb{E}[f(x,y^*(x;\\xi);\\eta,\\xi)]=\\mathbb{E}\\_\\xi\\mathbb{E}\\_{\\eta\\mid\\xi}[f(x,y^*(x;\\xi);\\eta,\\xi)].$$\nBy the fact that one can switch expectation and taking gradient under given assumptions, we have:\n$$\\nabla F(x)=\\mathbb{E}\\_\\xi\\nabla_x[\\mathbb{E}\\_{\\eta\\mid\\xi}[f(x,y^*(x;\\xi);\\eta,\\xi)]],$$\nNote that for a given $\\xi$, the gradient formulation follows exactly from that of the classical stochastic bilevel optimization without dependence structure using chain rule. See, for instance, [Ghadimi and Wang, 2018] for more details.\n$$\\nabla_x[\\mathbb{E}\\_{\\eta\\mid\\xi}[f(x,y^*(x;\\xi);\\eta,\\xi)]]=  \\mathbb{E}\\_{\\eta\\mid\\xi}[\\nabla_x  f(x,y^*(x;\\xi);\\eta,\\xi)]=\\mathbb{E}\\_{\\eta\\mid\\xi}[\\nabla_1 f(x,y^*(x;\\xi);\\eta,\\xi) + \\nabla_1 y^*(x;\\xi)^\\top \\nabla_2 f(x,y^*(x;\\xi);\\eta,\\xi)].$$\nFurther taking expectation with respect to $\\xi$ gives the expression of  $\\nabla F(x)$ after Line 127, i.e.,\n$$\\nabla F(x)=\\mathbb{E}\\_\\xi\\mathbb{E}\\_{\\eta\\mid\\xi}[\\nabla_1f(x,y^*(x;\\xi);\\eta,\\xi)+\\nabla_1 y^*(x;\\xi)^\\top\\nabla_2 f(x,y^*(x;\\xi);\\eta,\\xi)].$$\nPlugging in $\\nabla_1y^*(x;\\xi)$ obtained in Appendix B, we have the expression of $\\nabla F(x)$ after Line 129.  \n2. [Intuition to use EpochSGD] There are two reasons for using EpochSGD instead of classical SGD. Firstly, EpochSGD is faster than SGD by a logarithmic factor and is optimal in terms of sample and iteration complexity (see Hazan and Kale, 2014). Additionally, MLMC-based methods use the control variate technique between two neighboring approximations. If the difference between two neighboring approximations is too small, the approximation error decays too slowly. If the difference is too large, the variance reduction effect achieved by the control variate is not strong enough. To achieve a balance, MLMC usually uses a sequence of approximations that admits exponentially decaying approximation error. In this regard, the error of EpochSGD decays exactly exponentially in the number of epochs. Therefore, we can use the output at the end of each epoch for MLMC and do not need to manually select at which iteration of classical SGD we use the corresponding $y$ in the construction of MLMC (see Asi et al., 2021 for more discussion on EpochSGD).\n3. [Motivation of RT-MLMC estimator] Although DL-SGD is low-biased, the complexity of DL-SGD is too high. Thus we want to construct an unbiased estimator of DL-SGD at an even lower cost to achieve optimal complexity bounds. Note that even for stochastic bilevel optimization, existing double-loop methods cannot achieve optimal complexity.\n    - To ensure low bias, we use the equation after Line 163 and show that the RT-MLMC gradient estimator in Eq. (6) to be an unbiased estimator of DL-SGD in Eq.(5) and thus low-bias for the overall objective.\n    - To achieve low costs, RT-MLMC uses a randomized approach by incorporating a truncated geometric distribution $p_k\\propto 2^{-k}$ to generate an approximation level $k$. With a high probability, RT-MLMC generates a small $k$ and constructs $\\hat{v}^{k+1}$ and $\\hat{v}^k$, which only requires computing $y_{k+1}^0$ and $y_{k}^0$. This is inexpensive because $k$ is very small. With a low probability, RT-MLMC generates a large $k$ and computes $\\hat{v}^{k+1}$ and $\\hat{v}^k$, which is very expensive. The expected computational cost is thus mild by summation over high costs multiplied by low probability + low costs multiplied by high probability.\n    - Lastly, one might argue that dividing $p_k$ for small $k$ can cause high variance. This is mitigated as RT-MLMC incorporates a control variate technique by subtracting $\\hat{v}^{k+1}$ and $\\hat{v}^k$ that are highly correlated.\n    - In summary, RT-MLMC is unbiased of DL-SGD, thus admitting low-bias for $\\nabla F(x)$; it has much lower costs than DL-SGD and admits mild variance.\n\n4. [Compare to baseline in experiments] Note that even for a special case of CSBO, i.e., Problem (8) that is a stochastic bilevel optimization with multiple lower-level problems, many baseline approaches in bilevel optimization literature are not applicable. For meta-learning in form of Problem (8), many existing methods actually only solve a surrogate of the formulation by replacing the lower-level problem with gradient updates to obtain one-step or multi-step MAML or by averaging all lower-level problems so that there is only one lower-level constraint. However, the original formulation actually finds a lot important applications, such as personalized federated learning. To directly solve Problem (8), we provide detailed comparisons with baseline methods including BSVRB [Hu et al., 2023 Blockwise stochastic variance-reduced methods with parallel speedup for multi-block bilevel optimization. arXiv preprint arXiv:2305.18730, 2023] (the STOA that appears even after the NeurIPS deadline) and one-step/multi-step MAML in the general response with figures in the PDF file. In short, RT-MLMC outperforms these methods under various experimental setups. \n     - For Wasserstein DRO with side-information, the existing method in [Yang et al., 2022] heavily relies on convexity, linearity predictors, and linear loss assumptions to build reformulations that can be solved by convex solvers. However, when the hypothesis class from the covariate to the decision is parameterized by nonconvex neural networks, their method is not implementable. In contrast, our algorithm is the first implementable algorithm with a convergence guarantee. \n5. [Consistent notation in theory and experiments] We will keep the notation in the theory and the experiments consistent. \n6. [Include examples in the Introduction] We have added the applications into the introduction for a better illustration."}, {"review_id": "21DGxFiHt2", "reviewer": "Reviewer_NrrF", "summary": "Contextual stochastic bilevel optimization (CSBO) is introduced in this paper. An efficient double-loop gradient method based on the Multilevel Monte-Carlo (MLMC) is proposed. The proposed framework captures important applications such as meta-learning, Wasserstein distributionally robust optimization with side information (WDRO-SI), and instrumental variable regression (IV).", "questions": "1. I think it's necessary to emphasize the difficulty of solving the Contextual Stochastic Bilevel Optimization Problem compared with traditional bilevel optimization problems. The presentation of this work is poor. I believe this is an excellent work, I suggest that you should modify the presentation of this work to better clarify the contributions of this work.\n\n2. In the experiment, the results are limited. For example, in the meta-learning application, I suggest the authors compare the proposed method with the state-of-the-art bilevel optimization methods [1][2][3], which are shown to be able to address the meta-learning task. Alternatively, you need to specify why these methods [1][2][3] are not applicable to this application. Furthermore, the authors should conduct experiments on more datasets to better evaluate the proposed method, for example, Omniglot dataset.\n\n3. I suggest the author briefly introduce some existing bilevel optimization works in machine learning, for example, hyper-gradient-based methods [1, 4] and approximation-based methods [2], and then discuss why these methods fail to be applied to the Contextual Stochastic Bilevel Optimization Problem.\n\n[1] Bilevel optimization: Convergence analysis and enhanced design, ICML, 2021\n\n[2] Asynchronous Distributed Bilevel Optimization, ICLR 2023\n\n[3] Bilevel Programming for Hyperparameter Optimization and Meta-Learning, ICML 2018\n\n[4] Provably faster algorithms for bilevel optimization, NeurIPS 2021", "limitations": "NA", "rating": 5, "confidence": 2, "soundness": 2, "presentation": 1, "contribution": 2, "strengths": "An interesting problem, that is, Contextual Stochastic Bilevel Optimization Problem is proposed in this work. And an efficient algorithm is proposed to solve this problem. The proposed framework captures important applications such as meta-learning, Wasserstein distributionally robust optimization with side information (WDRO-SI), and instrumental variable regression (IV). However, I have some concerns as follows.", "weaknesses": "1. I think it's necessary to emphasize the difficulty of solving the Contextual Stochastic Bilevel Optimization Problem compared with traditional bilevel optimization problems. The presentation of this work is poor. I believe this is an excellent work, I suggest that you should modify the presentation of this work to better clarify the contributions of this work.\n\n2. In the experiment, the results are limited. For example, in the meta-learning application, I suggest the authors compare the proposed method with the state-of-the-art bilevel optimization methods [1][2][3], which are shown to be able to address the meta-learning task. Alternatively, you need to specify why these methods [1][2][3] are not applicable to this application. Furthermore, the authors should conduct experiments on more datasets to better evaluate the proposed method, for example, Omniglot dataset.\n\n3. I suggest the author briefly introduce some existing bilevel optimization works in machine learning, for example, hyper-gradient-based methods [1, 4] and approximation-based methods [2], and then discuss why these methods fail to be applied to the Contextual Stochastic Bilevel Optimization Problem.\n\n[1] Bilevel optimization: Convergence analysis and enhanced design, ICML, 2021\n\n[2] Asynchronous Distributed Bilevel Optimization, ICLR 2023\n\n[3] Bilevel Programming for Hyperparameter Optimization and Meta-Learning, ICML 2018\n\n[4] Provably faster algorithms for bilevel optimization, NeurIPS 2021", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your comments. Below, we address your questions.\n1. [Difficulty for solving CSBO] As discussed in the Introduction from Line 43 to Line 68, existing algorithms for traditional bilevel optimization problems with one single lower-level problem either cannot achieve optimal complexity bounds or does not apply to the CSBO problem due to potentially infinitely many lower-level problems parametrized by $\\xi$, each of which introduces a constraint involving solving a stochastic optimization problem. See more details in below. We highlight that the proposed RT-MLMC achieves the optimal complexity bounds for CSBO problems. It is also the first double-loop algorithm that achieves the optimal complexity bounds for classical stochastic bilevel optimization. \n    - Existing double-loop methods for classical bilevel optimization can be applied with a slight modification but admits a sub-optimal complexity bound, as we have shown in DL-SGD. \n    - As for the more widely investigated single-loop methods for classical bilevel optimization, these algorithms are specifically designed for the setting when there is only one lower-level constraint formulated as $y^*(x)$. In such a case, one can use a sequence of vectors $y^t$ to approximate $y^*(x^t)$ leveraging the intuition that $y^*(x^{t+1}) - y^*(x^t) \\approx \\nabla y^*(x^t) (x^{t+1} - x^t)$. However, when there are multiple lower-level constraints, i.e., $y^*(x;\\xi)$ for each realization of $\\xi$, we cannot no longer use only one sequence of $y^t$ to keep track of $y^*(x^t,\\xi^t)$ as $\\xi^t$ is randomly sampled at each iteration. For example, when $\\xi$ follows a normal distribution, it requires infinite number of sequence of $y$ to keep track of the lower level constraints.\n   - To our best knowledge, [Guo et al., 2021] is the only work that discussed the case when there are $M$ lower-level constraints, and they adopt $M$ sequences of $\\{y^t(i)\\}$ to keep track of $y^*(x;i)$  for $i\\in[M]$, respectively. Thus, the convergence rate of their algorithm is depends linearly on $M$. On the other hand, our proposed RT-MLMC algorithm gets rid of the dependence on the number of lower-level constraints and achieves an optimal complexity bound of $O(\\epsilon^{-4})$. This is particularly important for many applications that are special cases of the CSBO, including personalized federated learning [Xing, et al. \"Big-fed: Bilevel optimization enhanced graph-aided federated learning.\" *IEEE Transactions on Big Data* (2022)]. In this context, $M$ can be of size $O(10^9)$, meaning that each lower-level constraint represents a personalized keyboard usage preference. \n2. [Comparison to baselines in meta-learning] We respectfully point out that the references mentioned [1-4] address the traditional bilevel optimization problem, where there is only one lower-level constraint. When applied to meta-learning applications, they do not treat all tasks as multiple individual lower-level constraints. Instead, they use a surrogate aggregated lower-level objective to enforce one lower-level problem, for example, via averaging. In other words, they do not solve the bilevel optimization formulation for meta-learning proposed in [Rajeswaran et al., 2019] and studied in our paper. It is unclear how the proposed methods in [1-4] can be applied to the meta-learning formulation studied in our paper. Note that the meta-learning formulation in our paper is also used in personalized federated learning [Xing, Pengwei, et al., 2022]. If one adapts the surrogate aggregated lower-level constraint, it is not personalized at all. It implies that importance of solving the formulation (8) considered in our paper. Yet no algorithms in [1-4] can be applied.\n    - In the general response (also see the added PDF file), we have added comparison to two recent papers for solving a special of CSBO when there are $M$ lower-level problems [Guo and Yang, 2021] and [Hu et al. 2023, Blockwise stochastic variance-reduced methods with parallel speedup for multi-block bilevel optimization. arXiv preprint arXiv:2305.18730, 2023] (the second paper appeared on ArXiv after the NeurIPS submission deadline). The algorithm in [Guo and Yang, 2021] and the first algorithm in [Hu et al. 2023], both require computing the inverse of Hessian exactly in each iteration, which cannot be implemented efficiently especially for high-dimensional problems. We use the BSVRB, the second algorithm in [Hu et al. 2023], as a baseline to solve the special case of CSBO formulation. The comparison is in the general response with supporting Figure 5 in the PDF file. Note that the performance of BSVRB is worse than RT-MLMC. The reason is that the iteration complexity of the baseline BSVRB depends linearly on the number of lower-level problems, whereas the proposed RT-MLMC does not and achieves optimal complexity bounds.\n   -  We note that meta-learning with Omniglot dataset is a high-dimensional bilevel optimization problem with tremendously many lower-level problems, which is difficult to solve given the short rebuttal time period. We promise to add experiments either using this dataset or other types of datasets such as UCI Adult benchmark dataset *a8a* and web page classification dataset *w8a* in our revised paper.\n\n3. [More reference] Thanks for the reference. We have added a related discussion. The reasons why these algorithms fail for CSBO are discussed in the first bullet point.\n\nXing et al. \"Big-fed: Bilevel optimization enhanced graph-aided federated learning.\" IEEE Transactions on Big Data (2022).\n\nHu et al. \"Blockwise stochastic variance-reduced methods with parallel speedup for multi-block bilevel optimization\". arXiv preprint arXiv:2305.18730, 2023"}, {"review_id": "hxbq3sHr3s", "reviewer": "Reviewer_DnyS", "summary": "This work studies the so-called contextual stochastic bilevel optimization (CSBO) problem, in which the lower-level problem is a conditional expectation problem under some contextual information. Some applications in distributional robust optimization fall into this category. The paper proposes a double-loop Monte-Carlo based method, which leverage EpochSGD in Hazan and Kale, 2014 to approximate  the lower-level solution. A approximate hypergradient is provided based on an explicit form via implicit function theorem. Experiments on meta-learning and instrumental variable regression are provided.", "questions": "See Weakness part. ", "limitations": "See Weakness part. ", "rating": 5, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "1. This work considers a problem that is different from existing studies by introducing the contextual information. It covers some important cases in, e.g., DRO.\n\n2. The algorithms are reasonable to me. Using MLMC to further improve the performance of DL-SGD is good. Theoretical complexity and convergence are analyzed. Experiments seem to support the design principles. ", "weaknesses": "1. The studied problem seems a little bit artificial. The example in (3) seems to artificially change a single-level problem into a bilevel one. In addition, the hyper-gradient form is almost the same as in the non-contextual case and hence the tools therein may be used with some adaptations. \n\n2. Some important components such as MLMC, EpochSGD, Neumann series expansion are existing techniques. Some challenges like variance control (e.g., in $\\hat v$) and hyper-gradient computation can be well coped with by these techniques. Thus, the novelty is not that significant. \n\n3. In experiments, in Fig. 1, why does MAML stop at this large loss value? Is it because of a non-satisfactory hyperparameter tuning or something else. It seems that the stepsizes for MAML may be chosen to be too large. The comparison seems not fair. \n\n\n4. The experiments do not show the important of bilevel optimization. For example, no baselines other than the proposed CSBO solvers are provided. Some baselines (e.g., single-level ones, or some standard baselines) in Wasserstein DRO may need to be included. \n\nOverall, I am not fully convinced that this is an important problem, and given the above concerns, I lean toward the negative side. ", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for the comments. In below, we address your concerns.\n1. [Motivation of CSBO] The CSBO problem is not artificially made up. In addition to optimization-based meta-learning and Wasserstein DRO with side information (also known as the causal optimal transport) mentioned in the numerical study, there are other applications that are special cases of the CSBO problems but are not special cases of Equation (3).\n    - Personalized federated learning [Xing, Pengwei, et al. \"Big-fed: Bilevel optimization enhanced graph-aided federated learning.\" *IEEE Transactions on Big Data* (2022)] is a special case of CSBO, see Equation (2) in the reference. It is important to note that in personalized federated learning, the number of lower-level constraints can be on the order of $O(10^9)$ as each constraint represents one person. Therefore, the RT-MLMC method, which has convergence independent of the number of lower-level constraints, is crucial for achieving efficiency. Convergence of existing methods in [Guo et al., 2021] and [Hu et al., 2023 Blockwise stochastic variance-reduced methods with parallel speedup for multi-block bilevel optimization. arXiv preprint arXiv:2305.18730, 2023] all depend linearly on $M$.\n    - End-to-end learning/Contextual Optimization [Sadana, Utsav, et al. \"A Survey of Contextual Optimization Methods for Decision Making under Uncertainty.\" *arXiv preprint arXiv:2306.10374* (2023).] In this survey, there are three paradigms for studying end-to-end learning/contextual optimization. The third paradigm, which integrates learning and optimization, falls as a special case of CSBO problems (see Figure 3 in Sadana et al., 2023). Also, see Table 3 on Page 19 in in Sadana et al., 2023 for a list of more than 20 papers that fall into this category. Our work is the first to provide non-asymptotic optimal complexity bounds for such a problem.\n    \n    Indeed, we find that contextual stochastic bilevel optimization has various applications in machine learning and optimization. \n    \n2. [Novelty] Note that the contextual information $\\xi$ introduces additional challenges in designing an efficient optimal algorithm. In particular, most of the existing single-loop algorithms for classical bilevel optimization do not apply to our problem due to potentially infinitely many lower-level problems parametrized by $\\xi$, each of which introduces a constraint involving solving a stochastic optimization problem. Additionally, we have shown that the double-loop algorithm DL-SGD is not optimal and have proposed an optimal algorithm RT-MLMC. The closest work to ours [Guo et al., 2021] can only handle a finite number of lower-level constraints ($M$), and the complexity bounds depend linearly on $M$. Our algorithm does not depend on $M$ and sheds light on how to solve stochastic bilevel optimization with even infinitely many lower-level constraints. \n3. [Comparison to MAML] We respectfully point out that the step size of MAML has been fine-tuned in our experiment. In our general response (see newly added Figure 4 in the PDF file), we report the plot of MAML performance for different choices of step size from the list {5e-3, 1e-2, 5e-2, 1e-1, 2e-1}. From the plot we realize that for small step sizes the MAML tends to have similar performance, whereas MAML tends to diverge for too large step sizes. The poor performance of MAML on meta-learning problem (8) is that it solves a different formulation, i.e., it replaces the lower level problems with one-step gradient update. It is natural that MAML cannot achieve good performance on the CSBO objective, as the approximation gap is theoretically $O(1)$ unless one performs multi-step MAML.\n    \n    - Besides, in our general response, we provide the performance of multi-step step MAML, which replaces the lower-level problems in (8) with $m$-step gradient updates, with $m\\in$ {$1,4,8,12$}. From the plot we can see as $m$ increases, multi-step MAML tends to have better performance, but it still cannot outperform our proposed RT-MLMC algorithm.\n    \n4. [Experiments against baselines] Thank you for the suggestions. As we are the first to propose and solve the CSBO problem, we have not found any other baselines for CSBO. \n    \n    - Regarding Wasserstein DRO with side information, it is worth noting that the existing method in [Yang et al., 2022] heavily relies on convexity, linearity predictors, and linear loss assumptions to build reformulations that can be solved via convex solvers. However, when the hypothesis class from the covariate to the decision is parameterized by nonconvex neural networks, their method is not implementable. In contrast, our algorithm is the first implementable algorithm with a convergence guarantee. We compare our method to naively incorporating SAA and Wasserstein DRO methods that do not explicitly leverage side information in Figure 2 (right).\n    \n    - For multi-task stochastic bilevel optimization [Guo et al., 2021], which is a special case of CSBO with $M$ lower-level problems, we compare the performance of RT-MLMC with the state-of-the-art method BSVRB proposed in [Hu et al., 2023] (this manuscript appears on arXiv after the NeurIPS submission deadline) in the one-page PDF response. Please see Figure 5 for a comparison with BSVRB. Our RT-MLMC method converges much faster.\n\nXing, Pengwei, et al. \"Big-fed: Bilevel optimization enhanced graph-aided federated learning.\" *IEEE Transactions on Big Data* (2022)\n\nSadana, Utsav, et al. \"A Survey of Contextual Optimization Methods for Decision Making under Uncertainty.\" arXiv preprint arXiv:2306.10374 (2023).\n\nHu et al. \"Blockwise stochastic variance-reduced methods with parallel speedup for multi-block bilevel optimization\". arXiv preprint arXiv:2305.18730, 2023"}]

Author rebuttal:
We are grateful to the reviewers for constructive comments and suggestions, which significantly improve the quality of our paper. We are happy to engage in further discussion. We first restate the importance of CSBO problem in applications and the challenges for solving it. You may find added experiments in the attached PDF.

1.  In addition to the ones demonstrated in the numerical experiments, CSBO covers many other important applications, including personalized federated learning [Xing, et al. "Big-fed: Bilevel optimization enhanced graph-aided federated learning." IEEE Transactions on Big Data (2022)] and end-to-end learning that integrates learning and optimization (See the third paradigm in the survey [Sadana, Utsav, et al. "A Survey of Contextual Optimization Methods for Decision Making under Uncertainty." arXiv preprint arXiv:2306.10374 (2023).]) 

2.  Existing algorithms for traditional bilevel optimization problems with one single lower-level problem either cannot achieve optimal complexity bounds or does not apply to the CSBO problem due to potentially infinitely many lower-level problems parametrized by $\xi$, each of which introduces a constraint involving solving a stochastic optimization problem. We highlight that the proposed RT-MLMC achieves the optimal complexity bounds for CSBO problems. We illustrate in more details in the response to each reviewer.

In below, we add some extra experiments to address some concerns from reviewers.

1. Reviwer Dnys asked about the step size of MAML in meta-learning. We point out that the step size of MAML has been fine-tuned in our experiment. In our general response (Figure 4a), we report the plot of MAML performance for different choices of step size from the list {5e-3, 1e-2, 5e-2, 1e-1, 2e-1}. From the plot we can see that for small step sizes the MAML tends to have similar performance, whereas MAML tends to diverge for too large step sizes. The key issue of MAML is that it solves a different formulation than problem (8), i.e., it replaces the lower level problems with one-step gradient update. Thus, it is natural that MAML cannot achieve good performance on the CSBO objective function, as the approximation gap is theoretically $O(1)$ unless one performs multi-step MAML.
    - To further illustrate the performance of MAML, in our general response, we provide the performance of multi-step MAML in the plot of Figure 4b, which replaces the lower level problems in (8) with $m$-step gradient updates, with $m\in \\{1,4,8,12\\}$. From the plot we can see as $m$ increases, multi-step MAML tends to have better performance, but it still cannot outperform our proposed RT-MLMC algorithm.

2. Reviewers Z7Go and NrrF pointed out that we should add more baseline comparison for solving the meta-learning formulation (8), which is a special case of CSBO problem. Note that it is a stochastic bilevel optimization problem with multiple lower-level constraints. Many baseline approaches can actually only solve a surrogate of the formulation by replacing the lower-level problem with gradient updates to obtain one-step or multi-step MAML or by averaging all lower-level problems so that there is only one lower-level constraint. 
    - The performance of MAML is discussed in the previous bullet points.
   - Only two recent papers proposed algorithms for directly solving the formulation (8): [Guo and Yang, 2021] and [Hu et al. 2023] (note that [Hu et al. 2023] appears after the NeurIPS deadline. We still use it as a baseline). The algorithm in [Guo and Yang, 2021] and the first algorithm in [Hu et al. 2023], both require computing the inverse of Hessian exactly in each iteration, which cannot be implemented efficiently especially for high-dimensional problems (For instance, in our meta-learning experiment the dimension of decision variable is $512\times10=5120$, so these two algorithms will require inverting a $5120\times5120$ size matrix in each iteration). We use the $\mathrm{BSVRB}^{v2}$, the second algorithm in [Hu et al. 2023], as a baseline to solve the special case of CSBO formulation. Figure 5 in the PDF file compare the performance of $\mathrm{BSVRB}^{v2}$ and the proposed RT-MLMC. Note that the performance of $\mathrm{BSVRB}^{v2}$ is worse than RT-MLMC. The reason is that the iteration complexity of $\mathrm{BSVRB}^{v2}$  depends linearly on the number of lower-level problems, whereas the proposed RT-MLMC does not and achieves optimal complexity bounds.
3. Reviewer Z7Go asked to add more baseline methods for Wasserstein DRO with side information. Note that the existing method in [Yang et al., 2022] heavily relies on convexity, linearity predictors, and linear loss assumptions to build reformulations that can be solved via convex solvers. However, when the hypothesis class from the covariate to the decision is parameterized by nonconvex neural networks, their method is not implementable. In contrast, our algorithm is the first implementable algorithm with a convergence guarantee. We compare our method to baselines such that naively incorporating SAA and Wasserstein DRO methods that do not explicitly leverage side information in Figure 2 (c). We again summarize the numerical results in Figure 2(c) as Table 5 in our PDF file.

References:
- Zhishuai Guo and Tianbao Yang. Randomized stochastic variance-reduced methods for stochastic bilevel optimization. arXiv preprint arXiv:2105.02266, 2021
- Quanqi Hu, Zi-Hao Qiu, Zhishuai Guo, Lijun Zhang, and Tianbao Yang. Blockwise stochastic variance-reduced methods with parallel speedup for multi-block bilevel optimization. arXiv preprint arXiv:2305.18730, 2023. (Appears on ArXiv after the NeurIPS submission deadline.)

Meta-review: The paper investigates a generalization of the stochastic bilevel optimization model, in which the lower and upper optimization levels share a random variable. The authors design two gradient-based approaches named RT-MLMC and DL-SGD for solving this problem and analyze their performance. Finally, numerical examples are presented to demonstrate the generality and performance of the proposed model and methods.

I agree with the reviewers' assessment that this is a good contribution that addresses a common practical problem in bilevel optimization. I encourage the authors to take into account the reviewers' feedback for the final version of the paper.
Decision: Accept (poster)