SlowGuess/NeurIPS_2023_OCR huggingface dataset:
$$
\begin{array}{l} \mathbb {E} \left[ \mathbf {z} ^ {\top} \mathbf {M} \mathbf {z} \mathbf {z} ^ {\top} \mathbf {M} ^ {\top} \mathbf {z} \right] = \sum_ {i = 1} ^ {n} \sum_ {j = 1} ^ {n} t _ {i} t _ {j} (\pi_ {\mathbf {X}}) _ {i i} (\pi_ {\mathbf {X}}) _ {j j} \\ + 2 \sum_ {i = 1} ^ {n - 1} \sum_ {j = i + 1} ^ {n} t _ {i} ^ {2} (\pi_ {\mathbf {X}}) _ {i j} ^ {2} \\ + 2 \sum_ {i = 1} ^ {n - 1} \sum_ {j = i + 1} ^ {n} t _ {i} t _ {j} (\pi_ {\mathbf {X}}) _ {i j} ^ {2} \\ \leq \left\| \mathbf {t} \right\| _ {\infty} ^ {2} \cdot \left(\left(\sum_ {i = 1} ^ {n} \ell_ {i}\right) ^ {2} + 2 \left\| \pi_ {\mathbf {X}} \right\| _ {F} ^ {2}\right) \\ \leq \| \mathbf {t} \| _ {\infty} ^ {2} \cdot \left(d ^ {2} + 2 \operatorname {t r} \left(\pi_ {\mathbf {X}}\right)\right) \\ \leq \| \mathbf {t} \| _ {\infty} ^ {2} \cdot (d ^ {2} + 2 d), \\ \end{array}
$$
where the last two inequalities follow from the fact that the sum of leverage scores (which are the diagonal entries of the projection matrix  $\pi_{\mathbf{X}}$ ) is at most the rank of the matrix.
We are now equipped to analyze the ATE estimation obtained by purely using random vector regression and regression adjustment.
Theorem 16 (ATE with full observation). 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 either equal to  $y_{i}^{(0)}$  or  $y_{i}^{(1)}$ . Moreover let  $Z^{+}$  be the set of entries  $i$  such that  $y_{i}$  is equal to  $y_{i}^{(1)}$  and  $Z^{-} = [n]\setminus Z^{+}$ . Let  $\widetilde{\mathbf{b}} = \arg \min_{\mathbf{b}}\| \mathbf{X}\mathbf{b} - \mathbf{y}\| _2$ . Let  $\widetilde{\mathbf{y}}^{(0)} = \mathbf{y}^{(0)} - \mathbf{X}\widehat{\mathbf{b}}$  and  $\widetilde{\mathbf{y}}^{(1)} = \mathbf{y}^{(1)} - \mathbf{X}\widehat{\mathbf{b}}$ . Let  $\widehat{\tau}$  be the following Horvitz-Thompson estimator, i.e.,
$$
\widehat {\tau} = \frac {2}{n} (\sum_ {i \in Z ^ {+}} \widetilde {y} _ {i} ^ {(1)} - \sum_ {i \in Z ^ {-}} \widetilde {y} _ {i} ^ {(0)}).
$$
Then  $|\mathbb{E}[\widehat{\tau} -\tau ]|\leq \frac{d}{n}\| \mathbf{y}^{(1)} - \mathbf{y}^{(0)}\|_{\infty},$  and
$$
\mathbb {E} \left[ (\widehat {\tau} - \tau) ^ {2} \right] \leq \frac {d ^ {2} + 2 d}{n ^ {2}} \cdot \left\| \mathbf {y} ^ {(1)} - \mathbf {y} ^ {(0)} \right\| _ {\infty} ^ {2} + \frac {1}{n ^ {2}} \min  _ {\mathbf {b}} \| \mathbf {X b} - \boldsymbol {\mu} \| _ {2} ^ {2},
$$
where  $\pmb {\mu} = \mathbf{y}^{(1)} + \mathbf{y}^{(0)}$
Proof. Let  $\mathbf{z} \in \{-1, +1\}^n$ , where  $z_i = +1$  if  $i \in Z^+$ , and  $z_i = -1$ , otherwise. Moreover let  $\widetilde{\boldsymbol{\mu}} := \widetilde{\mathbf{y}}^{(0)} + \widetilde{\mathbf{y}}^{(1)} = \mathbf{y}^{(0)} + \mathbf{y}^{(1)} - 2\mathbf{X}\widehat{\mathbf{b}}$ . First, note that  $2\mathbf{y} = \boldsymbol{\mu} + \mathbf{z} \odot \mathbf{t}$ . Therefore
$$
\widetilde {\boldsymbol {\mu}} = \boldsymbol {\mu} - 2 \mathbf {X} \widehat {\mathbf {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}).
$$
$$
\begin{array}{l} \widehat {\tau} - \tau = \frac {1}{n} \left(\sum_ {i \in Z ^ {+}} \left(2 \widetilde {y} _ {i} ^ {(1)} - y _ {i} ^ {(1)} + y _ {i} ^ {(0)}\right) - \sum_ {i \in Z ^ {-}} \left(2 \widetilde {y} _ {i} ^ {(0)} + y _ {i} ^ {(1)} - y _ {i} ^ {(0)}\right)\right) \tag {11} \\ = \frac {1}{n} \left(\sum_ {i \in Z ^ {+}} \left(y _ {i} ^ {(1)} + y _ {i} ^ {(0)} - 2 \mathbf {X} \widehat {\mathbf {b}}\right) - \sum_ {i \in Z ^ {-}} \left(y _ {i} ^ {(1)} + y _ {i} ^ {(0)} - 2 \mathbf {x} _ {i} ^ {\top} \widehat {\mathbf {b}}\right)\right) \\ = \frac {1}{n} \widetilde {\boldsymbol {\mu}} ^ {\top} \mathbf {z}. \\ \end{array}
$$
$$
\mathbb {E} [ \widehat {\tau} - \tau ] = \frac {1}{n} \mathbb {E} \left[ (\boldsymbol {\mu} - \pi_ {\mathbf {X}} \boldsymbol {\mu}) ^ {\top} \mathbf {z} - (\mathbf {z} \odot \mathbf {t}) ^ {\top} \pi_ {\mathbf {X}} \mathbf {z} \right]. \tag {12}
$$
Thus by the linearity of expectation and since  $z_{i}$  is equal to  $+1$  and  $-1$  with an equal probability,  $\mathbb{E}\left[(\boldsymbol {\mu} - \pi_{\mathbf{X}}\boldsymbol {\mu})^{\top}\mathbf{z}\right] = 0$ . Let  $\mathbf{T}$  be a diagonal matrix associated with  $\mathbf{t}$ . Then  $\mathbf{z}\odot \mathbf{t} = \mathbf{T}\mathbf{z}$ . Therefore
$$
\mathbb {E} \left[ \widehat {\boldsymbol {\tau}} - \boldsymbol {\tau} \right] = - \frac {1}{n} \mathbb {E} \left[ \mathbf {z} ^ {\top} \mathbf {T} \boldsymbol {\pi} _ {\mathbf {X}} \mathbf {z} \right]
$$
Therefore since  $z_{i}$  and  $z_{j}$  are independent for  $i \neq j$ , and  $z_{i}^{2} = 1$  with probability one.
$$
\mathbb {E} \left[ \widehat {\tau} - \tau \right] = - \frac {1}{n} \sum_ {i = 1} ^ {n} t _ {i} \left(\pi_ {\mathbf {X}}\right) _ {i i} = - \frac {1}{n} \sum_ {i = 1} ^ {n} t _ {i} \ell_ {i}
$$
Therefore by the triangle inequality
$$
| \mathbb {E} [ \widehat {\tau} - \tau ] | \leq \frac {1}{n} \sum_ {i = 1} ^ {n} | \mathbf {t} _ {i} \ell_ {i} | \leq \frac {\| \mathbf {t} \| _ {\infty}}{n} \sum_ {i = 1} ^ {n} | \ell_ {i} |.
$$
Then the first part of the theorem follows by the fact that the sum of leverage scores is less than  $d$ , and noting that by definition,  $\mathbf{t} = \mathbf{y}^{(1)} - \mathbf{y}^{(0)}$ .
We now prove the second part of the theorem. First, note that since  $z_{i}$  and  $z_{j}$  are independent for  $i \neq j$  and  $z_{i} = \pm 1$  with equal probability,
$$
\mathbb {E} \left[ (\boldsymbol {\mu} - \pi_ {\mathbf {X}} \boldsymbol {\mu}) ^ {\top} \mathbf {z} \mathbf {z} ^ {\top} \mathbf {T} \pi_ {\mathbf {X}} \mathbf {z} \right] = \mathbb {E} \left[ (\boldsymbol {\mu} - \pi_ {\mathbf {X}} \boldsymbol {\mu}) ^ {\top} \mathbf {T} \pi_ {\mathbf {X}} \mathbf {z} \right] = 0.
$$
Therefore by (11) and Lemma 15,
$$
\begin{array}{l} \mathbb {E} \left[ (\widehat {\tau} - \tau) ^ {2} \right] = \frac {1}{n ^ {2}} \mathbb {E} \left[ \mathbf {z} ^ {\top} (\pmb {\mu} - \pi_ {\mathbf {X}} \pmb {\mu}) (\pmb {\mu} - \pi_ {\mathbf {X}} \pmb {\mu}) ^ {\top} \mathbf {z} \right] \\ + \frac {1}{n ^ {2}} \mathbb {E} \left[ \mathbf {z} ^ {\top} \mathbf {T} \boldsymbol {\pi_ {\mathbf {X}}} \mathbf {z} \mathbf {z} ^ {\top} \boldsymbol {\pi_ {\mathbf {X}}} \mathbf {T} \mathbf {z} \right] \\ \leq \frac {1}{n ^ {2}} \left(\| \boldsymbol {\mu} - \pi_ {\mathbf {X}} \boldsymbol {\mu} \| _ {2} ^ {2} + (d ^ {2} + 2 d) \cdot \| \mathbf {t} \| _ {\infty} ^ {2}\right). \\ \end{array}
$$
The next step is two use a subset of samples for regression and the rest of the samples for estimating the ATE.
Since we want to have a guarantee on the regression error, we need to make sure that the sampling (i.e., partitioning) to two groups is done according to leverage scores. However, if the leverage score of a row is greater than half, then it cannot have a probability larger than its leverage score for being assigned to each group. To reduce this probability, we use the ridge leverage score. The following lemma states that by picking a large enough ridge parameter, we can achieve a small enough ridge leverage score for all rows (therefore, small enough sampling probabilities). This is at the expense of a regularization term of the form  $\lambda \cdot \| \mathbf{b}\| _2^2$  in the regression error.
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]$ .
Proof. Recall that  $\zeta \coloneqq \max_{i\in [n]}\| \mathbf{x}_i\| _2$ . Therefore,  $\| \mathbf{x}_i\| ^2\leq \zeta^2$  for all  $i\in [n]$ . We have
$$
\ell_ {i} ^ {(\lambda)} (\mathbf {X}) = \mathbf {x} _ {i} ^ {\top} (\mathbf {X} ^ {\top} \mathbf {X} + \lambda \mathbf {I}) ^ {- 1} \mathbf {x} _ {i} \leq \mathbf {x} _ {i} ^ {\top} (\lambda \mathbf {I}) ^ {- 1} \mathbf {x} _ {i} = \| \mathbf {x} _ {i} \| ^ {2} / \lambda \leq 1 / c.
$$
Equipped with this, we present the main result of this section.
$$
\mathbb {E} \left[ (\widehat {\tau} - \tau) ^ {2} | \mathcal {E} \right] \leq \frac {8 (1 + \epsilon)}{n ^ {2}} \min  _ {\mathbf {b} \in \mathbb {R} ^ {d}} \left(\| \mathbf {X b} - \boldsymbol {\mu} \| _ {2} ^ {2} + 1 0 0 \log (n / \delta) \cdot \zeta^ {2} \cdot \| \mathbf {b} \| _ {2} ^ {2}\right) + \frac {3 2 d}{n ^ {2}} \cdot \| \mathbf {y} ^ {(1)} - \mathbf {y} ^ {(0)} \| _ {\infty} ^ {2},
$$
Proof. 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  $\widehat{\tau_{\overline{S}}}$ . Then 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_ {\overline {{S}}}} - \tau_ {\overline {{S}}}) ^ {2} \right] \right] \\ + \mathbb {E} _ {S} \left[ \frac {| S | ^ {2}}{n ^ {2}} \cdot \mathbb {E} _ {\mathbf {z}} \left[ \left(\widehat {\tau} _ {S} - \tau_ {S}\right) ^ {2} \right] \right] \\ + \mathbb {E} _ {S} \left[ 2 \cdot \frac {| S | \cdot | \overline {{S}} |}{n ^ {2}} \cdot \mathbb {E} _ {\mathbf {z}} \left[ (\widehat {\tau} _ {\overline {{S}}} - \tau_ {\overline {{S}}}) \cdot (\widehat {\tau} _ {S} - \tau_ {S}) \right] \right]. \\ \end{array}
$$
We now bound all of the terms above. Let
$$
\widetilde {\boldsymbol {\mu}} ^ {(1)} := \widetilde {\mathbf {y}} ^ {(1, 0)} + \widetilde {\mathbf {y}} ^ {(1, 1)} = \mathbf {y} ^ {(0)} + \mathbf {y} ^ {(1)} - 2 \mathbf {X} \widehat {\mathbf {b}} ^ {(1)}
$$
$$
\widetilde {\boldsymbol {\mu}} ^ {(2)} := \widetilde {\mathbf {y}} ^ {(2, 0)} + \widetilde {\mathbf {y}} ^ {(2, 1)} = \mathbf {y} ^ {(0)} + \mathbf {y} ^ {(1)} - 2 \mathbf {X} \widehat {\mathbf {b}} ^ {(2)}
$$
$$
\begin{array}{l} \widehat {\tau} _ {\overline {{S}}} - \tau_ {\overline {{S}}} = \frac {1}{| \overline {{S}} |} \cdot \left(\sum_ {i \in Z ^ {+} \cap \overline {{S}}} \left(2 \widetilde {y} _ {i} ^ {(1, 1)} - y _ {i} ^ {(1)} + y _ {i} ^ {(0)}\right) - \sum_ {i \in Z ^ {-} \cap \overline {{S}}} \left(2 \widetilde {y} _ {i} ^ {(1, 0)} + y _ {i} ^ {(1)} - y _ {i} ^ {(0)}\right)\right) \\ = \frac {1}{| \overline {{S}} |} \cdot \left(\sum_ {i \in Z ^ {+} \cap \overline {{S}}} \left(y _ {i} ^ {(1)} + y _ {i} ^ {(0)} - 2 \mathbf {x} _ {i} ^ {\top} \widehat {\mathbf {b}} ^ {(1)}\right) - \sum_ {i \in Z ^ {-} \cap \overline {{S}}} \left(y _ {i} ^ {(1)} + y _ {i} ^ {(0)} - 2 \mathbf {x} _ {i} ^ {\top} \widehat {\mathbf {b}} ^ {(1)}\right)\right) \\ = \frac {1}{| \overline {{S}} |} \cdot (\widetilde {\boldsymbol {\mu}} _ {\overline {{S}}} ^ {(1)}) ^ {\top} \mathbf {z} _ {\overline {{S}}}. \\ \end{array}
$$
$$
\widehat {\tau} _ {S} - \tau_ {S} = \frac {1}{| S |} \cdot \left(\widetilde {\boldsymbol {\mu}} _ {S} ^ {(2)}\right) ^ {\top} \mathbf {z} _ {S}
$$
$$
\begin{array}{l} \frac {| S | \cdot | \bar {S} |}{n ^ {2}} \cdot \mathbb {E} _ {\mathbf {z}} \left[ \left(\widehat {\tau_ {S}} - \tau_ {\bar {S}}\right) \cdot \left(\widehat {\tau} _ {S} - \tau_ {S}\right) \right] = \frac {1}{n ^ {2}} \cdot \mathbb {E} _ {\mathbf {z}} \left[ \mathbf {z} _ {\bar {S}} ^ {\top} \widetilde {\boldsymbol {\mu}} _ {\bar {S}} ^ {(1)} \left(\widetilde {\boldsymbol {\mu}} _ {S} ^ {(2)}\right) ^ {\top} \mathbf {z} _ {S} \right] \\ \leq \frac {1}{n ^ {2}} \cdot \left(\mathbb {E} _ {\mathbf {z}} \left[ \mathbf {z} _ {S} ^ {\top} \widetilde {\boldsymbol {\mu}} _ {S} ^ {(1)} (\widetilde {\boldsymbol {\mu}} _ {S} ^ {(1)}) ^ {\top} \mathbf {z} _ {\overline {{S}}} \right] \cdot \mathbb {E} _ {\mathbf {z}} \left[ \mathbf {z} _ {S} ^ {\top} \widetilde {\boldsymbol {\mu}} _ {S} ^ {(2)} (\widetilde {\boldsymbol {\mu}} _ {S} ^ {(2)}) ^ {\top} \mathbf {z} _ {S} \right]\right) ^ {1 / 2}, \\ \end{array}
$$
where the inequality follows from Cauchy-Schwarz inequality. Moreover
$$
\frac {| \bar {S} | ^ {2}}{n ^ {2}} \cdot \mathbb {E} _ {\mathbf {z}} \left[ (\widehat {\tau_ {\bar {S}}} - \tau_ {\bar {S}}) ^ {2} \right] = \frac {1}{n ^ {2}} \cdot \mathbb {E} _ {\mathbf {z}} \left[ \mathbf {z} _ {\bar {S}} ^ {\top} \widetilde {\boldsymbol {\mu}} _ {\bar {S}} ^ {(1)} (\widetilde {\boldsymbol {\mu}} _ {\bar {S}} ^ {(1)}) ^ {\top} \mathbf {z} _ {\bar {S}} \right], \text {a n d}
$$
$$
\frac {| S | ^ {2}}{n ^ {2}} \cdot \mathbb {E} _ {\mathbf {z}} \left[ (\widehat {\tau} _ {S} - \tau_ {S}) ^ {2} \right] = \frac {1}{n ^ {2}} \cdot \mathbb {E} _ {\mathbf {z}} \left[ \mathbf {z} _ {S} ^ {\top} \widetilde {\boldsymbol {\mu}} _ {S} ^ {(2)} (\widetilde {\boldsymbol {\mu}} _ {S} ^ {(2)}) ^ {\top} \mathbf {z} _ {S} \right]
$$
Therefore we only need to bound  $\mathbb{E}_{\mathbf{z}}\left[\mathbf{z}_S^\top \widetilde{\boldsymbol{\mu}}_S^{(2)}(\widetilde{\boldsymbol{\mu}}_S^{(2)})^\top \mathbf{z}_S\right]$  and  $\mathbb{E}_{\mathbf{z}}\left[\mathbf{z}_{\overline{S}}^{\top}\widetilde{\boldsymbol{\mu}}_{\overline{S}}^{(1)}(\widetilde{\boldsymbol{\mu}}_{\overline{S}}^{(1)})^\top \mathbf{z}_{\overline{S}}\right]$ . Since  $S$  and  $\overline{S}$  are disjoint,  $\widetilde{\boldsymbol{\mu}}_{\overline{S}}^{(1)}$  is independent of  $\mathbf{z}_{\overline{S}}$ . Moreover, the entries of  $\mathbf{z}$  are independent from each other. Therefore, since  $2\mathbf{y} = \boldsymbol {\mu} + \mathbf{z}\odot \mathbf{t}$ ,
$$
\begin{array}{l} \mathbb {E} _ {\mathbf {z}} \left[ \mathbf {z} _ {\bar {S}} ^ {\top} \tilde {\boldsymbol {\mu}} _ {\bar {S}} ^ {(1)} (\tilde {\boldsymbol {\mu}} _ {\bar {S}} ^ {(1)}) ^ {\top} \mathbf {z} _ {\bar {S}} \right] = \mathbb {E} _ {\mathbf {z}} \left[ \sum_ {i \in \bar {S}} (\tilde {\boldsymbol {\mu}} _ {i} ^ {(1)}) ^ {2} \right] \\ = \mathbb {E} _ {\mathbf {z}} \left[ \left\| (\boldsymbol {\mu} - \mathbf {X} (\mathbf {X} ^ {\top} \mathbf {S} \mathbf {S} \mathbf {X} + \lambda \cdot \mathbf {I}) ^ {- 1} \mathbf {X} ^ {\top} \mathbf {S} \mathbf {S} (\boldsymbol {\mu} + \mathbf {z} \odot \mathbf {t})) _ {\overline {{S}}} \right\| _ {2} ^ {2} \right], \\ \end{array}
$$
where  $\lambda = 100\zeta^2\log (n)$ . We have
$$
\begin{array}{l} \mathbb {E} _ {\mathbf {z}} \left[ \left\| (\boldsymbol {\mu} - \mathbf {X} (\mathbf {X} ^ {\top} \mathbf {S} \mathbf {S} \mathbf {X} + \lambda \cdot \mathbf {I}) ^ {- 1} \mathbf {X} ^ {\top} \mathbf {S} \mathbf {S} (\boldsymbol {\mu} + \mathbf {z} \odot \mathbf {t})) _ {\overline {{S}}} \right\| _ {2} ^ {2} \right] \\ \leq 2 \mathbb {E} _ {\mathbf {z}} \left[ \left\| (\boldsymbol {\mu} - \mathbf {X} (\mathbf {X} ^ {\top} \mathbf {S} \mathbf {S} \mathbf {X} + \lambda \cdot \mathbf {I}) ^ {- 1} \mathbf {X} ^ {\top} \mathbf {S} \mathbf {S} \boldsymbol {\mu}) _ {\overline {{S}}} \right\| _ {2} ^ {2} \right] \\ + 2 \mathbb {E} _ {\mathbf {z}} \left[ \left\| (\mathbf {X} (\mathbf {X} ^ {\top} \mathbf {S} \mathbf {S} \mathbf {X} + \lambda \cdot \mathbf {I}) ^ {- 1} \mathbf {X} ^ {\top} \mathbf {S} \mathbf {S} \mathbf {T} \mathbf {z}) _ {\overline {{S}}} \right\| _ {2} ^ {2} \right]. \\ \end{array}
$$
Note that due to guarantees of leverage score sampling,
$$
\begin{array}{l} \left\| (\boldsymbol {\mu} - \mathbf {X} (\mathbf {X} ^ {\top} \mathbf {S} \mathbf {S} \mathbf {X} + \lambda \cdot \mathbf {I}) ^ {- 1} \mathbf {X} ^ {\top} \mathbf {S} \mathbf {S} \boldsymbol {\mu}) _ {\overline {{S}}} \right\| _ {2} ^ {2} \leq \left\| \boldsymbol {\mu} - \mathbf {X} (\mathbf {X} ^ {\top} \mathbf {S} \mathbf {S} \mathbf {X} + \lambda \cdot \mathbf {I}) ^ {- 1} \mathbf {X} ^ {\top} \mathbf {S} \mathbf {S} \boldsymbol {\mu} \right\| _ {2} ^ {2} \\ \leq (1 + \epsilon) \min  _ {\mathbf {b}} \left(\| \mathbf {X} \mathbf {b} - \boldsymbol {\mu} \| + \lambda \| \mathbf {b} \| _ {2} ^ {2}\right) \\ \end{array}
$$
$$
\mathbb {E} _ {\mathbf {z}} \left[ \| (\mathbf {X} (\mathbf {X} ^ {\top} \mathbf {S} \mathbf {S} \mathbf {X} + \lambda \cdot \mathbf {I}) ^ {- 1} \mathbf {X} ^ {\top} \mathbf {S} \mathbf {S} \mathbf {T} \mathbf {z}) _ {\overline {{S}}} \| _ {2} ^ {2} \right] \leq \mathbb {E} _ {\mathbf {z}} \left[ \| \mathbf {X} (\mathbf {X} ^ {\top} \mathbf {S} \mathbf {S} \mathbf {X} + \lambda \cdot \mathbf {I}) ^ {- 1} \mathbf {X} ^ {\top} \mathbf {S} \mathbf {S} \mathbf {T} \mathbf {z} \| _ {2} ^ {2} \right].
$$
$$
\mathbf {X} ^ {\top} \mathbf {X} \preceq (1 + \epsilon) \mathbf {X} ^ {\top} \mathbf {S S X} + \lambda \cdot \mathbf {I},
$$
$$
\begin{array}{l} \mathbb {E} _ {\mathbf {z}} \left[ \left\| \mathbf {X} \left(\mathbf {X} ^ {\top} \mathbf {S S X} + \lambda \cdot \mathbf {I}\right) ^ {- 1} \mathbf {X} ^ {\top} \mathbf {S S T z} \right\| _ {2} ^ {2} \right] \leq \mathbb {E} _ {\mathbf {z}} \left[ \mathbf {z} ^ {\top} \mathbf {T S S X} \left(\mathbf {X} ^ {\top} \mathbf {S S X} + \lambda \cdot \mathbf {I}\right) ^ {- 1} \mathbf {X} ^ {\top} \mathbf {S S T z} \right] \\ = \sum_ {i \in S} t _ {i} ^ {2} \mathbf {S} _ {i i} ^ {2} (\mathbf {S X} (\mathbf {X} ^ {\top} \mathbf {S S X} + \lambda \cdot \mathbf {I}) ^ {- 1} \mathbf {X} ^ {\top} \mathbf {S}) _ {i i} ^ {2} \\ \leq 4 \left\| \mathbf {t} \right\| _ {\infty} ^ {2} \cdot d. \\ \end{array}
$$
$$
\mathbb {E} _ {\mathbf {z}} \left[ \mathbf {z} _ {\overline {{S}}} ^ {\top} \widetilde {\boldsymbol {\mu}} _ {\overline {{S}}} ^ {(1)} \left(\widetilde {\boldsymbol {\mu}} _ {\overline {{S}}} ^ {(1)}\right) ^ {\top} \mathbf {z} _ {\overline {{S}}} \right] \leq 8 \cdot \| \mathbf {t} \| _ {\infty} ^ {2} \cdot d + 2 \cdot (1 + \epsilon) \min  _ {\mathbf {b}} \left(\| \mathbf {X b} - \boldsymbol {\mu} \| + \lambda \| \mathbf {b} \| _ {2} ^ {2}\right)
$$
With a similar argument, one can show that
$$
\mathbb {E} _ {\mathbf {z}} \left[ \mathbf {z} _ {S} ^ {\top} \widetilde {\boldsymbol {\mu}} _ {S} ^ {(2)} (\widetilde {\boldsymbol {\mu}} _ {S} ^ {(2)}) ^ {\top} \mathbf {z} _ {S} \right] \leq 8 \cdot \| \mathbf {t} \| _ {\infty} ^ {2} \cdot d + 2 \cdot (1 + \epsilon) \min  _ {\mathbf {b}} \left(\| \mathbf {X b} - \boldsymbol {\mu} \| + \lambda \| \mathbf {b} \| _ {2} ^ {2}\right).
$$
$$
\mathbb {E} _ {\mathbf {z}} \left[ \mathbb {E} _ {S} \left[ (\widehat {\tau} - \tau) ^ {2} \right] \right] \leq \frac {3 2 d}{n ^ {2}} \cdot \left\| \mathbf {y} ^ {(1)} - \mathbf {y} ^ {(0)} \right\| _ {\infty} ^ {2} + \frac {8 (1 + \epsilon)}{n ^ {2}} \min  _ {\mathbf {b}} \left(\left\| \mathbf {X b} - \boldsymbol {\mu} \right\| _ {2} ^ {2} + 1 0 0 \log (n) \zeta^ {2} \cdot \left\| \mathbf {b} \right\| _ {2} ^ {2}\right).
$$
We can get a high probability bound for the above result using the Hanson-Wright inequality [31]. Adjusting the vectors  $\mathbf{y}^{(1)}$  and  $\mathbf{y}^{(0)}$  using  $\mathbf{Xb}^*$  as  $\widetilde{\mathbf{y}}^{(0)} = \mathbf{y}^{(0)} - \mathbf{X}\widehat{\mathbf{b}}$  and  $\widetilde{\mathbf{y}}^{(1)} = \mathbf{y}^{(1)} - \mathbf{X}\widehat{\mathbf{b}}$ , and using the Horvitz-Thompson estimator on the adjusted outcomes  $\widetilde{\mathbf{y}}^{(0)}$  and  $\widetilde{\mathbf{y}}^{(1)}$ , we achieve the following result.
C.1 Unconditional Bound for ATE
The only reason the  $1 - \delta$  probability appears in our bound for the variance of ATE estimation (Theorem 2) is that leverage score sampling only gives a high probability guarantee for the regression error. Moreover, our variance only depends on the regression error, and the estimator is always unbiased, even if the leverage score sampling does not give the  $(1 + \epsilon)$  approximation guarantee on the regression. This "bad event" happens with probability at most  $\delta$ . Therefore, it is enough to bound the regression error when the bad event happens. Here we present such an approach. However, to keep our argument simple, we just work with a single solution vector. To actually apply this to our ATE estimation approach, one needs to consider both solution vectors for groups  $S$  and  $\overline{S}$  and consider the bad event for both. Such an argument follows similar to the following and provides the same result (up to constant factors). This approach essentially can make all of our results unconditional (i.e., remove the  $1 - \delta$  probability) at the expense of increasing the variance/err bounds by a constant factor and adding an additive term proportional to  $\delta$ .
If  $\mathbf{X}^{\top}\mathbf{S}^{\top}\mathbf{S}\mathbf{X}$  is not a spectral approximation of  $\mathbf{X}^{\top}\mathbf{X}$ , i.e.,  $(1 - \epsilon)\mathbf{X}^{\top}\mathbf{X} \preceq \mathbf{X}^{\top}\mathbf{S}^{\top}\mathbf{S}\mathbf{X} \preceq (1 + \epsilon)\mathbf{X}^{\top}\mathbf{X}$  does not hold (this can be checked since we have access to matrix  $\mathbf{X}$  and only happens with probability at most  $\delta$ ), we just set  $\mathbf{b} = 0$ . In this case,  $\| \mathbf{X}\mathbf{b} - \boldsymbol{\mu}\|_2^2 + \lambda \| \mathbf{b}\|_2^2 = \| \boldsymbol{\mu}\|_2^2$ . Otherwise (i.e.,  $(1 - \epsilon)\mathbf{X}^{\top}\mathbf{X} \preceq \mathbf{X}^{\top}\mathbf{S}^{\top}\mathbf{S}\mathbf{X} \preceq (1 + \epsilon)\mathbf{X}^{\top}\mathbf{X})$  let  $\widehat{\mathbf{X}} \in \mathbb{R}^{(n + d)\times d}$  be the matrix  $\mathbf{X}$  concatenated with factor  $\lambda$  of the identity matrix and  $\widehat{\boldsymbol{\mu}} \in \mathbb{R}^{n + d}$  be the vector  $\boldsymbol{\mu}$  concatenated with a zero vector and
$$
\mathbf {b} ^ {*} := \underset {\mathbf {b}} {\arg \min } \| \widehat {\mathbf {X}} \mathbf {b} - \widehat {\boldsymbol {\mu}} \| _ {2} ^ {2} = \underset {\mathbf {b}} {\arg \min } \| \mathbf {X} \mathbf {b} - \boldsymbol {\mu} \| _ {2} ^ {2} + \lambda \| \mathbf {b} \| _ {2} ^ {2}.
$$
Let  $\widetilde{\mathbf{X}}\in \mathbb{R}^{(n + d)\times d}$  be the matrix  $\mathbf{S}\mathbf{X}$  concatenated with the factor  $\lambda$  of the identity matrix and  $\widetilde{\pmb{\mu}}\in \mathbb{R}^{(n + d)}$  be the vector  $\mathbf{S}\pmb{\mu}$  concatenated with a zero vector and
$$
\widetilde {\mathbf {b}} := \underset {\mathbf {b}} {\arg \min} \| \widetilde {\mathbf {X}} \mathbf {b} - \widetilde {\boldsymbol {\mu}} \| _ {2} ^ {2} = \underset {\mathbf {b}} {\arg \min} \| \mathbf {S X b} - \mathbf {S} \boldsymbol {\mu} \| _ {2} ^ {2} + \lambda \| \mathbf {b} \| _ {2} ^ {2}.
$$
Then since  $\widetilde{\mathbf{X}}$  is a spectral approximation of  $\widehat{\mathbf{X}}$  (i.e.,  $(1 - \epsilon)\widehat{\mathbf{X}}^{\top}\widehat{\mathbf{X}}\preceq \widetilde{\mathbf{X}}^{\top}\widetilde{\mathbf{X}}\preceq (1 + \epsilon)\widehat{\mathbf{X}}^{\top}\widehat{\mathbf{X}}$ ),
$$
\left\| \widehat {\mathbf {X}} \widetilde {\mathbf {b}} - \widehat {\boldsymbol {\mu}} \right\| _ {2} \leq \left\| \widehat {\mathbf {X}} \mathbf {b} ^ {*} - \widehat {\boldsymbol {\mu}} \right\| _ {2} + \left\| \widehat {\mathbf {X}} \widetilde {\mathbf {b}} - \widehat {\mathbf {X}} \mathbf {b} ^ {*} \right\| _ {2} \leq \left\| \widehat {\mathbf {X}} \mathbf {b} ^ {*} - \widehat {\boldsymbol {\mu}} \right\| _ {2} + \frac {1}{1 - \epsilon} \left\| \widetilde {\mathbf {X}} \widetilde {\mathbf {b}} - \widetilde {\mathbf {X}} \mathbf {b} ^ {*} \right\| _ {2}, \tag {13}
$$
where the first inequality follows by triangle inequality. Moreover, by triangle inequality and optimality of  $\tilde{\mathbf{b}}$
$$
\left\| \widetilde {\mathbf {X}} \widetilde {\mathbf {b}} - \widetilde {\mathbf {X}} \mathbf {b} ^ {*} \right\| _ {2} \leq \left\| \widetilde {\mathbf {X}} \widetilde {\mathbf {b}} - \widetilde {\boldsymbol {\mu}} \right\| _ {2} + \left\| \widetilde {\mathbf {X}} \mathbf {b} ^ {*} - \widetilde {\boldsymbol {\mu}} \right\| _ {2} \leq 2 \cdot \left\| \widetilde {\mathbf {X}} \mathbf {b} ^ {*} - \widetilde {\boldsymbol {\mu}} \right\| _ {2}. \tag {14}
$$