SlowGuess/NeurIPS_2023_OCR huggingface dataset:
$$
\| \phi (z) \| _ {\mathcal {H}} \leq \epsilon^ {- 1} \| k (\cdot , Y) \| _ {\mathcal {H}} + \epsilon^ {- 1} \max  \left(\| \beta_ {1} (x) \| _ {\mathcal {H}}, \| \beta_ {0} (x) \| _ {\mathcal {H}}\right) + \| \beta_ {1} (x) \| _ {\mathcal {H}} + \| \beta_ {0} (x) \| _ {\mathcal {H}},
$$
which is upper bounded by  $2(\epsilon^{-1} + 1)M$ . Consequently,  $\mathbb{E}\left[\| \phi (Z)\|_{\mathcal{H}}^4\right]^2\leq \left[2(\epsilon^{-1} + 1)M\right]^8$
Furthermore, conditions (iv), (vi), (vii) and (3) deal with the proper behaviour of the AIPW estimator; they are standard in the causal inference scenario. Condition (iv) is equivalent to two-fold cross-fitting i.e., training  $\hat{\phi}^{(r)}$  on only half of the data and evaluating such an estimator on the remaining half. Condition (v) replaces the Donsker class condition from the finite dimensional setting.
We emphasize the importance of double-robustness in the test; normality of the statistic is achieved due to the  $o_{\mathbb{P}}(1 / \sqrt{n})$  rate, which is possible in view of the doubly robust nature of the estimators. We also highlight the fact that IPW estimators of the form  $\hat{\phi}^{(r)}(z) = \left\{\frac{a}{\hat{\pi}(x)} - \frac{1 - a}{1 - \hat{\pi}(x)}\right\} k(\cdot, y)$  can be embedded in the framework considering  $\beta_0, \beta_1 = 0$ . In fact, the doubly robust kernel mean embedding estimator may be viewed as an augmented version of the KTE (which is a kernelized IPW) using regression approaches to kernel mean embeddings (Singh et al., 2020), just as AIPW augments IPW with regression approaches. Furthermore, (3) is always attained when the propensity scores  $\pi$  are known (i.e. experimental setting), given that  $\| \hat{\pi}^{(r)} - \pi \| = 0$ .
Based on the normal asymptotic behaviour of  $T_{h}^{\dagger}$ , we propose to test the null hypothesis  $H_0: P_{Y_1^*} = P_{Y_0^*}$  given the p-value  $p = 1 - \Phi (T_h^\dagger)$ , where  $\Phi$  is the cdf of a standard normal. For an  $\alpha$ -level test, the test rejects the null if  $p \leq \alpha$ . We consider a one-sided test, rather than studying the two-sided p-value  $1 - \Phi (|T_h^\dagger|)$ , given that positive values of  $T_{h}^{\dagger}$  are expected for  $\mathbb{E}[\phi(Z)] \neq 0$ . The next algorithm illustrates the full procedure of the test, which we call AIPW-xKTE (Augmented Inverse Propensity Weighted cross Kernel Treatment Effect).
Algorithm 1 AIPW-xKTE
1: input Data  $\mathcal{D} = (X_i, A_i, Y_i)_{i=1}^{2n}$ .
2: output The p-value of the test.
3: Choose kernel  $k$  and estimators  $\hat{\beta}_a^{(r)}$ ,  $\hat{\pi}^{(r)}$  for  $r \in \{1, 2\}$ .
4: Split data in two sets  $\mathcal{D}_1 = (X_i, A_i, Y_i)_{i=1}^n$ ,  $\mathcal{D}_2 = (X_i, A_i, Y_i)_{i=n+1}^{2n}$ .
5: If  $\hat{\beta}_a^{(r)},\hat{\pi}^{(r)}$  are such that condition (v) from Theorem 4.1 is attained, train them on  $\mathcal{D}$ . Otherwise, train them on  $\mathcal{D}_{1 - r}$ .
6: Define  $\hat{\phi}^{(r)}(z) = \{\frac{a}{\hat{\pi}^{(r)}(x)} -\frac{1 - a}{1 - \hat{\pi}^{(r)}(x)}\} \{k(\cdot ,y) - \hat{\beta}_a^{(r)}(x)\} +\hat{\beta}_1^{(r)}(x) - \hat{\beta}_0^{(r)}(x).$
7: Define  $f_{h}^{\dagger}(Z_{i}) = \frac{1}{n}\sum_{j = n + 1}^{2n}\langle \hat{\phi}^{(1)}(Z_{i}),\hat{\phi}^{(2)}(Z_{j})\rangle$  for  $i = 1,\dots ,n$
8: Calculate  $T_{h}^{\dagger} \coloneqq \frac{\sqrt{n}\bar{f}_{h}^{\dagger}}{S_{h}^{\dagger}}$ , where  $\bar{f}_h^\dagger = \frac{1}{n}\sum_{i = 1}^{n}f_h^\dagger (Z_i)$  and  $S_{h}^{\dagger} = \sqrt{\frac{1}{n}\sum_{i = 1}^{n}(f_{h}^{\dagger}(Z_{i}) - \bar{f}_{h}^{\dagger})^{2}}$ .
9: return p-value  $p = 1 - \Phi (T_h^\dagger)$
Note that the proposed statistic is, at heart, a two sample test (with a nontrivial causal twist); in contrast to Shekhar et al. (2022), the two samples are not independent and are potentially confounded.
Extensive literature focuses on designing estimator  $\hat{\pi}^{(r)}$ , logistic regression being the most common choice. At this time, not so many choices exist for estimators  $\hat{\beta}_a^{(r)}$ , given that it involves a regression task in a Hilbert space. Conditional mean embeddings are the most popular regressor, although other choices exist (Cevid et al., 2022).
Note that we have motivated the proposed procedure for testing distributional treatment effects with characteristic kernels. However, the actual null hypothesis being tested is  $H_0: \mathbb{E}[k(\cdot, Y_1^*)] = \mathbb{E}[k(\cdot, Y_0^*)]$ . If the kernel chosen is not characteristic, the test would continue to be valid for  $H_0$ , although it would not be valid to test equality between  $\mathbb{P}_{Y_0^*}$  and  $\mathbb{P}_{Y_1^*}$ . For instance, AIPW-xKTE with a linear kernel could be used to test equality in means of counterfactuals.
Furthermore, the proposed test is permutation-free, as the statistic  $T_h^\dagger$  ought to be computed only once. This permutation-free nature is crucial, as it avoids the repeated estimation of  $\hat{\pi}^{(r)}, \hat{\beta}_a^{(r)}$ . For instance, conditional mean embeddings involve the inversion of a matrix, which scales at least at  $O(n^{\omega})$ , with practical values being  $\omega = 2.87$  by Strassen's algorithm (Strassen et al., 1969). Calculating the conditional mean embedding for every permutation would imply  $O(Bn^{\omega})$ , where  $B$  is the number of permutations. Furthermore, regressors for mean embeddings of different nature might involve a higher complexity, hence avoiding permutations becomes even more important in the approach.
If the actual embedding  $\phi$  was known, the power of AIPW-xKTE could not be improved beyond a constant factor (by minimax optimality of cross U-statistics in  $L^2$  distance). Further, every procedure will suffer from the error in estimation of  $\phi$ . This means that we are potentially incurring in a loss of power by avoiding a permutation-based approach, however such a loss is controlled by a small factor. Nonetheless, this potential loss is inherited from splitting the data in our estimator (only half of the data is used on each side of the inner product). We highlight that sample splitting is needed when using flexible doubly-robust estimators, hence we expect no loss in power compared to other potential doubly-robust approaches in that case.
In this section, we explore the empirical calibration and power of the proposed test AIPW-xKTE. For this, we assume that we observe  $(x_{i},a_{i},y_{i})_{i = 1}^{n}\sim (X,A,Y)$  and that (causal inference assumptions) consistency, no unmeasured confounding, and overlap hold. Both synthetic data and real data are evaluated. All the tests are considered at a 0.05 level. For an exhaustive description of the simulations and outcomes, including additional experiments, we direct the reader to Appendix B.
Synthetic data. All data (covariates, treatments and responses) are artificially generated. We define four scenarios:
- Scenario I: There is no treatment effect; thus,  $\mathbb{P}_{Y_0^*} = \mathbb{P}_{Y_1^*}$ .
- Scenario II: There exists a treatment effect that only affects the means of  $\mathbb{P}_{Y_0^*},\mathbb{P}_{Y_1^*}$
- Scenario III and Scenario IV: There exists a treatment effect that does not affect the means but only affects the higher moments of  $\mathbb{P}_{Y_0^*}$  and  $\mathbb{P}_{Y_1^*}$ , differently for each scenario.
For all four scenarios, we consider the usual observational study setting, where the propensity scores  $\pi(X)$  are treated as unknown and hence they must be estimated. We define the proposed AIPW-xKTE test with the mean embedding regressions fitted as conditional mean embeddings and the propensity scores estimated by logistic regression.
We first study the empirical calibration of AIPW-xKTE and the Gaussian behaviour of  $T_{h}^{\dagger}$  under the null. Figure 1 exhibits the performance of AIPW-xKTE in Scenario I. Both a standard normal behaviour and proper calibration are empirically attained in the simulations.
Due to the fact that the KTE (Muandet et al., 2021) may not be used in the observational setting, where the propensity scores are not known, there is no natural benchmark for the proposed test. In particular, we were unable to control the type-1 error of the test presented in Fawkes et al. (2022), and hence omitted from our simulations. Consequently, we compare the power of the proposed AIPW-xKTE and IPW-xKTE with respect to three methods that are widely used while conducting inference on the average treatment effect: Causal Forests (Wager and Athey, 2018), Bayesian Additive Regression Trees (BART) (Hahn et al., 2020), and a linear regression based AIPW estimator (Baseline-AIPW).
Figure 1: Illustration of 500 simulations of the AIPW-xKTE under the null: (A) Histogram of AIPW-xKTE alongside the pdf of a standard normal for  $n = 500$ , (B) Normal Q-Q plot of AIPW-xKTE for  $n = 500$ , (C) Empirical size of AIPW-xKTE against different sample sizes. The figures show the Gaussian behaviour of the statistic under the null, which leads to a well calibrated test.
Figure 2: True positive rates of 500 simulations of the tests in Scenarios II, III, and IV. AIPW-xKTE shows notable true positive rates in every scenario, unlike competitors.
Figure 2 exhibits the performance of such tests in Scenario II, Scenario III, and Scenario IV. The three methods dominate AIPW-xKTE in Scenario II, where there exists a mean shift in counterfactuals. However, and as expected, such methods show no power if the distributions differ but have equal means. In contrast, AIPW-xKTE detects distributional changes beyond the mean, exhibiting power in all scenarios.
Remark: While this work focuses on the observational setting, where double robustness is crucial, we highlight that the proposed AIPW-xKTE test may also be used in experiments (where propensity scores are known) for computational gains. The proposed test avoids permutations, which makes it more computationally efficient than the KTE (Muandet et al., 2021). We refer the reader to Appendix A for a comparison between the proposed AIPW-xKTE and the KTE.
Real data. We use data obtained from the Infant Health and Development Program (IHDP) and compiled by Hill (2011), in which the covariates come from a randomized experiment studying the effects of specialist home visits on cognitive test scores. This data has seen extensive use in causal inference (Johansson et al., 2016; Louizos et al., 2017; Shalit et al., 2017). We work with 18 variables of the covariate set and unknown propensity scores. We highlight that the propensity score model is likely misspecified in this real life scenario.
We consider six scenarios with the IHDP data. For Scenarios I, II, III and IV, we generate the response variables similarly to the previous experiments. In Scenario V, we take the IQ test (Stanford Binet) score measured at the end of the intervention (age 3) as our response variable. For Scenario VI, we calculate the average treatment effect in Scenario V using Causal Forests (obtaining 0.003 i.e. a positive shift) and subtract it from the IQ test score of those who are treated, thus obtaining a distribution of the IQ test scores with zero average treatment effect.
For each of the scenarios, we consider the tests AIPW-xKTE, Causal Forests, BART and Baseline-AIPW on 500 bootstrapped subsets to estimate the rejection rates. All of them showed expected levels of rejection under the null i.e. Scenario I. The results for the remaining scenarios can be found in Table 1. We note that the performance of the tests on Scenarios I, II, III, and IV is similar to the analogous scenarios on synthetic data. While AIPW-xKTE exhibits a loss in power with respect to
Table 1: True positive rates (± std) for the different scenarios and tests using the IHDP data. While AIPW-xKTE shows less power in Scenario II, it outperforms its competitors in Scenarios III and IV.
Test	Scenario
II	III	IV	V	VI
AIPW-xKTE	0.44 ± 0.05	0.34 ± 0.05	0.53 ± 0.05	0.99 ± 0.01	0.03 ± 0.02
Baseline-AIPW	0.95 ± 0.02	0.00 ± 0.00	0.00 ± 0.00	1.00 ± 0.00	0.00 ± 0.00
BART	0.77 ± 0.04	0.15 ± 0.04	0.24 ± 0.04	0.10 ± 0.03	0.04 ± 0.02
CausalForest	1.00 ± 0.00	0.04 ± 0.02	0.05 ± 0.02	1.00 ± 0.00	0.00 ± 0.00
the other tests when the average treatment effect is non zero, it detects distributional changes beyond the mean.
Besides BART, all of the tests reject the null in almost every simulation of Scenario V. As expected, Causal Forest, BART and Baseline-AIPW barely reject the null in Scenario VI, which considers the data of Scenario V with zero average treatment effect. Interestingly, the proposed distributional test has a rejection rate below 0.05, which supports the fact that specialist home visits have no effect on the distribution of cognitive test scores beyond an increase in the mean.
We have developed a computationally efficient kernel-based test for distributional treatment effects. It is, to our knowledge, the first kernel-based test for distributional treatment effects that allows for a doubly-robust approach with provably valid type-I error. Furthermore, it does not suffer from the computational costs inherent to permutation-based tests. The proposed test empirically proves valid in the observational setting, where its predecessor KTE may not be used: the test is well calibrated and shows power in a variety of scenarios. Procedures designed to test for average treatment effects only outperform the proposed test if there is a mean shift between counterfactuals.
There are several possible avenues for future work. We highlighted that our procedure holds if consistency, no unmeasured confounding, and overlap hold. However, this may often not be the case in observational studies. Generalizing the work to other causal inference frameworks, for example by considering instrumental variables, would be of interest. Exploring the extension of this work to test for conditional treatment effects could also be a natural direction to follow. Lastly, we expect state-of-the-art regressors for kernel mean embeddings, such as distributional random forests (Cecvid et al., 2022), to find outstanding use in our estimator. We envisage that this may motivate the development of flexible mean embedding estimators designed for complex data, such as image or text.
The project that gave rise to these results received the support of a fellowship from 'la Caixa' Foundation (ID 100010434). The fellowship code is LCF/BQ/EU22/11930075.
Abadie, A. (2002). Bootstrap tests for distributional treatment effects in instrumental variable models. Journal of the American Statistical Association, 97(457):284-292.
Bosq, D. (2000). Linear processes in function spaces: theory and applications, volume 149. Springer Science & Business Media.
Chernozhukov, V., Fernandez-Val, I., and Melly, B. (2013). Inference on counterfactual distributions. Econometrica, 81(6):2205-2268.
Díaz, I. (2017). Efficient estimation of quantiles in missing data models. Journal of Statistical Planning and Inference, 190:39-51.
Fawkes, J., Hu, R., Evans, R. J., and Sejdinovic, D. (2022). Doubly robust kernel statistics for testing distributional treatment effects even under one-sided overlap. arXiv preprint arXiv:2212.04922.
Funk, M. J., Westreich, D., Wiesen, C., Stürmer, T., Brookhart, M. A., and Davidian, M. (2011). Doubly robust estimation of causal effects. American Journal of Epidemiology, 173(7):761-767.
Gretton, A., Borgwardt, K. M., Rasch, M. J., Scholkopf, B., and Smola, A. (2012). A kernel two-sample test. The Journal of Machine Learning Research, 13(1):723-773.
Hahn, P. R., Murray, J. S., and Carvalho, C. M. (2020). Bayesian regression tree models for causal inference: Regularization, confounding, and heterogeneous effects (with discussion). *Bayesian Analysis*, 15(3):965-1056.
Richardson, A., Hudgens, M. G., Gilbert, P. B., and Fine, J. P. (2014). Nonparametric bounds and sensitivity analysis of treatment effects. Statistical Science: A Review Journal of the Institute of Mathematical Statistics, 29(4):596.
Robins, J., Li, L., Tchetgen, E., van der Vaart, A., et al. (2008). Higher order influence functions and minimax estimation of nonlinear functionals. *Probability and Statistics: Essays in Honor of David A. Freedman*, 2:335-421.
Robins, J. M. and Rotnitzky, A. (2001). Comments on: Inference for semiparametric models: Some questions and an answer. Statistica Sinica.
Rubin, D. B. (2005). Causal inference using potential outcomes: Design, modeling, decisions. Journal of the American Statistical Association, 100(469):322-331.
Shalit, U., Johansson, F. D., and Sontag, D. (2017). Estimating individual treatment effect: generalization bounds and algorithms. In International Conference on Machine Learning, pages 3076-3085. PMLR.
Shekhar, S., Kim, I., and Ramdas, A. (2022). A permutation-free kernel two-sample test. In Advances in Neural Information Processing Systems.
Shen, S. (2019). Estimation and inference of distributional partial effects: theory and application. Journal of Business & Economic Statistics, 37(1):54-66.
Singh, R., Sahani, M., and Gretton, A. (2019). Kernel instrumental variable regression. Advances in Neural Information Processing Systems, 32.
Singh, R., Xu, L., and Gretton, A. (2020). Kernel methods for causal functions: Dose, heterogeneous, and incremental response curves. arXiv preprint arXiv:2010.04855.
Singh, R., Xu, L., and Gretton, A. (2021). Kernel methods for multistage causal inference: Mediation analysis and dynamic treatment effects. arXiv preprint arXiv:2111.03950.
Song, L., Huang, J., Smola, A., and Fukumizu, K. (2009). Hilbert space embeddings of conditional distributions with applications to dynamical systems. In Proceedings of the 26th Annual International Conference on Machine Learning, pages 961-968.
Strassen, V. et al. (1969). Gaussian elimination is not optimal. Numerische mathematik, 13(4):354-356.
Van der Vaart, A. W. (2000). Asymptotic statistics, volume 3. Cambridge University Press.
Wager, S. and Athey, S. (2018). Estimation and inference of heterogeneous treatment effects using random forests. Journal of the American Statistical Association, 113(523):1228-1242.
Westling, T. and Carone, M. (2020). A unified study of nonparametric inference for monotone functions. The Annals of Statistics, 48(2):1001.
Zhu, Y., Gultchin, L., Gretton, A., Kusner, M. J., and Silva, R. (2022). Causal inference with treatment measurement error: a nonparametric instrumental variable approach. volume 180 of Proceedings of Machine Learning Research, pages 2414-2424. PMLR.
Cevid, D., Michel, L., Naf, J., Meinshausen, N., and Buhlmann, P. (2022). Distributional random forests: Heterogeneity adjustment and multivariate distributional regression. Journal of Machine Learning Research, 23(333):1-79.
A Comparison between KTE and AIPW-xKTE in the experimental setting
In Section 5, we exhibited the performance of the proposed AIPW-xKTE test in the observational setting. We emphasize that the proposed AIPW-xKTE test is designed to target the observational setting, where the KTE may not be used given that the propensity scores are unknown. In such a setting, the double-robustness of the proposed test is fully exploited. However, the proposed test may also be used in the experimental setting for its computational efficiency: the AIPW-xKTE test avoids permutations, which makes it more computationally efficient than the KTE.
Consequently, we explore in this appendix the performance of the AIPW-xKTE in the experimental setting. We consider a usual experimental design where  $\frac{n}{2}$  units are treated and  $\frac{n}{2}$  units are not treated so that  $\pi(X) = \frac{1}{2}$ . We compare the power of the proposed AIPW-xKTE with respect to the KTE test. Further, we include an IPW version of the proposed test (by taking  $\beta_0 = \beta_1 = 0$ , see comments in Section 4) in order to elucidate the interest in considering the doubly robust version. We refer the reader to Appendix B for an exhaustive description of the experiments.
Figure 4 exhibits the performance of the tests in Scenario II, Scenario III, and Scenario IV. We note that IPW-xKTE has less power than the KTE, which is its analogous permutation-based version. This is due to the controlled loss in power of the cross U-statistic approach following the comments in Section 4. However, we see that the AIPW-xKTE competes with the KTE, despite showing slightly less power. The loss in power inherited by the permutation-free approach is compensated by the gain in power due to using an AIPW estimator. Furthermore, AIPW-xKTE shows a drastic improvement in computational costs with respect to KTE, as exhibited in Table 2. Lastly, we see that AIPW-xKTE clearly outperforms IPW-xKTE, which illustrates the significant benefits of the doubly robust approach.
Table 2: Average times (in milliseconds) for different tests and sample sizes. AIPW-xKTE is almost 20 times faster than KTE for  $n = 350$  with only 100 permutations used.
Test	n
150	250	350
AIPW-xKTE	1.770	2.994	4.940
IPW-xKTE	1.496	2.297	2.331
KTE	16.903	44.483	89.875
Figure 3: Illustration of 500 simulations of the AIPW-xKTE under the null in the experimental setting: (A) Histogram of AIPW-xKTE alongside the pdf of a standard normal for  $n = 500$ , (B) Normal Q-Q plot of AIPW-xKTE for  $n = 500$ , (C) Empirical size of AIPW-xKTE and IPW-xKTE against different sample sizes. The figures show the Gaussian behaviour of the statistic under the null, which leads to a well calibrated test.
Figure 4: True positive rates of 500 simulations of the tests in Scenarios II, III, and IV in the experimental setting. AIPW-xKTE demonstrates comparable power to KTE, and improves appreciably over IPW-xKTE.
We present in this appendix a comprehensive display of the simulations conducted. Subsection B.1 includes an exhaustive description of the experiments introduced in Section 5 and Appendix A. Additional simulations and results may be found in Subsection B.2.
B.1 Exhaustive description of experiments
We assume that we observe  $(x_{i},a_{i},y_{i})_{i = 1}^{n}\sim (X,A,Y)$  and that (causal inference assumptions) consistency, no unmeasured confounding, and overlap hold. All the tests are considered at a 0.05 level. Both synthetic data and real data are evaluated.
Synthetic data. All the data (covariates, treatments and responses) is artificially generated. We define
$$
Y _ {0} ^ {*} = \beta^ {T} X + \epsilon_ {0}, Y _ {1} ^ {*} = \beta^ {T} X + b + \epsilon_ {1}, \tag {4}
$$
such that  $\epsilon_0, \epsilon_1 \sim N(0, 0.5)$  are independent noises,  $X \sim N(0, I_5)$  and  $\beta = [0.1, 0.2, 0.3, 0.4, 0.5]^T$ . We set  $\mathsf{b} = 0$  and  $\mathsf{b} = 2$  for Scenario I and Scenario II respectively. For Scenario III, we set  $b = 2Z - 1$ , where  $Z$  is an independent Bernoulli random variable  $Z \sim \text{Bernoulli}(0.5)$ . In Scenario IV,  $b \sim \text{Uniform}(-2, 2)$ .
In the experimental setting, we consider a usual experimental design where  $\frac{n}{2}$  units are treated and  $\frac{n}{2}$  units are not treated such that  $\pi(X) = \frac{1}{2}$  almost surely. In the observational setting, we define  $\pi(X) = s(\alpha^T X + \alpha_0)$ , where  $s(z) = 1 / (1 + \exp(-z))$  (sigmoid function),  $\alpha_0 = 0.05$  and  $\alpha = [0.05, 0.04, 0.03, 0.02, 0.01]^T$ . In such setting, we estimate  $\hat{\pi}$  using an L2 regularized logistic regression with the regularization term set to 1e-6. For AIPW-xKTE, the mean embedding regressions are fitted as conditional mean embeddings.
In the experimental setting, we consider the tests IPW-xKTE, AIPW-xKTE, and KTE (Muandet et al., 2021). In the observational setting, we consider IPW-xKTE, AIPW-xKTE and permutation-based tests based on Causal Forests (Wager and Athey, 2018), Bayesian Additive Regression Trees (BART) (Hahn et al., 2020), and a linear regression based AIPW estimator (Baseline-AIPW). For the latter three, we recalculate the respective statistics for every permutation and reject the null if the original statistic is above the 0.95 empirical quantile. For all scenarios and settings, we consider 500
Figure 5: Illustration of one simulation of the simulated data in the experimental setting with  $n = 350$  for Scenarios I, II, III, and IV. We display scatter plots of the outcome against the first covariate and histograms of the outcomes, both grouped by treatment.