Calibrating doubly-robust estimators with unbalanced treatment assignment

Extracted main text — title through conclusion, appendix excluded. This is what our citation measures are computed over, published so the extraction can be checked by eye.

35,138 characters · 8 sections · 44 citation commands

Rendered from LaTeX for readability, not typeset faithfully. Citation keys are highlighted; maths is left as source; figures, tables and equation environments are summarised rather than reproduced; unrecognised commands are greyed out so nothing is silently dropped. Email addresses are removed.

abstractMachine learning methods, particularly the double machine learning (DML) estimator Chernozhukov2018, are increasingly popular for the estimation of the average treatment effect (ATE). However, datasets often exhibit unbalanced treatment assignments where only a few observations are treated, leading to unstable propensity score estimations. We propose a simple extension of the DML estimator which undersamples data for propensity score modeling and calibrates scores to match the original distribution. The paper provides theoretical results showing that the estimator retains the DML estimator's asymptotic properties. A simulation study illustrates the finite sample performance of the estimator.

\jelcodes{C14 \and C21 \and C52 \and C55}

Introduction

Estimation of the average treatment effect (ATE) is of central importance in empirical research. The interest generally lies in the effect that a binary treatment has on an outcome variable. For example, we might be interested in the effect that a training program has on the unemployment duration. With the increasing availability of large datasets, the use of machine learning (ML) methods to estimate the ATE has become popular Athey2019,Liuyi2021. The double machine learning (DML) estimator Chernozhukov2018 is a widely adopted ATE estimator that relies on ML-estimated nuisance functions. The DML estimator has been shown to be consistent and asymptotically normal, even when using ML methods that converge at a slower rate than the parametric rate Chernozhukov2018.

While dataset sizes have increased, the number of treated units often remains small. The treatment is often costly and/or time-consuming. For example, the submission of a new drug to a patient might take several months, or a training program for the unemployed is very costly. Control outcomes and covariates, on the other hand, are more easily collected from, e.g., administrative agencies, medical records, or financial markets Kunzel2019,Bouchaud2022,Hujer2006. The dataset might then consist of only a few treated, but many control observations. This unbalancedness can lead to unstable propensity score estimations Huber2013. Even though the DML estimator relies on a doubly robust approach that combines the conditional outcome expectations with the propensity score, instability in the propensity score estimation can lead to high variability in the ATE estimate.

In this paper, we propose a simple extension of the DML estimator that addresses the issue of unbalanced treatment assignment. Inspired by the ML classification literature Japkowicz2002, the proposed approach undersamples the data used for fitting the ML model for the propensity score. Using the relation between the true and the undersampled propensity score, we calibrate the propensity scores to match the original data distribution. We show that the proposed estimator has the same asymptotic distribution as the DML estimator, attains the parametric rate of convergence $\sqrt{N}$ and its variance achieves the semi-parametric efficiency bound Hahn1998. We illustrate the finite sample performance of the estimator in a simulation study. While we present the results in the context of the DML estimator, the proposed approach applies to any ATE estimator that relies on the efficient score function.

Calibration estimator

Notation and causal identification

We define causal effects using Rubin's Rubin1972 potential outcome framework. The interest lies in the effect of a binary treatment variable\footnote{The estimator presented in this paper can be directly generalized to multivalued treatments Farrell2015,Knaus2020.} $D$ on an outcome variable $Y$. We denote the potential outcome for treatment $D=d$, that is the outcome one would observe if the treatment was $d$, as $Y^d$. The effect of interest is the average treatment effect (ATE) defined as:

equation[equation omitted — 64 chars of source]

The potential outcomes are not directly observed and the parameter of interest has to be identified from observational data. The researcher observes a sample of i.i.d. random variables $\{Z_1, \dots, Z_N\}$ where $Z_i := \left(X_i, D_i, Y_i\right)$, with $X_i$ being a $p$-dimensional vector of exogenous control variables with support $\mathcal{X}$, $D_i$ the binary treatment random variable, and $Y_i$ the observed outcome. Define the conditional outcome expectations as $\mu_d(X) = \mathbb{E}[Y\vert D=d, X]$ and the propensity score as $p(X) = \mathbb{E}[D|X] = \mathbb{P}[D=1|X]$. The respective estimated quantities are denoted as $\widehat{\mu}_d(X)$ and $\widehat{p}(X)$. Finally, we denote by $\tau(Z)$ the efficient score function:

