Statistical Inference Under Constrained Selection Bias

Decision-makers in the public and private sectors increasingly seek to use machine learning or statistical models built on top of large-scale datasets in order to inform policy, operational decisions, individualized treatment rules, and more. However, these administrative datasets are typically purely observational, meaning that they are not carefully designed to sample from a true distribution of interest. Consequently, such efforts have been hindered by sampling biases and other distributional shifts between the observed data and the target distribution. Such selection biases have presented severe issues for past algorithmic systems [10, 14], policy analysis [18, 28], epidemiological studies [17, 25], among others.

As a motivating example, decision-makers in public health may wish to estimate the risk factors contributing to adverse outcomes related to COVID-19, e.g., hospitalizations, in order to target preventative interventions. Nevertheless, they usually only have data from people who engaged with some specific services in the past such as insurance claims or patients who visited the health system. This scenario poses the difficulty that some demographic groups will be underrepresented in the data, influenced by some latent factors (for instance, high healthcare costs) and the outcome itself. It is a known phenomenon that individuals with worse healthcare access will seek care primarily when facing severe diseases, thus increasing observed hospitalizations for this group. Hence, creating a problem akin to confounding when the relationship between these two quantities is estimated.

We consider the task of accurately estimating some functional of the ground-truth distribution using samples from an observed, potentially shifted, distribution. For instance, our goal might be to estimate the expectation of a measured covariate or the treatment effect from an intervention. In general, such inference is intractable without assumptions on the relationship between the observed and target distributions. The simplest setting is where we can directly observe some samples from the target distribution. However, we are interested in the scenarios where this is not the case, for example, the scenario introduced in the previous paragraph; given that the data is conditioned among other things on the outcome, it presents the insidious problem that it will never be representative of the population as whole. Absent target-distribution samples, other frameworks such as those related to distributionally robust optimization (DRO) [7, 16, 46, 39] or sensitivity analysis [12, 37, 43] allow users to specify the total magnitude of distribution shift allowed. Such magnitude is usually modeled by imposing a limit on the maximum distance (e.g., KL or $\chi^{2}$ divergence) between the target distribution and the observed one. However, since the target distribution (or samples of it) is unavailable, this assumption cannot be empirically verified, and often cannot even be set in an informed manner. Our goal is to provide provable guarantees without introducing untestable assumptions on the magnitude of shift, by using only observable quantities to set bounds on the possible difference in distributions.

Intuitively, this may be possible because decision makers often have aggregate knowledge about the target distribution. For instance, a policymaker may know, via census data, the distribution of demographic characteristics such as age, race, or income from the population as a whole. Moreover, in the public health setting, serosurveys may provide ground-truth estimates of exposure to infectious diseases in specific locations or population groups [19]. This aggregate information is typically not sufficient by itself for the original task because it fails to measure the key outcome or covariates of interest (e.g., knowing the demographic distribution of a population is by itself unhelpful for estimating a patient’s risk of hospitalization). However, it imposes constraints on the nature of the distribution shift between the observed distribution and the underlying population: any valid shift must respect these known quantities. Concretely, our contributions are as follows:

• •

  We introduce a framework that allows user-specified constraints on distribution shift via known expectations. Our framework incorporates such external information into an optimization program whose value gives valid max/min bounds on a statistic of interest using samples from an observed distribution. As a result, we are able to provide estimates by adding more observables in lieu of assumptions that are not empirically testable (e.g., the KL divergence between the target and sample distributions). We provide conditions under which this optimization problem is convex and hence, efficiently solvable.

• •

  We analyze the statistical properties of estimating these bounds using a sample average approximation for the optimization problem. We show that our estimators are asymptotically normal, allowing us to provide valid confidence intervals.

• •

  We extend our framework to accommodate estimands without a closed form (e.g., a regression coefficient or an estimated model parameter), provide statistical guarantees for this setting, and propose computational approaches to solve the resulting optimization problem.

• •

  We perform experiments on synthetic and semi-synthetic data to test our methods. The results empirically confirm that our framework provides valid bounds for the target estimand and allows effective use of domain knowledge: incorporation of more informative constraints produces tighter bounds, which can be strongly informative about the estimand of interest. All the code we used for our experiments is available at hidden.

• •

  We showcase a real-world case of study of assessing disparities in COVID-19 disease burden using a large-scale insurance claims dataset where our method allows us to produce policy-relevant results under the presence of selection bias.

