Accelerated failure time model in XGBoost

Likelihood function taking account of three censoring types.

XGBoost optimizes a twice-differentiable convex loss function $l(y_i,{\hat{y}}_i)$ in its second-order method of gradient boosting.

Now let’s define a loss function $l_{AFT}$ for the AFT model. Let $\mathcal{D}={(x_i,y_i)}_{i=1}^n$ denote the training data, and $Y_1,...,Y_n$ denote random variables i.i.d. with distribution for $Y$. The likelihood for $\mathcal{D}$ is the product of probability densities $f_Y$ for individual data points.

$$L(\mathcal{D})=\mathbb{P}[Y_1=y_1,\dots ,Y_n=y_n]=\prod _{i=1}^n\mathbb{P}[Y_i=y_i]=\prod _{i=1}^n{f}_Y\left(y_i\right)$$

As usual, we change it to log scale, the goal is changed to maximize log likelihood.

$$\ln L(\mathcal{D})=\sum _{i=1}^n\ln \mathbb{P}[Y_i=y_i]=\sum _{i=1}^n\ln f_Y\left(y_i\right)$$

Under censoring, we don’t know $y_i$ for some individuals. Therefore we revise the likelihood function by using definition of probability to take account of the censored data:

$$\ln L(\mathcal{D})=\underset{\text{uncensored label }}{\underbrace{\sum \ln \mathbb{P}[Y_i=y_i]}}+\underset{\text{censored label with }{y}_i\in [{\underset{―}{y}}_i,\overline{y_i}]}{\underbrace{\sum \ln \mathbb{P}[{\underset{―}{y}}_i\le Y_i\le \overline{y_i}]}}=\underset{\text{uncensored label }}{\underbrace{\sum \ln f_Y\left(y_i\right)}}+\underset{\text{censored label with }{y}_i\in [\underset{―}{y_i},\overline{y_i}]}{\underbrace{\sum \ln (F_Y\left(\overline{y_i}\right)-F_Y\left(\underset{―}{y_i}\right)}}$$

where $\underset{―}{y_i}$ and $\overline{y_i}$ are lower and upper bounds for $y_i$, respectively. $F_Y$ is the cumulative distribution function (CDF). When $\overline{y_i}$ is infinity, it is right-censored data while when $\underset{―}{y_i}$ is 0, it indicates left-censored data.

Likelihood function in cooperation with AFT

AFT model:

$$ln(y_i)={\hat{y}}_i+\sigma z_i$$ where ${\hat{y}}_i=\mathcal{T}(x)$. and NOTE $y_i$ indicates time to event and its variance comes from $z_i$. $Z$ is a random variable of a known probability distribution.

Suppose $Y=g(Z)$ and $g(\cdot )$ is a monotone increasing function. The pdf and cdf of $Y$ can be expressed in terms of pdf and cdf of $Z$:

$$f_Y(y)=f_Z(g^{-1}(y))\cdot \frac{d}{dy}{g}^{-1}(y)F_Y(y)=F_Z(g^{-1}(y))$$

Therefore, the loss function of AFT model will be:

$${\ell }_{\mathrm{AFT}}(y,\hat{y})=\{\begin{array}{ll}-\ln [f_Z(s(y))\cdot \frac{1}{\sigma y}]& \text{if }y\text{ is not censored }\\ -\ln [F_Z(s(\overline{y}))-F_Z(s(\underset{―}{y}))]& \text{if }y\text{ is censored with }y\in [\underset{―}{y},\overline{y}]\end{array}$$

where $s(y)=(\ln y-\hat{y})/\sigma$

Gradient and hessian of the AFT loss

The gradient boosting algorithm in XGBoost uses the gradient and hessian of the loss function, which are first and second partial derivatives of $\ell$ with respect to $\hat{y}$. The gradient and hessian of the AFT loss function are as follows:

$${\frac{\partial {\ell }_{\mathrm{AFT}}}{\partial \hat{y}}|}_{y,\hat{y}}=\{\begin{array}{ll}\frac{f_Z^{\prime }(s(y))}{\sigma f_Z(s(y))}& \text{if }y\text{ is not censored }\\ \frac{f_Z(s(\overline{y}))-f_Z(s(\underset{―}{y}))}{\sigma [F_Z(s(\overline{y}))-F_Z(s(\underset{―}{y}))]}& \text{ if }y\text{ is censored with }y\in [\underset{―}{y},\overline{y}]\end{array}$$

$${\frac{{\partial }^2{\ell }_{\mathrm{AFT}}}{\partial {\hat{y}}^2}|}_{y,\hat{y}}=\{\begin{array}{ll}-\frac{f_Z(s(y))f_Z^{\prime \prime }(s(y))-f_Z^{\prime }(s(y){)}^2}{{\sigma }^2{f}_Z(s(y){)}^2}& \text{if }y\text{ is not censored}\\ -[F_Z(s(\overline{y}))-F_Z(s(\underset{―}{y}))][f_Z^{\prime }(s(\overline{y}))-f_Z^{\prime }(s(\underset{―}{y}))]\\ \frac{+{[f_Z(s(\overline{y}))-f_Z(s(\underset{―}{y}))]}^2}{{\sigma }^2{[F_Z(s(\overline{y}))-F_Z(s(\underset{―}{y}))]}^2}& \text{if }y\text{ is censored }\end{array}$$

Cite1: Survival regression with accelerated failure time model in XGBoost

Cite2: XGBoost Documentation