Statistical Inference with Distribution Shifts

Modern complex datasets often consist of various sub-populations. To develop robust and generalizable methods in the presence of sub-population heterogeneity, it is important to guarantee a uniform learning performance instead of an average one. In many applications, prior information is often available on which sub-population or group the data points belong to. Given the observed groups of data, we develop a min-max-regret (MMR) learning framework for general supervised learning, which targets to minimize the worst-group regret. Motivated from the regret-based decision theoretic framework, the proposed MMR is distinguished from the value-based or risk-based robust learning methods in the existing literature. The regret criterion features several robustness and invariance properties simultaneously. In terms of generalizability, we develop the theoretical guarantee for the worst-case regret over a super-population of the meta data, which incorporates the observed sub-populations, their mixtures, as well as other unseen sub-populations that could be approximated by the observed ones. We demonstrate the effectiveness of our method through extensive simulation studies and an application to kidney transplantation data from hundreds of transplant centers. 

Should I gather cheap, low-quality data or expensive, high-quality data? For example, in the social sciences, running behavioral experiments on Amazon Mechanical Turk is often easier than carefully recruiting participants from the target population. However, data from Amazon Mechanical Turk may suffer from bias issues. This leads to a tradeoff between data quantity and data quality. We formalize this decision problem from a distribution shift perspective, taking into account (a) data quality, (b) data quantity, and (c) problem difficulty. We demonstrate that it is possible to predict the usefulness of data based on summary statistics. More specifically, our proposed notion of data usefulness allows us to predict how much the mean squared error (MSE) of estimation and prediction procedures would improve with additional data from a particular candidate distribution, without having access to individual-level data from the candidate distribution. We illustrate the effectiveness of our approach on both estimation and prediction tasks. 

The instrumental variable (IV) approach is commonly used to infer causal effects in the presence of unmeasured confounding. Existing methods typically aim to estimate mean causal effects, whereas a few focus on quantile treatment effects. This work aims to estimate the entire interventional distribution, which yields classical causal estimands as functionals. We propose a method called Distributional Instrumental Variable (DIV), which leverages generative modeling in a nonlinear IV setting. We establish identifiability of the interventional distribution under general assumptions and illustrate an "under-identified" case where DIV can identify causal effects while two-stage least squares fail. Empirical results show that DIV performs well across a broad range of simulated data, outperforming existing IV approaches in identifiability and estimation error for mean or quantile treatment effects. Furthermore, we apply DIV to an economic dataset to examine the causal relationship between institutional quality and economic development, finding results that align with the original study. We also apply DIV to a single-cell dataset to assess generalizability and stability in predicting gene expression under unseen interventions.

Co-authors: Anastasiia Holovchak (ETH Zurich), Sorawit Saengkyongam (ETH Zurich), Nicolai Meinshausen (ETH Zurich) 

Estimation of heterogeneous treatment effects plays an essential role in various scientific fields, industries, and policy-making settings. Although many existing methods estimate conditional average treatment effects (CATEs) based on the data from a single population, in practice, researchers often have access to the data collected from multiple sites, whose characteristics differ in unknown ways. While a naive pooled model risks bias, the estimates from site-specific models lack external validity. In this paper, we introduce a robust CATE estimation methodology with multi-site data from heterogeneous populations. We propose a minimax-regret framework that learns a generalizable CATE model by minimizing the worst-case regret over a class of distributions, represented as convex combinations of observed distributions across sites. Using robust optimization, the proposed methodology accounts for distribution shifts in both individual covariates and CATEs across sites. We show that the resulting CATE model has an interpretable closed-form solution, expressed as a weighted average of site-specific CATE models. This result allows researchers to leverage flexible CATE estimation methods across each site, which can then be aggregated to produce the final model. Through simulations and a real-world application, we show that the proposed methodology improves the robustness and generalizability of the existing approaches.