equation*[equation* omitted — 114 chars of source]

Identification of the ATE from observable outcomes is achieved under the following assumption.

assumption[Identification] For observation $Z = \left(X, D, Y\right)$ assume that Rosenbaum1983: \begin{enumerate}[label=(\roman*)] • $Y^0, Y^1 \perp D \vert X = x $ for any $x \in \mathcal{X}$ (conditional independence assumption). • The observed outcome is $Y = D Y^1 + (1-D) Y^0$ (stable unit treatment value assumption). • For any $x \in \mathcal{X}$ it holds that $\eta < p(x) < 1-\eta$ for some $\eta>0$ (common support). • $X = X^1 = X^0$ where $X^d$ denotes the random covariate vector under treatment $d$ (exogenity of covariates). \end{enumerate}

Under Assumption (ref) the ATE can be characterized as a functional of the joint distribution of the observed data $(X,D,Y)$ Athey2019:

align[align omitted — 131 chars of source]

Double machine learning estimator

The estimator proposed in this study extends the popular ATE estimator developed by Chernozhukov2018 generally referred to as the double machine learning (DML) estimator. DML builds on two key ingredients. First, it uses the efficient score function $\tau(Z)$ to construct an estimator of the ATE. Second, it uses cross-fitting to estimate the nuisance functions $\mu_d(X)$ and $p(X)$. In more detail, the DML estimator for the ATE is defined as follows:

enumerate[label*=Step \arabic*] • For some fixed $K \in \{2, \dots, N\}$, randomly partition the observation indices into $K$ sets $\mathcal{I}_1, \dots, \mathcal{I}_K$ of equal size. Denote the complement of $\mathcal{I}_k$ by $\mathcal{I}_{-k} = \{1, \dots, N\}\setminus\mathcal{I}_k$. Denote the cardinality of each set of indices by $\vert \mathcal{I}_k \vert$. • for $k=1$ to $K$ do:\\[0.2cm] Estimate the nuisance functions $\mu_d(x)$ and $p(x)$ on the sample defined by indices $\mathcal{I}_{-k}$ and denote the estimated functions by $\widehat{\mu}_d^{\mathcal{I}_{-k}}(x)$ and $\widehat{p}^{\mathcal{I}_{-k}}(x)$.\\[0.2cm] end for • Estimate the ATE using the estimator: \begin{equation} \widehat{\theta}^{DML} = \frac{1}{N}\sum_{k=1}^{K} \sum_{i\in\mathcal{I}_k} \widehat\tau^{\mathcal{I}_{-k}}(Z_i) \end{equation} where: \begin{equation*} \widehat\tau^{\mathcal{I}_{-k}}(Z_i) = \widehat\mu^{\mathcal{I}_{-k}}_1(X_i) - \widehat\mu^{\mathcal{I}_{-k}}_0(X_i) + \frac{D_i}{\widehat{p}^{\mathcal{I}_{-k}}(X_i)}(Y_i-\widehat\mu^{\mathcal{I}_{-k}}_1(X_i)) - \frac{1-D_i}{1-\widehat{p}^{\mathcal{I}_{-k}}(X_i)}(Y_i-\widehat\mu^{\mathcal{I}_{-k}}_0(X_i)) \end{equation*}

Chernozhukov2018 provide an asymptotic theory for the DML estimator. In particular, they show that the estimator is consistent and asymptotically normal, even when using machine learning (ML) methods that converge at a slower rate than the parametric rate.

Calibration estimator

A shortcoming of the DML estimator is its poor finite sample performance when the treatment assignment is unbalanced, i.e. when either very few or very many observations are treated. In the following, without loss of generality, we consider only the case where very few are treated. ML models perform poorly when the data is unbalanced Japkowicz2002 and predict propensity scores that are potentially close to zero or one. A common approach to improving the performance of ML models is to undersample the observations that are overrepresented, i.e. the observations that are not treated He2009.\footnote{Another common technique used in the ML literature is to oversample the minority class Japkowicz2002. This approach might be problematic in the context of DML since it introduces dependence in the data, complicating the asymptotic theory of the proposed estimator. The convergence rates of common ML methods needed for the asymptotic results presented in Section (ref) have been proven for the case of independent and identically distributed data; see, e.g. Belloni2013 for the Lasso, Luo2016 for $L_2$ boosting, Wager2016 for random forests, and Chen1999 for neural networks.}

