SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Justin Weltz* Tanner Fiez† Eric Laber‡ Alexander Volfovsky§ Blake Mason ¶ Houssam Nassif ¶ Lalit Jain **
Most linear experimental design problems assume homogeneous variance, even though heteroskedastic noise is present in many realistic settings. Let a learner have access to a finite set of measurement vectors  $\mathcal{X} \subset \mathbb{R}^d$  that can be probed to receive noisy linear responses of the form  $y = x^\top \theta^* + \eta$ . Here  $\theta^* \in \mathbb{R}^d$  is an unknown parameter vector, and  $\eta$  is independent mean-zero  $\sigma_x^2$ -strictly-sub-Gaussian noise defined by a flexible heteroskedastic variance model,  $\sigma_x^2 = x^\top \Sigma^* x$ . Assuming that  $\Sigma^* \in \mathbb{R}^{d \times d}$  is an unknown matrix, we propose, analyze and empirically evaluate a novel design for uniformly bounding estimation error of the variance parameters,  $\sigma_x^2$ . We demonstrate the benefits of this method with two adaptive experimental design problems under heteroskedastic noise, fixed confidence transductive best-arm identification, and level-set identification; proving the first instance-dependent lower bounds in these settings. Lastly, we construct near-optimal algorithms and empirically demonstrate the large improvements in sample complexity gained from accounting for heteroskedastic variance in these designs.
Accurate data analysis pipelines require decomposing observed data into signal and noise components. To facilitate this, the noise is commonly assumed to be additive and homogeneous. However, in applied fields this assumption is often untenable, necessitating the development of novel design and analysis tools. Specific examples of deviations from homogeneity include the variability in sales figures of different products in e-commerce [24, 31], the volatility of stocks across sectors [5], spatial data [58, 43], and genomics [48, 10]. This heteroskedasticity is frequently structured as a function of observable features. In this work, we develop two near-optimal adaptive experimental designs that take advantage of such structure.
There are many approaches, primarily developed in the statistics and econometrics literature, for efficient mean estimation in the presence of heteroskedastic noise (see [56, 16] for a survey). These methods concentrate on settings where data have already been collected and do not consider sampling to better understand the underlying heteroskedasticity. Some early work exists on experimental design in the presence of heteroskedastic noise, however these methods are non-adaptive (see Section 2 for details). Importantly, using naive data collection procedures designed for homoskedastic environments may lead to sub-optimal inference downstream (see Section 1.2 for examples). A scientist may fail to take samples that allow them to observe the effect of interest with a fixed sampling budget, thus wasting time and money. This work demonstrates how to efficiently and adaptively collect data
*Department of Statistical Science, Duke University, justin.weltz@duke.edu, work conducted at Amazon
$^{\dagger}$ Amazon.com, USA, fieztann@amazon.com
$^{\ddagger}$ Department of Statistical Science, Duke University, eric.laber@duke.edu
$^{\S}$ Department of Statistical Science, Duke University, alexander.volfovsky@duke.edu
Amazon.com, USA, bjmason@amazon.com
Meta, USA, houssamn@meta.com
**Michael G. Foster School of Business, University of Washington, lalitj@uw.edu, work conducted at Amazon
in the presence of heteroskedastic noise. It bridges the gap between heteroskedasticity-robust (but potentially inefficient) passive data collection techniques, and powerful (but at times brittle) active algorithms that largely ignore the presence of differing variances.
1.1 Problem Setting
To ground our discussion of heteroskedasticity in active data collection, we focus on adaptive experimentation with covariates; also known as pure-exploration linear bandits [7, 45]. Let a learner be given a finite set of measurement vectors  $\mathcal{X} \subset \mathbb{R}^d$ ,  $\mathrm{span}(\mathcal{X}) = \mathbb{R}^d$ , which can be probed to receive linear responses with additive Gaussian noise of the form:
$$
y = x ^ {\top} \theta^ {*} + \eta \quad \text {w i t h} \quad \eta \sim \mathcal {N} \left(0, \sigma_ {x} ^ {2}\right), \quad \sigma_ {x} ^ {2} := x ^ {\top} \Sigma^ {*} x. \tag {1}
$$
Here,  $\theta^{*} \in \mathbb{R}^{d}$  is an unknown parameter vector,  $\eta$  is independent mean-zero Gaussian noise with variance  $\sigma_{x}^{2} = x^{\top}\Sigma^{*}x$ , and  $\Sigma^{*} \in \mathbb{R}^{d\times d}$  is an unknown positive semi-definite matrix with the assumption that  $\sigma_{\mathrm{min}}^{2} \leq \sigma_{x}^{2} \leq \sigma_{\mathrm{max}}^{2}$  for all  $x \in \mathcal{X}$ , where  $\sigma_{\mathrm{max}}^{2}, \sigma_{\mathrm{min}}^{2} > 0$  are known. We focus on Gaussian error for ease of exposition but extend to strictly sub-Gaussian noise in the Appendix. The linear model in Eq. (1) with unknown noise parameter  $\Sigma^{*}$  directly generalizes the standard multi-armed bandit to the case where the arms have different variances [4]. Additionally, this heteroskedastic variance structure arises naturally in linear mixed effects models [20], which are used in a variety of experimental settings to account for multi-level heterogeneity (patient, spatial and block design variation [41]). Throughout,  $\mathcal{Z} \subset \mathbb{R}^{d}$  will be a set (that need not be equal to  $\mathcal{X}$ ) over which our adaptive designs will be evaluated. While the theory holds true for  $\sigma_{\mathrm{min}}^{2} \leq \sigma_{x}^{2} \leq \sigma_{\mathrm{max}}^{2}$  in general, we consider the case where  $\sigma_{\mathrm{max}}^{2} = \max_{x \in \mathcal{X}} \sigma_{x}^{2}$  and  $\sigma_{\mathrm{min}}^{2} = \min_{x \in \mathcal{X}} \sigma_{x}^{2}$  so that  $\kappa := \sigma_{\mathrm{max}}^{2} / \sigma_{\mathrm{min}}^{2}$  summarizes the level of heteroskedasticity.
1.2 Heteroskedasticity Changes Optimal Adaptive Experimentation
To underscore the importance of adapting to heteroskedasticity, we consider the performance of RAGE [12], a best arm identification algorithm designed for homoskedastic settings. This algorithm does not account for differing variances, so it is necessary to upper bound  $\sigma_x^2 \leq \sigma_{\max}^2$  for each  $x$  to identify the best arm with probability  $1 - \delta$  for  $\delta \in (0,1)$ . In Fig. 1, we compare this approach to the optimal sampling allocation accounting for heteroskedasticity in a setting that is a twist on a standard benchmark [46] for best-arm identification algorithms. Let  $\mathcal{X} = \mathcal{Z} \subset \mathbb{R}^2$  and  $|\mathcal{X}| = 3$ . We take  $x_1 = e_1$  and  $x_2 = e_2$  to be the first and second standard bases and set  $x_3 = \{\cos(0.5), \sin(0.5)\}$ . Furthermore, we take  $\theta^* = e_1$  such that  $x_1$  has the highest reward and  $x_3$  has the second highest. In the homoskedastic case, an optimal algorithm will focus on sampling  $x_2$  because it is highly informative for distinguishing between  $x_1$  and  $x_3$  as shown in Fig 1b. However, if the errors are heteroskedastic, such a sampling strategy may not be optimal. If  $\sigma_2 \gg \sigma_1, \sigma_3$ , the optimal allocation for heteroskedastic noise focuses nearly all samples on  $x_1$  and  $x_3$  (as shown in Fig. 1a), which are more informative than  $x_2$  because they have less noise. Hence, ignoring heteroskedasticity and upper-bounding the variances leads to a suboptimal sampling allocation and inefficient performance. This effect worsens as the amount of heteroskedasticity increases, leading to an additional multiplicative factor of  $\kappa$  (cf., [12], Thm. 1) as demonstrated in Fig. 1c. In this important example, we see that the sample complexity of an algorithm that ignores heteroskedasticity suffers a linear dependence on  $\kappa$ , while a method that adapts to the heteroskedasticity in the data does not.
1.3 Paper Contributions
1. Finite Sample Noise Estimation Guarantees. We use an experimental design technique to construct estimates of the heteroskedastic noise variances with bounded maximum absolute error. This adaptive experimental design is a two-phase sample splitting procedure based on a pair of distinct G-optimal experimental designs [28], with bounds that scale favorably with the underlying problem dimension.
2. Near-Optimal Experimental Design Algorithms. We propose adaptive experimental design algorithms for transductive linear best-arm and level-set identification with heteroskedastic noise, that combine our variance estimation method with optimal experimental design techniques. We prove sample complexity upper bounds for unknown  $\Sigma^{*}$  that nearly match instance dependent lower bounds when  $\Sigma^{*}$  is known. Importantly, we show that in contrast to naive methods which
(a) Heteroskedastic design
(c) Sample complexity as  $\kappa$  increases
Figure 1: We compare the sampling strategies of algorithms that do and do not account for heteroskedastic variance in Figs. 1a and 1b respectively, when  $\kappa = 20$ . Figure 1c shows that the gap in the optimal sample complexity between the two algorithms becomes linear as  $\kappa$  increases.
(b) Homoskedastic design
have a multiplicative dependence on  $\kappa$ , we only suffer an additive dependence. Lastly, we show that the theoretical improvements translate to stronger empirical performance.
Experimental Design and Active Learning for Heteroskedastic Regression. Many experimental design methods for heteroskedastic regression assume known weights (or efficiency functions) [55, 54]. Procedures that estimate the heteroskedastic variance focus on optimizing the design for the weighted least squares estimator given a burn-in period used to estimate the noise parameters [37, 53, 52]. These contributions assume that the variance is unstructured, which leads to sample complexities that scale with the number of arms. Similarly, past active learning work on linear bandits with heteroskedastic noise either assume the variances are known [30, 62], or that the variances are unstructured [14, 8, 4]. The works that assume structured heteroskedastic noise in a linear bandit framework focus on active maximum likelihood estimation and G-optimal design [7, 45], but do not exploit the structure of the heteroskedastic noise to improve performance, leading to complexities scaling poorly in the problem dimension. Lastly, Mukherjee et al. [38] study policy value estimation under the reward model described in Eq. 1. However, their variance estimator is different from ours.
Best-Arm and Level Set Identification in Linear and Logistic Bandits. Best-arm identification is a fundamental problem in the bandit literature [13, 46, 18]. Soare [44] constructed the first passive and adaptive algorithms for pure-exploration in linear bandits using G-optimal and  $\mathcal{X}\mathcal{V}$ -allocation experimental designs. Fiez et al. [12] built on these results, developing the first pure exploration algorithm (RAGE) for the transductive linear bandit setting with nearly matching upper and lower bounds. Recent work specializes these methods to different settings [47, 57, 25, 26], improves time complexity [60, 50, 33], and targets different (minimax, asymptotically optimal, etc.) lower bound guarantees [59, 22, 9]. In the level set literature, Mason et al. [35] provides the first instance optimal algorithm with matching upper and lower bounds. Critically, none of these contributions account for heteroskedastic variance. Finally, we note the connection between this work and the logistic bandits literature [23, 11, 1, 36]. In that setting, the observations are Bernoulli random variables with a probability given by  $\mathbb{P}(y = 1|x) = \left\{1 + \exp (-x^{\top}\theta^{*})\right\}^{-1}$  for  $x\in \mathbb{R}^d$ ,  $\theta^{*}\in \mathbb{R}^{d}$  and  $y\in \{0,1\}$ . Because of the mean-variance relationship for Bernoulli random variables, points  $x$  for which  $\mathbb{P}(y = 1|x)$  is near 0 or 1 have lower variance than other  $x$ 's. While the tools are different in these two cases, we highlight two common ideas: First, Mason et al. [36] present a method that explicitly handles the heteroskedasticity in logistic bandits. Second, a core focus of optimal logistic bandit papers is reducing the dependence on the worst case variance to only an additive penalty, similar to the guarantee of  $\kappa$  herein, which has been shown to be unavoidable in general [23].
3 Heteroskedastic Variance Estimation
In this section, we describe an algorithm for heteroskedastic noise variance estimation with error bounds that scale favorably in the problem dimension. We then evaluate this method empirically.
Given a sampling budget  $\Gamma$ , our goal is to choose  $\{x_{t}\}_{t = 1}^{\Gamma}\subset \mathcal{X}$  to control the tail of the maximum absolute error (MAE):  $\max_{x\in \mathcal{X}}|\widehat{\sigma}_x^2 -\sigma_x^2 |$  by exploiting the heteroskedastic variance model. Alg. 1, HEAD (Heteroskedasticity Estimation by Adaptive Designs), provides an experimental design for finite-time bounds on structured heteroskedastic variance estimation.
Notation. Let  $M \coloneqq d(d + 1) / 2$ . Define  $\mathcal{W} = \{\phi_x, \forall x \in \mathcal{X}\} \subset \mathbb{R}^M$  where  $\phi_x = \operatorname{vec}(xx^\top)G$  and  $G \in \mathbb{R}^{d^2 \times M}$  is the duplication matrix such that  $G\operatorname{vech}(xx^\top) = \operatorname{vec}(xx^\top)$ , with  $\operatorname{vech}(\cdot)$  denoting the half-vectorization. For a set  $\mathcal{V}$ , let  $P_{\mathcal{V}} \coloneqq \{\lambda \in \mathbb{R}^{|\mathcal{V}|} : \sum_{v \in \mathcal{V}} \lambda_v = 1, \lambda_v \geq 0 \, \forall v\}$  denote distributions over  $\mathcal{V}$  and  $\mathcal{V}(\mathcal{V}) \coloneqq \{v' - v : v, v' \in \mathcal{V}, v \neq v'\}$  represent the differences between members of set  $\mathcal{V}$ .
General Approach. Given observations from  $y_{t} = x_{t}^{\top}\theta^{*} + \eta_{t}$  for  $t\in \{1,2,\dots ,\Gamma \}$  and the true value of  $\theta^{*}$ , we can estimate  $\sigma_x^2$  for each  $x\in \mathcal{X}$  from the squared residuals  $\eta_t^2 = (y_t - x_t^\top \theta^*)^2$  (e.g., using method of moments or another estimating equation approach [15]). However, since  $\theta^{*}$  is not known, we cannot directly observe  $\eta_t^2$ , making estimation of  $\sigma_x^2$  more difficult. Consider a general procedure that first uses  $\Gamma /2$  samples from the budget to construct a least-squares estimator for  $\theta^{*}$ ,  $\widehat{\theta}_{\Gamma /2}$ , and then uses the final  $\Gamma /2$  samples to collect error estimates of the form  $\widehat{\eta}_t^2 = (y_t - x_t^\top \widehat{\theta}_{\Gamma /2})^2$ . Define  $\Phi_{\Gamma}\in \mathbb{R}^{\Gamma /2\times M}$  as a matrix with rows  $\{\phi_{x_t}\}_{t = \Gamma /2 + 1}^{\Gamma}$  and  $\widehat{\eta}\in \mathbb{R}^{\Gamma /2}$  as a vector with values  $\{\widehat{\eta}_t\}_{t = \Gamma /2 + 1}^{\Gamma}$ . With  $\widehat{\eta}$ , we can estimate  $\Sigma^{*}$  by least squares as
$$
\operatorname {v e c h}(\widehat{\Sigma}_{\Gamma}) = \operatorname *{arg  min}_{\mathbf{s}\in \mathbb{R}^{M}}\| \Phi_{\Gamma}\mathbf{s} - \widehat{\eta}^{2}\|^{2},
$$
where we minimize over the space of symmetric matrices. Splitting samples ensures that the errors in estimating  $\theta^{*}$  and  $\Sigma^{*}$  are independent. Note that  $(\widehat{\Sigma}_{\Gamma},\widehat{\theta}_{\Gamma /2})$  is an M-estimator of  $(\Sigma^{*},\theta^{*})$  derived from unbiased estimating functions, a common approach in heteroskedastic regression [21, 51, 34]. Finally, by plugging in the estimate of  $\Sigma^{*}$ , one can obtain an estimate of the variance,  $\widehat{\sigma}_x^2 = \min \{\max \{x^\top \widehat{\Sigma_\Gamma} x,\sigma_{\min}^2\} ,\sigma_{\max}^2\}$ . We show that this approach leads to the following general error decomposition for any  $x\in \mathcal{X}$  with  $z\coloneqq \phi_x^\top (\Phi_\Gamma^\top \Phi_\Gamma)^{-1}\Phi_\Gamma^\top$
$$
|x^{\top}(\Sigma^{*} - \widehat{\Sigma}_{\Gamma})x|\leq \underbrace{\left\{x^{\top}\bigl(\widehat{\theta}_{\Gamma / 2} - \theta^{*}\bigr)\right\}^{2}}_{\mathbf{A}} + \underbrace{\left|\sum_{t = \Gamma / 2}^{\Gamma}z_{t}\left(\eta_{t}^{2} - x_{t}^{\top}\Sigma^{*}x_{t}\right)\right|}_{\mathbf{B}} + \underbrace{\left|\sum_{t = \Gamma / 2}^{\Gamma}2z_{t}x_{t}^{\top}(\widehat{\theta}_{\Gamma /2} - \theta^{*})\eta_{t}\right|}_{\mathbf{C}}.
$$
Experimental Design for Precise Estimation. The analysis above gives an error decomposition for any phased non-adaptive data collection method that splits samples between estimation of  $\theta^{*}$  and  $\Sigma^{*}$ . Intuitively, we wish to sample  $\{x_{1},\ldots ,x_{\Gamma /2}\}$  and  $\{x_{\Gamma /2 + 1},\dots,x_{\Gamma}\}$  to control the maximum absolute error by minimizing the expression above. We achieve this via two phases of G-optimal experimental design [28]- first over the  $\mathcal{X}$  space to estimate  $\theta^{*}$ , and then over the  $\mathcal{W}$  space to estimate  $\Sigma^{*}$ . Precisely, we control quantity A via stage 1 of Alg. 1, which draws from a G-optimal design  $\lambda^{*}$  to minimize the maximum predictive variance
$$
\max  _ {x \in \mathcal {X}} \mathbb {E} \left[ \left\{x ^ {\top} \left(\widehat {\theta} _ {\Gamma / 2} - \theta^ {*}\right) \right\} ^ {2} \right].
$$
Quantity  $\mathbf{B}$  is another sub-exponential quantity (Definition B.2 in Appendix B), which we control by drawing from a G-optimal design  $\lambda^{\dagger}$  over  $\mathcal{W}$  to minimize  $\max_{\phi \in \mathcal{W}}\phi^{\top}(\Phi_{\Gamma}^{\top}\Phi_{\Gamma})^{-1}\phi$ . Finally,  $\mathbf{C}$  is controlled by a combination of the guarantees associated with  $\lambda^{*}$  and  $\lambda^{\dagger}$ . This leads to the following error bound on  $\widehat{\sigma}_x^2 = \min \left\{\max \{x^\top \widehat{\Sigma_\Gamma} x,\sigma_{\min}^2\} ,\sigma_{\max}^2\right\}$ .
Theorem 3.1. Assume  $\Gamma = \Omega$  [max  $\{\sigma_{\mathrm{max}}^2\log (|\mathcal{X}| / \delta)d^2,d^2\} ]$  . For any  $x\in \mathcal{X}$  and  $\delta \in (0,1)$  Alg. 1 guarantees the following:
$$
\mathbb {P} \left(| \widehat {\sigma} _ {x} ^ {2} - \sigma_ {x} ^ {2} | \leq C _ {\Gamma , \delta}\right) \geq 1 - \delta / 2 \quad \text {w h e r e} \quad C _ {\Gamma , \delta} = \mathcal {O} \left\{\sqrt {\log (| \mathcal {X} | / \delta) \sigma_ {\max } ^ {2} d ^ {2} / \Gamma} \right\}.
$$
The proof of Theorem 3.1 is in Appendix D and follows from the decomposition in Eq. 3; treating  $\mathbf{A}$ ,  $\mathbf{B}$  and  $\mathbf{C}$  as sub-exponential random variables with bounds that leverage the experimental designs of Alg. 1. The details on the constants involved in Theorem 3.1 are also included in Appendix D.
Algorithm 1: HEAD (Heteroskedasticity Estimation by Adaptive Designs)
Result: Find  $\widehat{\Sigma}_{\Gamma}$   
1 Input: Arms  $\mathcal{X} \in \mathbb{R}^d, \Gamma \in \mathbb{N}$   
2 //Stage 1: Take half the samples to estimate  $\theta^{*}$   
3 Determine  $\lambda^{*}$  according to  $\lambda^{*} = \arg \min_{\lambda \in P_{\mathcal{X}}} \max_{x \in \mathcal{X}} x^{\top} \left( \sum_{x' \in \mathcal{X}} \lambda_{x'} x' x'^{\top} \right)^{-1} x$   
4 Pull arm  $x \in \mathcal{X}[\lambda_x^*\Gamma/2]$  times and collect observations  $\{x_t, y_t\}_{t=1}^{\Gamma/2}$   
5 Define  $A^{*} = \sum_{t=1}^{\Gamma/2} x_t x_t^{\top}$  and  $b^{*} = \sum_{t=1}^{\Gamma/2} x_t y_t$  and estimate  $\widehat{\theta}_{\Gamma/2} = A^{*-1} b$   
6 //Stage 2: Take half the samples to estimate  $\Sigma^{*}$  given  $\widehat{\theta}_{\Gamma/2}$   
7 Determine  $\lambda^{\dagger} \gets \lambda_x^{\dagger} = \arg \min_{\lambda \in P_{\mathcal{X}}} \max_{x \in \mathcal{X}} \phi_x^{\top} \left( \sum_{x' \in \mathcal{X}} \lambda_{x'} \phi_{x'} \phi_{x'}^{\top} \right)^{-1} \phi_x$   
8 Pull arm  $x \in \mathcal{X}[\lambda_x^{\dagger}\Gamma/2]$  times and collect observations  $\{x_t, y_t\}_{t=\Gamma/2+1}^{\Gamma}$   
9 Let  $A^{\dagger} = \sum_{t=\Gamma/2+1}^{\Gamma} \phi_{x_t} \phi_{x_t}^{\top}, b^{\dagger} = \sum_{t=\Gamma/2+1}^{\Gamma} \phi_{x_t} (y_t - x_t^{\top} \widehat{\theta}_{\Gamma/2})^2$   
10 Output:  $\text{vech}(\widehat{\Sigma}_{\Gamma}) = A^{\dagger^{-1}} b^{\dagger}$ .
Note that we design to directly control the MAE of  $\widehat{\sigma}_x^2$ . This is in contrast to an inefficient alternative procedure that allocates samples to minimize the error of  $\widehat{\Sigma}_{\Gamma}$  in general, and then extends this bound to  $x^\top \widehat{\Sigma}_{\Gamma}x$  for  $x \in \mathcal{X}$ .
Theoretical Comparison of Variance Estimation Methods. In contrast to previous approaches that naively scale in the size of  $\mathcal{X}$  [14], the above result has a dimension dependence of  $d^2$ , which intuitively scales with the degrees of freedom of  $\Sigma^{*}$ . To highlight the tightness of this result, consider estimating the noise parameters for each arm,  $x \in \mathcal{X}$ , only using the samples of  $x$  (like the method in [14]). We improve this approach by adapting it to our heteroskedastic variance structure and strategically pick  $d^2$  points to avoid the dependence on  $|\mathcal{X}|$ . We refer to this method as the Separate Arm Estimator, and in Appendix D, we show that it suffers a dependence of  $d^4$ . This comparison point, along with our empirical results below, suggest that the error bounds established in Theorem 3.1 scale favorably with the problem dimension.
Empirical Comparison of Variance Estimation Methods. Define two sets of arms  $\mathcal{X}_1, \mathcal{X}_2$  such that  $x \in \mathcal{X}_1$  is drawn uniformly from a unit sphere and  $x \in \mathcal{X}_2$  is drawn uniformly from a sphere with radius 0.1; with  $|\mathcal{X}_1| = 200$  and  $|\mathcal{X}_2| = 1800$ . Let  $\Sigma^* = \mathrm{diag}(1, 0.1, 1, 0.1, \ldots) \in \mathbb{R}^{d \times d}$ , and  $\theta^* = (1, 1, \ldots, 1) \in \mathbb{R}^d$ . This setting provides the heteroskedastic variation needed to illustrate the advantages of using HEAD (Alg. 1). An optimal algorithm will tend to target orthogonal vectors in  $\mathcal{X}_1$  because these arms have higher magnitudes in informative directions. Defining  $\mathcal{X} = \mathcal{X}_1 \cup \mathcal{X}_2$ , we perform 32 simulations in which we randomly sample the arm set and construct estimators for each  $\sigma_x^2$ ,  $x \in \mathcal{X}$ . We compare HEAD, the Separate Arm Estimator, and the Uniform Estimator (see Alg. 3 in Appendix D). The Uniform Estimator is based on Alg. 1 but samples arms uniformly from  $\mathcal{X}$  and does not split samples between estimation of  $\theta^*$  and  $\Sigma^*$ . HEAD should outperform its competitors in two ways: 1) optimally allocating samples, and 2) efficiently sharing information across arms in estimation. The Uniform Estimator exploits the problem structure for estimation but not when sampling. In contrast, the Seperate Arm Estimator optimally samples, but does not use the relationships between arms for estimation. Fig. 2 depicts the average MAE and standard errors for each estimator. In Fig. 2a, we see that HEAD outperforms its competitors over a series of sample sizes for  $d = 15$ . Fig. 2b compares the estimators over a range of dimensions for a sample size of  $95k$ . While HEAD continues to accurately estimate the heteroskedastic noise parameters at high dimensions, the Separate Arm Estimator error scales poorly with dimension as the analysis suggests.
4 Adaptive Experimentation with Covariates in Heteroskedastic Settings
As an application of the novel sampling procedure and resulting variance estimation bounds developed in section 3, we now study adaptive experimentation with covariates in the presence of heteroskedastic
As  $\lambda_x\Gamma /2\notin \mathbb{N}$ , it is not possible in reality to take  $\lambda_{x}\Gamma /2$  samples. This problem can easily be handled by rounding procedures detailed in Appendix A.
(b)
Figure 2: The proposed experimental design estimator outperforms competitors in terms of maximum absolute error  $\max_{x\in \mathcal{X}}|\sigma_x^2 -\widehat{\sigma}_x^2 |$  over a range of sample sizes (a) and dimensions (b).