Title: Experimental Designs for Heteroskedastic Variance
Venue: NeurIPS 2023
Accepted: yes
Keywords: Heteroskedastic Variance

Linear Bandits

Experimental design
Abstract: Most linear experimental design problems assume homogeneous variance, while the presence of 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^{\ast}+\eta$. 
Here $\theta^{\ast}\in \mathbb{R}^d$ is an unknown parameter vector, and $\eta$ is independent mean-zero $\sigma_x^2$-sub-Gaussian noise defined by a flexible heteroskedastic variance model, $\sigma_x^2 = x^{\top}\Sigma^{\ast}x$. Assuming that $\Sigma^{\ast}\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 this method on two adaptive experimental design problems under heteroskedastic noise, fixed confidence transductive best-arm identification and level-set identification and prove the first instance-dependent lower bounds in these settings.
Lastly, we construct near-optimal algorithms and demonstrate the large improvements in sample complexity gained from accounting for heteroskedastic variance in these designs empirically.

Experimental Designs for Heteroskedastic Variance

Abstract

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 X ⊂ R d that can be probed to receive noisy linear responses of the form y = x ⊤ θ ∗ + η . Here θ ∗ ∈ R d is an unknown parameter vector, and η is independent mean-zero σ 2 x -strictly-sub-Gaussian noise defined by a flexible heteroskedastic variance model, σ 2 x = x ⊤ Σ ∗ x . Assuming that Σ ∗ ∈ R d × d is an unknown matrix, we propose, analyze and empirically evaluate a novel design for uniformly bounding estimation error of the variance parameters, σ 2 x . We demonstrate the benefits of this method with two adaptive experimental design problems under heteroskedastic noise, fixed confidence transductive bestarm 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.

1 Introduction

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

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 X ⊂ R d , span( X ) = R d , which can be probed to receive linear responses with additive Gaussian noise of the form:

Here, θ ∗ ∈ R d is an unknown parameter vector, η is independent mean-zero Gaussian noise with variance σ 2 x = x ⊤ Σ ∗ x , and Σ ∗ ∈ R d × d is an unknown positive semi-definite matrix with the assumption that σ 2 min ≤ σ 2 x ≤ σ 2 max for all x ∈ X , where σ 2 max , σ 2 min > 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 Σ ∗ 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, Z ⊂ R d will be a set (that need not be equal to X ) over which our adaptive designs will be evaluated. While the theory holds true for σ 2 min ≤ σ 2 x ≤ σ 2 max in general, we consider the case where σ 2 max = max x ∈X σ 2 x and σ 2 min = min x ∈X σ 2 x so that κ := σ 2 max /σ 2 min 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 σ 2 x ≤ σ 2 max for each x to identify the best arm with probability 1 -δ for δ ∈ (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 X = Z ⊂ R 2 and |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 θ ∗ = 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 σ 2 ≫ σ 1 , σ 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 upperbounding 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 κ (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 κ , while a method that adapts to the heteroskedasticity in the data does not.

1.3 Paper Contributions

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.

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 Σ ∗ that nearly match instance dependent lower bounds when Σ ∗ is known. Importantly, we show that in contrast to naive methods which

x

1

x

2

x

3

















(a) Heteroskedastic design

x

1

x

2

x

3























(b) Homoskedastic design































(c) Sample complexity as

κ

increases

have a multiplicative dependence on κ , we only suffer an additive dependence. Lastly, we show that the theoretical improvements translate to stronger empirical performance.

2 Related Work

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 XY -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 P ( y = 1 | x ) = { 1 + exp( -x ⊤ θ ∗ ) } -1 for x ∈ R d , θ ∗ ∈ R d and y ∈ { 0 , 1 } . Because of the mean-variance relationship for Bernoulli random variables, points x for which 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 κ 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 Γ , our goal is to choose { x t } Γ t =1 ⊂ X to control the tail of the maximum absolute error (MAE): max x ∈X | ̂ σ 2 x -σ 2 x | 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 := d ( d +1) / 2 . Define W = { ϕ x , ∀ x ∈ X} ⊂ R M where ϕ x = vec( xx ⊤ ) G and G ∈ R d 2 × M is the duplication matrix such that G vech( xx ⊤ ) = vec( xx ⊤ ) , with vech( · ) denoting the half-vectorization. For a set V , let P V := { λ ∈ R |V| : ∑ v ∈V λ v = 1 , λ v ≥ 0 ∀ v } denote distributions over V and Y ( V ) := { v ′ -v : v, v ′ ∈ V , v = v ′ } represent the differences between members of set V .

General Approach. Given observations from y t = x ⊤ t θ ∗ + η t for t ∈ { 1 , 2 , . . . , Γ } and the true value of θ ∗ , we can estimate σ 2 x for each x ∈ X from the squared residuals η 2 t = ( y t -x ⊤ t θ ∗ ) 2 (e.g., using method of moments or another estimating equation approach [15]). However, since θ ∗ is not known, we cannot directly observe η 2 t , making estimation of σ 2 x more difficult. Consider a general procedure that first uses Γ / 2 samples from the budget to construct a least-squares estimator for θ ∗ , ̂ θ Γ / 2 , and then uses the final Γ / 2 samples to collect error estimates of the form ̂ η 2 t = ( y t -x ⊤ t ̂ θ Γ / 2 ) 2 . Define Φ Γ ∈ R Γ / 2 × M as a matrix with rows { ϕ x t } Γ t =Γ / 2+1 and ̂ η ∈ R Γ / 2 as a vector with values { ̂ η t } Γ t =Γ / 2+1 . With ̂ η , we can estimate Σ ∗ by least squares as

where we minimize over the space of symmetric matrices. Splitting samples ensures that the errors in estimating θ ∗ and Σ ∗ are independent. Note that ( ̂ Σ Γ , ̂ θ Γ / 2 ) is an M-estimator of (Σ ∗ , θ ∗ ) derived from unbiased estimating functions, a common approach in heteroskedastic regression [21, 51, 34]. Finally, by plugging in the estimate of Σ ∗ , one can obtain an estimate of the variance, ̂ σ 2 x = min { max { x ⊤ ̂ Σ Γ x, σ 2 min } , σ 2 max } . We show that this approach leads to the following general error decomposition for any x ∈ X with z := ϕ ⊤ x (Φ ⊤ Γ Φ Γ ) -1 Φ ⊤ Γ :

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 θ ∗ and Σ ∗ . Intuitively, we wish to sample { x 1 , . . . , x Γ / 2 } and { x Γ / 2+1 , . . . , x Γ } 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 X space to estimate θ ∗ , and then over the W space to estimate Σ ∗ . Precisely, we control quantity A via stage 1 of Alg. 1, which draws from a G-optimal design λ ∗ to minimize the maximum predictive variance

Quantity B is another sub-exponential quantity (Definition B.2 in Appendix B), which we control by drawing from a G-optimal design λ † over W to minimize max ϕ ∈W ϕ ⊤ (Φ ⊤ Γ Φ Γ ) -1 ϕ . Finally, C is controlled by a combination of the guarantees associated with λ ∗ and λ † . This leads to the following error bound on ̂ σ 2 x = min { max { x ⊤ ̂ Σ Γ x, σ 2 min } , σ 2 max } .

Theorem 3.1. Assume Γ = Ω [ max { σ 2 max log( |X| /δ ) d 2 , d 2 }] . For any x ∈ X and δ ∈ (0 , 1) , Alg. 1 guarantees the following:

The proof of Theorem 3.1 is in Appendix D and follows from the decomposition in Eq. 3; treating A , B and 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.

Note that we design to directly control the MAE of ̂ σ 2 x . This is in contrast to an inefficient alternative procedure that allocates samples to minimize the error of ̂ Σ Γ in general, and then extends this bound to x ⊤ ̂ Σ Γ x for x ∈ X .

Theoretical Comparison of Variance Estimation Methods. In contrast to previous approaches that naively scale in the size of X [14], the above result has a dimension dependence of d 2 , which intuitively scales with the degrees of freedom of Σ ∗ . To highlight the tightness of this result, consider estimating the noise parameters for each arm, x ∈ 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 |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 X 1 , X 2 such that x ∈ X 1 is drawn uniformly from a unit sphere and x ∈ X 2 is drawn uniformly from a sphere with radius 0 . 1 ; with |X 1 | = 200 and |X 2 | = 1800 . Let Σ ∗ = diag(1 , 0 . 1 , 1 , 0 . 1 , . . . ) ∈ R d × d , and θ ∗ = (1 , 1 , . . . , 1) ∈ 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 X 1 because these arms have higher magnitudes in informative directions. Defining X = X 1 ∪ X 2 , we perform 32 simulations in which we randomly sample the arm set and construct estimators for each σ 2 x , x ∈ 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 X and does not split samples between estimation of θ ∗ and Σ ∗ . 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 95 k . 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

(a)

(b)

HEADEstimator

SeparateArmEstimator

UniformEstimator

0.6

0.4

Error

Error

0.5

MaxAbsolute

MaxAbsolute

0.3

0.4

0.2

0.3

0.2

0.1

0.1

0.0

40000

50000

60000

70000

80000

90000

3

6

9

12

15

Sample Count

Dimension

noise through the lens of transductive linear identification problems [12]. Recalling the problem setting from Section 1.1, we consider the following fixed confidence identification objectives:

Best-Arm Identification ( BAI ). Identify the singleton set { z ∗ } where z ∗ = arg max z ∈Z z ⊤ θ ∗ .

Level-Set Identification ( LS ). Identify the set G α = { z : z ⊤ θ ∗ > α } given a threshold α ∈ R .

This set of objectives, together with a learning protocol, allow us to capture a number of practical problems of interest where heteroskedastic noise naturally arises in the assumed form. Consider multivariate testing problems in e-commerce advertisement applications [17, 39]. The expected feedback is modeled through a linear combination of primary effects from content variation dimensions (e.g., features or locations that partition the advertisement), and interaction effects between content variation in pairs of dimensions. The noise structure is naturally heteroskedastic and dependent on the combination of variation choices. Moreover, the transductive objectives allow for flexible experimentation such as identifying the optimal combination of variations or the primary effects that exceed some threshold. We return to this setting in Experiment 3 of Section 5. In the remainder of this section, we characterize the optimal sample complexities for the heteroskedastic transductive linear identification problems and design algorithms that nearly match these lower bounds.

4.1 Linear Estimation Methods with Heteroskedastic Noise

To draw inferences in the experimentation problems of interest, it is necessary to consider estimation methods for the unknown parameter vector θ ∗ ∈ R d . Given a matrix of covariates X T ∈ R T × d , a vector of observations Y T ∈ R T , and a user-defined weighting matrix W T = diag( w 1 , . . . , w T ) ∈ R T × T , the weighted least squares (WLS) estimator is defined as:

Let Σ T = diag( σ 2 x 1 , σ 2 x 2 , . . . , σ 2 x T ) ∈ R T × T denote the variance parameters of the noise observations. The WLS estimator is unbiased regardless of the weight selections, with variance

The WLS estimator with W T = Σ -1 T is the minimum variance linear unbiased estimator [2]. In contrast, the ordinary least squares (OLS) estimator, obtained by taking W T = I in the WLS estimator, is unbiased but with a higher variance,

Consequently, WLS is a natural choice with heteroskedastic noise for sample efficient estimation.

4.2 Optimal Sample Complexities

We now motivate the optimal sample complexity for any algorithm that returns the correct quantity with a fixed probability greater than 1 -δ , δ ∈ (0 , 1) , in transductive identification problems with

heteroskedastic noise. We seek to identify a set of vectors H OBJ , where

Algorithms consist of a sampling strategy and a rule to return a set ̂ H at time τ .

Definition 4.1. An algorithm is δ -PAC for OBJ ∈ { BAI , LS } if ∀ θ ∈ Θ , P θ ( ̂ H = H OBJ ) ≥ 1 -δ .

Throughout, we will only consider algorithms that are δ -PAC. For drawing parallels between best-arm and level set identification, it will be useful to define a set of values Q OBJ , where

̸

and a constant b OBJ such that b BAI = 0 and b LS = α . Both problems amount to verifying that for all h ∈ H OBJ and q ∈ Q OBJ , ( h -q ) ⊤ ̂ θ WLS > b OBJ , or equivalently ( h -q ) ⊤ θ ∗ -b OBJ -( h -q ) ⊤ ( θ ∗ -̂ θ WLS ) > 0 . Using a sub-Gaussian tail bound and a union bound [32], this verification requires the following to hold for all h ∈ H OBJ and q ∈ Q OBJ for the data collected:

Let λ x denote the proportion of T samples given to a measurement vector x ∈ X so that when W T = Σ -1 T in Eq. (4.1),

as in Eq. (2). Now, we reformulate the condition in Eq. (3) and minimize over designs λ ∈ P X to see that δ -PAC verification requires

This motivating analysis gives rise to the following sample complexity lower bound.

Theorem 4.1. Consider an objective, OBJ , of best-arm identification ( BAI ) or level-set identification ( LS ). For any δ ∈ (0 , 1) , a δ -PAC algorithm must satisfy

Remark 4.1. This lower bound assumes that the noise variance parameters are known . We compare to the existing lower bounds for the identification objectives with homoskedastic noise [12, 44, 35] in Appendix G. These apply in this setting by taking the variance parameter σ 2 = σ 2 max and are characterized by an instance-dependent parameter, ρ ∗ OBJ , such that κ -1 ≤ ψ ∗ OBJ /ρ ∗ OBJ ≤ 1 (recall that κ := σ 2 max /σ 2 min ). In general, κ can be quite large in many problems with heteroskedastic noise. Consequently, this implies that near-optimal algorithms for linear identification with homoskedastic noise can be highly suboptimal (by up to a multiplicative factor of κ ) in this setting.

4.3 Adaptive Designs with Unknown Heteroskedastic Variance

We now present Alg. 2, H-RAGE (Heteroskedastic Randomized Adaptive Gap Elimination), as a solution for adaptive experimentation with covariates in heteroskedastic settings. The approach is composed of two components: 1) Obtain accurate estimates of the unknown variance parameters { σ 2 x } x ∈X using HEAD ; 2) Use adaptive experimental design methods for minimizing the uncertainty in directions of interest given the estimated variance parameters. Denote by ∆ = min h ∈H OBJ min q ∈Q OBJ ( h -q ) ⊤ θ ∗ -b OBJ the minimum gap for the objective and assume max z ∈Z ℓ |⟨ z -z ∗ , θ ∗ ⟩| ≤ 2 . We prove the following guarantees for Alg. 2.

