Title: Differentiable sorting for censored time-to-event data.
Venue: NeurIPS 2023
Accepted: yes
Keywords: Survival Analysis

Censored Data

Semi-supervised Learning

Time-to-event-data

Algorithmic Supervision

Sorting

Risk Prediction

Weakly-supervised Learning

Machine Learning

Cox's Partial Likelihood

Differentiable Sorting Networks

Transitive Inductive Bias

Ranking Losses

Listwise Ranking

Healthcare Applications

Deep Learning

Neural Networks

Top-k Risk Prediction
Abstract: Survival analysis is a crucial semi-supervised task in machine learning with significant real-world applications, especially in healthcare. The most common approach to survival analysis, Cox’s partial likelihood, can be interpreted as a ranking model optimized on a lower bound of the concordance index.  We follow these connections further, with listwise ranking losses that allow for a relaxation of the pairwise independence assumption. Given the inherent transitivity of ranking, we explore differentiable sorting networks as a means to introduce a stronger transitive inductive bias during optimization. Despite their potential, current differentiable sorting methods cannot account for censoring, a crucial aspect of many real-world datasets. We propose a novel method, Diffsurv, to overcome this limitation by extending differentiable sorting methods to handle censored tasks. Diffsurv predicts matrices of possible permutations that accommodate the label uncertainty introduced by censored samples. Our experiments reveal that Diffsurv outperforms established baselines in various simulated and real-world risk prediction scenarios. Furthermore, we demonstrate the algorithmic advantages of Diffsurv by presenting a novel method for top-k risk prediction that surpasses current methods.

Abstract

Survival analysis is a crucial semi-supervised task in machine learning with significant real-world applications, especially in healthcare. The most common approach to survival analysis, Cox's partial likelihood, can be interpreted as a ranking model optimized on a lower bound of the concordance index. We follow these connections further, with listwise ranking losses that allow for a relaxation of the pairwise independence assumption. Given the inherent transitivity of ranking, we explore differentiable sorting networks as a means to introduce a stronger transitive inductive bias during optimization. Despite their potential, current differentiable sorting methods cannot account for censoring, a crucial aspect of many real-world datasets. We propose a novel method, Diffsurv, to overcome this limitation by extending differentiable sorting methods to handle censored tasks. Diffsurv predicts matrices of possible permutations that accommodate the label uncertainty introduced by censored samples. Our experiments reveal that Diffsurv outperforms established baselines in various simulated and real-world risk prediction scenarios. Furthermore, we demonstrate the algorithmic advantages of Diffsurv by presenting a novel method for top-k risk prediction that surpasses current methods. In conclusion, Diffsurv not only provides a novel framework for survival analysis through differentiable sorting, but also significantly impacts real-world applications by improving risk stratification and offering a methodological foundation for developing predictive models in healthcare and beyond.

1 Introduction

Survival analysis plays a pivotal role in many realworld machine learning applications, spanning fields such as reliability engineering, marketing, and insurance, with a particularly significant impact in healthcare. The goal of survival analysis is to predict the time until the occurrence of an event of interest, such as death, based on a set of covariates. In clinical studies, these include demographic variables such as sex and age, but may also encompass more complex data modalities such as medical images.

The concept of censoring is a distinguishing characteristic that sets survival analysis apart from 28 conventional machine learning approaches. Particularly prevalent in observational datasets, it refers 29 to situations where event times remain unobserved because a patient might not have undergone the 30 event by the time of data collection. This can be due to a variety of reasons such as the study period 31 ending before all events of interest have occurred or subjects leaving the study. 32

33

Overlooking censoring can skew predictions towards the censoring event, rather than the event of

interest. This bias becomes particularly noticeable when the study's endpoint can be inferred from the 34

observed covariates - age being a notable example. In such cases, the predicted event times are likely 35

to biased towards the censoring event time, thereby neglecting the actual event of interest [Kvamme 36 and Borgan, 2019]. 37

The Cox Proportional Hazards model is widely used for handling censored data in survival analysis 38 [Cox, 1972]. The model optimizes a partial likelihood function over ranked data, considering only 39 the order of events, not their exact time of occurrence. As such, Cox's partial likelihood serves as a 40 ranking loss, learning from the order of patients based on their hazard of experiencing an event, not 41 their exact survival time. 42

Raykar et al. [2007] showed that Cox's partial likelihood (CPL) and ranking losses can be directly 43 equated, with both providing lower bounds to the concordance index, the primary evaluation metric 44 used in survival analysis. Both losses are foundational to many survival deep learning methodologies 45 like DeepSurv Katzman et al. [2018] and DeepHit Lee et al. [2018]. 46

47

48

49

50

51

52

53

However, this relation operates under the assumption of pairwise independence. This simplification, while practical, can de-emphasize the transitive properties inherent in survival data. As shown by Goldstein and Langholz [1992], larger risk set sizes can lead to more efficient estimators, suggesting potential benefits in considering listwise ranking losses Cao et al. [2007]. These losses optimize over lists of values rather than individual pairs, thereby better capturing the transitive dynamics of the data. Despite the similarities, listwise losses have remained largely unexplored within the field of survival analysis. This could be partly due to an uncertainty around how to handle censoring.

54

55

56

57

58

59

60

61

62

63

64

65

66

67

68

69

70

71

72

We propose a new approach that takes advantage of recent developments in continuous relaxations of sorting operations, allowing end-to-end training of neural networks with ordering supervision [Grover et al., 2019, Blondel et al., 2020, Petersen et al., 2021]. This method incorporates a sorting algorithm into the network architecture, where the order of the samples is known, but their exact values are unsupervised. With this, we introduce Diffsurv , an extension of differentiable sorting methods that enables end-to-end training of survival ranking models with censored data.

Briefly, our contributions are summarised:

Our primary contribution is the extension of differentiable sorting methods to account for censoring by introducing the concept of possible permutation matrices. (Section 3.1)

We empirically demonstrate that our new differentiable sorting method matches or improves risk ranking performance across multiple semi-simulated and real-worlds censored datasets. (Section 4)

We investigate the role of transitivity in survival analysis and show, through experiments with semi-simulated data, that differentiable sorting networks can benefit from this inherent property of the data. (Section 4)

We demonstrate that differentiable sorting of censored data enables the development of new methods with practical applications, using the example of end-to-end learning for top-k risk stratification. (Section 3.2)

2 Survival Analysis and its Relation to Ranking

A dataset with censored event times is summarized as D = { t i , x i , δ i } N i =1 , where N is the total 73 number of patients. For a patient i , the time-to-event t i is the minimum of the true survival time 74 t ∗ i and the censoring time c ∗ i , with δ i indicating whether an event ( t ∗ i ≤ c ∗ i , δ i = 1 ) or censoring 75 ( t ∗ i > c ∗ i , δ i = 0 ) was observed. Covariates are x i ∈ R d representing a d -dimensional vector but the 76 methods discussed here also generalise to higher dimensional tensors such as image data. 77

The widely-used method for addressing censoring in survival analysis is the Cox Partial Likelihood 78 (CPL) model, introduced by Cox [1972]. The CPL is designed to maximize the following general 79 form: 80

CNN

z

7

z

6

z

5

z

4

z

3

z

2

z

1

P

1

·

P

2

·

P

3

·

P

4

·

P

5

·

P

6

·

P

7

=

P

P

=

















.

5

.

3

.

1

.

1

.

0

.

0

.

0

.

4

.

3

.

1

.

1

.

1

.

0

.

0

.

1

.

2

.

4

.

1

.

1

.