Let $X$ be a random variable over a space $\mathcal{X}$, with distribution $\mathbb{P}$ in a population of interest. Our goal is to estimate a real-valued functional $f(\mathbb{P})$. One prominent example is estimating the expectation $f(\mathbb{P})=\mathbb{E}_{X\sim\mathbb{P}}[h(X)]$ for some function $h$, but we will consider a range of examples for $f$. If we observe iid samples from $\mathbb{P}$, estimating $f(\mathbb{P})$ is a standard inference task. For instance, we may use the plug-in estimate $f(\hat{\mathbb{P}}_{N})$, (where $\hat{\mathbb{P}}_{N}$ is the empirical distribution) or any number of other strategies. However, we consider a setting where we instead observe samples drawn iid from a different distribution $\mathbb{Q}$. For illustration purposes, recall the example mentioned in the introduction: policymakers are interested in information relative to the population as a whole ($f(\mathbb{P})$) but they only have access to observational data from the hospitals or insurance claims, i.e., a sample from $\mathbb{Q}$.

We assume that both $\mathbb{P}$ and $\mathbb{Q}$ have densities $p$ and $q$, respectively, and that $q(X)>0$ whenever $p(X)>0$. More formally, we require that $\mathbb{P}$ is absolutely continuous with respect to $\mathbb{Q}$ so that the ground truth density ratio $\theta_{0}(X)=\frac{p(X)}{q(X)}$ is well defined. Intuitively, inference on $\mathbb{P}$ is impossible if some portions of the distribution can never be observed. Accordingly, similar overlap assumptions are near-universal in causal inference [24], domain adaptation [9], selection bias [38], distributionally robust optimization [35], etc. We expect this assumption to be satisfied in many application domains because selection bias is not strong enough to make some observations literally impossible. Our running COVID-19 example illustrates this: a patient with any covariates $X$ can in principle seek care (and some such patients do), even if they have worse access than patients with a different set of covariates $X^{\prime}$.

With slight abuse of notation, for any $\mathbb{P}$, we let $\theta\mathbb{P}$ denote a distribution with density $\theta(X)p(X)$ so that $\theta_{0}\mathbb{Q}=\mathbb{P}$. If $\theta_{0}$ were known, estimating $f(\mathbb{P})$ using samples from $\mathbb{Q}$ would be easily accomplished using standard importance sampling methods: in the case of estimating an expectation, we have that $\mathbb{E}_{X\sim\mathbb{P}}[h(X)]=\mathbb{E}_{X\sim\mathbb{Q}}[\theta_{0}(X)h(X)]$. However, our focus is on the setting where $\theta_{0}$ is not known, and we don’t have access to samples from $\mathbb{P}$ from which to estimate it.

Without any information about the relationship between $\mathbb{P}$ and $\mathbb{Q}$ this is a hopeless task. However, in many settings of importance some information is available, in the form of auxiliary functions whose true expectation under $\mathbb{P}$ is known. In policy settings, this may come from census data, or well-design population surveys that estimate the ground-truth prevalence of specific health conditions (c.f. [26, 19, 34]). Concretely, we are able to observe $\mathbb{E}_{\mathbb{P}}[g_{j}(X)]$ for a collection of functions $g_{1}...g_{m}$. Thus, we know that $\theta_{0}$ must satisfy

Let $\Theta$ to be the set of density ratios $\theta$ which satisfy the above constraints, all of which are consistent with our observations. In general, $\Theta$ will not be a singleton and so $f(\mathbb{P})$ will not be point-identified. However, it is partially identified with the following bounds:

In the following, we provide computationally and statistically efficient methods for estimating these upper and lower bounds, each of which are defined via an optimization problem over $\theta$. $\Theta$ encapsulates the constraints on potential distribution shifts that are known in a particular domain, allowing an analyst to translate additional domain knowledge into tighter identification of the estimand. Intuitively, such identification will be close to point-wise when the constraints are informative enough about the behaviour of the estimand of interest under the true distribution. Our aim is to use the observed sample $X_{1},\dots,X_{N}$ drawn iid from $\mathbb{Q}$ to estimate the value of the optimization problems defining our lower and upper bounds. Note that in order to do so, we have to tackle two different types of uncertainty. First, the uncertainty derived from partial identification where even the population level quantities $\min_{\theta\in\Theta}/\max_{\theta\in\Theta}f(\theta\mathbb{Q})$ do not exactly identify $f(\mathbb{P})$ even if $\mathbb{Q}$ were known exactly. Second, finite-sample uncertainty from the estimation of the population level quantity $f(\theta\mathbb{Q})$ using a sample $X_{1},\dots,X_{N}$: the statistical uncertainty from taking $f(\theta\hat{\mathbb{Q}}_{N})$ as an approximation of $f(\theta\mathbb{Q})$. We will provide confidence bounds for $f(\mathbb{P})$ which account for both partial identification and finite-sample uncertainty in estimation. Throughout, we will assume for convenience that the optimization problems defining each side of the bound have a unique solution (if this is not satisfied, we could, e.g., modify the objective function to select a minimum-norm solution).

