SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Yifan Hu*
EPFL & ETH Zürich
Jie Wang
Yao Xie
Daniel Kuhn
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.
A contextual stochastic bilevel optimization (CSBO) problem differs from a classical stochastic bilevel optimization problem only in that its lower-level problem is conditioned on a given context  $\xi$ .
$$
\min  _ {x \in \mathbb {R} ^ {d _ {x}}} F (x) := \mathbb {E} _ {\xi \sim \mathbb {P} _ {\xi}, \eta \sim \mathbb {P} _ {\eta | \xi}} [ f (x, y ^ {*} (x; \xi); \eta , \xi) ] \quad (\text {u p p e r l e v e l}) \tag {1}
$$
where  $y^{*}(x;\xi)\coloneqq \mathrm{argmin}_{y\in \mathbb{R}^{d_{y}}}\mathbb{E}_{\eta \sim \mathbb{P}_{\eta |\xi}}[g(x,y;\eta ,\xi)]\quad \forall \xi$  and  $x$ . (lower level)
Here  $\xi \sim \mathbb{P}_{\xi}$  and  $\eta \sim \mathbb{P}_{\eta|\xi}$  are random vectors, with  $\mathbb{P}_{\eta|\xi}$  denoting the conditional distribution of  $\eta$  for a given  $\xi$ . The dimensions of the upper-level decision variable  $x$  and the lower-level decision variable  $y$  are  $d_x$  and  $d_y$ , respectively. The functions  $f$  and  $g$  are continuously differentiable in  $(x,y)$  for any given sample pair  $(\xi,\eta)$ . The function  $f(x,y;\eta,\xi)$  can be nonconvex in  $x$ , but the function  $g(x,y;\eta,\xi)$  must be strongly convex in  $y$  for any given  $x$ ,  $\eta$  and  $\xi$ . Thus,  $y^*(x;\xi)$  is the unique minimizer of the strongly convex lower-level problem for any given  $x$  and  $\xi$ . Note that, on its own, the lower-level problem can be viewed as a contextual stochastic optimization problem [Bertsimas and Kallus, 2020] parametrized in  $x$ . We assume that the joint distribution of  $\xi$  and  $\eta$  is unknown. However, we assume that we have access to any number of independent and
*Correspondence: yifan.hu@epfl.ch.
identically distributed (i.i.d.) samples from  $\mathbb{P}_{\xi}$ , and for any given realization of  $\xi$ , we can generate any number of i.i.d. samples from the conditional distribution  $\mathbb{P}_{\eta|\xi}$ . The bilevel structure generally makes the objective function  $F(x)$  nonconvex in the decision variable  $x$ , except for few special cases. Thus we aim to develop efficient gradient-based algorithms for finding an  $\epsilon$ -stationary point of the nonconvex objective function  $F$ , i.e., a point  $\widehat{x}$  satisfying the inequality  $\mathbb{E}\| \nabla F(\widehat{x})\|^2 \leq \epsilon^2$ .
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  _ {x \in \mathbb {R} ^ {d _ {x}}} \quad \mathbb {E} _ {\xi \sim \mathbb {P} _ {\xi}} [ f (x, y ^ {*} (x); \xi) ] \tag {2}
$$
$$
\text {w h e r e} \quad y ^ {*} (x) := \operatorname {a r g m i n} _ {y \in \mathbb {R} ^ {d _ {y}}} \mathbb {E} _ {\eta \sim \mathbb {P} _ {\eta}} [ g (x, y; \eta) ].
$$
Indeed, (2) is a special case of CSBO if the upper- and lower-level objective functions are stochastically independent. Another special case of CSBO is the conditional stochastic optimization (CSO) problem [Hu et al., 2020a,b, He and Kasiviswanathan, 2023, Goda and Kitade, 2023] representable as
$$
\min  _ {x \in \mathbb {R} ^ {d _ {x}}} \mathbb {E} _ {\xi \sim \mathbb {P} _ {\xi}} [ f (x, \mathbb {E} _ {\eta \sim \mathbb {P} _ {\eta | \xi}} [ h (x; \eta , \xi) ]; \xi) ]. \tag {3}
$$
Indeed, (3) is a special case of CSBO if we set  $g(x,y;\eta ,\xi) = \| y - h(x;\eta ,\xi)\| ^2$ , in which case the lower-level problem admits the unique closed-form solution  $y^{*}(x,\xi) = \mathbb{E}_{\mathbb{P}_{\eta |\xi}}[h(x;\eta ,\xi)]$ .
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  $\xi$  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 as special cases, including meta-learning [Rajeswaran et al., 2019], personalized federated learning [Shamsian et al., 2021, Xing et al., 2022, Wang et al., 2023a], hierarchical representation learning [Yao et al., 2019], end-to-end learning [Donti et al., 2017, Sadana et al., 2023, Rychener et al., 2023, Grigas et al., 2021], Sinkhorn distributionally robust optimization (DRO) [Wang et al., 2023b], Wasserstein DRO with side information [Yang et al., 2022], information retrieval [Qiu et al., 2022], contrastive learning [Qiu et al., 2023], and instrumental variable regression [Muandet et al., 2020]. Below we 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 a common regularization center  $\theta$  shared by all tasks or all users.
$$
\min  _ {x} \mathbb {E} _ {i \sim \mu} \mathbb {E} _ {D _ {i} ^ {\text {t e s t}} \sim \rho_ {i}} \left[ l _ {i} \left(y _ {i} ^ {*} (x), D _ {i} ^ {\text {t e s t}}\right) \right] \quad (\text {u p p e r l e v e l})
$$
$$
\text {w h e r e} y _ {i} ^ {*} (x) = \operatorname {a r g m i n} _ {y _ {i}} \mathbb {E} _ {D _ {i} ^ {\text {t r a i n}} \sim \rho_ {i}} \left[ l _ {i} \left(y _ {i}, D _ {i} ^ {\text {t r a i n}}\right) + \frac {\lambda}{2} \| y _ {i} - x \| ^ {2} \right], \forall i \in [ M ], x. \quad (\text {l o w e r l e v e l}) \tag {4}
$$
Here,  $\mu$  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  $x$ , and the lower-level problem finds an optimal regularization parameter  $x_{i}$  close to  $x$  for each individual task or user. Note that  $M$  may be as large as  $\mathcal{O}(10^{3})$  in meta-learning and as large as  $\mathcal{O}(10^{9})$  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 click-through rate, and then making decisions based on this estimation. End-to-end learning streamlines this into a single-step method, allowing the optimization to account for estimation errors, thereby enabling more informed decisions. This can be framed as a 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  $\xi$ . For example, in online advertising,  $x$  represents the click-through rate estimator, and  $y^{*}(x;\xi)$  denotes the optimal advertisement display for a customer characterized by the feature vector  $\xi$ . For a comprehensive review, see 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-optimal convergence rates or are entirely unable to handle the contextual information. Indeed, a major challenge of CSBO is to estimate  $y^{*}(x;\xi)$  for (typically) uncountably many realizations of  $\xi$ . 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., 2021, Chen et al., 2022a, 2021, Hong et al., 2023, Yang et al., 2021]. Indeed, as the unique minimizer  $y^{*}(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^t\}_t$  approximates  $\{y^{*}(x^{t})\}_{t}$ . Specifically, these approaches leverage the approximation
$$
y ^ {*} (x ^ {t + 1}) - y ^ {*} (x ^ {t}) \approx \nabla y ^ {*} (x ^ {t}) ^ {\top} (x ^ {t + 1} - x ^ {t}),
$$
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;\xi)$  of the lower-level problem in (1) additionally depends on the context  $\xi$ , i.e., each realization of  $\xi$  corresponds to a lower-level constraint. Consequently, there can be infinitely many lower-level constraints. It is unclear how samples from  $\mathbb{P}_{\eta|\xi}$  corresponding to a fixed context  $\xi$  can be reused to estimate the minimizer  $y^{*}(x;\xi')$  corresponding to a different context  $\xi'$ . Since gradient-based methods sample  $\xi^t$  independently in each iteration  $t$ , no single sequence  $\{y^t\}_t$  can approximate the function  $\{y^{*}(x^{t},\xi^{t})\}_{t}$ . Guo and Yang [2021] and Hu et al. [2023] analyze a special case of the CSBO problem (1), in which  $\xi$  is supported on  $M$  points as shown in (4). 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  $\xi$ .
SBO problems can also be addressed with double-loop stochastic gradient 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  $\epsilon$ -optimality for a fixed  $\xi$  requires  $\mathcal{O}(\epsilon^{-2})$  samples from  $\mathbb{P}_{\eta|\xi}$  and gradient estimators for the function  $g$ , which leads to a  $\widetilde{\mathcal{O}}(\epsilon^{-6})$  total sample and computational complexity to obtain an  $\epsilon$ -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 a 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 [Glynn and Quinn, 2010]. We refer to the proposed method as random truncated MLMC (RT-MLMC) and demonstrate that the RT-MLMC estimator for  $\nabla F$  requires only  $\mathcal{O}(1)$  samples from  $\mathbb{P}_{\eta|\xi}$ . This is a significant improvement vis-à-vis DL-SGD, which requires  $\mathcal{O}(\epsilon^{-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  $\epsilon$ -stationary point of  $F$  is given by  $\widetilde{\mathcal{O}}(\epsilon^{-4})$ .
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.
- We introduce CSBO as a unifying framework for a broad range of machine learning tasks and optimization problems. We propose two methods, DL-SGD and RT-MLMC, and analyze their complexities; see Table 1 for 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].
Table 1: Complexity of the RT-MLMC and DL-SGD algorithms for finding an  $\epsilon$ -stationary point of  $F$ . The sample complexity refers to the total number of samples from  $\mathbb{P}_{\xi}$  as well as  $\mathbb{P}_{\eta|\xi}$ .
Nonconvex CSBO	Sample Complexity	Gradient Complexity of g and f	Per-iteration Memory Cost
RT-MLMC	O(ε-4)	O(ε-4) | O(ε-4)	O(dx + dy)
DL-SGD	O(ε-6)	O(ε-6) | O(ε-4)	O(dx + dy)
- For meta-learning with  $M$  tasks, the complexity bounds of RT-MLMC are constant in  $M$ . Thus, RT-MLMC outperforms the methods by Guo and Yang [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 non-parametric decision rules. In contrast, RT-MLMC 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.
Preliminaries For any function  $\psi : \mathbb{R}^{d_x} \times \mathbb{R}^{d_y}$  with arguments  $x \in \mathbb{R}^{d_x}$  and  $y \in \mathbb{R}^{d_y}$ , we use  $\nabla \psi$ ,  $\nabla_1 \psi$  and  $\nabla_2 \psi$  to denote the gradients of  $\psi$  with respect to  $(x, y)$ ,  $x$  and  $y$ , respectively. Similarly, we use  $\nabla^2 \psi$ ,  $\nabla_{11}^2 \psi$  and  $\nabla_{22}^2$  to denote Hessians of  $\psi$  with respect to  $(x, y)$ ,  $x$  and  $y$ , respectively. In addition,  $\nabla_{12}^2 \psi$  stands for the  $(d_x \times d_y)$ -matrix with entries  $\partial_{x_i y_j}^2 \psi$ . A function  $\varphi : \mathbb{R}^d \to \mathbb{R}$  is  $L$ -Lipschitz continuous if  $|\varphi(x) - \varphi(x')| \leq L \| x - x' \|$  for all  $x, x' \in \mathbb{R}^d$ , and it is  $S$ -Lipschitz smooth if it is continuously differentiable and satisfies  $\| \nabla \varphi(x) - \nabla \varphi(x') \| \leq S \| x - x' \|$  for all  $x, x' \in \mathbb{R}^d$ . In addition,  $\varphi$  is called  $\mu$ -strongly convex if it is continuously differentiable and if  $\varphi(x) - \varphi(x') - \nabla \varphi(x')^\top (x - x') \geq \frac{\mu}{2} \| x - x' \|^2$  for all  $x, x' \in \mathbb{R}^d$ . The identity matrix is denoted by  $I$ . Finally, we use  $\widetilde{\mathcal{O}}(\cdot)$  as a variant of the classical  $\mathcal{O}(\cdot)$  symbol that hides logarithmic factors.
2 Algorithms for Contextual Stochastic Bilevel Optimization
Throughout the paper, we 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 (1) satisfies the following regularity conditions:
(i)  $f$  is continuously differentiable in  $x$  and  $y$  for any fixed  $\eta$  and  $\xi$ , and  $g$  is twice continuously differentiable in  $x$  and  $y$  for any fixed  $\eta$  and  $\xi$ .
(ii)  $g$  is  $\mu_g$ -strongly convex in  $y$  for any fixed  $x$ ,  $\eta$  and  $\xi$ .
(iii)  $f, g, \nabla f, \nabla g$  and  $\nabla^2 g$  are  $L_{f,0}, L_{g,0}, L_{f,1}, L_{g,1}$  and  $L_{g,2}$ -Lipschitz continuous in  $(x,y)$  for any fixed  $\eta$  and  $\xi$ , respectively.
(iv) If  $(\eta, \xi) \sim \mathbb{P}_{(\eta, \xi)}$ , then  $\nabla f(x, y; \eta, \xi)$  is an unbiased estimator for  $\nabla \mathbb{E}_{(\eta, \xi) \sim \mathbb{P}_{(\eta, \xi)}}[f(x, y; \eta, \xi)]$  with variance  $\sigma_f^2$  uniformly across all  $x$  and  $y$ . Also, if  $\eta \sim \mathbb{P}_{\eta|\xi}$ , then  $\nabla g(x, y; \eta, \xi)$  is an unbiased estimator for  $\nabla \mathbb{E}_{\eta \sim \mathbb{P}_{\eta|\xi}}[g(x, y; \eta, \xi)]$  with variance  $\sigma_{g,1}^2$ , and  $\nabla^2 g(x, y; \eta, \xi)$  is an unbiased estimator for  $\nabla^2 \mathbb{E}_{\eta \sim \mathbb{P}_{\eta|\xi}}[g(x, y; \eta, \xi)]$  with variance  $\sigma_{g,1}^2$  uniformly across all  $x, y$  and  $\xi$ .
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  $L_{F}$ -Lipschitz continuous as well as  $S_{F}$ -Lipschitz smooth for some  $L_{F}, S_{F} > 0$ . Assumptions 2.1 (i-iii) also imply that the gradients of  $f$  and  $g$  with respect  $(x,y)$  can be interchanged with the expectations with respect to  $(\eta,\xi) \sim \mathbb{P}_{\eta,\xi}$  and  $\eta \sim \mathbb{P}_{\eta|\xi}$ . Hence, Assumptions 2.1 (i-iii) 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 2.1 (i-iii).
In order to design SGD-type algorithms for problem (1), we first construct gradient estimators for  $F$ . To this end, we observe that the Jacobian  $\nabla_1 y^*(x; \xi) \in \mathbb{R}^{d_x \times d_y}$  exists and is Lipschitz continuous in  $x$  for any fixed  $\xi$  thanks to [Chen et al., 2021, Lemma 2]. By the chain rule, we therefore have
$$
\nabla F (x) = \mathbb {E} _ {(\eta , \xi) \sim \mathbb {P} _ {(\eta , \xi)}} \left[ \nabla_ {1} f (x, y ^ {*} (x; \xi); \eta , \xi) + \nabla_ {1} y ^ {*} (x; \xi) ^ {\top} \nabla_ {2} f (x, y ^ {*} (x; \xi); \eta , \xi) \right].
$$
Algorithm 1 EpochsGD  $(K,x,\xi ,y_1^0)$
Input: # of epochs  $K$ , sample  $\xi$ , upper-level decision  $x$ , initial iterate  $y_1^0$ .
1: for  $k = 1$  to  $K$  do
2: for  $j = 0$  to  $2^k - 1$  do
3: Sample  $\eta_k^j$  from  $\mathbb{P}_{\eta|\xi}$  and update  $y_k^{j+1} = y_k^j - \beta_0 2^{-k} \nabla_2 g(x, y_k^j; \eta_k^j, \xi)$ .
4: end for
5: Update  $y_{k+1}^{0} = 2^{-k} \sum_{j=0}^{2^{k}-1} y_{k}^{j}$ .
6: end for
Output:  $y_1^0, y_K^0$ , and  $y_{K + 1}^0$ .
By following a similar procedure as in [Ghadimi and Wang, 2018], we can derive an explicit formula for  $\nabla_1y^* (x;\xi)$  (for details we refer to Appendix B) and substitute it into the above equation to obtain
$$
\begin{array}{l} \nabla F (x) = \mathbb {E} _ {(\eta , \xi) \sim \mathbb {P} _ {(\eta , \xi)}} \left[ \nabla_ {1} f (x, y ^ {*} (x; \xi); \eta , \xi) \right. \\ \left. - \left(\mathbb {E} _ {\eta^ {\prime} \sim \mathbb {P} _ {\eta | \xi}} \nabla_ {1 2} ^ {2} g (x, y ^ {*} (x; \xi); \eta^ {\prime}, \xi)\right) \Lambda (x, y ^ {*} (x; \xi); \xi) \nabla_ {2} f (x, y ^ {*} (x; \xi); \eta , \xi) \right], \\ \end{array}
$$
where  $\Lambda (x,y;\xi) = (\mathbb{E}_{\eta \sim \mathbb{P}_{\eta |\xi}}\nabla_{22}^{2}g(x,y;\eta ,\xi))^{-1}$ . Thus, the main challenges of constructing a gradient estimator are to compute and store the minimizer  $y^{*}(x,\xi)$  as well as the inverse expected Hessian matrix  $\Lambda (x,y;\xi)$  for all (potentially uncountably many) realizations of  $\xi$ . Computing these two objects exactly would be too expensive. In the remainder of this section, we thus derive estimators for  $y^{*}(x;\xi)$  and  $\Lambda (x,y;\xi)$ , and we combine these two estimators to construct an estimator for  $\nabla F(x)$ .
Estimating  $y^{*}(x;\xi)$ . We estimate  $y^{*}(x;\xi)$  using the gradient-based method EpochsGD by Hazan and Kale [2014], Asi et al. [2021], which involves  $K$  epochs of increasing lengths. Each epoch  $k = 1,\dots ,K$  starts from the average of the iterates computed in epoch  $k - 1$  and then applies  $2^{k}$  stochastic gradient steps to the lower-level problem with stepsize  $2^{-k}$  (see Algorithm 1). In the following we use the output  $y_{K + 1}^{0}$  of Algorithm 1 with inputs  $K, x, \xi$  and  $y_0$  as an estimator for the minimizer  $y^{*}(x;\xi)$  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;\xi)$  could also be estimated via classical SGD. Even though this would lead to similar complexity results, the analysis would become more cumbersome.
Estimating  $\Lambda (x,y;\xi)$ . Following [Ghadimi and Wang, 2018], one can estimate the inverse of an expected random matrix  $A$  with  $0\prec A\prec I$  using a Neumann series argument. Specifically, we have
$$
[ \mathbb {E} _ {A \sim \mathbb {P} _ {A}} A ] ^ {- 1} = \sum_ {n ^ {\prime} = 0} ^ {\infty} (I - \mathbb {E} _ {A \sim \mathbb {P} _ {A}} A) ^ {n} = \sum_ {n ^ {\prime} = 0} ^ {\infty} \prod_ {n = 1} ^ {n ^ {\prime}} \mathbb {E} _ {A _ {n} \sim \mathbb {P} _ {A}} (I - A _ {n}) \approx \sum_ {n ^ {\prime} = 0} ^ {N} \prod_ {n = 1} ^ {n ^ {\prime}} \mathbb {E} _ {A _ {n} \sim \mathbb {P} _ {A}} (I - A _ {n}).
$$
The truncated series on the right hand side provides a good approximation if  $N \gg 1$ . Assumption 2.1 (iii) implies that  $0 \prec \nabla_{22}^{2} g(x, y; \eta_{n}, \xi) \prec 2L_{g,1}I$ . Hence, the above formula can be applied to  $A = \frac{1}{2L_{g,1}} \nabla_{22}^{2} g(x, y; \eta_{n}, \xi)$ , which gives rise to an attractive estimator for  $\Lambda(x, y; \xi)$  of the form
$$
\widehat {\Lambda} (x, y; \xi) := \left\{ \begin{array}{l l} \frac {N}{2 L _ {g , 1}} I & \text {i f} \widehat {N} = 0, \\ \frac {N}{2 L _ {g , 1}} \prod_ {n = 1} ^ {\widehat {N}} \left(I - \frac {1}{2 L _ {g , 1}} \nabla_ {2 2} ^ {2} g (x, y; \eta_ {n}, \xi)\right) & \text {i f} \widehat {N} \geq 1. \end{array} \right. \tag {5}
$$
Here,  $\widehat{N}$  is a random integer drawn uniformly from  $\{0,1,\dots ,N - 1\}$  that is independent of the i.i.d. samples  $\eta_{1},\ldots ,\eta_{\widehat{N}}$  from  $\mathbb{P}_{\eta |\xi}$ . Chen et al. [2022b] showed that the estimator (5) displays the following properties. Its bias decays exponentially with  $N$ , its variance grows quadratically with  $N$ , and its sampling cost grows linearly with  $N$ . Below we call  $N$  the approximation number.