Undersampling the observations changes the underlying distribution of the data and the predicted propensity scores will be biased for the true propensity scores from the unbalanced distribution Pozzolo2015. To formalize this concept, let $S_i$ be a random variable equal to 1 if observation $i$ is sampled and 0 otherwise. It then follows that $\mathbb{P}[S_i=1|D_i=1]=1$ since all treated observations are kept in the sample. For the untreated observations we have that $\mathbb{P}[S_i=1|D_i=0] = \mathbb{E}[D_i]/(1-\mathbb{E}[D_i]) =: \gamma$. Moreover, since the undersampling strategy is not dependent on the control variables $X$, we have that $\mathbb{P}[S_i=1\vert D_i=d, X_i] = \mathbb{P}[S_i=1\vert D_i=d]$. Using Bayes' rule it can be shown that the propensity score for the undersampled data $p_S(X_i)$ is given by Pozzolo2015:

equation[equation omitted — 123 chars of source]

from where it follows that $p_S(X_i)\neq p(X_i)$ for $\gamma < 1$.

A na\"ive solution to address this issue would be to not only undersample the data used for estimating the propensity score but also the data on which the efficient score function is computed. In other words, a valid strategy would be to undersample the entire dataset and apply DML to the undersampled data (hereafter referred to as U-DML). However, this approach reduces the number of observations from $N$ to $2\cdot \mathbb{E}[D_i] \cdot N$. In cases where only 5% of the observations are treated, we would discard 90% of the observations.

We propose a calibration estimator that uses the entire sample and corrects for the bias in the propensity score. The main idea is to only undersample the data used for fitting the ML model for the propensity score. Using the relation between the true and the undersampled propensity score in Equation (ref), we calibrate the propensity scores to match the original data distribution. In more detail, the calibrated-undersampled DML (CU-DML) estimator is defined as follows:

enumerate[label*=Step \arabic*] • Estimate $\gamma$ as: \begin{equation} \widehat{\gamma} = \frac{\sum_{i=1}^{N} D_i}{\sum_{i=1}^{N} (1-D_i)} \end{equation} • For some fixed $K \in \{2, \dots, N\}$, randomly partition the observation indices into $K$ sets $\mathcal{I}_1, \dots, \mathcal{I}_K$ of equal size. Denote the complement of $\mathcal{I}_k$ by $\mathcal{I}_{-k} = \{1, \dots, N\}\setminus\mathcal{I}_k$. Denote the cardinality of each set of indices by $\vert \mathcal{I}_k \vert$. • for $k=1$ to $K$ do:\\[0.2cm] \begin{enumerate}[label*=.\arabic*] • Estimate the nuisance functions $\mu_d(X)$ on the sample defined by indices $\mathcal{I}_{-k}$ and denote the estimated functions by $\widehat{\mu}_d^{\mathcal{I}_{-k}}(X)$. • Draw random variables $S_i$ for $i \in \mathcal{I}_{-k}$ from a Bernoulli distribution with probabilitiy $\mathbb{P}[S_i=1|D_i=1] = 1$ and $\mathbb{P}[S_i=1|D_i=0] = \widehat\gamma$. Define the undersampled indices as: \begin{equation*} \mathcal{I}_{-k}^S = \{ i \in \mathcal{I}_{-k} \vert S_i=1 \} \end{equation*} Estimate the nuisance function $p_S(X)$ on the sample defined by indices $\mathcal{I}^S_{-k}$ and denote the estimated function by $\widehat{p}_S^{\mathcal{I}^S_{-k}}(X)$. Calibrate the propensity score estimated on the undersampled data to match the original data distribution: \begin{equation} \widehat{p}^{\mathcal{I}_{-k}}(X) = \frac{\widehat{\gamma}\cdot \widehat{p}_S^{\mathcal{I}^S_{-k}}(X)}{\widehat{\gamma} \cdot \widehat{p}_S^{\mathcal{I}^S_{-k}}(X) + \left(1-\widehat{p}_S^{\mathcal{I}^S_{-k}}(X)\right)} \end{equation} \end{enumerate} end for • Estimate the ATE using the estimator: \begin{equation} \widehat{\theta}^{CU-DML} = \frac{1}{N}\sum_{k=1}^{K} \sum_{i\in\mathcal{I}_k} \widehat\tau^{\mathcal{I}_{-k}}(Z_i) \end{equation} where: \begin{equation*} \widehat\tau^{\mathcal{I}_{-k}}(Z_i) = \widehat\mu^{\mathcal{I}_{-k}}_1(X_i) - \widehat\mu^{\mathcal{I}_{-k}}_0(X_i) + \frac{D_i}{\widehat{p}^{\mathcal{I}_{-k}}(X_i)}(Y_i-\widehat\mu^{\mathcal{I}_{-k}}_1(X_i)) - \frac{1-D_i}{1-\widehat{p}^{\mathcal{I}_{-k}}(X_i)}(Y_i-\widehat\mu^{\mathcal{I}_{-k}}_0(X_i)) \end{equation*}

