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.
54,478 characters · 24 sections · 121 citation commands
Type 2 Tobit Sample Selection Models with Bayesian Additive Regression Trees
Treatment effect estimates obtained from observational data are often biased by non-random sample selection. heckman1974shadow proposed to jointly model the outcome and the selection of variables to produce unbiased estimates.{\footnote{The Heckman selection model was named the Type 2 Tobit model by {amemiya1984tobit} . }} However, this model assumes a linear functions of covariates in the selection and outcome equations, and joint normality of the error terms.
To address these issues, we introduce a type 2 Tobit model with separate sums of trees in the selection and outcome equations instead of linear combinations of covariates. This model is henceforth referred to as TOBART-2. This builds on the Bayesian formulation of type 2 Tobit models introduced by omori2007efficient, chib2009estimation, and van2011bayesian. The Markov chain Monte Carlo implementation is most similar to that of van2011bayesian. However, Type 2 Tobit models that impose a normality assumption can produce biased estimates when applied to data with small deviations from normality.\footnote{Many other papers have provided alternatives to the assumption of joint normality of the errors in sample selection models. gallant1987semi use Hermite series, marchenko2012heckman use a t-distribution, ding2014bayesian applies a t-distribution in a Bayesian model, and ogundimu2016sample consider a skew-normal distribution. zhelonkin2016robust propose an outlier-robust estimator.} A more robust method, TOBART-2-NP, contains a Dirichlet Process Mixture (DPM) of bivariate Gaussian distributions for the error term, as described by van2011bayesian.\footnote{A similar DPM of normal distributions is used to model the error terms in the instrumental variable BART model of mcculloch2021causal, which is based on conley2008semi.}
The advantages of TOBART-2 are that it accounts for nonlinearity, variable selection, and model uncertainty, while flexibly modelling the errors to allow for non-normality. In addition to point estimates, the method produces credible intervals for outcome predictions and treatment effect estimates. As noted by van2011bayesian, posterior inference can be applied to the covariance of error terms to indicate the extent of sample selection. Variable importance measures can be obtained for the selection and outcome equations.
Data analysis competitions and other simulation studies have demonstrated the impressive performance of BART hahn2019atlantic, dorie2019automated, wendling2018comparing in estimating treatment effects. BART models have been formulated to account for confounding hill2011bayesian, hahn2020bayesian, kim2022bayesian, sparsity linero2018bayesianB, smoothness linero2018bayesianA, and unobserved confounders with instrumental variables mcculloch2021causal. To the best of our knowledge, the only other paper that describes a BART model that allows for a form of selection on unobservables is the multivariate missing not at random outcome model of goh2024joint, created concurrently with TOBART-2.\footnote{Some papers describe BART-based methods for data with selection on observables, e.g. wang2023improving, gao2021treatment, elliott2023improving, bisbee2019barp. } However, goh2024joint model missingness as a function of covariates and latent outcomes, and do not model selection through a joint distribution of error terms as in standard Heckman selection model formulations.\footnote{TOBART-2 is arguably a more interpretable model with a more clearly specified form of selection on unobservables, and it allows for direct specification of the prior on selection and allows for inference on the level of selection on unobservables. The method of goh2024joint is not generalized to non-normal errors, although such an extension would be straightforward. The general approach to selection modelling of TOBART-2 is more common in the econometrics literature, whereas the approach of goh2024joint is more often applied in other disciplines.}
Bayesian Model Averaging of linear Type 2 Tobit models has been considered by jordan2012tobit and eicher2012robust. Nonlinear approaches to Type 2 Tobit models include kernel regression ahn1993semiparametric, newey1990semiparametric, lee1994semiparametric, chen2010semiparametric, splines marra2013estimation,das2003nonparametric, newey2009two, and neural networks zhang2021deep.\footnote{However, zhang2021deep do not evaluate their methods using simulations with non-zero correlation of the error terms in the selection and outcome equations. i.e. there is no selection on unobservables.} chib2009estimation formulate a semiparametric Bayesian model with sample selection and endogeneity by using a second-order Markov Process prior. However, chib2009estimation apply the restrictive assumption of jointly normally distributed errors. kim2019bayesian formulate a Bayesian sample selection model with a scale mixture of bivariate normal distributions for the error term, and a Berstein polynomial regression model for the outcome. However, the selection equation contains a linear function of covariates. wiesenfarth2010bayesian make use of penalized splines in the outcome and selection equations of a Bayesian sample selection model, although the errors are restrictively modelled as being jointly normally distributed. wiemann2022correcting introduce a Bayesian method with splines and a bivariate copula.
fan2007sample apply trees and model averaging to data with sample selection. zadrozny2004learning, tran2017selection, and many others combine machine learning methods with inverse probability weighting approaches that cannot account for selection on unobservables. alaimo20222 apply gradient boosted trees to tax revenue data with sample selection, and account for model uncertainty by using a bootstrap procedure. However, this approach is also based on the inverse propensity methods of zadrozny2004learning and therefore does not account for selection on unobservables.
brewer2024addressing describe control function approaches that allow for machine learning methods to be combined with a Heckman-style two-step framework. This method leads to improved predictive accuracy in simulations with moderate to high levels of selection. However, this approach relies on the assumption of joint normality and brewer2024addressing do not discuss efficiency or methods for inference. zadrozny2001learning similarly apply a control function approach with decision trees. zhu2017nonasymptotic apply a LASSO penalty using a control function approach and semiparametric framework. ogundimu2022lasso, ogundimu2024lasso applies LASSO and adaptive LASSO in both the selection and outcome equations. iqbal2023bayesian apply spike-and-slab priors to the coefficients in the selection and outcome equations to produce a Bayesian sample selection model for high dimensional data. bia2024double describe a double machine learning approach to estimation of the average treatment effect for high-dimensional data with sample selection.
In addition to TOBART-2, this paper introduces the following:
An R package containing implementations of TOBART-2 and the Bayesian linear model described by van2011bayesian is available at \url{https://github.com/EoghanONeill/TobitBART}. The remainder of the paper is structured as follows: In section (ref) we describe the TOBART-2 model and MCMC implementation, in section (ref) we discuss the choice of estimand and evaluation metrics, section (ref) contains a simulation study, section (ref) contains an application to a real world data, and section (ref) concludes the paper.
Description of Model and Priors
Suppose there are $n$ observations, and the $n \times p$ matrix of explanatory variables, $X$, has $i^{th}$ row $x_i=[x_{i1},...,x_{ip}]$. Following the notation of chipman2010bart, let $T$ be a binary tree consisting of a set of interior node decision rules and a set of terminal nodes, and let $M = \{ \mu_1 , ..., \mu_b \}$ denote a set of parameter values associated with each of the $b$ terminal nodes of $T$. The interior node decision rules are binary splits of the predictor space into the sets $\{ x_{is} \le c \}$ and $\{ x_{is} > c \}$ for continuous $x_{s}$. Each observation's $x_i$ vector is associated with a single terminal node of $T$, and is assigned the $\mu$ value associated with this terminal node. For a given $T$ and $M$, the function $g(x_i;T,M)$ assigns a $\mu \in M$ to $x_i$. This gives the single tree model $Y \sim g(x_i;T,M) + \varepsilon \ , \ \varepsilon \sim N(0,\sigma^2)$ chipman1998bayesian. For the standard BART model, the outcome is determined by a sum of trees, $$Y_i = f(x_i) +\varepsilon_i = \sum_{j=1}^m g(x_i ; T_j, M_j)+\varepsilon_i \ \ , \ \ \varepsilon_i \overset{i.i.d}{\sim} N(0, \sigma^2)$$ where $g(x_i;T_j,M_j)$ is the output of a decision tree. $T_j$ refers to a decision tree indexed by $j=1,...,m$, where $m$ is the total number of trees in the model. $M_j$ is the set of terminal node parameters of $T_j$.
Prior independence is assumed across trees $T_j$ and across terminal node means $M_j = (\mu_{1j}...\mu_{b_j j})$, where $1,...,b_j$ indexes the terminal nodes of tree $j$. The prior of chipman2010bart has the form: $$p(M_1,...,M_m,T_1,...,T_m,\sigma) \propto \left[ \prod_j \left[ \prod_k p(\mu_{kj}|T_j) \right] p(T_j)\right]p(\sigma) $$ In standard BART, $\mu_{kj} | T_j \overset{i.i.d}{\sim} N(0,\sigma_0^2)$ where $\sigma_0 = \frac{0.5}{e \sqrt{m}}$ and $e$ is a user-specified hyper-parameter.
chipman2010bart set a regularization prior on the tree size and shape $p(T_j)$. The probability that a given node within a tree $T_j$ is split into two child nodes is $\alpha (1+d_h)^{-\beta}$, where $d_h$ is the depth of internal node $h$, and the parameters $\alpha$ and $\beta$ determine the size and shape of $T_j$ respectively. chipman2010bart use uniform priors on available splitting variables and splitting points. chipman2010bart assume that the model precision $\sigma^{-2}$ has a conjugate prior distribution $\sigma^{-2} \sim Ga(\frac{v}{2}, \frac{v \lambda}{2})$ with degrees of freedom $v$ and scale $\lambda$.
BART predictions are averages of sum-of-tree models. Therefore model uncertainty is taken into account and there are two levels of regularization. Firstly, greater prior probability is placed on models with shallower trees with fewer splitting points. Secondly, over-fitting is further avoided through the prior on the terminal node parameters $\mu_{kj}$, as in standard Bayesian linear regression.
BART Implementation
Samples can be taken from the posterior distribution $p((T_1, M_1),...,(T_m,M_m), \sigma | y)$ by a Bayesian backfitting MCMC algorithm. This algorithm is a Gibbs sampler, involving $m$ successive draws from $(T_j , M_j )| T_{(j)} , M_{(j)} , \sigma , y $ for $j=1,...,m$, where $T_{(j)} , M_{(j)} $ are the trees and parameters for all trees except the $j^{th}$ tree, followed by a draw of $\sigma $ from the full conditional $\sigma | T_1,...,T_m,M_1,...,M_m,y$.
The $q^{th}$ Gibbs iteration defines the sum of trees function $f_q^*(.)= \sum_{j=1}^m g_q(. \ ; T_{j}^*, M_{j}^*)$. After burn-in, the sequence of $f^*$ draws, $f_1^*,...,f_Q^*$ may be regarded as an approximate, dependent sample of size $Q$ from $p(f|y)$. To estimate the unknown function $f(x)$, the expectation $E(f(x)|y)$ is approximated by $\frac{1}{Q} \sum_{q=1}^Q f_{q}^* (x)$. Posterior credible intervals can be obtained from quantiles of the draws $f_{q}^* (x)$.
In the standard (linear) Type 2 Tobit model omori2007efficient, separate parameters and variables enter the selection and outcome equations: $$ y_i =
, \ z_i^* = \bm{w}_i'\bm{\theta} + \xi_i , \ y_i^* = \bm{x}_i'\bm{\beta} + \eta_i , \
\sim \mathcal{N}(\bm{0}, \Sigma) where \Sigma =
$$ In Type 2 TOBART, the linear combinations $\bm{w}_i'\bm{\theta}$ and $\bm{x}_i'\bm{\beta} $ are replaced by sums-of-trees denoted by $f_z(\bm{w}_i)$ and $f_y(\bm{x}_i)$ respectively. $$z_i^* = f_z(\bm{w}_i) + \xi_i \ , \ y_i^* = f_y(\bm{x}_i) + \eta_i $$ For linear Type 2 Tobit, the coefficient priors are $ \bm{\theta} \sim \mathcal{N}(\bm{\theta}_0, \bm{\Theta}_0) \ , \ \bm{\beta} \sim \mathcal{N} (\bm{\beta}_0, \bm{B}_0 ) $, and for TOBART the priors are $f_z \sim BART $ and $f_y \sim BART $. The covariance matrix parameters have the priors $ \gamma \sim \mathcal{N} (\gamma_0, G_0) $ and $ \phi \sim IG(\frac{n_0}{2}, \frac{S_0}{2}) $. An alternative prior, $\gamma | \phi \sim \mathcal{N} (\gamma_0, \tau \phi)$, is introduced by \cite{van2011bayesian}, who notes that dependence between $\gamma$ and $\phi$ through $\tau>0$ allows various shapes of the prior on correlation between $\xi$ and $\nu$, and an induced prior correlation with less probability mass near $\pm 1 $. Details of the Type 2 Tobit and TOBART Gibbs sampler are provided in appendix (ref).
We also consider a prior for the distribution of the error terms similar to the prior introduced by ding2014bayesian. ding2014bayesian assumes that the errors have a bivariate t-distribution: $$
\sim \ \mathcal{N}(\bm{0}, \alpha \Omega/q_i) \ , \ where \ q_i \sim \alpha \chi_{\nu}^2/\nu \ , \ i=1,\dots,N $$ where $ \Omega =
$ and the prior is $ \alpha \sim b / \xi_{c}^2 $ and $ \tilde{\Sigma} =
\Omega
\sim W_2^{-1} (\nu_0, I_2) $, which is equivalent to $p(\Omega) \sim (1 - \rho^2)^{-3/2} \tilde{\sigma}_2^{- (\nu_0 +3 )} \exp \Big\{ - \frac{1}{2 \tilde{\sigma}_2^2 (1- \rho^2) } \Big\}$ and $\tilde{\sigma}_1^2|\Omega \sim \{ (1 - \rho^2) \chi_{\nu_0}^2 \}^{-1} $ , where $\tilde{\sigma}_1 = \sqrt{\tilde{\Sigma}_{1,1}} $ , and $\tilde{\sigma}_2 = \sqrt{\tilde{\Sigma}_{2,2}} = \sqrt{\Omega_{2,2}} $.
However, instead of specifying a t-distribution, we maintain a normal specification for the error term, and allow for departures from normality through a mixture of normal distributions. $$
\sim \ \mathcal{N}(\bm{0}, \Omega) \ , \ \tilde{\Sigma} =
\Omega
\sim W_2^{-1} (\nu_0, I_2) $$ As noted by \cite{ding2014bayesian}, this prior imposes a $U([-1,1])$ marginal prior on $\rho = corr (\xi_i, \eta_i) $ when $\nu_0 = 3$. In contrast, the priors of \cite{van2011bayesian} and \cite{omori2007efficient} can be notably non-uniform.\footnote{The prior of \cite{van2011bayesian} can place most mass around zero, or closer to $\pm 1$ depending on the value of $\tau$. See Appendix (ref) and the appendix of iqbal2023bayesian for further discussion.}
However, this prior does not allow for adequate calibration of the (marginal) prior on $\tilde{\sigma}_2 $ relative to the data-informed prior calibration described by chipman2010bart and mcculloch2021causal. In particular, it can be shown that (see appendix (ref) for derivation) $$ p(\tilde{\sigma}_2^2 ) \propto \left( \frac{1}{\tilde{\sigma}_2^2 } \right)^{\frac{\nu_0-1}{2} + 1} \exp \left( - \frac{1}{ 2 \tilde{\sigma}_2 ^2} \right) $$ therefore the marginal prior is $\tilde{\sigma}_2^2 \sim \Gamma^{-1}( \frac{\nu_0-1}{2} , \frac{1}{2})$. This prior does not allow for both a uniform marginal prior on $\rho$ and a data-informed marginal prior on $\tilde{\sigma}_2^2 $. Therefore, we specify the following prior $$
\sim \ \mathcal{N}(\bm{0}, \Omega) \ , \ \tilde{\Sigma} =
\Omega
\sim W_2^{-1} (\nu_0, c I_2) $$ where we have introduced the hyperparameter $c$. Similarly to \cite{ding2014bayesian}, we can re-express this prior as $p(\Omega) \sim (1 - \rho^2)^{-3/2} \tilde{\sigma}_2^{- (\nu_0 +3 )} \exp \Big\{ - \frac{c}{2 \tilde{\sigma}_2^2 (1- \rho^2) } \Big\}$ and $\tilde{\sigma}_1^2|\Omega \sim \{ \frac{(1 - \rho^2)}{c} \chi_{\nu_0}^2 \}^{-1} $ or $\tilde{\sigma}_1^2|\Omega \sim \Gamma^{-1} \left( \frac{\nu_0}{2}, \frac{c}{2(1-\rho^2)} \right) $. It can be shown (see Appendix \ref{marginal_prior_app}) that this implies the following marginal priors: $$ p(\rho) \propto (1 - \rho^2)^{\frac{\nu - 3}{2} } \ and \ p(\tilde{\sigma}_2^2) \propto \left( \frac{1}{\tilde{\sigma}_2^2} \right)^{\frac{\nu_0-1}{2} + 1} \exp \left( - \frac{ c }{ 2 \tilde{\sigma}_2^2 } \right) $$ This prior allows us to set $\nu_0 =3$ for a uniform marginal $\rho$ distribution and set $c$ such that the $q^{th}$ quantile of the marginal prior on $\tilde{\sigma}_2^2$, i.e. $\Gamma^{-1}( \frac{\nu_0-1}{2}, \frac{c}{2}) $, is equal to an estimated variance of $\eta_i$, denoted by $\hat{\sigma}_2^2$. This is similar to the prior calibration approach of \cite{chipman2010bart} and \cite{mcculloch2021causal}. We set $q = 0.95$. Setting $\hat{\sigma}_2^2$ equal to the sample variance of the observed outcomes might not give a well-calibrated prior if the variance of the observed outcomes is far from the variance of the full set of outcomes, therefore it is preferable to use the estimated outcome variance from a linear Tobit model estimated by maximum likelihood.\footnote{The value of $c$ is $\hat{\sigma}_2^2$ multiplied by the $(1-q)^{th}$ quantile of the $\chi_{\nu_0}^2$ distribution. We also considered a number of data-informed prior calibration methods for the priors of omori2007efficient and van2011bayesian. See Appendix (ref).}
The models presented here are based on the semiparametric Bayesian Type 2 Tobit model implementation introduced by van2011bayesian.\footnote{The notation is more similar to that of omori2007efficient.} In the standard (linear) Type 2 Tobit model, separate parameters and variables enter the selection and outcome equations: $$ y_i =
, \ z_i^* = \bm{w}_i'\bm{\theta} + \xi_i , \ y_i^* = \bm{x}_i'\bm{\beta} + \eta_i,
\overset{i.i.d.}{\sim} \mathcal{N}\Bigg(
, \Sigma_i =
\Bigg) $$ $$\{
' , \gamma_i, \phi_i \} | H \sim H \ , \ H | \alpha, H_0 \sim \mathcal{DP} (\alpha, H_0)$$ In Type 2 TOBART, the linear combinations $\bm{w}_i'\bm{\theta}$ and $\bm{x}_i'\bm{\beta} $ are replaced by sums-of-trees denoted by $f_z(\bm{w}_i)$ and $f_y(\bm{x}_i)$ respectively. $$z_i^* = f_z(\bm{w}_i) + \xi_i \ , \ y_i^* = f_y(\bm{x}_i) + \eta_i \ , \ $$ For linear Tobit, the priors are $ \bm{\theta} \sim \mathcal{N}(\bm{\theta}_0, \bm{\Theta}_0) \ , \ \bm{\beta} \sim \mathcal{N} (\bm{\beta}_0, \bm{B}_0 ) $, and for TOBART $f_z \sim BART $ and $f_y \sim BART $. $H$ is a discrete distribution of the parameters $ \vartheta_i = \{
' , \gamma_i, \phi_i \} $ . The base measure $H_0$ is defined by $
' \sim \mathcal{N} (\bm{0}, \Omega) \ , \ \gamma_i \sim \mathcal{N} (\gamma_0, G_0) \ , \ \phi_i \sim IG(\frac{n_0}{2}, \frac{S_0}{2}) $ and \cite{van2011bayesian} sets $\Omega = 10 \bm{I}_2$.\footnote{An alternative would be to set it equal to $\frac{\hat{\sigma}_y^2}{0.016} \bm{I}_2$ such that the prior probability that a $\mu_{i}$ value lies in $(-10,10)$ is $0.8$. This is similar to the hyperparameter settings described by \cite{chipman1998bayesian} and \cite{conley2008semi}, albeit without a conjugate prior. The other base distribution hyperparameters are set to the same values as for standard TOBART-2. Alternatively, as suggested by \cite{george2019fully}, scale and shape of the prior variance of the outcome error can be set to values that imply a tighter base distribution, as the mixture distribution creates additional spread for the overall prior.} However, since $ f_z(\bm{w}_i) $ is non-linear, we can set $\mu_{i1}$ to 0 and still have a flexible selection model, $\Pr(z_i^* >0) = \Phi ( f_z(\bm{w}_i))$, with a standard normal marginal distribution of $\xi_i$ for identification of $f_z(\bm{w}_i)$. The parameter $\alpha$ determines the distribution of the weights given to each element of $H$. As $\alpha$ increases, more elements receive non-negligible weight. As in the parametric error model, \cite{van2011bayesian} uses the alternative prior $\gamma_i | \phi_i \sim \mathcal{N} (\gamma_0, \tau \phi_i)$.
george2019fully and van2011bayesian apply different priors to $\alpha$. george2019fully choose the maximum and minimum number of components in $G$, denoted by $I_{min}$ and $I_{max}$, then solve for $\alpha_{min}$ such that the mode of the number of elements (denoted by $I$) of $H$ is $I_{min}$. Similarly $\alpha_{max}$ is derived from $I_{max}$. Then $p(\alpha) \propto (1- \frac{ \alpha - \alpha_{min}}{\alpha_{max} - \alpha_{min}} )^{\psi} $ . The default values are $I_{min} = 1$, $I_{max} = [(0.1)*n]$, where $[\bullet]$ denotes the floor function, and $\psi = 0.5$. mcculloch2021causal also specify a DPM prior on the errors of an IV-BART model and set $I_{min} = 2$, $I_{max} = [(0.1)*n] + 1$. van2011bayesian applies the prior from escobar1994estimating, $\alpha \sim \text{Gamma}(c_1, c_2)$, with $c_1=2$ and $c_2=2$. Details of the sampler introduced by van2011bayesian are provided in appendix (ref). We apply the sampler of van2011bayesian by default and provide an option for the prior of george2019fully.
Samples of $y_i^*$ and $z_i^*$ in the Gibbs sampler for TOBART are the same as in the standard Tobit sampler. However, in contrast to the draws of coefficient parameters for linear Tobit, each tree must be drawn separately conditional on all other trees. Detailed outlines of the full TOBART-2 and TOBART-2-NP samplers are included in appendices (ref) and (ref) respectively. An alternative sampler, in which the outcome equation tree draws are made unconditional on the terminal nodes and covariance term $\gamma$, and in which the outcome equation leaf parameters and $\gamma$ are sampled jointly, is described in appendix (ref). This alternative sampler is intended to reduce the dependence between the outcome equation tree samples and $\gamma$ samples. The sampler of van2011bayesian is similarly motivated.
As noted by van2011bayesian, the samples of the error terms allow for inference on the degree of dependence between the errors in the selection and outcome equations. At each iteration $t$ of the TOBART-2-NP Gibbs sampler, we may generate a pseudo-sample $\{u_{i,t} \}_{i=1}^n$ where $u_{i,t} \sim \mathcal{N}(\mu_{i,t}, \Sigma_{i,t})$ and calculate a dependence measure. The set of values of this measure across all iterations of the Gibbs sampler are an approximate sample from its posterior distribution.
iqbal2023bayesian apply the van2011bayesian prior with $n_0 = 2$ and $S_0= 2$. It is claimed that these values give an induced prior on the outcome variance $ \phi + \gamma^2$ that is similar to an Inverse Gamma. However, the induced prior mean is $ E[ \phi + \gamma^2 ] = (1 + \tau)\frac{S_0}{n_0-2} + g_0^2 $ and therefore is undefined if $n_0=2$, which it would also be for the Inverse Gamma. We note in Appendix (ref) that the prior distribution of $ \phi + \gamma^2$ is a non-trivial sum of an inverse gamma distributed variable and a (dependent) square of a location and scale shifted t-distributed variable.
Appendix (ref) contains an additional discussion of some methods for calibration of $n_0$ and $S_0$, and of potential disadvantages of the recommendation of iqbal2023bayesian to set $\tau=5$ in general. In particular, large values of $\tau$ increase the prior mean of of $\phi + \gamma^2$. Moreover, we observe that larger values of $\tau$ can result in a more peaked bimodal prior for $\rho$ with modes closer to $\pm1$.\footnote{An alternative would be to apply a hyperprior to $\tau$. We did not find this to be useful.}
We assume that the data generating process is as described for the Type 2 TOBART model. To predict the observed outcome, $Y_i$ for selected observations, we may use the following estimand: $$ E[Y_i | Z_i^* \ge 0 , \bm{x}_i, \bm{w}_i ] = f_{y}(\bm{x}_i) + E[\eta_i |f_z(\bm{w}_i) + \xi_i \ge 0 ] = f_{y}(\bm{x}_i) + \gamma \frac{\phi(f_z(\bm{w}_i) )}{\Phi(f_z(\bm{w}_i) )} $$ We expect any nonlinear machine learning method trained on the selected observations to estimate this equally as well as TOBART-2. For the nonparametric TOBART-2 DGP, we instead have $$ E[Y_i | Z_i^* \ge 0 , \bm{x}_i, \bm{w}_i, \mu_{i1} , \mu_{i2} , \Sigma_i ] = $$ $$ f_{y}(\bm{x}_i) + E[\eta_i |f_z(\bm{w}_i) + \xi_i \ge 0 , \mu_{i1} , \mu_{i2} , \Sigma_i ] = f_{y}(\bm{x}_i) + \mu_{i2} + \gamma_i \frac{\phi(f_z(\bm{w}_i) + \mu_{i1})}{\Phi(f_z(\bm{w}_i) + \mu_{i1})} $$ Then we predict the outcome by averaging over MCMC draws indexed by $d=1,\dots,D$ $$ E[Y_i | Z_i^* \ge 0, \bm{x}_i, \bm{w}_i ] \approx \frac{1}{D} \sum_{d=1}^{D} E[Y_i | Z_i^* \ge 0 , f_{y}^{(d)}, f_{z}^{(d)}, \mu_{i1}^{(d)} , \mu_{i2}^{(d)} , \Sigma_i^{(d)}, \bm{x}_i, \bm{w}_i ] = $$ $$ \frac{1}{D} \sum_{d=1}^{D} \Bigg\{ f_{y}^{(d)}(\bm{x}_i) + \mu_{i2}^{(d)} + \gamma_i^{(d)} \frac{\phi(f_z^{(d)}(\bm{w}_i) + \mu_{i1}^{(d)})}{\Phi(f_z^{(d)}(\bm{w}_i) + \mu_{i1}^{(d)})} \Bigg\} $$
If the task is to predict the latent outcome regardless of whether the test observation will be selected or not, or if there will be no selection mechanism in the test data, and instead all outcomes will be observed, then we predict $Y_i$ with the following estimand for the TOBART-2 DGP: $$ E[Y_i^* | \bm{x}_i, \bm{w}_i ] = f_{y}(\bm{x}_i) $$ Therefore, for this DGP, a machine learning method trained on the selected data would target an estimand containing a bias term $ \gamma \frac{\phi(f_z(\bm{w}_i) )}{\Phi(f_z(\bm{w}_i) )} $. In contrast, a Tobit-2-based approach, such as TOBART-2, directly models $f_y$, and we can use the estimate $\frac{1}{D} \sum_{d=1}^{D} f_{y}^{(d)}(\bm{x}_i) $ for standard TOBART-2 and $\frac{1}{D} \sum_{d=1}^{D} \Big\{ f_{y}^{(d)}(\bm{x}_i) + \mu_{i2}^{(d)} \Big\} $ for TOBART-2-NP.
Potential outcomes under treatment and control allocation are denoted by $Y_i(1)$ and $Y_i(0)$ respectively. Similarly, $Y_i^*(1)$ and $Y_i^*(0)$ denote potential latent outcomes, and $Z_i^*(1)$ and $Z_i^*(0)$ denote potential selection equation latent variable values. The binary treatment variable is denoted by $T_i$.
Generally, the estimand of interest is an effect on the latent outcome $Y_i^*$, not an effect on the observed outcome $Y_i$. If treatment affects the latent outcome, but has no effect on selection, i.e. we now have $f_y(\bm{x}_i, T_i)$ and $f_z(\bm{w}_i)$, then the estimand of interest for the TOBART-2 DGP is: $$ E[Y_i^*(1) | \bm{x}_i, \bm{w}_i ] - E[Y_i^*(0) | \bm{x}_i, \bm{w}_i ] $$ Assume that there is unconfoundedness of treatment, possibly conditional on covariates, i.e. $ Y_i^*(1), Y_i^*(0) \perp T_i | \bm{x}_i$. Then TOBART-2 directly models the above estimand as $f_y(\bm{x}_i, 1) - f_y(\bm{x}_i, 0) \approx \frac{1}{D} \sum_{d=1}^{D} f_y^{(d)}(\bm{x}_i, 1) - f_y^{(d)}(\bm{x}_i, 0) $. A nonlinear machine learning method naively trained only on selected observations targets the same estimand because \footnote{The naive method can be an “S-learner” with one trained model for all selected observations, with treatment included as a covariate, or a “T-learner” with two models separately trained on selected treatment and control observations kunzel2019metalearners.} $$ E[Y_i^*(1) | Z_i^* \ge 0, \bm{x}_i, \bm{w}_i ] - E[Y_i^*(0) | Z_i^* \ge 0, \bm{x}_i, \bm{w}_i ] = f_y (\bm{x}_i, 1) - f_y (\bm{x}_i, 0) + \gamma \Bigg( \frac{\phi(f_z(\bm{w}_i) )}{\Phi(f_z(\bm{w}_i) )} - \frac{\phi(f_z(\bm{w}_i) )}{\Phi(f_z(\bm{w}_i) )} \Bigg) $$ $$ = f_y (\bm{x}_i, 1) - f_y (\bm{x}_i, 0) . $$ Now suppose that treatment enters both the outcome and selection equations, so that we have $f_y(\bm{x}_i, T_i)$ and $f_z(\bm{w}_i, T_i)$, and assume $ Y_i^*(1), Y_i^*(0), Z_i^*(1), Z_i^*(0), \perp T_i | \bm{x}_i,\bm{w}_i$. TOBART-2 still directly models the estimand of interest $f_y (\bm{x}_i, 1) - f_y (\bm{x}_i, 0)$. However, the treatment effect estimates produced by a machine learning method naively trained only on the selected observations would target an estimand containing a selection bias term: $$ E[Y_i^*(1) | Z_i^* (1) \ge 0, \bm{x}_i, \bm{w}_i ] - E[Y_i^*(0) | Z_i^*(0) \ge 0, \bm{x}_i, \bm{w}_i ] = $$ $$ f_y (\bm{x}_i, 1) - f_y (\bm{x}_i, 0) + \gamma \Bigg( \frac{\phi(f_z(\bm{w}_i,1) )}{\Phi(f_z(\bm{w}_i,1) )} - \frac{\phi(f_z(\bm{w}_i,0) )}{\Phi(f_z(\bm{w}_i,0) )} \Bigg) $$ Therefore TOBART-2 is expected to produce more accurate ITE estimates for DGPs in which treatment affects sample selection.
In addition to the simulations described here, we provide a simulation study with the linear DGP of iqbal2023bayesian in Appendix (ref), although this DGP does not include an excluded instrument.
We simulate the data generating processes described by brewer2024addressing, with less observations and variables for computational feasibility. We also consider non-normally distributed errors. Covariates $X_1,...,X_{10}$ are generated from a multivariate normal distribution with $\text{Cov}(X_{ki},X_{ji}) = 0.3^{|k-j|}$. The excluded instrument, $W_{exc,i}$, that appears in the selection equation, but not the outcome equation is generated by $W_{exc,i} = \sum_{j=1}^{10} 0.05 X_{ji} + e_i \ , \ e_{w,i} \sim \mathcal{N}(0,0.75) $.
We consider the following DGPs, and two more DGPs in Appendix (ref):
Appendix (ref) includes results for the above DGPs with the following non-normal distributions for the errors. \footnote{TOBART-2-NP results will be added to an updated version of this paper at a later date.}
The numbers of training and test observations are set to $2500$ and $500$ respective;y, and the number of repetitions for each simulation scenario is set to $10$.
When designing a simulation study to evaluate heterogeneous treatment effect estimation, we consider the following:
We generate $p$ independent standard normally distributed covariates $X_{1},\dots, X_p \sim \mathcal{N}(0,1)$. A binary treatment variable is generated with selection on observables. The number of covariates is $p=50$, and $X_1, \dots, X_5$ are excluded from the outcome equation data, i.e. they are in $\bm{w}_i$, but not $\bm{x}_i$. There are 4000 training observations and 4000 test observations. The propensity score function is: $$\pi(\bm{x}_i) = 0.7 \Phi(0.5 x_{6,i}x_{7,i} + 0.3 x_{8,i}^2) +0.1 $$ The treatment variable is denoted by $T_i \sim \text{Bernoulli}(\pi(\bm{x}_i))$. The selection and outcome equations are: $$ Y_i^* = 0.5x_{6,i} + 0.5 x_{6,i}x_{7,i} + 0.25 (x_{3,i}-1)^2 T_i + \eta_i $$ $$ Z_{i}^* = 0.5 x_{1,i}x_{6,i} + x_{1,i}x_{6,i}T_i + \xi_i $$ This implies that the heterogeneous effects on the inverse Mills ratio are equal to: $$ \frac{\phi(f_z(\bm{w}_i,1) )}{\Phi(f_z(\bm{w}_i,1) )} - \frac{\phi(f_z(\bm{w}_i,0) )}{\Phi(f_z(\bm{w}_i,0) )} = \frac{\phi (0.5 x_{1,i}x_{6,i} + x_{1,i}x_{6,i} ) }{ \Phi(0.5 x_{1,i}x_{6,i} + x_{1,i}x_{6,i} ) } - \frac{\phi (0.5 x_{1,i}x_{6,i} ) }{ \Phi(0.5 x_{1,i}x_{6,i} ) } $$ The heterogeneous effects on the latent outcome are equal to: $$ f_y (\bm{x}_i, 1) - f_y (\bm{x}_i, 0) = 0.25 (x_{3,i}-1)^2 $$ Let the potential outcomes of the selection indicator variable under treatment and control allocation be denoted by $S_i(1)$ and $S_i(0)$ respectively. We use the following measures to evaluate all methods:
$\bullet$ MSE of estimated effect on the probability of censoring $$ \frac{1}{n_{test}} \sum_{i=1}^{n_{test}} \Big( \big(\hat{\Pr}(S_i(1)=0|X) - \hat{\Pr}(S_i(0)=0|X) \big) - \big({\Pr}(S_i(1)=0|X) - {\Pr}(S_i(0)=0 |X) \big)\Big)^2 $$ $\bullet$ RMSE of estimated treatment effect on the latent outcome: $$ \sqrt{\frac{1}{n_{test}} \sum_{i=1}^{n_{test}} \Big( \big(\hat{\mathbb{E}}(Y_i(1)=0|X) - \hat{\mathbb{E}}(Y_i(0)=0|X) \big) - \big({\mathbb{E}}(Y_i(1)=0|X) - {\mathbb{E}}(Y_i(0)=0 |X) \big)\Big)^2 }$$
Table (ref) contains the results for the simulation study from brewer2024addressing.\footnote{TOBART-2-NP results will be added to an updated version of this paper at a later date.} As anticipated, the sparse linear Bayesian model described by iqbal2023bayesian produces better results for DGP 1, because a linear model is appropriate for a DGP with linear selection and outcome equations. BART and TOBART do not improve estimates of sample selection probabilities relative to the sparse linear model, because for both DGPs the true model for the latent selection variable is linear.
Latent outcome predictions for the nonlinear DGP2 are notably more accurate for Soft BART and TOBART models than for linear Tobit models or BART. When there is no sample selection, there is no notable difference between the SoftBART and Soft TOBART models. When there is sample selection in unobservables, TOBART models produce much more accurate nonlinear DGP predictions than the other models. Moreover, confidence interval coverage is close to 95% and posterior means of the correlation between selection and outcome equation errors are very close to the true values.
The best performing TOBART models are marginalized TOBART-2 and Soft TOBART-2. A combination of the two models, i.e. marginalized soft TOBART-2, produces similar estimated to unmarginalized Soft TOBART-2, although it produces more accurate estimates of $\rho$ when $\rho=0.45$.\footnote{Marginalized SoftBART surprisingly produces much less accurate esitmates of selection probabilities. Perhaps this is explainable by differences in the implementations of unmarginalized and marginalized tree samplers.} We do not observe large differences in results for different priors on the distribution of the error terms, and different calibrations of these priors. Perhaps the methods would be more sensitive to the prior if the sample size were much smaller than 2000.
Table (ref) contains the results for the treatment effect simulation study. Nonlinear methods outperform linear methods for this simple quadratic DGP with interactions. Tree-based methods with soft splitting are more accurate than their counterparts with hard splitting rules. Marginalized TOBART outperforms vanilla implementations of TOBART, suggesting that without marginalization TOBART sampler did not converge or mix well.
Surprisingly, SoftBART with the inverse Mills ratio as a covariate, following the general approach of brewer2024addressing, is almost as accurate as Soft TOBART-2 across all measures except PEHE of effects on selection and outcome CATE posterior interval length. Soft marginalized TOBART-2 produces the best estimates of outcome treatment effects, and surprisingly the worst predictions of selection treatment effects. It also produces the smallest posterior intervals, albeit with over-coverage.
\FloatBarrier
We estimate the heterogeneous effects of allocation to a deductible plan for health insurance. The dataset has been studied by deb2002structure, cameron2005microeconometrics, and others. The outcome variable is the log of medical expenses in dollars. Selected observations are those with non-zero medical expenditures. There are 15 other covariates.\footnote{The other covariates are: the log of one plus the coinsurance rate, the log of the participation incentive payment, physical limitations (binary), number of chronic diseases, good/fair/poor self-rated health (binary variables), log of family income, log of family size, education of household head in years, age in years, female (binary), child (1 if has child below 18 years), interaction of female and child variables, black (binary).} None of the covariates are excluded from the outcome equation. Further applications to female earning data mroz1987sensitivity and a job training experiment for disadvantaged youths schochet2001national are included in Appendix (ref).
The data application results are presented below.\footnote{TOBART-2-NP results will be added to an updated version of this paper at a later date.} Table (ref) shows that the posterior mean and 95% posterior interval for the correlation parameter $\rho$ are comparable to the estimates produced by van2011bayesian. Table (ref) contains estimates of average treatment effects of individual deductible plans on probabilities of zero expenditures and the log of medical expenditures. The average treatment effect estimates are larger in magnitude for the Bayesian methods than for standard Tobit. The average treatment effect on the latent outcome estimate produced by Soft-TOBART-2 is remarkably close to the effect estimated by van2011bayesian with a Bayesian linear Tobit model.
Appendix (ref) contains 5-fold cross validated Brier scores of selection probabilities, MSE of observed outcome predictions, and coverage of observed outcome posterior intervals. All models produce similar results. This is unsurprising because sample selection models are not expected to produce better predictions of the observed outcomes, but rather better predictions of the mean of latent outcomes, which cannot be evaluated using observed sample selected data.
\FloatBarrier
\FloatBarrier
\FloatBarrier
The results of this paper suggest that TOBART-2 can produce more accurate outcome predictions and treatment effect estimates than linear sample selection models when the DGP is non-linear. However, for some DGPs, a naive application of BART or softBART to the selected data, with an estimate of the inverse Mills ratio as a covariate, also provides estimates with comparable accuracy.
This paper also describes alternative MCMC implementations and prior settings. Improvements in predictive accuracy can be obtained by the use of soft splitting rules and marginalization in MCMC draws of trees. Prior sensitivity, particularly for small sample sizes, would be an interesting area for future research. Possible extensions of the model include heteroskedastic type 2 Tobit models donald1995two, de2022generalized, with modelling of variances by Bayesian trees pratola2020heteroscedastic.