Theorem 4.2. Consider an objective, OBJ , of best-arm identification ( BAI ) or level-set identification ( LS ). The set returned from Alg. 2 at a time τ is δ -PAC and

Remark 4.2. This result shows that Alg. 2 is nearly instance-optimal. Critically, the impact of not knowing the noise variances ahead of time only has an additive impact on the sample complexity relative to the lower bound. In contrast, existing near-optimal methods for identification problems with homoskedastic noise would be suboptimal by a multiplicative factor depending on κ relative to the lower bound. Given that the lower bound assumes the variances are known, an interesting question for future work is a tighter lower bound assuming the variance parameters are unknown.

Algorithm Overview. Observe that for ψ ∗ OBJ , the optimal allocation depends on both the unknown gaps (( h -q ) ⊤ θ ∗ -b OBJ ) for h ∈ H OBJ and q ∈ Q OBJ , and the unknown noise variance parameters { σ 2 x } x ∈X . To handle this, the algorithm begins with an initial burn-in phase using the procedure from Alg. 1 to produce estimates { ̂ σ 2 x } x ∈X of the unknown parameters { σ 2 x } x ∈X , achieving a multiplicative error bound of the form | σ 2 x -̂ σ 2 x | ≤ σ 2 x / 2 for all x ∈ X . From this point, the estimates { ̂ σ 2 x } x ∈X are used as plug-ins to the experimental designs and the weighted least squares estimates. In each round ℓ ∈ N of the procedure, we maintain a set of uncertain items Z ℓ , and obtain the sampling distribution ̂ λ ℓ , by minimizing the maximum predictive variance of the WLS estimator over all directions y = h -q for h ∈ H OBJ and q ∈ Q OBJ . This is known as an XY -allocation and G -optimal-allocation in the best-arm and level set identification problems respectively. Critically, the number of samples taken in round ℓ guarantees that the error in all estimates is bounded by ϵ ℓ , such that for any h ∈ H OBJ and q ∈ Q OBJ , ( h -q ) ⊤ θ ∗ > 2 ϵ ℓ is eliminated from the active set at the end of the round. By progressively honing in on the optimal item set H OBJ , the sampling allocation gets closer to approximating the oracle allocation.

5 Experiments

We now present experiments focusing on the transductive linear bandit setting. We compare H-RAGE and its corresponding oracle to their homoskedastic counterparts, RAGE [12] and the oracle allocation

with noise equal to σ 2 max . RAGE is provably near-optimal in the homoskedastic setting and known to perform well empirically. The algorithm is similar in concept to Alg. 2, but lacks the initial exploratory phase to estimate the variance parameters. We include the pseudo-code of RAGE in Appendix H. The homoskedastic and heteroskedastic oracles sample with respect to the optimal design with a sub-Gaussian interval stopping condition. All algorithms are run at a confidence level of δ = 0 . 05 , and we use the Franke-Wolfe method to compute the designs. We demonstrate that H-RAGE has superior empirical performance over a range of settings.

Linear bandits: Experiment 1 (Signal-to-Noise). We begin by presenting an experiment that is a variation of the standard benchmark problem for BAI in linear bandits [46] introduced in Sec. 1.2. Define e i ∈ R d as the standard basis vectors in d dimensions. In the standard example, X = Z = { e 1 , e 2 , x ′ } where x ′ = { e 1 cos( ω ) + e 2 sin( ω ) } for ω → 0 , and θ ∗ = e 1 . Thus x 1 = e 1 is optimal, x 3 = x ′ is near-optimal, and x 2 = e 2 is informative to discriminate between x 1 and x 3 . We consider an extension of this benchmark in multiple dimensions that highlights the performance gains from accounting for heteroskedastic variance. First, define an arm set H that is comprised of the standard example but with varied magnitudes. For q ∈ (0 , 1) ,

Then define G to be a series of vectors such that the rank of span ( { ϕ x : x ∈ G ∪ H} ) = R M . G allows for variance estimation. Finally, let the arm and item set be X = Z = G ∪ H and θ ∗ = e 1 .

We define Σ ∗ as the d -dimensional identity matrix. Observe that x 1 = e 1 is optimal, the d -1 arms given by { e 1 cos( ω ) + e i sin( ω ) } d i =2 are near-optimal with equal gaps ∆ = 1 -cos( ω ) , and the d -1 arms given by { e 2 } ∪ q ∗ { e i } d i =3 are informative to sample for discriminating between x 1 = e 1 and x i = e 1 cos( ω ) + e i sin( ω ) for each i ∈ { 2 , . . . , d } , respectively. To identify x 1 , we must estimate the ∆ gap adequately in every dimension with arms { e 2 } ∪ q ∗ { e i } d i =3 . Therefore, performing BAI without accounting for heteroskedastic variances will tend toward a design that prioritizes sampling informative arms { q × e i } d i =3 , since these have smaller magnitudes and seem less informative about the near-optimal arms in dimensions i ∈ { 3 , 4 . . . d } . However, the signal to noise ratio is actually constant across arms (because the variance of q ∗ { e i } d i =3 is equal to q 2 ) leading to an oracle distribution that equally samples across { e 2 } ∪ q ∗ { e i } d i =3 when we account for heteroskedastic variances. This change in allocation is illustrated in Appendix H for d = 4 , ω = 0 . 01 , q = 0 . 4 and G = { 0 . 5( e 1 + e 2 ) + 0 . 1 e 3 , 0 . 5( e 1 + e 2 ) + 0 . 1( e 3 + e 4 ) } . In the same setting, Fig. 3a shows that the heteroskedastic variance experimental design and the tighter predictive confidence intervals of the weighted least squares estimator result in a large decrease in the sample size needed to identify the best arm.

Experiment 2: Benchmark. This experiment is also similar to the benchmark example introduced in Sec. 1.2. In this variation, the informative arms have the same magnitude but are bent at a 45 ◦ angle,

Let the arm and item set be X = Z and θ ∗ = e 1 . In Experiment 1, we have many near-optimal arms and H-RAGE reveals that each of these are equally difficult to distinguish from the best. In contrast, Experiment 2 contains one near-optimal arm along with many potentially informative arms, and we are interested in identifying the one that is most helpful. Intuitively, Experiment 1 uses heteroskedastic variance estimates to assess the difficulty of many problems, while Experiment 2 uses the same information to assess the benefits of many solutions. Let the unknown noise matrix be given by Σ ∗ = diag( σ 2 1 , σ 2 2 , . . . , σ 2 d ) where σ 2 1 = α 2 , ( σ 2 2 , σ 2 3 ) = β 2 and ( σ 2 4 , σ 2 5 . . . , σ 2 d ) = α 2 . x 1 = e 1 is again optimal and x 2 = { e 1 cos( ω ) + e 2 sin( ω ) } is a near optimal arm. In this case, { e 2 / √ 2 + e 1 / √ 2 } ∪ { e 2 / √ 2 + e i / √ 2 } d i =3 are informative for distinguishing between x 1 and x 2 , and the oracle allocation assuming homoskedastic noise (Homoskedastic Oracle) picks three of these vectors to sample. However, if α 2 ≫ β 2 , then it is optimal to sample { e 2 / √ 2 + e 3 / √ 2 , e 3 } over other potential informative arm combinations. Appendix H shows this contrast in allocations for d = 3 , ω = 0 . 01 , α 2 = 1 and β 2 = 0 . 2 . In the same setting, Fig. 3b shows that estimating and accounting for heteroskedastic noise contributes to a reduction in sample complexity.