Notice that, in contrast to the nuisance functions, $\gamma$ can be estimated from the entire sample. For $\widehat{\gamma} = 1$ the CU-DML estimator reduces to the classical DML estimator. Furthermore, the results presented in the next section show that asymptotically $\widehat{\theta}^{\text{CU-DML}}$ converges in probability to $\widehat{\theta}^{\text{DML}}$, also for $\widehat{\gamma} < 1$.

Asymptotic results

We now present the asymptotic results for the CU-DML estimator. We start by introducing the following assumptions.

assumption[Boundedness of conditional variances] For the conditional variance of the outcome it holds that: \begin{equation*} \sup_{x\in\mathcal{X}} \mathrm{Var}\left[Y|D=d,X=x\right] < \zeta \end{equation*} for some $\zeta < \infty$.

assumption[Boundedness of the propensity score]\phantom{empty} \begin{enumerate}[label=(\roman*)] • For the unconditional propensity score $\lambda:=\mathbb{E}[p(X)]=\mathbb{E}[D]$ and its estimator $\widehat{\lambda} = N^{-1} \sum_{i=1}^N D_i$ it holds that: \begin{equation*} \epsilon < \lambda < 1 - \epsilon \qquad \epsilon < \widehat{\lambda} < 1 - \epsilon \end{equation*} for some $\epsilon > 0$. • For all $x\in\mathcal{X}$ it holds that: \begin{equation*} \eta < p_S(x) < 1 - \eta \end{equation*} for some $\eta > 0$. \end{enumerate}

assumption[Convergence of the ML estimators]\label[type]{assumption:convergence_ml}\phantom{empty} \begin{enumerate}[label=(\roman*)] • The ML methods are sup-norm consistent: \begin{equation*} \sup_{x\in\mathcal{X}} \vert \widehat\mu_d(x) - \mu_d(x) \vert \overset{p}{\rightarrow} 0 \qquad \sup_{x\in\mathcal{X}} \vert \widehat{p}_S(x) - p_S(x) \vert \overset{p}{\rightarrow} 0 \end{equation*} • The ML methods have risk-decay rates that satisfy (risk-decay assumption): \begin{equation*} \mathbb{E}\left[\left(\widehat\mu_d(x) - \mu_d(x)\right)^2 \right] \mathbb{E}\left[\left(\widehat{p}_S(x) - p_S(x)\right)^2 \right] = o(N^{-1}) \end{equation*} \end{enumerate}

These assumptions closely resemble the ones required for the asymptotic results of the DML estimator Chernozhukov2018,Wager2022. In addition to the boundedness of the propensity score, Assumption (ref) (ref) requires the unconditional propensity score to be bounded away from zero and one. This assumption is natural, as in a situation where the expected propensity score equals zero (one), there are no (only) treated. Importantly, Assumption (ref) (ref) relaxes the usual bounds imposed on the propensity score, since $\tilde\eta<p(X)<1-\tilde\eta$ with $\tilde\eta<\eta$ whenever $\gamma<1$.\footnote{From Equation (ref) it follows that $\tilde\eta<p(X)<1-\tilde\eta$ with $\tilde\eta = \gamma\eta/(1-\eta+\gamma\eta)$.} Assumption (ref) states the convergence rate requirements of the ML estimators in terms of the undersampled propensity score $p_S(X)$. The theoretical result is then given by the following theorem.

