FastSurvival: Hidden Computational Blessings in Training Cox Proportional Hazards Models

Thank you very much for reviewing our paper. Please see below for our answers to your questions.

1. Runtime reduction results for computing derivatives are trivial.

Please look at Line 146. The result is surprising especially for computing the exact 2nd order partial derivatives.

Even for the quasi and proximal Newton methods which make diagonal approximation of Hessian (Line 93-94), the complexity is still high. We need to calculate the full Hessian in the space of $\boldsymbol{\beta}$ via the formula $\boldsymbol{X}^T H(\boldsymbol{\eta}) \boldsymbol{X}$, whose complexity is $O(n p^2)$ when $H(\boldsymbol{\eta})$ is diagonal. Plus, the Newton method requires $O(p^3)$ to solve a linear system via the LDL method. In total, it is $O(np^2 + p^3)$.

In contrast, we optimize each coordinate with complexity $O(n)$, and in total $O(np)$ for $p$ coordinates. This reduction is not at all trivial, and we support this claim through empirical results on real-world datasets.

2. Experimental results on runtime are misleading

Your preprocessing of the dataset is different from ours; from running on the wrong preprocessed datasets, one can't conclude our experimental results are wrong. As written in Line 257-259, we perform binary thresholding to preprocess the data. You didn't do thresholding for continuous variables, which is what makes the variables correlated and challenging. In Appendix C2, our preprocessed Flchain dataset has 333 encoded binary features. Since you have only 38 features, you're using a totally different dataset.

The sksurvCPH you used is also not the exact Newton method. They used a trick on Line 463 of the file coxph.py on the GitHub repo. They take half the exact Newton step size if the loss goes up. This trick does not work in general. When $\lambda_2 = 0.001$, sksurvCPH runs over 100 seconds without converging. The loss is 14522.97. In contrast, our method runs under 40 seconds, and the loss is 10659.91 (high precision solution).

We uploaded to Anonymous GitHub the complete code (existing Newton methods and preprocessing steps) to support our claims. Due to the rebuttal policy, we only shared the link with the Area Chair, who can verify results in Section 4.1 and this response by running our codes.

3. Eq. (6), how do we get from $\boldsymbol{\beta}$ to $\boldsymbol{\eta}?$

As mentioned in Line 89, we use an intermediate variable $\boldsymbol{\eta}$ with $\boldsymbol{\eta = X \beta}$. By the chain rule, we have $\nabla_{\boldsymbol{\beta}} \ell(\boldsymbol{\beta}) = \boldsymbol{X}^T \nabla_{\boldsymbol{\eta}} \ell(\boldsymbol{\eta})$ and $\nabla_{\boldsymbol{\beta}}^2 \ell(\boldsymbol{\beta}) = \boldsymbol{X}^T \nabla_{\boldsymbol{\eta}}^2 \ell(\boldsymbol{\eta}) \boldsymbol{X}$.

4. Probabilistic interpretations are not novel

Thank you for pointing out this paper; it does not overlap with the novel contributions in this submission. We will cite this reference during revision. The reference you cited has only basics. The equation at the bottom of page 1 boils down to the optimality condition when the gradient is equal to 0. Our contributions are substantially different. We show that

• The 2nd- and 3rd-order partial derivatives have probabilistic interpretations, which is unexpected.
• The 2nd- and 3rd-order partial derivatives are in the same formulas as the 2nd- and 3rd-moment calculations.
• The Lipschitz constants can be computed explicitly for the 1st- and 2nd-order partial derivatives.

5. Why is Eq.17 an improvement over and different from the existing Newton methods?

See the Abstract (Line 7-9), Introduction (Line 33-39), and Preliminaries (Line 95-104). To reiterate:

Problems with existing Newton methods

a) They all have trouble converging due to the vanishing Hessian. This is a well-known issue for Newton-based methods (see left two plots in Fig 1).

b) Even if they converge, none of them can converge with high precision fast enough (Line 100-104; see right two plots in Fig 1).

Difference over previous Newton methods

a) We always converge to optimal solutions with high precision. For highly correlated features, other methods could have the loss blow up.

b) Without our formula to calculate $L_2$ in Eq.13, it is impossible to implement Eq.17.

c) Without our formula to calculate $L_3$ in Eq.14, it is impossible to implement for Eq.18.

6. The cubic equation and formula have existed for several hundred years

The cubic equation has almost nothing to do with our cubic surrogate function, despite sharing the same word ``cubic''. The tasks are completely different. In our work, we minimize a special polynomial whose highest order is 3. One core step is to find roots to the second-order polynomial (see Appendix A4 for details). In contrast, the cubic equation and cubic formula are used to find the roots to the third-order polynomial. We never used the cubic formula in our work. We also solve $\ell_1$-regularized cubic surrogate problem, which has not been studied before.