1

.

0

.

0

.

1

.

2

.

4

.

2

.

1

.

0

.

0

.

1

.

1

.

2

.

3

.

2

.

1

.

0

.

0

.

1

.

1

.

2

.

3

.

3

.

0

.

0

.

0

.

0

.

1

.

3

.

6















Q

p

=

















1

1

1

1

1

1

1

1

1

0

0

0

0

0

0

1

1

1

1

1

1

0

1

1

1

1

1

1

0

1

1

1

1

0

0

0

0

1

1

1

1

1

0

0

1

1

1

1

1















p

=

∑

n

j

=1

(

Q

p

◦

P

)

i,j

L

=

∑

n

i

=1

y

i

log(

p

i

)

Survival

TopK

· · ·

where f θ is the hazard function, a score prediction function estimating the probability of an event at a 81 particular time, given input features x i . The product only includes uncensored patients, whereas the 82 denominator term also includes censored patients with t j > t i . 83

84

85

86

87

Reflecting the structure of survival data, the Cox Partial Likelihood (CPL) model compares individuals still "at risk" at each time point, similar to a nested case-control study. This directly shapes the likelihood equation in CPL, with the numerator representing the hazard function for the eventexperiencing individual, and the denominator summing over all individuals still at risk.

88

89

90

91

92

Extensions of the Cox model, like [Katzman et al., 2018] and [Kvamme et al., 2019], have modified f θ = exp ( θ · x i ) to relax the linear covariate interaction and proportional hazards assumptions. Introducing neural networks h θ to handle the non-linearity, adjusting f θ to be f θ = exp ( h θ ( x i )) , and to manage non-proportional hazards, they set f θ = exp ( h θ ( x i , t i )) .

2.1 Pairwise Independence

Both of these previous works note that the risk set R = { j : t j > t i } is intractable for deep learning 93 applications as it considers all comparable patients. To mitigate memory constraints, we can sample 94 a fixed-size risk set, denoted as ˜ R , such that | ˜ R| = n < N . Kvamme et al. [2019] go further, arguing 95 it is reasonable to take a constant sample size of 1 and include the individual i in the risk set (such 96 that n = 2 ). This leads to the simplified loss of the form 97

Further, take the mean log partial likelihood to be 98

where n e is the number of non-censored events. In this simplified form, it can be seen that the partial 99 likelihood only considers the pairwise relative ordering or ranking of survival times. 100

The concordance index or c-index Harrell et al. [1982] is a commonly used as an evaluation for survival 101 analysis methods and is a generalization of the Area Under the Receiver Operating Characteristic 102 Curve (AUROC) that handles right-censored data. It is defined as 103

Raykar et al. [2007] first showed that the Cox's partial likelihood is approximately equivalent to 104 maximizing the concordance index or c-index and that closer bounds can be found by minimizing the 105

general ranking loss, with acceptable pairs A = { ( i, j ) : δ i = 1 ∧ t j > t i } and 106

where ϕ is a function that relaxes the non-differentiable ✶ of the c-index. From Equation 3, it can be seen that ϕ : x →-log(1+exp( -x )) = log( σ ( x )) . Here, we have shown that the simplifications to the partial likelihood made by Kvamme et al. [2019] are equivalent to using the log-sigmoid ranking loss.

The key difference between ranking and partial likelihood losses comes when considering the assumption that it is reasonable to take a constant sample size of 2 (one pair in the risk set) in the partial likelihood. This effectively introduces the assumption that each pair (i, j) is independent of any other pair. However, this assumption seems puzzling given the inherent transitivity of ranking (if i > j and j > k then i > k ).

2.2 Listwise Ranking and Differentiable Sorting Networks

[Cao et al., 2007] first proposed the notion of a listwise ranking loss, arguing its benefits in situations where the entire order of items is of importance. This approach treats the ranking problem as a permutation problem, with the aim of learning a model that can provide the optimal permutation of an entire list of items, rather than considering independent pairs. Indeed, this idea is closely related to the partial likelihood from Cox's original work. As pointed out by Wang and Yang [2022], the Top-1 probability method originally proposed by Cao et al. [2007] and extended to ListMLE by Xia et al. [2008] takes the same form as CPL.

More recent works closely related to listwise ranking have begun to explore combining traditional sorting algorithms with differentiable sorting functions [Petersen et al., 2021]. Sorting networks are a family of sorting algorithms that consist of two basic components: wires and conditional swaps [Batcher, 1968, Knuth, 1998]. Wires carry values to be compared at conditional swaps. If one value is bigger than the other then the values carried forward are swapped around. This allows construction of sorting networks that can provably sort a list of values. Conditional swaps are exactly the min and max operators that ensure that with inputs { a, b } and outputs a ∗ ≤ b ∗ , a ∗ = min( a, b ) and b ∗ = max( a, b ) . Examples of odd even and bitonic networks are shown in Appendix C.

In order to train models based on ordering information alone, differences between predicted and true orderings must be backpropagated through sorting algorithms. However, they often require the use of non-differentiable max and min operators. These are analogous to the non-differentiable indicator function that was discussed earlier in the c-index equation 4. Differentiable sorting methods, similar to ranking losses, rely on approximating these operators with smooth alternatives [Grover et al., 2019].

Petersen et al. [2021] propose combining traditional sorting networks and differentiable sorting functions. Consider that when a asymptotically approaches b , the transition point where it surpasses b is non-continuous and therefore non-differentiable. Just as previously shown in the ranking loss, such operations can be made differentiable using the logistic relaxation

If an inverse temperature parameter β > 0 is introduced such that σ : x → 1 1+ e -βx , then as β →∞ the functions tend to the exact min and max functions. Other relaxations of the step function can also be considered, Petersen et al. [2022a] show that the Cauchy distribution preserves monotonicity which is desirable for optimization. Given this, we use the Cauchy distribution as our relaxation for all experiments, where σ : x → 1 π arctan ( βx ) + 1 2 .

For an input list to be ordered, each layer of the sorting network can be considered an independent permutation matrix P l with elements given by

107

108

109

110

111

112

113

114

115

116

117

118

119

120

121

122

123

124

125

126

127

128

129

130

131

132

133

134

135

136

137

138

139

140

141

142

143

144

145

146

147

148

where a signifies intermediate values being compared. The first layer is input with z i = h θ ( x i ) , 149 each vector of covariates or images being processed independently by the same neural network. The 150

indices being compared at each layer are determined by the sorting network and the final predicted 151 probability matrix is the product of each layer of sorting operations, 152

Where P , is the final doubly-stochastic permutation matrix, doubly-stochastic meaning that the rows and columns both sum to 1. It is possible to interpret each element P ij of the predicted permutation matrix as the predicted probability of permuting from a randomly assigned rank i to a true rank j . Finally, we can define a loss by minimizing the cross-entropy between the ground truth orders represented by true permutation matrix Q and predicted permutation matrix P as

where P c and Q c denote the c-th columns of their respective matrices.

2.3 Differentiable Sorting Networks Relation to Ranking and Partial Likelihood

It is possible to directly relate differentiable sorting networks with ranking losses and partial likelihood. Expanding out the cross entropy loss, we find

where q ic = 1 only when i is the true rank otherwise 0 . Each p ic is always a function of the difference 162 in pairs of inputs x i and x j . This is complicated by the products of intermediate values a introduced 163 by the sorting network but denoted as 164

where P l to denotes the set of comparisons to be made at each layer of the sorting network. With n = 2 and β = 1 , a sorting network only requires a single relaxed conditional swap and the loss returns to the same recognisable log-sigmoid ranking loss in Equation 5, and Cox negative log partial likelihood in Equation 3.