theoremUnder Assumptions (ref), (ref) and (ref) it holds that: \begin{equation*} \sqrt{N} \left( \widehat{\theta}^{CU-DML} - \theta \right) \overset{d}{\longrightarrow} \mathcal{N}\left(0, V^*\right) \end{equation*} where: \begin{equation*} V^* = \mathrm{Var}[\mu_1(X_i)-\mu_0(X_i)] + \mathbb{E}\left[\frac{\sigma_1^2(X_i)}{p(X_i)}\right] + \mathbb{E}\left[\frac{\sigma_0^2(X_i)}{1-p(X_i)}\right] \end{equation*} with $\sigma^2_d(X_i) = \mathrm{Var}[Y_i^d\vert X_i]$.

The proof of Theorem (ref) is relegated to Appendix (ref). Theorem (ref) shows that the CU-DML estimator has the same asymptotic distribution as the DML estimator. In particular, the estimator attains the parametric rate of convergence $\sqrt{N}$ and its variance achieves the semi-parametric efficiency bound Hahn1998.

Simulation study

In this section, we investigate the finite sample performance of the CU-DML estimator in a simulation study. We consider two simulation strategies. In the first strategy, we generate data from a synthetic data generating process (DGP). In more detail, we generate data from the DGP used by Nie2020:

align[align omitted — 289 chars of source]

The baseline main effect is the scaled Friedman1991 function $b(X_i) = \sin\left(\pi X_{i,1}X_{i,2}\right) + 2\left(X_{i,3}-0.5\right)^2 + X_{i,4}+ 0.5X_{i,5}$. For the propensity score we follow Kunzel2019 and set $p(X_i) = \alpha \left(1+\beta_{2,4}\left(\min(X_{i,1},X_{i,2}) \right)\right)$ where $\beta_{2,4}(\cdot)$ is the beta cumulative distribution function with shape parameters 2 and 4. The share of treated is $\mathbb{E}[D_i] = (31/21) \alpha$ and the true ATE is 1. The innovation standard deviation is set to $\sigma=1$, and results for $\sigma=5$ are relegated to Appendix (ref).

In the second strategy, we analyze the CU-DML estimator's performance using an Empirical Monte Carlo Study (EMCS) approach Huber2013,Lechner2013. Our EMCS, following Knaus2022, uses a dataset from Lechner2020 of Swiss unemployed individuals in 2003. The dataset includes data on the effect of a job search program (treatment) on the number of months employed in the first six months after the start of the program (outcome), and 49 covariates providing information on individual characteristics of the unemployed, the regional employment agency, and regional labour market characteristics. For more data details, see Knaus2022. The EMCS proceeds as follows. We estimate the propensity score using a logistic regression on the entire sample of 91'339 unemployed individuals. Then, we define a new sample consisting of only the non-treated individuals whose fitted propensity score $\hat{p}_i$ lies between $0.05$ and $0.95$. In each simulation, we draw a random sample with replacement of size $N$ from this new sample and randomly assign treatments as $D_i = \mathds{1}{\{V_i < \hat{p}_i/\lambda\}}$, where $\hat{p}_i$ is the fitted propensity score from the logistic regression, $V_i \sim \text{Unif}[0,1]$ and $\lambda$ controls the share of treated $\mathbb{E}[D_i]$. The true ATE is 0 by construction.