Experiment 3: Multivariate Testing. In this experiment, we return to the motivating example of Section 4, multivariate testing in e-commerce advertising [17, 39]. We divide an advertisement into natural locations or features called dimensions, each of which has different content options or variations. We induce heteroskedastic variance by allowing both the expected value and the variance

(a) Experiment 1

(b) Experiment 2

(c) Experiment 3

Heteroskedastic-RAGE

RAGE

HeteroskedasticOracle

HomoskedasticOracle

1e7

2.00

1e7

1e7

SampleComplexity

8

SampleComplexity

1.75

SampleComplexity

5

1.50

4

6

1.25

3

1.00

4

0.75

0.50

2

1

0.25

0.00

Algorithm

Algorithm

Algorithm

of an advertisement's reward to depend on the n variations in each of the m dimensions. Let X = Z denote the set of layouts corresponding to the combinations of variant choices in the dimensions so that |X| = n m . The multivariate testing problem is often modeled in the linear bandit framework using the expected feedback for any layout x ∈ X ⊂ { 0 , 1 } d , where d = 1 + mn + n 2 m ( m +1) / 2 , given by

In this model, x i,j ∈ { 0 , 1 } is an indicator for variation j ∈ { 1 , . . . , n } being placed in dimension i ∈ { 1 , . . . , m } , and ∑ n j =1 x i,j = 1 for all i ∈ { 1 , . . . , n } for any layout x ∈ X . Observe that the expected feedback for a layout is modeled as a linear combination of a common bias term, primary effects from the variation choices within each dimension, and secondary interaction effects from the combinations of variation choices between dimensions.

We consider an environment where the effect of variation changes in one dimension, call it dimension j , has much higher variance than others; resulting in the best variation for dimension j being harder to identify. An algorithm that accounts for heteroskedastic variance will devote a greater number of samples to compare variations in dimension j , whereas an algorithm that upper bounds the variance will split samples equally between dimensions. For a simulation setting with 2 variations in 3 dimensions, we define Σ ∗ = diag(0 . 3 , 0 . 7 , 10 -3 , 10 -3 , . . . ) ∈ R 7 × 7 and θ ∗ = (0 , 0 . 005 , 0 . 0075 , 0 . 01 , -0 . 1 , -0 . 1 , . . . ) ∈ R 7 . This simulation setting implies that the second variation in each of the dimensions is better than the first, but the positive effect in the first dimension is hardest to identify. In Appendix H we can see that this causes the heteroskedastic oracle to devote more weight than the homoskedastic oracle to vectors that include the second variation in the first dimension. Fig. 3c depicts the improvement in sample complexity resulting from accounting for the heteroskedastic variances.

6 Conclusion

This paper presents an investigation of online linear experimental design problems with heteroskedastic noise. We propose a two-phase sample splitting procedure for estimating the unknown noise variance parameters based on G-optimal experimental designs and show error bounds that scale efficiently with the dimension of the problem. The proposed approach is then applied to fixed confidence transductive best-arm and level set identification with heteroskedastic noise, where we present instance-dependent lower bounds with provably near-optimal algorithms.

This paper leads to a number of interesting future research directions. For example, there may be situations where the noise structure is not fully characterized by Σ ∗ . Future work can explore the development of methods that handle more general noise structures (including correlated or heavytailed noise) while remaining near-optimal. Furthermore, it is possible that an integrated approach to variance estimation and adaptive experimentation with covariates would be more sample efficient than our two-step method. Lastly, it would be interesting to extend our algorithm to nonlinear response models.

References

Marc Abeille, Louis Faury, and Clément Calauzènes. Instance-wise minimax-optimal algorithms for logistic bandits. In International Conference on Artificial Intelligence and Statistics , pages 3691-3699. PMLR, 2021.

Alexander C Aitken. IV.- On least squares and linear combination of observations. Proceedings of the Royal Society of Edinburgh , 55:42-48, 1936.

Zeyuan Allen-Zhu, Yuanzhi Li, Aarti Singh, and Yining Wang. Near-optimal discrete optimization for experimental design: A regret minimization approach. Mathematical Programming , 186:439-478, 2021.

András Antos, Varun Grover, and Csaba Szepesvári. Active learning in heteroscedastic noise. Theoretical Computer Science , 411(29-30):2712-2728, 2010.

David Blitz and Pim Van Vliet. The volatility effect: Lower risk without lower return. Journal of portfolio management , pages 102-113, 2007.

Romain Camilleri, Kevin Jamieson, and Julian Katz-Samuels. High-dimensional experimental design and kernel bandits. In International Conference on Machine Learning , pages 1227-1237. PMLR, 2021.

Kamalika Chaudhuri, Prateek Jain, and Nagarajan Natarajan. Active heteroscedastic regression. In International Conference on Machine Learning , pages 694-702. PMLR, 2017.

Wesley Cowan, Junya Honda, and Michael N Katehakis. Normal bandits of unknown means and variances. Journal of Machine Learning Research , 18(154):1-28, 2018.

Rémy Degenne, Pierre Ménard, Xuedong Shang, and Michal Valko. Gamification of pure exploration for linear bandits. In International Conference on Machine Learning , pages 24322442. PMLR, 2020.

Bianca Dumitrascu, Gregory Darnell, Julien Ayroles, and Barbara E Engelhardt. Statistical tests for detecting variance effects in quantitative trait studies. Bioinformatics , 35(2):200-210, 2019.

Louis Faury, Marc Abeille, Clément Calauzènes, and Olivier Fercoq. Improved optimistic algorithms for logistic bandits. In International Conference on Machine Learning , pages 3052-3060. PMLR, 2020.

Tanner Fiez, Lalit Jain, Kevin G Jamieson, and Lillian Ratliff. Sequential experimental design for transductive linear bandits. Advances in neural information processing systems , 32:10667-10677, 2019.

Tanner Fiez, Sergio Gamez, Arick Chen, Houssam Nassif, and Lalit Jain. Adaptive experimental design and counterfactual inference. In Workshops of Conference on Recommender Systems , 2022.

Xavier Fontaine, Pierre Perrault, Michal Valko, and Vianney Perchet. Online a-optimal design and active linear regression. In International Conference on Machine Learning , pages 33743383. PMLR, 2021.

VP Godambe and ME Thompson. Some aspects of the theory of estimating equations. Journal of Statistical Planning and Inference , 2(1):95-104, 1978.

William H Greene. Econometric analysis . Pearson Education India, 2003.

Daniel N Hill, Houssam Nassif, Yi Liu, Anand Iyer, and SVN Vishwanathan. An efficient bandit algorithm for realtime multivariate optimization. In International Conference on Knowledge Discovery and Data Mining , pages 1813-1821, 2017.

Matthew Hoffman, Bobak Shahriari, and Nando Freitas. On correlation and budget constraints in model-based bandit optimization with application to automatic machine learning. In Artificial Intelligence and Statistics , pages 365-374. PMLR, 2014.

Jean Honorio and Tommi Jaakkola. Tight bounds for the expected risk of linear classifiers and PAC-Bayes finite-sample guarantees. In Artificial Intelligence and Statistics , pages 384-392. PMLR, 2014.

Francis L Huang, Wolfgang Wiedermann, and Bixi Zhang. Accounting for heteroskedasticity resulting from between-group differences in multilevel models. Multivariate Behavioral Research , pages 1-21, 2022.

Peter J Huber. The behavior of maximum likelihood estimates under nonstandard conditions. In Berkeley Symposium on Mathematical Statistics and Probability , pages 221-233. University of California Press, 1967.

Yassir Jedra and Alexandre Proutiere. Optimal best-arm identification in linear bandits. Advances in Neural Information Processing Systems , 33:10007-10017, 2020.

Kwang-Sung Jun, Lalit Jain, Blake Mason, and Houssam Nassif. Improved confidence bounds for the linear logistic model and applications to bandits. In International Conference on Machine Learning , pages 5148-5157. PMLR, 2021.

Maurits Kaptein and Petri Parvinen. Advancing e-commerce personalization: Process framework and case study. International Journal of Electronic Commerce , 19(3):7-33, 2015.

Zohar S Karnin. Verification based solution for structured mab problems. Advances in Neural Information Processing Systems , 29, 2016.

Julian Katz-Samuels, Lalit Jain, Kevin G Jamieson, et al. An empirical process approach to the union bound: Practical algorithms for combinatorial and linear bandits. Advances in Neural Information Processing Systems , 33:10371-10382, 2020.

Emilie Kaufmann, Olivier Cappé, and Aurélien Garivier. On the complexity of best-arm identification in multi-armed bandit models. The Journal of Machine Learning Research , 17(1): 1-42, 2016.

Jack Kiefer. Optimum experimental designs. Journal of the Royal Statistical Society: Series B (Methodological) , 21(2):272-304, 1959.

Jack Kiefer and Jacob Wolfowitz. The equivalence of two extremum problems. Canadian Journal of Mathematics , 12:363-366, 1960.

Johannes Kirschner and Andreas Krause. Information directed sampling and bandits with heteroscedastic noise. In Conference On Learning Theory , pages 358-384. PMLR, 2018.

Fanjie Kong, Yuan Li, Houssam Nassif, Tanner Fiez, Shreya Chakrabarti, and Ricardo Henao. Neural insights for digital marketing content design. In International Conference on Knowledge Discovery and Data Mining , pages 4320-4332, 2023.

Tor Lattimore and Csaba Szepesvári. Bandit algorithms . Cambridge University Press, 2020.

Zhaoqi Li, Lillian Ratliff, Houssam Nassif, Kevin Jamieson, and Lalit Jain. Instance-optimal pac algorithms for contextual bandits. In Conference on Neural Information Processing Systems , pages 37590-37603, 2022.

James G MacKinnon and Halbert White. Some heteroskedasticity-consistent covariance matrix estimators with improved finite sample properties. Journal of econometrics , 29(3):305-325, 1985.

Blake Mason, Lalit Jain, Subhojyoti Mukherjee, Romain Camilleri, Kevin Jamieson, and Robert Nowak. Nearly optimal algorithms for level set estimation. In International Conference on Artificial Intelligence and Statistics , pages 7625-7658. PMLR, 2022.

Blake Mason, Kwang-Sung Jun, and Lalit Jain. An experimental design approach for regret minimization in logistic bandits. In AAAI Conference on Artificial Intelligence , volume 36, pages 7736-7743, 2022.

