SlowGuess/NeurIPS_2023_OCR huggingface dataset:
<sup>5</sup>We use  $\mathcal{L}_{\mathcal{D}_{12m}}^{(1)}, \mathcal{L}_{\mathcal{D}_{6m}}^{(2)}, \mathcal{L}_{\mathcal{D}_m}^{(3)}$  to ensure that each estimate uses same amount of samples (12m).
and  $a_4 = 0$  for quadratic  $q$ . In contrast, biased first order method only converges to a neighborhood. In Fig. 1b, we consider  $q(s) = s^4$  and  $p(\delta) = \delta /2$  where  $\delta \in [0,2]$ . All three methods are biased and their biases decrease with  $m$ , i.e.  $\mathcal{O}(m^{-k})$  for  $k$ th order method. Depending on the constants (e.g.  $a_i$ ,  $\sigma_{i}$ ), a higher-order extrapolation method may need decently large  $m$  (burn-in phase) to outperform lower-order methods.
3 Applying Stochastic Extrapolation in the CSO Problem
In this section, we apply the extrapolation-based scheme from the previous section to reduce the bias in the CSO problem. We focus on variants of BSGD and their accelerated version BSpiderBoost based on our second-order approximation operator (Definition 4). Let  $H_{\xi}$ ,  $\tilde{H}_{\xi}$ , and  $H_{\xi}'$  indicate different sets, each of which contains  $m$  i.i.d. random variables/samples drawn from the conditional distribution  $\mathbb{P}(\eta|\xi)$ . Remember that, as mentioned earlier, we use  $g_{\eta}(\boldsymbol{x})$  to represent  $g_{\eta}(\boldsymbol{x};\xi)$ .
Extrapolated BSGD. At time  $t$ , BSGD constructs a biased estimator of  $\nabla F(\pmb{x}^t)$  using one sample  $\xi$  and  $2m$  i.i.d. samples from the conditional distribution as in (1)
$$
G _ {\mathrm {B S G D}} ^ {t + 1} = \left(\frac {1}{m} \sum_ {\tilde {\eta} \in \tilde {H} _ {\xi}} \nabla g _ {\tilde {\eta}} \left(\boldsymbol {x} ^ {t}\right)\right) ^ {\top} \nabla f _ {\xi} \left(\frac {1}{m} \sum_ {\eta \in H _ {\xi}} g _ {\eta} \left(\boldsymbol {x} ^ {t}\right)\right). \tag {5}
$$
To reduce this bias, we apply the second-order extrapolation operator from (4). At time  $t$ , we define  $\mathcal{D}_{\pmb{g},\xi}^{t + 1}$  to be the distribution of the random variable  $\frac{1}{m}\sum_{\eta \in H_{\xi}}g_{\eta}(\pmb{x}^t)$ . Then we apply  $\mathcal{L}_{\mathcal{D}_{g,\xi}^{t + 1}}^{(2)}$  by setting  $q$  to  $\nabla f_{\xi}$  and  $s = 0$ , i.e.
$$
\begin{array}{l} \mathcal {L} _ {\mathcal {D} _ {\boldsymbol {g}, \xi} ^ {t + 1}} ^ {(2)} \nabla f _ {\xi} (0) := 2 \nabla f _ {\xi} \left(\frac {1}{2 m} \sum_ {\eta \in H _ {\xi}} g _ {\eta} (\boldsymbol {x} ^ {t}) + \frac {1}{2 m} \sum_ {\eta^ {\prime} \in H _ {\xi} ^ {\prime}} g _ {\eta^ {\prime}} (\boldsymbol {x} ^ {t}))\right) \\ \left. - \frac {1}{2} \left(\nabla f _ {\xi} \left(\frac {1}{m} \sum_ {\eta \in H _ {\xi}} g _ {\eta} \left(\boldsymbol {x} ^ {t}\right)\right) + \nabla f _ {\xi} \left(\frac {1}{m} \sum_ {\eta^ {\prime} \in H _ {\xi} ^ {\prime}} g _ {\eta^ {\prime}} \left(\boldsymbol {x} ^ {t}\right)\right)\right), \right. \tag {6} \\ \end{array}
$$
where  $\frac{1}{m}\sum_{\eta \in H_{\xi}}g_{\eta}(\pmb{x}^{t})$  and  $\frac{1}{m}\sum_{\eta^{\prime}\in H_{\xi}^{\prime}}g_{\eta^{\prime}}(\pmb{x}^{t})$  are i.i.d. drawn from  $\mathcal{D}_{\pmb{g},\xi}^{t + 1}$ . In Algorithm 2 (Appendix A), we present our extrapolated BSGD (E-BSGD) scheme, where we replace  $\nabla f_{\xi}\bigl (\frac{1}{m}\sum_{\eta \in H_{\xi}}g_{\eta}(\pmb{x}^{t})\bigr)$  in (5) by  $\mathcal{L}_{\mathcal{D}_{\pmb{g},\xi}^{t + 1}}^{(2)}\nabla f_{\xi}(0)$  resulting in this following gradient estimate:
$$
G _ {\mathrm {E} - \mathrm {B S G D}} ^ {t + 1} = \left(\frac {1}{m} \sum_ {\tilde {\eta} \in \tilde {H} _ {\xi}} \nabla g _ {\tilde {\eta}} (\boldsymbol {x} ^ {t})\right) ^ {\top} \mathcal {L} _ {\mathcal {D} _ {g, \xi} ^ {t + 1}} ^ {(2)} \nabla f _ {\xi} (0). \tag {7}
$$
Extrapolated BSpiderBoost. BSpiderBoost, proposed by Hu et al. [16], uses the variance reduction methods for nonconvex smooth stochastic optimization developed by Fang et al. [10], Wang et al. [38]. BSpiderBoost builds upon BSGD and has two kinds of updates: a large batch and a small batch update. In each step, it decides which update to apply based on a random coin. With probability  $p_{\mathrm{out}}$ , it selects a large batch update with  $B_{1}$  outer samples of  $\xi$ . With remaining probability  $1 - p_{\mathrm{out}}$ , it selects a small batch update where the gradient estimator will be updated with gradient information in the current iteration generated with  $B_{2}$  outer samples of  $\xi$  and the information from the last iteration. Formally, it constructs a gradient estimate as follows,
$$
G _ {\mathrm {B S B}} ^ {t + 1} = \left\{ \begin{array}{l l} G _ {\mathrm {B S B}} ^ {t} + \frac {1}{B _ {2}} \sum_ {\xi \in \mathcal {B} _ {2}, | \mathcal {B} _ {2} | = B _ {2}} \left(G _ {\mathrm {B S G D}} ^ {t + 1} - G _ {\mathrm {B S G D}} ^ {t}\right) & \text {w i t h p r o b .} 1 - p _ {\text {o u t}} \\ \frac {1}{B _ {1}} \sum_ {\xi \in \mathcal {B} _ {1}, | \mathcal {B} _ {1} | = B _ {1}} G _ {\mathrm {B S G D}} ^ {t + 1} & \text {w i t h p r o b .} p _ {\text {o u t}}. \end{array} \right. \tag {8}
$$
We propose our extrapolated BSpiderBoost scheme (formally defined in Algorithm 3, Appendix A) by replacing the BSGD gradient estimates in (8) with E-BSGD.
$$
G _ {\mathrm {E} - \mathrm {B S B}} ^ {t + 1} = \left\{ \begin{array}{l l} G _ {\mathrm {E} - \mathrm {B S B}} ^ {t} + \frac {1}{B _ {2}} \sum_ {\xi \in \mathcal {B} _ {2}, | \mathcal {B} _ {2} | = B _ {2}} \left(G _ {\mathrm {E} - \mathrm {B S G D}} ^ {t + 1} - G _ {\mathrm {E} - \mathrm {B S G D}} ^ {t}\right) & \text {w i t h p r o b .} 1 - p _ {\text {o u t}} \\ \frac {1}{B _ {1}} \sum_ {\xi \in \mathcal {B} _ {1}, | \mathcal {B} _ {1} | = B _ {1}} G _ {\mathrm {E} - \mathrm {B S G D}} ^ {t + 1} & \text {w i t h p r o b .} p _ {\text {o u t}}. \end{array} \right. \tag {9}
$$
Sample Complexity Analyses of E-BSGD and E-BSpiderBoost. We adopt the standard assumptions used in the literature [27, 35, 33, 41]. All proofs are deferred to Appendix D.
Assumption 3 (Lower bound)  $F$  is lower bounded by  $F^{\star}$ .
Assumption 4 (Bounded variance) Assume that  $g_{\eta}$  and  $\nabla g_{\eta}$  have bounded variances, i.e., for all  $\xi$  in the support of  $\mathbb{P}(\xi)$  and  $\pmb {x}\in \mathbb{R}^p$ ,  $\sigma_g^2\coloneqq \mathbb{E}_{\eta |\xi}[\| g_\eta (\pmb {x};\xi) - \mathbb{E}_{\eta |\xi}[g_\eta (\pmb {x};\xi)]\| _2^2 ] < \infty$  and  $\zeta_g^2\coloneqq \mathbb{E}_{\eta |\xi}[\| \nabla g_\eta (\pmb {x};\xi) - \mathbb{E}_{\eta |\xi}[\nabla g_\eta (\pmb {x};\xi)]\| _2^2 ] < \infty$
Assumption 5 (Lipschitz continuity/smoothness of  $f_{\xi}$  and  $g_{\eta}$ ) For all  $\xi$  in the support of  $\mathbb{P}(\xi)$ ,  $f_{\xi}(\cdot)$  is  $C_f$ -Lipschitz continuous (i.e.,  $\| f_{\xi}(\pmb{x}) - f_{\xi}(\pmb{x}')\|_2 \leq C_f\| \pmb{x} - \pmb{x}'\|_2 \forall \pmb{x}, \pmb{x}' \in \mathbb{R}^p$ ) and  $L_f$ -Lipschitz smooth (i.e.,  $\| \nabla f_{\xi}(\pmb{x}) - \nabla f_{\xi}(\pmb{x}')\|_2 \leq L_f\| \pmb{x} - \pmb{x}'\|_2, \forall \pmb{x}, \pmb{x}' \in \mathbb{R}^p$ ) for any  $\xi$ . Similarly, for all  $\xi$  in the support of  $\mathbb{P}(\xi)$  and  $\eta$  in the support of  $\mathbb{P}(\eta|\xi)$ ,  $g_{\eta}(\cdot; \xi)$  is  $C_g$ -Lipschitz continuous and  $L_g$ -Lipschitz smooth.
The smoothness of  $f_{\xi}$  and  $g_{\eta}$  naturally implies the smoothness of  $F$ . Zhang and Xiao [41, Lemma 4.2] show that Assumption 5 ensures  $F$  is: 1)  $C_F$ -Lipschitz continuous with  $C_F = C_fC_g$ ; and 2)  $L_{F}$ -Lipschitz smooth with  $L_{F} = L_{g}C_{f} + C_{g}^{2}L_{f}$ . We denote  $\tilde{L}_F = \zeta_gC_f + \sigma_gCgL_f$ . Moreover, Assumption 5 also guarantees that  $f_{\xi}$  and  $g_{\eta}$  have bounded gradients. In addition,  $f_{\xi}$  and  $g_{\eta}$  are assumed to satisfy the following regularity condition in order to apply our extrapolation-based scheme from Section 2.
Assumption 6 (Regularity) For all  $\xi$  in the support of  $\mathbb{P}(\xi)$ ,  $\nabla f_{\xi}$  is 4th-order differentiable with bounded derivatives (i.e.,  $a_{l} := \sup_{\pmb{g} \in \mathbb{R}^{p}} \left\| \nabla^{(l)} f_{\xi}(\pmb{g}) \right\|_{2} < \infty$  for  $l = 1, 2, 3, 4$ ,  $\forall \pmb{x} \in \mathbb{R}^{p}$ ) and  $g_{\eta}$  has bounded moments up to 4th-order (i.e.,  $\sigma_{k} = \sup_{\pmb{x} \in \mathbb{R}^{d}} \sup_{\xi} \mathbb{E}_{\eta|\xi} \left[ \sum_{i=1}^{p} \left[ g_{\eta}(\pmb{x}) - \mathbb{E}_{\eta|\xi} [g_{\eta}(\pmb{x})] \right]_{i}^{k} \right] < \infty, k = 1, 2, 3, 4$ ).
Remark 2 The core piece of Assumption 6 is the 4th order differentiability of  $\nabla f_{\xi}$  as other parts can be easily satisfied through appropriate boundedness assumptions. This condition though is satisfied by common instantiations of CSO/FCCO. We discuss some examples including invariant logistic regression, instrumental variable regression, first-order MAML for sine-wave few-shot regression task, deep average precision maximization in Section 5. Therefore, our improvements in sample complexity apply to all these problems.
Consider some time  $t > 0$ . Let  $G^{t + 1}$  be a stochastic estimate of  $\nabla F(\pmb{x}^t)$  where  $\pmb{x}^t$  is the current iterate. The next iterate  $\pmb{x}^{t + 1} \coloneqq \pmb{x}^t - \gamma G^t$ . Let  $\mathbb{E}[\cdot]$  denote the conditional expectation, where we condition on all the randomness until time  $t$ . We consider the bias and variance terms coming from our gradient estimate. Formally, we define the following two quantities
$$
\mathcal {E} _ {\mathrm {b i a s}} ^ {t + 1} = \left\| \nabla F (\boldsymbol {x} ^ {t}) - \mathbb {E} [ G ^ {t + 1} ] \right\| _ {2} ^ {2}, \quad \mathcal {E} _ {\mathrm {v a r}} ^ {t + 1} = \mathbb {E} [ \| G ^ {t + 1} - \mathbb {E} [ G ^ {t + 1} ] \| _ {2} ^ {2} ].
$$
Our idea of getting to an  $\varepsilon$ -stationary point (Definition 1) will be to ensure that  $\mathcal{E}_{\mathrm{bias}}^{t + 1}$  and  $\mathcal{E}_{\mathrm{var}}^{t + 1}$  are bounded. The main technical component of our analyses is in fact analyzing these bias and variance terms for the various gradient estimates considered. For this purpose, we first analyze the bias and variance terms for the (original) BSGD (Lemma 5) and BSpiderBoost (Lemma 7) algorithms, which are then used to get the corresponding bounds for our E-BSGD (Lemma 6) and E-BSpiderBoost (Lemma 8) algorithms. Through these bias and variance bounds, we establish the following main results of this section.
Theorem 3 [E-BSGD Convergence] Consider the (CSO) problem. Suppose Assumptions 3, 4, 5, 6 hold true and  $L_{F}, C_{F}, \tilde{L}_{F}, C_{g}, F^{\star}$  are constants and  $C_{e}(f; g) := \frac{8a_{3}\sigma_{3} + 18a_{4}\sigma_{2}^{2} + 5a_{4}\sigma_{4}}{96}$  defined in Corollary 1 are associated with second order extrapolation in the CSO problem. Let step size  $\gamma \leq 1/(2L_{F})$ . Then the output  $\mathbf{x}^{s}$  of E-BSGD (Algorithm 2) satisfies:  $\mathbb{E}[\|\nabla F(\mathbf{x}^{s})\|_{2}^{2}] \leq \varepsilon^{2}$ , for nonconvex  $F$ , if the inner batch size  $m = \Omega(C_{e}C_{g}\varepsilon^{-1/2})$ , and the number of iterations
$$
T = \Omega (L _ {F} (F (\pmb {x} ^ {0}) - F ^ {\star}) (\tilde {L} _ {F} ^ {2} / _ {m} + C _ {F} ^ {2}) \varepsilon^ {- 4}).
$$
The E-BSGD takes  $\mathcal{O}(\varepsilon^{-4})$  iterations to converge and compute  $\mathcal{O}(\varepsilon^{-0.5})$  gradients per iteration. Therefore, its resulting sample complexity is  $\mathcal{O}(\varepsilon^{-4.5})$  which is more efficient than  $\mathcal{O}(\varepsilon^{-6})$  of BSGD. Similar improvements can be observed for E-BSpiderBoost in Theorem 4.
Theorem 4 [E-BSpiderBoost Convergence] Consider the (CSO) problem under the same assumptions as Theorem 3. Let step size  $\gamma \leq 1 / (13L_F)$ . Then the output  $\mathbf{x}^s$  of E-BSpiderBoost (Algorithm 3)
satisfies:  $\mathbb{E}[\| \nabla F(\pmb{x}^s)\|_2^2] \leq \varepsilon^2$ , for nonconvex  $F$ , if the inner batch size  $m = \mathcal{O}(C_e C_g \varepsilon^{-0.5})$ , the hyperparameters of the outer loop of  $E$ -BSpiderBoost  $B_1 = (\tilde{L}_F^2 / m + C_F^2) \varepsilon^{-2}$ ,  $B_2 = \sqrt{B_1}$ ,  $p_{out} = 1 / B_2$ , and the number of iterations
$$
T = \Omega \left(L _ {F} \left(F \left(\boldsymbol {x} ^ {0}\right) - F ^ {\star}\right) \varepsilon^ {- 2}\right).
$$
The resulting sample complexity of E-BSpiderBoost is  $\mathcal{O}(\varepsilon^{-3.5})$ , which improves  $\mathcal{O}(\varepsilon^{-5})$  bound of BSpiderBoost [16] and  $\mathcal{O}(\varepsilon^{-4})$  bound of V-MLMC/RT-MLMC [17].
4 Applying Stochastic Extrapolation in the FCCO Problem
In this section, we apply the extrapolation-based scheme from Section 2 to the FCCO problem. We focus on case where  $n = O(\varepsilon^{-2})$ . For larger  $n$ , we can treat the FCCO problem as a CSO problem and get an  $\mathcal{O}(\varepsilon^{-3.5})$  bound from Theorem 4. All missing details are presented in Appendix E.
Now, a straightforward algorithm for FCCO is to use the finite-sum variant of SpiderBoost (or SPIDER) [10, 38] in Algorithm 3. In this case, if we choose the outer batch sizes to be  $B_{1} = n$ ,  $B_{2} = \sqrt{n}$  and the inner batch size to be  $m = \max \{\varepsilon^{-2}/n, \varepsilon^{-1/2}\}$ . The resulting sample complexity of E-BSpiderBoost now becomes,  $\mathcal{O}(\max \{\sqrt{n}/\varepsilon^{2.5}, 1/\sqrt{n}\varepsilon^4\})$ , which recovers  $\mathcal{O}(\varepsilon^{-3.5})$  bound as in Theorem 4 for  $n = \Theta(\varepsilon^{-2})$ . However, when  $n$  is small, such as  $n = \mathcal{O}(1)$ , the sample complexity degenerates to  $\mathcal{O}(\varepsilon^{-4})$  which is worse than the  $\Omega(\varepsilon^{-3})$  lower bound of stochastic optimization [3]. We leave the details to Theorem 8. We still use Assumptions 3, 4, 5, 6 for the analysis of FCCO problem, replacing the role of  $\xi$  with  $i$ .
Algorithm 1 E-NestedVR
1: Input:  $\pmb{x}^0 \in \mathbb{R}^d$ , step-size  $\gamma$ , batch sizes  $S_1, S_2, B_1, B_2$ , Probability  $p_{\mathrm{in}}, p_{\mathrm{out}}$   
2: for  $t = 0, 1, \dots, T - 1$  do  
3: if  $(t = 0)$  or (with prob.  $p_{\mathrm{out}}$ ) then ▷ Large outer batch  
4: for  $i \in \mathcal{B}_1 \sim [n]$  with  $|\mathcal{B}_1| = B_1$  do  
5: draw  $y_i^{t+1}$  from distribution  $\mathcal{D}_{y,i}^{t+1}$  defined in (10)  
6: compute  $z_i^{t+1}$  using (11) and define  $\phi_i^t = x^t$   
7: end for  
8:  $G_{\mathrm{E-NVR}}^{t+1} = \frac{1}{B_1} \sum_{i \in \mathcal{B}_1} (z_i^{t+1})^\top \mathcal{L}_{\mathcal{D}_{y,i}^{t+1}}^{(2)} \nabla f_i(0)$   
9: else ▷ Small outer batch  
10: for  $i \in \mathcal{B}_2$  with  $|\mathcal{B}_2| = B_2$  do  
11: draw  $y_i^{t+1}$  and  $y_i^t$  from distribution  $\mathcal{D}_{y,i}^{t+1}$  and  $\mathcal{D}_{y,i}^t$  defined in (10)  
12: compute  $z_i^{t+1}$  using (11) and define  $\phi_i^t = x^t$   
13: end for  
14:  $G_{\mathrm{E-NVR}}^{t+1} = G_{\mathrm{E-NVR}}^t + \frac{1}{B_2} \sum_{i \in \mathcal{B}_2} (z_i^{t+1})^\top (\mathcal{L}_{\mathcal{D}_{y,i}^{t+1}}^{(2)} \nabla f_i(0) - \mathcal{L}_{\mathcal{D}_{y,i}^t}^{(2)} \nabla f_i(0))$   
15: end if  
16:  $x^{t+1} = x^t - \gamma G_{\mathrm{E-NVR}}^{t+1}$   
17: end for  
18: Output:  $x^s$  picked uniformly at random from  $\{x^t\}_{t=0}^{T-1}$
Extrapolated Nested VR. We now introduce a nested variance reduction algorithm E-NestedVR which reaches low sample complexity for all choices of  $n$ . Missing proofs from this section are presented in Appendix E. For the stochasticities in the FCCO problem, our idea is to use two nested SpiderBoost variance reduction components: one for the outer random variable  $i$  and the other for the inner random variable  $\eta |i$ . In each outer (resp. inner) SpiderBoost step, we choose large batch  $B_{1}$  (resp.  $S_{1}$ ) with probability  $p_{\mathrm{out}}$  (resp.  $p_{\mathrm{in}}$ ); otherwise we choose small batch. Let  $H_{i}$  denote a set of  $m$  i.i.d. samples drawn from the conditional distribution  $\mathbb{P}(\eta |i)$ . Similarly, let  $\tilde{H}_i$  denote another set of  $m$  i.i.d. samples drawn from the same conditional distribution. For each given  $i$ , we approximate  $\mathbb{E}_{\eta |i}[g_{\eta}(\boldsymbol{x}^t)]$  with  $y_{i}^{t + 1}$  from distribution  $\mathcal{D}_{\boldsymbol{y},i}^{t + 1}$  where,
$$
\boldsymbol {y} _ {i} ^ {t + 1} = \left\{ \begin{array}{l l} \frac {1}{S _ {1}} \sum_ {\eta \in H _ {i}} g _ {\eta} \left(\boldsymbol {x} ^ {t}\right) & \text {w i t h p r o b .} p _ {\text {i n}} \text {o r} t = 0 \\ \boldsymbol {y} _ {i} ^ {t} + \frac {1}{S _ {2}} \sum_ {\eta \in H _ {i}} \left(g _ {\eta} \left(\boldsymbol {x} ^ {t}\right) - g _ {\eta} \left(\phi_ {i} ^ {t}\right)\right) & \text {w i t h p r o b .} 1 - p _ {\text {i n}}. \end{array} \right. \tag {10}
$$
Similarly, we approximate  $\mathbb{E}_{\bar{\eta} |i}[\nabla g_{\bar{\eta}}(\pmb{x}^t)]$  with  $\pmb{z}_i^{t + 1}$  defined as follows
$$
\boldsymbol {z} _ {i} ^ {t + 1} = \left\{ \begin{array}{l l} \frac {1}{S _ {1}} \sum_ {\tilde {\eta} \in \tilde {H} _ {i}} \nabla g _ {\tilde {\eta}} (\boldsymbol {x} ^ {t}) & \text {w i t h p r o b .} p _ {\text {i n}} \text {o r} t = 0 \\ \boldsymbol {z} _ {i} ^ {t} + \frac {1}{S _ {2}} \sum_ {\tilde {\eta} \in \tilde {H} _ {i}} (\nabla g _ {\tilde {\eta}} (\boldsymbol {x} ^ {t}) - \nabla g _ {\tilde {\eta}} (\phi_ {i} ^ {t})) & \text {w i t h p r o b .} 1 - p _ {\text {i n}}, \end{array} \right. \tag {11}
$$
(a) CSO: Large  $n = 50000$
(b) FCCO: small  $n = 50$
Figure 2: Performances of algorithms and their extrapolated versions on the invariant logistic regression task. Algorithms in each subplot use the same amount of inner batch size in each iteration. The shaded region represents the  $95\%$  -confidence interval computed over 10 runs.
where  $\phi_i^t$  is the last time  $i$  is visited before time  $t$ . If  $i$  is not selected at time  $t$ , then  $\pmb{y}_i^{t + 1} = \pmb{y}_i^t$  and  $\pmb{z}_i^{t + 1} = \pmb{z}_i^t$ . Note that we use independent samples for  $\pmb{y}_i^{t + 1}$  and  $\pmb{z}_i^{t + 1}$ .
Finally, we present E-NestedVR in Algorithm 1 where second-order extrapolation operator  $\mathcal{L}^{(2)}$  is applied to each occurrence of  $\nabla f_{i}$ . We now analyze its convergence guarantee. Our analysis works by first looking at the effect of multi-level variance reduction without the extrapolation (that we refer to as NestedVR, Theorem 10, Appendix E.2), and then showing how extrapolation could further help to drive down the sample complexity.
Theorem 5 [E-NestedVR Convergence] Consider the (FCCO) problem. Under the same assumptions as Theorem 3.
- If  $n = \mathcal{O}(\varepsilon^{-2/3})$ , then we choose the hyperparameters of  $E$ -Nested VR (Algorithm 1) as  $B_1 = B_2 = n, p_{out} = 1, S_1 = \tilde{L}_F^2 \varepsilon^{-2}, S_2 = \tilde{L}_F \varepsilon^{-1}, p_{in} = \tilde{L}_F^{-1} \varepsilon, \gamma = \mathcal{O}\left(\frac{1}{L_F}\right)$ .
- If  $n = \Omega(\varepsilon^{-2/3})$ , then we choose the hyperparameters of  $E$ -Nested VR as  $B_1 = n$ ,  $B_2 = \sqrt{n}$ ,  $p_{out} = 1 / \sqrt{n}$ ,  $S_1 = S_2 = \max \left\{ C_e C_g \varepsilon^{-1/2}, \tilde{L}_F^2 / (n \varepsilon^2) \right\}$ ,  $p_{in} = 1$ ,  $\gamma = \mathcal{O}\left(\frac{1}{L_F}\right)$ .
Then the output  $\pmb{x}^s$  of  $E$ -Nested VR satisfies:  $\mathbb{E}[\|\nabla F(\pmb{x}^s)\|_2^2] \leq \varepsilon^2$ , for nonconvex  $F$  with iterations
From Theorem 5, E-NestedVR has a sample complexity of  $\mathcal{O}(n\varepsilon^{-3})$  in the small  $n$  regime ( $n = \mathcal{O}(\varepsilon^{-2/3}))$  and  $\mathcal{O}(\max\{\sqrt{n}/\varepsilon^{2.5}, 1/\sqrt{n}\varepsilon^4\})$  in the large  $n$  regime ( $n = \Omega(\varepsilon^{-2/3}))$ ). Therefore, in the large  $n$  regime, this improves the  $\mathcal{O}(n\varepsilon^{-3})$  sample complexity of MSVR-V2 [19].
In this section, we demonstrate the numerical performance of our proposed algorithms. We focus on the application of invariant logistic regression here. In Appendix F, we discuss performance of our proposed algorithms on other common CSO/FCCO applications.
5.1 Application of Invariant Risk Minimization
Invariant learning has wide applications in machine learning and related areas [22, 1]. Invariant logistic regression [16] is formulated as follows:
$$
\min _ {\boldsymbol {x}} \mathbb {E} _ {\boldsymbol {\xi} = (\boldsymbol {a}, b)} [ \log (1 + \exp (- b \mathbb {E} _ {\boldsymbol {\eta} | \boldsymbol {\xi}} [ \boldsymbol {\eta} ] ^ {\top} \boldsymbol {x}) ],
$$
where  $\pmb{a}$  and  $b$  represent a sample and its corresponding label, and  $\eta$  is a noisy observation of the sample  $\pmb{a}$ . This first part can be considered as a CSO objective, with  $f_{\xi}(y) \coloneqq \log (1 + \exp (-by))$  and  $g_{\eta}(\pmb{x};\xi) \coloneqq \eta^{\top}\pmb{x}$ . As the loss  $f_{\xi} \in \mathcal{C}^{\infty}$  is smooth, our results from Sections 3 and 4 are applicable.
An  $\ell_2$ -regularizer is added to ensure the existence of an unique minimizer. Since the gradient of the penalization term is unbiased, we only have to consider the bias of the data-dependent term. We generate a synthetic dataset with  $d = 10$  dimensions. The minimizer is drawn from Gaussian distribution  $\pmb{x}^{\star} \sim \mathcal{N}(0,1) \in \mathbb{R}^{d}$ . We draw invariant samples  $\{(a_i, b_i)\}_i$  where  $a_i \sim \mathcal{N}(0,1) \in \mathbb{R}^d$  and compute  $b_i = \mathrm{sgn}(a_i^\top x^{\star})$  and its perturbed observation  $\eta \sim \mathcal{N}(a_i, 100) \in \mathbb{R}^d$ .
We consider drawing  $\xi$  from a large set  $(n = 50000)$  and a small set  $(n = 50)$  as CSO and FCCO problems respectively. As baselines, we implemented the BSGD and BSpiderBoost methods
(a) CSO: Large  $n = 10000$
(b) FCCO: small  $n = 50$
Figure 3: Performances of algorithms and their extrapolated versions on the instrumental variable regression task. The shaded region represents the  $95\%$  -confidence interval.
from [16], V-MLMC approach from [17], and NestedVR approach from Appendix E.2 which achieves the same complexity as MSVR-V2 [19] for the FCCO problem. RT-MLMC shares the same sample complexity as V-MLMC, and is thus omitted from the experiment [17]. The results are shown in Fig. 2. In the CSO setting, we compare biased gradient methods with their extrapolated variants (BSGD vs. E-BSGD, BSpiderBoost vs. E-BSpiderBoost, and NestedVR vs. E-NestedVR). The extrapolated versions of BSGD, BSpiderBoost, and NestedVR consistently reach lower error than their non-extrapolated counterparts, as is evident in Figure 2a. In this case, the performance of BSpiderBoost is similar to BSGD as also noted by the authors of these techniques [16], and a drawback of BSpiderBoost seems to be that it is much harder to tune in practice. However, it is clear that E-BSGD outperforms BSGD, and E-BSpiderBoost outperforms BSpiderBoost, respectively. In the FCCO setting, we compare extrapolation based methods and MLMC based methods. Figure 2a, shows that E-NestedVR outperforms all other extrapolated algorithms, including the V-MLMC approach of [17], matching our theoretical findings.
5.2 Application of Instrumental Variable Regression
Instrumental variable regression is a popular technique in econometrics that aims to estimate the causal effect of input  $X$  on a confounded outcome  $Y$  with an instrument variable  $Z$ , which, for a given  $X$ , is conditionally independent of  $Y$ . As noted by [23], the instrumental variable regression is a special case of the CSO problem. The instrumental variable regression problem is phrased as:
$$
\min _ {w} \mathbb {E} _ {\xi = (Y, Z)} \left[ \ell (Y, \mathbb {E} _ {X | Z} [ g _ {X} (w) ]) \right]
$$
where  $\xi = (Y,Z),\eta = X$ . This can be viewed in the CSO framework with  $f_{\xi}(\cdot) = \ell (Y,\cdot)$ . We choose  $\ell (y,\hat{y}) = \log \cosh (y - \hat{y})$  as regression loss function and  $g_{X}(w) = w^{\top}X$  to be a linear regressor. In this case,  $f_{\xi}\in \mathcal{C}^{\infty}$  with  $\nabla f_{\xi}(\hat{y}) = \tanh (\hat{y} -Y)$ , and our results from Sections 3 and 4 apply. We generate the data similar to [31]
$$
Z \sim \operatorname {U n i f o r m} \left(\left[ - 3, 3 \right] ^ {2}\right), \quad e \sim \mathcal {N} (0, 1), \quad \delta \sim \mathcal {N} (0, 0. 1), \quad \gamma \sim \text {E x p o n e n t i a l} (1 0)
$$
$$
X = \frac {1}{2} z _ {1} + \frac {1}{2} e + \gamma , Y = X + e + \delta
$$