We compare the CU-DML estimator with the DML estimator Chernozhukov2018, and three popular adjustments of the DML estimator that are designed to improve its finite sample performance when the treatment assignment is unbalanced.\footnote{Over the past two decades, these adjustments have been primarily employed to improve the finite sample performance of the inverse-probability-weighted estimator when the treatment assignment is unbalanced; see, among others, the extensive simulation study of Huber2013. More recently, these approaches have also been employed to improve the performance of the DML estimator.} The W-DML estimator uses winsorized propensity scores at 0.01 and 0.99 Imbens2004. In the N-DML estimator we normalize the weights $w_i^{\mathcal{I}_{-k}} := D_i/\widehat{p}^{\mathcal{I}_{-k}}(X_i)$ to sum to unity Wooldridge2018. Following Huber2013, the T-DML estimator normalizes the weights $w_i^{\mathcal{I}_{-k}}$ and then truncates them to not exceed 0.04, before normalizing them again. Finally, we also consider the U-DML estimator, where the entire sample is undersampled and the DML estimator is applied to the undersampled data. All estimators use 5-fold cross-fitting. The nuisance functions are estimated using random forests Breiman2001 with 500 trees. The maximal depth of the trees and the minimal number of observations in the leaf nodes are determined by 5-fold cross-validation over 20 simulation replications and set to the most frequently selected values.\footnote{Properely tuning the hyperparameters of the random forests is crucial, especially for DML when the treatment assignment is unbalanced. The popular Python implementation of random forest, scikit-learn scikit-learn, sets the default minimal number of observations in the leaf nodes to 1. This may lead to highly unstable propensity scores when the treatment assignment is unbalanced. See Bach2024 for a recent simulation study on the importance of hyperparameter-tuning for the DML estimator.} Details on the software used for the simulations are provided in Appendix (ref).

table[table omitted — 7,608 chars of source]

The results are presented in Table (ref). The table reports the root mean squared error (RMSE), the absolute value of the average bias, the standard deviation, and the coverage of the 95% confidence interval for each ATE estimator. The simulation is repeated 1'000 times for each of the following tuples of sample size and share of treated: $\left(N, \mathbb{E}[D_i]\right)\allowbreak \in \allowbreak\{(2'000, 0.05), (2'000, 0.1), (4'000, 0.025), (4'000, 0.05),(4'000, 0.1), (8'000, 0.025), (8'000, 0.05), (8'000, 0.1)\}$. The case where $N=2'000$ and $\mathbb{E}[D_i]=0.025$ is excluded since we would only have 50 treated observations. The best performance in terms of RMSE in each simulation is highlighted in bold. For the synthetic DGP (Panel A), the lowest RMSE is achieved by the CU-DML estimator in 5 out of 8 settings. Only when at least 400 observations are treated, the CU-DML estimator is outperformed in terms of RMSE by the normalized and trimmed DML estimators. However, the RMSE differences between N-DML, T-DML, and CU-DML estimators are considerably small in these cases. These results are unaffected by a higher innovation standard deviation $\sigma=5$, see Table (ref). For the EMCS (Panel B), CU-DML achieves the smallest RMSE for all combinations of $\left(N, \mathbb{E}[D_i]\right)$. Also for the EMCS, we observe that the outperformance of the CU-DML estimator decreases for larger samples and/or less imbalanced samples. These patterns are in line with the findings from the machine learning classification literature, showing that class imbalancedness is a relative problem, related to both the degree of imbalancedness and the sample size Japkowicz2002.

The bias and standard deviation of the estimators unveil that, while the analyzed estimators do not differ substantially in terms of their bias, the CU-DML estimator has the lowest standard deviation across all simulation strategies and settings. The simulation study also confirms the theoretical result presented in the previous section. First, the coverage of the proposed estimator is in all strategies and settings close to 95%. Second, its RMSE decreases at rate $\sqrt{N}$.

Conclusion

This paper addresses a common finite sample problem of the double-machine learning ATE estimator (DML) Chernozhukov2018: in settings with unbalanced treatment assignment estimations of the propensity scores become either too close to zero or one. This causes the DML estimator to become unstable, especially in small samples. We propose a simple yet effective adjustment of the DML estimator (CU-DML) where the machine learning models for the propensity scores are estimated over an undersampled dataset. The resulting propensity score predictions are then calibrated to adjust for the undersampling.

We provide theoretical results for the CU-DML estimator, showing that it has the same asymptotic distribution as the DML estimator. In particular, the estimator attains the parametric rate of convergence $\sqrt{N}$ and its variance achieves the semi-parametric efficiency bound Hahn1998. Furthermore, a small simulation study provides evidence for the finite sample performance of the proposed estimator, showing that it is of particular use in settings with highly unbalanced treatment assignments or small samples.

Future research could adapt the proposed approach to other estimators that rely on the estimation of the propensity score, such as the inverse probability-weighted estimator. Furthermore, CU-DML could be extended to estimators of the conditional average treatment effect Fan2022,Zimmert2019.

appendix