3 Methods

3.1 Diffsurv: Handling Censoring with Possible Permutation Matrices

For risk sets of size 2, given proper case-control sampling, it will always be possible to define a single ground truth permutation matrix Q . However, when venturing to higher risk set sizes, differentiable sorting methods can no longer handle censoring since there is not a single ground truth permutation matrix Q . We cannot determine the exact rank of patients who are censored before another who experienced an event. It is only possible to know the range of possible ranks to which a patient should belong. In Figure 2, we provide an illustration demonstrating the possible ranks for a number of censored and uncensored events.

Though we no longer have access to a single permutation matrix, we may instead consider the set of all possible permutation matrices, Q = { Q 1 , Q 1 , . . . , Q κ } . In the best case, all values are uncensored and |Q| = 1 and in the worse case, when all patients are censored |Q| = n ! . Our primary contribution is to extend differentiable sorting methods to censored ranks by discriminating between possible and impossible permutations.

We introduce a more computationally tractable representation of Q by defining the possible permutation matrix , Q p , which is the element-wise maximum of every permutation in Q ,

For survival analysis, it is possible to determine Q p in linear time given a sorted list of event times t i and event indicators δ i . We will consider higher ranks to correspond with a smaller time-to-event.

Let us consider two scenarios:

153

154

155

156

157

158

159

160

161

165

166

167

168

169

170

171

172

173

174

175

176

177

178

179

180

181

182

183

184

185

186

187

188

189

190

191

192

193

194

195

196

197

198

199

200

201

202

203

204

205

207

c

1

e

2

c

3

c

4

e

5

c

6

c

7

c

1

e

2

c

3

c

4

e

5

c

6

c

7

a)

b)

time

possible permuted ranks

For a right-censored individual i (i.e., δ i = 0 ), the possible ranks must be lower than the ranks of preceding uncensored events. We can express the set of possible ranks R i as:

This set includes all ranks r that are less than the rank of any uncensored patient j with an event time t j preceding the censoring time t i of the individual i .

For an uncensored individual i (i.e., δ i = 1 ), the possible rank must be lower than all preceding uncensored events and higher than the ranks of all subsequent events. We can define the set of possible ranks as:

This set includes all ranks r that satisfy both conditions: being less than the rank of any preceding uncensored patient j with t j < t i and greater than the rank of any subsequent patient j with t j > t i .

With these observations, it's straightforward to construct Q p . If it's feasible for patient i to permute to rank j , i. e. j ∈ R i , then Q pij = 1 , otherwise Q pij = 0 . See Figure 2 (c) for a visual representation of Q p .

Given the possible permutation matrix Q p and the predicted permutation matrix P , the vector of probabilities p of a value being ranked within the set of possible ranks can be computed. Although the ground truth probabilities are unknown, the range of possible ranks is known, and the model can be optimized to maximize the sum of the predicted probabilities of all possible ranks for each sample. Noted here as the column-sum of the element-wise product ◦ , between Q p and P .

The binary cross-entropy loss can then be easily applied 206

where y i indicates whether set of predicted ranks is possible or impossible.

Equation 15 accounts for the potential challenges of incorporating right-censored samples. The binary 208 cross-entropy remains identical whether the model predicts uniform probability for all possible ranks 209 or concentrates the probability mass on a single rank possible rank. 210

The introduction of the possible permutation matrix can be used in conjunction with any differentiable 211 sorting method that outputs a doubly-stochastic permutation matrix. This includes methods such 212

213

214

215

216

217

218

219

220

221

222

223

224

225

226

227

228

229

230

231

232

233

234

235

236

237

238

239

240

241

242

243

244

245

4 Experiments

In our experiments, we aim to assess the performance of Diffsurv and compare it against the con246 ventional Cox Partial Likelihood (CPL) methods. Initially, we focus on confirming the importance 247

as SinkhornSort from Cuturi et al. [2019]. Though, in this paper, we will restrict our focus to the discussed differentiable sorting networks. We refer to the use of differentiable sorting networks and the possible permutation matrix as Diffsurv .

3.2 Top-K risk prediction

Finally, we demonstrate how the algorithmic supervision of sorting algorithms enables the development of novel methods in survival analysis, using the example of top-k risk prediction. In practical settings, it is often not necessary to rank all samples correctly. Rather, it is essential to identify the samples with the highest risk, such as by a healthcare provider, to prioritize care and interventions.

With Diffsurv, top-k risk prediction is straightforward to implement by modifying the loss such that the negative log-likelihood of predicted top-k ranks in P is maximised for individuals with a possible permutation to any top-k rank according to Q p .

First, let's denote T k as the set of values with a possible permutation to a top-k rank, derived from the ground truth possible permutation matrix Q p :

Importantly, due to the uncertainty introduced by censoring, the set of individuals with a possible permutation to a top k rank T k can be arbitrarily large. For example, in case all individuals are censored, T k is the set of all individuals. Then, the top-k loss is described as:

This loss is minimized when the model correctly predicts a top-k rank for the indices in T k . This represents the individuals with possible permutations to the top-k highest risk ranks. Importantly, this loss function is optimized for the identification of potential top-k high-risk individuals, without considering the specific order within these top-k ranks. To establish a baseline for comparison with Diffsurv's top-k risk prediction, we also introduce two variants of the Cox Partial Likelihood method. In the first variant, we adjust the likelihood term so that the product considers only the set of patients who have a potential permutation to a top-k rank, according to the matrix of possible permutations Q p :

In the second variant, we further limit the set of patients to those who have both experienced an event and have a possible permutation to a top-k rank:

Note that the denominator term is unchanged in both variants and considers only comparable pairs and includes censored patients T j > Ti . Evaluation of top-k risk prediction is also complicated by the uncertainty due to censoring. For Diffsurv and both variants of the Cox Partial Likelihood, we can first define the set of individuals predicted to be within the top-k highest risk:

We can then define the fraction of how many of these individuals are in the set of possible top-k highest ranks T k to evaluate the top-k risk prediction performance:

248

249

250

251

252

253

254

255

256

257

258

259

260

261

262

263

264

265

266

267

268

269

270

271

272

273

274

275

276

277

278

279

280

281

282

283

284

285

286

287

288

289

h

0

=1

h

1

=16

h

2

=28

Images from SVHN

Sample from

time function

Parameterize time function

with log normal house numbers

Sample censoring

time

Replace 30% of

true times with

censoring time

λ

= 1

.

5

λ

= 1

λ

= 0

.

5

f

λ

=

λe

-

λx

∼

t

∗

0

=61.2

t

∗

1

=32.5

t

∗

2

=12.9

c

∗

∼

Uniform

(0

, t

∗

)

c

∗

0

=23.2

c

∗

1

=29.3

c

∗

2

=10.1

t

i

= min(

t

∗

i

, c

∗

i

)

t

0

=23.2

t

1

=32.5

t

2

=12.9

of taking a listwise approach and evaluating the ability of differentiable sorting networks to better capture the inherent transitivity in semi-simulated data. Subsequently, we extend our analysis to compare Diffsurv and its top-k extension across multiple publicly available real-world datasets.

Baselines : We compare Diffsurv primarily against Cox's Partial Likelihood, using the Ranked List implementation from pycox [Kvamme et al., 2019]. We include Efron and Breslow estimates of CPL from Yang et al. [2022] for survSVHN. For smaller datasets, we add non-deep learning baselines: Lifelines' Cox Regression [Davidson-Pilon, 2019] and sksurv's Random Survival Forests [Pölsterl, 2020]. We do not compare with DeepHit [Lee et al., 2018] since we do not model non-proportional hazards. For an extended discussion, see Appendix A.