We now proceed to estimate the aforementioned bounds. Proposition 1 reduces the problem of computing such bounds to solving an optimization program. Throughout, we consider the problem of estimating the lower bound (i.e., solving the minimization problem), as the upper bound is symmetric. To start, we impose no restrictions on $\theta(X)$ by representing $\theta(X)$ as an arbitrary function. This yields the following optimization problem:

where the constraint $\theta(X)\geqslant 0$ ensures that $\theta$ is a valid density ratio. However, neither $f(\theta\mathbb{Q})$ nor the population level restrictions can be observed directly. Nevertheless, having access to a sample from $\mathbb{Q}$, it is only natural to use plug-in estimators to estimate the unavailable quantities from the previous program. This yields the sample optimization problem:

where $\hat{\mathbb{Q}}_{N}$ is the empirical distribution of $\mathbb{Q}$. We refer to the first problem (1) as the population optimization problem ($\min_{\theta}f(\theta\mathbb{Q})$) and to the second one (2) as the plug-in approximation ($\min_{\theta}f(\theta\hat{\mathbb{Q}}_{N})$).

Let $\nu$ be the optimal value of the population problem and $\hat{\nu}_{N}$ the optimal value of the plug-in approximation problem. Suppose $\sqrt{n}(\hat{\nu}_{N}-\nu)$ converges in distribution to a normally distributed random variable, such that, it has mean zero and variance that can be estimated from the data. If this is indeed the case, then the standard procedure to derive confidence intervals for $\nu$ can be applied. Furthermore, as corollary of Proposition 1, we can produce a high probability lower bound for $f(\mathbb{P})$. As the maximization problem is symmetric, we can use the same argument to upper bound $f(\mathbb{P})$.

We now turn to analyze several scenarios where this recipe can be applied. We use the theory of sample average approximation in optimization [40, 41] to describe the conditions to guarantee, both convergence and asymptotic normality, for the statistic $\sqrt{n}(\hat{\nu}_{N}-\nu)$. However, before developing this mathematical scaffolding, it is necessary to make one more assumption: we will assume a, differentiable, finite-dimensional parametrization of $\theta(X)$, denoted by $\theta_{\alpha}(X)$, for some $\alpha\in\mathbb{R}^{d}$. This class is still quite expressive; e.g., a standard basis for smooth functions (e.g., a set of polynomial basis functions) allows us to represent any smooth function in this framework, or the analyst could use a highly expressive class such as neural networks [21].

We conduct experiments to show how our method allows users to specify domain knowledge in order to obtain informative bounds on the estimand of interest. We simulate inference for a range of different $f(\mathbb{P})$ across various scenarios by testing our method with different choices of constraints and datasets. In addition, we present a real-world use case where we were able to produce policy-relevant conclusions using our method.

In each experiment, we start with samples from a ground truth distribution $\mathbb{P}$ and then simulate the observed distribution $\mathbb{Q}$ using sampling probabilities which depend on the covariates (i.e., simulating selection bias). This ensures that the ground-truth value $f(\mathbb{P})$ is known (to high precision), allowing us to verify if our bounds contain the true value. We consider two classes of estimands. First, estimating the conditional mean $\mathbb{E}_{\mathbb{P}}[Y|A=1]$ for an outcome $Y$ and covariate $A$. Second, estimating the coefficient of a linear regression model as an example of the $m$-estimation setting from above.

We show how the amount of specified domain knowledge can result in tighter bounds along two axes. First, we vary the parametric form for $\theta$, ranging from an arbitrary function of $X$ (i.e., a separate parameter for each discrete covariate stratum) to one that imposes specific assumptions (e.g., separability across specific groups of covariates, or that the covariates which the selection probabilities depend on are known). Second, we vary the number and informativeness of the constraints $\{g_{j}\}$ used to form the set $\Theta$; modeling the ability of users to impose an increasing degree of constraints on possible distribution shifts. The Appendix shows experiments with additional kinds of constraints, e.g., on the sign of the covariance between pairs of variables.

