Model selection for estimation of causal parameters

Model selection has a long history in statistics and machine learning. For optimizing loss-based estimators, the most commonly used methods include cross-validation, the Akaike information criterion, and the Bayesian information criterion (Akaike, 1974; Schwarz, 1978; Friedman et al., 2001; Arlot and Celisse, 2010).

The focused information criterion is a model selection criterion which, for a given focus parameter, estimates the mean-squared error of submodels (Claeskens and Hjort, 2003, 2008). It relies on knowing an asymptotically unbiased estimator of the parameter of interest. Its theoretical justification is given in a local misspecification framework.

More recently, Cui and Tchetgen-Tchetgen (2019) introduce a model selection tool for finite-dimensional functionals in a semiparametric model if a doubly robust estimation function is available. It is based on a pseudo-risk criterion that has a robustness property if one of the estimators is biased.

For the task of model selection when estimating heterogeneous treatment effects, several methods have been developed (Kapelner et al., 2014; Rolling and Yang, 2014; Athey and Imbens, 2016; Nie and Wager, 2017; Zhao et al., 2017; Powers et al., 2018). Most of the methodologies are specific to the considered model class. A comparison of this line of work for individual treatment effects can be found in Schuler et al. (2018).

Van der Laan and Robins (2003) propose a loss-based approach for parameter-specific model selection. In this work, the authors recommend minimizing an empirical estimate of the overall risk $R(\hat{\theta}^{(g)},\hat{\eta})$, where $\hat{\theta}^{(g)}$, $g=1,\ldots,G$ are candidate estimators and $\hat{\eta}$ is an efficient estimator of the nuisance parameter, computed on the training data set. Our approach is more generic in the sense that we do not assume that parameter of interest minimizes a known loss function.

Closest to our work is the sample-splitting criterion developed by Brookhart and Van Der Laan (2006). Roughly speaking, the data is split into a training and a test data set. Then, estimators are computed on the training and the test data set, and the squared deviation of estimators is aggregated across several splits. The criterion developed by Brookhart and Van Der Laan (2006) can be seen as a form of Monte Carlo cross-validation. In the following, we discuss a variant of this approach that splits the data into $k$ folds and thus mimics $k$-fold cross-validation procedures, which are popular in practice. The data $D=(D_{1},\ldots,D_{n})$ is randomly split into $K$ disjoint roughly equally-sized folds $D^{0,1},\ldots,D^{0,K}$. Define $D^{1,k}=D\setminus D^{0,k}$. Assuming that the data are i.i.d., $D^{1,k}$ and $D^{0,k}$ are independent for each $k$. Let $\hat{\theta}^{(0)}$ be an unbiased estimator of the parameter of interest $\theta^{(0)}\in\mathbb{R}^{d}$. If several unbiased estimators are available, aggregation procedures such as inverse variance weighting can be used in a pre-processing step to obtain $\hat{\theta}^{(0)}$. Let $\hat{\theta}^{(g)}$ be candidate estimators, $g=0,\ldots,G$. Then, we can compute the risk criterion

Using independence of $D^{1,k}$ and $D^{0,k}$,

As $\text{Var}(\hat{\theta}^{(0)}(D^{0,1}))$ is constant in $g$, the criterion in equation (1) can be used to select an estimator $\hat{\theta}^{(g)}$ with low mean-squared error for estimating $\theta^{(0)}$ among $\hat{\theta}^{(g)},g=0,\ldots,G$. The criterion in equation (1) is attractive as it is simple and widely applicable. We will compare the proposed model selection criterion to the criterion in equation (1), both from a theoretical perspective and in simulations.

In this paper, we work towards making parameter-specific model selection more reliable. We derive a model selection criterion that estimates the squared $\ell_{2}$-deviation of an estimator from its target. We show that the selected model has equal or lower variance than the baseline estimator asymptotically. Compared to the model selection procedure (1), we show that the proposed criterion has equal or lower variance asymptotically. The proposed criterion is flexible in the sense that it can be used for any low-dimensional estimator in both parametric and semi-parametric settings. Theoretical justification of the method is given under the assumption of asymptotic linearity.

Even if the candidate estimators are asymptotically linear, the model selection procedure is discontinuous and will not result in a regular estimator, even for $n\rightarrow\infty$. Thus, the final estimator does not have a Gaussian limit distribution. We derive asymptotically valid confidence intervals for the resulting estimator that takes into account the model selection step.

