Sample Constrained Treatment Effect Estimation
Abstract
Treatment effect estimation is a fundamental problem in causal inference.
We focus on designing efﬁcient randomized controlled trials, to accurately estimate the effect of some treatment on a population of n individuals.
In particular, we study sample-constrained treatment effect estimation , where we must select a subset of s ⌧ n individuals from the population to experiment on.
This subset must be further partitioned into treatment and control groups.
Algorithms for partitioning the full population into treatment and control groups, or for choosing a single representative subset, have been well-studied.
The key challenge in our setting is jointing choosing a representative subset and a partition for that set.
We focus on both individual and average treatment effect estimation, under a linear effects model.
We give provably efﬁcient experimental designs and corresponding estimators, by identifying connections to discrepancy minimization and leveragescore-based sampling used in randomized numerical linear algebra.
Our theoretical results obtain a smooth transition to known guarantees when s equals the population size.
We also empirically demonstrate the performance of our algorithms.
Introduction
Experimentation has long been held as a gold standard for inferring causal effects since one can explicitly enforce independence between treatment assignment and other variables which inﬂuence the outcome of interest.
We consider the potential outcomes framework [37, 40], where each individual is associated with a control and treatment value (also called the potential outcomes ) and based on the treatment assignment, we can observe only one of these values.
Efﬁcient designs of experimentation for estimating individual treatment effects which measure the difference between treatment and control values for each individual, and the average treatment effect which measures the average individual treatment effect has been well-studied [35].
In the absence of assumptions on the functional form of the potential outcomes, the minimax optimal approach for conducting an experiment is to assign individuals to treatment or control completely at random, without consideration of baseline covariates (features) [25].
However, by considering covariates for each individual, and using additional assumptions of smoothness, substantial gains can be made in terms of the variance of the treatment effect estimate via alternative assignment procedures.
The most common approach attempts to minimize imbalance, i.e., the difference between the baseline covariates in the treatment and control groups [6, 25, 35].
While experimental designs that minimize imbalance increase the power of an experiment for a given pool of subjects, there are many practical applications where the experimenter wishes to minimize the total number of subjects who are placed into the experiment.
For example, in medicine, clinical trials may carry nontrivial risk to patients.
Within industrial applications, experiments may carry substantial costs in terms of testing changes, which decrease the quality of the user experience, or have direct monetary costs.
In this paper, we examine the problem of selecting a subset of s individuals from a larger population and assigning treatments such that the estimated treatment effect has a small error.
We consider two different estimands: individual treatment effect (ITE) and average treatment effect (ATE).
A bit more formally, we represent the d -covariates of a population of n individuals using X 2 R n ⇥ d .
We assume that the treatment and control values, denoted by y 1 , y 0 2 R n , are functions of the covariates, i.e., y 1 = f
( X , ⇣⇣⇣ 1 ) and y
= g ( X , ⇣⇣⇣ 0 ) where ⇣⇣⇣ 0 , ⇣⇣⇣ 1 2 R n are noise vectors.
The ITE for the i th individual is y 1 i  y
i and ATE is the average of all the ITE values.
We further assume a linear model, i.e., the functions f, g are linear in X and ⇣
⇣⇣ 1 , ⇣⇣⇣ 0 .
The goal is to pick a subset of s individuals and partition this subset into control and treatment groups.
For an individual i in the treatment group, we measure y 1 i , and for an individual j in the control, we measure y
j .
From this small set of measurements, we seek to estimate the ITE or ATE over the full population.
Without parametric assumptions, ITE estimation is not feasible [43].
We focus on linear models in particular, since they are important in developing theory.
E.g., in the literature on optimal designs in active learning, much of the foundational theory is built around linear models.
Identifying estimators based on linearity assumptions is an active area of study in the causal inference literature [20, 50].
Our setup is similar to active learning [42], where the goal is to minimize the number of individual labels that we access for solving linear regression or other downstream tasks.
The key difference is that we must select both a subset of individuals, and for each i , can measure only one of two labels: y 1 i or y
i .
In particular, ITE estimation can be thought of as solving two simultaneous active linear regression problems – one for the treatment outcomes and one for the control outcomes.
Thus, standard active learning-based approaches, such as [11, 12, 34], fall short.
Even when s equals the population size n , i.e., when active learning becomes trivial, our problem does not.
We must still pick a partition of the full population into treatment and control groups.
Overall, sample constrained treatment effect estimation by designing efﬁcient randomized controlled trials has received little attention, compared to various approaches that use observational data, such as [24, 39, 46].
Our Contributions.
For ITE estimation, we propose an algorithm using leverage score sampling [51], which is a popular approach to subset selection for fast linear algebraic computation.
For ATE estimation, we employ a recursive application of a covariate balancing design [20].
We provide a theoretical analysis in terms of root mean squared error (ITE) and deviation error (ATE).
Recall that we assume the treatment and control values are linear functions of the covariates plus
Gaussian noise, i.e., y 1 =
X  1 + ⇣⇣⇣ 1 and y
= X  0 + ⇣⇣⇣ 0 where ⇣⇣⇣ 1 , ⇣⇣⇣ 0 2 R n have i.i.d.
mean zero, variance  2 Gaussian entries, and  1 ,  0 2 R d are coefﬁcient vectors.
For ITE estimation, we give a randomized algorithm that selects ⇥ ( d log d ) individuals in expectation, using leverage scores, which measure the importance of an individual based on their covariates.
Our algorithm obtains, with high probability, root mean squared error O ⇣p log d/n · (     +     )
+  ⌘ (see Corollary 3.7).
We argue that this is optimal up to constants and a p log d factor, even for approaches that experiment on the full population .
The key challenge in achieving this bound is to extend leverage scores to our simultaneous linear regression setting, ensuring that we do not share samples across the treatment and control effect estimation problems.
To do this, we introduce a smoothed covariate matrix, whose leverage scores are bounded.
This ensures that, when applying independent leverage score sampling, with high probability few individuals are randomly assigned to both control and treatment, and thus removing such individuals from one of the groups does not introduce too much error.
For ATE estimation we give a randomized algorithm that selects at most s individuals for treatment/control assignment and obtains an error of e
O   / p s + (     +     ) /s
, where e
O ( · ) hides logarithmic factors (see Theorem 4.3).
The error decreases with increasing values of s and when s = n , it matches state-of-the-art guarantees due to Harshaw et al.
Our algorithm for ATE estimation is based on covariate balancing .
This is a popular approach where one attempts to assign similar individuals to the treatment and control groups, to ensure that the observed effect is attributed to the administered treatment alone.
Harshaw et al.
[20] designed an algorithm by minimizing the discrepancy of an augmented covariate matrix, which achieves low
ATE estimation error.
To extend their approach to our setting, ﬁrst, we need to select a subset of s individuals that are representative of the entire population, and then balance the covariates.
Uniform sampling or importance sampling techniques give high error here.
Instead, we employ a recursive strategy, which repeatedly partitions the individuals into two subsets by balancing covariates, and selects the smaller subset to recurse on, until we have selected at most s individuals.
We observe that our techniques for ITE and ATE estimation should extend to the setting when the outcomes are non-linear functions of the covariates, which are linear in some higher-dimensional kernel space.
This is immediate for our discrepancy minimization design for ATE, which only requires knowing the pairwise inner products of the covariate vectors.
For ITE estimation, leverage score sampling for kernel ridge regression [3] is most likely applicable.
Extensions to broader classes of non-linear models are beyond the scope of this work, but they are an interesting future direction.
Finally, in Section 5, we provide an empirical evaluation of the performance of our ITE and ATE estimation methods, comparing against uniform sampling and other baselines on several datasets.
Our results suggest that our techniques can help reduce the costs associated with running a randomized controlled trials substantially using only a small fraction of the population.
Other Related Work.
For ATE estimation, the most well-studied approaches to experiment design are covariate balancing and randomization.
A variety of design techniques have been studied based on these approaches, such as blocking [18], matching [23, 45], rerandomization [30, 35], and optimization [25].
Using observational data, treatment effect estimation using covariate regression adjustment [31] and various active learning-based sampling techniques have gained recent attention [24, 38, 46].
Compared to ATE, estimating ITE is signiﬁcantly harder and has received attention only recently using machine learning methods [7, 43, 49].
There has been a lot of recent work on efﬁcient experimental designs to minimize experimental costs, in various domains, such as causal discovery [1, 2, 16, 17, 28, 44], multi-arm bandits [4, 27, 36], and group testing [9, 10, 15].
Preliminaries
Notation.
We use bold capital letters, e.g., X to denote matrices and bold lowercase letters, e.g., y to denote vectors.
We use X
[ i, :] and X [: , j ] to denote the i th row and j th column of X respectively, which we always view as column vectors.
The i th largest singular value of X is denoted by
i ( X ) .
For any vector x , the Euclidean norm or the ` 2 -norm is denoted by k x k .
For a population of n individuals, we represent each with an integer in [ n ] where we denote [ n ] def = { 1 , 2 , · · · , n } .
Each individual j 2 [ n ] is associated with a treatment and a control value, denoted yj , y
j 2 R + , respectively.
The vectors associated with all n treatment and control values are denoted y 1 and y 0 .
Additionally, each individual is associated with a d -dimensional covariate vector.
Combined, they comprise the rows of the covariate matrix X 2 R n ⇥ d .
In this paper, we consider the ﬁnite population framework, where the potential outcomes of individuals are ﬁxed and the randomness is only due to treatment assignment [13].
We make the
SUTVA assumption, i.e., the treatment outcome value of any individual is independent of treatment assignments of others in the population [48].
Assumption 2.1 (Linearity Assumption) .
Under the linearity assumption, the treatment and control values are a linear function of the covariates.
Formally, for some  0 ,  1 2 R d ,
[Equation]
and
[Equation]
where ⇣⇣⇣ 1 , ⇣⇣⇣ 0 2 R n are noise vectors, with each coordinate drawn independently from the Gaussian distribution with zero mean and variance  2 , i.e., N (0 ,  2 ) .
We further assume that X is rownormalized, i.e., k X
[ i, :] k
 1 8 i 2 [ n ] .
Deﬁnition 2.2 (Individual Treatment Effect) .
Given a population of n individuals, the individual treatment effect (ITE) of j 2 [ n ] is the difference between the treatment and control values:
[Equation]
Deﬁnition 2.3 (Average Treatment Effect) .
Given a population of n individuals, the average treatment effect (ATE), denoted by ⌧ , is the average individual treatment effect:
[Equation]
Deﬁnition 2.4 (Root Mean Squared Error) .
For a set of estimated individual treatment effects, d
ITE ( j ) for j 2 [ n ] , the root mean squared error ( RMSE ) is deﬁned as:
[Equation]
Deﬁnition 2.5 (Leverage Score) .
Given a matrix X 2 R n ⇥ d , the leverage score of j th row X [ j, :] , denoted by ` j ( X ) , is deﬁned as:
[Equation]
where + denotes the Moore–Penrose pseudo-inverse.
Individual Treatment Effect Estimation
We now describe our algorithm for ITE estimation.
The algorithm identiﬁes a subset of the population to experiment on, using importance based sampling techniques, that are well-studied in randomized numerical linear algebra [51].
Missing proof details in this section are presented in Appendix A.1.
Overview of our approach.
Under the linearity assumption (Assumption 2.1), we can reformulate the problem of estimating the ITE for every individual as simultaneously solving two linear regression instances: one for control and one for treatment, i.e., we regress y 0 , y 1 on X .
However, there are two challenges: 1) we would like to solve these regression problems using measurements from just a small subset of s individuals and 2) we only have access to either the control or treatment measurement y
j or y 1 j for any individual in this set.
To tackle the ﬁrst challenge, we use a sampling technique based on the importance of each row in X , captured via its leverage score (Defn. 2.5).
Intuitively, we want to select s individuals (or equivalently rows) that capture the entire row space of X and use them to estimate the ITE of all other individuals.
Leverage scores capture the importance of a row in making up the row space.
E.g., if a row is orthogonal to all the other rows, it’s leverage score will be the maximum value of 1 .
Unfortunately, if we apply leverage score sampling independently to the regression problems for yand y 1 , rows with high leverage leverage scores may be sampled for both instances.
This presents a problem, since we can only read at most one of y
j or y 1 j .
To mitigate this issue, we construct a smoothed matrix X ⇤ , which consists of X projected onto its singular vectors with high singular values.
Intuitively, this dampens the effects of high leverage score ‘outlier’ rows that don’t contribute signiﬁcantly to the spectrum of X .
Formally, we prove that the maximum leverage score of X ⇤ is bounded, which let’s us solve our two regression problems via independent sampling.
There will be few repeated samples across our subsets, which introduce minimal error.
3.1 Leverage Score Sampling
For some   0 , to be ﬁxed later, we deﬁne a smoothed matrix for X , the projection onto singular vectors with high singular values, as follows:
Deﬁnition 3.1 (Smoothed matrix) .
Given X 2 R n ⇥ d with singular value decomposition X =
U ⌃ V T , let  ⇤ be the set of indices corresponding to singular values greater than p  , i.e.,  ⇤ def
= { i |  i ( X )  p  } ; we denote d 0 def = |  ⇤ | .
Let ⌃ ⇤ =
⌃ (  ⇤ ,  ⇤ ) denote the principal sub-matrix of ⌃ associated with these large singular values.
Similarly, let U ⇤ 2 R n ⇥ d 0 ,
V ⇤ 2 R d ⇥ d 0 be the associated column sub-matrices of U and V .
Then, we deﬁne:
[Equation]
Sampling Matrix.
Our algorithm will sample individuals, corresponding to rows of the smoothed matrix of X , i.e., X ⇤ , independently – the i th row is included in the sample with some probability ⇡⇡⇡ i .
Let the set of rows sampled be denoted by S .
We can associate a sampling matrix W with S .
The j th row of W is associated with the j th element in the set S (under some ﬁxed order).
If the j th element in S is the row for individual i for some i 2 [ n ] , then, W [ j, :] is equal to e i / p ⇡⇡⇡ i .
Here, e i 2 R n denotes the i th standard basis vector.
In this way, WX ⇤ consists of the subset of rows sampled in
S , reweighted by the inverse squareroot of their sampling probabilities, which is necessary to keep expectations correct in solving the linear regression.
[Table] [Table]
d
Algorithm S
AMPLING -ITE.
We perform row sampling twice, with probabilities proportional to the leverage scores of X ⇤ , to construct two sets S 0 , S 1 .
See the discussion below for the exact deﬁnition of the sampling probabilities ⇡⇡⇡ i , which are proportional to the leverage scores of X ⇤ .
These two sets are used to estimate the vectors y 0 and y 1 , respectively.
It is possible that a row gets included in both
S 0 and S 1 .
In that case, we simply remove the row from S 1 .
As a result, j th row is included in Swith probability ⇡⇡⇡ j · (1  ⇡⇡⇡ j ) for every j 2 [ n ] .
We construct sampling matrices W 0 and W 1 using probabilities ⇡⇡⇡ and ⇡⇡⇡ (1  ⇡⇡⇡ ) respectively.
Finally, in Algorithm 1, we solve the following linear regressions, for i = 0 , 1 separately:
[Equation]
[Equation]
Our estimate for each ITE ( j ) , denoted by d
ITE ( j ) is set to j th entry of the vector X ⇤
[Equation]
Observe that by construction, S 0 \ S 1 is empty.
This ensures that we have access to only one of y
j or y 1 j for any individual j in solving the above two subsampled regression problems.
We note that in Algorithm 1, we could remove j from one of S 0 , S 1 , or with equal probability from either of the two sets, and obtain the exact same guarantees.
3.2 Theoretical Guarantees
First, we bound the error due to sampling.
Critically, we show that the leverage scores of X ⇤ , and in turn the probabilities ⇡⇡⇡ , are bounded by 1 /  .
Thus, the sampling probabilities for S 1 , ⇡⇡⇡ (1  ⇡⇡⇡ ) are not too far from ⇡⇡⇡ itself.
As we assume the row norms of X are bounded by 1 , the row norms of X ⇤ are also bounded.
Thus, there can be no rows in X ⇤ that are nearly orthogonal to all other rows – i.e., there can be no rows with very high leverage scores.
Such rows would lead to small singular values.
However, we know that the smallest singular value of X ⇤ is at least p  .
In particular, we prove:
Claim
[Equation]
[ n ] .
Setting ⇡⇡⇡ .
It is well known that if we sample rows of X ⇤ with probabilities ⇡⇡⇡ proportional to the leverage scores, we will obtain a (1 ± ✏ ) relative error approximation for linear regression [41].
The result of Sarlos [41] applies to sampling s rows with replacement , each equal to j with probability ⇡ ⇡ ⇡ j / k ⇡ ⇡ ⇡ k .
It is not hard to observe that it extends to the variant where each row is included in the sample independently with similar probability.
Therefore, we have:
Lemma 3.3 (Follows from [41]) .
For X 2 R n ⇥ d , y 2 R n , let S ✓ [ n ] include each j 2 [ n ] independently with probability ⇡ ⇡ ⇡ j satisfying ⇡ ⇡ ⇡ j  min  1 , ` j ( X ) · c · [log(rank( X ))
+ 1 ✏ ] for some large enough constant c .
Let W 2 R | S | ⇥ n be a sampling matrix that includes row e j / p ⇡⇡⇡ j if j 2 S , where e j 2 R n is the j th standard basis vector.
Let
[Equation]
Then, E [ | S | ] =
P nj =1 ⇡⇡⇡ j and with probability  1   :
[Equation]
If the ⇡⇡⇡ j ’s are within constants of the required bound, E [ | S | ]
= O
[Equation]
✏  .
Note that the bound on E [ | S | ] follows from the well known fact that the sum of leverage scores, is equal to the rank, i.e.,
[Equation]
The sampling probabilities are set to ⇡⇡⇡ j = min { 1 , ` j ( X ⇤ ) · c 0 · [log(rank( X ⇤ ))
+ 30 / ✏ ] } for some constant c 0  2 c , where c is the constant in Lemma 3.3.
Thus, by the lemma, we will have, with probability  29 / 30 ,
[Equation]
It remains to show that we will have a similar guarantee for the control group.
The rows in S 1 are included independently with probability ⇡⇡⇡ j · (1  ⇡⇡⇡ j ) .
If we can prove that ⇡⇡⇡ j · (1  ⇡⇡⇡ j )  ⇡⇡⇡ j 2 , then Lemma 3.3 will still apply, since we have set c 0 = 2 c .
To do so, it sufﬁces to argue that ⇡⇡⇡ j  1 / 2 by setting the parameters appropriately.
Claim 3.4.
If  = 4 c 0 max { log(rank( X ⇤ )) , 30 / ✏ } and ⇡⇡⇡ j = min { 1 , ` j ( X ⇤ ) · c 0 · [log(rank( X ⇤ ))
+ 30 / ✏ ] } , we have ⇡⇡⇡ j  1 / 2 for every j 2 [ n ] .
Proof.
⇡⇡⇡ j 
[Equation]
log
[Equation]
X
[Equation] [Equation]
✏
[Equation]

[Equation]
[Equation]
c
[Equation]
log(rank( X ⇤ ))
+ 30 / ✏ ] (Claim 3.2)
[Equation]
In Appendix A.1, we argue that using the smoothed matrix X ⇤ introduces an error of p  .
Combining all of them, we have the following corollary:
Corollary 3.5.
Suppose  and ⇡⇡⇡ j are set as in Claim 3.4, for some sufﬁciently large constant c 0 .
Then, Algorithm S
AMPLING -ITE satisﬁes, for i = 0 , 1 , with probability at least 14 / 15 :
[Equation]
RMSE Guarantees.
The root mean squared error (Defn. 2.4) for the ITE estimates is given by:
[Equation]
By setting ✏ = 120 c 0 d log d/s in Corollary 3.5, we get the following theorem for our Algorithm 1:
Theorem 3.6.
Suppose s  120 c 0 d log d .
There is a randomized algorithm that selects a subset
S ✓ [ n ] of the population with E [ | S | ]  s , and, with probability at least 9 / 10 , returns
ITE estimates d
ITE ( j ) for all j 2 [ n ] with error:
[Equation]
For the sake of simplicity of analysis, we used a constant success probability in Theorem 3.6.
All our claims can easily be updated with a general failure probability of  , with a dependence of p 1 /  , using Lemma 3.3.
The corollary below follows immediately from Theorem 3.6.
Corollary 3.7 (Main ITE Error Bound) .
The root mean squared error obtained by Algorithm 1 is minimized when s = ⇥
( d log d ) and is given by:
[Equation]
Our upper bound on RMSE for Algorithm 1 increases with s , if s grows strictly faster than d log d asymptotically, i.e., s = !
( d log d ) .
Therefore, to obtain low error, we set s = c · d log d for some constant c , even if the sample constraint allows for larger values.
We believe this is an artifact of our analysis.
In Section 5, we observe empirically that the error decreases with increasing s .
Remark.
We observe that the RMSE bound in Corollary 3.7 is nearly optimal, even for algorithms that experiment on the full population .
The O (  ) term cannot be improved by more than constants, as a consequence of our noise model (see Assumption 2.1).
Even if we knew the true  1 and  0 , our
RMSE would be O (  ) .
The term (   +   ) / p n is also necessary.
Suppose the matrix X is such that all rows, except row j , are zero vectors.
Row j is a standard basis vector, i.e., its i th entry is 1 for some i .
Suppose also that  1 and  0 are both independently set to the same standard basis vector with probability 1 / 2 , and set to zero otherwise.
Then, with probability 1 / 2 , ITE ( j ) =
and with probability 1 / 2 ,
ITE ( j ) = ± 1 .
No algorithm which observes just one of y 1 j or y
j can obtain expected error o(1) in estimating ITE ( j ) .
That is, no algorithm can obtain RMSE o (1 / p n )
=
o  (   +   ) / p n  .
Average Treatment Effect Estimation
In this section, we describe our approach for estimating the average treatment effect, under the sample constraint, by building upon a recent work on efﬁcient experimental design by Harshaw et al.
Missing details from this section are collected in Appendix A.2.
Horvitz-Thompson Estimator.
Suppose S + ✓ [ n ] is the population assigned to the treatment group and S  = [ n ] \ S
+ is the remaining population, i.e., the control group.
A well-studied estimator for estimating the average treatment effect is the Horvitz-Thompson estimator [22], denoted by b ⌧ .
If every individual is assigned to S + (or S  ) with probability 0 .
5 , then, b ⌧ is deﬁned as follows:
[Equation]
Input: Covariate matrix X 2 R n ⇥ d , number of experiments to be run s .
Output: Estimate for ATE.
1: Set t = 1 , Z t := X , n t = n .
2: while True do 3:
Z + t , Z  t
G RAM -S CHMIDT -W
ALK ( Z t ,  0 ) where  0 = log(16 log( n/s )) .
4: if n t  s then 5: break 6: else if size ( Z + t )  size ( Z  t ) then 7:
Set Z t +
Z  t and n t + size ( Z  t ) .
8: else 9:
Set Z t +
Z + t and n t + size ( Z + t ) .
10: end if 11: t t +12: end while 13: Use Z + t , Z  t to construct the ATE estimator as: b ⌧ s = 2 t /n · ⇣P j 2 Z + t y 1 j
P
[Equation] [Equation]
[Equation]
y
[Equation]
⌘ .
14: return b ⌧ s .
Harshaw et al.
[20] present an experimental design based on the Gram-Schmidt-Walk algorithm for discrepancy minimization [8].
Their Gram-Schmidt-Walk design produces a random partition of the population with a good balance in every dimension, i.e., control and treatment groups have similar covariates.
For the Horvitz-Thompson estimator, they give a tradeoff between covariate balancing and robustness (estimation error).
Formally, they obtain:
Lemma 4.1 (Proposition 3 in [20]) .
For all  > 0 , with probability at least 1
[Equation]
⌘ , the Gram-Schmidt-Walk design satisﬁes: | b ⌧  ⌧
[Equation]
  , where L =
[Equation]
✓
[Equation]
Overview of R ECURSIVE -C OVARIATE -B
ALANCING .
Our main idea in Algorithm 2 is to partition the population using the Gram-Schmidt-Walk design (GSW) recursively until the total size of population that we can experiment on reduces to s .
In each recursive call, we start by partitioning the available individuals Z t into treatment and control groups, denoted by Z + t , Z  t using GSW.
Next, we identify the smaller of these two subsets, say Z + t and recurse on Z + t .
We stop after k recursive calls when there are only s individuals to experiment on, i.e., | Z + k
[ Z  k |  s .
Finally, we construct our estimator b ⌧ s , similar to the Horvitz-Thompson estimator, by scaling the treatment and control contributions due to Z + k and Z  k using a factor 2 k .
We note that our experimental design ensures that every individual is assigned to treatment or control with equal probability.
This implies that on expectation, the sizes of the treatment and control groups are equal (for every partitioning).
However, when we consider a particular assignment, it is possible that the size of the smaller partition is not exactly half of the population.
As a result, the total number of samples used might be smaller by a factor of at most 2.
Theoretical Guarantees.
Our analysis approach, inspired by the coreset construction for discrepancy minimization [26], is based on the observation that if we can obtain good estimates for the contributions
P
[Equation]
i , we obtain a good estimate for ATE ( ⌧ ).
Using the next lemma, we argue that after a call to GSW algorithm that partitions [ n ] into the sets S + and S  , we can obtain additive approximations of
[Equation]
and
P
[Equation]
[ n ] y 0 i .
Our approximations are the contributions of treatment and control values in S + and S  scaled appropriately, i.e.,
[Equation]
Lemma 4.2.
Suppose the Gram-Schmidt-Walk design [20] partitions the population [ n ] into two disjoint groups S + and S  .
Under the linearity assumption, with probability 1  1 / 3 log( n/s ) , for both the control and treatment groups, the following holds:
[Equation]
Building upon the previous lemma, we argue in Theorem 4.3 that the additive approximation errors obtained from repeated use of GSW in our algorithm R ECURSIVE -C OVARIATE -B
ALANCING result in a low estimation error.
Theorem 4.3 (Main ATE Error Bound) .
The estimator b ⌧ s in Algorithm R ECURSIVE -C OVARIATEB
ALANCING obtains the following guarantee, with probability at least 2 / 3 :
[Equation]
Remark.
When s = n , the above theorem matches the guarantees obtained by GSW design described in Lemma 4.1.
Moreover, we obtain a better dependence compared to sampling s rows uniformly at random and using the y 1 , y 0 values of the sampled rows to estimate the population mean of treatment and control groups in ATE.
An application of standard concentration inequalities or the central limit theorem, will yield a multiplicative factor increase in one of the error terms, with a dependence of e
O  1 /s · k X k 2 (     +     )  , instead of the e
O  1 /s · (     +     )  obtained by our algorithm, where k X k 2 denotes the spectral norm of X and e
O ( · ) hides the logarithmic factors.
Experimental Evaluation
In this section, we provide an evaluation of our algorithms on various semi-synthetic datasets.
Missing details about data generation and additional results are collected in Appendix A.3.
Data Generation.
We evaluate our approaches on ﬁve datasets: (i) IHDP.
This contains data regarding the cognitive development of children, and consists of 747 samples with 25 covariates describing properties of the children and their mothers, and whose outcome values are simulated [21, 14].
(ii)
Twins.
This contains data regarding the mortality rate in twin births in the USA between 19891991 [5].
Following the work of [32], we select twins belonging to same-sex, with weight less than 2kg, resulting in about 11984 twin pairs, each with 48 covariates.
We use the post-treatment mortality outcomes of the twins as potential outcomes.
(iii) LaLonde.
This contains data regarding the effectiveness of a job training program on the real earnings of an individual after completion of the program [29], which is also the outcome value.
The corresponding covariate matrix contains 445 rows and 10 covariates per row.
(iv) Boston.
This is constructed based on the housing prices in the Boston area [19].
The outcome value for each sample represents the median house price.
The corresponding covariate matrix contains 506 rows and 12 covariates per row.
(v) Synthetic.
We construct a covariate matrix X 2 R 2000 ⇥ 25 , using an approach due to [33].
There is a high disparity in leverage score values in X , similar to what we observe in other datasets.
Using a random linear function on X and adding Gaussian noise, we generate the potential outcomes.
For IHDP and Twins datasets, we use the simulated values for potential outcomes, similar to Shalit et al.
[43] and Louizos et al.
For the Synthetic dataset, we simulate values for the outcomes using linear functions of the covariate matrix.
For Boston, Lalonde datasets, as we have access to only one of the outcome values, we chose to compare our algorithms for a ﬁxed shift in treatment effect (i.e., the true treatment effect is equal to a constant), similar to Arbour et al.
[Figure] [Figure]
Baselines.
(i) ATE.
We compare the performance of our Algorithm R ECURSIVE -C OVARIATEB
ALANCING (referred to as ‘Recursive-GSW’ ) to three baselines: (i) Uniform .
We sample s rows uniformly at random and assign them to treatment and control groups with equal probability.
By scaling the total sum of treatment values from the sampled set by the inverse sampling probability, we estimate the contribution of treatment values in ATE and follow a similar procedure for the control group.
(ii) GSW-pop .
We use the GSW algorithm to partition the full population and return the estimate obtained using the Horvitz-Thompson estimator for ATE.
(iii) Complete Randomization .
We partition the population into treatment and control using complete randomization, i.e., with equal probability, and return the estimate obtained using the Horvitz-Thompson estimator for ATE.
The last two baselines are overall n individuals rather than a subset of size s .
(ii) ITE.
We compare the performance of our Algorithm S
AMPLING -ITE (referred to as ‘Leverage’ ) with respect to three baselines: (i) Uniform .
We run Algorithm 1 on X and uniform sampling distribution given by ⇡⇡⇡ j = s/n 8 j .
(ii) Leverage-nothresh .
We run Algorithm 1 on X , instead of
X ⇤ with the probability distribution ⇡⇡⇡ j / ` j ( X ) 8 j .
(iii) Lin-regression .
This captures the best linear ﬁt regression error, i.e., assuming we have access to both y 1 , y 0 , we regress these vectors on X to obtain  1 ,  0 , and use the resultant ITE estimates X  1  X  0 .
Evaluation.
To evaluate the performance of average treatment effect estimation ( ⌧ ) on the datasets, we compare the deviation error of the estimator b ⌧ s , given by | b ⌧ s  ⌧ | for different sample sizes.
To evaluate the performance of individual treatment effect estimates, we compare the root mean squared error RMSE (see Defn. 2.4) for different sample sizes.
Results.
For every dataset, we run each experiment for 1000 trials and plot the mean using a colored line.
Also, we shade the region between 30 and 70 percentile around the mean to signify the conﬁdence interval as shown in Figures 1, 2 representing ATE and ITE results respectively.
(i) ATE .
For all datasets, we observe that the deviation error obtained by our algorithm labeled as
Recursive-GSW in Figure 1, is signiﬁcantly smaller than that of Uniform baseline.
Surprisingly, for the IHDP dataset, our approach is signiﬁcantly better than Complete-randomization , for all sample sizes, including using just 10% of data.
For all the remaining datasets using a sample of size 30% , we achieve the same error (up to the conﬁdence interval) as that of Complete-randomization .
Complete randomization is one of the most commonly used methods for experimental design and our results indicate a substantial reduction in experimental costs.
For IHDP dataset, a sample size of about 10% of the population is sufﬁcient to achieve a similar error as that of GSW-pop .
For the remaining
datasets
, we observe that for sample sizes of about 30% of the population, the deviation error obtained by our algorithm is
within the shaded conﬁdence interval of the error obtained by GSW-pop .
Therefore, for a speciﬁed error tolerance level for ATE, we
can
reduce
the
associated experimental costs using just a small subset of the dataset using our algorithm
(ii) ITE .
For all sample sizes, we observe that the RMSE obtained by our algorithm labeled as Leverage in Figure 2, is signiﬁcantly smaller than that of all the other baselines, including Uniform and Leverage-nothresh .
E.g., we observe that when the sample size is 20% of the population in
IHDP
dataset, the error obtained
by
Leverage is at least 50% times smaller than that of Uniform
and
Leverage-nothresh .
For the Synthetic dataset, the error obtained by Leverage is extremely close to that of the error due to the best linear ﬁt, Lin-regression (
see
the
zoomed
in part of the
ﬁgure
Similar to ATE results, our algorithms result in a reduction of experimental costs for ITE estimation using
only a fraction
of the dataset.
Conclusion
We study the sample constrained treatment effect estimation problem and give efﬁcient algorithms for both ITE and ATE estimation.
Our empirical evaluation shows that our algorithms, using only a fraction of the data, perform well compared to popular baselines that are widely used and require running experiments on the entire population.
There are several interesting directions for future work.
It would be interesting to study sample constrained treatment effect estimation under interference [47].
For ITE estimation, we leave it as an open question to extend our approach to give an algorithm with an error growing smaller with s for all values of s  n .
Moreover, the p log d factor in Corollary 3.
likely can be improved, using recent work which improves standard leverage score sampling bounds by a log d factor [11], yielding a bound that is optimal up to constants.
For ATE estimation, the
Horvitz-Thompson estimator can include sampling probabilities different from 0 .
It is an interesting open question to extend our recursive balancing approach when the estimator contains arbitrary probabilities.
Acknowledgements
Most of this work was done while R. Addanki was a student at the University of Massachusetts
Amherst and an intern at Adobe.
Part of this work was done while R. Addanki was a visiting student at the Simons Institute for the Theory of Computing.
This work was supported by a Dissertation Writing
Fellowship awarded by the Manning College of Information and Computer Sciences, University of
Massachusetts Amherst to R. Addanki.
R. Addanki and C. Musco were additionally supported in part by NSF grants CCF-2046235, IIS-1763618, CCF-1934846, CCF- and CCF-1637536, along with Adobe and Google Research Grants.
Checklist
The checklist follows the references.
Please read the checklist guidelines carefully for information on how to answer these questions.
For each question, change the default [TODO] to [Yes] , [No] , or [N/A] .
You are strongly encouraged to include a justiﬁcation to your answer , either by referencing the appropriate section of your paper or providing a brief inline description.
Please do not modify the questions and only use the provided macros for your answers.
Note that the
Checklist section does not count towards the page limit.
In your paper, please delete this instructions block and only keep the Checklist section heading above along with the questions/answers below.
For all authors... (a) Do the main claims made in the abstract and introduction accurately reﬂect the paper’s contributions and scope?
[Yes] (b) Did you describe the limitations of your work?
[Yes] (c) Did you discuss any potential negative societal impacts of your work?
[N/A]
We do not foresee any direct negative outcomes of our work.
As with all theoretical work, our results are based on simpliﬁed models of the real world, and this is important to keep in mind.
In designing randomized controlled trials using our algorithms, e.g., a medical study, attempting to select individuals or samples based on an error threshold, while ignoring other complex factors would be irresponsible.
(d) Have you read the ethics review guidelines and ensured that your paper conforms to them?
[Yes] 2.
If you are including theoretical results... (a) Did you state the full set of assumptions of all theoretical results?
[Yes] (b) Did you include complete proofs of all theoretical results?
[Yes] 3.
If you ran experiments... (a) Did you include the code, data, and instructions needed to reproduce the main experimental results (either in the supplemental material or as a URL)?
[Yes] Our code is publicly accessible using the following github repository.
All the datasets are already public, and we have provided instructions on all the parameters along with the experimental setup in Sections 5 and Appendix A.3.
(b) Did you specify all the training details (e.g., data splits, hyperparameters, how they were chosen)?
[Yes] (c) Did you report error bars (e.g., with respect to the random seed after running experiments multiple times)?
[Yes] (d) Did you include the total amount of compute and the type of resources used (e.g., type of GPUs, internal cluster, or cloud provider)?
[Yes] 4.
If you are using existing assets (e.g., code, data, models) or curating/releasing new assets... (a) If your work uses existing assets, did you cite the creators?
[Yes] (b) Did you mention the license of the assets?
[Yes] (c) Did you include any new assets either in the supplemental material or as a URL?
[Yes] (d) Did you discuss whether and how consent was obtained from people whose data you’re using/curating?
[N/A] All our datasets are already public.
(e) Did you discuss whether the data you are using/curating contains personally identiﬁable information or offensive content?
[N/A] 5.
If you used crowdsourcing or conducted research with human subjects... (a) Did you include the full text of instructions given to participants and screenshots, if applicable?
[N/A] (b) Did you describe any potential participant risks, with links to Institutional Review Board (IRB) approvals, if applicable?
[N/A] (c) Did you include the estimated hourly wage paid to participants and the total amount spent on participant compensation?
[N/A]