The closest baselines to compare our methods against are sensitivity analyses based on distributionally robust optimization (DRO). However, DRO requires the distance between distributions under some divergence (the unobservable parameter $\rho$) as input. Conversely, our method requires observable aggregate data to provide partial identification, making a head-to-head comparison difficult. In order to provide baselines for comparison, we implement DRO on two artificial settings where the maximum distance $\rho$ is provided. The first baseline (omniscient) optimistically assumes that the $\chi^{2}$ divergence ($\rho$) between $\hat{\mathbb{Q}}_{n}$ and $\mathbb{P}$ is known. This gives an ”unfair advantage” to DRO because not only is a valid value of $\rho$ provided but, furthermore, it is the smallest of such feasible values (which is not actually observable when $\mathbb{P}$ is unknown). Hence, we also propose a second scenario (observable) where the tightest possible value of $\rho$ is inferred from the same data used by our method. In this scenario $\rho$ is estimated by solving for $\max_{\theta\in\Theta}D_{\chi^{2}}(\theta\hat{\mathbb{Q}}_{n},\hat{\mathbb{Q}}_{n})$. Thus, the solution of the program will be the radius of the smallest ball centered at $\hat{\mathbb{Q}}_{n}$ that contains all distributions that satisfy all the observable constraints. Of course the previous program is not convex in $\theta$; hence we settle for the solution obtained from a saddle point obtained via gradient descent (which errs in favor of DRO by selecting a smaller value of $\rho$ than the one that may be justified by the data).

We now proceed to summarize how these experiments were instantiated in each setting, providing detailed information in the Appendix. Additionally, a guide on how to replicate each of these experiments is available in the code repository at hidden.

Synthetic data experiments. For the first set of experiments, we simulate a distribution over binary variables $X=(Y,Y_{2},A,X_{1},X_{2})\sim\mathbb{P}$ used to evaluate previous causal inference methods [27]. We add a selection bias scenario to this process by simulating an indicator variable $R\sim Ber(logit^{-1}(X_{1}-X_{2}))$. The observed distribution $\mathbb{Q}$ consists of those samples for which $R=1$. In this domain, we consider the task of estimating bounds for $\mathbb{E}_{\mathbb{P}}[Y|A=1]$ setting restrictions on $\mathbb{E}_{\mathbb{Q}}[\theta YX_{2}]$ (i.e., $\mathbb{E}_{\mathbb{P}}[YX_{2}]$ is known). A total of six experiments were run. In the first three experiments, we vary the functional form for $\theta$ while leaving the constraints defining $\Theta$ fixed. In the first experiment (unrestricted), $\theta$ is an arbitrary function of $X$. In the second experiment (separable), we specify $\theta(X):=\theta_{1}(A)+\theta_{2}(X_{1},X_{2})$ where $\theta_{1}$ is an arbitrary function of $A$ and $\theta_{2}$ is an arbitrary function of $X_{1}$ and $X_{2}$. Finally, in the third experiment (targeted), we fix $\theta(X)=\theta(X_{1},X_{2})$, i.e., $\theta$ is a function only of the variables which determine $R$, thus simulating the scenario where the variables driving selection bias are known (even if the exact selection probabilities are not).

For experiments four to six, we vary $\Theta$ instead by adding more informative constraints in each successive experiment. As our intention is to isolate the effect of adding more constraints, the parametric form of $\theta$ is set as flexible as possible, i.e., as an arbitrary function of $X$ (the unrestricted experiment above). In experiment four, we constrain two of the four parameters of the joint distribution of $Y_{2}$ and $X_{2}$ via constraints on $\mathbb{E}_{\mathbb{Q}}[\theta Y_{2}X_{2}]$ and $\mathbb{E}_{\mathbb{Q}}[\theta Y_{2}(1-X_{2})]$. In experiment five, we add constraints on the remaining components of the joint distribution $\mathbb{E}_{\mathbb{Q}}[\theta(1-Y_{2})X_{2}]$ and $\mathbb{E}_{\mathbb{Q}}[\theta(1-Y_{2})(1-X_{2})]$. Experiment six, add constraints on the outcome ($\mathbb{E}_{\mathbb{P}}[YX_{2}]$) as well. Through out, we will refer to these experimental setups as (partial) race + income, (full) race + income and race + income + outcome respectively (the names were chosen to keep consistency with the semi-synthetic experiments).

