SlowGuess/NeurIPS_2023_OCR huggingface dataset:
ITE Estimation with Partial Observation. Our next step is to extend the above to the partial observation setting. To do so, we apply leverage score sampling, which is a standard technique for approximate regression via subsampling [10, 26, 35]. It is well known that if we sample  $O(d\log (d / \delta) / \epsilon^2)$  rows of  $\mathbf{X}$  and the corresponding entries of  $\mathbf{v}$  according to the leverage scores of  $\mathbf{X}$ , then we can solve a reweighted regression problem on just these rows to find  $\widehat{\mathbf{b}}$  satisfying  $\left\| \mathbf{X}\widehat{\mathbf{b}} -\mathbf{v}\right\| _2^2\leq (1 + \epsilon)\cdot \min_{\mathbf{b}\in \mathbb{R}^d}\left\| \mathbf{X}\mathbf{b} - \mathbf{v}\right\| _2^2$ . Observe that to compute  $\widehat{\mathbf{b}}$ , we only need to observe potential outcomes for the units (i.e., the entries of  $\mathbf{v}$ ) that are sampled. Thus, this almost allows us to recover a similar bound to (1), up to a multiplicative  $(1 + \epsilon)$  factor, in the partial observation setting. However, the leverage score of any one row may be very small, and if sampled, that row may be scaled up by a large factor in our reweighted regression problem (to preserve expectation). It thus becomes difficult to control to error introduced by approximating  $\mathbf{t}$  by  $\mathbf{v}$ .
To handle this issue, we mix leverage score sampling, with uniform sampling at rate  $\alpha$ . This increases our sample complexity to  $O(d\log (d / \delta) / \epsilon^2 +\alpha n)$ , but allows us to bound the error due to approximating  $\mathbf{t}$  by  $\mathbf{v}$  by  $\frac{(1 + \epsilon)d}{\alpha}\| \boldsymbol {\mu}\| _2^2$ . This yields the final sample complexity and error bound of Theorem 1. For a full proof, see Appendix B.
2.2 ATE with Full Observation
We build on our techniques for ITE estimation to tackle the ATE estimation problem. We start with the classic Horvitz-Thompson (HT) estimator for ATE, defined below.
Definition 4 (Horvitz-Thompson estimator). Given two outcome vectors  $\mathbf{y}^{(1)},\mathbf{y}^{(0)}\in \mathbb{R}^n$  , let  $Z^{+}$  and  $Z^{-} = [n]\setminus Z^{+}$  be a random partitioning of the units to two groups under a distribution  $\mathcal{P}$  . Then the Horvitz-Thompson (HT) estimator is  $\widehat{\tau} := \frac{1}{n}\bigl (\sum_{i\in Z^{+}}\frac{y_{i}^{(1)}}{\mathbb{P}_{\mathcal{P}}[i\in Z^{+}]}- \sum_{i\in Z^{-}}\frac{y_{i}^{(0)}}{\mathbb{P}_{\mathcal{P}}[i\in Z^{-}]}\bigr)$  where  $\mathbb{P}_{\mathcal{P}}[i\in Z^{+}]$  denotes the probability of placing i in  $Z^{+}$  when the partitioning is performed according to distribution  $\mathcal{P}$
Note that if units are assigned independently with equal probability to treatment and control, the HT estimator is equivalent to outputting the average of the random vector  $\mathbf{v}$  used in our regression-based estimate for ITE. In this case, the HT estimator is unbiased and well known to have variance bounded by  $\frac{1}{n^2} \| \boldsymbol{\mu} \|_2^2$ , where  $\boldsymbol{\mu} \coloneqq \mathbf{y}^{(1)} + \mathbf{y}^{(0)}$ . Our goal is to improve this variance to instead depend on the error of the best linear fit to  $\boldsymbol{\mu}$ ,  $\min_{\mathbf{b} \in \mathbb{R}^d} \| \mathbf{X}\mathbf{b} - \boldsymbol{\mu} \|_2^2$ , as in Theorem 2. To do so, we will employ the classic technique of regression adjustment.
Horvitz-Thompson on Regression Adjusted Outcomes. Suppose we were given a vector  $\mathbf{b} \in \mathbb{R}^d$  determined independently of the randomness in the HT estimator. Then applying the HT estimator directly on the regression adjusted vectors  $\widetilde{\mathbf{y}}^{(i)} \coloneqq \mathbf{y}^{(i)} - \mathbf{X}\mathbf{b}$  (which have the same ATE as the original vectors), would yield an unbiased estimate for the ATE with variance at most  $\frac{1}{n^2} \| \mathbf{X}\mathbf{b} - \boldsymbol{\mu} \|_2^2$ . Naturally, we would like to find such a  $\mathbf{b}$  making this variance bound as small as possible.
To do so, we will apply the same approach as in ITE estimation. We cannot directly regress  $\pmb{\mu}$  onto our covariates to find an optimal  $\mathbf{b}$ , as we cannot observe any entry of  $\pmb{\mu}$  since it is the sum of the two potential outcomes. Instead, we can construct a vector  $\mathbf{u}$  with  $\mathbb{E}[\mathbf{u}] = \pmb{\mu}$ , and use it as a surrogate in our regression problem. In particular, we let  $u_{i} = 2y_{i}^{(1)}$  with probability  $1/2$  and  $u_{i} = 2y_{i}^{(0)}$  with probability  $1/2$ . As in the ITE setting, we can bound the error incurred by approximating  $\pmb{\mu}$  by  $\mathbf{u}$  in our regression problem, and thus bound the variance of our regression adjusted HT estimator. We note that to the best of our knowledge, this idea has not been leveraged in ATE estimation.
Unfortunately, a difficulty arises in this approach. We need to assign each unit to a treatment or control group both when computing  $\mathbf{u}$  to solve our regression problem, and later when applying the HT estimator. If we use the same assignment for both steps (as is required since we can only observe one of  $y_{i}^{(1)}$  or  $y_{i}^{(0)}$  for each unit), we introduce bias in our estimator.
Avoiding Estimator Bias. To alleviate this issue, we use ridge leverage score sampling to first partition the units into two groups. Ridge leverage score sampling is a technique similar to leverage score sampling that gives similar approximation guarantees to linear regression but for ridge regression problems (i.e., linear regression problems with an  $\ell_2$  regularization term). We pick a ridge (regularization) parameter that guarantees that each unit only needs to be in the first or second group with a probability of at most 0.5 while guaranteeing that the solution obtained from this sampling is a  $(1 + \epsilon)$  approximation of the optimal ridge regression solution. Then we solve ridge regression problems on the groups separately to obtain two vectors  $\widehat{\mathbf{b}}^{(1)}$  and  $\widehat{\mathbf{b}}^{(2)}$ , for the first and the second group, respectively. We then use  $\widehat{\mathbf{b}}^{(1)}$  to adjust the outcomes of the second group and use  $\widehat{\mathbf{b}}^{(2)}$  to adjust the outcomes of the first group. Since the vector of each group is not used to adjust the outcomes of itself, this gives an unbiased estimator.
2.3 ATE with Partial Observation
We now consider ATE estimation where we desire to only observe  $y_{i}^{(1)}$  or  $y_{i}^{(0)}$  for a small subset of the population. If we draw a subset of size  $m$  uniformly at random and observe the outcomes of the samples in the subset independently and with equal probability and apply the Horvitz-Thompson estimator, the variance is  $\mathbb{E}\left[(\widehat{\tau} -\tau)^2\right]\leq \frac{\|\pmb{\mu}\|_2^2}{mn} +\frac{\|\mathbf{t} - \overline{\mathbf{t}}\|_2^2}{m}$ . We would like to replace terms  $\| \pmb {\mu}\| _2^2$  and  $\| \mathbf{t} - \overline{\mathbf{t}}\| _2^2$  with  $\min_{\mathbf{b}}\| \pmb {\mu} - \mathbf{X}\mathbf{b}\| ^2$  and  $\min_{\mathbf{b}}\| \mathbf{t} - \overline{\mathbf{t}} -(\mathbf{X} - \overline{\mathbf{X}})\mathbf{b}\| ^2$ , respectively. To deal with the  $\| \pmb {\mu}\| _2^2$  term, we use the HT estimator applied to the GSW design [18] on the uniformly selected subset of size  $m$  of the data. Then for  $\phi \in (0,1)$  that the experiment designer selects, the variance reduces to  $\frac{1}{m^2}\min_{b\in \mathbb{R}^d}\left[\frac{m}{n}\frac{1}{\phi}\| \mathbf{X}\mathbf{b} - \pmb {\mu}\| _2^2 +\frac{\zeta^2}{1 - \phi}\| \mathbf{b}\| _2^2\right] + \frac{\|t - \overline{t}\|_2^2}{mn}$ . To replace the  $\frac{\|\mathbf{t} - \overline{\mathbf{t}}\|_2^2}{mn}$  term with the error
of the best linear fit, we use the partial observation RAHT estimator that given a solution vector  $\widetilde{\mathbf{b}}$  adjusts the estimate obtained by the GSW design on the sample of size  $m$ . Similar to our approach for partial observation ITE estimation, we compute  $\widetilde{\mathbf{b}}$  by sampling about  $d$  units according to leverage scores of the matrix and performing the random vector regression (pseudo-outcome regression) on a vector that is equal to  $\mathbf{t} / 2$  in expectation. Let  $S$  be the set of units selected by leverage score sampling and  $\overline{S} = [n] \setminus S$ . We denote the estimate obtained by sampling  $m$  units from  $\overline{S}$  (we denote this set with  $T$ ) and applying GSW design with  $\widehat{\tau}_{\overline{S}}$ . Then, the RAHT estimator estimates the ATE as  $\widehat{\tau}_{\overline{S}} - \frac{1}{m} \sum_{i \in T} (\mathbf{x}_i - \overline{\mathbf{x}}^{\overline{S}}) \widetilde{\mathbf{b}}$ . This estimate is biased because some units have a higher probability of being included in  $S$  and, therefore, would contribute to the estimate with a lower probability. To resolve this issue, we also estimate the ATE on the set  $S$  with the same random vector of outcomes that is used in regression to learn  $\widetilde{\mathbf{b}}$ . We use HT estimator for this estimate, and we denote it by  $\widehat{\tau}_{S}$ . Then, our final estimate is a convex combination of the two estimates as the following.
$$
\frac {| \overline {{S}} |}{n} \cdot \left(\widehat {\tau} _ {\overline {{S}}} - \frac {1}{m} \sum_ {i \in T} (\mathbf {x} _ {i} - \overline {{\mathbf {x}}} ^ {\overline {{S}}}) \widetilde {\mathbf {b}}\right) + \frac {| S |}{n} \cdot \widehat {\tau} _ {S}.
$$
This estimator is unbiased and achieves the variance bounds stated in Theorem 3.
Notation. Vectors and matrices are denoted with bold lowercase and capital letters, respectively. the Hadamard (entrywise) product of  $\mathbf{x},\mathbf{y}$  is denoted by  $\mathbf{x}\odot \mathbf{y}$ . We denote the vectors of all ones and all zeros in  $\mathbb{R}^n$  by  $\mathbf{1}_n$  and  $\mathbf{0}_n$ , respectively. We drop the subscript if the dimension is clear. We denote the identity matrix of dimension  $n$  by  $\mathbf{I}_n$ . For  $\mathbf{X}\in \mathbb{R}^{n\times d}$  we denote the projection matrix of  $\mathbf{X}$  by  $\pi_{\mathbf{X}}\coloneqq \mathbf{X}(\mathbf{X}^{\top}\mathbf{X})^{+}\mathbf{X}$ , where  $+$  denotes the pseudoinverse. For  $\mathbf{x}\in \mathbb{R}^n$  and  $i\leq n$ ,  $x_{i}$  denotes its  $i^{th}$  entry and  $\mathbf{x}_{1:i}$  denotes a vector in  $\mathbb{R}^i$  that is equal to its first  $i$  entries. We use  $\mathbf{X}_{i,:}$  and  $\mathbf{x}_i$  interchangeably to denote row  $i$  of  $\mathbf{X}$ . Similarly for  $S\subseteq [n]$ ,  $\mathbf{X}_{S,:}\in \mathbb{R}^{|S|\times d}$  denotes the submatrix with row indices in  $S$  and for  $\mathbf{y}\in \mathbb{R}^n$ ,  $\mathbf{y}_S\in \mathbb{R}^{|S|}$  denotes the vector restricted to indices in  $S$ .
Definition 5 (Ridge leverage scores [2]). Let  $\mathbf{X} \in \mathbb{R}^{n \times d}$ ,  $\lambda \geq 0$ , and  $\mathbf{x}_i \in \mathbb{R}^d$  be the  $i^{th}$  row of  $\mathbf{X}$ . Then the  $\lambda$ -ridge leverage score of row  $i$  of  $\mathbf{X}$  is defined as  $\ell_i^{(\lambda)} := \mathbf{x}_i^\top (\mathbf{X}^\top \mathbf{X} + \lambda \mathbf{I})^{-1} \mathbf{x}_i$ . We denote the leverage scores (i.e.,  $\lambda = 0$ ) with  $\ell_i$ . If the corresponding matrix  $\mathbf{X}$  is not clear from the context, we denote the  $\lambda$ -ridge leverage scores with  $\ell_i^{(\lambda)}(\mathbf{X})$ .
Note that  $\ell_{i}$  is the  $i$ 'th diagonal entry of the projection matrix  $\pi_{\mathbf{X}} = \mathbf{X}(\mathbf{X}^{\top}\mathbf{X})^{-1}\mathbf{X}^{\top}$ . Thus, since  $\operatorname{tr}(\pi_{\mathbf{X}}) = \operatorname{rank}(\mathbf{X})$ ,  $\sum_{i=1}^{n} \ell_{i} = \operatorname{rank}(\mathbf{X})$ . Since in our approach, we need to sample different sets of population for different purposes, we require a bound on the ridge leverage scores that can be achieved by taking the regularization factor  $\lambda$  to be large enough.
Theorem 6. Let  $\zeta \coloneqq \max_{i\in [n]}\| \mathbf{x}_i\| _2$ . Then for  $\lambda \geq c\cdot \zeta^2$ , for  $c\geq 1$ ,  $\ell_i^{(\lambda)}(\mathbf{X})\leq \frac{1}{c}$  for all  $i\in [n]$ .
The following theorem provides a sampling procedure based on  $\lambda$ -ridge leverage scores to solve a linear ridge regression problem approximately.
Theorem 7 (6, 12). Let  $\epsilon, \delta \in (0,1]$ , and  $\mathbf{u} \in \mathbb{R}^n$  be a vector that upper bounds the  $\lambda$ -ridge leverage scores of a matrix  $\mathbf{X} \in \mathbb{R}^{n \times d}$ , i.e.,  $\ell_i^{(\lambda)} \leq u_i$ . Let  $\mathbf{S} = \mathbb{R}^{n \times n}$  be a random diagonal matrix in which  $\mathbf{S}_{ii} = 1 / \sqrt{p_i}$  with probability  $p_i = \min \{1,3 \cdot u_i \cdot \log(d / \delta) / \epsilon\}$ , and  $\mathbf{S}_{ii} = 0$ , otherwise. Then for any  $\mathbf{y} \in \mathbb{R}^n$ , letting  $\widehat{\mathbf{b}} = \arg \min_{\mathbf{b} \in \mathbb{R}^d} \left[ \| \mathbf{S}\mathbf{X}\mathbf{b} - \mathbf{S}\mathbf{y} \|_2^2 + \lambda \| \mathbf{b} \|_2^2 \right]$ , with probably at least  $1 - \delta$ ,
$$
\left\| \mathbf {X} \widehat {\mathbf {b}} - \mathbf {y} \right\| _ {2} ^ {2} + \lambda \left\| \widehat {\mathbf {b}} \right\| _ {2} ^ {2} \leq (1 + \epsilon) \min  _ {\mathbf {b} \in \mathbb {R} ^ {d}} \left[ \left\| \mathbf {X} \mathbf {b} - \mathbf {y} \right\| _ {2} ^ {2} + \lambda \left\| \mathbf {b} \right\| _ {2} ^ {2} \right].
$$
4 Average Treatment Effect Estimation
In this section, we first analyze the random vector regression (i.e., pseudo-outcome regression) approach which is used both for our ITE and ATE estimation. We then discuss our full observation ATE estimation using this approach and provide a sketch of the proof of Theorem 2.
4.1 Random Vector Regression
For both ATE and ITE estimation, regression adjustment allows us to reduce the variance of estimation to the error of the best linear fit on vectors  $\mu$  and  $\mathbf{t}$ , respectively. Since we do not have access to the
entries of these vectors, we cannot compute a solution vector directly. Therefore, for ATE (or ITE), our approach is to instead compute a solution vector  $\mathbf{b}$  by performing regression on a random vector that, in expectation, is equal to  $\mu$  (or  $\mathbf{t}$ ). The following theorem characterizes the error of the solution vector obtained in this way.
The following is a more general result that combines this random assignment technique with leverage score sampling to allow us to observe only a small subset of the population. Algorithm 3 and the proof of the following theorem are presented in the appendix.
Theorem 8 (Random vector regression). Let  $\mathbf{y}^{(0)},\mathbf{y}^{(1)}\in \mathbb{R}^n$  and  $\mathbf{y}\in \mathbb{R}^n$  be a random vector such that for each  $i\in [n]$ ,  $y_{i}$  is independently and with equal probability is either equal to  $y_{i}^{(0)}$  or  $y_{i}^{(1)}$ . Moreover, let  $\mathbf{b}^{*} = \arg \min_{\mathbf{b}}\| \mathbf{X}\mathbf{b} - \mathbf{y}\| _2^2$ . Let  $\pmb {\mu}\coloneqq \mathbf{y}^{(1)} + \mathbf{y}^{(0)}$ . Then
$$
\mathbb {E} \left[ \| 2 \mathbf {X b} ^ {*} - \boldsymbol {\mu} \| _ {2} ^ {2} \right] \leq d \| \mathbf {y} ^ {(1)} - \mathbf {y} ^ {(0)} \| _ {\infty} ^ {2} + \min  _ {\mathbf {b}} \| \mathbf {X b} - \boldsymbol {\mu} \| _ {2} ^ {2}.
$$
Proof. Let  $\mathbf{z} \in \{-1, +1\}^n$ , where  $z_i = +1$  if  $y_i = y_i^{(1)}$ , and  $z_i = -1$ , otherwise. First, note that  $2\mathbf{y} = \boldsymbol{\mu} + \mathbf{z} \odot \mathbf{t}$ , where  $\boldsymbol{\mu} = \mathbf{y}^{(1)} + \mathbf{y}^{(0)}$ , and  $\mathbf{t} = \mathbf{y}^{(1)} - \mathbf{y}^{(0)}$ . Therefore
$$
\boldsymbol {\mu} - 2 \mathbf {X b} ^ {*} = \boldsymbol {\mu} - 2 \pi_ {\mathbf {X}} \mathbf {y} = \boldsymbol {\mu} - \pi_ {\mathbf {X}} (\boldsymbol {\mu} + \mathbf {z} \odot \mathbf {t}) = (\boldsymbol {\mu} - \pi_ {\mathbf {X}} \boldsymbol {\mu}) - \pi_ {\mathbf {X}} (\mathbf {z} \odot \mathbf {t}).
$$
Moreover since  $\mathbf{I} - \pi_{\mathbf{X}}$  and  $\pi_{\mathbf{X}}$  are orthogonal to each other,
$$
\left\| \boldsymbol {\mu} - 2 \mathbf {X b} ^ {*} \right\| _ {2} ^ {2} = \left\| (\boldsymbol {\mu} - \pi_ {\mathbf {X}} \boldsymbol {\mu}) \right\| _ {2} ^ {2} + \left\| \pi_ {\mathbf {X}} (\mathbf {z} \odot \mathbf {t}) \right\| _ {2} ^ {2}.
$$
Now note that  $\| (\pmb {\mu} - \pi_{\mathbf{X}}\pmb {\mu})\| _2^2 = \min_{\mathbf{b}}\| \mathbf{X}\mathbf{b} - \pmb {\mu}\| _2^2$  . So we only need to bound  $\mathbb{E}\left[\| \pi_{\mathbf{X}}(\mathbf{z}\odot \mathbf{t})\| _2^2\right]$  We have  $\| \pi_{\mathbf{X}}(\mathbf{z}\odot \mathbf{t})\| _2^2 = (\mathbf{z}\odot \mathbf{t})^\top \pi_{\mathbf{X}}(\mathbf{z}\odot \mathbf{t}) = \mathbf{z}^\top \mathbf{T}\pi_{\mathbf{X}}\mathbf{T}\mathbf{z},$  where  $\mathbf{T}$  is the diagonal matrix associated with the vector t. Since  $z_{i}$  and  $\mathbf{z}_j$  are independent for  $i\neq j$  , we have
$$
\mathbb {E} \left[ \mathbf {z} ^ {\top} \mathbf {T} \pi_ {\mathbf {X}} \mathbf {T} \mathbf {z} \right] = \sum_ {i = 1} ^ {n} t _ {i} ^ {2} (\pi_ {\mathbf {X}}) _ {i i} \leq \| \mathbf {t} \| _ {\infty} ^ {2} (\pi_ {\mathbf {X}}) _ {i i}.
$$
Then the result follows by noting that  $(\pi_{\mathbf{X}})_{ii}$  is equal to the leverage score of row  $i$ , and the sum of leverage scores is less than or equal to  $d$ .
We can use the same theorem to characterize the regression error for  $\mathbf{t}$  by changing the sign of  $\mathbf{y}^{(0)}$ .
4.2 Full Observation ATE
We now sketch a proof of our result for ATE estimation with full observation.
Proof sketch of Theorem 2. Here, we only sketch our proof for the variance bound. For  $S \subseteq [n]$  and  $\mathbf{z} \in \{-1, +1\}^n$  (where  $Z^{+} = \{i : z_{i} = +1\}$  and  $Z^{-} = \{i : z_{i} = -1\}$ ), we define
$$
\tau_ {S} := \frac {1}{| S |} \cdot \sum_ {i \in S} (y _ {i} ^ {(1)} - y _ {i} ^ {(0)}), \text {a n d} \widehat {\tau} _ {S} := \frac {2}{| S |} \cdot \left(\sum_ {i \in Z ^ {+} \cap S} \widetilde {y} _ {i} ^ {(2, 1)} - \sum_ {i \in Z ^ {-} \cap S} \widetilde {y} _ {i} ^ {(2, 0)}\right).
$$
Similarly, define  $\tau_{\overline{S}}$  and  $\widehat{\tau}_{\overline{S}}$ . Then for  $\widehat{\tau}$  as defined in Algorithm 1, we have,
$$
\begin{array}{l} \mathbb {E} _ {\mathbf {z}} \left[ \mathbb {E} _ {S} \left[ (\widehat {\tau} - \tau) ^ {2} \right] \right] = \mathbb {E} _ {\mathbf {z}} \left[ \mathbb {E} _ {S} \left[ \left(\frac {| S | \cdot \widehat {\tau} _ {S} + | \overline {{S}} | \cdot \widehat {\tau} _ {\overline {{S}}}}{n} - \frac {| S | \cdot \tau_ {S} + | \overline {{S}} | \cdot \tau_ {\overline {{S}}}}{n}\right) ^ {2} \right] \right] \\ = \mathbb {E} _ {S} \left[ \frac {| \overline {{S}} | ^ {2}}{n ^ {2}} \cdot \mathbb {E} _ {\mathbf {z}} \left[ (\widehat {\tau_ {S}} - \tau_ {\overline {{S}}}) ^ {2} \right] + \frac {| S | ^ {2}}{n ^ {2}} \cdot \mathbb {E} _ {\mathbf {z}} \left[ (\widehat {\tau} _ {S} - \tau_ {S}) ^ {2} \right] \right] + \mathbb {E} _ {S} \left[ 2 \cdot \frac {| S | \cdot | \overline {{S}} |}{n ^ {2}} \cdot \mathbb {E} _ {\mathbf {z}} \left[ (\widehat {\tau_ {S}} - \tau_ {\overline {{S}}}) \cdot (\widehat {\tau} _ {S} - \tau_ {S}) \right] \right]. \\ \end{array}
$$
By Cauchy-Schwarz inequality,  $\mathbb{E}_{\mathbf{z}}[(\widehat{\tau}_{\overline{S}} - \tau_{\overline{S}})\cdot (\widehat{\tau}_{S} - \tau_{S})]\leq \sqrt{\mathbb{E}_{\mathbf{z}}[(\widehat{\tau}_{\overline{S}} - \tau_{\overline{S}})^{2}]\cdot\mathbb{E}_{\mathbf{z}}[(\widehat{\tau}_{S} - \tau_{S})^{2}]}$ . Therefore we only need to bound  $\mathbb{E}_{\mathbf{z}}[(\widehat{\tau}_{\overline{S}} - \tau_{\overline{S}})^{2}]$  and  $\mathbb{E}_{\mathbf{z}}[(\widehat{\tau}_S - \tau_S)^2]$ . We can bound  $\mathbb{E}_{\mathbf{z}}[(\widehat{\tau}_S - \tau_S)^2]$  by Theorem 7 (ridge leverage score sampling) and using the fact that  $\widehat{\tau}_S$  is obtained by  $\widetilde{\mathbf{y}}^{(2,1)}$  and  $\widetilde{\mathbf{y}}^{(2,0)}$  which are adjusted by  $\widehat{\mathbf{b}}^{(2)}$ , a vector that is learned from the units in the set  $S$ . Since  $S$  and  $S$  are disjoint,  $\widehat{\mathbf{b}}^{(2)}$  is independent of  $\mathbf{z}_S$  and this allows us to bound  $\mathbb{E}_{\mathbf{z}}[(\widehat{\tau}_S - \tau_S)^2]$ . The bound on  $\mathbb{E}_{\mathbf{z}}[(\widehat{\tau}_{\overline{S}} - \tau_{\overline{S}})^2]$  follows from a similar argument.
In this section, we compare our method for estimating ATE and ITE with full observation with prior art and baselines on synthetic and real-world datasets. For ATE estimation, our experiments demonstrate that our approach either outperforms other unbiased estimation approaches or has comparable performance in terms of the achieved variance.
Algorithm 1: ATE estimation with leverage score sampling and cross adjustment
1 Input:  $\mathbf{X} \in \mathbb{R}^{n \times d}$ ,  $\mathbf{y}^{(0)}, \mathbf{y}^{(1)} \in \mathbb{R}^n$ ,  $\epsilon, \delta \in (0, 1]$ .
2 Set  $\mathbf{y} \in \mathbb{R}^n$  to a random vector such that for each  $i \in [n]$ ,  $y_i$  is independently and with equal probability either equal to  $y_i^{(0)}$  or  $y_i^{(1)}$ .
3 Set  $Z^{+}$  to the set of entries  $i$  such that  $y_{i}$  is equal to  $y_{i}^{(1)}$  and  $Z^{-} = [n]\setminus Z^{+}$
4 Set  $\lambda = \frac{6}{\epsilon^2} \cdot \zeta^2 \log (n / \delta)$ , where  $\zeta = \max_{i \in [n]} \| \mathbf{x}_i \|_2$ .
5 Set  $S \subseteq [n]$  to a random set of indices where each  $i \in [n]$  is included in  $S$  independently and with probability 0.5. Set  $\overline{S} = [n] \setminus S$ .
6 Set  $\mathbf{S}$  to an  $n \times n$  diagonal matrix corresponding to  $S$  where each entry  $\mathbf{S}_{ii} = 2$  if  $i \in S$ , and  $\mathbf{S}_{ii} = 2$ , otherwise. Similarly, set  $\overline{\mathbf{S}}$ .
7 Set  $\widehat{\mathbf{b}}^{(1)} = \underset {\mathbf{b}}{\arg \min}\| \mathbf{S}\mathbf{X}\mathbf{b} - \mathbf{S}\mathbf{y}\| _2^2 +\lambda \cdot \| \mathbf{b}\| _2^2,\widehat{\mathbf{b}}^{(2)} = \underset {\mathbf{b}}{\arg \min}\| \overline{\mathbf{S}}\mathbf{X}\mathbf{b} - \overline{\mathbf{S}}\mathbf{y}\| _2^2 +\lambda \cdot \| \mathbf{b}\| _2^2.$
8 Set  $\widetilde{\mathbf{y}}^{(1,0)} = \mathbf{y}^{(0)} - \mathbf{X}\widehat{\mathbf{b}}^{(1)},\widetilde{\mathbf{y}}^{(1,1)} = \mathbf{y}^{(1)} - \mathbf{X}\widehat{\mathbf{b}}^{(1)},\widetilde{\mathbf{y}}^{(2,0)} = \mathbf{y}^{(0)} - \mathbf{X}\widehat{\mathbf{b}}^{(2)},\widetilde{\mathbf{y}}^{(2,1)} = \mathbf{y}^{(1)} - \mathbf{X}\widehat{\mathbf{b}}^{(2)}$
9 Return  $\widehat{\tau} := \frac{2}{n} (\sum_{i\in Z^{+}\cap \overline{S}}\widetilde{y}_{i}^{(1,1)} + \sum_{i\in Z^{+}\cap S}\widetilde{y}_{i}^{(2,1)} - \sum_{i\in Z^{-}\cap \overline{S}}\widetilde{y}_{i}^{(1,0)} - \sum_{i\in Z^{-}\cap S}\widetilde{y}_{i}^{(2,0)})$
ATE methods. Our ATE estimation method is compared against five other approaches. (i) HT Uniform: This employs the Hurvitz Thompson estimator, using a uniform assignment of units to treatment and control groups. (ii) GSW: This design follows the methodology in [18]. (iii) Classic Regression Adjustment: As described in [24], this estimator is inherently biased. (iv) Leverage Score-Based: In this method, we learn two distinct vectors for  $\mathbf{y}^{(1)}$  and  $\mathbf{y}^{(0)}$  through leverage score sampling. The difference between these vectors predicts the treatment effect. Though this method has its advantages, it is biased and was previously employed by Addanki et al. [1]. (v) 4 Vectors: Obtained by specializing Mou et al. [28] to linear regressors, this approach adopts the same cross-adjustment mechanism to mitigate bias as in this paper. However, instead of random vector regression, it determines separate vectors for  $\mathbf{y}^{(1)}$  and  $\mathbf{y}^{(0)}$  for each group, leading to four vectors obtained by linear regressions.