In previous work, the goal of model selection for estimation is usually to select a nuisance parameter model from a set of candidate models, but the identification strategy is held fixed across models (Van der Laan and Robins, 2003; Brookhart and Van Der Laan, 2006; Cui and Tchetgen-Tchetgen, 2019). In this paper, the goal is to select among different estimands, or identification strategies. Mathematically, this correspond to selecting among different functionals of the underlying data generating distribution. In some situations, this results in dramatic improvements in the mean-squared error. However, there is no free lunch. Compared to the baseline procedure, for fixed $n$, model selection can lead to increased risk in parts of the parameter space. We provide a finite-sample bound that reveals that the excess risk due to model selection becomes negligible as the dimension of the target parameter grows.

In simulations, the proposed criterion exhibits reliable performance in a variety of scenarios, including experiments, instrumental variables settings and data with selection on observables. The code can be found at github.com/rothenhaeusler/tms.

We observe data $D=(D_{i}$, $i=1,\ldots,n)$, where the $D_{i}$ are independently drawn from some unknown distribution $\mathbb{P}$. Suppose we have access to estimators $\hat{\theta}^{(g)}(D)$, $g=0,\ldots,G$, of some unknown parameter $\theta^{(0)}$. In the following, to simplify notation, we will write $\hat{\theta}^{(g)}$ instead of $\hat{\theta}^{(g)}(D)$. We assume that the baseline estimator $\hat{\theta}^{(0)}$ is asymptotically unbiased for $\theta^{(0)}$, i.e. that $\mathbb{E}[\hat{\theta}^{(0)}]=\theta^{(0)}+o(n^{-1/2})$. In practice, the data scientist may know several estimators that are asymptotically unbiased for the parameter of interest. In this case, one can use aggregation procedures such as inverse variance weighting to construct an optimally weighted aggregated estimator $\hat{\theta}^{(0)}$.

In addition, the data scientist may have access to estimators $\hat{\theta}^{(g)}$ for which the data scientist is not sure whether they are approximately unbiased for the effect of interest. The goal is to select among the set of estimators, minimizing the mean-squared error with respect to the target of interest $\theta^{(0)}$. We assume that $\mathbb{E}[\hat{\theta}^{(g)}]=\theta^{(g)}+o(n^{-1/2})$ for some unknown $\theta^{(g)}$ and that $\sqrt{n}(\hat{\theta}^{(g)}-\theta^{(g)})$ converges to a non-degenerate random variable. We write $\sigma_{j}^{(g)}$ for the asymptotic standard deviation of $\sqrt{n}(\hat{\theta}^{(g)}_{j}-\theta^{(g)}_{j})$. Similarly, we assume that $\sqrt{n}(\hat{\theta}^{(g)}-\hat{\theta}^{(0)}-(\theta^{(g)}-\theta^{(0)}))$ converges to a non-degenerate random variable for $g\neq 0$ and write $\tau_{i}^{(g)}$ for the asymptotic standard deviation of $\sqrt{n}(\hat{\theta}^{(g)}_{j}-\theta^{(g)}_{j}-\hat{\theta}^{(0)}_{j}+\theta^{(0)}_{j})$.

We aim to find an estimator $g$ that minimizes

Here and in the following, we suppress the dependence of $R(g)$ and $\hat{\theta}^{(g)}$ on $n$. As bias and variance of $\hat{\theta}^{(g)}$ are unknown, the function $R(g)$ is unknown and one has to minimize a proxy of the risk $R(g)$ instead. We propose to estimate $R(g)$ in equation (2) via

where $\hat{\sigma}_{j}^{(g)}$ is an estimator of the asymptotic standard deviation of $\sqrt{n}(\hat{\theta}^{(g)}_{j}-\theta^{(g)}_{j})$ and $\hat{\tau}_{j}^{(g)}$ is an estimator of the asymptotic standard deviation of $\sqrt{n}(\hat{\theta}^{(g)}_{j}-\theta^{(g)}_{j}-\hat{\theta}^{(0)}_{j}+\theta^{(0)}_{j})$. If the estimators are asymptotically linear (i.e. in some semi-parametric or low-dimensional parametric settings), consistent estimators $\hat{\tau}^{(g)}$ and ${\hat{\sigma}^{(g)}}$ are usually available via plug-in estimators of the variance of the influence function (Van der Vaart, 2000; Tsiatis, 2007). An example will be discussed below. We propose to choose a final estimate $\hat{\theta}^{(\bar{g})}$ by solving

Let us consider a linear regression example. This example was mainly chosen for expository simplicity; the main motivating examples for this method are drawn from causal inference. The causal inference examples need more discussion and will be explained in detail in Section 4.

We will compare the risk proxy in equation (3) to sample-splitting based criteria in Section 3. The method is evaluated on simulated data sets in Section 4.

We make two major assumptions, in addition to the assumptions outlined in Section 2.1. The first major assumption is a slightly stronger version of asymptotic linearity. Asymptotic linearity is an assumption that is commonly made to justify asymptotic normality of an estimator (Van der Vaart, 2000; Tsiatis, 2007). As our goal is to estimate the mean-squared error of an estimator, we use a slightly stronger version that guarantees convergence of second moments. The second major assumption is that the variance estimates are consistent.