Darcy P Mays and Raymond H Myers. Design and analysis for a two-level factorial experiment in the presence of variance heterogeneity. Computational statistics & data analysis , 26(2): 219-233, 1997.

Subhojyoti Mukherjee, Qiaomin Xie, Josiah Hanna, and Robert Nowak. SPEED: Experimental design for policy evaluation in linear heteroscedastic bandits. arXiv preprint arXiv:2301.12357 , 2023.

Sareh Nabi, Houssam Nassif, Joseph Hong, Hamed Mamani, and Guido Imbens. Bayesian meta-prior learning using Empirical Bayes. Management Science , 68(3):1737-1755, 2022.

H Neudecker. Symmetry, 0-1 matrices and jacobians. Econometric Theory , 2:157-190, 1986.

Ann L Oberg and Douglas W Mahoney. Linear mixed effects models. Topics in biostatistics , pages 213-234, 2007.

Friedrich Pukelsheim. Optimal design of experiments . SIAM, 2006.

Takaki Sato and Yasumasa Matsuda. Spatial extension of generalized autoregressive conditional heteroskedasticity models. Spatial Economic Analysis , 16(2):148-160, 2021.

Marta Soare. Sequential resource allocation in linear stochastic bandits . PhD thesis, Université Lille 1-Sciences et Technologies, 2015.

Marta Soare, Alessandro Lazaric, and Rémi Munos. Active learning in linear stochastic bandits. Bayesian Optimization in Theory and Practice , 2013.

Marta Soare, Alessandro Lazaric, and Rémi Munos. Best-arm identification in linear bandits. In Advances in Neural Information Processing Systems , 2014.

Chao Tao, Saúl Blanco, and Yuan Zhou. Best arm identification in linear bandits with linear dimension dependency. In International Conference on Machine Learning , pages 4877-4886. PMLR, 2018.

Bryan W Ting, Fred A Wright, and Yi-Hui Zhou. Simultaneous modeling of multivariate heterogeneous responses and heteroskedasticity via a two-stage composite likelihood. Biometrical Journal , page 2200029, 2023.

Martin J Wainwright. High-dimensional statistics: A non-asymptotic viewpoint , volume 48. Cambridge university press, 2019.

Po-An Wang, Ruo-Chun Tzeng, and Alexandre Proutiere. Fast pure exploration via Frank-Wolfe. Advances in Neural Information Processing Systems , 34:5810-5821, 2021.

Halbert White. A heteroskedasticity-consistent covariance matrix estimator and a direct test for heteroskedasticity. Econometrica: journal of the Econometric Society , pages 817-838, 1980.

Douglas P Wiens. Designs for weighted least squares regression, with estimated weights. Statistics and Computing , 23:391-401, 2013.

Douglas P Wiens and Pengfei Li. V-optimal designs for heteroscedastic regression. Journal of Statistical Planning and Inference , 145:125-138, 2014.

Weng Kee Wong and R Dennis Cook. Heteroscedastic g-optimal designs. Journal of the Royal Statistical Society: Series B (Methodological) , 55(4):871-880, 1993.

Weng Kee Wong and Wei Zhu. Optimum treatment allocation rules under a variance heterogeneity model. Statistics in Medicine , 27(22):4581-4595, 2008.

Jeffrey M Wooldridge. Econometric analysis of cross section and panel data . MIT press, 2010.

Liyuan Xu, Junya Honda, and Masashi Sugiyama. A fully adaptive algorithm for pure exploration in linear bandits. In International Conference on Artificial Intelligence and Statistics , pages 843-851. PMLR, 2018.

Dan Yan, Xiaohang Ren, Wanli Zhang, Yiying Li, and Yang Miao. Exploring the real contribution of socioeconomic variation to urban PM 2 . 5 pollution: New evidence from spatial heteroscedasticity. Science of the Total Environment , 806:150929, 2022.

Junwen Yang and Vincent Tan. Minimax optimal fixed-budget best arm identification in linear bandits. Advances in Neural Information Processing Systems , 35:12253-12266, 2022.

Mohammadi Zaki, Avinash Mohan, and Aditya Gopalan. Towards optimal and efficient best arm identification in linear bandits. arXiv preprint arXiv:1911.01695 , 2019.

Huiming Zhang and Song Xi Chen. Concentration inequalities for statistical inference. Communications in Mathematical Research , 37(1):1-85, 2021.

Heyang Zhao, Dongruo Zhou, Jiafan He, and Quanquan Gu. Bandit learning with general function classes: Heteroscedastic noise and variance-dependent regret bounds. arXiv preprint arXiv:2202.13603 , 2022.

Contents

A Background on Experimental Designs and Rounding

We use optimal experiment design [29, 42] to minimize uncertainty in the estimator of interest. Consider the WLS estimator with the weight matrix W T = Σ -1 T so that V ( ̂ θ WLS ) = ( X ⊤ T Σ -1 T X T ) -1 = ( ∑ ⊤ t =1 σ -2 x t x t x ⊤ t ) -1 . Let λ x denote the proportion of the total samples T given to the measurement vector x ∈ X so that V ( ̂ θ WLS ) = ( ∑ x ∈X λ x σ -2 x xx ⊤ ) -1 /T . We wish to take T samples so as to construct the design ∑ x ∈X Tλ x σ -2 x xx ⊤ , however Tλ x is often not an integer. As has been discussed at depth [46, 12, 35, 42, 3], efficient rounding schemes can solve this problem. Given some constant tolerance threshold ϵ , a design λ ∈ P X , and a minimum sample size r ( ϵ ) , efficient rounding procedures return a fixed allocation { x t } T t =1 that yields a (1 + ϵ ) approximation to the intended design. A simple, well known procedure with r ( ϵ ) ≤ { d ( d +1) + 2 } /ϵ comes from Pukelsheim [42], while the scheme of Allen-Zhu et al. [3] yields r ( ϵ ) ≤ O ( d/ϵ 2 ) . In our experimental design algorithms, we leverage the aforementioned rounding schemes.

We adapt the guarantees of [3] for Alg. 1 as an example in the following proposition. For sample size N , define X ∗ as a set composed of the elements of X ⊂ R d duplicated N times and S N = { s ∈ { 0 , 1 , . . . , N } N |X| : ∑ N |X| t =1 s i ≤ N } as the set of discrete sample allocations.

Proposition A.1. Suppose ϵ ∈ (0 , 1 / 3] and N ≥ 5 d/ϵ 2 . Let λ ∈ P X , then in polynomial-time (in n and d ), we can round λN to an integral solution ̂ λ ∈ S N satisfying

Proof. This follows directly from Theorem 2.1 in [3] letting C N = { s ∈ [0 , N ] N |X| : ∑ N |X| i =1 s i ≤ N } be a continuous relaxation of S N and setting k = N and n = N |X| . ■

Alternatively, one could use the robust inverse propensity score (RIPS) estimator [6], which avoids the need for rounding schemes via robust mean estimation.

B Technical Preliminaries

For our analysis, we do not assume that the error, η , in Eq. 1 is Gaussian. We generalize by extending to strictly sub-Gaussian noise, in which the variance σ 2 x is equal to the optimal variance proxy for η t , ∀ t ∈ N . This paradigm allows for noise distributions such as the symmetrized Beta, symmetrized Gamma, and Uniform.

The following definitions and propositions will be used in the proofs of Appendices D-F.

Definition B.1. A real-valued random variable X is sub-Gaussian if there exists a positive constant σ 2 such that for all λ ≥ 0 .

Proposition B.1 (Equation 2.9, [49]) . Let X be a sub-Gaussian random variable with sub-Gaussian parameter σ 2 and mean E ( X ) = µ . Then, for all t > 0 , we have

Definition B.2. A real-valued random variable X with mean E ( X ) = µ is sub-exponential if there are non-negative parameters ( ν, α ) such that

Proposition B.2 (Proposition 2.9, [49]) . Suppose that X is sub-exponential with parameters ( ν, α ) and mean E ( X ) = µ . Then, for t > 0 ,

This implies that for δ ∈ (0 , 1) ,

Proposition B.3 (Equation 37 in Appendix B, [19]) . Let X be a sub-Gaussian random variable with sub-Gaussian parameter σ 2 . Then X 2 is sub-exponential with parameters ( ν = 4 σ 2 √ 2 , α = 4 σ 2 ) .

Proposition B.4 (Equation 2.18, [49]) . Suppose that X i for i ∈ { 1 , 2 , . . . , n } are sub-exponential with parameters ( ν i , α i ) such that E ( X i ) = µ i . Then, ∑ n i =1 ( X i -µ i ) is sub-exponential with parameters ( ν ∗ , α ∗ ) such that

Definition B.3. A G-optimal design [28], λ ∗ ∈ P X , for a set of arms X ⊂ R d is such that

Proposition B.5 (Lemma 4.6, [28]) . For any finite X ⊂ R d with span( X ) = R d ,

This implies that if we sample each arm x ∈ X ⌈ τλ ∗ x ⌉ times, then

Definition B.4. The Forbenius norm of A ∈ R n × m is

Definition B.5. The Spectral norm of A ∈ R n × m is

̸

Proposition B.6 (Matrix Norm Properties) . For A ∈ R n × m , we have sub-multiplicativity for the Frobenius norm:

and a bound on the Spectral norm:

Proposition B.7 (Corollary 4.7, [61]) . Let ξ 1 , . . . , ξ n be zero-mean σ 2 -sub-Gaussian and ∆ ∈ R n × n . For any t > 0 ,

which implies that with probability 1 -δ for δ ∈ (0 , 1) ,

Proposition B.8 (Equation 20.3, [32]) . Assume the setup of Eq. 1 and that a data set of size Γ , D = { x t , y t } Γ t =1 , has been collected through a fixed design. Define V Γ = ∑ Γ t =1 x t x ⊤ t and construct the least squares estimator for D , ̂ θ Γ = V -1 Γ ( ∑ Γ t =1 x t y t ) . Then, with probability 1 -δ for δ ∈ (0 , 1) ,

C Baseline Variance Estimation Procedures

These are the algorithms we reference in the theoretical and empirical comparisons of Section 3.

The Uniform Estimator. The Uniform Estimator draws arms uniformly at random and uses all samples to construct estimators ̂ θ Γ of θ ∗ and ̂ Σ Γ of Σ ∗ .

The Separate Arm Estimator. Define U as the set of subsets of W of size M = d ( d +1) / 2 , and for U ∈ U construct Φ U such that its rows are composed of ϕ x ∈ U . Let ζ min ( A ) be the minimum singular value of matrix A ∈ R d × d . The Seperate Arm Estimator picks a set of M arms such that

It then splits Γ samples evenly between these arms, estimates the sample variance for each and finds the least squares estimator ̂ Σ SA Γ .

