SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Next, we consider the per-iteration sampling costs and the average number of iterations for the EpochsGd. DL-SGD runs EpochsSGD as a subroutine for  $2^{K + 1} - 1$  number of iterations. Thus the sampling cost on average is  $N + 2^{K + 1} - 1$ .
Consider the variance of the DL-SGD method. Note that
$$
\widehat {v} ^ {K} (x) = H _ {K + 1} (1) - H _ {K + 1} (2)
$$
Following a similar decomposition as we did for the third term in bounding the variance of RT-MLMC estimators, we have
$$
\begin{array}{l} \mathbb {E} \| \widehat {v} ^ {K} (x) - \mathbb {E} \widehat {v} ^ {K} (x) \| ^ {2} \\ \leq \mathbb {E} \| \widehat {v} ^ {K} (x) - V (x) \| ^ {2} \\ \leq 2 \mathbb {E} \| H _ {K + 1} (1) - \mathbb {E} H ^ {*} (1) \| ^ {2} + 2 \mathbb {E} \| H _ {K + 1} (2) - \mathbb {E} H ^ {*} (2) \| ^ {2} \\ \leq 2 \mathbb {E} \| H _ {K + 1} (1) - \mathbb {E} H _ {K + 1} (1) \| ^ {2} + 2 \mathbb {E} \| \mathbb {E} H _ {K + 1} (1) - \mathbb {E} H ^ {*} (1) \| ^ {2} + 2 \mathbb {E} \| H _ {K + 1} (2) \| ^ {2} \\ + 2 \mathbb {E} \| \mathbb {E} H _ {K + 1} (2) - \mathbb {E} H ^ {*} (2) \| ^ {2} \\ \leq 2 \sigma_ {f} ^ {2} + \frac {L _ {f , 1} ^ {2} L _ {g , 0} ^ {2}}{2 ^ {K} \mu_ {g} ^ {2}} + L _ {g, 1} ^ {2} L _ {f, 0} ^ {2} \frac {N ^ {2}}{2 L _ {g , 1}} + 1 2 \left(L _ {f, 1} ^ {2} N ^ {2} + \frac {L _ {f , 0} ^ {2} L _ {g , 2} ^ {2} N ^ {2}}{L _ {g , 1} ^ {2}} + L _ {f, 0} ^ {2} N ^ {4}\right) \frac {L _ {g , 0} ^ {2}}{2 ^ {K + 1} \mu_ {g} ^ {2}} \\ = \mathcal {O} \left(N ^ {2} + N ^ {4} 2 ^ {- K}\right), \\ \end{array}
$$
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$ ,  $\nabla f$  and  $\nabla g$  and follows a similar argument as in bounding the third term for the variance of RT-MLMC estimators. Note that to control the bias, we let both  $N$  and  $K$  to be of order  $\mathcal{O}(\log (\epsilon^{-1}))$ . Thus  $N^2$  is the dominating term in the variance of DL-SGD gradient estimator.
Next, we demonstrate the proof of Theorem 1.
Proof of Theorem 1. We first demonstrate the analysis for RT-MLMC.
RT-MLMC gradient method: combining Lemmas 1 and 2, we know that
$$
\begin{array}{l} \mathbb {E} \| \nabla F (\widehat {x} _ {T}) \| ^ {2} \\ \leq \frac {2 (\mathbb {E} F (x _ {1}) - \min  _ {x} F (x))}{\alpha T} + \frac {2}{T} \sum_ {t = 1} ^ {T} \mathbb {E} \| \nabla F (x _ {t}) \| \| \mathbb {E} [ \nabla F (x _ {t}) - \widehat {v} (x _ {t}) | x _ {t} ] \| \\ + \frac {2 S _ {F} \alpha}{T} \sum_ {t = 1} ^ {T} \mathbb {E} \| \widehat {v} (x _ {t}) - \nabla F (x _ {t}) \| ^ {2} \\ \end{array}
$$
$$
\begin{array}{l} \leq \frac {2 (\mathbb {E} F (x _ {1}) - \min  _ {x} F (x))}{\alpha T} + \frac {2}{T} \sum_ {t = 1} ^ {T} \mathbb {E} \| \nabla F (x _ {t}) \| \| \mathbb {E} [ \nabla F (x _ {t}) - \widehat {v} (x _ {t}) | x _ {t} ] \| \\ + \frac {4 S _ {F} \alpha}{T} \sum_ {t = 1} ^ {T} \mathbb {E} \| \widehat {v} (x _ {t}) - V (x _ {t}) \| ^ {2} + \frac {4 S _ {F} \alpha}{T} \sum_ {t = 1} ^ {T} \mathbb {E} \| V (x _ {t}) - \nabla F (x _ {t}) \| ^ {2} \\ \end{array}
$$
$$
= \mathcal {O} \Big (\frac {1}{\alpha T} + S _ {F} \frac {1}{\mu_ {g}} \Big (1 - \frac {\mu_ {g}}{2 L _ {g , 1}} \Big) ^ {N} + S _ {F} \frac {N ^ {2}}{\mu_ {g} 2 ^ {K / 2}} + \alpha S _ {F} \frac {K N ^ {4}}{\mu_ {g} ^ {2}} + \alpha \frac {1}{\mu_ {g} ^ {2}} \Big (1 - \frac {\mu_ {g}}{2 L _ {g , 1}} \Big) ^ {2 N} \Big).
$$
To ensure that  $\widehat{x}_T$  is an  $\epsilon$ -stationarity point, it suffices to let the right hand side of the inequality above to be  $\mathcal{O}(\epsilon^2)$ . Correspondingly, we set  $N = \mathcal{O}(\log (\epsilon^{-1}))$ ,  $K = \mathcal{O}(\log (\epsilon^{-2}))$ ,  $\alpha = \mathcal{O}(T^{-1 / 2})$ , and  $T = \widetilde{\mathcal{O}} (\epsilon^{-4})$ . As a result, the total sampling and the gradient complexity over  $g$  are of order  $3KT = \widetilde{\mathcal{O}} (\epsilon^{-4})$ . Since at each upper iteration, we only compute one gradient of  $f$ , the gradient complexity over  $f$  is  $T = \widetilde{\mathcal{O}} (\epsilon^{-4})$ .
Next, we demonstrate the analysis for DL-SGD.
DL-SGD gradient method: combining Lemmas 1 and 2, we have
$$
\mathbb {E} \| \nabla F (\widehat {x} _ {T}) \| ^ {2}
$$
$$
\begin{array}{l} \leq \frac {2 \left(\mathbb {E} F \left(x _ {1}\right) - \min  _ {x} F (x)\right)}{\alpha T} + \frac {2}{T} \sum_ {t = 1} ^ {T} \mathbb {E} \| \nabla F (x _ {t}) \| \| \mathbb {E} \left[ \nabla F (x _ {t}) - \widehat {v} ^ {K} (x _ {t}) \mid x _ {t} \right] \| \\ + \frac {2 S _ {F} \alpha}{T} \sum_ {t = 1} ^ {T} \mathbb {E} \| \widehat {v} ^ {K} (x _ {t}) - \nabla F (x _ {t}) \| ^ {2} \\ \end{array}
$$
$$
\begin{array}{l} \leq \frac {2 \left(\mathbb {E} F \left(x _ {1}\right) - \min  _ {x} F (x)\right)}{\alpha T} + \frac {2}{T} \sum_ {t = 1} ^ {T} \mathbb {E} \| \nabla F (x _ {t}) \| \| \mathbb {E} \left[ \nabla F (x _ {t}) - \widehat {v} ^ {K} (x _ {t}) \mid x _ {t} \right] \| \\ + \frac {4 S _ {F} \alpha}{T} \sum_ {t = 1} ^ {T} \mathbb {E} \| \widehat {v} ^ {K} (x _ {t}) - V (x _ {t}) \| ^ {2} + \frac {4 S _ {F} \alpha}{T} \sum_ {t = 1} ^ {T} \mathbb {E} \| V (x _ {t}) - \nabla F (x _ {t}) \| ^ {2} \\ = \mathcal {O} \left(\frac {1}{\alpha T} + S _ {F} \frac {1}{\mu_ {g}} \left(1 - \frac {\mu_ {g}}{2 L _ {g , 1}}\right) ^ {N} + S _ {F} \frac {N ^ {2}}{\mu_ {g} 2 ^ {K / 2}} + \alpha S _ {F} N ^ {2} + \alpha \frac {1}{\mu_ {g} ^ {2}} \left(1 - \frac {\mu_ {g}}{2 L _ {g , 1}}\right) ^ {2 N}\right). \\ \end{array}
$$
To ensure that  $\widehat{x}_T$  is an  $\epsilon$ -stationarity point, it suffices to let  $N = \mathcal{O}(\log (\epsilon^{-1}))$ ,  $K = \mathcal{O}(\log (\epsilon^{-2}))$ ,  $\alpha = \mathcal{O}(T^{-1 / 2})$ , and  $T = \widetilde{\mathcal{O}} (\epsilon^{-4})$ . As a result, the sample complexity and the gradient complexity over  $g$  is of order  $\mathcal{O}(2^K T) = \widetilde{\mathcal{O}} (\epsilon^{-6})$ . At each upper iteration, we compute one gradient of  $f$ , and thus the gradient complexity over  $f$  is  $T = \widetilde{\mathcal{O}} (\epsilon^{-4})$ .
B Computing  $\nabla_1y^* (x;\xi)$
To derive an explicit formula for  $\nabla_1y^* (x;\xi)$ , we use the first-order optimality condition of the unconstrained lower-level problem. Indeed, as  $g(x,y;\eta ,\xi)$  is strongly convex in  $y$ , for given  $x$  and  $\xi$ ,  $y^{*}(x;\xi)$  is the unique solution of the equation
$$
\mathbb {E} _ {\mathbb {P} _ {\eta | \xi}} \nabla_ {2} g (x, y ^ {*} (x; \xi); \eta , \xi) = 0.
$$
Taking gradients with respect to  $x$  on both sides and using the chain rule, we obtain
$$
\mathbb {E} _ {\mathbb {P} _ {\eta | \xi}} \left[ \nabla_ {2 1} ^ {2} g (x, y ^ {*} (x; \xi); \eta , \xi) + \nabla_ {2 2} ^ {2} g (x, y ^ {*} (x; \xi); \eta , \xi) \nabla_ {1} y ^ {*} (x; \xi) \right] = 0.
$$
Since  $g(x,y;\eta ,\xi)$  is  $\mu_g$ -strongly convex in  $y$  for any given  $x,\eta$ , and  $\xi$ , the expected Hessian matrix  $\mathbb{E}_{\mathbb{P}_{\eta |\xi}}\nabla_{22}^{2}g(x,y;\eta ,\xi)\in \mathbb{S}_{+}^{d_y}$  is invertible, and thus we find
$$
\nabla_ {1} y ^ {*} (x; \xi) ^ {\top} = - \Big (\mathbb {E} _ {\mathbb {P} _ {\eta | \xi}} \nabla_ {1 2} ^ {2} g (x, y ^ {*} (x; \xi); \eta , \xi) \Big) \Big (\mathbb {E} _ {\mathbb {P} _ {\eta | \xi}} \nabla_ {2 2} ^ {2} g (x, y ^ {*} (x; \xi); \eta , \xi) \Big) ^ {- 1},
$$
where  $\mathbb{E}_{\mathbb{P}_{\eta|\xi}} \nabla_{12}^2 g(x, y^*(x; \xi); \eta, \xi)$  is in fact the transpose of  $\mathbb{E}_{\mathbb{P}_{\eta|\xi}} \nabla_{21}^2 g(x, y^*(x; \xi); \eta, \xi)$ .
C Efficient Implementation for Hessian Vector Products
Our goal is to compute  $\widehat{c}(x, y; \xi) = \nabla_{12}^{2} g(x, y; \eta', \xi) \widehat{\Lambda}(x, y; \xi) \nabla_{2} f(x, y; \eta'', \xi)$  in (8) efficiently.
Algorithm 4 Hessian Vector Implementation for Computing  $\widehat{c}(x, y; \xi)$ .
Input: Iteration points  $(x,y)$ , samples  $\xi$  and  $\eta', \{\eta_n\}_{n=1}^{\hat{N}}, \eta'' \sim \mathbb{P}_{\eta|\xi}, \hat{N} \geq 1$ .  
1: Compute  $r_0 = \nabla_2 f(x,y;\eta'',\xi)$   
2: for  $n = 1$  to  $\hat{N}$  do  
3: Get Hessian-vector product  $v_n = \nabla_{22}^2 g(x,y;\eta_n,\xi)r_{n-1}$  via automatic differentiation on  $y$   
4: Update  $r_n = r_{n-1} - \frac{1}{L_{g,1}} v_n$   
5: end for  
6: Update  $r_{N_3}' = \frac{N}{L_{g,1}} r_{N_3}$   
7: Get Hessian-vector product  $c = \nabla_{12}^2 g(x,y;\eta',\xi)r_{\hat{N}}'$  via automatic differentiation on  $y$   
Output:  $c$   
Remark: Step 4 and 9 can be implemented at cost  $\mathcal{O}(d)$  by as we compute  $\nabla_2 g(x,y;\eta_n,\xi)r_{n-1}$  first and then differentiate over  $y$  to avoid calculating Hessian matrices directly.
In this section, we provide all relevant implementation details. We fine-tune the stepsize for all approaches using the following strategy: we pick a fixed breakpoint denoted as  $t_0$  and adjust the stepsize accordingly. Specifically, for the  $t$ -th outer iteration when  $t \leq t_0$ , the stepsize is set as  $1 / \sqrt{t}$ , while for iterations beyond  $t_0$ , the stepsize is set as  $1 / t$ . We choose the breakpoint  $t_0$  such that training loss remains relatively stable as the number of outer iterations approaches  $t_0$ .
D.1 Meta-Learning
Let  $D = (a, b)$  denote the training or testing data, where  $a \in \mathbb{R}^d$  represents the feature vector and  $b \in [C]$  represents the label with  $C$  categories. In this section, the loss  $\ell_i(x, D)$  is defined as a multi-class logistic loss, given by:
$$
\mathcal {L} (x; D) = - \boldsymbol {b} ^ {T} x ^ {\top} a + \log \left(1 ^ {\top} e ^ {x ^ {\top} a}\right),
$$
where the parameter  $x \in \mathbb{R}^{d \times C}$  stands for the linear classifier,  $b \in \{0,1\}^C$  stands for the corresponding one-hot vector of the label  $b$ . The experiment utilizes the tinyImageNet dataset, which consists of 100,000 images belonging to 200 classes. After data pre-processing, each image has a dimension of 512. For dataset splitting, we divide the dataset such that  $90\%$  of the images from each class are assigned to the training set, while the remaining  $10\%$  belong to the testing set. The meta-learning task comprises 20 tasks, with each task involving the classification of 10 classes. Additionally, we set the hyper-parameter  $\lambda = 2$ .
Finally, we provide additional experiments for meta-learning in Figure 4. In the left plot, we examine the performance of MAML by varying the stepsize of inner-level gradient update from
$\{5\mathrm{e} - 3,1\mathrm{e} - 2,5\mathrm{e} - 2,1\mathrm{e} - 1,2\mathrm{e} - 1\}$ . From the plot we can tell that when using small stepsize, MAML tends to have similar performance whereas the performance of MAML tends to degrade when using large stepsize. In the right plot, we examine the performance of  $m$ -step MAML against upper-level iterations. The  $m$ -step MAML is an approximation of problem (4) via replacing  $y_{i}(x)$  with the  $m$ -step gradient update  $\widehat{y}_{i,m}(x)$ , which is defined recursively:
$$
\widehat {y} _ {i, 0} (x) = x, \quad \widehat {y} _ {i, k} (x) = \widehat {y} _ {i, k - 1} (x) - \gamma \nabla \left[ l _ {i} (\widehat {y} _ {i, k - 1} (x), D _ {i, k - 1} ^ {t r a i n}) + \frac {\lambda}{2} \| \widehat {y} _ {i, k - 1} (x) - x \| ^ {2} \right],
$$
with stepsize  $\gamma$  and  $D_{i,k - 1}^{train}\sim \rho_{i}$  for  $k = 1,\ldots ,m$ . Here we take the number of gradient updates at inner level  $m$  from  $\{1,4,8,12\}$ . From the plot, we realize that multi-step MAML tends to have better performance, but it still cannot outperform the RT-MLMC method.
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, we take the loss function
$$
l (w, \eta) = h (w - \eta) _ {+} + b (\eta - w) _ {+},
$$
where  $h > 0$  is a constant representing the per-unit holding cost and  $b > 0$  is a constant representing the per-unit backlog cost. Since Assumption 2.1 does not hold for the objective function due to its non-smooth structure, we approximate the loss function with the smoothed version
$$
l _ {\beta} (w, \eta) = \frac {h}{\beta} \log (1 + e ^ {\beta (w - \eta)}) + \frac {b}{\beta} \log (1 + e ^ {\beta (\eta - w)}),
$$
where we specify the hyper-parameter  $\beta = 5$  to balance the trade-off between loss function approximation and smoothness. The synthetic dataset in this part is generated in the following procedure: the covariate  $\xi$  is sampled from the 100-dimensional uniform distribution supported on  $[-15, 15]^{100}$ . The demand  $\eta$  depends on  $\xi$  in a nonlinear way:
$$
\eta = f _ {\mathrm {N N}} (x; \xi) + \epsilon , \quad \epsilon \sim \mathcal {N} (0, 1),
$$
$$
f _ {\mathrm {N N}} (x; \xi) = 1 0 * \operatorname {S i g m o i d} \left(W _ {3} \cdot \operatorname {R e L U} \left(W _ {2} \cdot \operatorname {R e L U} \left(W _ {1} \xi + b _ {1}\right) + b _ {2}\right) + b _ {3}\right),
$$
where the neural network parameter  $x \coloneqq (W_{1}, b_{1}, \ldots, W_{3}, b_{3})$ . In particular, the ground-truth neural network parameter is generated using the uniform initialization procedure by calling torch(nn.init.uniform_in pytorch. We quantify the performance of a given  $\theta$  using the testing loss  $\mathbb{E}_{(\xi, \eta) \sim \mathbb{P}^{*}}[l(f(x; \xi), \eta)]$ , where the expectation is estimated using sample average approximation based on  $2 \cdot 10^{5}$  sample points. When creating training dataset, we generate  $M = 50$  samples of  $\xi$ , denoted as  $\{\xi_{i}\}_{i \in [M]}$  and for each  $x_{i}$ , we generate  $m \in \{10, 30, 50, 100\}$  samples of  $\eta$  from the conditional distribution  $\mathbb{P}_{\eta | \xi_{i}}$ . When generating the left two plots of Fig. 3, we specify  $m = 100$ .
When solving the WDRO baseline, we consider the formulation
$$
\min  _ {x} \max  _ {\mathbb {P}} \left\{\mathbb {E} _ {(\xi , \eta) \sim \mathbb {P}} [ l _ {\beta} (f (x; \xi); \eta) ] - \lambda \mathcal {W} (\mathbb {P}, \mathbb {P} ^ {0}) \right\}, \tag {13}
$$
with  $\mathcal{W}(\cdot, \cdot)$  being the standard Wasserstein distance using the same cost function as 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  $\lambda$  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  $T$  and # of inner iterations  $T_{\mathrm{in}}$ , stepsizes  $\{\alpha_t\}_{t=1}^T$  and  $\{\beta_s\}_{s=1}^{T_{\mathrm{in}}}$ , initial iterate  $x_1$ .  
1: for  $t = 1$  to  $T$  do  
2: Sample  $(\xi^t, \eta^t)$  from  $\mathbb{P}$  and initialize  $\xi_1^t = \xi^t$ .  
3: for  $s = 1$  to  $T_{\mathrm{in}}$  do  
4: Generate  $g_s = \nabla_\xi [l_\beta (f(x_t; \xi_s^t); \eta^t) - \lambda \| \xi_s^t - \xi^t\|^2]$ .  
5: Update  $\xi_{s+1}^t = \xi_s^t + \beta_s \mathrm{Adam}(\xi_s^t, g_s)$ .  
6: end for  
7: Compute  $G_t = \nabla_x [l_\beta (f(x_t; \xi_{T_{\mathrm{in}} + 1}^t); \eta^t)]$ .  
8: Update  $x_{t+1} = x_t - \alpha_t G_t$ .  
9: end for  
Output:  $x_T$