Let us compare the first part of the assumption to asymptotic linearity. Asymptotic linearity assumes that $\|e_{g}(n)^{2}\|_{2}^{2}=o_{P}(1/n)$ while we assume that $\mathbb{E}[\|e_{g}(n)^{2}\|_{2}^{2}]=o(1/n)$. Thus, our assumption is stronger than asymptotic linearity. Let us now turn to our theoretical results.

Our first result shows that the proposed criterion is asymptotically unbiased for the mean-squared error of $\hat{\theta}^{(g)}$. The convergence rate depends on whether $\hat{\theta}^{(g)}$ is asymptotically biased for estimating $\theta^{(0)}$. If the estimator $\hat{\theta}^{(g)}$ is unbiased, the convergence rate is faster. The proof of the following result can be found in the supplement.

This result shows that the estimator is asymptotically unbiased, which is important for the theoretical justification of the approach. The major strength of the proposed approach will become apparant in the next section, where we compare asymptotic bias and variance to the cross-validation criterion (1).

In this section, we will show that the proposed criterion has asymptotically lower variance than the cross-validation criterion (1) if $\theta^{(g)}=\theta^{(0)}$. Furthermore, we show that cross-validation is generally biased for estimating the mean-squared error if $\theta^{(g)}=\theta^{(0)}$. Let us first discuss asymptotic bias of the naive approach. The short proof of this result can be found in the supplement.

Comparing this result with Theorem 1, we see a major difference how the methods behave when $\theta^{(g)}=\theta^{(0)}$. The proposed criterion is asymptotically unbiased for the mean-squared error, while cross-validation is not (even modulo additive constants). The cross-validation approach computes the estimator on randomly chosen subsets of the data. Thus, intuitively, this approach overestimates the asymptotic variance of unbiased estimators compared to the mean-squared error of biased estimators. In practice, this means that the approach tends to choose biased estimators with low variance over unbiased estimators with slighly higher variance. This effect can also be observed in Section 4.

Now, let us turn to the asymptotic variance of the model selection criteria. The proof of the following result can be found in the supplement.

Roughly speaking, this theorem shows that the proposed criterion $\hat{R}(g)$ has equal or lower asymptotic variance than the cross-validation criterion $\tilde{R}(g)$. This means that risk estimates based on the cross-validation criterion are harder to interpret than the risk estimates of the proposed criterion, since their variation will be larger. The difference in variance is particularly large if the splits are imbalanced, i.e. if the test data sets have much smaller or much larger sample size than the training data sets. Intuitively, if the validation data sets have small sample size, validation becomes unstable. However, if the validation data set is large, the estimator on the training data sets becomes unstable. This tradeoff can be avoided by estimating bias and variance separately, as done in the proposed approach.

Note that for understanding the excess risk of the resulting estimator it is not sufficient to study the asymptotic variance of the criterion itself. Thus, we study the asymptotic risk in Section 3.4 and give a finite-sample bound for the excess risk in Section 3.6.

Here and in the following, we focus on the case where the number of models $G$ is small and fixed and $n\rightarrow\infty$. First, we investigate the asymptotic behaviour of the proposed procedure in the case where the number of models is fixed and $n\rightarrow\infty$. The proof of the following result can be found in the supplement.

In words, for $n\rightarrow\infty$, the proposed method selects models with lower or equal risk than the baseline estimator $\hat{\theta}^{(0)}$. Interestingly, an analogous result does not hold for the cross-validation procedure (1). In Section 4 we will see that even for relatively large $n$, the selected model by cross-validation may have risk that is significantly larger than the risk of the baseline procedure.

Deriving confidence intervals that are valid in conjuction with a model selection step is a challenging topic and has attracted substantial interest in recent years, see for example Berk et al. (2013); Taylor and Tibshirani (2015). Generally speaking, statistical inference after a model selection step can be unreliable if the uncertainty induced by the model selection step is ignored. In this section, we describe how to construct confidence intervals that take into account the uncertainty induced by model selection. Intuitively speaking, the challenge is that the final estimator is a discontinuous function of the data. To be more precise, the final estimator $\hat{\theta}^{(\bar{g})}$ is not regular and not asymptotically normal.

The goal in this section is to find $I_{1}$ and $I_{2}$ as a function of the data $D_{1},\ldots,D_{n}$ such that

for some pre-determined $\alpha>0$ and $j$ and where

The following theorem shows how to construct confidence intervals in low-dimensional settings; i.e. in settings where the number of models $G$ is fixed and the sample size goes to infinity.