Network Architectures : For both CPL and Diffsurv, we use a fixed neural network architecture depending on the dataset. Small datasets utilize a single-layer neural network, survSVHN uses a ConvNet architecture as in Petersen et al. [2021], and MIMIC IV CXR uses EfficientNet-B0 [Tan and Le, 2020].

Training and Evaluation : We employ AdamW for optimization. Validation approach varies: for smaller datasets, we apply nested 5-fold cross-validation, while for imaging datasets we use train:val:test splits. We performed hyperparameter tuning for learning rate, weight decay, batch size, and risk set size. In the case of imaging datasets, we maintained fixed values for learning rate and weight decay. As in Petersen et al. [2021], we determine steepness as a function of the risk set size n , β = 2 n for odd-even and β = (log 2 n )(1 + log 2 n ) for bitonic. The type of sorting network can either be bitonic or odd-even and is determined during hyperparameter tuning. Further details on the experimental setup, including compute time, are provided in Appendix B and at https://github.com/andre-vauvelle/diffsurv .

Semi-synthetic survSVHN : Based on the Street View House Numbers (SVHN) dataset [Netzer et al., 2011], we simulate survival times akin to survMNIST Pölsterl [2019]. The increased complexity of SVHN over MNIST offers a testbed which is better able to discern the performance differences between methods. Each house number parameterizes a beta-exponential time function for survival times. Risk parameters or hazards λ i are calculated as the logarithm of house numbers, standardized and scaled for a mean survival time of 30. We introduce censoring by randomly selecting 30% of house numbers and replacing true times with values sampled uniformly between (0 , t i ] (See Figure 3).

We can examine the implications of inherent transitivity within the data. Instead of parameterizing a time function based on unique hazards derived from house numbers, we group λ i into distinct hazard quantiles. Each quantile encompasses a set of house numbers associated with a similar hazard level. We then calculate the transitivity ratio, defined as # of transitive triplets # of triplets , where a sampled triplet is considered transitive if ( λ i > λ j > λ k ).

This methodology provides us with a means to control the degree of transitivity in our data. At one extreme, we might categorize data into only two groups, representing the lower and upper halves of house numbers, which results in a transitivity ratio of 0. At the other extreme, each house number could constitute its own unique category (indicated by ∞ ), leading to high transitivity.

Our results, summarized in Table 1, align with the expectations laid out in Section 2.3: both Diffsurv and CPL methods perform similarly when the risk set size is at its minimum ( n = 2 ). However, with the expansion of the risk set size, the performance of the two methods diverges, with Diffsurv consistently outperforming CPL. Table 2 sheds light on a potential reason for this divergence. As

290

291

292

293

294

295

296

297

298

299

300

301

302

303

304

305

306

307

308

309

310

311

312

313

314

315

the number of quantiles is increased, thereby enhancing the degree of transitivity within the data, Diffsurv-based methods start to surpass CPL methods. This finding underscores the role of transitivity in survival data and validates Diffsurv's effectiveness in encapsulating this inherent property. Despite the strong performance of Diffsurv, the C-index for the ground truth risks is 0.980, which is still far above 0.943 for Diffsurv, highlighting the challenging nature of the survSVHN dataset.

Real-world datasets : We assess our methods on several public datasets: Four small, popular realworld survival datasets (FLCHAIN, NWTCO, SUPPORT, METABRIC) [Kvamme et al., 2019] and the MIMIC IV Chest X-Ray dataset (CXR) with death as the event [Johnson et al., 2019]. Further details in Appendix B.1.

The results presented in Table 3 demonstrate that Diffsurv achieves equal to or better performance on all datasets analyzed. Additionally, when Diffsurv is optimized for predicting the top 10% of highest-risk individuals, it matches or outperforms Cox's partial likelihood on the real-world datasets.

5 Conclusion

Diffsurv introduces a new perspective in survival analysis with censored data, highlighting the relations between survival analysis and the listwise ranking. Our experiments show the effectiveness of differentiable sorting methods for improving survival analysis predictions. Notably, Diffsurv matches or surpasses the performance of the established CPL methods across all examined datasets.

Crucially, Diffsurv sheds light on the importance of transitivity in ranking and survival data, revealing that methods sensitive to this inherent property, such as Diffsurv, show improved performance over those that are not. This insight underscores the value of a listwise approach in dealing with survival data and encourages further exploration for methods that promote a transitive inductive bias.

Moreover, Diffsurv provides a foundation for the development of innovative methods, including the top-k risk stratification method introduced in this work. Beyond survival analysis, the introduction of the possible permutations carries potential for other tasks that involve ranking based on limited order information. The utilization of specialized sorting networks, such as splitter selection networks as in Petersen et al. [2022b], could further leverage partial order information.

Though promising, this work is not without limitations. Future research could focus on extending 316 its applicability to non-proportional hazards and understanding the impact of ties. Moreover, in317 vestigating how well it can recover survival functions using approaches like Breslow's estimator 318 and evaluating with Brier scores would provide valuable insights into its potential and limitations. 319

Furthermore, it is important to note that Diffsurv is a survival ranking method and thus can not be 320 used to directly estimate the expected duration until an event occurs. 321

322

323

324

325

Overall, Diffsurv constitutes a meaningful advancement in survival analysis, showcasing its significant potential for enhancing risk prediction in real-world use cases. It not only demonstrates promising performance improvements, but also introduces new directions for future research, thereby making a valuable contribution to the field.

326

References

Laura Antolini, Patrizia Boracchi, and Elia Biganzoli. A time-dependent discrimination index for survival data. Statistics in Medicine , 24(24):3927-3944, 2005. ISSN 1097-0258. doi: 10.1002/sim. 2427. URL https://onlinelibrary.wiley.com/doi/abs/10.1002/sim.2427 . _eprint: https://onlinelibrary.wiley.com/doi/pdf/10.1002/sim.2427.

E. Batcher. Sorting networks and their applications. In Proceedings of the April 30-May 2, 1968, spring joint computer conference on - AFIPS '68 (Spring) , page 307, Atlantic City, New Jersey, 1968. ACM Press. doi: 10.1145/1468075.1468121. URL http://portal.acm.org/citation. cfm?doid=1468075.1468121 .

Mathieu Blondel, Olivier Teboul, Quentin Berthet, and Josip Djolonga. Fast Differentiable Sorting and Ranking, June 2020. URL http://arxiv.org/abs/2002.08871 . arXiv:2002.08871 [cs, stat].

Zhe Cao, Tao Qin, Tie-Yan Liu, Ming-Feng Tsai, and Hang Li. Learning to rank: from pairwise approach to listwise approach. In Proceedings of the 24th international conference on Machine learning , pages 129-136, Corvalis Oregon USA, June 2007. ACM. ISBN 978-1-59593-793-3. doi: 10.1145/1273496.1273513. URL https://dl.acm.org/doi/10.1145/1273496.1273513 .

R. Cox. Regression Models and Life-Tables. Journal of the Royal Statistical Society. Series B (Methodological) , 34(2):187-220, 1972. ISSN 0035-9246. URL https://www.jstor.org/ stable/2985181 . Publisher: [Royal Statistical Society, Wiley].