Algorithm 3: Uniform Estimator of Heteroskedastic Variance

Result:

Find ̂ Σ UE Γ

1 Input: Arms X ⊆ R d and Γ ∈

N

2 Sample arms uniformly at random from X

3 Define A ∗ = ∑ Γ t =1 x t x ⊤ t and b ∗ = ∑ Γ t =1 x t y t and estimate ̂ θ Γ = A ∗- 1 b ∗ .

4 Define A † = ∑ Γ t =1 ϕ x t ϕ ⊤ x t and b † = ∑ Γ t =1 ϕ x t ( y t - x ⊤ t ̂ θ Γ ) 2 .

5 Output: ̂ Σ UE Γ = A b †

† - 1 .

Algorithm 4: Seperate Arm Estimator of Heteroskedastic Variance

Result:

Find ̂ Σ SA Γ

1 Input: Arms X ⊆ R d and Γ ∈

N

2 Define U ∗ = arg max U ∈U ζ min (Φ -1 U ) .

3 Equally distribute Γ samples among arms X U ∗ = { m : ϕ m ∈ U ∗ } and observe rewards { y m,t } Γ /M t =1 for m ∈ X U ∗

4 Define the sample average for each arm ¯ y m = ( M/ Γ) ∑ Γ /M t =1 y m,t

5 Output: ̂ Σ SA Γ ← arg min s ∈ R M ∑ m ∈X U ∗ ∑ Γ /M t =1 { ϕ m ⊤ s -(¯ y m -y m,t ) 2 } 2

D Proofs of Variance Estimation

In this appendix, we analyze the maximum absolute error of estimates associated with the Seperate Arm Estimator and the HEAD Estimator. Recall that M := d ( d +1) / 2 . Define W = { ϕ x , ∀ x ∈ X} ⊂ R M where ϕ x = vec( xx ⊤ ) G and G ∈ R d 2 × M is the duplication matrix such that G vech( xx ⊤ ) = vec( xx ⊤ ) with vech( · ) denoting the half-vectorization.

D.1 Analysis of the Separate Arm Estimator

We begin by proving a general result bounding the absolute error of the sample variance for any sub-Gaussian random variable with parameter γ 2 .

Proposition D.1. Let X be a γ 2 -sub-Gaussian random variable with mean µ . Assume that we have collected n independent copies of X , { X 1 , X 2 , . . . , X n } , and label the sample mean as ¯ X = ∑ n i =1 X i /n . Defining the sample variance as ̂ σ 2 = ∑ n i =1 ( X i -¯ X ) 2 /n , the true variance as σ 2 = E ( X 2 i ) -µ 2 and δ ∈ (0 , 1) , we find that with probability greater than 1 -δ ,

Proof. We begin by upper bounding the absolute error using the triangle inequality,

We first bound quantity A . Leveraging Proposition B.3, we establish that X 2 i -E ( X 2 i ) is subexponential with parameters (4 γ 2 √ 2 , 4 γ 2 ) . We can then use Proposition B.4 to analyze the sum of independent sub-exponential random variables, finding that ∑ n i =1 { X 2 i -E ( X 2 i ) } /n is subexponential with parameters (4 γ 2 √ 2 / √ n, 4 γ 2 /n ) . Finally, we invoke Proposition B.2 to bound quantity A with probability 1 -δ/ 2 ,

Now we bound quantity B . Appealing to Proposition B.3, we find that ¯ X 2 -σ 2 n -µ 2 is subexponential with parameters (4 √ 2 γ 2 /n, 4 γ 2 /n ) . We then use Proposition B.2 to bound B with probability 1 -δ/ 2 ,

where ( a ) follows because √ 2 / √ log(4 /δ ) ≤ √ 2 . Finally, we combine quantities A , B and C , apply a union bound, and conclude that with probability 1 -δ for δ ∈ (0 , 1) ,

■

We can now use Proposition D.1 to control the sample variance estimates of the arms chosen by the Separate Arm Estimator. Theorem D.1 derives a concentration bound of the maximum absolute error of ˜ σ 2 x = min { max { σ 2 min , x ⊤ ̂ Σ SA Γ x } , σ 2 max } . This bound, which is O ( d 4 ) , scales unfavorably with dimension as highlighted in the theoretical and empirical comparisons conducted in Section 3.

Theorem D.1. Set ̂ Σ SA Γ as the output of Alg. 4 and define B such that for all x ∈ X , ∥ x ∥ 2 ≤ B . Defining ˜ σ 2 x = min { max { σ 2 min , x ⊤ ̂ Σ SA Γ x } , σ 2 max } , σ 2 x = x ⊤ Σ ∗ x and δ ∈ [0 , 1] , for any x ∈ X ,

Note that we set the probability bound to 1 -δ/ 2 instead of 1 -δ because we want to be able to leverage these variance estimation bounds with the identification algorithms of Section 4 (and so are anticipating another union bound).

Proof. Label M = d ( d + 1) / 2 and X U ∗ = { m : ϕ m ∈ U ∗ } , where U ∗ is defined in Alg. 4. Additionally, in Alg. 4 we observe { y m,t } Γ /M t =1 for each arm m ∈ X U ∗ . Define the sample average for m ∈ X U ∗ as

We obtain an estimate of Σ ∗ via

Label z m,t = (¯ y m -y m,t ) 2 ∈ R , Ψ ∗ = vech(Σ ∗ ) ∈ R M , ˜ Ψ Γ = vech( ̂ Σ SA Γ ) ∈ R M and .

We reformulate Eq. 5 as follows

where ( a ) follows because ∑ Γ /M t =1 z 2 m,t / (Γ /M ) and ¯ z 2 m do not depend on s . Defining ¯ z = [¯ z m ] m ∈X U ∗ ∈ R M , Eq. 5 becomes

We find that

Consequently, we can focus on bounding ∥ Φ U ∗ Ψ ∗ -¯ z ∥ 2 in order to control ∥ Φ U ∗ ( ˜ Ψ Γ -Ψ ∗ ) ∥ 2 . Note that the m th entry of (Φ U ∗ Ψ ∗ -¯ z ) is ( m ⊤ Σ ∗ m -¯ z m ) and that ¯ z m is the sample variance of observations { y m,t } Γ /M t =1 . Using Proposition D.1 and the fact that the y m,t are σ 2 max -sub-Gaussian, we show that with probability 1 -δ/ (2 M ) ,

We apply Eq. 6 to quantity ∥ Φ U ∗ Ψ ∗ -¯ z ∥ 2 and a union bound to find that with probability 1 -δ/ 2 ,

Consequently,

We now analyze the absolute error of the estimate ˜ σ x for an arbitrary action x ∈ X . Since Φ U ∗ is non-singular by construction, we know that for q := (Φ U ∗ -1 ) ⊤ ϕ x ∈ R M ,

where ( a ) follows from the definition of the matrix norm. We now leverage the nature of the duplication matrix, G , and the fact that ∥ x ∥ 2 ≤ B to bound ∥ ϕ x ∥ 2 2 .

where ( a ) follows from the fact that G ⊤ G ⪯ 2 I as shown in Equation (61) of [40], and ( b ) uses the fact that G ( G ⊤ G ) -1 G ⊤ vec( xx ⊤ ) = vec( xx ⊤ ) as shown in Equations (54) and (52) of [40]. Consequently, with probability greater than 1 -δ/ 2 , we have that for any x ∈ X ,

■

D.2 Analysis of the HEAD Estimator

We now present the maximum absolute error guarantees of the HEAD procedure, Alg. 1.

Proof of Theorem 3.1. In this proof, we bound the estimation error | ̂ σ 2 x -σ 2 x | , ∀ x ∈ X . Define σ 2 max = max x ∈X σ 2 x , Ψ ∗ = vech(Σ ∗ ) , N 1 = { 1 , . . . , Γ / 2 } , N 2 = { Γ / 2 + 1 , . . . , Γ } and

where ̂ Ψ Γ = vech( ̂ Σ Γ ) . Furthermore, define vector e with entries e t = ( y t -x ⊤ t ̂ θ Γ / 2 ) 2 and Φ such that the rows of Φ are composed of { ϕ x t } t ∈ N 2 . We find that

For any ϕ x ∈ W , we are interested in bounding ∣ ∣ ∣ ϕ ⊤ x ( ̂ Ψ Γ -Ψ ∗ ) ∣ ∣ ∣ ,

Define z ⊤ = ϕ ⊤ x (Φ ⊤ Φ) -1 Φ ⊤ and X such that the rows of X are composed of { x t } t ∈ N 2 . We can bound this expression with the triangle inequality as

We can simplify the first term as,

Additionally, the third term is equivalent to 2 | ( η ◦ z ) ⊤ X ( ̂ θ Γ / 2 -θ ∗ ) | , where ' ◦ " is element-wise multiplication. Consequently, the absolute error is equal to

Note that to bound | ϕ ⊤ x ( ̂ Ψ Γ -Ψ ∗ ) | with probability 1 -δ/ 2 , it is necessary to bound quantities A , B and C with probability 1 -δ/ 6 in anticipation of a union bound.

Quantity A . We sample { x t , t ∈ N 1 } according to the rounded G-optimal design over X constructed in Stage 1 of Alg. 1. We use the concentration bounds of Proposition B.1, to state that with probability 1 -δ/ 6 ,

We use Proposition B.5 in combination with the rounding procedure guarantees of Proposition A.1 to bound quantity G . For ϵ ∈ (0 , 1 / 3] , with probability 1 -δ/ 6 ,

Quantity B . According to Proposition B.3, η 2 t is sub-exponential with parameters (4 √ 2 σ 2 max , 4 σ 2 max ) . Consequently, by Proposition B.4, we know that B is sub-exponential with parameters

We now invoke the G-optimal design over W space conducted in Stage 2 of Alg. 1 and Propositions B.5 and A.1 to find that for ϵ ∈ (0 , 1 / 3] ,

Therefore, B has sub-exponential parameters ( √ 64 σ 4 max d 2 (1 + 6 ϵ ) / Γ , 8 σ 2 max d 2 (1 + 6 ϵ ) / Γ) . Using Proposition B.2, with probability 1 -δ/ 6 ,

Quantity C . Construct V ∈ R Γ / 2 × d such that the rows of V correspond to { x t : t ∈ N 1 } . We can upper bound quantity C using the Cauchy-Schwartz inequality in the following manner,

By Proposition B.8, with probability 1 -δ/ 12 ,

Now construct D z = diag( { z t : t ∈ N 2 } ) and express,

Consequently, and

We first establish that by the G-optimal designs in both X and W space along with Propositions B.5 and A.1,

where ϵ ∈ (0 , 1 / 3] and ( a ) follows from Eq. 7. Consequently, we find that

This inequality will be useful in the following derivation. Labeling ∆ = D z X ( V ⊤ V ) -1 X ⊤ D z , we use Proposition B.7 to bound η ⊤ ∆ η with probability 1 -δ/ 12 and explain below.