7. Why is the method robust to highly correlated features in variable selection, and is there any math theory?

At the core, our methods work because they allow solutions to converge with high precision quickly. During beam search, having high-precision solutions allows us to compare solutions and pass only the best ones to the next stage of support expansion. If one doesn't get high-precision solutions, some medium- or low-quality solutions can be passed to the next stage, which hurts the overall solution quality. For theories, there are two relevant papers ([13] and [47] in our reference). We can combine their proofs to show our beam search on the Cox model is an approximation algorithm, which means there is some quality guarantee on the final solution.

Hopefully we have answered all your questions and cleared all your doubts!

Thank you very much for reviewing our paper. Please see below for our answers to your questions.

1. Proof of convexity of negative log likelihood (nll)

The essential part of the nll is the logSumExp function, which is defined as $f(\boldsymbol{x}) = \log(\sum_{i=1}^m \exp(x_m))$. It is well known the logSumExp function is convex (see Optimization Models and Applications by Laurent El Ghaoui, 2017).

Moreover, the composition of a convex function with an affine (also linear) function is still convex (see the proof on Math Stack Exchange with question ID 654201). Therefore, $\log(\sum_{j \in R_i} e^{\boldsymbol{x}_j^T \boldsymbol{\beta}})$ is convex.

Lastly, the sum of two convex functions is still a convex function. Therefore, $\sum_{i=1}^n \delta_i \left[ \log(\sum_{j \in R_i} e^{\boldsymbol{x}_j^T \boldsymbol{\beta}}) - \boldsymbol{x}_i^T \boldsymbol{\beta} \right]$ is convex.

2. Is sorting needed during optimization?

Please look at our answer in the general response.

3. Page 15, math typos associated with Appendix A.1.1

Please look at our answer in the general response.

4. Minor writing typo

Thank you for your careful reading. We will fix this (wise $\rightarrow$ wide) during revision.

Thank you very much for reviewing our paper. Please see below for our answers to your questions.

1. Utility of Cox Model

We want to emphasize that we agree with what the reviewer has commented and simply want to continue this conversation. We think there are three aspects that make the Cox model still worthwhile and relevant:

(1) The Cox model is relevant in high stakes domains such as medicine. In these domains, professionals need to interpret the results and understand feature importance. One way to achieve this is to obtain as small number of coefficients as possible without losing predictive performance. In Section 4.2, we have pushed the frontier of this variable selection task.

(2) The Cox model is useful when served as a baseline to compare against more advanced and complicated models. If the more complicated model doesn't outperform the Cox model, we probably want to use the simpler Cox model, according to the principle of Occam's Razor.

(3) Even if we do not deploy the Cox model in the end, it is still beneficial to use during the data exploration stage. When data scientists deal with messy real-world datasets for the first time, they want to perform assessments to see whether there are any abnormalities. It is cumbersome to explore if they use a much more complicated model. Moreover, fixing the abnormalities found by the Cox model could potentially boost the performance of the complicated one.

We totally respect the reviewer's expertise in survival analysis. We agree with the reviewer that there exist survival models that are not based on the proportional harzard assumption. In Appendeix B Related Work, we have discussed some SOTA survival models. For example, OST and OSST are two decision tree methods. Besides the Cox model, there are also the accelerated failure time (AFT) model and the Aalen’s additive model. Lastly, other models include the ensemble method (SurvivalQuilts) and neural networks (DeepSurv). If you think there are other SOTA survival models we missed, please let us know, we are happy to cite them during the revision.

2. Line 93-94, formula for quasi-Newton method

Thank you for this suggestion. We are afraid our current notation cannot be further simplified. In our definition, $\text{diag}(\cdot)$ (see Line 94) takes a vector as input and outputs a square matrix whose diagonal equals to this input and other entries should be $0$. However, $\nabla^2_{\boldsymbol{\eta}} \ell(\boldsymbol{\eta})$ is a square matrix. Thus, we cannot send $\nabla_{\boldsymbol{\eta}}^2 \ell(\boldsymbol{\eta})$ as input to $\text{diag}(\cdot)$. In the paper, we had some comments (in green color) to help clarify things. This was the most succinct notation we could come up. If you have another idea to simplify Line 93-94, please let us know, and we are very happy to incorporate that during revision.

3. Line 257, definition of feature correlation

The correlation used in our work is neither based on Pearson nor Spearman. Please see Line 709-714 in Appendix C2 for the context where the feature correlation is used. For the synthetic dataset, the feature $\boldsymbol{x}_i$ is sampled from a Gaussian distribution:

$\boldsymbol{x}_i \sim \mathcal{N}(0, \boldsymbol{\Sigma})$,