Marco Cuturi, Olivier Teboul, and Jean-Philippe Vert. Differentiable Ranking and Sorting using Optimal Transport. In Advances in Neural Information Processing Systems , volume 32. Curran Associates, Inc., 2019. URL https://papers.nips.cc/paper_files/paper/2019/hash/ d8c24ca8f23c562a5600876ca2a550ce-Abstract.html .

Cameron Davidson-Pilon. lifelines: survival analysis in Python. Journal of Open Source Software , 4(40):1317, August 2019. ISSN 2475-9066. doi: 10.21105/joss.01317. URL https://joss. theoj.org/papers/10.21105/joss.01317 .

Larry Goldstein and Bryan Langholz. Asymptotic Theory for Nested Case-Control Sampling in the Cox Regression Model. The Annals of Statistics , 20(4):1903-1928, December 1992. ISSN 0090-5364, 2168-8966. doi: 10.1214/aos/1176348895. URL https://projecteuclid.org/journals/annals-of-statistics/volume-20/ issue-4/Asymptotic-Theory-for-Nested-Case-Control-Sampling-in-the-Cox/ 10.1214/aos/1176348895.full . Publisher: Institute of Mathematical Statistics.

Aditya Grover, Eric Wang, Aaron Zweig, and Stefano Ermon. Stochastic Optimization of Sorting Networks via Continuous Relaxations. arXiv:1903.08850 [cs, stat] , April 2019. URL http: //arxiv.org/abs/1903.08850 . arXiv: 1903.08850.

Frank E. Harrell, Jr, Robert M. Califf, David B. Pryor, Kerry L. Lee, and Robert A. Rosati. Evaluating the Yield of Medical Tests. JAMA , 247(18):2543-2546, May 1982. ISSN 00987484. doi: 10.1001/jama.1982.03320430047030. URL https://doi.org/10.1001/jama. 1982.03320430047030 .

Alistair E. W. Johnson, Tom J. Pollard, Seth J. Berkowitz, Nathaniel R. Greenbaum, Matthew P. Lungren, Chih-ying Deng, Roger G. Mark, and Steven Horng. MIMIC-CXR, a de-identified publicly available database of chest radiographs with free-text reports. Scientific Data , 6(1): 317, December 2019. ISSN 2052-4463. doi: 10.1038/s41597-019-0322-0. URL https:// www.nature.com/articles/s41597-019-0322-0 . Number: 1 Publisher: Nature Publishing Group.

Jared L. Katzman, Uri Shaham, Alexander Cloninger, Jonathan Bates, Tingting Jiang, and Yuval Kluger. DeepSurv: personalized treatment recommender system using a Cox proportional hazards deep neural network. BMC Medical Research Methodology , 18(1):24, February 2018. ISSN 1471-2288. doi: 10.1186/s12874-018-0482-1. URL https://doi.org/10.1186/ s12874-018-0482-1 .

381

382

383

384

385

386

387

388

389

390

391

392

393

394

395

396

397

398

399

400

401

402

403

404

405

406

407

408

409

410

411

412

413

414

415

416

417

418

419

420

421

422

423

424

425

Donald E Knuth. The Art of Computer Programming. Addison Wesley Longman Publishing Co., Inc , Volume 3: (2nd Ed.) Sorting and Searching, 1998.

Håvard Kvamme and Ørnulf Borgan. The Brier Score under Administrative Censoring: Problems and Solutions. Technical Report arXiv:1g912.08581, arXiv, December 2019. URL http://arxiv. org/abs/1912.08581 . arXiv:1912.08581 [cs, stat] type: article.

Håvard Kvamme, Ørnulf Borgan, and Ida Scheel. Time-to-Event Prediction with Neural Networks and Cox Regression. arXiv:1907.00825 [cs, stat] , September 2019. URL http://arxiv.org/ abs/1907.00825 . arXiv: 1907.00825.

Changhee Lee, William Zame, Jinsung Yoon, and Mihaela Van Der Schaar. DeepHit: A Deep Learning Approach to Survival Analysis with Competing Risks. Proceedings of the AAAI Conference on Artificial Intelligence , page 8, 2018.

Yuval Netzer, Tao Wang, Adam Coates, Alessandro Bissacco, Bo Wu, and Andrew Y. Ng. Reading digits in natural images with unsupervised feature learning. In NIPS Workshop on Deep Learning and Unsupervised Feature Learning 2011 , 2011. URL http://ufldl.stanford.edu/ housenumbers/nips2011_housenumbers.pdf .

Felix Petersen, Christian Borgelt, Hilde Kuehne, and Oliver Deussen. Differentiable Sorting Networks for Scalable Sorting and Ranking Supervision. arXiv:2105.04019 [cs] , July 2021. URL http: //arxiv.org/abs/2105.04019 . arXiv: 2105.04019.

Felix Petersen, Christian Borgelt, Hilde Kuehne, and Oliver Deussen. Monotonic Differentiable Sorting Networks. arXiv:2203.09630 [cs, stat] , March 2022a. URL http://arxiv.org/abs/ 2203.09630 . arXiv: 2203.09630.

Felix Petersen, Hilde Kuehne, Christian Borgelt, and Oliver Deussen. Differentiable Top-k Classification Learning, June 2022b. URL http://arxiv.org/abs/2206.07290 . Number: arXiv:2206.07290 arXiv:2206.07290 [cs].

Sebastian Pölsterl. scikit-survival: A library for time-to-event analysis built on top of scikit-learn. Journal of Machine Learning Research , 21(212):1-6, 2020. URL http://jmlr.org/papers/ v21/20-729.html .

Sebastian Pölsterl. Survival Analysis for Deep Learning, July 2019. URL https://k-d-w.org/ blog/2019/07/survival-analysis-for-deep-learning/ .

Vikas C Raykar, Harald Steck, Balaji Krishnapuram, Cary Dehing-oberije, and Philippe Lambin. On Ranking in Survival Analysis: Bounds on the Concordance Index. In Advances in Neural Information Processing Systems , volume 20. Curran Associates, Inc., 2007. URL https://papers. nips.cc/paper/2007/hash/33e8075e9970de0cfea955afd4644bb2-Abstract.html .

Mingxing Tan and Quoc V. Le. EfficientNet: Rethinking Model Scaling for Convolutional Neural Networks, September 2020. URL http://arxiv.org/abs/1905.11946 . arXiv:1905.11946 [cs, stat].

Bokun Wang and Tianbao Yang. Finite-Sum Coupled Compositional Stochastic Optimization: Theory and Applications, September 2022. URL http://arxiv.org/abs/2202.12396 . arXiv:2202.12396 [cs, math].

Fen Xia, Tie-Yan Liu, Jue Wang, Wensheng Zhang, and Hang Li. Listwise Approach to Learning to Rank - Theory and Algorithm. Proceedings of the 25th International Conference on Machine learning , 2008.

Xuelin Yang, Louis Abraham, Sejin Kim, Petr Smirnov, Feng Ruan, Benjamin Haibe-Kains, and Robert Tibshirani. FastCPH: Efficient Survival Analysis for Neural Networks, August 2022. URL http://arxiv.org/abs/2208.09793 . arXiv:2208.09793 [cs, stat].

426

427

428

429

430

431

432

433

434

435

436

437

438

439

440

441

442

443

444

445

446

447

448

449

450

451

452

453

454

455

456

457

458

459

460

461

462

463

464

465

466

467

468

469

470

471

472

473

A Non-proportional Hazards

Our current implementation of Diffsurv operates under the proportional hazards assumption. While this may not fully capture the intricacies of some survival analysis problems-particularly those involving non-proportional hazards-it does not necessarily limit the model's effectiveness in scenarios where the goal is to assess cumulative risk from a fixed index date or from the date of an imaging study. This aligns with the necessity of time-dependent modifications to the C-index for non-proportional models as indicated by Antolini et al. [2005].