where line (9) follows from Eq. 8, Proposition B.6 and Proposition B.7. In lines (10) and (11) , we use Proposition B.6 again and then the definition of the Frobenius norm in line (12) . Consequently, with probability 1 -δ/ 6 ,

Combining quantities A , B and C together and applying a union bound, we have with probability 1 -δ/ 2 ,

We can exploit the dependencies of V and W to simplify this bound in two ways under different conditions. Note that we only have dependency on |X| through term V . If d -1 > log( |X| ) / log(1 /δ ) and Γ > 16 log(12 /δ ) σ 2 max d 2 (1 + 6 ϵ ) , then with probability 1 -δ/ 2 ,

Alternatively, if just Γ > 16 log(12 |X| /δ ) σ 2 max d 2 (1 + 6 ϵ ) , then

Since σ 2 max ≥ σ 2 x ≥ σ 2 min for all x ∈ X ,

In summary, if Γ > 16 log(12 |X| /δ ) σ 2 max d 2 (1 + 6 ϵ ) ,

where C ′ = 256(1 + 6 ϵ ) .

■

We now prove a simple lemma relating the additive bound on | ̂ σ 2 x -σ 2 x | to a multiplicative bound 1 -C Γ ,δ < σ 2 x / ̂ σ 2 x < 1 + C Γ ,δ .

Lemma D.1. Define ̂ σ 2 x = min { max { x ⊤ ̂ Σ Γ x, σ 2 min } , σ 2 max } . Letting σ 2 min = min x ∈X σ 2 x , κ = σ 2 max /σ 2 min , and

Proof. Using Lemma D.1, we know that with probability greater than 1 -δ/ 2 the following event holds.

■

E Proofs of Best-Arm Identification

We divide the proof of Theorem 4.1 and 4.2 between Appendices E and F for ease of exposition. In Appendix E, we consider the transductive linear best-arm identification problem ( BAI ) in which we are interested in identifying z ∗ = max z ∈Z z ⊤ θ ∗ with probability 1 -δ for δ ∈ (0 , 1) . In Appendix F, we consider the level set ( LS ) counterpart. Recall that for a set V , let P V := { λ ∈ R |V| : ∑ v ∈V λ v = 1 , λ v ≥ 0 ∀ v } denote distributions over V and Y ( V ) := { v ′ -v : v, v ′ ∈ V , v = v ′ } be an operator giving differences. We begin with the lower bound.

E.1 Lower Bound

Proof of Theorem 4.1 for transductive best-arm identification. This proof is similar to the lower bound proofs of [35] and [12]. Let C := { ˜ θ ∈ R d : ∃ i s . t . ˜ θ ⊤ ( z ∗ -z i ) ≤ 0 } , i.e., ˜ θ ∈ C if and only if z ∗ is not the best arm in the linear bandit instance ( X , Z , ˜ θ ) . Additionally, label K = |X| .

We first invoke Lemma 1 of [27]. Under a δ -PAC strategy for finding the best arm for the bandit instance ( X , Z , θ ∗ ) , let T i denote the random variable that is the number of times arm i is pulled. Defining σ 2 x = x ⊤ Σ ∗ x for all x ∈ X , we let ν θ,i denote the reward distribution of the i -th arm of X , i.e., ν θ,i = N ( x ⊤ i θ, σ 2 x i ) . For any ˜ θ ∈ C we know from [27] that

Following the steps of [12], we find that

̸

Label Z = { z 1 , z 2 , . . . , z n } , where z 1 = z ∗ without loss of generality. For j = 1 , λ ∈ P X and ϵ > 0 , define A ( λ ) := ∑ K i =1 λ i x i x ⊤ i σ 2 x i , w j = z ∗ -z j and

̸

Note that w ⊤ j θ j ( ϵ, λ ) = -ϵ < 0 , implying that θ j ∈ C . We find that the KL divergence between ν θ,i and ν θ j ( ϵ,λ ) ,i is given by:

Returning to Eq. 17,

where ( a ) uses the fact that ∑ K i =1 λ i x i x ⊤ i σ 2 x i = A ( λ ) . Letting ϵ → 0 establishes the result. ■

E.2 Upper Bound

We now prove Theorem 4.2 by first establishing that Alg. 2 is δ -PAC for the transductive linear bandits best-arm identification problem.

E.2.1 Proof of Correctness

Overview of Correctness Proof. Lemma D.1 along with Lemmas E.1 through E.3 are combined to prove that Algorithm 2 is δ -PAC for the best-arm identification problem. Lemma E.1 leverages Lemma D.1 and begins by establishing that in any round ℓ ∈ N , the estimation error for the difference in mean rewards between z ∈ Z and z ∗ = arg max z ∈Z z ⊤ θ ∗ is bounded by ϵ ℓ = 2 -ℓ with probability 1 -δ for δ ∈ (0 , 1) . Lemma E.2 uses this result to prove that starting at round ℓ , z ∗ ∈ Z ℓ +1 with probability 1 -δ . This establishes that z ∗ is not eliminated from the set of items we are considering at any step with high probability. Lastly, Lemma E.3 uses Lemma E.1 to show that Algorithm 2 will eventually identify z ∗ . This is done by proving that Alg. 2 eliminates items with mean rewards that are outside a shrinking radius (halves at every step of the algorithm) of the highest reward, z ∗⊤ θ ∗ . In other words, with probability 1 -δ for δ ∈ (0 , 1) , we prove that items in Z ℓ with expected rewards that are 2 ϵ ℓ away from z ∗⊤ θ ∗ will not be in Z ℓ +1 . Since ϵ ℓ → 0 as ℓ →∞ , we know that lim ℓ →∞ P ( Z ℓ = { z ∗ } ) = 1 -δ .

Lemma E.1. With probability at least 1 -δ for δ ∈ (0 , 1) ,

for all ℓ ∈ N and z ∈ Z .

Proof. Define σ 2 x = x ⊤ Σ ∗ x , ̂ θ ℓ as the weighted least squares estimator constructed in Algorithm 2 at stage ℓ , and events

and

Label n ℓ as the number of data points collected during stage ℓ , Υ = diag( { σ 2 i } n ℓ i =1 ) , ̂ Υ = diag( { ̂ σ 2 i } n ℓ i =1 ) , X ℓ as a matrix whose rows are composed of the arms x i,ℓ drawn during round ℓ , Y ℓ as the vector of rewards drawn during round ℓ , A ℓ := ∑ n ℓ i =1 ̂ σ -2 i x ℓ,i x ⊤ ℓ,i , and ˜ Ω = A -1 ℓ X ⊤ ℓ ̂ Υ -1 Υ ̂ Υ -1 X ℓ A -1 ℓ as the weighted least squares estimator's covariance. We first establish that for z ∈ R d , and

Furthermore, defining u ⊤ = z ⊤ A -1 ℓ X ⊤ ℓ ̂ Υ -1 / 2 ∈ R n ℓ and conditioning on event J ,

where ( a ) is the result of event J . Using this result and Proposition B.1, with probability at least 1 -δ 4 ℓ 2 |Z| ,

Conditioned on event J , we have with probability 1 -δ 4 ℓ 2 |Z| ,

Applying a union bound, we find that

which proves the result.

■

Now we establish that, with probability 1 -δ for δ ∈ (0 , 1) , the best arm z ∗ will not be eliminated from Z ℓ for any round ℓ ∈ N .

Lemma E.2. With probability at least 1 -δ for δ ∈ (0 , 1) , for any ℓ ∈ N and z ′ ∈ Z ℓ ,

Proof. We know that with probability 1 -δ for all ℓ ∈ N and z ′ ∈ Z ℓ ,

where the last inequality follows from Lemma E.1.

■

Then, we establish that in round ℓ , items in Z ℓ that have mean rewards farther than 2 ϵ ℓ from the highest reward will be eliminated.

Lemma E.3. With probability at least 1 -δ for δ ∈ (0 , 1) , for any ℓ ∈ N and z ∈ Z ℓ such that ⟨ z ∗ -z, θ ∗ ⟩ > 2 ϵ ℓ ,

Proof. We know that with probability 1 -δ for all ℓ ∈ N ,

where ( a ) follows from Lemma E.1.

E.2.2

Proof of Sample Complexity

Overview of Sample Complexity Proof. In this section, we prove an upper bound on the sample complexity of Algorithm 2 for transductive best-arm identification with linear bandits. We begin by establishing some basic facts about the partial ordering of Hermetian positive definite matrices in Lemmas E.4 and E.5. Let A and B be Hermitian positive definite matrices in R d × d and define A ⪰ B if ( A -B ) is positive definite.

Lemma E.4. Let A and B be Hermetian Positive Definite matrices in R d × d ,

■

Proof. For any x ∈ R d we define y = A 1 / 2 x and

Lemma E.5. Let A and B be Hermitian positive definite matrices in R d × d ,

Proof. We use Lemma E.4 twice in the following proof.

where line 2 follows because A -1 / 2 BA -1 / 2 is similar to B 1 / 2 A -1 / 2 A -1 / 2 B 1 / 2 and both are Hermetian positive definite. ■

We now use Lemmas E.4 and E.5 in our problem context. For V ⊂ R d , define g : ( V , P X ) → R such that g ( v, λ ) = ∥ v ∥ ( ∑ x λ x xx ⊤ / ̂ σ 2 x ) -1 , f : ( V , P X ) → R such that f ( v, λ ) = ∥ v ∥ ( ∑ x λ x xx ⊤ /σ 2 x ) -1 and

Note that the function, q , is also defined and leveraged in Algorithm 2. Lemma E.6 establishes a multiplicative bound of the form

using Lemmas D.1 and E.5.

Lemma E.6. Define ̂ σ 2 x = min { max { x ⊤ ̂ Σ Γ x, σ 2 min } , σ 2 max } . For a given λ , define

Assume C Γ ,δ < 1 and v ∈ R d , then with probability 1 -δ/ 2 for δ ∈ (0 , 1) the following holds true:

or equivalently

Proof. Using Lemma D.1, we establish that

and

■

Leveraging the partial ordering of Hermitian positive definite matrices and Lemma E.5, we can then represent these inequalities as

This implies that and

By combining these inequalities we arrive at the result.

■

We now want to use the multiplicative bound of Lemma E.6 to define the relationship between q ( V ) and q ∗ ( V ) and eventually relate the sample complexity upper bound of Alg. 2 to the lower bound established in Theorem 4.1. In order to accomplish this goal, we prove some general results concerning minimax problems in Lemmas E.7 and E.8.

Lemma E.7. For c > 0 , v ∈ R d and λ ∈ P X , if

Proof. For a given λ ∈ P X , define v † ( λ ) = arg max v ∈V g ( v, λ ) and v ∗ ( λ ) = arg max v ∈V f ( v, λ ) , then

