Sample Constrained Treatment Effect Estimation

Extracted main text — title through conclusion, appendix excluded. This is what our citation measures are computed over, published so the extraction can be checked by eye.

50,893 characters · 12 sections · 49 citation commands

Rendered from LaTeX for readability, not typeset faithfully. Citation keys are highlighted; maths is left as source; figures, tables and equation environments are summarised rather than reproduced; unrecognised commands are greyed out so nothing is silently dropped. Email addresses are removed.

abstractTreatment effect estimation is a fundamental problem in causal inference. We focus on designing efficient 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 \ll n$ individuals from the population to experiment on. This subset must be further partitioned into treatment and control groups. Algorithms for partitioning the entire population into treatment and control groups, or for choosing a single representative subset, have been well-studied. The key challenge in our setting is jointly 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 efficient experimental designs and corresponding estimators, by identifying connections to discrepancy minimization and leverage-score-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. \let\thefootnote\relax\footnotetext{Author ordering is alphabetical. A preliminary version of this work appeared in the proceedings of 36th Conference on Neural Information Processing Systems (NeurIPS 2022).}

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 influence the outcome of interest. We consider the potential outcomes framework neyman, rubin1980randomization, 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. Efficient 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 morgan2012rerandomization. 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) kallusapriori. 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 pmlr-v130-arbour21a, kallusapriori, morgan2012rerandomization.

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 ${\mathbf X} \in \mathbb{R}^{n \times d}$. We assume that the treatment and control values, denoted by $\mathbf{y}^1, \mathbf{y}^0 \in \mathbb{R}^n$, are functions of the covariates, i.e., $\mathbf{y}^1 = f({\mathbf X}, \pmb{\zeta}^1)$ and $\mathbf{y}^0 = g({\mathbf X}, \pmb{\zeta}^0)$ where $\pmb{\zeta}^0, \pmb{\zeta}^1 \in \mathbb{R}^n$ are noise vectors. The ITE for the $i^{th}$ individual is $\mathbf{y}^1_i - \mathbf{y}^0_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 ${\mathbf X}$ and $\pmb{\zeta}^1, \pmb{\zeta}^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 $\mathbf{y}^1_i$, and for an individual $j$ in the control, we measure $\mathbf{y}^0_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 shalit2017estimating. 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 harshaw2019balancing,wager2016high.

Our setup is similar to active learning settles2009active, 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: $\mathbf{y}^1_i$ or $\mathbf{y}^0_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 chen2019active, cohen2013stability, mahoney2011randomized, 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 efficient randomized controlled trials has received little attention, compared to various approaches that use observational data, such as jesson2021causal, qin2021budgeted, sundin2019active.

Our Contributions

For ITE estimation, we propose an algorithm using leverage score sampling woodruff2014sketching, 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 harshaw2019balancing. 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., $\mathbf{y}^1 = {\mathbf X} \pmb{\beta}^1 + \pmb{\zeta}^1$ and $\mathbf{y}^0 = {\mathbf X} \pmb{\beta}^0 + \pmb{\zeta}^0$ where $\pmb{\zeta}^1, \pmb{\zeta}^0 \in \mathbb{R}^n$ have i.i.d. mean zero, variance $\sigma^2$ Gaussian entries, and $\pmb{\beta}^1,\pmb{\beta}^0 \in \mathbb{R}^d$ are coefficient vectors.