If we are primarily interested in understanding the cumulative hazard of an event occurring rather than tracking changes in the hazard over time, the assumption of proportional hazards becomes less pivotal. As such, Diffsurv and CPL remain valuable tools for these cases.

Despite the current limitation of Diffsurv to proportional hazards, it is conceivable that an extension to accommodate non-proportional hazards could be developed, similar to adaptations made for the CPL method.

For instance, as we briefly mentioned earlier, continuous-time extensions of partial likelihood can be used to enable non-proportional hazards Kvamme et al. [2019]. Implemented by directly modeling temporal covariates as f θ = exp ( h θ ( x i , T i )) .

Another class of methods focuses on discretizing the time-to-event variable and modeling the probability mass function (PMF) of event times. For instance, the DeepHit model Lee et al. [2018] employs a neural network architecture to learn the relationships between input features and discretized time-to-event outcomes. Time discretization facilitates modeling of non-proportional hazards but introduces two significant challenges: 1) sensitivity to the choice of time intervals, which can affect the model's accuracy and interpretability, and 2) increased computational complexity, as predictions must be made for each time interval. These models can be computationally expensive, especially for deep learning-based models like DeepHit, making them less suitable for high-dimensional and large-scale datasets, such as the imaging dataset used in this study.

Several future work proposals arise from these observations. First, differentiable sorting could explore the approach of directly modeling temporal covariates, resulting in a time-parameterized predicted permutation matrix. Second, extending Diffsurv to discrete time could be achieved by parameterizing a predicted permutation matrix for each time discretization.

B Training and evaluation

B.1 Datasets and Preprocessing

As in Goldstein and Langholz [1992] and Kvamme et al. [2019], we ensure that each risk set contains a valid risk set by sampling controls for a given case. Each batch consists of a number of risk sets such that the input data has shape (batch size, risk set size, covariate shape).

We provide an additional description of each small realworld dataset:

FLCHAIN dataset: A dataset containing information on patients with monoclonal gammopathy of undetermined significance (MGUS), focusing on serum free light chain (FLC) levels to study their prognostic significance in predicting disease progression. Number of covariates: 8.

NWTCOdataset: A dataset from a series of clinical trials on the treatment and outcomes of children with Wilms' tumor, a type of kidney cancer, aiming to improve understanding of tumor biology and optimize treatment strategies. Number of covariates: 9.

SUPPORT dataset: A dataset from a multi-center study investigating the prognosis and treatment preferences of seriously ill hospitalized adults, with the goal of improving end-oflife care and informing decision-making processes. Number of covariates: 22.

METABRIC dataset: A dataset comprising genomic and clinical data on breast cancer patients, focused on uncovering novel molecular subtypes for more precise prognostication and personalized treatment strategies. Number of covariates: 9.

474

475

476

477

478

479

480

481

482

483

484

485

486

487

488

489

490

491

492

493

494

495

496

497

498

499

500

501

502

503

504

505

506

507

508

509

510

511

512

513

514

515

516

Covariate preprocessing follows [Kvamme et al., 2019], and includes standardising continuous variables and one-hot encoding categorical variables.

MIMIC IV CXR : For the survival task, we extract death events from the MIMIC IV dataset. This is done by merging the data on "subject_id", "study_id", and "dicom_id" with the patient table from MIMIC IV, and on "subject_id" with the admission table. For patients without a recorded date of death, censoring dates are be determined as 1 year after the last recorded discharge date for each patient. We exclude 29,345 images without any matches in the MIMIC IV patient table, 19,337 images taken after the latest found discharge date and 55 images taken after a recorded date of death. Time to event is calculated as the number of days from the image study date to either date of death or the censoring date. For MIMIC IV CXR, images undergo several standard transformations: a random horizontal flip and a 15-degree rotation, resizing to 230 x 230 pixels, a 224 x 224 pixel center crop, and conversion to grayscale with three output channels. The data is then transformed into tensors and normalized using ImageNet's mean and standard deviation values. Finally, the train:val:test split of 8:1:1 is done at the patient level ensuring no images from a patient in the test set was found in the training data.

survSVHN : In this semi-synthetic dataset, we sample survival times using the beta-exponential distribution. The beta distribution used to sample from the exponential uses a fixed value of 500 for both shape parameters. We follow Petersen et al. [2021] by cropping the centered multi-digit numbers with a boundary of 30%, resizing it to a resolution of 64×64, and then selecting 54 × 54 pixels at a random location. For survSVHN the train:val:test split is provided by Netzer et al. [2011] as is 230,755:5,000:13,068.

B.2 Model Architecture and Hyperparameters

For the smaller real-world datasets, the hazard function f θ is a small fixed Multi-layer Perceptron network with 1 hidden layer and 64 hidden nodes. We also apply a fixed dropout rate of 0.1. Learning rate, weight decay, batch set size and risk set size were found using a grid search across the possible values in Table 4.

For imaging datasets, we fix learning rate and weight decay for both CPL and Diffsurv. For both survSVHN and MIMIC IV CXR, we use a fixed learning rate of 10 -4 and weight decay of 10 -5 . We also used early stopping with a patience of 20 epochs and a maximum of 100,000 training steps.

survSVHN : As per Petersen et al. [2021], the model consists of four convolutional layers (with a 5x5 kernel size and 32, 64, 128, 256 filters), each followed by ReLU and max-pooling (2x2 stride). The architecture concludes with a fully connected layer of 64 units, another ReLU, and a one-unit output layer.

MIMIC IV CXR : Here, we use EfficientNet-B0 with an added linear layer for single output. We first train the linear prediction layer alone for the initial 2,000 steps. After this, we continue training, this time including both the EfficientNet-B0 and the linear prediction layer.

For the results in Table 1, we keep a fixed batch size of 100. We also provide a comparison where the number of values is fixed in each batch in Table 9.

Note that during evaluation, the sorting network is not used since we only need to evaluate the ranks of the trained risk scores. Similarly, case-control sampling is not used. We measure the ranking performance of the models using the concordance index.

Further implementation details and the best hyperparameters for each dataset are provided at anon@ git.com .

517

518

519

520

521

522

523

524

525

526

527

528

529

530

531

532

533

534

535

536

537

538

539

540

541

542

543

544

545

546

547

548

549

550

551

552

553

554

555

B.3 Compute Requirements

Experiments on smaller real-world datasets are compact enough to facilitate effective training on a CPU, with each variant, including the CPH baselines, completing per experiment in less than 20 minutes. However, the larger imaging datasets require more significant computational power. In the most demanding case, the MIMIC IV CXR experiments, run on an 11GB NVIDIA GeForce GTX 1080 Ti, took roughly 18.5 hours per experiment. Both Diffsurv and CPH methods exhibited comparable run times; however, the Bitonic variant was the fastest, with a lead of approximately 6 minutes. All neural network baselines were implemented using PyTorch and PyTorch Lightning. Although measures were taken to reduce compute time and complexity, such as using half-precision and distributed data parallel (DDP) training strategies, the overall training times are far from optimized.