Lemma E.8. For c > 0 , v ∈ R d and λ ∈ P X , if

then then

Proof. For a given λ ∈ P X , define v † ( λ ) = arg max v ∈V g ( v, λ ) , v ∗ ( λ ) = arg max v ∈V f ( v, λ ) , λ ∗ = arg min λ ∈ P X max v ∈V f ( v, λ ) and λ † = arg min λ ∈ P X max v ∈V g ( v, λ ) , then

where lines ( a ) and ( b ) follows from Lemma E.7.

Leveraging Lemma E.8, we bound q ∗ ( V ) in terms of q ( V ) in Lemma E.9.

.

■

■

Lemma E.9. With probability 1 -δ/ 2 for δ ∈ (0 , 1) ,

Proof. From Lemma E.6, we know that

According to Lemma E.8, this implies

In round ℓ of Algorithm 2, we sample arms in X according to a given design, ̂ λ ℓ . Although other rounding procedures are possible (see Proposition A.1), Theorem 4.2 assumes that we take ⌈ τ ℓ ̂ λ ℓ,x ⌉ samples of arm x ∈ X to accomplish this goal. This results in more than τ ℓ samples being taken in round ℓ . We define the sparsity of any ̂ λ ℓ design using Caratheodory's Theorem to motivate the number of extra samples needed by sampling in this manner.

Lemma E.10. [Appendix 8, [44]] In Algorithm 2, the support of ̂ λ ℓ for all ℓ ∈ N is d ( d +1) / 2 + 1 .

Proof. Any design matrix in R d × d is symmetric and can consequently be expressed as a point in R M , M = d ( d +1) / 2 . The design matrices leveraged in this paper belong to a convex hull of a subset of points since they are a convex combination of xx ⊤ for x ∈ X . According to Caratheodory's theorem, any point in the convex hull of any subset of points in R D can be defined as a convex combination of D +1 points. The optimal experimental designs in this paper can therefore be represented using only M +1 points [44]. This sparsity holds regardless of the design's form and so applies to both the homoskedastic and heteroskedastic variance settings. ■

Finally, we connect the sample complexity upper bound of Algorithm 2 to the lower bound for δ -PAC algorithms established in Theorem 4.2 for transductive best-arm identification.

Proof of Theorem 4.2 for transductive best-arm identification. Assume that max z ∈Z ℓ |⟨ z -z ∗ , θ ∗ ⟩| ≤ 2 . Define S ℓ = { z ∈ Z ℓ : ⟨ z ∗ -z, θ ∗ ⟩ ≤ 4 ϵ ℓ } . Note that by assumption Z = Z 1 = S 1 . Lemma E.3 implies that with probability at least 1 -δ for δ ∈ (0 , 1) , we know Z ℓ ⊂ S ℓ for all ℓ ∈ N . This means that

For ℓ ≥ ⌈ log 2 (4∆ -1 ) ⌉ , we know that S ℓ = { z ∗ } . Thus, the sample size needed to identify z ∗ is upper bounded as follows.

where ( a ) follows from Lemma E.10.

We relate this quantity to the lower bound in Theorem 4.1 below. With probability 1 -δ ,

here ( a ) follows from the fact that the maximum of a set of numbers is always greater than the average and ( b ) by the fact that the minimum of a sum is greater than the sum of the minimums. We used the fact that

by the triangle inequality in ( c ) . Finally, ( d ) follows from Lemma E.9.

Consequently, putting the previous pieces together we find that the sample complexity can be upper bounded as

where the dependency on d ( d +1) / 2 can reduced to d by Proposition A.1. This proves the result. ■

F Proofs of Level Set Identification

This appendix contains the proofs for the results on transductive linear bandit level set identification. In this problem, we are interested in identifying the set G α ( θ ∗ ) = { z ∈ Z : z ⊤ θ ∗ > α } , i.e., the items with mean rewards that exceed threshold α . We prove the lower bound on the sample complexity from Theorem 4.1 in App. F.1. The proof of correctness and the sample complexity upper bound for Alg. 2 are presented in App. F.2. These proofs resemble the transductive linear bandit identification proofs with slight modifications accounting for the the problem structure. Recalling B α ( θ ∗ ) = { z ∈ Z : z ⊤ θ ∗ < α } , we assume that Z = G α ( θ ∗ ) ∪ B α ( θ ∗ ) , which is equivalent to the condition { z ∈ Z : z ⊤ θ ∗ = α } = ∅ . Let ∆ = min z ∈Z | z ⊤ θ ∗ -α | denote the minimum absolute gap for the threshold α .

F.1 Lower Bound

Proof of Theorem 4.1 for transductive level set identification. Label X ⊂ R d as the measurement vectors such that K = |X| and Z as the decision item set. This proof is similar to the lower bound proofs of [35] and [12]. Let C be the set of alternatives such that for ˜ θ ∈ C , G α ( θ ∗ ) = G α ( ˜ θ ) . This set of alternatives is decomposed by [35] as,

̸

We now recall Lemma 1 of [27]. Under a δ -PAC strategy for bandit instance ( X , Z , θ ∗ ) , let T i denote the random variable that is the number of times arm i is pulled. Defining σ 2 x = x ⊤ Σ ∗ x for all x ∈ X , we denote ν θ,i as the reward distribution of the i -th arm of X , i.e., ν θ,i = N ( x ⊤ i θ, σ 2 x i ) . For any ˜ θ ∈ C , we know from [27] that

Following the steps of [12], we find that

For λ ∈ P X , define A ( λ ) := ∑ K i =1 λ i x i x ⊤ i σ 2 x i . For ϵ > 0 and z ∈ G α ( θ ∗ ) , define

where z ⊤ θ z ( ϵ, λ ) = α -ϵ , implying that θ z ( ϵ, λ ) ∈ C . For z ∈ G c α ( θ ∗ ) , define

where z ⊤ θ z ( ϵ, λ ) = α + ϵ , implying that θ z ( ϵ, λ ) ∈ C .

Defining ϵ z = ϵ if z ∈ G α ( θ ∗ ) and ϵ z = -ϵ if z ∈ G c α ( θ ∗ ) , we find that the KL divergence between ν θ,i and ν θ z ( ϵ,λ ) ,i is given by:

Returning to Eq. 18,

where ( a ) uses the fact that ∑ K i =1 λ i x i x ⊤ i σ 2 x i = A ( λ ) . Letting ϵ z → 0 establishes the result. ■

F.2 Upper Bound

We now prove the upper bound on the sample complexity of Alg. 2 for level set estimation. This proof follows similarly to the upper bound on the sample complexity of Alg. 2 for best-arm identification shown in App. E.2. Note that in each round ℓ ∈ N , Alg. 2 maintains a set Z ℓ of undecided items, a set G ℓ of items estimated to have value exceeding the threshold α , and a set B ℓ of items estimated to have value falling below the threshold α .

F.2.1 Proof of Correctness

To begin, we show that that for all rounds ℓ ∈ N of Alg. 2, the error in the estimates of z ⊤ θ ∗ for all z ∈ Z are bounded by ϵ ℓ with probability at least 1 -δ for δ ∈ (0 , 1) .

Lemma F.1. For all ℓ ∈ N and z ∈ Z , with probability at least 1 -δ for δ ∈ (0 , 1) ,

Proof. This proof follows identically to Lemma E.1 by replacing z ∗ with the zero vector.

■

We now apply the preceding result to show that with probability 1 -δ for δ ∈ (0 , 1) , Alg. 2 does not place items in the wrong category.

Lemma F.2. With probability at least 1 -δ for δ ∈ (0 , 1) , for any z ∈ Z ℓ , ℓ ∈ N , such that z ⊤ θ ∗ > α ,

Proof. Consider z ∈ G α so that z ⊤ θ ∗ > α . Using Lemma F.1, with probability at least 1 -δ for δ ∈ (0 , 1) we have

By the procedure of Alg. 2, this guarantees G ℓ +1 ⊆ G α .

and for any z ∈ Z ℓ such that z ⊤ θ ∗ < α ,

Now, consider z ∈ B α so that z ⊤ θ ∗ < α . Using Lemma F.1, with probability at least 1 -δ for δ ∈ (0 , 1) we have

By the procedure of Alg. 2, this guarantees B ℓ +1 ⊆ B α .

■

We now show that any z ∈ Z such that | z ⊤ θ ∗ -α | > 2 ϵ ℓ is removed from the active set, Z ℓ , in round ℓ ∈ N of Alg. 2 with probability at least 1 -δ for δ ∈ (0 , 1) . Since ϵ ℓ → 0 , this proves that we will eventually identify G α correctly.

Lemma F.3. With probability at least 1 -δ for δ ∈ (0 , 1) , for any ℓ ∈ N and z ∈ Z ℓ such that |⟨ z, θ ∗ ⟩ -α | > 2 ϵ ℓ ,

Proof. Set δ ∈ (0 , 1) . Let us begin by considering any z ∈ Z ℓ such that ⟨ z, θ ∗ ⟩ -α > 2 ϵ ℓ . Using Lemma F.1, with probability at least 1 -δ ,

Now, consider any z ∈ Z ℓ such that ⟨ z, θ ∗ ⟩ -α < -2 ϵ ℓ . Using Lemma F.1, with probability at least 1 -δ ,

Hence, for z ∈ Z ℓ such that |⟨ z, θ ∗ ⟩ -α | > 2 ϵ ℓ , we have that |⟨ z, ̂ θ ℓ ⟩ -ϵ ℓ | > α with probability at least 1 -δ . ■

By the elimination conditions of Alg. 2, this guarantees that the z ∈ Z ℓ such that |⟨ z, θ ∗ ⟩ -α | > 2 ϵ ℓ are not included in Z ℓ +1 .

F.2.2 Proof of Upper Bound on Sample Complexity

Now we prove the upper bound on the sample complexity of Alg. 2 for the transductive level set identification problem.

Proof of Theorem 4.2 for transductive level set identification. Set δ ∈ (0 , 1) . Assume that max z ∈Z ℓ |⟨ z, θ ∗ ⟩ -α | ≤ 2 . Define S ℓ = { z ∈ Z ℓ : |⟨ z, θ ∗ ⟩ -α | ≤ 4 ϵ ℓ } for ℓ ∈ N . Note that by assumption Z = S 1 = Z 1 . Lemma F.3 implies that with probability at least 1 -δ , we know Z ℓ ⊂ S ℓ for all ℓ ∈ N . This means

For ℓ ≥ ⌈ log 2 (4∆ -1 ) ⌉ , we know that S ℓ = ∅ , thus the sample complexity to identify the set G α is upper bounded as follows.

where ( a ) follows from Lemma E.10. We now relate this quantity to the lower bound and explain below. With probability 1 -δ ,

■

