H AZARD G RADIENT P ENALTY FOR S URVIVAL A NALYSIS
Anonymous authors
Paper under double-blind review
A BSTRACT
Survival analysis appears in various fields such as medicine, economics, engineering, and business.
Recent studies showed that the Ordinary Differential Equation (ODE) modeling framework integrates many existing survival models while the framework is flexible and widely applicable.
However, naively applying the ODE framework to survival analysis problems may model fiercely changing density function with respect to covariates which may worsen the model’s performance.
Though we can apply L1 or L2 regularizers to the ODE model, their effect on the ODE modeling framework is barely known.
In this paper, we propose hazard gradient penalty (HGP) to enhance the performance of a survival analysis model.
Our method imposes constraints on local data points by regularizing the gradient of hazard function with respect to the data point.
Our method applies to any survival analysis model including the ODE modeling framework and is easy to implement.
We theoretically show that our method is related to minimizing the
KL divergence between the density function at a data point and that of the neighborhood points.
Experimental results on three public benchmarks show that our approach outperforms other regularization methods.
I NTRODUCTION
Survival analysis (a.k.a time-to-event modeling) is a branch of statistics that predicts the duration of time until an event occurs (Kleinbaum & Klein, 2012).
Survival analysis appears in various fields such as medicine (Schwab et al., 2021), economics (Meyer, 1988), engineering (O’Connor &
Kleyner, 2011), and business (Jing & Smola, 2017; Li et al., 2021).
Due to the presence of rightcensored data, which is data whose event has not occurred yet, survival analysis models require special considerations.
Cox proportional hazard model (CoxPH)
(Cox, 1972; Katzman et al., 2018) and accelerated time failure model (AFT) (Wei, 1992) are widely used to handle right-censored data.
Yet the assumptions made by these models are frequently violated in the real world (Lee et al., 2018;
Tang et al., 2022a).
Recent studies showed that the Ordinary Differential Equation (ODE) modeling framework integrates many existing survival analysis models including CoxPH and AFT (Groha et al., 2020; Tang et al., 2022a;b).
They also showed that the ODE modeling framework is flexible and widely applicable.
However, naively applying the ODE framework to survival analysis problems may result in wildly oscillating density function that may worsen the model’s performance.
Regularization techniques that can regularize this undesirable behavior are understudied.
Though applying L1 or L2 regularizers to the ODE model is one option, their effects on the ODE modeling framework are barely known.
The cluster assumption from semi-supervised learning states that the decision boundaries should not cross high-density regions (Chapelle et al., 2006).
Likewise, survival analysis models need hazard functions that slowly change in high-density regions.
Suppose we attempt to predict the time to death of three individuals A, B, and C. Assume the traits of A and B are similar and the traits of B and C are dissimilar.
It is natural to expect that the probability distribution of time-to-death of A should be close to that of B while far from that of C.
The expectation aligns with the cluster assumption.
Explicitly modeling the assumption enhances the performance as long as it holds.
In this paper, we propose hazard gradient penalty to make a slowly changing (with respect to covariates) survival analysis model in high-density regions.
In a nutshell, the hazard gradient penalty regularizes the gradient of the hazard function with respect to the data point from the real data distribution.
Our method has several advantages.
The method is computationally efficient.
The method is theoretically sound.
The method is applicable to any survival analysis model including the ODE modeling framework as long as it models hazard function.
It is easy to implement.
We theoretically show that our method is related to minimizing the KL divergence between the density function at a data point and that of the neighborhood points of the data point.
Experimental results on three public benchmarks show that our approach outperforms other regularization methods.
P RELIMINARIES
Survival analysis data comprises of an observed covariate x , a failure event time t , and an event indicator e .
If an event is observed, t corresponds to the duration time from the beginning of the follow-up of an individual until the event occurs.
In this case, the event indicator e = 1 .
If an event is unobserved, t corresponds to the duration time from the beginning of follow-up of an individual until the last follow-up.
In this case, we cannot know the exact time of the event occur and event indicator e = 0 .
An individual is said to be right-censored if e = 0 .
The presence of right-censored data differentiates survival analysis from regression problems.
In this paper, we only focus on the single-risk problem where event e is a binary-valued variable.
Given a set of triplet D = { ( x i , t i , e i ) } Ni =1 , the goal of survival analysis is to predict the likelihood of an event occur p ( t | x ) or the survival probability S ( t | x ) .
The likelihood and the survival probability have the following relationship:
[Equation]
Modeling p ( t | x ) or S ( t | x ) directly is challenging as those have the following constraints:
[Equation]
Previous works instead modeled the hazard function (a.k.a conditional failure rate) h ( t | x ) (Cox, 1972; Katzman et al., 2018;
Wei, 1992; Zhong et al., 2021).
[Equation]
As the hazard function is a probability per unit time, it is unbounded upwards.
Hence, the only constraint of the hazard function is that the function is non-negative: h ( t | x ) ≥
2. T HE ODE M ODELING F RAMEWORK
We can obtain an ODE which explains the relationship between the hazard function and the survival function by putting derivative of equation 1 into equation 2 (Kleinbaum & Klein, 2012).
[Equation]
Starting from initial value log S (0 | x ) = 0 , we can define log S ( t | x ) as the solution of the ODE initial value problem where the ODE is defined as equation 3 1 .
[Equation]
We can train the ODE model by minimizing the negative log-likelihood.
[Equation]
Following Groha et al.
(2020), we update the model parameters using Neural ODEs (Chen et al., 2018).
The hazard function h θ ( t | x ) is modeled using a neural network followed by the softplus activation function to ensure that the output is always non-negative.
2. N EURAL ODE S
Neural ODEs model the continuous dynamics of variables (Chen et al., 2018).
Starting from z (0) , we can define the output z ( T ) to be the solution of the following ordinary differential equation (ODE) initial value problem.
[Equation]
Naively applying an ODE solver to an ODE initial value problem leads to practical difficulties.
An ODE solver builds a big computation graph which incurs high memory cost and additional numerical errors may occur in backpropagation steps.
Chen et al.
(2018) showed that we can obtain the gradients of a scalar-valued loss w.r.t all inputs of any ODE solver with constant memory cost.
We can calculate the gradients without backpropagating through the operations of the solver but with another call to an ODE solver.
M ETHODS
Our method minimize the hazard gradient penalty
[Equation]
and the negative log-likelihood in equation 4 at the same time.
Intuitively speaking, we regularize the model so that the hazard function does not vary much when a small noise ν is added or subtracted to the data point x .
In section 3.2, we show that minimizing the hazard gradient penalty is connected to minimizing the KL divergence between the density at x and the density at x ′ ∈ B ( x , ϵ ) .
In this section, we introduce the hazard gradient penalty and show that it is related to minimizing the
KL divergence between the density function at a data point and that of its neighbours.
See Figure 1 for the graphical overview of our method.
The cluster assumption from semi-supervised learning states that the decision boundaries should not cross high-density regions (Chapelle et al., 2006).
In a similar vein, hazard functions of survival analysis models should change slowly in high-density regions.
Consider a case where two data points x 1 , x 2 ∈ R d failed at t 1 , t 2 ( t 1 < t 2 ) each.
Under the cluster assumption, a point x ′ 2 ∈ B ( x 2 , ϵ ) 2 should fail at t ′
≈ t 2 .
If ˆ p ( t | x ′ 2 ) is skewed for some reason and puts high density at t ′ 1 < t 1 , it worsens the model’s performance.
To evade such situation, ˆ p ( t | x ′ 2 ) should not deviate too much from ˆ p ( t | x 2 ) .
To achieve this, we propose the following regularizer.
[Equation]
In practice, we implement h θ ( t | x ) using a neural network whose input is a combination (concatenation, addition or both) of t and x .
Hence we can write h θ ( t | x ) and h θ ( t, x ) interchangeably.
The gradient ∇ x h θ ( t, x ) is naturally defined and so is ∇ x h θ ( t | x ) .
3. E FFICIENT S AMPLING FROM THE S URVIVAL D ENSITY
The sampling operation t ∼ S ( t | x ) 4 in equationmay induce computational overhead.
To boost the sampling operation, we use log S θ ( t | x ) which was computed during the negative log-likelihood calculation in equation 4.
Let [ t 1 , . . .
, t K ] be the union of the time points in minibatch.
The time points are sorted in increasing order.
The adaptive time stepping in ODE solvers are sensitive to the time interval t K − t 1 rather than the number of time points (Rubanova et al., 2019).
We can access log S θ ( t k | x ) with negligible overhead as long as t 1 < t k < t K .
We sample t k from a categorical distribution whose k -th weight is defined as S θ ( t k | x ) = exp( − (cid:82)
t k 0 h ( τ | x ) dτ ) .
We finalize the sampling process by sampling t from the uniform distribution U ([ t k , t k +1 ]) .
In this way, we don’t have to calculate S ( t | x ) again for sampling t .
See
Algorithm 1 in Appendix for the pseudo code of ODE based survival analysis with the hazard gradient penalty.
3. C ONNECTION TO KL D IVERGENCE
We now show that the hazard gradient penalty in equation 5 is equivalent to minimizing the approximation of the upper bound of the KL divergence between the density function at a data point and that of the neighborhood points of the data point.
Henceforth, we denote X by the subset of d -dimensional real space R d .
Theorem 1 Suppose the hazard function is strictly positive function for all data point x , x ′ ∈ X .
The KL divergence
[Equation]
is upper bounded by
[Equation]
To prove Theorem 1, we need the following lemma.
Lemma
The expectation of survival densities difference under the density is the negative of the expectation of hazard functions difference under the survival density.
In other words,
[Equation]
for all x , x ′ ∈ X
Proof)
We use the fact that E S ( t | x ) (log S ( t | x ) − log S ( t | x ′ )) is constant with respect to t .
[Equation]
Hence,
E
[Equation]
S ( t | we cannot
, we
We now go back to Theorem 1 and prove the theorem.
[Equation]
Theorem 2 An approximation of the upper bound of the KL divergence given in equation 6 is upper bounded by 2 ϵ E S ( t | x ) ∥∇ x h ( t | x ) ∥ 2 if x ′ is in the epsilon ball centered at x , i.e. x ′ ∈ B ( x , ϵ ) .
To prove the theorem
[Equation]
we first find the approximation.
Lemma
[Equation]
x h
[Equation]
t | x ) T ( x ′ − x ) ∥ 2 is an approximation of the upper bound of the KL divergence which is given in equation 6.
[Equation]
and
[Equation]
x
[Equation]
− x ) (cid:13)(cid:13) 2 ■
Obviously, 2 E S ( t | x ) (cid:13)(cid:13) ∇ x h ( t | x ) T ( x ′ − x ) (cid:13)(cid:13) 2 ≤ 2 E S ( t | x ) max x ′ ∈X (cid:13)(cid:13) ∇ x h
( t | x ) T ( x ′ − x ) (cid:13)(cid:13) 2 .
As we assumed x ′ ∈ B ( x , ϵ ) in Theorem 2, max x ′ ∈X (cid:13)(cid:13) ∇ x h
( t | x ) T ( x ′ − x ) (cid:13)(cid:13) 2 is achieved when x ′ − x = ϵ ∇ x h ( t | x ) / ∥∇ x h ( t | x ) ∥ 2 .
Hence,
[Equation]
and this concludes the proof.
■
Theorem 2 shows that regularizating the hazard gradient penalty in equation 5 is equivalent to minimizing the approximation of the upper bound of the KL divergence E p ( t | x ) [log p ( t | x ) − log p ( t | x ′ )] .
To incorporate the regularizer into the negative log-likelihood loss, we minimize the Lagrange multiplier defined as the sum of the negative log-likelihood and the hazard gradient penalty regularizer.
[Equation]
Here, λ is a coefficient that balances the negative log-likelihood and the regularizer.
See Appendix
B for the code snippet of our JAX implementation (Bradbury et al., 2018).
Minimizing the hazard gradient penalty in equation 5 has two advantages over minimizing the KL divergence directly: a) computational efficiency and b) reduced burden of hyperparameter tuning.
To compute the KL divergence, we first sample x ′ ∈ B ( x , ϵ ) .
We then need to compute four values: h ( t | x ) , S ( t | x ) , h ( t | x ′ ) and S ( t | x ′ ) .
In this case, we have to compute hazard values of every t ∼ S ( t | x ) .
Further, we need one more hazard function integration S ( t | x ′ ) =
exp( − (cid:82) h ( t | x ′ )) .
On the other hand, regularizing the hazard gradient penalty only need to calculate the gradient of the hazard function.
When it comes to regularizing the KL divergence, we have to set the appropriate value of the regularizing coefficient λ ′ and the size of the ball ϵ .
On the other hand, if we regularize the hazard gradient penalty, we don’t need to tune ϵ as λ in equation 5 incorporates ϵ .
E XPERIMENTS
In this section, we experimentally show that the hazard gradient penalty outperforms other regularizers.
Further, we check the hyperparameter sensitivity of hazard gradient penalty .
Throughout the experiments, we use three public datasets: Study to Understand Prognoses Preferences Outcomes and Risks of Treatment (SUPPORT) 5 , the Molecular Taxonomy of Breast Cancer International Consortium (METABRIC) 6 , and the Rotterdam tumor bank and German Breast Cancer Study Group (RotGBSG) 7 .
Table 4 summarizes the statistics of the datasets.
See Appendix A for evaluation metrics and experimental details.
4. M ETHODS C OMPARED
We compare ODE + HGP with four methods: vanilla ODE , ODE + L1 , ODE + L2 ODE + LCI .
Vanilla ODE minimizes the expectation of the negative log-likelihood in equation 4.
ODE + L1 minimizes the Lagrange multiplier defined as the sum of the expectation of the negative log-likelihood and the L1 penalty term: E ( x ,t,e )
∼D
L x + α (cid:80) Pp =1 | w p |
Here, w p s are model parameters and α is a coefficient that balances the negative log-likelihood and the L1 penalty term.
This is an extension of Lasso-Cox (Tibshirani, 1997) to the ODE modeling framework.
ODE + L2 minimizes the Lagrange multiplier defined as the sum of the expectation of the negative log-likelihood and the L2 penalty term: E ( x ,t,e )
∼D
L x + α (cid:80) Pp =1 w 2 p
Here, w p s are model parameters and α is a coefficient that balances the negative log-likelihood and the L2 penalty term.
This is an extension of Ridge-Cox (Verweij & Van Houwelingen, 1994) to the
ODE modeling framework.
ODE + LCI minimizes the Lagrange multiplier defined as the sum of the expectation of the negative log-likelihood and the negative of the lower bound of a simplified version of time-dependent C-index.
The regularizer is defined as − (cid:88)
[Equation]
[Table]
This is equivalent to time dependent concordance index in Section A.1.1 if we don’t take the KaplanMeier estimator into account.
The regularizer is a reminiscent of the lower bound of C-index (Steck et al., 2007).
Although the lower bound of C-index was originally proposed as a substitute of the negative log-likelihood, Chapfuwa et al.
(2018) used the lower bound (Steck et al., 2007) as a regularizer of the AFT model (Wei, 1992).
4. R ESULTS
Table 1 shows the mC td , mAUC , and iNBLL scores 8 .
The hazard gradient penalty outperforms other methods across almost all metrics and datasets.
The interesting point is that both L1 and L penalties do not affect the ODE model’s performance in most cases.
We speculate that regularizing the weight norm is effective in CoxPH as the model is simple and has a strong assumption that the hazard rate is constant.
On the contrary, regularizing the norm of the weight may not be able to affect the ODE model’s performance as ODE models are much more complex than CoxPH.
Also, the experimental results highlight the possibility that the performance of the survival analysis models is more related to the local information such as the gradient at each data point rather than the global information such as the weight norm of the model.
Figure 2 shows that the ODE + HGP effectively regularized the variation of the density with respect to the input while other methods could not.
Table 1 also shows that regularizing the lower bound of the C-index is not effective in many cases.
We conjecture that the method is ineffective as the ODE modeling framework is flexible and optimizing the negative log-likelihood can discriminate each data point’s rank.
Furthermore, regularizing the lower bound of the C-index does not harness the information of neighbors of data points.
We also compare HGP against a Neural ODEs specific regularizer.
See Appendix D for details.
[Table] [Table]
Table 2 shows the results by varying the number of samples M in the sampling process t ∼ S ( t | x ) in equation 5.
As long as the regularizer is applied, the number of samples M does not affect the performance.
Even when M = 1 , the regularizer works well.
Figure 3 shows the results on
SUPPORT and RotGBSG datasets by varying the coefficient λ in equation 7.
Since the performance variation by λ is stable, the hyperparameter λ can be tuned without much difficulty in practical setups.
Table 3 shows the time taken for training/evaluation step for each methods.
The time it takes to train ODE + L1 and to train ODE + L2 is slightly faster than vanilla ODE training time.
It is straightforward to see that the L1 and L2 penalty makes smooth dynamics as the number of function
[Figure]
datasets
We
D e by the of uncensored
We choose x ′ =
x
+ ϵ g / ∥ g ∥ 2 where g = ∇ x log p θ ( t | x ) so that log p ( t | x ) − log p θ ( t | x ′ )
∥ 2 is maximized under x ′ ∈ B ( x , ϵ ) constraint.
We set ϵ = 1 e −across all experiments.
The hazard gradient penalty effectively regularizes the variation.
As the density of probability distribution p ( · | x ) should be concentrated at t for uncensored data ( x , t ) the figures show that the hazard gradient penalty effectively regularizes the KL divergence between the density function at a point x and that of neighborhood points x ∈ B ( x
[Figure]
evaluations (NFE) in the training time of ODE + L1 and ODE + L2 are lower than that of vanilla
ODE.
The decrease in NFE compensates for the overheads of calculating L1 and L2 penalties, which makes the training time of ODE + L1 and ODE + L2 faster than that of vanilla ODE.
The same applies to ODE + HGP .
Despite the additional sampling process and gradient calculation, the training time of ODE + HGP is on par with vanilla ODE.
The decreased NFE thanks to smooth dynamics made by the hazard gradient penalty compensates for the additional computations.
The smooth dynamics made by L1, L2, and the hazard gradient penalty can also be observed in the evaluation phase.
The NFE of ODE + L1 , ODE + L2 , and ODE + HGP is much smaller than that of vanilla ODE.
Not all regularizers make smooth dynamics.
The NFE of ODE + LCI is comparable to the NFE of vanilla ODE.
Clearly, there’s no point in regularizing LCI in terms of inference speed.
[Table]
R ELATED W ORKS
A line of research integrated deep neural networks to CoxPH (Faraggi & Simon, 1995; Katzman et al., 2018) and Extended Hazards (Zhong et al., 2021) for more model flexibility.
Another line of research proposed distribution-free survival analysis models via the time domain discretization (Lee et al., 2018) or adversarial learning approach (Chapfuwa et al., 2018).
Previous works (Goldstein et al., 2020; Han et al., 2021) proposed new objectives to optimize Brier score (Graf et al., 1999),
Binomial log-likelihood, or distributional calibration directly.
Yet to the best of our knowledge, none of the previous works focused on the effect of gradient penalty on survival analysis models.
Previous works proposed L1 and L2 regularization in the survival analysis literature (Tibshirani, 1997; Verweij & Van Houwelingen, 1994).
Those methods regularize the survival analysis models so that the L1 or L2 norm of the model parameters does not increase so much.
Our method is different from those methods in that we penalize the norm of the gradient on each local data point.
Our method is closely related to semi-supervised learning (Chapelle et al., 2006).
Among many semi-supervised learning methods, our method is germane to virtual adversarial training (Miyato et al., 2018) in that it regularizes function variation between a local data point and its neighbours.
However, virtual adversarial training is different from ours in that the method was demonstrated in the classification setting and the output is a discrete distribution.
In Generative Adversarial Nets (GANs) literature (Goodfellow et al., 2014), the gradient penalty had been studied actively.
Gulrajani et al.
(2017) proposed the gradient penalty to satisfy the 1-Lipschitz function constraint in Kantrovich-Rubinstein duality.
Mescheder et al.
(2018) proposed the gradient penalty to penalize the discriminator for deviating from the Nash equilibrium.
Ours is different from these works in that we propose gradient penalty so that the density at x does not deviate much from that of x ’s neighborhood points.
C ONCLUSION
In this paper, we introduced a novel regularizer for survival analysis.
Unlike previous methods, we focus on individual local data point rather than global information.
We theoretically showed that regularizing the norm of the gradient of hazard function with respect to the data point is related to minimizing the KL divergence between the data point and that of its neighbours.
Empirically, we showed that the proposed regularizer outperforms other regularizers and it is not sensitive to hyperparameters.
Furthermore, the proposed regularizer is computationally efficient and incurs an ignorable overhead.
Nonetheless, as minimizing the proposed regularizer may conflict with optimizing the negative log-likelihood, practitioners should tune the balancing coefficient λ for each dataset.
The paper highlights the new possibility that the recent advancements in semi-supervised learning could enhance the performance of survival analysis models.
A E VALUATION M ETRICS AND E XPERIMENTAL D ETAILS A. E VALUATION M ETRICS
Throughout this subsection, we denote ˆ S ( t | x ) as the estimate of S ( t | x ) , I ( · ) as the indicator function, ( x i , T i , e i ) as the i th covariate, time, event indicator of the dataset, ˆ G ( t ) as the KaplanMeier estimator for censoring distribution (Kaplan & Meier, 1958), and ω i as 1 / ˆ G ( T i ) .
A.1. T IME D EPENDENT C ONCORDANCE I NDEX ( C td )
The concordance index, or C-index is defined as the proportion of correctly ordered pairs among all comparable pairs.
We use time dependent variant of C-index that truncates pairs within the prespecified time point Uno et al.
The time dependent concordance index at t , C td ( t ) , is defined as
[Equation]
T
[Equation]
<
T
[Equation]
T
[Equation]
< t )
To evaluate C td at [ t 1 , . .
, t L ] at the same time, we take its mean mC td = 1 L (cid:80)
Ll =1 C td ( t l ) .
A.1. T IME D EPENDENT A REA U NDER C URVE (AUC)
is an extension of the ROC-AUC to survival data Hung & Chiang (2010).
It measures how well a model can distinguish individuals who fail before the given time ( T i < t ) and who fail after the given time ( T j > t ).
The AUC at time t , AUC ( t ) , is defined as
[Equation]
To evaluate AUC at [ t 1 , . .
, t L ] at the same time, we take its mean mAUC =
1 L (cid:80) Ll =1 AUC ( t l ) .
A.1. N EGATIVE B INOMIAL L OG -L IKELIHOOD
We can evaluate the negative binomial log-likelihood (NBLL) to measure both discrimination and calibration performance Kvamme et al.
The negative binomial log-likelihood at t measures how close the survival probability is to 1 if the given data survived at t and how close the survival probability is to 0 if the given data failed before t .
The NBLL at t , NBLL ( t ) , is defined as
[Equation]
For the convenience of evaluation, we integrate the NBLL, iNBLL =
[Equation]
A. E XPERIMENTAL D ETAILS
Across all datasets, we split training set, validation set and test set into 70%, 10% and 20% each using PyTorch ’s random split function Paszke et al.
We set seed = 42 when splitting.
Across all experiments, we use an MLP with two hidden layers where each layer has 64 hidden units.
Across all layers, we apply Layer normalization Ba et al.
Instead of naively feeding time t into the neural network, we feed scaled time ˜ t
= ( t − t 2 ) /
( t 3 − t 1 ) where t 1 , t 2 , and tare first, second, and third quartile of failure event times.
We found that this strategy enhances the
ODE model’s performance and boosts training time.
To incorporate time t into the survival analysis model, we project the time into an eight dimensional vector using a single layer MLP and then concatenate it to the input data.
The time t is also specified by adding projected output into each layer output.
We use the AdamW optimizer Loshchilov & Hutter (2019) and clipped the gradient norm so that it does not exceed 1.
We set the learning rate to 0.001.
We have implemented the code using JAX Bradbury et al.
and Diffrax Kidger (2021) 9 .
To find the best λ in equation 7, we run experiments with λ = 1 , 5 , 10 , 50 and report the results at λ = 10 as it shows decent performance across all metrics and datasets.
We also have to set the number of samples M from the time sampling process t ∼ p θ ( t | x ) in equation 5.
We set
M = 5 across all the hazard gradient penalty experiments.
To find the best coefficient α in ODE + L1 , ODE + L2 , and ODE + LCI experiments, we set α = 1 e − 1 , 1 e − 2 , 1 e − 3 and run the experiments.
We report the best α in terms of mAUC .
To report mC td and mAUC , we calculate C td and AUC at 10%, 20%, . .
, 90% event quantiles and average them.
To report iNBLL , we integrate from the minimum time of the test set to the maximum time of the test set.
We use scikit-survival P ¨ olsterl (2020) to report mC td and mAUC .
We use pycox Kvamme et al.
(2019) to report iNBLL .
Across all experiments, we run 7 experiments with different seeds and report their mean and the standard deviation.
B C ODE S NIPPET OF HGP
16return jax.grad(
, argnums params,
) 18 19 @jax.jit 20 def loss_func( params: optax.
Params, X: jnp.ndarray, event: jnp.ndarray, t: jnp.ndarray, timestamp: jnp.ndarray, unique_idx: jnp.ndarray, key: jnp.ndarray 24 ): batch_size =
X.shape[0] time_size = jnp.size(timestamp) 27timestamp
= jnp.insert(timestamp, 0, 0) int_hazard_t0 = jnp.zeros((X.shape[0], 1)) states_t0 = jnp.concatenate((
_t
, axis
1) ys = odeint(timestamp, states_t0, params) int_hazard =
ys[1:, :, -1]
33log_surv = - int_hazard log_surv_t = log_surv[unique_idx, jnp.arange(X.shape[0])] hazard_t = hazard_func(t, states_t0, args=params)
[:, -1] 37log_prob_t = jnp.log(hazard_t + 1e-6)
+ log_surv_t 39assert log_surv_t.shape ==
log_prob_t.shape assert event.shape =
=
log_prob_t.shape 42nll = - (event * log_prob_t + (1. - event)
* log_surv_t).mean() 44# now [batch_size, time_size] logits = jnp.transpose(jax.lax.stop_gradient(log_surv), (1, 0))
time_idx = jrandom.categorical( key, jnp.tile(logits[:, :, None], (1, 1, args.sample_size)), axis=  ) 51t0 = timestamp[0:][time_idx] t1 = timestamp[1:][time_idx] t_prime = t0 + (t1-t0)
* jrandom.uniform(key, shape=t0.shape) assert t_prime.shape
==
(batch_size, args.sample_size)
# E_{S(t | x)} || \nabla_x h(t | x) || dh_dx = jax.vmap(calc_dh_dx, in_axes=(None, 1, None))(params, t_prime , X) assert dh_dx.shape
==
(args.sample_size, batch_size, X.shape[1]) dh_dx_norm =
jnp.linalg.norm(dh_dx, ord=2, axis=-1).mean() 61loss = nll + args.lambda_
* dh_dx_norm 63return loss, (nll, dh_dx_norm)
C T HE E XISTENCE OF THE P ROBABILITY D ISTRIBUTION s θ ( t | x )
In general, we cannot guarantee the existence of probability distribution s θ ( t | x ) =
S θ ( t | x ) (cid:82)
S θ ( t | x ) dt as the integration (cid:82) S θ ( t | x ) dt may not exist.
To ensure the existence of (cid:82) S θ ( t | x ) dt , we simply add a constraint: h θ ( t | x ) ≥ ϵ .
Here, ϵ is a very small constant (e.g. ϵ = 1 e − 8 ).
As ϵ is a very small constant, it has a negligible impact on the algorithm.
The
constraint can be achieved easily by adding ϵ to the softplus output of the hazard function.
[Table]
[Table] [Table]
Under the constraint, S θ ( t | x ) is bounded by an exponential function:
[Equation]
Also, (cid:82) S θ ( t | x ) dt is bounded:
[Equation]
D C OMPARISON A GAINST N EURAL ODE S S PECIFIC R EGULARIZER
In this section, we compare HGP against Neural ODEs specific regularizer: STEER (Ghosh et al., 2020).
STEER regularizes Neural ODEs by perturbing the final time of the integration.
Tablecompares the performance of HGP against STEER.
Overall, HGP outperforms STEER on various setups.
E T HE TAKEN TIME AND THE NUMBER OF FUNCTION EVALUATIONS
In this section, we compare HGP against competing regularizers in terms of the taken time and the number of function evaluations on SUPPORT dataset.
See Table 6 for the details.
Overall, the result aligns with that of Table 3.