To understand the runtime differences of Diffsurv and Cox Partial Likelihood variants in isolation, we also run an additional experiment. For each method, we compute and time over 100 trials for a forward and backward pass on a NVIDIA GeForce GTX 1080 Ti using randomly generated logits. For diffsurv, this includes computing the predicted permutation matrix with differentiable sorting networks and applying the masking (Equation 15) and binary cross entropy (Equation 16). The possible permutation matrix generation is possible to precompute off GPU on the dataloader, so is not included in the timing. In Table 5, the Diffsurv methods, particularly Bitonic, consistently outperform CPL methods in compute time almost across all batch sizes and risk set sizes. As risk set size increases, CPL methods exhibit a decreasing trend in compute time, while Diffsurv's Odd-Even method experiences a notable rise, especially from risk set sizes of 32 to 128.

It is worth noting that CPL methods are currently computed over batches using a simple for loop, as the present implementations do not support batch parallel computation. However, there's potential for further optimization. For similar batch sizes and risk set sizes, we noted very similar overall convergence times with Diffsurv Bitonic variant being marginally faster than other methods. In the context of full model runs, the difference between Diffsurv and CPL in terms of training times is minimal; it is the model architecture that have a dominant effect on compute time.

C Sorting Networks

There are multiple different types of sorting networks each with varying complexity. The ability to implement networks with the divide-and-conquer paradigm allows for sorting networks that scale more efficiently. Examples for Odd-Even and Bitonic sorting networks with n = 8 are shown in Figure 4. The latter allows construction of networks with size complexity O ( nlog 2 n ) verses the O ( n 2 ) in Odd-Even networks.

It is worth emphasising these are not neural networks. They are called "networks" because they are typically represented as diagrams that show how the items are compared and swapped as they are being sorted. Differentiable sorting networks do not introduce any additional parameters that need to be updated during optimization.

D Calibration of Predicted Permutations

Model calibration in survival analysis models is essential for ensuring that the predicted probabilities of outcomes align closely with the true probabilities. An improperly calibrated model may lead

556

557

558

559

560

561

562

563

564

565

566

567

568

(a)

(b)

to incorrect risk assessments and treatment decisions, potentially resulting in suboptimal patient care and even adverse clinical consequences. We focus on the calibration of predicted individual rankings, as demonstrated in Figure 5 and Table 6. Specifically, we qualitatively illustrate in Figure 5 that for a model with a risk set size of 8, both discrete predicted ranks and ranking probabilities are accurately calibrated for the Diffsurv approach. To perform a quantitative comparison with baseline methods, we need to derive ranking probabilities for the CPL model. Based on the assumptions in Raykar et al. [2007], we assume that the probability of correct pairwise ordering for the CPL adheres to the logistic function. We thus compute permutation matrices using differential sorting networks, employing predicted partial log hazards as inputs and the logistic sigmoid function as the differentiable sorting operator. By subsequently calculating Brier scores for the rank probabilities in the predicted permutation matrices survSVHN, we analyze various combinations of batch size and risk set size. Our findings show that the Diffsurv models consistently exhibit the lowest Brier scores across all settings (refer to Table 6).

0

1

2

3

4

5

6

7

Ground truth rank

0

1

2

3

4

5

6

7

Predicted rank

Perfect calibration

0

1

2

3

4

5

6

7

Ground truth rank

0

.

10

0

.

15

0

.

20

Permutation probability

A

B

T

569

570

571

572

573

574

575

576

577

578

579

580

E Effect of top-k loss variants on predictive performance

The Top-K risk prediction variants for Diffsurv and the Cox Partial Likelihood, as introduced in Section 3.2, inherently bias the model to recognize only the top-k individuals with the highest risk. In this section, we explore the extent to which these specific variants influence the predictive performance of the models across the entire set of individuals. These results are shown in Table 7.

F Additional Results

In Table 8 and Table 9, additional results for the MIMIC IV CXR and survSVHN imaging datasets are provided. Here, we maintain a constant total number of samples in each batch, which means that an increase in risk set size is compensated by a decrease in batch size. These results offer further understanding of the balance required between these two variables. We observed, while larger risk set sizes generally improve performance for both Diffsurv and CPH, the benefits tend to taper off as training can become more unstable and noisy with smaller batch sizes.

