SlowGuess/NeurIPS_2023_OCR huggingface dataset:
The optimal solution is  $\frac{1}{n}\mathbf{1}^{\top}(\mathbf{y} - (\mathbf{X} - \overline{\mathbf{X}})\widehat{\mathbf{b}}) = \frac{1}{n}\mathbf{1}^{\top}\mathbf{y} = \overline{y}$ . Therefore
$$
\left\| \widetilde {\mathbf {X}} \widetilde {\mathbf {b}} - \mathbf {y} \right\| _ {2} ^ {2} = \left\| (\mathbf {X} - \overline {{\mathbf {X}}}) \widehat {\mathbf {b}} - \mathbf {y} + \widetilde {b} _ {d + 1} \cdot \mathbf {1} \right\| _ {2} ^ {2} \geq \left\| (\mathbf {X} - \overline {{\mathbf {X}}}) \widehat {\mathbf {b}} - \mathbf {y} + \overline {{\mathbf {y}}} \right\| _ {2} ^ {2}. \tag {5}
$$
Thus by (4) and (5),
$$
\left\| (\mathbf {X} - \overline {{\mathbf {X}}}) \widehat {\mathbf {b}} - (\mathbf {y} - \overline {{\mathbf {y}}}) \right\| _ {2} ^ {2} \leq (1 + \epsilon) \min _ {\mathbf {b} \in \mathbb {R} ^ {d}} \left\| (\mathbf {X} - \overline {{\mathbf {X}}}) \mathbf {b} - (\mathbf {y} - \overline {{\mathbf {y}}}) \right\| _ {2} ^ {2},
$$
and the result follows by Lemma 11.
Note that the number of samples in Theorem 10 depends on the number of covariates (also called features). This dependence can be alleviated by using ridge leverage score sampling at the expense of a higher variance bound. If we want to sample  $m < d\log(d / \delta) / \epsilon^2$ , we can use ridge leverage score sampling, which gives a bound on the ridge regression error instead of the regression error. Using  $\lambda$ -ridge leverage scores, for  $\lambda = \zeta^2 \cdot \frac{n}{m}$ , where  $\zeta$  is the maximum  $\ell_2$  norm over rows of  $\mathbf{X}$ , gives the following result. Before providing this result, note that the block leverage score sampling approach of [12] shows that to solve
$$
\min _ {\mathbf {b} \in \mathbb {R} ^ {d}} \left[ \| \mathbf {X b} - \mathbf {y} \| _ {2} ^ {2} + \lambda \| \mathbf {b} _ {1: d - 1} \| _ {2} ^ {2} \right],
$$
we can use the same ridge leverage sampling as for  $\min_{\mathbf{b} \in \mathbb{R}^d} \left[ \| \mathbf{X}\mathbf{b} - \mathbf{y} \|_2^2 + \lambda \| \mathbf{b} \|_2^2 \right]$ , and the number of samples only needs to be increased by a constant factor. We require this for our ridge leverage score sampling-based mean estimation result.
Theorem 12 (Ridge leverage score sampling based regression adjustment). Let  $0 < \epsilon < 1$ ,  $0 < \delta < 0.5$ ,  $m \in \mathbb{N}$ , and  $\mathbf{X}$  be a covariate matrix. Then Algorithm 2 computes an unbiased estimator of the mean of an outcome vector  $\mathbf{y}$  (on a population of size  $n$ ) by using  $O(m \log (m / \delta) / \epsilon^2)$  samples. Moreover, there is an event  $\mathcal{E}$  with  $\operatorname*{Pr}(\mathcal{E}) \geq 1 - \delta$  (over the randomness of the algorithm) such that the variance of the estimator conditioned on  $\mathcal{E}$  is bounded by
$$
(1 + \epsilon) \frac {1}{m ^ {2}} \min  _ {\mathbf {b} \in \mathbb {R} ^ {d}} \left[ \left\| (\mathbf {X} - \overline {{\mathbf {X}}}) \mathbf {b} - (\mathbf {y} - \overline {{\mathbf {y}}}) \right\| _ {2} ^ {2} + \zeta^ {2} \| \mathbf {b} \| _ {2} ^ {2} \right]. \tag {6}
$$
Proof. Let  $\widetilde{\mathbf{X}} := \left[(\mathbf{X} - \overline{\mathbf{X}}) \quad \mathbf{1}\right]$ , and
$$
\mathbf {Z} = \left[ \begin{array}{c} (\mathbf {X} - \overline {{\mathbf {X}}}) \\ \sqrt {\lambda} \mathbf {I} _ {d} \end{array} \right], \widetilde {\mathbf {Z}} = \left[ \begin{array}{c c} (\mathbf {X} - \overline {{\mathbf {X}}}) & \mathbf {1} \\ \sqrt {\lambda} \mathbf {I} _ {d} & \mathbf {0} \end{array} \right].
$$
Matrix  $\widetilde{\mathbf{Z}}$  corresponds to the regression problem
$$
\min  _ {\mathbf {b} \in \mathbb {R} ^ {d + 1}} \left[ \left\| \widetilde {\mathbf {X}} \mathbf {b} - \mathbf {y} \right\| _ {2} ^ {2} + \lambda \left\| \mathbf {b} _ {1: d} \right\| _ {2} ^ {2} \right].
$$
In other words, defining the vector  $\widetilde{\mathbf{y}}\in \mathbb{R}^{n + d}$  as a vector that is equal to  $\mathbf{y}$  on the first  $n$  entries and is zero on other entries, we have
$$
\left\| \widetilde {\mathbf {Z}} \mathbf {b} - \widetilde {\mathbf {y}} \right\| _ {2} ^ {2} = \left\| \widetilde {\mathbf {X}} \mathbf {b} - \mathbf {y} \right\| _ {2} ^ {2} + \lambda \left\| \mathbf {b} _ {1: d} \right\| _ {2} ^ {2},
$$
for any vector  $\mathbf{b} \in \mathbb{R}^{d+1}$ . Now note that the solution to  $\arg \min_{\mathbf{b} \in \mathbb{R}^{d+1}} \left\| \widetilde{\mathbf{Z}} \mathbf{b} - \widetilde{\mathbf{y}} \right\|_2^2$  is
$$
(\widetilde {\mathbf {Z}} ^ {\top} \widetilde {\mathbf {Z}}) ^ {- 1} \widetilde {\mathbf {Z}} ^ {\top} \widetilde {\mathbf {y}} = (\widetilde {\mathbf {Z}} ^ {\top} \widetilde {\mathbf {Z}}) ^ {- 1} \widetilde {\mathbf {X}} ^ {\top} \mathbf {y},
$$
by definition of  $\widetilde{\mathbf{y}}$ . Moreover by definition,
$$
\widetilde {\mathbf {Z}} ^ {\top} \widetilde {\mathbf {Z}} = \widetilde {\mathbf {X}} ^ {\top} \widetilde {\mathbf {X}} + \lambda \left[ \begin{array}{c c} \mathbf {I} _ {d} & \mathbf {0} \\ \mathbf {0} ^ {\top} & 0 \end{array} \right].
$$
Therefore since  $\mathbf{1}^{\top}(\mathbf{X} - \overline{\mathbf{X}}) = \mathbf{0}^{\top}$ , we have
$$
\widetilde {\mathbf {Z}} ^ {\top} \widetilde {\mathbf {Z}} = \left[ \begin{array}{c c} (\mathbf {X} - \overline {{\mathbf {X}}}) ^ {\top} (\mathbf {X} - \overline {{\mathbf {X}}}) + \lambda \mathbf {I} _ {d} & \mathbf {0} \\ \mathbf {0} ^ {\top} & n \end{array} \right].
$$
Since this is a block diagonal matrix
$$
(\widetilde {\mathbf {Z}} ^ {\top} \widetilde {\mathbf {Z}}) ^ {- 1} = \left[ \begin{array}{c c} \big ((\mathbf {X} - \overline {{\mathbf {X}}}) ^ {\top} (\mathbf {X} - \overline {{\mathbf {X}}}) + \lambda \mathbf {I} _ {d} \big) ^ {- 1} & \mathbf {0} \\ \mathbf {0} ^ {\top} & n ^ {- 1} \end{array} \right]
$$
$$
(\widetilde {\mathbf {Z}} ^ {\top} \widetilde {\mathbf {Z}}) ^ {- 1} \widetilde {\mathbf {Z}} ^ {\top} \widetilde {\mathbf {y}} = \left[ \binom {(\mathbf {X} - \overline {{\mathbf {X}}}) ^ {\top} (\mathbf {X} - \overline {{\mathbf {X}}}) + \lambda \mathbf {I} _ {d}) ^ {- 1} (\mathbf {X} - \overline {{\mathbf {X}}}) ^ {\top} \mathbf {y}} {\overline {{y}}} \right].
$$
$$
\left\| \pi_ {\widetilde {\mathbf {Z}}} \widetilde {\mathbf {y}} - \widetilde {\mathbf {y}} \right\| _ {2} ^ {2} = \left\| \left[ \begin{array}{c} (\mathbf {X} - \overline {{\mathbf {X}}}) \\ \sqrt {\lambda} \mathbf {I} _ {d} \end{array} \right] \left((\mathbf {X} - \overline {{\mathbf {X}}}) ^ {\top} (\mathbf {X} - \overline {{\mathbf {X}}}) + \lambda \mathbf {I} _ {d}\right) ^ {- 1} (\mathbf {X} - \overline {{\mathbf {X}}}) ^ {\top} \mathbf {y} + \left[ \begin{array}{c} \overline {{\mathbf {y}}} - \mathbf {y} \\ \mathbf {0} _ {d} \end{array} \right] \right\| _ {2} ^ {2}.
$$
Note that  $\left((\mathbf{X} - \overline{\mathbf{X}})^{\top}(\mathbf{X} - \overline{\mathbf{X}}) + \lambda \mathbf{I}_d\right)^{-1}(\mathbf{X} - \overline{\mathbf{X}})^{\top}(\mathbf{y} - \overline{\mathbf{y}})$  is the solution to
$$
\underset {\mathbf {b} \in \mathbb {R} ^ {d}} {\arg \min } \left[ \left\| (\mathbf {X} - \overline {{\mathbf {X}}}) \mathbf {b} - (\mathbf {y} - \overline {{\mathbf {y}}}) \right\| _ {2} ^ {2} + \lambda \left\| \mathbf {b} \right\| _ {2} ^ {2} \right],
$$
and because 1 is in the null space of  $\mathbf{X} - \overline{\mathbf{X}}$
$$
\left(\left(\mathbf {X} - \overline {{\mathbf {X}}}\right) ^ {\top} (\mathbf {X} - \overline {{\mathbf {X}}}) + \lambda \mathbf {I} _ {d}\right) ^ {- 1} \left(\mathbf {X} - \overline {{\mathbf {X}}}\right) ^ {\top} (\mathbf {y} - \overline {{\mathbf {y}}}) = \left(\left(\mathbf {X} - \overline {{\mathbf {X}}}\right) ^ {\top} (\mathbf {X} - \overline {{\mathbf {X}}}) + \lambda \mathbf {I} _ {d}\right) ^ {- 1} \left(\mathbf {X} - \overline {{\mathbf {X}}}\right) ^ {\top} \mathbf {y}.
$$
$$
\left\| (\pi_ {\widetilde {\mathbf {Z}}} - \mathbf {I}) \widetilde {\mathbf {y}} \right\| _ {2} ^ {2} = \left\| (\pi_ {\mathbf {Z}} - \mathbf {I}) \left[ \begin{array}{c} \mathbf {y} - \overline {{\mathbf {y}}} \\ \mathbf {0} _ {d} \end{array} \right] \right\| _ {2} ^ {2},
$$
$$
\min  _ {\mathbf {b} \in \mathbb {R} ^ {d}} \left[ \left\| (\mathbf {X} - \overline {{\mathbf {X}}}) \mathbf {b} - (\mathbf {y} - \overline {{\mathbf {y}}}) \right\| _ {2} ^ {2} + \lambda \| \mathbf {b} \| _ {2} ^ {2} \right] = \min  _ {\mathbf {b} \in \mathbb {R} ^ {d + 1}} \left[ \left\| \widetilde {\mathbf {X}} \mathbf {b} - \mathbf {y} \right\| _ {2} ^ {2} + \lambda \| \mathbf {b} _ {1: d} \| _ {2} ^ {2} \right]. \tag {7}
$$
Now note that by ridge leverage score sampling, we find a vector  $\mathbf{b} \in \mathbb{R}^{d+1}$  such that, with high probability,
$$
\left\| \widetilde {\mathbf {X}} \widetilde {\mathbf {b}} - \mathbf {y} \right\| _ {2} ^ {2} + \lambda \left\| \widetilde {\mathbf {b}} _ {1: d} \right\| _ {2} ^ {2} \leq (1 + \epsilon) \min  _ {\mathbf {b} \in \mathbb {R} ^ {d + 1}} \left[ \left\| \widetilde {\mathbf {X}} \mathbf {b} - \mathbf {y} \right\| _ {2} ^ {2} + \lambda \| \mathbf {b} _ {1: d} \| _ {2} ^ {2} \right].
$$
$$
\left\| \widetilde {\mathbf {X}} \widetilde {\mathbf {b}} - \mathbf {y} \right\| _ {2} ^ {2} + \lambda \left\| \widetilde {\mathbf {b}} _ {1: d} \right\| _ {2} ^ {2} \leq (1 + \epsilon) \min  _ {\mathbf {b} \in \mathbb {R} ^ {d}} \left[ \| (\mathbf {X} - \overline {{\mathbf {X}}}) \mathbf {b} - (\mathbf {y} - \overline {{\mathbf {y}}}) \| _ {2} ^ {2} + \lambda \| \mathbf {b} \| _ {2} ^ {2} \right]. \tag {8}
$$
Now let  $\widehat{\mathbf{b}}\in \mathbb{R}^d$  be the vector that is equal to the first  $d$  entries of  $\widetilde{\mathbf{b}}$ . Then we have
$$
\left\| \widetilde {\mathbf {X}} \widetilde {\mathbf {b}} - \mathbf {y} \right\| _ {2} ^ {2} + \lambda \left\| \widetilde {\mathbf {b}} _ {1: d} \right\| _ {2} ^ {2} = \left\| (\mathbf {X} - \overline {{\mathbf {X}}}) \widehat {\mathbf {b}} + \widetilde {\mathbf {b}} _ {d + 1} \cdot \mathbf {1} - \mathbf {y} \right\| _ {2} ^ {2} + \lambda \left\| \widehat {\mathbf {b}} \right\| _ {2} ^ {2}.
$$
Now consider a regression problem of the following form.
The optimal solution to this problem is  $\frac{1}{n}\mathbf{1}^{\top}((\mathbf{y} - (\mathbf{X} - \overline{\mathbf{X}})\widehat{\mathbf{b}})) = \frac{1}{n}\mathbf{1}^{\top}\mathbf{y} = \overline{y}$ . Therefore
$$
\left\| \widetilde {\mathbf {X}} \widetilde {\mathbf {b}} - \mathbf {y} \right\| _ {2} ^ {2} = \left\| (\mathbf {X} - \overline {{\mathbf {X}}}) \widehat {\mathbf {b}} - \mathbf {y} + \widetilde {b} _ {d + 1} \cdot \mathbf {1} \right\| _ {2} ^ {2} \geq \left\| (\mathbf {X} - \overline {{\mathbf {X}}}) \widehat {\mathbf {b}} - \mathbf {y} + \overline {{\mathbf {y}}} \right\| _ {2} ^ {2}. \tag {9}
$$
Thus by (8) and (9), and because  $\lambda \left\| \widehat{\mathbf{b}} \right\|_2^2 \geq 0$
$$
\left\| (\mathbf {X} - \overline {{\mathbf {X}}}) \widehat {\mathbf {b}} - (\mathbf {y} - \overline {{\mathbf {y}}}) \right\| _ {2} ^ {2} \leq (1 + \epsilon) \min  _ {\mathbf {b} \in \mathbb {R} ^ {d}} \left[ \| (\mathbf {X} - \overline {{\mathbf {X}}}) \mathbf {b} - (\mathbf {y} - \overline {{\mathbf {y}}}) \| _ {2} ^ {2} + \lambda \| \mathbf {b} \| _ {2} ^ {2} \right].
$$
Therefore by Lemma 11, the variance of the regression-adjusted estimator with  $m$  samples is less than,
$$
\frac {(1 + \epsilon)}{m n} \min  _ {\mathbf {b} \in \mathbb {R} ^ {d}} \left[ \left\| (\mathbf {X} - \overline {{\mathbf {X}}}) \mathbf {b} - (\mathbf {y} - \overline {{\mathbf {y}}}) \right\| _ {2} ^ {2} + \lambda \| \mathbf {b} \| _ {2} ^ {2} \right]
$$
Finally by definition  $\frac{\lambda}{mn} = \frac{\zeta^2n / m}{nm} = \frac{\zeta^2}{m^2}$ , and the result follows.
We now explain previous approaches to mean estimation. The classical approach to regression adjustment [9] is to learn a vector  $\widehat{\mathbf{b}}$  from the same set of samples that are used for estimation.
$$
\widehat {\mathbf {b}} = \underset {\mathbf {b}} {\arg \min} \| \mathbf {y} _ {S ^ {\prime}} - \overline {{\mathbf {y}}} _ {S ^ {\prime}} - (\mathbf {X} _ {S ^ {\prime}} - \overline {{\mathbf {X}}} _ {S ^ {\prime}}) \mathbf {b} \| _ {2} ^ {2}.
$$
The Gram-Schmidt walk design of Harshaw et al. [18] has originally been developed for the average treatment effect estimation and achieves the following bound using the Horvitz-Thompson estimator — see Algorithm 5.
Theorem 13 (Gram-Schmidt walk design [18]). Let  $0 < \phi < 1$ . Given  $\mathbf{y}^{(1)}, \mathbf{y}^{(0)} \in \mathbb{R}^n$  and  $\mathbf{X} \in \mathbb{R}^{n \times d}$ , the Horvitz-Thompson estimator under the Gram-Schmidt walk design is unbiased and has a variance of
$$
\mathbb {E} \left[ (\widehat {\tau} - \tau) ^ {2} \right] \leq \frac {1}{n ^ {2}} \min  _ {\mathbf {b} \in \mathbb {R} ^ {n}} \left[ \frac {1}{\phi} \| \mathbf {X b} - \boldsymbol {\mu} \| _ {2} ^ {2} + \frac {\zeta^ {2}}{1 - \phi} \| \mathbf {b} \| _ {2} ^ {2} \right],
$$
where  $\pmb{\mu} = \mathbf{y}^{(1)} + \mathbf{y}^{(0)}$ .
A special setting of ATE estimation problem reduces to mean estimation. Suppose  $\mathbf{y}^{(1)} = \mathbf{y}$  and  $\mathbf{y}^{(0)} = \mathbf{0}$ . Then the average treatment effect for this problem equals the mean of  $\mathbf{y}$ . This approach partitions the samples into two groups. The first outcome is observed for the first group, and the second outcome is observed for the second group. The estimate is then produced by the Horvitz-Thompson estimator. Note that in our special case, only the samples in the first group amount to observation since  $\mathbf{y}^{(0)}$  is the fixed vector of all zeros. This yields the following guarantee for mean estimation
Theorem 14 (Gram-Schmidt walk based design [18]). The Horvitz-Thompson estimator applied to the output of Algorithm 5 gives an unbiased estimator of the mean of  $\mathbf{y}$  with a variance of less than or equal to
$$
\frac {1}{m ^ {2}} \min  _ {\mathbf {b} \in \mathbb {R} ^ {d}} \left[ \left\| (\mathbf {X} - \overline {{\mathbf {X}}}) \mathbf {b} - (\mathbf {y} - \overline {{\mathbf {y}}}) \right\| _ {2} ^ {2} + \zeta^ {2} \| \mathbf {b} \| _ {2} ^ {2} \right] \tag {10}
$$
with  $O(m)$  samples.
In comparison, the bound of Theorem 12 is better by a factor of  $m / n$  on the  $\| \mathbf{X}\mathbf{b} - \mathbf{y}\| _2^2$  term.
B Individual Treatment Effect Estimation
In this section, we give an ITE estimation algorithm with partial observation (Algorithm 3). Our approach is to learn a vector  $\mathbf{b}$  that nearly optimally fits  $\mathbf{t} := \mathbf{y}^{(1)} - \mathbf{y}^{(0)}$  linearly onto  $\mathbf{X}$ , i.e.,  $\| \mathbf{X}\mathbf{b} - \mathbf{t}\|_2$  is small. Since we do not have access to  $\mathbf{t}/2$ , we perform regression on a random vector for which the expectation is  $\mathbf{t}$ . To this end, we construct the random vector by setting  $\mathbf{y}_i$  independently to either  $\mathbf{y}_i^{(1)}$  or  $-\mathbf{y}_i^{(0)}$  with equal probability. If we observe all samples, and use Theorem 8 on vectors  $\mathbf{y}^{(1)}$  and  $-\mathbf{y}^{(0)}$ , this approach gives the following expected error bound.
$$
\mathbb {E} \left[ \| 2 \mathbf {X b} ^ {*} - \mathbf {t} \| _ {2} ^ {2} \right] \leq d \left\| \mathbf {y} ^ {(1)} + \mathbf {y} ^ {(0)} \right\| _ {\infty} ^ {2} + \min  _ {\mathbf {b}} \| \mathbf {X b} - \mathbf {t} \| _ {2} ^ {2}.
$$
We can achieve a similar bound using leverage score sampling by thresholding the probability of sampling rows from below while we guarantee that the number of samples does not increase too much with high probability. The following is the main technical contribution of this part.
Theorem 1 (ITE Estimation). For any  $\epsilon, \delta, \alpha \in (0,1]$ , there exists a randomized algorithm (Algorithm 3) that observes a potential outcome for  $O(d\log (d / \delta) / \epsilon^2 + \alpha \cdot n)$  units and outputs a vector  $\hat{\mathbf{t}}$ . Moreover there is an event  $\mathcal{E}$  such that  $\operatorname*{Pr}(\mathcal{E}) \geq 1 - \delta$  (over the randomness of the algorithm) and
Proof. Note that  $2\mathbf{y} = \mathbf{t} + \mathbf{z}\odot \boldsymbol{\mu}$ . Then
$$
\begin{array}{l} 2 \mathbf {X} \widehat {\mathbf {b}} = 2 \mathbf {X} (\mathbf {X} ^ {\top} \mathbf {S} \mathbf {S} \mathbf {X}) ^ {- 1} \mathbf {X} ^ {\top} \mathbf {S} \mathbf {S} \mathbf {y} \\ = \mathbf {X} \left(\mathbf {X} ^ {\top} \mathbf {S S X}\right) ^ {- 1} \mathbf {X} ^ {\top} \mathbf {S S t} + \mathbf {X} \left(\mathbf {X} ^ {\top} \mathbf {S S X}\right) ^ {- 1} \mathbf {X} ^ {\top} \mathbf {S S} (\mathbf {z} \odot \boldsymbol {\mu}). \\ \end{array}
$$
Algorithm 3: ITE estimation with random vector regression and leverage score sampling
1 Input:  $\mathbf{X} \in \mathbb{R}^{n \times d}$ ,  $\mathbf{y}^{(1)}, \mathbf{y}^{(2)} \in \mathbb{R}^n$ ,  $0 < \epsilon < 1$ ,  $0 < \alpha \leq 1$ .
2 Compute the leverage scores  $\ell_{i}$  of  $\mathbf{X}$ .
3 Set  $p_i = \min \{1, \max \{\alpha, 3 \cdot \ell_i \cdot \log(d / \delta) / \epsilon^2\}\}$  for  $i \in [n]$ .
4 Let  $\mathbf{S} \in \mathbb{R}^{n \times n}$  be a random diagonal matrix with  $\mathbf{S}_{ii} = 1 / \sqrt{p_i}$  with probability  $p_i$ , and  $\mathbf{S}_{ii} = 0$ , otherwise.
5 Let  $\widehat{\mathbf{b}} = \arg \min_{\mathbf{b} \in \mathbb{R}^{d + 1}} \| \mathbf{S}\mathbf{X}\mathbf{b} - \mathbf{S}\mathbf{y} \|_2^2$ .
6 Return  $\widehat{\mathbf{t}} := 2\mathbf{X}\widehat{\mathbf{b}}$ .
$$
(\mathbf{X}^{\top}\mathbf{S}\mathbf{S}\mathbf{X})^{-1}\mathbf{X}^{\top}\mathbf{S}\mathbf{S}\mathbf{t} = \operatorname *{arg  min}_{\mathbf{b}\in \mathbb{R}^{d}}\| \mathbf{S}\mathbf{X}\mathbf{b} - \mathbf{S}\mathbf{t}\|_{2}^{2}.
$$
Therefore by Theorem 7, and since  $\| \pi_{\mathbf{X}}\mathbf{t} - \mathbf{t}\| _2^2 = \min_{\mathbf{b}}\| \mathbf{X}\mathbf{b} - \mathbf{t}\| _2^2$
$$
\left\| \mathbf {X} \left(\mathbf {X} ^ {\top} \mathbf {S S X}\right) ^ {- 1} \mathbf {X} ^ {\top} \mathbf {S S t} - \mathbf {t} \right\| _ {2} ^ {2} \leq (1 + \epsilon) \| \pi_ {\mathbf {X}} \mathbf {t} - \mathbf {t} \| _ {2} ^ {2}.
$$
$$
\begin{array}{l} \mathbb {E} \left[ \left\| \mathbf {X} \left(\mathbf {X} ^ {\top} \mathbf {S S X}\right) ^ {- 1} \mathbf {X} ^ {\top} \mathbf {S S} (\mathbf {z} \odot \boldsymbol {\mu}) \right\| _ {2} ^ {2} \right] \leq (1 + \epsilon) \mathbb {E} \left[ (\mathbf {z} \odot \boldsymbol {\mu}) ^ {\top} \mathbf {S S X} \left(\mathbf {X} ^ {\top} \mathbf {S S X}\right) ^ {- 1} \mathbf {X} ^ {\top} \mathbf {S S} (\mathbf {z} \odot \boldsymbol {\mu}) \right] \\ = (1 + \epsilon) \sum_ {i: \mathbf {S} _ {i i} \neq 0} \frac {\ell_ {i} (\mathbf {S X}) \cdot \boldsymbol {\mu} _ {i} ^ {2}}{p _ {i}} \\ \leq \frac {1}{\alpha} \cdot (1 + \epsilon) \| \boldsymbol {\mu} \| _ {\infty} ^ {2} \sum_ {i: \mathbf {S} _ {i i} \neq 0} \ell_ {i} (\mathbf {S X}) \\ \leq \frac {d}{\alpha} \cdot (1 + \epsilon) \| \boldsymbol {\mu} \| _ {\infty} ^ {2}, \\ \end{array}
$$
where the last inequality follows from the fact that the sum of leverage scores is less than  $d$ , and the penultimate inequality follows from the definition of  $\mathbf{p}_i$ . Finally, the result follows by noting that,
$$
\mathbb {E} \left[ \left(\mathbf {X} \left(\mathbf {X} ^ {\top} \mathbf {S S X}\right) ^ {- 1} \mathbf {X} ^ {\top} \mathbf {S S t} - \mathbf {t}\right) ^ {\top} \mathbf {X} \left(\mathbf {X} ^ {\top} \mathbf {S S X}\right) ^ {- 1} \mathbf {X} ^ {\top} \mathbf {S S} (\mathbf {z} \odot \boldsymbol {\mu}) \right] = 0.
$$
C ATE with Full Observation
Our approach for ATE with full observation is similar to ITE in the sense that we perform regression on a random vector to learn a solution vector  $\widehat{\mathbf{b}}$ . Then, for estimation, we apply regression adjustment based on  $\widehat{\mathbf{b}}$  to adjust the outcomes. We first present a naive version of this approach that is biased and achieves a factor proportional to  $d^{2} / n^{2}$  in one of the terms bounding the variance. We improve this to  $d / n^{2}$  by partitioning the units into two groups and using the learned vector for each group to adjust the estimation in the other group. With this approach, our estimator is unbiased.
Before discussing our results, we discuss the other approaches in the literature. The Gram-Schmidt walk design of [18] that uses the Horvitz-Thompson estimator gives the following variance for the ATE estimation.
$$
\frac {1}{n ^ {2}} \min  _ {\mathbf {b} \in \mathbb {R} ^ {n}} \left[ \frac {1}{\phi} \| \mathbf {X b} - \boldsymbol {\mu} \| _ {2} ^ {2} + \frac {\zeta^ {2}}{1 - \phi} \| \mathbf {b} \| _ {2} ^ {2} \right],
$$
where  $0 < \phi < 1$  is an input parameter. The classic regression adjustment with interaction term gives an asymptotic variance proportional to  $\frac{1}{m} (\sigma_1^2 + 2\sigma_0^2 + \sigma_{1,0})$ , where  $\sigma_1^2$  is the variance of  $\widetilde{y}_i^{(1)} = y_i^{(1)} - \overline{y}^{(1)} + (\mathbf{X}_i - \overline{\mathbf{x}})\mathbf{b}^{(1)}$  [24], where  $\mathbf{b}^{(1)}$  is the best linear fit for  $\mathbf{y}^{(1)} - \overline{\mathbf{y}}^{(1)}$  using  $\mathbf{X} - \overline{\mathbf{X}}$  when  $n \to \infty$  [24]. Similarly,  $\sigma_1^2$  is the variance of  $\widetilde{y}_i^{(0)} = y_i^{(0)} - \overline{y}^{(0)} + (\mathbf{X}_i - \overline{\mathbf{x}})\mathbf{b}^{(0)}$ , and  $\sigma_{1,0}$  is the covariance of  $\widetilde{y}_i^{(1)} \cdot \widetilde{y}_i^{(0)}$ . Note that the above mentioned bound is achieved by learning two separate vectors for regression over  $y_i^{(1)} - \overline{y}^{(1)}$  and  $y_i^{(0)} - \overline{y}^{(0)}$ . Moreover, note that because these bounds are asymptotic, we need to assume a model of data generation.
For our naive approach, we first need the following technical lemma.
Lemma 15. Let  $\mathbf{X} \in \mathbb{R}^{n \times d}$ ,  $\mathbf{T} \in \mathbb{R}^{n \times n}$  be a diagonal matrix and  $\mathbf{z} \in \{-1, +1\}^n$  be a random vector where each entry is picked independently and uniformly at random. Then
$$
\mathbb {E} \left[ \mathbf {z} ^ {\top} \mathbf {T} \pi_ {\mathbf {X}} \mathbf {z} \mathbf {z} ^ {\top} \pi_ {\mathbf {X}} \mathbf {T} \mathbf {z} \right] \leq (d ^ {2} + 2 d) \cdot \| \mathbf {t} \| _ {\infty} ^ {2},
$$
where  $\mathbf{t} \in \mathbb{R}^n$  is the vector corresponding to the diagonal matrix  $\mathbf{T}$ .
Proof. Let  $\mathbf{M} \in \mathbb{R}^{n,n}$  and, Then
$$
\begin{array}{l} \mathbb {E} \left[ \mathbf {z} ^ {\top} \mathbf {M} \mathbf {z} \mathbf {z} ^ {\top} \mathbf {M} ^ {\top} \mathbf {z} \right] = \mathbb {E} \left[ (\sum_ {i = 1} ^ {n} \sum_ {j = 1} ^ {n} z _ {i} z _ {j} \mathbf {M} _ {i j}) ^ {2} \right] \\ = \mathbb {E} \left[ \sum_ {i = 1} ^ {n} \sum_ {j = 1} ^ {n} \sum_ {k = 1} ^ {n} \sum_ {l = 1} ^ {n} z _ {i} z _ {j} {\bf z} _ {k} {\bf z} _ {l} {\bf M} _ {i j} {\bf M} _ {k l} \right] \\ = \sum_ {i = 1} ^ {n} \sum_ {j = 1} ^ {n} \sum_ {k = 1} ^ {n} \sum_ {l = 1} ^ {n} \mathbb {E} \left[ z _ {i} z _ {j} \mathbf {z} _ {k} \mathbf {z} _ {l} \mathbf {M} _ {i j} \mathbf {M} _ {k l} \right]. \\ \end{array}
$$
Note that for any term such that there exists one of the  $i,j,k,l$  that is not equal to any of the other three, the expectation is equal to zero. Therefore we need to look at terms where for each  $i,j,k,l$ , at least one of the other three is equal to it. Therefore we have four cases:
1.  $i = j = k = l$ . This gives an expectation of  $\mathbf{M}_{ii}^2$ . So total from these terms is  $\sum_{i=1}^{n} \mathbf{M}_{ii}^2$ .
2.  $i = j$  and  $k = l$  but  $i \neq k$ . This gives an expectation of  $\mathbf{M}_{ii}\mathbf{M}_{kk}$ . So total from these terms is  $2\sum_{i=1}^{n-1}\sum_{k=i+1}^{n}\mathbf{M}_{ii}\mathbf{M}_{kk}$ .
3.  $i = k$ ,  $j = l$  and  $i \neq j$ . This gives an expectation of  $\mathbf{M}_{ij}^2$ . So total from these terms is  $2\sum_{i=1}^{n-1}\sum_{j=i+1}^{n}\mathbf{M}_{ij}^2$ .
4.  $i = l$ ,  $j = k$ ,  $i \neq j$ . This gives an expectation of  $\mathbf{M}_{ij} \mathbf{M}_{ji}$ . So total from these terms is  $2 \sum_{i=1}^{n-1} \sum_{j=i+1}^{n} \mathbf{M}_{ij} \mathbf{M}_{ji}$ .
Recall that we denote the leverage scores of matrix  $\mathbf{X}$  with  $\ell_{i}$ . Therefore taking the above over our matrix  $\mathbf{M} = \mathbf{T}\pi_{\mathbf{X}}$ , since  $\mathbf{M}_{ij} = t_i(\pi_{\mathbf{X}})_{ij}$ , we have