SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Ruofan Wu\*†, Jiawei Qiao\*‡, Mingzhe Wu§, Wen Yu‡, Ming Zheng‡, Tengfei Liu†, Tianyi Zhang†, and Weiqiang Wang†
†Ant Group ‡Fudan University §Coupang
{ruofan.wrf, aaron.ltf, zty113091, weiqiang.wwq}@antgroup.com jeremyqjw@163.com, wumingzhe.darcy@gmail.com, {wenyu, mingzheng}@fudan.edu.cn
We present neural frailty machine (NFM), a powerful and flexible neural modeling framework for survival regressions. The NFM framework utilizes the classical idea of multiplicative frailty in survival analysis as a principled way of extending the proportional hazard assumption, at the same time being able to leverage the strong approximation power of neural architectures for handling nonlinear covariate dependence. Two concrete models are derived under the framework that extends neural proportional hazard models and nonparametric hazard regression models. Both models allow efficient training under the likelihood objective. Theoretically, for both proposed models, we establish statistical guarantees of neural function approximation with respect to nonparametric components via characterizing their rate of convergence. Empirically, we provide synthetic experiments that verify our theoretical statements. We also conduct experimental evaluations over 6 benchmark datasets of different scales, showing that the proposed NFM models achieve predictive performance comparable to or sometimes surpassing state-of-the-art survival models. Our code is publicly availabel at https://github.com/Rorschach1989/nfm
Regression analysis of time-to-event data [40] has been among the most important modeling tools for clinical studies and has witnessed a growing interest in areas like corporate finance [22], recommendation systems [38], and computational advertising [71]. The key feature that differentiates time-to-event data from other types of data is that they are often incompletely observed, with the most prevailing form of incompleteness being the right censoring mechanism [40]. In the right censoring mechanism, the duration time of a sampled subject is (sometimes) only known to be larger than the observation time instead of being recorded precisely. It is well known in the community of survival analysis that even in the case of linear regression, naively discarding the censored observations produces estimation results that are statistically biased [7], at the same time losses sample efficiency if the censoring proportion is high.
Cox's proportional hazard (CoxPH) model [14] using the convex objective of negative partial likelihood [15] is the de facto choice in modeling right censored time-to-event data (hereafter abbreviated as censored data without misunderstandings). The model is semiparametric [4] in the sense that the baseline hazard function needs no parametric assumptions. The original formulation of CoxPH model
assumes a linear form and therefore has limited flexibility since the truth is not necessarily linear. Subsequent studies extended CoxPH model to nonlinear variants using ideas from nonparametric regression [36, 8, 9], ensemble learning [37], and neural networks [25, 42]. While such extensions allowed a more flexible nonlinear dependence structure with the covariates, the learning objectives were still derived under the proportional hazards (PH) assumption, which was shown to be inadequate in many real-world scenarios [31]. The most notable case was the failure of modeling the phenomenon of crossing hazards [61]. It is thus of significant interest to explore extensions of CoxPH that both allow nonlinear dependence over covariates and relaxations of the PH assumption.
Frailty models [69, 21] are among the most important research topics in modern survival analysis, in that they provide a principled way of extending CoxPH model via incorporating a multiplicative random effect to capture unobserved heterogeneity. The resulting parameterization contains many useful variants of CoxPH like the proportional odds model [3], under specific choices of frailty families. While the theory of frailty models has been well-established [48, 49, 53, 45], most of them focused on the linear case. Recent developments on applying neural approaches to survival analysis [42, 46, 64, 56] have shown promising results in terms of empirical predictive performance, with most of them lacking theoretical discussions. Therefore, it is of significant interest to build more powerful frailty models via adopting techniques in modern deep learning [29] with provable statistical guarantees.
In this paper, we present a general framework for neural extensions of frailty models called the neural frailty machine (NFM). Two concrete neural architectures are derived under the framework: The first one adopts the proportional frailty assumption, allowing an intuitive interpretation of the neural CoxPH model with a multiplicative random effect. The second one further relaxes the proportional frailty assumption and could be viewed as an extension of nonparametric hazard regression (NHR) [13, 44], sometimes referred to as "fully neural" models under the context of neural survival analysis [52]. We summarize our contributions as follows.
- We propose the neural frailty machine (NFM) framework as a principled way of incorporating unobserved heterogeneity into neural survival regression models. The framework includes many commonly used survival regression models as special cases.
- We derive two model architectures based on the NFM framework that extend neural CoxPH models and neural NHR models. Both models allow stochastic training and scale to large datasets.
- Theoretically, we show statistical correctness of the two proposed models via characterizing the rates of convergence of the proposed nonparametric function estimators. The proof technique is different from previous theoretical studies on neural survival analysis and is applicable to many other types of neural survival models.
- Empirically, we verify the empirical efficacy of the proposed framework via conducting extensive studies on various benchmark datasets at different scales. Under standard performance metrics, both models are empirically shown to perform competitively, matching or sometimes outperforming state-of-the-art neural survival models.
2.1 Nonlinear extensions of CoxPH
Most nonlinear extensions of CoxPH model stem from the equivalence of partial likelihood and semiparametric profile likelihood [50] of CoxPH model, resulting in nonlinear variants that essentially replaces the linear term in partial likelihood with nonlinear variants: [36] used smoothing splines, [8, 9] used local polynomial regression [24]. The empirical success of tree-based models inspired subsequent developments like [37] that equip tree-based models such as gradient boosting trees and random forests with losses in the form of negative log partial likelihood. Early developments of neural survival analysis [25] adopted similar extension strategies and obtained neural versions of partial likelihood. Later attempts [42] suggested using the successful practice of stochastic training which is believed to be at the heart of the empirical success of modern neural methods [34]. However, stochastic training under the partial likelihood objective is highly non-trivial, as mini-batch versions of log partial likelihood [42] are no longer valid stochastic gradients of the full-sample log partial likelihood [64].
2.2 Beyond CoxPH in survival analysis
In linear survival modeling, there are standard alternatives to CoxPH such as the accelerated failure time (AFT) model [7, 73], the extended hazard regression model [23], and the family of linear transformation models [74]. While these models allow certain types of nonlinear extensions, the resulting form of (conditional) hazard function is still restricted to be of a specific form. The idea of nonparametric hazard regression (NHR) [13, 44, 63] further improves the flexibility of nonparametric survival analysis via directly modeling the conditional hazard function by nonparametric regression techniques such as spline approximation. Neural versions of NHR have been developed lately such as the CoxTime model [46]. [56] used a neural network to approximate the conditional survival function and could be thus viewed as another trivial extension of NHR.
Aside from developments in NHR, [47] proposed a discrete-time model with its objective being a mix of the discrete likelihood and a rank-based score; [75] proposed a neural version of the extended hazard model, unifying both neural CoxPH and neural AFT model; [64] used an ODE approach to model the hazard and cumulative hazard functions.
2.3 Theoretical justification of neural survival models
Despite the abundance of neural survival models, assessment of their theoretical properties remains nascent. In [76], the authors developed minimax theories of partially linear cox model using neural networks as the functional approximator. [75] provided convergence guarantees of neural estimates under the extended hazard model. The theoretical developments therein rely on specific forms of objective function (partial likelihood and kernel pseudo-likelihood) and are not directly applicable to the standard likelihood-based objective which is frequently used in survival analysis.
3.1 The neural frailty machine framework
Let  $\tilde{T} \geq 0$  be the interested event time with survival function denoted by  $S(t) = \mathbb{P}(\tilde{T} > t)$  associated with a feature(covariate) vector  $Z \in \mathbb{R}^d$ . Suppose that  $\tilde{T}$  is a continuous random variable and let  $f(t)$  be its density function. Then  $\lambda(t) = f(t) / S(t)$  is the hazard function and  $\Lambda(t) = \int_0^t \lambda(s) ds$  is the cumulative hazard function. Aside from the covariate  $Z$ , we use a positive scalar random variable  $\omega \in \mathbb{R}^+$  to express the unobserved heterogeneity corresponding to individuals, or frailty. In this paper we will assume the following generating scheme of  $\tilde{T}$  via specifying its conditional hazard function:
$$
\lambda (t | Z, \omega) = \omega \widetilde {\nu} (t, Z). \tag {1}
$$
Here  $\widetilde{\nu}$  is an unspecified non-negative function, and we let the distribution of  $\omega$  be parameterized by a one-dimensional parameter  $\theta \in \mathbb{R}$ . The formulation (1) is quite general and contains several important models in both traditional and neural survival analysis:
1. When  $\omega$  follows parametric distributional assumptions, and  $\widetilde{\nu}(t,Z) = \lambda(t)e^{\beta^\top Z}$ , (1) reduces to the standard proportional frailty model [45]. A special case is when  $\omega$  is degenerate, i.e., it has no randomness, then the model corresponds to the classic CoxPH model.
2. When  $\omega$  is degenerate and  $\widetilde{\nu}$  is arbitrary, the model becomes equivalent to nonparametric hazard regression (NHR) [13, 44]. In NHR, the function parameter of interest is usually the logarithm of the (conditional) hazard function.
In this paper we construct neural approximations to the logarithm of  $\widetilde{\nu}$ , i.e.,  $\nu(t,Z) = \log \widetilde{\nu}(t,Z)$ . The resulting models are called Neural Frailty Machines (NFM). Depending on the prior knowledge of the function  $\nu$ , we propose two function approximation schemes:
For example in medical biology, it was observed that genetically identical animals kept in as similar an environment as possible will typically not behave the same upon exposure to environmental carcinogens [6]
3The choice of one-dimensional frailty family is mostly for simplicity and cleanness of theoretical derivations. Note that there exist multi-dimensional frailty families like the PVF family [69]. Generalizing our theoretical results to such kinds of families would require additional sets of regularity conditions, and will be left to future explorations.
The proportional frailty (PF) scheme assumes the dependence of  $\nu$  on event time and covariates to be completely decoupled, i.e.,
$$
\nu (t, Z) = h (t) + m (Z). \tag {2}
$$
Proportional-style assumption over hazard functions has been shown to be a useful inductive bias in survival analysis. We will treat both  $h$  and  $m$  in (2) as function parameters, and device two multi-layer perceptrons (MLP) to approximate them separately.
The fully neural (FN) scheme imposes no a priori assumptions over  $\nu$  and is the most general version of NFM. It is straightforward to see that the most commonly used survival models, such as CoxPH, AFT[73], EH[75], or PF models are included in the proposed model space as special cases. We treat  $\nu = \nu(t,Z)$  as the function parameter with input dimension  $d + 1$  and use a multi-layer perceptron (MLP) as the function approximator to  $\nu$ . Similar approximation schemes with respect to the hazard function have been proposed in some recent works [52, 56], referred to as "fully neural approaches" without theoretical characterizations.
The choice of frailty family There are many commonly used families of frailty distributions [45, 21, 69], among which the most popular one is the gamma frailty, where  $\omega$  follows a gamma distribution with mean 1 and variance  $\theta$ . We briefly introduce some other types of frailty families in appendix A.
3.2 Parameter learning under censored observations
In time-to-event modeling scenarios, the event times are typically observed under right censoring. Let  $C$  be the right censoring time which is assumed to be conditionally independent of the event time  $\tilde{T}$  given  $Z$ , i.e.,  $\tilde{T} \perp C|Z$ . In data collection, one can observe the minimum of the survival time and the censoring time, that is, observe  $T = \tilde{T} \wedge C$  as well as the censoring indicator  $\delta = I(\tilde{T} \leqslant C)$ , where  $a \wedge b = \min(a, b)$  for constants  $a$  and  $b$  and  $I(\cdot)$  stands for the indicator function. We assume  $n$  independent and identically distributed (i.i.d.) copies of  $(T, \delta, Z)$  are used as the training sample  $(T_i, \delta_i, Z_i), i \in [n]$ , where we use  $[n]$  to denote the set  $\{1, 2, \dots, n\}$ . Additionally, we assume the unobserved frailties are independent and identically distributed, i.e.,  $\omega_i \stackrel{\text{i.i.d.}}{\sim} f_\theta(\omega), i \in [n]$ . Next, we derive the learning procedure based on the observed log-likelihood (OLL) objective under both PF and FN scheme. To obtain the observed likelihood, we first integrate the conditional survival function given the frailty:
$$
S (t | Z) = \mathbb {E} _ {\omega \sim f _ {\theta}} \left[ e ^ {- \omega \int_ {0} ^ {t} e ^ {\nu (s, Z)} d s} \right] =: e ^ {- G _ {\theta} \left(\int_ {0} ^ {t} e ^ {\nu (s, Z)} d s\right)}. \tag {3}
$$
Here the frailty transform  $G_{\theta}(x) = -\log \left(\mathbb{E}_{\omega \sim f_{\theta}}[e^{-\omega x}]\right)$  is defined as the negative of the logarithm of the Laplace transform of the frailty distribution. The conditional cumulative hazard function is thus  $\Lambda (t|Z) = G_{\theta}(\int_0^t e^{\nu (s,Z)}ds)$ . For the PF scheme of NFM, we use two MLPs  $\widehat{h} = \widehat{h} (t;\mathbf{W}^h,\mathbf{b}^h)$  and  $\widehat{m} = \widehat{m} (Z;\mathbf{W}^{m},\mathbf{b}^{m})$  as function approximators to  $\nu$  and  $m$ , parameterized by  $(\mathbf{W}^h,\mathbf{b}^h)$  and  $(\mathbf{W}^m,\mathbf{b}^m)$ , respectively. 4 According to standard results on censored data likelihood [40], we write the learning objective under the PF scheme as:
$$
\begin{array}{l} \mathcal {L} \left(\mathbf {W} ^ {h}, \mathbf {b} ^ {h}, \mathbf {W} ^ {m}, \mathbf {b} ^ {m}, \theta\right) \\ = \frac {1}{n} \left[ \sum_ {i \in [ n ]} \delta_ {i} \log g _ {\theta} \left(e ^ {\widehat {m} (Z _ {i})} \int_ {0} ^ {T _ {i}} e ^ {\widehat {h} (s)} d s\right) + \delta_ {i} \widehat {h} (T _ {i}) + \delta_ {i} \widehat {m} (Z _ {i}) - G _ {\theta} \left(e ^ {\widehat {m} (Z _ {i})} \int_ {0} ^ {T _ {i}} e ^ {\widehat {h} (s)} d s\right) \right]. \tag {4} \\ \end{array}
$$
Here we define  $g_{\theta}(x) = \frac{\partial}{\partial x} G_{\theta}(x)$ . Let  $(\widehat{\mathbf{W}}_n^h, \widehat{\mathbf{b}}_n^h, \widehat{\mathbf{W}}_n^m, \widehat{\mathbf{b}}_n^m, \widehat{\boldsymbol{\theta}}_n)$  be the maximizer of (4) and further denote  $\widehat{h}_n(t) = \widehat{h}(t; \widehat{\mathbf{W}}_n^h, \widehat{\mathbf{b}}_n^h)$  and  $\widehat{m}_n(Z) = \widehat{m}(Z; \widehat{\mathbf{W}}_n^m, \widehat{\mathbf{b}}_n^m)$ . The resulting estimators for conditional cumulative hazard and survival functions are:
$$
\widehat {\Lambda} _ {\mathrm {P F}} (t | Z) = G _ {\widehat {\theta} _ {n}} \left(\int_ {0} ^ {t} e ^ {\widehat {h} _ {n} (s) + \widehat {m} _ {n} (Z)} d s\right), \quad \widehat {S} _ {\mathrm {P F}} (t | Z) = e ^ {- \widehat {\Lambda} _ {\mathrm {P F}} (t | Z)}, \tag {5}
$$
Here we adopt the conventional notation that  $\mathbf{W}$  is the collection of the weight matrices of the MLP in all layers, and  $\mathbf{b}$  corresponds to the collection of the bias vectors in all layers.
For the FN scheme, we use  $\widehat{\nu} = \widehat{\nu}(t,Z;\mathbf{W}^{\nu},\mathbf{b}^{\nu})$  to approximate  $\nu(t,Z)$  parameterized by  $(\mathbf{W}^{\nu},\mathbf{b}^{\nu})$ . The OLL objective is written as:
$$
\begin{array}{l} \mathcal {L} (\mathbf {W} ^ {\nu}, \mathbf {b} ^ {\nu}, \theta) \\ = \frac {1}{n} \left[ \sum_ {i \in [ n ]} \delta_ {i} \log g _ {\theta} \left(\int_ {0} ^ {T _ {i}} e ^ {\hat {\nu} (s, Z _ {i}; \mathbf {W} ^ {\nu}, \mathbf {b} ^ {\nu})} d s\right) + \delta_ {i} \hat {\nu} \left(T _ {i}, Z _ {i}; \mathbf {W} ^ {\nu}, \mathbf {b} ^ {\nu}\right) - G _ {\theta} \left(\int_ {0} ^ {T _ {i}} e ^ {\hat {\nu} (s, Z _ {i}; \mathbf {W} ^ {\nu}, \mathbf {b} ^ {\nu})} d s\right) \right]. \tag {6} \\ \end{array}
$$
Let  $(\widehat{\mathbf{W}}_n^\nu, \widehat{\mathbf{b}}_n^\nu, \widehat{\theta}_n)$  be the maximizer of (6), and further denote  $\widehat{\nu}_n(t,Z) = \widehat{\nu}(t,Z; \widehat{\mathbf{W}}_n^\nu, \widehat{\mathbf{b}}_n^\nu)$ . The conditional cumulative hazard and survival functions are therefore estimated as:
$$
\widehat {\Lambda} _ {\mathrm {F N}} (t | Z) = G _ {\widehat {\theta} _ {n}} \left(\int_ {0} ^ {t} e ^ {\widehat {\nu} _ {n} (s, Z)} d s\right), \quad \widehat {S} _ {\mathrm {F N}} (t | Z) = e ^ {- \widehat {\Lambda} _ {\mathrm {F N}} (t | Z)}. \tag {7}
$$
The evaluation of objectives like (6) and its gradient requires computing a definite integral of an exponentially transformed MLP function. Instead of using exact computations that are available for only a restricted type of activation functions and network structures, we use numerical integration for such kinds of evaluations, using the method of Clenshaw-Curtis quadrature [5], which has shown competitive performance and efficiency in recent applications to monotonic neural networks [68].
Remark 3.1. The interpretation of frailty terms differs in the two schemes. In the PF scheme, introducing the frailty effect strictly increases the modeling capability (i.e., the capability of modeling crossing hazard) in comparison to CoxPH or neural variants of CoxPH [45]. In the FN scheme, it is arguable that in the i.i.d. case, the marginal hazard function is a reparameterization of the hazard function in the context of NHR. Therefore, we view the incorporation of frailty effect as injecting a domain-specific inductive bias that has proven to be useful in survival analysis and time-to-event regression modeling and verify this claim empirically in section 5.2. Moreover, frailty becomes especially helpful when handling correlated or clustered data where the frailty term is assumed to be shared among certain groups of individuals [53]. Extending NFM to such scenarios is valuable and we left it to future explorations.
4 Statistical guarantees
In this section, we present statistical guarantees regarding both NFM estimates in the sense of nonparametric regression [65], where we obtain rates of convergence to the ground truth function parameters (which is frequently referred to as the true parameter in statistics literature). The results in this section is interpreted as showing the statistical correctness of our approach.
Proof strategy Our proof technique is based on the method of sieves [60, 59, 11] that views neural networks as a special kind of nonlinear sieve [11] that satisfies desirable approximation properties [72]. Our strategy is different from previous theoretical works on neural survival models [75, 76] where the developments implicitly requires the loss function to be well-controlled by the  $L_{2}$  loss and is therefore not directly applicable to our model due to the flexibility in choosing the frailty transform. Since both models produce estimates of function parameters, we need to specify a suitable function space to work with. Here we choose the following Hölder ball as was also used in previous works on nonparametric estimation using neural networks [57, 26, 76]
$$
\mathcal {W} _ {M} ^ {\beta} (\mathcal {X}) = \left\{f: \max  _ {\alpha : | \alpha | \leq \beta} \operatorname {e s s s u p} _ {x \in \mathcal {X}} \left| D ^ {\alpha} (f (x)) \right| \leq M \right\}, \tag {8}
$$
where the domain  $\mathcal{X}$  is assumed to be a subset of  $d$ -dimensional euclidean space.  $\alpha = (\alpha_{1},\ldots ,\alpha_{d})$  is a  $d$ -dimensional tuple of nonnegative integers satisfying  $|\alpha | = \alpha_{1} + \dots +\alpha_{d}$  and  $D^{\alpha}f = \frac{\partial^{|\alpha|}f}{\partial x_1^{\alpha 1}\cdots x_d^{\alpha d}}$  is the weak derivative of  $f$ . Now assume that  $M$  is a reasonably large constant, and let  $\Theta$  be a closed interval over the real line. We make the following assumptions for the true parameters under both schemes:
Condition 4.1 (True parameter, PF scheme). The euclidean parameter  $\theta_0 \in \Theta \subset \mathbb{R}$ , and the two function parameters  $m_0 \in \mathcal{W}_M^\beta ([-1,1]^d)$ ,  $h_0 \in \mathcal{W}_M^\beta ([0,\tau])$ , and  $\tau > 0$  is the ending time of the study duration, which is usually adopted in the theoretical studies in survival analysis [67].
Condition 4.2 (True parameter, FN scheme). The euclidean parameter  $\theta_0 \in \Theta \subset \mathbb{R}$ , and the function parameter  $\nu_0 \in \mathcal{W}_M^\beta ([0,\tau ] \times [-1,1]^d)$ ,
Next, we construct sieve spaces for function parameter approximation via restricting the complexity of the MLPs to "scale" with the sample size  $n$ .
Condition 4.3 (Sieve space, PF scheme). The sieve space  $\mathcal{H}_n$  is constructed as a set of MLPs satisfying  $\widehat{h} \in \mathcal{W}_{M_h}^\beta ([0,\tau])$ , with depth of order  $O(\log n)$  and total number of parameters of order  $O(n^{\frac{1}{\beta + d}}\log n)$ . The sieve space  $\mathcal{M}_n$  is constructed as a set of MLPs satisfying  $\widehat{m} \in \mathcal{W}_{M_m}^\beta ([-1,1]^d)$ , with depth of order  $O(\log n)$  and total number of parameters of order  $O(n^{\frac{d}{\beta + d}}\log n)$ . Here  $M_h$  and  $M_{m}$  are sufficiently large constants such that every function in  $\mathcal{W}_M^\beta ([-1,1]^d)$  and  $\mathcal{W}_M^\beta ([0,\tau])$  could be accurately approximated by functions inside  $\mathcal{H}_n$  and  $\mathcal{M}_n$ , according to [72, Theorem 1].
Condition 4.4 (Sieve space, FN scheme). The sieve space  $\mathcal{V}_n$  is constructed as a set of MLPs satisfying  $\widehat{\nu} \in \mathcal{W}_{M_{\nu}}^{\beta}([0,\tau])$ , with depth of order  $O(\log n)$  and total number of parameters of order  $O(n^{\frac{d+1}{\beta+d+1}}\log n)$ . Here  $M_{\nu}$  is a sufficiently large constant such that  $\mathcal{V}_n$  satisfies approximation properties, analogous to condition 4.3.
For technical reasons, we will assume the nonparametric function estimators are constrained to fall inside the corresponding sieve spaces, i.e.,  $\widehat{h}_n\in \mathcal{H}_n$ ,  $\widehat{m}_n\in \mathcal{M}_n$  and  $\widehat{\nu}\in \mathcal{V}_n$ . This will not affect the implementation of optimization routines as was discussed in [26]. Furthermore, we restrict the estimate  $\widehat{\theta}_n\in \Theta$  in both PF and FN schemes.
Additionally, we need the following regularity condition on the function  $G_{\theta}(x)$ :