Peer reviews:
{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper concerns a novel way to solve risk stratification for time-to-event data in the presence of right-censoring. Particularly, the authors solved the proportional hazard model via the differentiable sorting algorithm. Experiments were done on simulated and real-world benchmark datasets.\n\nSTRENGTHS:\n- The paper is well-written and easy to follow.\n- The authors extended the differentiable sorting network to account for right-censoring by introducing a possible permutation matrix and applied it to survival analysis.\n- This paper investigated the role of transitivity in the proportional hazards survival model.\n\nWEAKNESSES:\n- Neither the experimental nor theoretical advantage of Diffsurv over CPL is clear enough. I do not find a specific reason to use Diffsurv rather than CPL for training a PH survival model. \n\n- Even if the authors showed some advantages of the proposed Diffsurv on the simulated dataset, it hardly shows a difference against CPL in real-world benchmark datasets. \n\n- I believe Diffsurv is more than just one interesting way of solving the proportional hazards model. I think the differentiable sorting is expected to do better than the partial likelihood approach; however, the experimental results do not strongly support the advantage. The potential problem is the proportional hazard assumption - we observe marginal improvement as the PH assumption is too strong. Applying the differentiable sorting to models with no PH assumptions, such as CoxTime or DeepHit (as discussed in the appendix), may result in clear improvement (moreover, the experimental results also suggest that the method should be extended to non-PH because, under extremely simple simulated data that follows Weibull distribution with mild independent censoring, Diffsurv outperforms CPL but in real-world cases, the results hardly differ from those of the CPL’s suggesting that real-world applicability is questionable.).\n\nQUESTIONS:\n- I am not so sure about the clinical context of the top-$k$ risk stratification task. Could the authors kindly explain why the top-$k$ risk stratification is important enough to be evaluated separately, even if C-index accounts for the risk stratification?\n\n- I do not think Kvamme et al. explicitly assumed pairwise independence, but they presented that it is sufficient to consider a risk set of size 1, which means having a risk set of size 1 is as good as having a larger risk set. Moreover, due to the nature of the random selection of SGD, the transitivity is implicitly considered. Also, in the paper, they investigated the impact of risk set size, which supports their claim “it is often sufficient to choose $n=1$.”\n\n- I am curious how emphasizing top-K prediction changes the discrimination and calibration performance of survival models. \n\n- Could authors kindly provide a computational cost analysis (either experimental or theoretical) of Diffsurv? It will be informative if the authors compare the computational costs of Diffsurv and CPL.\n\n- RSF is not based on PH assumption. If the authors did not include DeepHit or other SOTA deep survival models just because they are not based on PH assumption, why is RSF given as a baseline?\n\nLIMITATIONS:\nThe authors did not provide any theoretical analysis of Diffsurv. CPL might seem to be naive as it does not exhaustively consider all possible pairs, however, the estimator that maximizes the partial likelihood has nice statistical properties, including consistency and asymptotic normality. It would be nice if the authors could provide statistical/information theory behind Diffsurv and compare it with the existing ones.", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 5, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-07-05"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe authors propose a deep learning-based survival model that utilizes a novel differentiable sorting objective capable of handling censored time-to-event data. In contrast to previous pairwise sorting objectives, the proposed listwise sorting objective achieves the transitive property inherent in survival data. Through semi-synthetic and real-world experiments, the authors demonstrate that the proposed method achieves comparable or improved risk ranking performance by leveraging the advantages of the proposed differentiable sorting objective.\n\nSTRENGTHS:\nExtending a differentiable sorting method to account for censoring and applying it to train survival models is novel.\n\nWEAKNESSES:\n-\tThe writing of the paper should be improved: it has many typos and grammatical errors. Some of the examples are listed below (note that there are many others not listed here):  Line 61: “extension differentiable” -> “extension of differentiable” / Line 97: “n=2” should be in mathematical form / Line 154: “It possible” -> “It is possible” / Line 161: “loss out we find” -> “loss, we find” / Line 77: “1-dimensional vector of size d” ->  “d-dimensional vector” / Make notations for referring to “Equation” consistent throughout the paper. \n\n-\tThe motivation for using the “listwise” sorting method that achieves transitive property over the “pairwise” sorting method is not clearly stated throughout the paper (especially in the Introduction). It would be helpful to state the formal definition of transitive property in survival data with mathematical notations, and the limitations when this property is not maintained.  \n\n-\tThe design of semi-synthetic experiments does not elaborate why listwise sorting should outperform pairwise sorting. \n\n-\tThe experiments are performed only with limited benchmarks (mainly the variants of partial likelihood methods and the proposed method itself) and evaluated only in terms of the discriminative power. The authors should provide predictive power (such as Brier scores) of survival models and in-depth qualitative analysis on the benefit of using the proposed listwise sorting method (when it has pros and when not).\n\n-\tIt would be helpful if mathematical definitions are used for describing two scenarios of ranking right-censored samples in lines 187-191.\n\nQUESTIONS:\n-\tRegarding Weakness 2, what is the limitation of pairwise sorting over listwise sorting? If correct pairwise sorting is achieved for all acceptable pairs, doesn’t it necessarily achieve transitive property?\n\n-\tRegarding Weakness 3, why listwise sorting over pairwise sorting is important in the semi-synthetic survSVHN?\n\n-\tAcceptable pairs do not consider cases in which the risk of a pair of samples cannot be directly compared. However, the proposed permutation ranking takes into account the possibility that right-censored samples may have higher ranks. The question arises: what is the benefit of including such potential cases for which we cannot guarantee a direct comparison? Furthermore, wouldn't it be harmful if a right-censored sample is highly likely to experience an early time-to-event, which can be inferred from the covariates? Can this be supported qualitatively based on the semi-synthetic (or synthetic) experiments?\n\n-\tIn relation to Question 3, the semi-synthetic data generation process does not account for independent censoring, which assumes that the time-to-event and time-to-censoring are conditionally independent given the covariates. This assumption is commonly utilized in various survival analysis literature. Wouldn't the proposed semi-synthetic scenario be more favorable for utilizing the proposed permutation ranking method for sorting, rather than relying on acceptable pairs for sorting?\n\nLIMITATIONS:\n-\tIt seems that computing the Q matrix is computationally burdensome as it needs to compare all the sample pairs. \n\n- Please see details in Weakness 2.\n\n- Please see details in Weakness 4.", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-07-05"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper investigates the problem of survival analysis under the proportional hazards (PH) hypothesis with deep learning models.\n\nThe main issue with previous approaches (Katzman et al. 2018 and Kvamme et al. 2019 as main representatives) is that the Cox loss is only computed in small mini-batches, potentially of size 2 in the case of the last reference. As a consequence, only pairwise ordering is imposed, potentially breaking transitivity.\n\nThis paper proposes to leverage recent advances in differentiable sorting to take into account transitivity in training deep survival models under the PH hypothesis. The main issue comes from censoring, which only implies a partial ordering. The main contribution of the authors is to propose a method that takes into account this censoring, defining a possible permutation matrix. The resulting method is flexible and can be adapted to plain survival or other tasks, such as top-K risk prediction.\n\nExperiments on real data show that:\n1. the method improves over previous deep learning methods (table 1), showing a +1 c-index point\n2. the gap between both methods increases when the transitivity of the dataset increases (table 2)\n3. this gap is maintained for other datasets (Table 3, first part) and other tasks (table 3, last part)\n\nSTRENGTHS:\n- Potential large impact: this paper provides a more sound way to deal with survival data and minibatches in deep learning, which could have a large impact\n- Quality: the experiments are well conducted and yield convincing results\n- Originality: to the best of my knowledge, this is one of the first papers leveraging differentiable sorting for survival analysis\n- Clarity: the paper is well written, up to minor issues (see weaknesses)\n\nWEAKNESSES:\n- Clarity\n  - notations in Equations (2), (3) and (4) are fuzzy. How is $j \\in \\mathcal{R} \\ \\lbrace i \\rbrace$ chosen? It is unclear whether the sum applies to all such $j$ (as should be the case in eq. (4), since the c-index is computed over all acceptable pairs as defined L 106) or to only a random sample (as in Eq. (3)).\n\n- Top-k task\n\n  - As pointed out by the authors, due to the partial ordering the set $\\mathcal{T}_k$ can grow very large (L217). As a consequence, the top-k-score defined in Eq (20) can be arbitrarily good. How do the authors control for this potential effect in the results?\n\nQUESTIONS:\n- cf top-k task: how do the authors control for the effect of the size of $\\mathcal{T}_k$ on the top-k-score?\n- Is there a difference in run-time between previous approaches and the proposed differentiable sorting?\n\nLIMITATIONS:\n- This paper is limited to proportional hazards assumptions, but it is clearly acknowledged by the authors", "TITLE": "Official Review", "RECOMMENDATION": 7, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 4, "DATE": "2023-07-07"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe paper aims to address the survival problem as a ranking problem. The author asserts that by taking into account the transitivity property of ranking, sorting networks can be employed to tackle the current ranking problem effectively. Furthermore, the paper presents a solution to handle numerous potential permutation matrices in the presence of censoring. Additionally, it claims to outperform baseline methods in terms of top-k risk prediction. They tested their method on a semi-synthetic dataset based on the SVHN dataset and 5 real-world datasets.\n\nSTRENGTHS:\nBy utilizing sorting networks, the paper leverages the inherent transitivity of ranking to its advantage. Additionally, it offers a comprehensive solution to handle all conceivable permutation matrices.\n\nWEAKNESSES:\n1. Survival analysis encompasses more than just generating rankings. In many cases, it is crucial to estimate the timing of the occurrence of the event of interest. The paper could address this aspect by comparing the learned hazard with the ground truth hazard in the semisynthetic experiment.\n2. when it comes to top-k risk prediction, the method needs to provide a meaningful ordering, which, unfortunately, is not demonstrated in the paper.\n3. It is not immediately apparent why this particular method is necessary. Is its sole purpose to accommodate the limitation of fitting the entire at-risk set into memory?\n4. It would be beneficial for the author to provide a more detailed explanation of sorting networks and their functioning. The working principles of these networks are currently unclear to me. Furthermore, it would be helpful to know if sorting networks have any limitations regarding the size of the input list they can handle.\n5. Improvements are not significant enough, even though baseline is a very simple method.\n\nQUESTIONS:\n1. What is the relationship between batch size and risk set size? How is it possible for the risk set size to be larger than the batch size?\n2. I believe that as the risk set size in the Cox model increases, the performance of methods should be closer, but this does not seem to be the case (based on lines 111 and 112). Could you explain this to me?\n3. Cox proportional hazards regression (CPL) forms the basis of Cox regression, and the hazard function in both methods should be the same. Consequently, the c-index should be the same for both methods. Why is it different?\n4. According to the paper, the transitivity ratio should be between 0 and 1. However, it is reported as infinity. Why is that?\n\nLIMITATIONS:\n1. How does this method handle issues related to time-varying hazards?\n2. As the method only provides risk scores, it is unclear how it can be assessed using metrics like IBS and Calibration.", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 2, "DATE": "2023-07-08"}