\paragraph{ITE estimation.} For ITE estimation, we give a randomized algorithm that selects $\Theta(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\Bigl(\sqrt{{\log d}/{n}} \cdot (\left\lVert\pmb{\beta}^1\right\rVert + \left\lVert\pmb{\beta}^0\right\rVert ) + \sigma \Bigl)$ (see Corollary (ref)). We argue that this is optimal up to constants and a $\sqrt{\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.

\paragraph{ATE estimation.} For ATE estimation we give a randomized algorithm that selects at most $s$ individuals for treatment/control assignment and obtains an error of $\widetilde O\left({\sigma}/{\sqrt{s}} + (\left\lVert\pmb{\beta}^1\right\rVert + \left\lVert\pmb{\beta}^0\right\rVert)/{s} \right)$, where $\widetilde O(\cdot)$ hides logarithmic factors (see Theorem (ref)). The error decreases with increasing values of $s$ and when $s = n$, it matches state-of-the-art guarantees due to Harshaw et al. harshaw2019balancing.

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. harshaw2019balancing 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, first, 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 alaoui2015fast 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 (ref), 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 greevy2004optimal, matching imai2008variance, stuart2010matching, rerandomization li2018asymptotic, morgan2012rerandomization, and optimization kallusapriori. Using observational data, treatment effect estimation using covariate regression adjustment lin2013agnostic and various active learning-based sampling techniques have gained recent attention jesson2021causal, nikolaev2013balance, sundin2019active. Compared to ATE, estimating ITE is significantly harder and has received attention only recently using machine learning methods athey2016recursive, shalit2017estimating, wager2018estimation. There has been a lot of recent work on efficient experimental designs to minimize experimental costs, in various domains, such as causal discovery addanki2020efficient, addanki2021intervention, eberhardt2007causation, ghassami2018budgeted, kocaoglu2017cost, shanmugam2015learning, multi-arm bandits alieva2021robust,kazerouni2021best, nair2021budgeted, and group testing bondorf2020sublinear, chan2011non, du2000combinatorial.

Preliminaries

Notation. We use bold capital letters, e.g., ${\mathbf X}$ to denote matrices and bold lowercase letters, e.g., ${\mathbf y}$ to denote vectors. We use ${\mathbf X}[i, :]$ and ${\mathbf X}[:, j]$ to denote the $i^{th}$ row and $j^{th}$ column of ${\mathbf X}$ respectively, which we always view as column vectors. The $i^{th}$ largest singular value of ${\mathbf X}$ is denoted by $\sigma_i({\mathbf X})$. For any vector ${\mathbf x}$, the Euclidean norm or the $\ell_2$-norm is denoted by $\left\lVert{\mathbf x}\right\rVert$.

For a population of $n$ individuals, we represent each with an integer in $[n]$ where we denote $[n] \mathbin{\stackrel{\rm def}{=}} \{1, 2, \cdots, n\}$. Each individual $j \in [n]$ is associated with a treatment and a control value, denoted $\mathbf{y}^1_j, \mathbf{y}^0_j \in \mathbb{R}^+$, respectively. The vectors associated with all $n$ treatment and control values are denoted $\mathbf{y}^1$ and $\mathbf{y}^0$. Additionally, each individual is associated with a $d$-dimensional covariate vector. Combined, they comprise the rows of the covariate matrix ${\mathbf X} \in \mathbb{R}^{n \times d}$.

In this paper, we consider the finite population framework, where the potential outcomes of individuals are fixed and the randomness is only due to treatment assignment ding2017bridging. We make the {SUTVA} assumption, i.e., the treatment outcome value of any individual is independent of treatment assignments of others in the population wager2020stats.

assumption[Linearity Assumption] Under the linearity assumption, the treatment and control values are a linear function of the covariates. Formally, for some $\pmb{\beta}^0, \pmb{\beta}^1 \in \mathbb{R}^d$, \[ \mathbf{y}^1 = {\mathbf X} \pmb{\beta}^1 + \pmb{\zeta}^1 \mbox{ and } \mathbf{y}^0 = {\mathbf X} \pmb{\beta}^0 + \pmb{\zeta}^0, \] where $\pmb{\zeta}^1, \pmb{\zeta}^0 \in \mathbb{R}^n$ are noise vectors, with each coordinate drawn independently from the Gaussian distribution with zero mean and variance $\sigma^2$, i.e., $N(0, \sigma^2)$. We further assume that ${\mathbf X}$ is row-normalized, i.e., $\left\lVert{\mathbf X}[i, :]\right\rVert \le 1 \ \forall i \in [n]$.

definition[Individual Treatment Effect] Given a population of $n$ individuals, the individual treatment effect (ITE) of $j \in [n]$ is the difference between the treatment and control values: \[ \textsc{ITE}(j) \mathbin{\stackrel{\rm def}{=}} \mathbf{y}^1_j - \mathbf{y}^0_j.\]

definition[Average Treatment Effect] Given a population of $n$ individuals, the average treatment effect (ATE), denoted by $\tau$, is the average individual treatment effect: \[ \tau \mathbin{\stackrel{\rm def}{=}} \frac{1}{n} \sum_{j \in [n]} \textsc{ITE}(j) = \frac{1}{n} \sum_{j \in [n]}\mathbf{y}^1_j - \mathbf{y}^0_j.\]

definition[Root Mean Squared Error] For a set of estimated individual treatment effects, $\widehat{\textsc{ITE}}(j)$ for $j \in [n]$, the root mean squared error (RMSE\xspace) is defined as: \[ \textsc{RMSE}\xspace \mathbin{\stackrel{\rm def}{=}} \frac{1}{\sqrt{n}} \cdot \left\lVert\widehat{\textsc{ITE}}(j) - \textsc{ITE}(j)\right\rVert.\]

definition[Leverage Score] Given a matrix ${\mathbf X} \in \mathbb{R}^{n \times d}$, the leverage score of $j^{\text{th}}$ row ${\mathbf X}[j, :]$, denoted by $\ell_j({\mathbf X})$, is defined as: \[ \ell_j({\mathbf X}) \mathbin{\stackrel{\rm def}{=}} {\mathbf X}[j, :]^T {({\mathbf X}^T {\mathbf X})}^{+} {\mathbf X}[j,:],\] where ${}^{+}$ denotes the Moore–Penrose pseudo-inverse.

Individual Treatment Effect Estimation

In this section, we describe our algorithm for ITE estimation. The algorithm identifies a subset of the population to experiment on, using importance based sampling techniques, that are well-studied in randomized numerical linear algebra woodruff2014sketching. Missing proof details in this section are presented in Appendix (ref).

\paragraph{Overview of our approach.} Under the linearity assumption (Assumption (ref)), 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 $\mathbf{y}^0, \mathbf{y}^1$ on ${\mathbf 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 $\mathbf{y}^0_j$ or $\mathbf{y}^1_j$ for any individual in this set.

To tackle the first challenge, we use a sampling technique based on the importance of each row in ${\mathbf X}$, captured via its leverage score (Defn. (ref)). Intuitively, we want to select $s$ individuals (or equivalently rows) that capture the entire row space of ${\mathbf 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 $\mathbf{y}^0$ and $\mathbf{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 $\mathbf{y}^0_j$ or $\mathbf{y}^1_j$. To mitigate this issue, we construct a smoothed matrix ${\mathbf X}^*$, which consists of ${\mathbf 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 significantly to the spectrum of ${\mathbf X}$. Formally, we prove that the maximum leverage score of ${\mathbf 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.

Leverage Score Sampling

For some $\gamma \ge 0$, to be fixed later, we define a smoothed matrix for ${\mathbf X}$, the projection onto singular vectors with high singular values, as follows:

definition[{Smoothed} matrix] Given ${\mathbf X} \in \mathbb{R}^{n \times d}$ with singular value decomposition $ {\mathbf X} = {\mathbf U} {\mathbf \Sigma} {\mathbf V}^T$, let $\Gamma^*$ be the set of indices corresponding to singular values greater than $\sqrt{\gamma}$, i.e., $\Gamma^* \mathbin{\stackrel{\rm def}{=}} \{ i \mid \sigma_i({\mathbf X}) \ge \sqrt{\gamma} \}$; we denote $d' \mathbin{\stackrel{\rm def}{=}} |\Gamma^*|$. Let $ {\mathbf \Sigma}^* = {\mathbf \Sigma}(\Gamma^*,\Gamma^*)$ denote the principal sub-matrix of ${\mathbf \Sigma}$ associated with these large singular values. Similarly, let ${\mathbf U}^* \in \mathbb{R}^{n \times d'}, {\mathbf V}^* \in \mathbb{R}^{d \times d'}$ be the associated column sub-matrices of ${\mathbf U}$ and ${\mathbf V}$. Then, we define: \[ {\mathbf X}^* \mathbin{\stackrel{\rm def}{=}} {\mathbf U}^* {\mathbf \Sigma}^* {\mathbf V}^{*T}.\]

\paragraph{Sampling Matrix.} Our algorithm will sample individuals, corresponding to rows of the {smoothed} matrix of ${\mathbf X}$, i.e., ${\mathbf X}^*$, independently -- the $i^{th}$ row is included in the sample with some probability $\pmb{\pi}_i$. Let the set of rows sampled be denoted by $S$.

We can associate a sampling matrix ${\mathbf W}$ with $S$. The $j^{th}$ row of ${\mathbf W}$ is associated with the $j^{th}$ element in the set $S$ (under some fixed order). If the $j^{th}$ element in $S$ is the row for individual $i$ for some $i \in [n]$, then, ${\mathbf W}[j,:]$ is equal to $\mathbf{e}_i/\sqrt{ \pmb{\pi}_i}$. Here, $\mathbf{e}_i\in \mathbb{R}^n$ denotes the $i^{th}$ standard basis vector. In this way, ${\mathbf W} {\mathbf X}^* $ 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.

algorithm[algorithm omitted — 1,164 chars of source]

{2pt}

\paragraph{Algorithm Sampling-ITE\xspace.} We perform row sampling twice, with probabilities proportional to the leverage scores of ${\mathbf X}^*$, to construct two sets $S^0, S^1$. See the discussion below for the exact definition of the sampling probabilities $\pmb{\pi}_i$, which are proportional to the leverage scores of ${\mathbf X}^*$. These two sets are used to estimate the vectors $\mathbf{y}^0$ and $\mathbf{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 $S^1$ with probability $\pmb{\pi}_j \cdot (1-\pmb{\pi}_j)$ for every $j \in [n]$. We construct sampling matrices ${\mathbf W}^0$ and ${\mathbf W}^1$ using probabilities $\pmb{\pi}$ and $\pmb{\pi} (1-\pmb{\pi})$ respectively. Finally, in Algorithm (ref), we solve the following linear regressions, for $i = 0, 1$ separately: \[ \widetilde{\pmb{\beta}}^i = \arg \min_{\pmb{\beta} \in \mathbb{R}^d} \left\lVert{\mathbf W}^i {\mathbf X}^* \pmb{\beta} - {\mathbf W}^i \mathbf{y}^i\right\rVert^2\] Our estimate for each $\textsc{ITE}(j)$, denoted by $\widehat{\textsc{ITE}}(j)$ is set to $j^{th}$ entry of the vector ${\mathbf X}^*\widetilde{\pmb{\beta}}^1 - {\mathbf X}^*\widetilde{\pmb{\beta}}^0$. Observe that by construction, $S^0 \cap S^1$ is empty. This ensures that we have access to only one of $\mathbf{y}^0_j$ or $\mathbf{y}^1_j$ for any individual $j$ in solving the above two subsampled regression problems.

We note that in Algorithm (ref), 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.

Theoretical Guarantees

First, we bound the error due to sampling. Critically, we show that the leverage scores of ${\mathbf X}^*$, and in turn the probabilities $\pmb{\pi}$, are bounded by $1/\gamma$. Thus, the sampling probabilities for $S^1,$ $\pmb{\pi}(1-\pmb{\pi})$ are not too far from $\pmb{\pi}$ itself.

As we assume the row norms of ${\mathbf X}$ are bounded by $1$, the row norms of ${\mathbf X}^*$ are also bounded. Thus, there can be no rows in ${\mathbf 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 ${\mathbf X}^*$ is at least $\sqrt{\gamma}$. In particular, we prove:

claim$\ell_j({\mathbf X}^*) \le 1/\gamma$, for all $j \in [n]$.

\paragraph{Setting $\pmb{\pi}$.} It is well known that if we sample rows of ${\mathbf X}^*$ with probabilities $\pmb{\pi}$ proportional to the leverage scores, we will obtain a $(1\pm \epsilon)$ relative error approximation for linear regression sarlos2006improved. The result of Sarlos sarlos2006improved applies to sampling $s$ rows with replacement, each equal to $j$ with probability $\pmb{\pi}_j/\left\lVert\pmb{\pi}\right\rVert$. 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[Follows from sarlos2006improved] For ${\mathbf X} \in \mathbb{R}^{n \times d}$, ${\mathbf y} \in \mathbb{R}^{n}$, let $S \subseteq [n]$ include each $j \in [n]$ independently with probability $\pmb{\pi}_j$ satisfying $\pmb{\pi}_j \ge \min\left\{1, \ell_j({\mathbf X}) \cdot c \cdot [\log (\mathrm{rank}({\mathbf X})) + \frac{1}{\delta \epsilon} ]\right\}$ for some large enough constant $c$. Let ${\mathbf W} \in \mathbb{R}^{|S| \times n}$ be a sampling matrix that includes row $\mathbf{e}_j/\sqrt{\pmb{\pi}_j}$ if $j \in S$, where $\mathbf{e}_j \in \mathbb{R}^n$ is the $j^{th}$ standard basis vector. Let $\widetilde{\pmb{\beta}} = \arg \min_{\pmb{\beta} \in \mathbb{R}^d} \left\lVert{\mathbf W} {\mathbf X} \pmb{\beta} - {\mathbf W} {\mathbf y}\right\rVert^2.$ Then, $\mathbb{E}[|S|] = \sum_{j=1}^n \pmb{\pi}_j$ and with probability $\ge 1-\delta$: \[\left\lVert{\mathbf X} \widetilde{\pmb{\beta}} - {\mathbf y} \right\rVert \le (1+\epsilon) \cdot \min_{\pmb{\beta}} \left\lVert{\mathbf X} {\pmb{\beta}} - {\mathbf y}\right\rVert.\] If the $\pmb{\pi}_j$'s are within constants of the required bound, $\mathbb{E}[|S|] = O \left (d \log d + \frac{d}{\epsilon \delta} \right )$.

Note that the bound on $\mathbb{E}[|S|]$ follows from the well known fact that the sum of leverage scores, is equal to the rank, i.e., $\sum_{j=1}^n \ell_j({\mathbf X}) = \operatorname{rank}({\mathbf X}) \le d$ woodruff2014sketching.

The sampling probabilities are set to $\pmb{\pi}_j = \min \left\{1, \ell_j({\mathbf X}^*) \cdot c_0 \cdot [\log (\mathrm{rank}({\mathbf X}^*)) + {30}/{\epsilon} ] \right\}$ for some constant $c_0 \ge 2c$, where $c$ is the constant in Lemma (ref). Thus, by the lemma, we will have, with probability $\ge 29/30$, $\left\lVert{\mathbf X}^* \widetilde \pmb{\beta}^0 - {\mathbf y}^0\right\rVert \le (1+\epsilon) \left\lVert{\mathbf X}^* \pmb{\beta}^0 - {\mathbf y}^0\right\rVert.$ 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 $\pmb{\pi}_j \cdot (1-\pmb{\pi}_j)$. If we can prove that $\pmb{\pi}_j \cdot (1-\pmb{\pi}_j) \ge \frac{\pmb{\pi}_j}{2}$, then Lemma (ref) will still apply, since we have set $c_0 = 2c$. To do so, it suffices to argue that $\pmb{\pi}_j \le 1/2$ by setting the parameters appropriately.

claimIf $\gamma = \small{4 c_0 \max\left\{\log (\mathrm{rank}({\mathbf X}^*)),30/\epsilon \right\} }$ and $\pmb{\pi}_j = \min \{1, \ell_j({\mathbf X}^*) \cdot c_0 \cdot [\log (\mathrm{rank}({\mathbf X}^*)) + 30/\epsilon ] \}$, we have $\pmb{\pi}_j \le 1/2$ for every $j \in [n]$.

proof\begin{align*} \pmb{\pi}_j \le \ell_j({\mathbf X}^*) \cdot c_0 \cdot [\log (\mathrm{rank}({\mathbf X}^*)) + 30/\epsilon ] &\le 1/\gamma \cdot c_0 \cdot [\log (\mathrm{rank}({\mathbf X}^*)) + 30/\epsilon ] (Claim (ref))\\ &\le \frac{c_0 [\log (\mathrm{rank}({\mathbf X}^*)) + 30/\epsilon ]}{4 c_0 \max\left\{\log (\mathrm{rank}({\mathbf X}^*)),30/\epsilon \right\}} \le \frac{1}{2}. \end{align*}

In Appendix (ref), we argue that using the smoothed matrix ${\mathbf X}^*$ introduces an error of $\sqrt{ \gamma}$. Combining all of them, we have the following corollary:

corollarySuppose $\gamma$ and $\pmb{\pi}_j$ are set as in Claim (ref), for some sufficiently large constant $c_0$. Then, Algorithm Sampling-ITE\xspace satisfies, for $i=0,1$, with probability at least $14/15$: \[\left\lVert{\mathbf X}^* \widetilde{\pmb{\beta}}^i - \mathbf{y}^i\right\rVert \le \left(1+ \epsilon\right) \cdot \left(\sqrt{\gamma}\left\lVert\pmb{\beta}^i\right\rVert + \left\lVert\pmb{\zeta}^i\right\rVert \right) \mbox{ for $i = 0, 1$}.\]

\paragraph{RMSE Guarantees.} The root mean squared error (Defn. (ref)) for the ITE estimates is given by: $$\textsc{RMSE}\xspace= \frac{1}{\sqrt n} \left\lVert({\mathbf X}^* \widetilde{\pmb{\beta}}^1 - {\mathbf X}^* \widetilde{\pmb{\beta}}^0) - (\mathbf{y}^1 - \mathbf{y}^0) \right\rVert.$$

By setting $\epsilon = {120 c_0 d\log d/s}$ in Corollary (ref), we get the following theorem for our Algorithm (ref):

theoremSuppose $s \ge 120 c_0 d \log d$. There is a randomized algorithm that selects a subset $S \subseteq [n]$ of the population with $\operatorname*{{\bf {E}}}[|S|] \le s$, and, with probability at least $9/10$, returns ITE estimates $\widehat{\textsc{ITE}}(j)$ for all $j \in [n]$ with error: \begin{equation*} RMSE\xspace = O\Bigl(\sqrt{\frac{1}{n} \max\left\{{\frac{s}{d}}, \log d\right\}} \cdot (\left\lVert\pmb{\beta}^1\right\rVert + \left\lVert\pmb{\beta}^0\right\rVert ) + \sigma \Bigl). \end{equation*}

For the sake of simplicity of analysis, we used a constant success probability in Theorem (ref). All our claims can easily be updated with a general failure probability of $\delta$, with a dependence of $\sqrt{1/\delta}$, using Lemma (ref). The corollary below follows immediately from Theorem (ref).

corollary[Main ITE Error Bound] The root mean squared error obtained by Algorithm (ref) is minimized when $s = \Theta(d\log d)$ and is given by: $$\textsc{RMSE}\xspace = O\Bigl(\sqrt{\frac{\log d}{n}} \cdot (\left\lVert\pmb{\beta}^1\right\rVert + \left\lVert\pmb{\beta}^0\right\rVert ) + \sigma \Bigl).$$

Our upper bound on RMSE\xspace for Algorithm (ref) increases with $s$, if $s$ grows strictly faster than $d\log d$ asymptotically, i.e., $s = \omega(d\log d)$. Therefore, to obtain low error, we set $s = c \cdot 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 (ref), we observe empirically that the error decreases with increasing $s$.

\paragraph{Remark.} We observe that the RMSE bound in Corollary (ref) is nearly optimal, even for algorithms that experiment on the full population. The $O(\sigma)$ term cannot be improved by more than constants, as a consequence of our noise model (see Assumption (ref)). Even if we knew the true $\pmb{\beta}^1$ and $\pmb{\beta}^0$, our RMSE\xspace would be $O(\sigma)$.

The term $(\left\lVert\pmb{\beta}^0\right\rVert+\left\lVert\pmb{\beta}^1\right\rVert)/\sqrt{n}$ is also necessary. Suppose the matrix ${\mathbf 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 $\pmb{\beta}^1$ and $\pmb{\beta}^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$, $\textsc{ITE}(j) = 0$ and with probability $1/2$, $\textsc{ITE}(j) = \pm 1$. No algorithm which observes just one of $\mathbf{y}^1_j$ or $\mathbf{y}^0_j$ can obtain expected error o(1) in estimating $\textsc{ITE}(j)$. That is, no algorithm can obtain RMSE\xspace $o(1/\sqrt{n}) = o\left ((\left\lVert\pmb{\beta}^0\right\rVert+\left\lVert\pmb{\beta}^1\right\rVert)/\sqrt{n}\right )$.

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 efficient experimental design by Harshaw et al. harshaw2019balancing. Missing details from this section are collected in Appendix (ref).

\paragraph{Horvitz-Thompson Estimator.} Suppose ${\mathbf S}^+ \subseteq [n]$ is the population assigned to the treatment group and ${\mathbf S}^- = [n] \setminus {\mathbf 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 horvitz1952generalization, denoted by $\widehat \tau$. If every individual is assigned to ${\mathbf S}^+$ (or ${\mathbf S}^-$) with probability $0.5$, then, $\widehat \tau$ is defined as follows: \[\widehat \tau = \frac{2}{n} \left( \sum_{i \in {\mathbf S}^+} \mathbf{y}^1_i - \sum_{i \in {\mathbf S}^-} \mathbf{y}^0_i \right).\]

algorithm[algorithm omitted — 1,171 chars of source]

Harshaw et al. harshaw2019balancing present an experimental design based on the Gram-Schmidt-Walk algorithm for discrepancy minimization bansal2018gram. 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[Proposition 3 in harshaw2019balancing] For all $\Delta > 0$, with probability at least $1-2\exp{\left(- \frac{\Delta^2 n}{8L}\right)}$, the Gram-Schmidt-Walk design satisfies: $|\widehat \tau - \tau| \le \Delta$, where $L = \frac{2}{n} \min_{\pmb{\beta} \in \mathbb{R}^d}\left(\left\lVert\frac{\mathbf{y}^1 + \mathbf{y}^0}{2} - {\mathbf X} \pmb{\beta}\right\rVert^2 + \left\lVert\pmb{\beta}\right\rVert^2 \right)$.

\paragraph{Overview of Recursive-Covariate-Balancing\xspace.} Our main idea in Algorithm (ref) 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 ${\mathbf Z}_t$ into treatment and control groups, denoted by ${\mathbf Z}^+_t, {\mathbf Z}^-_t$ using GSW. Next, we identify the smaller of these two subsets, say ${\mathbf Z}^+_t$ and recurse on ${\mathbf Z}^+_t$. We stop after $k$ recursive calls when there are only $s$ individuals to experiment on, i.e., $|{\mathbf Z}^+_k \cup {\mathbf Z}^-_k| \le s$. Finally, we construct our estimator $\widehat \tau_s$, similar to the Horvitz-Thompson estimator, by scaling the treatment and control contributions due to ${\mathbf Z}^+_k$ and ${\mathbf 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 karnin2019discrepancy, is based on the observation that if we can obtain good estimates for the contributions $\sum_{i \in [n]} \mathbf{y}^1_i$ and $\sum_{i \in [n]} \mathbf{y}^0_i$, we obtain a good estimate for ATE ($\tau$). Using the next lemma, we argue that after a call to GSW algorithm that partitions $[n]$ into the sets ${\mathbf S}^+$ and ${\mathbf S}^-$, we can obtain additive approximations of $\sum_{i \in [n]} \mathbf{y}^1_i$ and $\sum_{i \in [n]} \mathbf{y}^0_i$. Our approximations are the contributions of treatment and control values in ${\mathbf S}^+$ and ${\mathbf S}^-$ scaled appropriately, i.e., $\sum_{i \in {\mathbf S}^+} 2 \cdot \mathbf{y}^1_i$ and $\sum_{i \in {\mathbf S}^-} 2 \cdot \mathbf{y}^0_i$.

lemmaSuppose the Gram-Schmidt-Walk design harshaw2019balancing partitions the population $[n]$ into two disjoint groups ${\mathbf S}^+$ and ${\mathbf S}^-$. Under the linearity assumption, with probability $1-1/3\log(n/s)$, for both the control and treatment groups, the following holds: \[\left|\sum_{j \in {\mathbf S}^+}2 \mathbf{y}^i_j - \sum_{j \in [n] } \mathbf{y}^i_j \right| \le 4 \sqrt{\log (16 \log (n/s))} \cdot \left(2\sigma \sqrt{n} + \left\lVert\pmb{\beta}^i\right\rVert \right) \quad \mbox{ for $i = 0, 1$}. \]

Building upon the previous lemma, we argue in Theorem (ref) that the additive approximation errors obtained from repeated use of GSW in our algorithm Recursive-Covariate-Balancing\xspace result in a low estimation error.

theorem[Main ATE Error Bound] The estimator $\widehat \tau_s$ in Algorithm Recursive-Covariate-Balancing\xspace obtains the following guarantee, with probability at least $2/3$: \[ \left|\widehat \tau_s - \tau \right| = O\left( \sqrt{\log \log (n/s)} \cdot \left(\frac{\sigma}{\sqrt{s}} + \frac{\left\lVert\pmb{\beta}^1\right\rVert + \left\lVert\pmb{\beta}^0\right\rVert}{s} \right)\right). \]

\paragraph{Remark.} When $s = n$, the above theorem matches the guarantees obtained by GSW design described in Lemma (ref). Moreover, we obtain a better dependence compared to sampling $s$ rows uniformly at random and using the $\mathbf{y}^1, \mathbf{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 $\widetilde{O}\left(1/s \cdot {\left\lVert{\mathbf X}\right\rVert_2 (\left\lVert\pmb{\beta}^1\right\rVert + \left\lVert\pmb{\beta}^0\right\rVert)}\right)$, instead of the $\widetilde{O}\left(1/s \cdot ({\left\lVert\pmb{\beta}^1\right\rVert + \left\lVert\pmb{\beta}^0\right\rVert})\right)$ obtained by our algorithm, where $\left\lVert{\mathbf X}\right\rVert_2$ denotes the spectral norm of ${\mathbf X}$ and $\widetilde{O}(\cdot)$ 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 (ref). Our code is publicly accessible using the following \href{https://github.com/raddanki/Sample-Constrained-Treatment-Effect-Estimation}{github repository}.

\paragraph{Data Generation.} We evaluate our approaches on five 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 hill2011bayesian, dorie2016npci. (ii) Twins. This contains data regarding the mortality rate in twin births in the USA between 1989-1991 almond2005costs. Following the work of louizos2017causal, 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 lalonde1986evaluating, 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 harrison1978hedonic. 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 ${\mathbf X} \in \mathbb{R}^{2000 \times 25}$, using an approach due to ma2014statistical. There is a high disparity in leverage score values in ${\mathbf X}$, similar to what we observe in other datasets. Using a random linear function on ${\mathbf 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. shalit2017estimating and Louizos et al. louizos2017causal. 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 fixed shift in treatment effect (i.e., the true treatment effect is equal to a constant), similar to Arbour et al. pmlr-v130-arbour21a.

\paragraph{Baselines.}

enumerate• ATE. We compare the performance of our Algorithm Recursive-Covariate-Balancing\xspace (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$. • \textbf{ITE.} We compare the performance of our Algorithm \textsc{Sampling-ITE}\xspace (referred to as \textit{`Leverage'}) with respect to three baselines: \emph{ (i) Uniform}. We run Algorithm (ref) on ${\mathbf X}$ and uniform sampling distribution given by $\pmb{\pi}_j = s/n \ \forall j$. \emph{(ii) Leverage-nothresh}. We run Algorithm (ref) on ${\mathbf X}$, instead of ${\mathbf X}^*$ with the probability distribution $\pmb{\pi}_j \propto \ell_j({\mathbf X}) \forall j$. \emph{(iii) Lin-regression}. This captures the best linear fit regression error, i.e., assuming we have access to both $\mathbf{y}^1, \mathbf{y}^0$, we regress these vectors on ${\mathbf X}$ to obtain $\pmb{\beta}^1, \pmb{\beta}^0$, and use the resultant ITE estimates ${\mathbf X} \pmb{\beta}^1 - {\mathbf X} \pmb{\beta}^0$.

figure[figure omitted — 1,172 chars of source]

\paragraph{Evaluation.} To evaluate the performance of average treatment effect estimation ($\tau$) on the datasets, we compare the deviation error of the estimator $\widehat \tau_s$, given by $|\widehat \tau_s - \tau|$ for different sample sizes. To evaluate the performance of individual treatment effect estimates, we compare the root mean squared error RMSE (see Defn. (ref)) for different sample sizes.

\paragraph{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 confidence interval as shown in Figures (ref), (ref) representing ATE and ITE results respectively.

enumerate• ATE. For all datasets, we observe that the deviation error obtained by our algorithm labeled as Recursive-GSW in Figure (ref), is significantly smaller than that of Uniform baseline. Surprisingly, for the IHDP dataset, our approach is significantly 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 {confidence 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 sufficient 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 {confidence interval} of the error obtained by \emph{GSW-pop}. Therefore, for a specified error tolerance level for ATE, we can reduce the associated experimental costs using just a small subset of the dataset using our algorithm. • \textbf{ITE}. For all sample sizes, we observe that the \textsc{RMSE}\xspace obtained by our algorithm labeled as \emph{Leverage} in Figure (ref), is significantly smaller than that of all the other baselines, including \emph{Uniform} and \emph{Leverage-nothresh}. E.g., we observe that when the sample size is $20\%$ of the population in {IHDP} dataset, the error obtained by \emph{Leverage} is at least $50\%$ times smaller than that of \emph{Uniform} and \emph{Leverage-nothresh}. For the Synthetic dataset, the error obtained by \emph{Leverage} is extremely close to that of the error due to the best linear fit, \emph{Lin-regression} (see the zoomed in part of the figure). 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 efficient 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 ugander2013graph. 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 \le n$. Moreover, the $\sqrt{\log d}$ factor in Corollary (ref) likely can be improved, using recent work which improves standard leverage score sampling bounds by a $\log d$ factor chen2019active, yielding a bound that is optimal up to constants. For ATE estimation, the Horvitz-Thompson estimator can include sampling probabilities different from $0.5$. 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-1908849, and CCF-1637536, along with Adobe and Google Research Grants.