Active Adaptive Experimental Design for Treatment Effect Estimation with Covariate Choices

Masahiro Kato, Akihiro Oga, Wataru Komatsubara, and Ryo Inokuchi. 2024. International Conference on Machine Learning (ICML 2024, oral presentation). Proceedings page (PMLR), arXiv:2403.03589. 日本語版

For an experiment that estimates the average treatment effect, we propose a design that adaptively optimizes the treatment assignment probability together with the choice of which covariates the subjects admitted to the experiment should have. Optimizing both pushes the theoretical limit on the estimation variance below what a design that optimizes the assignment probability alone can reach. ICML 2024 selected this study for oral presentation, which went to 1.5% of accepted papers.

The problem

On each round of the experiment, the experimenter picks one subject, assigns a treatment, and observes the outcome immediately. Once the experiment ends, the average treatment effect is estimated from the collected sample. The aim is to make the asymptotic variance of that estimate as small as possible. Earlier work on adaptive experiments optimized the treatment assignment probability, the propensity score, during the experiment. But how subjects are chosen for the experiment should also govern precision, and the question here is how far precision can be pushed once that choice becomes part of the design.

The proposed method

We first treat the covariate distribution and the assignment probability as design variables and derive the efficient covariate density and assignment probability that minimize the semiparametric efficiency bound, the theoretical limit on the estimation variance. We then design an experiment that estimates these two sequentially from the data arriving during the experiment and updates subject selection and treatment assignment accordingly. Finally, we propose an estimator of the average treatment effect from the sample collected after the experiment ends.

Main results and conditions

Optimizing the covariate density together with the assignment probability lowers the semiparametric efficiency bound below what optimizing the assignment probability alone can reach, and the asymptotic variance of the proposed estimator attains this lowered bound. The guarantee is asymptotic, holding as the sample size grows. It also assumes that the experimenter can see a subject's covariates before deciding whether to include that subject, and that the outcome is observed immediately after assignment.

Relation to earlier work

Hahn derived the theoretical limit on the variance of average treatment effect estimation (1998, On the Role of the Propensity Score in Efficient Semiparametric Estimation of Average Treatment Effects). Aiming at that limit, Hahn, Hirano, and Karlan proposed an experimental design that adaptively optimizes the treatment assignment probability (2011, Adaptive Experimental Design Using the Propensity Score). This study is the generalization that widens their design variables to include the covariate distribution.

Where it applies

The situation we have in mind is one where each sample costs a lot and the precision of the treatment effect has to improve within a fixed budget. Testing an advertisement or a policy where the subjects can be selected is a typical case, as is a clinical trial in which enrollment can be designed. For consulting on experimental design in practice, write to mkato-csecon@g.ecc.u-tokyo.ac.jp.

Paper and materials

BibTeX

@inproceedings{active-adaptive-experimental-design,
  author       = {Masahiro Kato and Akihiro Oga and Wataru Komatsubara and Ryo Inokuchi},
  title        = {Active Adaptive Experimental Design for Treatment Effect Estimation with Covariate Choices},
  year         = {2024},
  booktitle    = {International Conference on Machine Learning (ICML, Oral 1.5\%)},
  eprint       = {2403.03589},
  archivePrefix = {arXiv},
  url          = {https://proceedings.mlr.press/v235/kato24a.html},
}