Semi-synthetic data experiments. We use the Folkstables dataset [11] which provides an interface for the data from the US Census American Community Survey from 2014 to 2019. Based on the ACSEmployment task suggested in [11], we consider a binary outcome variable $Y$ indicating whether or not a person was employed at the time of the survey. We are interested in predicting the female unemployment rate (e.g., motivated by a policymaker who seeks to track or intervene on gender employment gaps [2] [32]). The variables driving the selection bias are citizenship and veteran status, as these variables can contribute to self-selection in a potential social services database.

We first consider estimating $\mathbb{E}_{\mathbb{P}}[Y|A=1]$ where $A$ indicates the sex of an individual. We conduct six experiments that exactly parallel the construction of the synthetic experiments, with the variable for income level playing the role of $Y_{2}$ and race/ethnicity playing the role of $X_{2}$ (see Appendix for details). This simulates a setting where information about other outcomes correlated with the variables of interest is leveraged for partial identification. In the last experiment race + income + outcome, the additional constrain $\mathbb{E}_{\mathbb{Q}}[\theta YX_{2}]$ simulates a scenario where information on the outcome of interest is available with respect to a different demographic grouping than our desired estimand (race/ethnicity instead of gender).

We then consider estimating the coefficient of the indicator variable $A=1$ in a linear regression model of $Y$ on 7 covariates. This second setting is an example of the $M$-estimator framework from above. We run the experiments (partial) race + income, (full) race + income and race + income + outcome using the settings of the $E_{\mathbb{P}}[Y|A=1]$ case. Details for each experiment are shown in the Appendix.

Real world case study. We apply our method to a dataset of over five million de-identified medical claims from COVID-19 patients. Medical claims provide the services that patients received, their diagnosis, and demographic attributes. We use this data to study disparities in hospitalization risk for US COVID-19 patients across race/ethnicity groups. We provide bounds for the relative excess risk of hospitalization for non-white patients after conditioning on clinically relevant covariates (a set of five comorbidities and age). Formally, our estimand is $f(\mathbb{P})=\sum_{x}[P(Y=1|X=x,Race)-P(Y=1|X=x,Race=White)]P(X=x|Race)/P(Y=1|Race=White,X)$, where $Y$ is whether a patient was hospitalized. As mentioned in the introduction, selection bias is particularly challenging for these data, because it may depend jointly on both the outcome and main covariate of interest. We model selection based on race, income, and hospitalization status to account for this potential interaction between disease severity and healthcare access in ascertainment. As constraints $g$, we use the total hospitalizations per racial group reported by the CDC and the total number of estimated infections per group from CDC serology studies.

To further investigate the stability of our estimand, every experiment was run across 5 different random restarts of the projected gradient descent algorithm.

Results. Figures 1 and 2 show our main results. Each figure plots the bounds output by our method for each experiment described above. In each plot, we also show the true value of the estimand $f(\mathbb{P})$, the naive estimate using samples from $\mathbb{Q}$ (without any attempt to account for selection bias) and the bounds produced by the two DRO baselines. In each plot (1 and 2), accounting for the different results across multiple bootstrap replicates, there is a box plot at the endpoints of every bound as well as a band to represent $f(\mathbb{Q})$ .

The results of our experiments are consistent across all six scenarios, showing how the estimated bounds vary depending on the constraints imposed and on the functional form assumed for $\theta$. Concordant with the theory, the bounds always contain the true value of the estimand. Additionally, as expected, when more external information is available the bounds become narrower. When additional constraints are added, it is possible to obtain narrow bounds for the estimand that rule out the naive estimates. Our method Pareto-improves over the observed-data DRO baseline across all settings. Interestingly, in the majority of settings, our method also produces tighter intervals than the omniscient DRO baseline. Thus, it can be highlighted the advantage of more nuanced descriptions of uncertainty: even in the case the true value of the divergence is known, moment constraints may provide a tighter bound than a divergence metric. We also find that in the targeted experiment where the variables driving selection bias are known, our method often provides close to point identification, indicating that this is a particularly powerful form of domain knowledge when available. These results demonstrate that our framework allows users to translate common forms of domain knowledge into robust and informative statistical inferences.

Figure 3 summarizes the results from the real-world case study. We find that Black, Asian, and Hispanic patients have a higher risk of hospitalization compared to white patients. Importantly, our method produces informative bounds for all groups, which are close to point identification for Black and Asian patients while highlighting greater sensitivity to potential selection biases for Hispanic patients. Since the true underlying target distribution is unknown, we cannot compare our results to the ‘ground truth’ in this domain; rather, this showcases that our methods produce informative and theoretically-guaranteed intervals for a policy-relevant question.