Note that by definition the conditional distribution of $Z$ given $D_{1},\ldots,D_{n}$ is known to the researcher. Also, the researcher often can construct estimators of the variances in parametric and semi-parametric settings via plug-in estimators of the influence function, see e.g. (Van der Vaart, 2000; Tsiatis, 2007). In these cases, $b_{i,\alpha}(D_{1},\ldots,D_{n})$ can be computed by the researcher, for example by Monte-Carlo simulation. Thus, Theorem 4 allows us to construct asymptotically valid confidence intervals for the final estimator $\hat{\theta}^{(\bar{g})}$ in many parametric and semi-parametric settings.

This result is different from standard asymptotic arguments in the sense that $\hat{\theta}^{(\bar{g})}-\theta^{(0)}$ is not asymptotically Gaussian. However, as the result shows, it is possible to recover the exact asymptotic distribution of $\hat{\theta}^{(\bar{g})}$ in low-dimensional scenarios and use this information to conduct asymptotically valid statistical inference. We will evaluate the empirical performance of confidence intervals constructed via Theorem 4 in Section 4.

As discussed in Section 3.4, the proposed method selects a model that is asymptotically no worse than the baseline estimator. However, there is no free lunch. In transitional regimes, for fixed $n$, the estimator can perform worse than the baseline estimator $\hat{\theta}^{(0)}$. This is to be expected from statistical theory, see for example the discussion of the Hodges-Le Cam estimator on page 110 in Van der Vaart (2000). This makes it important to understand in which cases we can expect reliable performance of the proposed model selection procedure. In the case $d=1$, a Bayesian bound (Gill and Levit, 1995) reveals that improving over the Cramér-Rao bound in some parts of the parameter space must lead to deteriorating performance in other parts of the parameter space. In the following we will provide a finite-sample bound that will show that for large $d$ the excess risk becomes negligible, uniformly over a set of distributions.

In particular, the excess risk due to model selection goes to zero as $\frac{d}{\log G}\rightarrow\infty$ and if $\sup_{g,j}|(\hat{\sigma}_{j}^{(g)})^{2}-\sigma_{j}^{(g)}|\rightarrow 0$ and $\sup_{g,j}|\hat{\tau}_{j}^{(g)}-\tau_{j}^{(g)}|\rightarrow 0$. In Section 4 we will see in an example that the excess risk declines for growing dimension of the estimand.

In observational studies, it is common practice to estimate causal effects under the assumption of unconfoundedness and under the overlap assumption. Roughly speaking, the overlap assumption states that treatment assignment probabilities are bounded away from zero and one, conditional on covariates $X$. If these assumptions are met, it is possible to identify the average treatment effect via matching, inverse probability weighting, regression adjustment, or doubly robust methods (Hernan and Robins, 2020; Imbens and Wooldridge, 2009). However, if the overlap is limited, estimating the average treatment effect can be unreliable.

To deal with the issue of limited overlap, researchers sometimes switch to different estimands such as the average effect on the treated (ATT) or the overlap weighted effect (Crump et al., 2006). In the following, we will focus on the overlap-weighted effect as it is the causal contrast that can be estimated with the lowest asymptotic variance in certain scenarios (Crump et al., 2006). The overlap-weighted effect is defined as

where $\tau(x)=\mathbb{E}[Y(1)-Y(0)|X=x]$. Note that if the treatment effect is homogeneous $\tau(x)\equiv\text{const.}$, then the overlap-weighted effect and the average treatment effect coincide, that means $\theta^{(1)}=\theta^{(0)}$. Thus, shrinkage towards an efficient estimator of the overlap effect is potentially beneficial under treatment effect homogeneity.

We investigate shrinking between estimators of the average treatment effect and the overlap-weighted effect in a data-driven way. The proposed model selection tool will be used to trade off bias and variance.

The instrumental variables approach is a widely-used method to estimate causal effect of a treatment $T$ on a target outcome $Y$ in the presence of confounding (Wright, 1928; Bowden and Turkington, 1990; Angrist et al., 1996). Roughly speaking, the method relies on a predictor $I$ (called the instrument) of the treatment $T$ that is not associated with the error term of the outcome $Y$. We will not discuss the assumptions behind instrumental variables in detail, but refer the interested reader to Hernan and Robins (2020). We will focus on the case, where $I$, $T$ and $Y$ are one-dimensional. Under IV assumptions and linearity, the target quantity can be re-written as

Estimating this quantity can be challenging if the instrument is weak, i.e. if $\text{Cov}(I,T)\approx 0$. In this case, the approach can benefit from shrinkage towards the ordinary least-squares solution (Nagar, 1959; Theil, 1961; Rothenhäusler et al., 2020; Jakobsen and Peters, 2020). Doing so may decrease the variance but generally introduces bias. We will focus on the case where we have some additional observational data, where we observe $T$ and $Y$, but where the instrument $I$ is unobserved.