where $\boldsymbol{\Sigma}$ is the the (Toeplitz) covariance matrix with $\Sigma_{jl} = \rho^{|j - l|}$, and $\rho \in (0, 1]$ is the correlation parameter. When $\rho$ is large, like $\rho=0.9$ in our paper, the features become highly correlated. This serve as a great benchmark in Section 4.2 to test how well different algorithms can recover the sparse coefficients. As you have seen, our new method significantly outperforms other methods.

4. Page 15, math typos associated with Appendix A.1.1

Please look at our answer in the general response.

5. Other minor writing typos

We will fix them during revision. Thank you very much for point these out!

Thank you very much for reviewing our paper. Please see below for our answers to your questions.

1. Comparison with fastCPH and BigSurvSGD on computational efficiency

Thank you for pointing out these two baselines. We will cite them during revision.

Regarding BigSurvSGD, below is a comparison between BigSurvSGD and our methods on Experiment 4.1. BigSurvSGD finishes running in a few seconds with low-quality solutions. To compare convergence speed, we set the time limit for our methods so that they run with the same or less amount of time than BigSurvSGD. We can see that our methods achieve much smaller training losses than bigSurvSGD within the same or less wall-clock time, especially using the cubic method. We also want to point out that we implement our methods in python (we loop through all coordinates sequentially) whereas BigSurvSGD is implemented in C++.

Note that bigSurvSGD uses stochastic gradient descent. Stochasticity introduces random noise and does not ensure monotonic loss decrease. This is problematic if we want to use SGD inside beamsearch for the variable selection task in presence of highly correlated features. Moreover, for gradient descent, we have mentioned in Line 84-85 that it is hard to pick the right step size, which significantly impacts the convergence speed.

Regarding FastCPH (LassoNetCoxRegressor from the lassonet package), this uses a neural network instead of a pure linear model. It is not allowed to specify no hidden layers for FastCPH (# of neural network layers $>$2; see Line 19 in the file model.py in the lassonet package on GitHub). Therefore, we cannot use this package to train a linear Cox model.

2. Comparison with lasso on solution quality

In Section 4.2, we have already included the lasso baseline. See the baseline sksurvCoxnet (the default $\ell_1$ ratio is $1.0$). During revision, we will make it clearer that skusrvCoxnet is the lasso baseline.

3. Tightness of bounds when optimizing the bound instead of the loss function

At each iteration, all current methods (existing Newton methods or our methods) try to minimize a function that approximates the original loss function.

Existing Newton methods They approximate the original loss function through a 2nd-order Taylor expansion (see Line 87-94). The exact Newton method computes the Hessian exactly while quasi and proximal Newton methods compute the Hessian approximately.

Our quadratic surrogate function Here, we also make a 2nd-order Taylor expansion but with two differences. Firstly, we only do the approximation in a single coordinate. Secondly, the expansion's 2nd-order coefficient is an upper bound on the exact 2nd-order derivative. For tightness, if we use the exact 2nd-order derivative instead of the upper bound, we can get a better approximation, which we discuss next.

Our cubic surrogate function Here, we make a 3rd-order Taylor expansion on a single coordinate. We use the exact 1st- and 2nd -order partial derivatives and also use an upper bound on the 3rd-order partial derivative. For tightness, we obtain a more accurate approximation of the Cox function than our quadratic surrogate function, at the neighborhood of $x$.

Both of our methods ensure monotonic loss decrease and global convergence while existing Newton methods do not. Please see the proof sketch in our general response.

4. Can we construct confidence intervals for coefficients?

This is not the goal of our paper, but yes, it is possible, after the optimal solution is obtained.

After the optimal solution $\boldsymbol{\beta}^*$ is obtained, we first calculate the Hessian $\boldsymbol{H}$ with respect to the loss function. Next, we calculate the standard error vector $\boldsymbol{s} = \text{sqrt}\left[ (\boldsymbol{H}^{-1}).\text{diag}() \right]$, where $\boldsymbol{Z}.\text{diag}()$ means taking the diagonal of the square matrix $\boldsymbol{Z}$ as a vector. The confidence interval is then $\boldsymbol{\beta}^* \pm \alpha \boldsymbol{s}$, where $\alpha \in \mathbb{R}_+$ corresponds to different confidence levels. Please see the book In All Likelihood by Yudi Pawitan for details. Also, please see an implementation in the lifelines package on GitHub.

5. Is sorting needed during optimization?

Please look at the our answer in the general response.

6. How can this work be extended to time-varying features?

It takes more than preprocessing and passing a new dataset to our algorithm. Time-varying features are not our focus in this work, but we will try our best to answer this question. For details, please see the lifelines documentation page on this topic.

For time-invariant features, we only care about time duration, and reverse cumulative sum works fine. However, for time-varying features, we cannot use the vanilla reverse cumulative sum. At different times, new features can enter the loss function (while old features exit). Our optimization techniques are still applicable, but we need to write a customized reverse cumulative sum function.