Where to Experiment? Site Selection Under Distribution Shift via Optimal Transport and Wasserstein DRO

TL;DR

This paper introduces a novel approach for selecting experimental sites under distribution shift using optimal transport and Wasserstein DRO, aiming to minimize estimation errors and improve robustness in experimental design.

Contribution

It formulates site selection as an optimal transport problem, develops theoretical bounds for estimation errors, and proposes a Wasserstein DRO-based robust site selection method with a data-driven uncertainty radius.

Findings

Methods outperform random and stratified sampling when covariates are prognostic.

Approach is effective for moderate-to-large problems and when covariates inform treatment effects.

Simulation and real data reanalysis validate the robustness and efficiency of the proposed methods.

Abstract

How should researchers select experimental sites when the deployment population differs from observed data? I formulate the problem of experimental site selection as an optimal transport problem, developing methods to minimize downstream estimation error by choosing sites that minimize the Wasserstein distance between population and sample covariate distributions. I develop new theoretical upper bounds on PATE and CATE estimation errors, and show that these different objectives lead to different site selection strategies. I extend this approach by using Wasserstein Distributionally Robust Optimization to develop a site selection procedure robust to adversarial perturbations of covariate information: a specific model of distribution shift. I also propose a novel data-driven procedure for selecting the uncertainty radius the Wasserstein DRO problem, which allows the user to benchmark robustness levels against observed variation in their data. Simulation evidence, and a reanalysis of a randomized microcredit experiment in Morocco (Cr\'epon et al.), show that these methods outperform random and stratified sampling of sites when covariates have prognostic R-squared > .5, and alternative optimization methods i) for moderate-to-large size problem instances ii) when covariates are moderately informative about treatment effects, and iii) under induced distribution shift.