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.
52,519 characters · 24 sections · 138 citation commands
Type I Tobit Bayesian Additive Regression Trees for Censored Outcome Regression
\email{[email removed]}
\affil*[1]{ \orgdiv{Econometric Institute}, \orgname{Erasmus University Rotterdam} }
\abstract{Censoring occurs when an outcome is unobserved beyond some threshold value. Methods that do not account for censoring produce biased predictions of the unobserved outcome. This paper introduces Type I Tobit Bayesian Additive Regression Tree (TOBART-1) models for censored outcomes. Simulation results and real data applications demonstrate that TOBART-1 produces accurate predictions of censored outcomes. TOBART-1 provides posterior intervals for the conditional expectation and other quantities of interest. The error term distribution can have a large impact on the expectation of the censored outcome. Therefore, the error is flexibly modeled as a Dirichlet process mixture of normal distributions. An R package is available at \url{https://github.com/EoghanONeill/TobitBART}.}
Censoring occurs when, beyond some threshold value, the observed outcome is equal to the threshold instead of the true latent outcome value. For example, scientific equipment can often only make accurate measurements within a known range of outcome values, and observations outside this range are set to its limits. Often the estimand of interest is the conditional expectation or conditional average treatment effect on the outcome before censoring. Estimation of a standard regression model using data without censored values, or with censored observations set equal to threshold values, results in biased estimates. Tobit models directly model the latent outcome and censoring process tobin1958estimation.
In this paper, we combine the Bayesian Type I Tobit model chib1992bayes with Bayesian Additive Regression Trees chipman2010bart. The latent outcome (before censoring) is modeled as a sum-of-trees, which allows for nonlinear functions of covariates. The error term is modeled as a Dirichlet process mixture of normal distributions, as in fully nonparametric BART george2019fully. Smooth data generating processes with sparsity are modelled by soft trees with a Dirichlet prior on splitting variable probabilities, as introduced by linero2018bayesianB.
In simulations and applications to real data, TOBART-1 outperforms a Tobit gradient boosted tree method, Grabit sigrist2019grabit, a Tobit Gaussian Process model groot2012gaussian, standard linear Tobit, and simple hurdle models based on standard machine learning methods. Unlike other methods, TOBART-1 accounts for model uncertainty and can non-parametrically model the error term. Posterior intervals are available for censored outcomes, uncensored outcomes, conditional expectations, and probabilities of censoring. Grabit, Gaussian Processes, and other methods rely on cross-validation for parameter tuning and are sensitive to the tuned variance of the error term, whereas TOBART-1 performs well without parameter tuning and accounts for uncertainty in the variance of the error term.
TOBART-1 with a Dirichlet process mixture of normal distributions for the error term (TOBART-1-NP) removes the restrictive normality assumption often imposed in censored outcome models. We observe that this can lead to more accurate outcome predictions in simulations with non-normally distributed errors, and in real data applications, which may involve non-normally distributed outcomes.\footnote{A Dirichlet process mixture for the error term distribution has previously been included in a censored outcome model by kottas2009bayesian.}
A variety of methods have been proposed for nonparametric and semiparametric censored outcome models. lewbel2002nonparametric describe a local linear kernel estimator for the setting in which both the uncensored outcome mean function of regressors and error distribution are unknown. fan1994censored describe a quantile-based local linear approximation method. huang2021novel introduces a semiparametric method involving B-splines. chen2005nonparametric use a local polynomial method. Other papers on the topic of semiparametric and nonparametric censored outcome regression include cheng2021semiparametric, heuchenne2007location, heuchenne2010estimation, huang2019estimation; and oganisian2021bayesian. Gaussian Process censored outcome regression methods are applied by groot2012gaussian, cao2018model,gammelli2020estimating, gammelli2022generalized and basson2023variational. zhang2021deep and wu2018deep implement censored outcome neural network methods.
A number of recent papers have considered Tobit model selection and regularization. zhang2012focused describe Focused Information Criteria based Tobit model selection and averaging. jacobson2022high provide theoretical and empirical results for Tobit with a Lasso penalty and a folded concave penalty (SCAD). muller2016censored and soret2018lasso describe a LASSO penalized censored outcome models. bradic2016robust study robust penalized estimators for censored outcome regression.
The Bayesian Tobit literature includes quantile regression methods ji2012model, yu2007bayesian, alhamzawi2016bayesian, and Bayesian elastic net Tobit alhamzawi2020new. ji2012model account for model uncertainty by implementing Tobit quantile regression with Stochastic Search Variable Selection. However, the outcome and latent variable are modeled as linear functions of covariates. TOBART-1 provides a competing approach to the methods referenced above that does not impose linearity.
The remainder of the paper is structured as follows: In section (ref) we describe the TOBART-1 model and Markov chain Monte Carlo (MCMC) implementation, section (ref) contains simulation studies for prediction and treatment effect estimation with censored data, section (ref) contains applications to real world data, and section (ref) concludes the paper.
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$.
For the standard BART model, the outcome is determined by a sum of trees, $$Y_i = \sum_{j=1}^m g(x_i ; T_j, M_j)+\varepsilon_i$$ 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$, and $\varepsilon_i \overset{i.i.d}{\sim} N(0, \sigma^2)$.
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 form of the prior used by chipman2010bart is $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) $ where $\mu_{kj} | T_j \overset{i.i.d}{\sim} N(0,\sigma_{\mu}^2)$ where $\sigma_{\mu} = \frac{0.5}{\kappa \sqrt{m}}$ and $\kappa$ 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. 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$.
Samples from $p((T_1, M_1),...,(T_m,M_m), \sigma | y)$ can be made by a Bayesian backfitting MCMC algorithm. This algorithm involves $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$. After burn-in, the sequence of $f^*$ draws, $f_1^*,...,f_Q^*$, where $f^*(.)= \sum_{j=1}^m g(. \ ; T_j^*, M_j^*)$, is an approximate sample of size $Q$ from $p(f|y)$.
In addition to the standard Bayesian tree model for $f(\bm{x}_i)$ described in section (ref), we also implement TOBART and TOBART-NP with soft trees and sparse splitting rules as described by linero2018bayesianB. Predictions from soft trees are weighted linear combinations of all terminal node parameter values, with the weights being functions of distances between covariates and splitting points. The prediction from a single tree function is $$ g(\bm{x}_i ; T_j, M_j) = \sum_{\ell = 1}^{L_j} \mu_{j,\ell} \xi (\bm{x}_i, T_j, \ell) $$ $$\xi (\bm{x}_i, T_j, \ell) = \prod_{b \in \mathcal{A}(\ell)} \zeta \left( \frac{ x_{{j_b}} - C_b}{\tau_b} \right)^{ \mathbb{I} \{ x_{j_b} > C_b \} } \times $$ $$\Big\{ 1 - \zeta \left( \frac{ x_{{j_b}} - C_b}{\tau_b} \right) \Big\}^{ \mathbb{I} \{ x_{j_b} \le C_b \} } $$ where $L_j$ is the number of leaves in the $j^{th}$ tree, $\mu_{j,\ell}$ is the $\ell^{th}$ terminal node parameter of the $j^{th}$ tree, $\mathcal{A}(\ell)$ denotes the set of ancestor nodes of terminal node $\ell$. The splitting variable, splitting point, and bandwidth parameter at internal node $b$ are denoted by $x_{{j_b}}$, $C_b$, and $\tau_b$ respectively. The gating function $\zeta $ is the logistic function $ \zeta(x) = (1+\exp(-x))^{-1}$.
Sparse splitting rules are introduced by placing a Dirichlet prior on the splitting probabilities $(s_1,\dots, s_p) \sim \mathcal{D} (\frac{a}{p},\dots, \frac{a}{p})$. The parameter $a$ controls the level of sparsity and has the prior $\text{Beta}(0.5,1)$. linero2018bayesianB demonstrate that soft trees allow BART to model smooth functions, and the Dirichlet prior on splitting probabilities adapts to unknown levels of sparsity to provide improved predictions on high dimensional data sets.
The Type I Tobit model with censoring from below at $a$ and censoring from above at $b$ is: $$ Y_i^* = \bm{x}_i \bm{\beta} + \varepsilon_i \ , \ \varepsilon_i \sim i.i.d. \ N(0, \sigma^2) $$ $$ Y_i =
$$ where a normal prior is placed on $\beta$, and an inverse gamma prior is placed on $\sigma^2$ chib1992bayes.
The Type I TOBART model replaces the linear combination $\bm{x}_i \bm{\beta}$ with the sum-of-trees function $f(\bm{x}_i)$ : $$ Y_i^* = f(\bm{x}_i) + \varepsilon_i \ , \ \varepsilon_i \sim i.i.d. \ N(0, \sigma^2) $$ $$ Y_i =
$$ where a BART prior is placed on $ f(\bm{x}_i)$ and an inverse gamma prior is placed on $\sigma^2$. \footnote{$\sigma^{-2} \sim Ga(\frac{v}{2}, \frac{v \lambda}{2})$. For standard BART, $\lambda$ is set such that the $q^{th}$ quantile of the prior distribution of $\sigma$ is the sample standard deviation of the residuals from a linear model. For censored outcomes, this may give poor calibration of the $\sigma$ prior. We consider four options in a simulation study in appendix \ref{tbarts_sims_app}. A sample standard deviation estimate from an intercept-only Tobit model generally gives good results, although often there is little difference across $\lambda$ values.}
Tobit can be implemented by MCMC with data augmentation chib1992bayes. The realization, $y_i^*$, of the variable $Y_i^*$ is observed for uncensored outcomes, and is sampled from its full conditional for censored outcomes. $$y_i^* = y_i \text{ if } y_i \in (a,b) \ , \ \text{ and } $$ $$y_i^* \sim
$$ where $\mathcal{TN}_{[l,u]}$ denotes a normal distribution truncated to the interval $[l,u]$. The full conditionals for $f(\bm{x}_i)$ and $\sigma^2$ are standard full conditionals for BART with $y_i^*$ as the dependent variable and $\bm{x}_i$ as the potential splitting variables. Appendix \ref{mcmc_app} contains a description of a sampler that produces draws $f^{(1)}(\bm{x}_i), \hdots,f^{(D)}(\bm{x}_i)$ and $\sigma^{(1)},\hdots,\sigma^{(D)}$.
The conditional mean of the latent variable is $f(\bm{x}_i)$. If censoring is also applied to the test data, then the outcomes are predicted by averaging the standard Tobit expectation formula across MCMC iterations:
For all MCMC iterations $d=1,...,D$ calculate $$ E[Y_i|X_i=\bm{x}_i, f^{(d)}, \sigma^{(d)}] = a \Phi \Big(\frac{a - f^{(d)}(\bm{x}_i)}{\sigma^{(d)}} \Big) + $$ $$ f^{(d)}(\bm{x}_i) \Bigg[ \Phi \Big(\frac{b - f^{(d)}(\bm{x}_i)}{\sigma^{(d)}} \Big) - \Phi \Big(\frac{a - f^{(d)}(\bm{x}_i)}{\sigma^{(d)}} \Big) \Bigg] $$ $$ + \sigma^{(d)} \Bigg( \phi \Big(\frac{a - f^{(d)}(\bm{x}_i)}{\sigma^{(d)}} \Big) - \phi \Big(\frac{b - f^{(d)}(\bm{x}_i)}{\sigma^{(d)}} \Big) \Bigg) + $$ $$ b \Bigg[ 1 - \Phi \Big(\frac{b - f^{(d)}(\bm{x}_i)}{\sigma^{(d)}} \Big) \Bigg]$$ The predicted outcome is $\frac{1}{D} \sum_{d=1}^D E[Y_i|X_i=\bm{x}_i, f^{(d)}, \sigma^{(d)}]$. The expectation conditional on the outcome not being in the censored range is: $$E[Y_i| a < Y_i < b, X_i=\bm{x}_i, f^{(d)}, \sigma^{(d)}] = $$ $$ f^{(d)}(\bm{x}_i) + \sigma^{(d)} \frac{ \phi \Big(\frac{a - f^{(d)}(\bm{x}_i)}{\sigma^{(d)}} \Big) - \phi \Big(\frac{b - f^{(d)}(\bm{x}_i)}{\sigma^{(d)}} \Big) }{ \Phi \Big(\frac{b - f^{(d)}(\bm{x}_i)}{\sigma^{(d)}} \Big) - \Phi \Big(\frac{a - f^{(d)}(\bm{x}_i)}{\sigma^{(d)}} \Big) } $$
The accuracy of the conditional expectation of the TOBART model depends on the validity of the assumption of normality of the errors. More general censored outcomes can be modelled by assuming a Dirichlet Process mixture distribution for the error terms. $$ y_i^* = f(\bm{x}_i) + \varepsilon_i \ ,\ y_i =
$$ $$\varepsilon_i \sim i.i.d. \ N(\gamma_i, \sigma_i^2) \ , \ \vartheta_i = (\gamma_i, \sigma_i) \sim G $$ $$ G \sim \mathcal{DP}(G_0, \alpha)$$ The distribution of the error term is specified similarly to \cite{george2019fully}. The base distribution $G_0$ is defined as follows: $$ p(\gamma, \sigma | \nu, \lambda_1, \gamma_0, k_0) = p(\sigma | \nu, \lambda) p(\gamma | \sigma, \gamma_0, k_0) $$ $$ \sigma^2 \sim \frac{\nu \lambda}{\chi_{\nu}^2} \ , \ \gamma | \sigma \sim \mathcal{N} \Big(\gamma_0, \frac{\sigma^2}{k_0} \Big) $$ where, in contrast to the standard BART prior of \cite{chipman2010bart}, $\nu$ is set to $10$ instead of $3$.\footnote{\cite{george2019fully} recommend $\nu=10$ as the spread of the error increases when there are many components and the spread of a single components can be reduced by increasing $\nu$. This gives better results than $\nu=3$ for some DGPs in a simulation study in Appendix \ref{tbarts_sims_app}.} The parameter $\lambda$ is set such that the $q^{th}$ quantile of the prior distribution of $\sigma$ is the sample standard deviation of the outcome, or of the residuals from a linear model. For TOBART-NP, $q=0.9$ instead of $0.95$. \footnote{This is complicated by the censoring of the outcome. Some options are: 1. Estimate the standard deviation assuming that censored outcome is normally distributed. 2. Estimate the standard deviation of a linear type I Tobit model (contains option 1 as a special case but not feasible when there are more regressors than observations). 3. Estimate the standard deviation of the censored outcome without accounting for censoring. We use option 2 for TOBART-NP. } The prior on $\alpha$ is the $\alpha \sim \Gamma(2,2)$ prior introduced by \cite{escobar1995bayesian} and applied by \cite{van2011bayesian}.\footnote{The TobitBART package also includes an option for the prior described by \cite{rossi2014bayesian} and \cite{george2019fully}, $p(\alpha) \propto \left( 1 - \frac{\alpha - \alpha_{min}}{\alpha_{max} - \alpha_{min}} \right)^{\psi}$, where $\alpha_{min}$ and $\alpha_{max}$ are set so that the modal numbers of components are $I_{min} = 1$ and $I_{max} = [(0.1)n]$ respectively, and $\psi = 0.5$.}
The outcome is scaled by subtracting the sample mean before applying the Gibbs sampler, therefore george2019fully set $\gamma_0 = 0$.\footnote{However, the mean cannot be estimated for censored data without making further assumptions. Options include: 1. Estimate the mean (and variance) of a censored normal distribution. 2. Calculate the sample mean of the censored outcome without accounting for censoring. We use option 1. } The parameter $k_0$ is scaled with the marginal distribution of $\gamma$ ( $\gamma \sim \frac{\sqrt{\lambda } }{ \sqrt{k_0 } } t_{\nu} $). Given $k_s$ (set to $10$ by default), $k_0$ is set such that $ \max_{i=1,...,n} |e_i| = k_s \frac{\sqrt{\lambda } }{ \sqrt{k_0 } } $ where $k_s = 10$. and $e_1,...,e_n$ are the residuals from a linear model.\footnote{The residuals likely underestimate the true errors for censored observations.} The Gibbs sampler for TOBART-NP is described in Appendix (ref).
For each MCMC iteration, $d$, and observation $i$, we obtain $\vartheta_i^{(d)} = (\gamma_i^{(d)}, \sigma_i^{(d)})$. The conditional expectation, $E[y_i|\bm{x}_i, f^{(d)}, \gamma_i^{(d)}, \sigma^{(d)}] $, is calculated as outlined in section (ref).
Let a binary variable $T_i$ equal $1$ if unit $i$ is assigned to treatment and $0$ if $i$ is assigned to the control group. The potential outcomes under treatment and control group allocation are denoted by $Y_i(1)$ and $Y_i(0)$ respectively. Similarly, the potential outcomes of the latent outcome are denoted by $Y_i^*(1), Y_i^*(0)$. Assume the data generating process is as follows: $$ Y_i^* = \mu(\bm{x}_i) + \tau(\bm{x}_i) T_i + \varepsilon_i \ , \ \varepsilon_i \sim \mathcal{N}(0,\sigma^2) $$ $$ Y_i =
$$ where $\mu(\bm{x}_i)$ and $\tau(\bm{x}_i)$ are possibly nonlinear functions of covariates. Assume conditional unconfoundedness, i.e. $Y_i^*(1),Y_i^*(0) \perp T_i | X_i$ . The estimand is the conditional average treatment effect on $Y_i^*$, i.e., $ E[Y_i^*(1) - Y_i^*(0) | X_i = \bm{x}_i] = \tau(\bm{x}_i) $. However, a model naively trained on only uncensored outcomes estimates the following effects \footnote{This bias occurs if all the uncensored observations are included in one regression and differences in predictions for $T_i=1$ and $T_i=0$ are obtained, i.e. an S-learner approach \citep{kunzel2019metalearners}, or if the two conditional expectations are obtained from separate regressions for treated and untreated uncensored observations, i.e. a T-Learner approach. In both cases, the conditional expectations are not equal to the expectation of the latent outcome.} $$ E[Y_i(1) | a < y_i < b, X_i = \bm{x}_i] - $$ $$ E[Y_i(0) | a < y_i < b, X_i = \bm{x}_i] = \tau(\bm{x}_i) + $$ $$ \sigma \Bigg( \frac{ \phi \Big(\frac{a - ( \mu(\bm{x}_i) + \tau(\bm{x}_i)))}{\sigma} \Big) - \phi \Big(\frac{b - ( \mu(\bm{x}_i) + \tau(\bm{x}_i)) }{\sigma} \Big) }{ \Phi \Big(\frac{b - ( \mu(\bm{x}_i) + \tau(\bm{x}_i)) }{\sigma} \Big) - \Phi \Big(\frac{a - ( \mu(\bm{x}_i) + \tau(\bm{x}_i)) }{\sigma} \Big) } - $$ $$\frac{ \phi \Big(\frac{a - \mu(\bm{x}_i)}{\sigma} \Big) - \phi \Big(\frac{b - \mu(\bm{x}_i)}{\sigma} \Big) }{ \Phi \Big(\frac{b - \mu(\bm{x}_i)}{\sigma} \Big) - \Phi \Big(\frac{a - \mu(\bm{x}_i)}{\sigma} \Big) } \Bigg) \ . $$ A sufficiently flexible nonparametric method, without restrictive assumptions on the error term, will produce estimates that approximate the expression above. A model naively trained on the full data set with censoring similarly gives biased estimates (see Appendix \ref{TEbias_app}). By directly modelling $Y_i^*$, censored outcome models avoid the bias described above. Similar biases occur if the error term is not normally distributed.
We adapt the data generating processes (DGPs) introduced by friedman1991multivariate to a censored regression setting. This DGP has often been applied in comparisons of semiparametric regression methods. We also make use of the censored outcome simulations described by groot2012gaussian, sigrist2019grabit, and jacobson2022high for fair comparison against competing methods with existing synthetic censored data.
The covariates $x_1,....,x_p$ are independently sampled from the uniform distribution on the unit interval. The outcome before censoring is generating from one of the following functions:
The variance of the error, $\sigma^2$, is set to 1. See the Supplementary Appendix (Online Resource 1) for the results obtained from simulations with $\sigma \in \{0.1, 2\}$. We also consider deviations from the assumption of normally distributed errors. In particular, we include results for simulations in which $\varepsilon$ is generated from Skew-t, and $\text{Weibull}(1/2, 1/5)$ distributions.\footnote{bradic2016robust considered $\text{Weibull}(1/2, 1/5)$ errors in a simulation study.} The number of covariates, $p$, is set to $30$. We generate 500 training and 500 test observations.
We compare the performance of TOBART-1, TOBART-1-NP, Soft TOBART-1, and Soft TOBART-1-NP against Grabit sigrist2019grabit, linear Tobit tobin1958estimation, BART chipman2010bart, Random Forests (RF) breiman2001random, Gaussian Processes, and a Tobit Gaussian Process model groot2012gaussian.\footnote{Standard BART for continuous outcomes is trained on censored outcomes. Probit BART is trained on a binary variable indicating censorship. Similarly, Random Forests are separately trained on continuous censored outcomes and a binary censorship indicator.} The results for a Gaussian Process (GP) with only 5 variables (always including all informative variables) are included because GPs were observed to produce inaccurate predictions when applied to data with 30 variables.\footnote{The GP Matlab code was obtained from \url{https://www.cs.ru.nl/ perry/software/tobit1.html}.} Censored outcome predictions are evaluated using Mean Squared Error (MSE), and predicted probabilities of censoring are evaluated using the Brier Score.\footnote{See the Supplementary Appendix (Online Resource 1) for implementation details and parameter settings.} \footnote{Latent outcome predictions similarly demonstrate that TOBART outperforms other methods, and these results are available on request. However, it is unsurprising that Tobit based latent outcome predictions outperform naive approaches due to the aforementioned censoring bias. } All results are averaged over 5 repetitions. \footnote{Computational times are included in Appendix (ref).}
The results for simulations with normally distributed errors are presented in Tables (ref) and (ref). The TOBART algorithms generally outperform competing methods across all DGPs, except unsurprisingly for the linear jacobson2022high simulations linear Tobit is outperformed only by Soft TOBART. TOBART-NP can slightly improve on TOBART in some cases, but generally the results are similar when errors are normally distributed. The differences in criteria across methods are small for the more linear DGPs from sigrist2019grabit and jacobson2022high, as linear Tobit is designed for a linear DGP, and the nonlinear methods BART and RF can model the relatively simple response surface well. It is worth noting that TOBART outperforms Grabit even though the true standard deviation, $\sigma=1$, is included as one of five possible Grabit hyperparameter values in cross-validation. The same pattern of results can be observed for simulations with $\sigma =0.1$ and $\sigma = 2$ in the Supplementary Appendix. The Supplementary Appendix contains comparisons of Area Under the Curve for all methods and DGPs, from which similar conclusions can be drawn.
The results for Skew-t and Weibull distributed errors are also presented in Tables (ref) and (ref).\footnote{Results for t-distributed errors with $\nu=3$ are in the Supplementary Appendix.} The TOBART models outperform all other methods for almost all DGPs and criteria. The results for the Weibull distribution generally favour TOBART-NP and Soft TOBART-NP, indicating that there is some improvement from the Dirichlet Process model when the errors are sufficiently non-Gaussian.
The average coverage and length of $95\%$ prediction intervals for the latent outcomes and the observed outcomes are given in the Supplementary Appendix (Online Resource 1). For most DGPs and error distributions, TOBART and Soft TOBART provides the closest to 95% coverage of prediction intervals for both latent and observed outcomes. For some DGPs with non-normal errors, the more conservative intervals produced by TOBART-NP and Soft TOBART-NP provide better coverage.
A number of recent simulation studies have demonstrated that BART is among the most accurate treatment effect estimation methods wendling2018comparing, mcconnell2019estimating, dorie2019automated,hahn2019atlantic. However, in practice many data sets, including randomized trial data sets, contain censored outcomes. For example, antibody concentrations or environmental levels of chemicals can only be measured accurately within a certain range as a result of limitations of measuring equipment. Often economic data is censored due to privacy considerations, for example income might be censored above a certain threshold. TOBART provides a machine learning treatment effect estimation method with uncertainty quantification that can be applied to this data while still making use of the information provided by censored observations. We demonstrate the effectiveness of TOBART by censoring the outcomes of DGPs from published studies of machine learning methods for treatment effect estimation. The chosen data generating processes contain linear and non-linear functions of covariates, constant and heterogeneous effects, and various degrees of confounding.
$P=10$ covariates are generated from a multivariate Gaussian distribution, $X_1,\hdots, X_{10} \sim \mathcal{MVN}(\bm{0}, \Sigma)$, with $\Sigma_{jk} = 0.6^{|j-k|} + 0.1 \mathbb{I}(j \neq k) $. The binary treatment variable is Bernoulli distributed, $Z_i \sim \text{Bern}(\pi(\bm{x}_i)) $, where $$ \pi(\bm{x}_i) = \Phi(-0.4 + 0.3 X_{i,1} + 0.2 X_{i,2} ) $$ and $\Phi(\cdot)$ is the cumulative distribution function of the standard normal distribution.
The prognostic score function, $\mu(\bm{x}_i)$, and CATE function, $\tau(\bm{x}_i)$, are defined as $$ \mu(\bm{x}_i) = 3 + X_{i,1} + 0.8 \sin ( X_{i,2} ) + 0.7 X_{i,3} X_{i,4} - X_{i,5}$$ $$ \tau(\bm{x}_i) = 2 + 0.8 X_{i,1} - 0.3 X_{i,2}^2 $$ The outcome before censoring is generated as: $$ Y_i^* = \mu(\bm{x}_i) + \tau(\bm{x}_i) Z_i + \varepsilon_i \ , \ \text{where} \ \varepsilon_i \sim \mathcal{N}(0,1) $$ The number of sampled observations is 200. The observed outcome $Y_i$ is censored from below at the $15^{th}$ percentile of the generated $Y_i^*$ values, and from above at the $85^{th}$ percentile.
$P=20$ covariates are generated from independent standard uniform distributions $X_1,...,X_{20} \sim \mathcal{U}[0,1]$. There is no confounding as $\pi(\bm{x}_i)=0.5$ and $Z_i \sim \text{Bern}(\pi(\bm{x}_i)) $. The prognostic score function, $\mu(\bm{x}_i)$, and CATE function, $\tau(\bm{x}_i)$, are defined as $ \mu(\bm{x}_i) = 0 $ and $$ \tau(\bm{x}_i) = \Bigg( 1 + \frac{1}{1 + \exp \Big(-20(X_{i,1} - \frac{1}{3}) \Big)} \Bigg) \times $$ $$\Bigg( 1 + \frac{1}{1 + \exp \Big(-20(X_{i,2} - \frac{1}{3}) \Big)} \Bigg). $$ The outcome before censoring is generated as: $$ Y_i^* = \mu(\bm{x}_i) + \tau(\bm{x}_i) Z_i + \varepsilon_i \ , \ \text{where} \ \varepsilon_i \sim \mathcal{N}(0,1) $$ The number of sampled observations is 200. The observed outcome $Y_i$ is censored from below at the $15^{th}$ percentile of the generated $Y_i^*$ values, and from above at the $85^{th}$ percentile.
The covariates are generated as follows across scenarios A to D. In simulation A, $X_1,...,X_{12} \sim \mathcal{U}[0,1]$. In simulations B to D, $X_1,...,X_{12} \sim \mathcal{N}(0,1)$.
$\pi(\bm{x}_i)$ is defined as follows across scenarios A to D: (A) $ \text{trim}_{0.1} \{ \sin(\pi X_{i,1} X_{i,2} ) \} $, (B) constant equal to $0.5$, (C) $1/\{1 + \exp(X_{i,2} + X_{i,3} )\}$, (D) $1/\{1 + \exp(-X_{i,1}) + \exp( - X_{i,2} )\}$.
$\mu(\bm{x}_i)$ is defined as follows across scenarios A to D: (A) $\sin(\pi X_{i,1} X_{i,2}) + 2 (X_{i,3}-0.5)^2 + X_{i,4} + 0.5 X_{i,5} $, (B) $\max \{X_{i,1} + X_{i,2}, X_{i,3} ,0 \} $, (C) $2 \log \{ 1 + \exp ( X_{i,1} + X_{i,2} + X_{i,3} ) \} $, (D) $\frac{1}{2} [ \max \{ X_{i,1} + X_{i,2} + X_{i,3} ,0 \} + \max \{ X_{i,4} + X_{i,5} ,0 \} ] $.
$\tau(\bm{x}_i)$ is defined as follows across scenarios A to D: (A) $ ( X_{i,1} + X_{i,2})/2$, (B) $ X_{i,1} + \log \{1 + \exp( X_{i,2}) \}$, (C) constant equal to $ 1 $ , (D) $\max \{ X_{i,1} + X_{i,2} + X_{i,3} ,0 \} - \max \{ X_{i,4} + X_{i,5} ,0 \}$.
The outcome before censoring is generated as: $$ Y_i^* = \mu(\bm{x}_i) + \tau(\bm{x}_i) (Z_i-0.5) + \varepsilon_i \ , \ \text{where} \ \varepsilon_i \sim \mathcal{N}(0,1) $$ The number of sampled observations is 200. The observed outcome $Y_i$ is censored from below at the $15^{th}$ percentile of the generated $Y_i^*$ values, and from above at the $85^{th}$ percentile.
All methods are evaluated in terms of Precision in Estimation of Heterogeneous Effects (PEHE), which is defined as $\frac{1}{N}\sum_{i=1}^N (\hat{\tau}(\bm{x}_i) - \tau(\bm{x}_i) )^2$ . Confidence intervals are evaluated in terms of average coverage of $95\%$ intervals and average length of intervals.
The results are presented in Table (ref). For all DGPs, at least one TOBART method attains lower PEHE than all other methods, often by a large margin. Local Linear Forests friedberg2020local attain similar PEHE to TOBART and TOBART-NP for nie2021quasi DGP D, which involves partly linear prognostic and treatment effect functions, although soft TOBART is notably more accurate. The average coverages of TOBART and soft TOBART credible intervals for $\tau(\bm{x}_i)$ are generally much closer to $95\%$ than the coverages of intervals produced by competing methods. TOBART-NP produces very wide credible intervals relative to TOBART. TOBART-NP produces better coverage than TOBART for four DGPs.
For the data application, we consider the same methods as in section (ref), excluding Gaussian Processes and adding a hurdle model combining linear regression and probit. For each data set, we average results over 10 training-test splits. Each split is defined by taking a random sample of $\text{floor}(0.7n)$ training observations stratified by censorship status. Categorical variables were encoded as sets of dummy variables. The numbers of observations, covariates, and proportions of censored observations are given in table (ref). Appendix (ref) contains descriptions of each data set with references to original sources.
The data application results are presented in Table (ref) . For most data sets the results are similar across methods, particularly when methods are evaluated in terms of Brier score for predicted probabilities of censoring. Similar results can be observed for the AUC in the Supplementary Appendix (Online Resource 1). TOBART can give notably lower MSE of outcome predictions relative to other methods for some data sets.
In contrast to the simulation studies above, there is not a clear winning method in Table (ref). Although censored outcome models have been applied to these data sets in previous work, perhaps other models are more suitable for some data sets. This is evidenced by the fact that for many data sets the combination of probit and a linear model outperforms Tobit. Therefore for some data sets zero inflated, hurdle, or sample selection models might be more appropriate. For the data sets on which Tobit outperforms probit and a linear model in terms of MSE, namely Recon and Atrazine, the best method is Soft TOBART. The TOBART models also notably outperform other methods when applied to the BostonHousing and Missouri data sets.
A lesson from this study is that it is important to select the appropriate model for the data set. The TOBART and Grabit methods are designed for the same form of DGPs, therefore it is arguably fairer to compare these two methods. Soft TOBART produces lower MSE predictions than Grabit across almost all data sets.\footnote{Potentially Grabit could produce better results with more hyperparameter tuning, although this would be computationally costly.} Nonetheless, the results are less impressive than those observed in the simulation study. Possible explanations for this include slow mixing of the TOBART Markov Chain, small sample sizes for some data sets, and very small or very large proportions of censored outcomes.
Type I TOBART produces accurate predictive probabilities of censoring, predictions of outcomes, and treatment effect estimates. TOBART-NP, gives better uncertainty quantification for some simulated DGPs. Advantages of TOBART over competing methods include the fact that hyperparameter tuning is not required, and the straightforward combination of the method with other variations on BART to allow for smooth DGPs and sparsity linero2018bayesianB.
\backmatter
\bmhead{Supplementary information}
The online supplementary appendix contains (A) additional simulation study results, (B) additional data application results, and (C) implementation details and parameter settings.
\bmhead{Acknowledgments}
The author gratefully acknowledges helpful comments from Mikhail Zhelonkin, Chen Zhou, and participants at the Econometric Institute internal seminar.
Conflict of interest: The authors declare no competing interests.
\FloatBarrier