where ( a ) follows from the fact that the maximum of a set of numbers is always greater than the average, ( b ) by the fact that the minimum of a sum is greater than the sum of the minimums, and ( c ) from Lemma E.9. Consequently, putting the previous pieces together we find that the sample

complexity can be upper bounded as

where the dependency on d ( d +1) / 2 can reduced to d by Proposition A.1. This proves the result.

G Comparing Identification Lower Bounds

In this section, we use the general notation for the level set and best-arm identification problems developed in Section 4. Define

which is the lower bound for both BAI and LS by [12] and [35] in the homoskedastic case when upper bounding the variance.

Lemma G.1. In the same setting as Theorem 4.1,

and both the upper and lower bounds are tight.

Proof. Consider semidefinite matrices A , B , and C in R d × d such that A ⪯ B ⪯ C . Then C -1 ⪯ B -1 ⪯ A -1 . Hence, for x ∈ R d , ∥ x ∥ C -1 ≤ ∥ x ∥ B -1 ≤ ∥ x ∥ A -1 . Defining

we note that

Hence,

Applying this to the result of Theorem 4.1 and invoking Lemma E.8, we see that

Recalling that

is the lower bound in both the best arm and level set problems for the homoskedastic variance algorithm in a heteroskedastic noise setting (when upper bounding the variance), we see that

Moreover, both the upper and lower bounds are tight by taking σ 2 x = σ 2 min , ∀ x ∈ X (to match the lower bound) and σ 2 x = σ 2 max , ∀ x ∈ X (to match the upper bound). ■

H Experiment Details

This Appendix gives more details on the experiments presented in Section 5.

H.1 Oracle Designs

We demonstrate the difference in oracle allocations referenced by Experiments 1-3 in Section 5. We begin with Experiment 1 in Fig. 4, where the oracle distribution that assumes homoskedastic variance devotes more attention to the informative directions with smaller magnitude. In contrast, the heteroskedastic oracle balances evenly because each informative arm has an equivalent signal-noise ratio when accounting for variance.

Next, we analyze Experiment 2 in Fig. 5. In this setting the oracle distribution that assumes homoskedastic variance allocates evenly between three informative arms, while the heteroskedastic oracle prioritizes the informative directions with less variance.























































Lastly, we examine the oracle distributions for the multivariate testing experiment in Fig. 6. Note that the first four arm have higher variance than the last four arms. We can see in Fig. 6 that the Heteroskedastic Oracle prioritizes the second group of arms because these will require more samples to identify their effect sizes.

H.2 Assuming Homosekedastic Noise

If we assume that there is homoskedastic noise in a heteroskedastic setting, then we need to upperbound the noise with σ 2 max in the sample size calculation of Alg. 2 in order to maintain correctness. In Section 5, we compare H-RAGE to RAGE ; the latter uses the σ 2 max upper-bound and is described in Alg. 5.



































Algorithm 5: ( RAGE ) Randomized Adaptive Gap Elimination

Peer reviews:
{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper studies linear experimental design in the presence of heteroskedastic noise. The model can observe noise linear measurements generated by unknown model parameters, with sub-Gaussian covariance depends on the measurement vector. The goal is to estimate the unknown covariance matrix and corresponding covariance wrt measurement $x$. The algorithm works as follows: it uses the first half of the samples to construct a least-squares estimator for the model parameter, then use the second half to estimate the noises and consequently, the covariance. This algorithm leads to a sample complexity depends quadratically on $d$. As applications, the proposed approach can be used for best arm identification and level set identification. Finally, experiments are conducted to empirically validate the algorithm.\n\nSTRENGTHS:\nThis paper studied an interesting problem, in practical setting, heteroskedastic noises usually present, so it is important to study the sample complexity and efficient algorithm for this setting.\n\nThe presentation of the paper is good, it is a good mixture of motivation, intuition and technicality. I enjoyed reading the paper.\n\nThe results obtained in this paper are strong.\n\nWEAKNESSES:\nIt will be better to also include the runtime complexity of the proposed algorithm. Linear bandits now admit highly efficient algorithms with curated data structures [Yang et al., ICML'22] and so does efficient experimental design [Allen-Zhu et al., ICML'17]. It will be interesting to study the runtime complexity of the proposed algorithm.\n\n[Yang el al., ICML'22]: S. Yang, T. Ren, S. Shakkottai, E. Price, I. Dhillion, S. Sanghavi. Linear Bandit Algorithms with Sublinear Time Complexity.\n\n[Allen-Zhu et al., ICML'17]: Z. Allen-Zhu, Y. Li, A. Singh, Y. Wang. Near-Optimal Discrete Optimization for Experimental Design: A Regret Minimization Approach.\n\nQUESTIONS:\nWhat is the runtime complexity of the proposed algorithm?\n\nLIMITATIONS:\nYes.", "TITLE": "Official Review", "RECOMMENDATION": 7, "REVIEWER_CONFIDENCE": 1, "SOUNDNESS_CORRECTNESS": 4, "DATE": "2023-06-23"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper investigates the problem of linear regression under heteroskedastic noise with adaptive data collection.\nThe proposed algorithm begins with estimating the covariance structure via Alg 1 and then using this in Alg 2 to construct an optimal design for estimating $\\theta$. The algorithm is instance-optimal up to an additive term depending on the conditioning $\\kappa$. They conclude with experiments demonstrating their algorithm's performance.\n\nSTRENGTHS:\n* This work analyzes a useful problem and proposes a nice adaptation of RAGE to their setup.\n\nWEAKNESSES:\n* I'm not sure how common or realistic this assumption is, but the heteroskedasticity model (i.e., that $\\sigma_x^2 = x^T\\Sigma^\\star x$) seems very restrictive to me. It would be nice to have some further discussion about this and possibly about extending the algorithm to more general noise settings.\n* As they remark in the paper, the lower bound is weakened by the fact that the noise variances are considered to be known, and so the lower bound essentially becomes the same as Fiez et al. The lower bound would be much more compelling to me if it could incorporate the difficulty associated with estimating $\\Sigma^\\star$ (which could also address the current $\\kappa$-dependent gap between the upper and lower bounds).\n\nQUESTIONS:\n* Regarding the optimality of Alg. 2, what is the appropriate regime to be thinking of for the quantities in Theorems 1 and 2? For example, how should the $\\log(\\Delta^{-1})\\log(|Z|/\\delta)$ multiplicative factor differing between the upper and lower bounds compare against $\\psi^\\star$. Also, how large will $d$ be, given that the additive suboptimality terms has a $d^2$ factor in it?\n\n\n\nMinor:\n* In Alg. 1, should it say $\\lambda^\\star_x$ in line 4 and $\\lambda^\\dag_x$ in line 8? Otherwise, I don't see how the quantities computed in lines 3 and 7 are used.\n\nLIMITATIONS:\nThe authors adequately addressed the limitations.", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 2, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-07-02"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe paper considers the problem of Experimental Designs for Heteroskedastic Variance. Here, a learner has access to finite set of measurement vectors $\\mathcal{X}$ and may query observations of the following form for chosen $x \\in \\mathcal{X}$: $y = x^\\top \\theta^* + \\eta$ where $\\eta$ is mean-zero and sub-Gaussian with variance $\\sigma^2_x = x^\\top \\Sigma x$. The paper first designs a procedure to estimate $\\Sigma$ and then utilizes this estimate for two related problems featuring linear objectives: Best-arm Identification and Level-set Identification. The paper then shows that when heteroskedastic variance is present, this procedure leads to improved results for these tasks. The paper then conducts empirical experiments to validate their claims.\n\nThe main technical contribution of the paper is the procedure to estimate $\\Sigma$. If $\\theta$ were known, $(y - \\theta^\\top x)^2$ provides an unbiased estimate of $\\sigma_x^2$. However, since $\\theta^*$ is not known, the paper first estimates it via a weighted least-squares estimator where the weights are chosen to ensure that each $x \\in \\mathcal{X}$ is well-represented in the spectrum of the weighted second-moment matrix. This ensures that $\\theta$ is accurate (at least along the $x$ in $\\mathcal{X}$). Then, a similar sampling procedure is used with the vectors $\\phi_x = \\mathrm{vec} (xx^\\top)$ and noting that $(y - x^\\top \\theta^*)^2$ is an unbiased estimate of $\\sigma_x^2 = \\mathrm{vec} (xx^\\top)^\\top \\mathrm{vec} (\\Sigma)$ which is a linear equation in $\\Sigma$. A similar weighting procedure is used to obtain an estimate of $\\Sigma$ such that $\\sigma_x^2$ is well-approximated for any $x \\in \\mathcal{X}$. These estimates are then used to design improved strategies for the aforementioned estimation tasks with linear objects.\n\nOverall this is a nice paper with strong technical results. The setting seems novel and interesting and the technical analysis is non-trivial. The main question I have for the authors is that the sampling procedure incurs large sample complexities in some extreme cases. Say, $\\mathcal{X}$ corresponded to the standard basis vectors. In this setting, the error bound on estimating $\\sigma_{e_1}$ depends on $\\max_j \\sigma_{e_j}$ while intuitively this dependence seems extreme. Furthermore, the fact that $d^2$ samples are needed to estimate $\\Sigma$ only along the directions in $\\mathcal{X}$ also seems too large if $\\mathcal{X}$ is not too large. It would be great if the authors can comment on these dependencies.\n\nSTRENGTHS:\nSee main review\n\nWEAKNESSES:\nSee main review\n\nQUESTIONS:\nSee main review\n\nLIMITATIONS:\nYes", "TITLE": "Official Review", "RECOMMENDATION": 7, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-07-06"}

comments: SUMMARY:
This paper considers the online linear experimental design problem that is affected by heteroskedastic noise. The authors propose a two-phase sample splitting method aimed at estimating the unknown parameters of noise variance, utilizing the principle of G-optimal experimental designs. 

Furthermore, the study extends the application of the proposed methodology to handle fixed confidence transductive best-arm and level set identification problems under conditions of heteroskedastic noise. The paper also deals with the instance-dependent lower bounds.

STRENGTHS:
This paper does solid work and is in relatively good shape.

WEAKNESSES:
1. The problem is not well motivated. It is unclear what are the practical applications of this problem. 

2. The numerical experiments only contain some toy examples. It is unclear how it performed in a large-scale set. 

3. Comparision between the previous works are not present. 

4. The problem setting looks weird. The sensing noise variance is $x^T \Sigma x$, why it depends on the sensing vector $x$? What are the motivations for this problem setting?

QUESTIONS:
1. How is the current approach compared to the previous works? Both in a theoretical sense and numerical experiments. 

2. What motivates the problem setting, that is, the sensing noise variance is $x^T \Sigma x$.

LIMITATIONS:
See above.
TITLE: Official Review
SOUNDNESS_CORRECTNESS: 2
DATE: 2023-07-07