Differentiable sorting for censored time-to-event data.
Authors: Andre Vauvelle
Benjamin Wild
Roland Eils
Spiros Denaxas
Venue: NeurIPS 2023 poster (2023)
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
TL;DR: Diffsurv introduces a novel extension of differentiable sorting methods to survival analysis, effectively handling censored data.

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.

Differentiable sorting for censored time-to-event data.
Andre Vauvelle Benjamin Wild"; Roland Eils', Spiros Denaxas' University College London', BenevolentAl: Berlin Institute of Health fandre. vauvelle. 19, denaxas]@ucl. ac uk (benjamin. wild, roland e eils) (bbh))bbihcchartt.....
Abstract
Survival analysis is a crucial semi-supervised task in machine learning with signifi- cant real-world applications, especially in healthcare. The most common approach to survival amalysis, Cox's partial likelihood, can be interpreted as ranking model optimized on a lower bound of the concordance index. We follow these connec- tions further, with listwise ranking losses that allow for relaxation of the pairwise independence assumption. Given the inherent transitivity of ranking, we explore differentiable sorting networks means to introduce stronger transitive in- ductive bias during optimization. Despite their potential, current differentiable sorting methods cannot account for censoring, crucial aspect of many real-world datasets. We propose a novel method, Diffsurv, to overcome this limitation by extending differentiable sorting methods handle censored tasks. Diffsurv pre- dicts 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 novel method for top-k risk prediction that surpasses current methods. In conclu- sion, Diffsurv not only provides a novel framework for survival analysis through differentiable sorting, but also significantly impacts real-world applications by im- proving risk stratification and offering a methodological foundation for developing predictive models in healthcare and beyond.
2
6
10
12
13
14
15
17
19
Introduction
21
22
Survival analysis plays pivotal role in many realworld machine learning applications, spanning fields such as reliability engineering, marketing, and insurance, with particularly significant impact healthcare. The goal survival analysis is to predict the time until the occurrence of an event of interest, such as death, based 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.
24
26
The concept of censoring distinguishing characteristic that sets survival analysis apart from 29 conventional machine learning approaches. Particularly prevalent in observational datasets, it refers
28
30 to situations where event times remain unobserved because patient might not have undergone the
31 event by the time of data collection. This can be due to variety of reasons such the study period
32 ending before events of interest have occurred or subjects leaving the study.
33 Overlooking censoring can skew predictions towards the censoring event, rather than the event of
34 interest. This bias becomes particularly noticeable when the study's endpoint can be inferred from the observed covariates age being notable example. In such cases, the predicted event times are likely
35
Equal contribution
37th Conference on Neural Information Processing Systems (NcurIPS 2023).

---

36
to biased towards the censoring event time, thereby neglecting the actual event of interest K vamme and Borgan 2019].
37
38
The Cox Proportional Hazards model widely used for handling censored data in survival analysis Cox [972]. The model optimizes a partial likelihood function over ranked data, considering only
39
40
the order of events, not their exact time of occurrence. As such, Cox's 's partial likelihood serves as ranking loss, learning from the order of patients based on their hazard of experiencing an event, not their exact survival time.
41
42
Raykar et al.) [007] showed that Cox's partial likelihood (CPL) and ranking losses be directly equated, with both providing lower bounds to the concordance index, the primary evaluation metric used in survival analysis. Both losses are foundational to many survival deep learning methodologies like DeepSurv Katzman et al. 2018) and DeepHit Siiiiteees Lee et al. 2018).
44
46
47
However, this relation operates under the assumption of pairwise independence. This simplification, while practical, can de-emphasize the transitive properties inherent survival data. As shown by Goldstein and Langholz [192], larger risk set sizes can lead to more efficient estimators, suggesting potential benefits in considering listwise ranking losses Cao et al. [2077]. These losses optimize over lists of values rather than individual pairs, thereby better capturing the transitive dynamics of the data.
48
49
50
51
52
Despite the similarities, listwise losses have remained largely unexplored within the field of survival analysis. This could be partly due to uncertainty around how to handle censoring.
54
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 al., 2019, Blondel et al. 2020, PPetersnn et al., 202T]. 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.
55
56
57
58
60
Briefly, our contributions are summarised:
61
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
63
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
64
65
66
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
67
68
89
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
70
71
2
Survival Analysis and its Relation to Ranking
72
73 A dataset with censored event times is summarized as D = (4,, 2; 6, 6, where N is the total
74 number of patients. For a patient i, the time-to-event f; is the minimum of the true survival time
75 tr and the censoring time ej. with ò; indicating whether an event (t² < c 8, = 1) or censoring
78 (f) > ct, d; 0) was observed. Covariates are x; E RRd representing a d-dimensional vector but the
77 methods discussed here also generalise to higher dimensional tensors such as image data.
78 The widely-used method for addressing censoring in survival analysis is the Cox Partial Likelihood (CPL) model, introduced by Cox [972]. The CPL is designed to maximize the following general
80 form:
2(0)= = II fo(a:) 41,=11 Ejn fo(
(1)

---

21
Diffsurv
P
Σラ=1 (Qp P)as Survival
23
CNN
24
10 10
TopK
25
26
Q2
L= Ei Vi log(p))
PrByPePPPPoPoPPiE
Figure 1:
Differentiable Sorting for Censored Time-to-Event Data. Inputs, in this case SVHN images, transformed into scalar values through neural network. A differentiable permutation matrix, P. is computed using sorting networks. The model can be optimized for downstream tasks, such as risk stratification and top-k highest risk prediction, by using the matrix Qp of possible permutations based on the observed events and censoring.
81 where fo is the hazard function, a score prediction function estimating the probability of an event at a particular time, given input features æ;, The product only includes uncensored patients, whereas the denominator term also includes censored patients with > ti.
83
Reflecting the structure of survival data, the Cox Partial Likelihood (CPL) model compares individuals still "at risk" 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 event- experiencing individual, and the denominator summing over all individuals still at risk.
85
87
Extensions of the Cox model, like Katzman et al. 2018 and [vvame al., 2019), have modified 89
fo exp(@ az;) to relax the linear covariate interaction and proportional hazards assumptions. Introducing neural networks hg to handle the non- non-linearity, adjusting fo to be fo = exp(ho(x:)), and to manage non-proponee non-proportional hazards, they set fe exp(hg (I,,
91
92
2.1
Pairwise Independence
:ty 4 intractable for deep learning 94 applications as it considers all comparable patients. To mitigate memory constraints, we can sample < N. Kvamme et al. [2019] go further, arguing 96 is reasonable to take a constant sample size of and include the individual in the risk set (such
93 Both of these previous works note that the risk
95 a fixed-size risk set, denoted as R, such that [R|
97 that n = 2). This leads to the simplified loss of the form
fo(x:) fe(x.) + fo(
((
II 26
where J(i) R {i}.
(2)
98 Further, take the mean log partial likelihood to be
1 Σ
loss
log(1 + exp(h+(2 ha(xx))), where J(i) E (i),
(3)
Rh. 2:0,=1
where n. is the number of non-censored events. In this simplified form, it can be that the partial likelihood only considers the pairwise relative ordering or ranking of survival times.
100
The concordance index or c-index Harrell et al. [1882 is a commonly used as an evaluation for survival analysis methods and a generalization the Area Under the Receiver Operating Characteristic
102
Curve (AUROC) that handles right-censored data. It is defined as
c-index = Σ {())) i6
104
Raykar et al. [007] first showed that the Cox's partial likelihood approximately equivalent to maximizing the concordance index c-index and that closer bounds can be found by minimizing the

---

106
general ranking loss, with acceptable pairs A = {(i,5) : 6, = 1 1 > t and
ranking-loss :-
Σ (iJ)EA
((o((x)) - fo(æ))),
(5)
107
where Φ is function that relaxes the non-differentiable 1 of the c-index. From Equation 3 can be
108
seen that Ø 2 + - -log(1 + exp( )x) = log(o (x)). Here, we have shown that the simplifications to the partial likelihood made by Kvamme et al. 2019 are equivalent to using the log sigm-idmoo ranking
109
110
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
112
113
114
115
i and j > k then k).
116
2.2
Listwise Ranking and Differentiable Sorting Networks
117
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
118
119
120
to the partial likelihood from Cox's s original work. As pointed out by WWang and Yang [2022] the
Top-1 1 probability method originally proposed by Cao et al. 2007) and extended to ListMLE by Xia 122
et al. 2008 takes the same form as CPL.
124
More recent works closely related to listwise ranking have begun explore combining traditional sorting algorithms with differentiable sorting functions (Petersen et al. [021]. Sorting networks are 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 bigger than the other then the values carried forward are swapped around. This allows construction of sorting networks that can provably sort list of values. Conditional swaps are exactly the min and max operators that ensure that with inputs {a, b) and outputs a* S b², a² min(a, b) and
125
128
127
128
129
130
b max(u, b). Examples of odd even and bitonic networks are shown in Appendix C
132
In order 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 -in-differentiabee max and min operators. These are analogous to the -on-differentiable indicator function that was discussed earlier in the c-index equation 4 Differentiable sorting methods, similar
133
134
135
136
to ranking losses, rely on approximating these operators with smooth alternatives [Grover al.
137
2019].
Petersen et al. [021 propose combining traditional sorting networks and differentiable sorting functions. Consider that when a asymptotically approaches 6, 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
139
140
141
min, (a, b) = a (b a) + b (a bb) and max, (a, b) = a (aa - b) + b. ob - a).
(6)
142
If an inverse temperature parameter 3 > 0 is introduced such that x 1+e-xz then as / oo 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 + arctan((22) + 2.
143
144
145
146
For input list to be ordered, each layer of the sorting network can be considered an independent permutation matrix P. with elements given by
148
Pi, = Pt.j5 = (a, - a;) and Pt.s5 = Pt.ja = 1 o(a - a.),
(7)
149 where a signifies intermediate values being compared. The first layer is input with 2; (ax), each vector of covariates or images being processed independently by the same neural network. The

---

indices being compared at each layer are determined by the sorting network and the final predicted probability matrix the product of each layer of sorting operations.
152
P
II P.
(8)
Where P. is the final doubly-stochastie permutation matrix, doubly-stochastic meaning that the rows and columns both sum to 1. It is possible to interpret each element P. of the predicted permutation matrix the predicted probability of permuting from randomly assigned rank i to true rank
154
156 j. Finally, we can define a loss by minimizing the cross-entropy between the ground truth orders
157
represented by true permutation matrix Q and predicted permutation matrix P as
Σ () CrossEntropy(Pe.... (Pa) Q.).
(9)
158
where Pa and Q. denote the e-th columns of their respective matrices.
2.3
Differentiable Sorting Networks Relation to Ranking and Partial Likelihood
160
is possible to directly relate differentiable sorting networks with ranking losses and partial likelihood.
Expanding out the cross entropy loss, find
C
the log(Pic)
(10)
162 where qic = only when is the true rank otherwise O. Each Ptc. is always a function of the difference
163 pairs of inputs 2; and zy. This complicated by the products of intermediate values a introduced
184 by the sorting network but denoted as
n II (at.a))EP:dz1
Pic
a(a - ag))
(11)
where Pi to denotes the set of comparisons to be made at each layer of the sorting network. With 186 n = 2 and 3 1, a sorting network only requires single relaxed conditional swap and the loss
returns to the same recognisable log-sigmoid ranking loss in Equation 5 and Cox negative log partial
167
168
likelihood Equation
Methods
169
170
3.1
Diffsurv: Handling Censoring with Possible Permutation Matrices
For risk of size 2, given proper case-control sampling, will always be possible define single ground truth permutation matrix Q. However, when venturing to higher risk set sizes, differentiable sorting methods can longer handle censoring since there not single ground truth permutation matrix Q. We cannot determine the exact rank of patients who are censored before another who experienced an event. 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.
172
173
174
175
178
177
178
Though we no longer have access to a single permutation matrix, we may instead consider the set of all possible permutation matrices, Q (Q1, Q .......... Qx In the best case, all values uncensored and [ and in the worse case, when all patients are censored [ = n!. Our primary contribution is to extend differentiable sorting methods censored ranks by discriminating between possible and impossible permutations.
180
182
We introduce more computationally tractable representation of Q by defining the possible permuta- tion matrix, Qp, which is the element-wise maximum of every permutation in Q.
184
Q pkj max Q1aj. Qzij...., )
(12)
For survival analysis, it is possible to determine Qp in linear time given a sorted list of event times I; and event indicators 5,. We will consider higher ranks to correspond with smaller time-to-event. Let us consider two scenarios:
186
5

---

a)
time
C)
C1
62
C3
C4
65
C6
1
1
1
C1
1
0
0
0
62
b)
possible permuted ranks
1
1
1
C3
Q2
1
1
C4
1
1
0
6G
0
1
1
C6
0
1
1
Figure 2: For an example case (a) with two events (o, e₁ and et) and multiple censored samples (o. C3, C4, Co e7) the uncertainty in the possible permuted rankings (b) due to censoring is taken into account to derive the possible permutation matrix Qp (c).
188
1. For a right-censored individual i (i.e., 6, = 0), the possible ranks must be lower than the ranks of preceding uncensored events. We can express the of possible ranks R₁ as:
189
R, = [r|r<< rank(j), Vj
1,4, < t;}
(13)
This set includes all ranks r that are less than the rank of any uncensored patient j with an event time 1 preceding the censoring time & of the individual
191
192
2. For an uncensored individual (i.e., d; 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:
193
194
R3 = (r rank((), Vj: 1,4 <4) fr rank Vj: > ())
(14)
195
This set includes all ranks r that satisfy both conditions: being less than rank any preceding uncensored patient j with i 1, and greater than the rank of any subsequent patient j with t, > t,.
197
With these observations, it's straightforward to construct Qp- It's's feasible for patient to permute to rank j. i.e. j E R,, then Qpij = otherwise Qpij See Figure (c) for visual representation
199
of Q
201
Given the possible permutation matrix Qp and the predicted permutation matrix P. the vector of probabilities of value being ranked within the set of possible ranks can be computed. Although
202
the ground truth probabilities are unknown, the range of possible ranks is known, and the model can 203
be optimized to maximize the sum of the predicted probabilities of all possible ranks for each sample.
205
Noted here as the column- of the element-win wise product o, between Qp and
p =
Σ O2p P)i.j- j=1
(15)
206
The binary cross-entropy loss can then easily applied
7
Di log(p.) - (1 - y1) log(1 - pt)
(16)
207
where ys indicates whether set of predicted ranks is possible or impossible.
Equation 15 accounts for the potential challenges of incorporating right-censored samples. The binary enoss-yntropy remains identical whether the model predicts uniform probability for all possible ranks concentrates the probability
209
210
single rank possible rank.
211
The introduction the possible permutation matrix can used in conjunction with any differentiable sorting method that outputs doubly-stochastie permutation matrix. This includes methods such

---

213
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.
214
215
216
3.2
Top-K risk prediction
Finally, we demonstrate how the algorithmic supervision of sorting algorithms enables the develop- ment 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 identify the samples with the highest risk, such as by healthcare provider, to prioritize care and interventions.
218
219
220
With Diffsurr, top-k risk prediction is straightforward to implement by modifying the loss such that the negative log- log-likelihood of predicted top-k ranks in P is maximised for individuals with a possible permutation to any top-k rank according to Qp.
222
224
First, let's denote Tx as the set of values with a possible permutation to top-k rank, derived from the ground truth possible permutation matrix Qp:
225
Ts = { Σ Qpi > 0}
(17)
j=l
Importantly, due the uncertainty introduced by censoring, the set of individuals with a possible 227
permutation to a top k rank Ta can be arbitrarily large. For example, in case all individuals are censored, Tk the set of all individuals. Then, the top-k loss is described as:
228
Lopp
---oo log BET
ΣS, Pd
(18)
229
This loss is minimized when the model correctly predicts a top-k rank for the indices in Ty. This represents the individuals with possible permutations to the top-k highest risk ranks. Importantly, this loss function is optimized for the identification potential top-k high-risk individuals, without 231
considering the specific order within these top-k ranks. To establish baseline for comparison with
233
Diffsurv's top-k risk prediction, we also introduce two variants of the Cox Partial Likelihood method. In the first variant, 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
235
236
p:
LOL = II fe(x.)
(19)
Ejn,,,, So(o))
237 In the second variant, we further limit the set of patients to those who have both experienced an event and have possible permutation to a top-k rank:
238
fo(ax)
LCPL II
(20)
=AAET ESTYYT Se(
Note that the denominator term is unchanged in both variants and considers only comparable pairs and includes censored patients T; >> Tí. Evaluation of top-k risk prediction is also complicated by the uncertainty due censoring. For Diffsurv and both variants of the Cox Partial Likelihood, can first define the set of individuals predicted to be within the top-k highest risk:
240
242
P (i|rank fo (x,)) 2 k}
(21)
We can then define the fraction of how many of these individuals are in the set of possible top-k highest ranks T5 to evaluate the top-k risk prediction performance:
244
top-k-score |P& 0 Tkl
(22)
IP
4 Experiments
246 In our experiments, we aim to assess the performance of Diffsurv and compare it against the con- ventional Cox Partial Likelihood (CPL) methods. Initially, we focus confirming the importance
7

---

ho=1
1
fx = Ae
17=61.2
cj-23.2
to=23.2
Uniform00
0.5
16
11=16
:1=32.5
c1=29.3
t1=32.5
28
h2=28
12=12.9
t2=12.9
Images from SVIIN
Parameterize time function with log normal house numbers
Sample from time function
Sample censoring time
Replace tre times with censoring time
Figure 3: Visual abstract of the SurvSVHN dataset.
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.
240
251
We compare Diffsuuv primarily against Cox's Partial Likelihood, using the Ranked List implementation from pycox Kvammee et al. [0199]. We include Efron and Breslow estimates of CPL
Baselines:
252
253
from Yang al. [2022] for SurvSVHN. For smaller datasets, we add non-deep learning baselines: Lifelines' Cox Regression [Davidson-Pilee
254
2019] and sksurv's Random Survival Forests [Pölterl| We do not compare with DeepHit [Lee et al. (018)) since we do not model non- non-proportional hazards. For an extended discussion, see Appendix
255
2020].
256
257
Network Architectures: For both CPL and Diffsurv, we use a fixed neural network architecture
258
depending on the dataset. Small datasets utilize a single-layer neural network, survSVHN uses a ConvNet architecture as in Petersen al. [202T]. and MIMIC IV CXR uses EfficientNet-BO Tan
259
280
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 valttst 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 al. [2021], determine steepness as a function of the risk set size nI 3 = 2n for odd-even and 3 (logg n)(1 + +log n) for bitonic. The type of sorting network can either be bitonic or odd-even and is determined during hyperparameter tuning. Further
262
264
265
266
267
268
https://githun com/ vauvee-vauvllle www.eeevauveleee/eeee difffur
details on the experimental setup, including compute time, are provided in Appendix
B and at
269
270
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 beta-exponential time function for survival
272
274
times. Risk parameters or hazards A; are calculated as the logarithm of house numbers, standardized and scaled for 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, 4:] (See Figure 3
275
276
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 A into distinct hazard quantiles. Each quantile encompasses a set of house numbers associated with similar hazard level. We then calculate the transitivity ratio, defined as
278
of transitive triplets where sampled triplet is of triplets
281
considered transitive if (A; > >j > A&).
282
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 oo), leading high transitivity.
284
286
Our results, summarized Table I align with the expectations laid out in Section 2.3 both Diffsury 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 Diffsury
288
consistently outperforming CPL. Table 2 sheds light a potential reason for this divergence. As
8

---

Table 1: Results for training on survSVHN with increasing risk set size. Mean (standard deviation)
c-index over 5 trails with different seeds. When n = 2 both methods are equivalent to the ranking loss to up continuous relaxation of swap operation.
Risk set size
2
4
8
16
32
Diffsurv
.918 (.003)
.934 (.002)
.940 (.001)
.943 (.002)
.911 (.002)
Cox Partial Likelihood
.913 (.002)
.925 (.002)
.931 (.002)
.933 (.002)
.900 (.003)
Table 2: Results for training on survSVHN while increasing transitivity. Metric is c-index. Mean performance over 3 trails with different seeds. Bold indicates significant improvement (t-test, p s 0.01). Restricted to a fixed batch size and risk set size of 32.
Number of Quantiles Transitivity Ratio
4
16
32
64
128
DC
.0
.374
.657
.819
-088
.954
.955
.991
Diffsurv: Bitonic
.643
.803
.882
.922
.933
.999
.939
.939
Diffsurv: Odd Even
.646
.802
.883
.923
.935
.939
.940
.941
CPL: Ranked List
.651
.803
.880
.909
.966
.920
.921
.920
CPL: Efron
.647
.801
.871
.898
.904
.905
.909
.908
CPL: Breslow
.648
.801
.871
.898
.904
.909
.907
.910
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 '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.
291
293
295
We assess our methods on several public datasets: Four small, popular real- world 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 al. 2019 Further details in Appendix B.1
Real-worke datasets:
296
297
299
The results presented in Table demonntratee that Diffsurv achieves equal to better performance
300
on all datasets analyzed. Additionally, when Diffsury is optimized for predicting the top 10% of
301
highest-risk risk individuals, matches or outperforms Cox's partial likelihood on the real-world datasets.
Conclusion
302
303
Diffsurv introduces 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, Diffsury matches or surpasses the performance of the established CPL methods across all examined datasets.
304
305
306
Crucially, Diffsurv sheds light 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 not. This insight underscores the value of listwise approach in dealing with survival data and encourages further exploration for methods that promote a transitive inductive bias.
308
309
310
Moreover, Diffssurv provides 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 al. [2022b), could further leverage partial order information.
312
314
316
Though promising, this work not without limitations. Future research could focus on extending its applicability to non-prope con-proportional hazards and understanding the impact of ties. Moreover, in- vestigating how well it can recover survival functions using approaches like Breslow's estimator and evaluating with Brier scores would provide valuable insights into its potential and limitations.
318

---

Table 3: Results for real-world and symi-syn datasets. Mean (standard deviation). Survival
metric is c-index, Top 10% metric is scp-e- Bold indicates significant improvement (t-test,
p S 0.01).
FLCHAIN
NWTCO
SUPPORT
METABRIC
MIMIC CXR
SUrVSVHN
Size
6,524
4.028 85.8%
8,873
1,904
377.110
248.823
Censored Proportion
69.9%
32.0%
42.1%
60.9%
30.0%
Survival
Cox Regression
.750 (.083)
.692 (.021)
.598 (.010)
.628 (.013)
Random Survival Forest
.789 (.011)
.691 (.024)
.614 (.009)
.641 (.012)
Cox Partial Likelihood
.794 (.013)
.709 (.01))
.642 (,006)
.698 (.011)
.760 (.002)
933 (.002)
Diffsurv
(.09))
.703 (.026)
.645 (,002)
684 (.011)
.763 (.001)
943 (.002)
Top 10% prediction
Cox Partial Likelihood
(.033)
.390 (.06))
.280 (.023)
.249 (.065)
.390 (.010)
CPL-TopK (Variant I)
.699 (.007)
(.061
.479 (,016)
.527 (.083)
.408 (.008)
CPL-TopK (Variant II)
460 (.009)
.413 (.054)
479 (.035)
487 (.058)
.066 (.006)
Diffsurv
.452 (.011)
(.08))
.296 (,015)
.331 (.102)
.422 (.002)
Diffsurv-Topk
.482 (.019)
.421 (.065)
.508 (.027)
.533 (.092)
.412 (.009)
320
Furthermore, it is important to note that Diffsurv is a survival ranking method and thus can not be used directly estimate the expected duration until an event occurs.
321
322
Overall, Diffsury constitutes 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 the field.
324
325
Acknowledgements
326
327
We extend our gratitude to Felix Peterson, Leon Sixt, and Samuel Holt for their insightful discussions and invaluable feedback, which significantly contributed to the quality of this work. We also thank
328
329
the anonymous reviewers for their thoughtful feedback and fruitful discussions, which have been
330
instrumental enhancing the manuscript.
10

---

References
332
Laura Antolini, Patrizia Boracchi, and Elia Biganzoli. time-dependent discrimination index for survival data. Statistics in Medicine, 24(24)3922-39444 2005. ISSN 1097-0258. doi: 10. 1002/si.
338
2427. URL t
334
10020sim...mmmmmmm 2427 _eprint:
335
336
K. E. Batcher. Sorting networks and their applications. In Proceedings of the April 30-May 2. 1968, spring joint computer conference AFIPS '68 (Spring), page 307, Atlantic City, New Jersey. 1968. ACM Press. doi: 10. 455514448875..1881822 http:/// acm org/ citg/ittiionn
338
fmm7doid-146800 1468121
340
Mathieu Blondel, Olivier Teboul, Quentin Berthet, and Josip Djolonga. Fast Differentiable Sorting and Ranking, June 2020. URL stat].
08877 arXiv:2002.08871 [cs,
342
343
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
344
learning, pages 129-136, Corvalis Oregon USA, June 2007. ACM. ISBN 978- 978-1-59593--93-33 doi: 10. 10.1145/1273496. 1273513. UtLptttps: //d1 acm. org /ttg//oi/11... 1141/12733990 127333
346
347 D. R. Cox. Regression Models and Life-Tables Journal of the Royal Statistical Society. Series (Methodological), 34(2):17-2200, 1972 ISSN 0035-9246. stable/298511 Publisher: [Royal Statistical Society, Wiley].
348
350
Marco Cuturi, Olivier Teboul, and Jean-Philippe Vert. Differentiable Ranking and Sorting using Optimal Transport. In Advances in Newral Information Processing Systems, volume 32. Curran Associates, Inc., 2019. URL https: /tttpss:/paper.....
351
352
353
(BB22488f22662255007666666666660 oat.a//e0....aa222200ooooeaaaa.. Abstract.htm
354
Cameron Davidson-Pilon. lifelines: survival analysis in Python. Journal of Open Source Software, 4(40): August 2019. ISSN 2475-9066. doi: 10.21 11//1 URL https: //joss
355
356
Sheoj.org/papers/t0..2 2ttp://01..111///osso 01311
357
Larry Goldstein and Bryan Langholz. pling in the Cox Regression Model. The Annals of Statistics, 20(4):1903-1928 De- cember 1992. ISSN 0090-5364, 2168-8966. doi: 11/11///0//1173488888 http:///ooogccaall..crc//o/////s
Asymptotic Theory for Nested Case-Control Sam-
359
URL
360
361
the- Cox/ Publisher: Institute of Mathematical Statistics.
382
10. 10.1214/200//11764488883
363
Aditya Grover, Eric Wang, Aaron Zweig, and Stefano Ermon. Stochastic Optimization of Sorting Networks via Continuous Relaxations. arXiv:/903.088. 08850 [cs, stat], April 2019. URL http
365
/aaxiiv.org 0tt////b/1993088500000 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)2254325446 May 1982 ISSN 0098-
367
7484. doi:
URL https://doi... org/ 10. 1001/101/// /0000/jaama.. jama
369
1982. 0332043044000
370
Alistain E. W. Johnson, Tom J. Pollard, Seth J. Berkowitz, Nathaniel R. Greenbaum, Matthew P. Lungren, Chih-ying Deng. Roger G. Mark, and Steven Horg. MIMIC-CXR, a de-identified
371
publicly available database of chest radiographs with free-text reports. Scientific Data, 6(1);
373
317, December 2019. ISSN 2052-4463. doi: 10. 11888/44597--11-02222 URL https://
nwtura. wateee.coccaaae cwmm/articlee//441577000000 01220009000000000009---00000 Group.
Number: 1 Publisher: Nature Publishing
375
376
Jared L. Katzman, Uri Shaham, Alexander Cloninger, Jonathan Bates, Tingting Jiang, and Yu- val Kluger. DeepSurv: personalized treatment recommender system using a Cox proportional hazards deep neural network. BMC Medical Reseurch Methodology, 18(1)22,, February 201 8. URL httpp:://dot org/10. 1186/
377
379
ISSN 1471-2288. doi: 10.1
112874-018-04821

---

381
Donald E Knuth. The Art of Computer Programming. Addison Wesley Longman Publishing Co., Inc, Volume 3: (2nd Ed.) Sorting and Searching, 1998.
382
385
Håvard Kvamme and Ørnulf Borgan. The Brier Score under Administrative Censoring: Problems and Solutions. Technical Report 1 20012.0.880 08581, arXiv, December 2019, http:////aa org/abs/1912 10tpa/881110008111111 arxiv: 1912.08581 [es, stat] type: article.
385
386
Hávard Kvamme, Ørnulf Borgan, and Ida Scheel. Time-to-Even Prediction with Neural Networks and Cox Regression. arXiv: 1907.00825 [cs, star), September 2019. URL stttp://arxix org/
388
abs/1907.008//
arXiv: 000..00
389
Changhee Lee, William Zame, Jinsung Yoon, and Mihaela Van Der Schaar. DeepHit: A Deep Learn- ing Approach to Survival Analysis with Competing Risks. Proceedings of the AAAI Conference on Artificial Intelligence, page 8, 2018.
390
392
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 Learn- ing and Unsupervised Feature Learning 2011, 2011. URL httpht/tp///o/// stanford edu/
383
395
housenumbers h htw.re0uue..uuuuuuuebeeeeee. pdf
396
Felix Petersen, Christian Borgelt, Hilde Kuchne, and Oliver Deussen. Differentiable Sorting Networks for Scalable Sorting and Ranking Supervision. arXiv:2105.04011 [cs], July 2021. | arXiv: 2105.04019.
398 // //arxiv.org abs/ 2105. 04019
389
Felix Petersen, Christian Borgelt, Hilde Kuehne, and Oliver Deussen. Monotonic Differentiable Sorting Networks, arXiv:2203.0966 [cs, stat]. March 2022a. URL httpp:////axx org/ abs/ 2203.09630 arXiv: 2203.09630.
400
401
Felix Petersen, Hilde Kuehne, Christian Borgelt, and Oliver Deussen. Differentiable Top-k Classification Learning, June 2022b. URL http:ttt///arxivorrr/////// wrwp.//wrr//bb///22....... 07290 Number:
403
arXiv:2206.07299 arXiv:2206.07290 [cs].
Sebastian Pölsterl. scikit-survival: A library for time-to- analysis built top of seakit-leam. Journal of Machine Learning Research, 21(212): 2020. URL http://jmlr org waaggppprs/////
406
v21 20- 729. html
408
Sebastian Pölsterl Survival Analysis for Deep Learning, July 2019. URL httthtttt/s-d-uu http:////////....iaaaaaaaaaaaaaa
409
410
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 Neural Infor- mation Processing Systems, volume 20. Curran Associates, Inc., 2007. URL https::::/// //tppe/pppps.pe
411
412
413
Absrtrat.bsttrc.t..omm .tmll
414
Mingxing Tan and Quoe V. Le. EfficientNet: Rethinking Model Scaling for Convolutional Neural Networks, September 2020. URL http://aaxxv..oo//aae///000 11946 arXiv:1905.1194 [cs, stat].
416
417
Bokun Wang and Tianbao Yang. Finite-Sum Coupled Compositional Stochastic Optimization: Theory and Applications, September 2022. URL http::/arxiv.. www.orggg/a/////a//eea/e/ee/ 12396 arXiv:2202.122 12396 [cs, math].
419
420
Fen Xia, Tin-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
421
422
Xuelin Yang, Louis Abraham, Sejin Kim, Petr Smirnov, Feng Ruan, Benjamin Kaine-Kains, and Robert Tibshirani. FasICPH: Efficient Survival Analysis for Neural Networks, August 2022. URL http://arxiv.cr org/ htrg/abs/22003 09793 arXiv:2208.0077 [cs, stat].
424
425
12

---

428 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-particulel involving con-propottional hazards it does not necessarily limit the model's effectiveness in sce- narios 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 prn-orrpot models as indicated by/ Antollni et al. 22055.
428
those
430
432
433
If we are primarily interested in understanding the cumulative hazard of an event occurring rather than tracking changes the hazard over time, the assumption of proportional hazards becomes less pivotal. As such, Diffsurv and CPL remain valuable tools for these cases.
434
436
Despite the current limitation of Diffsurv to proportional hazards, it is conceivable that an extension to accommodate non proppoortiioal hazards could be developed, similar to adaptations made for the CPL method.
437
438
439
For instance, as we briefly mentioned earlier, continuous- time extensions of partial likelihood can be used to enable non- con-proportional hazards Kvamme et al. [219]] Implemented by directly modeling temporal covariates as fo = exp(ho (z;, 7.)).
441
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 et al. (2018) employs neural network architecture to learn the relationships between input features and discretized time-to to-eent outcomes. Time discretization facilitates modeling of non-propottional 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-base based models like DeepHit, making them less suitable for high-dimensional and lcrge-scale datasets, such as the imaging dataset used in this study.
445
445
447
449
451
Several future work proposals arise from these observations. First, differentiable sorting could explore the approach of directly modeling temporal covariates, resulting in paramerrrmeteri predicted permutation matrix. Second, extending Diffsurv to discrete time could be achieved by parameterizing predicted permutation matrix for each time discretization.
453
455
B
Training and evaluation
456
B.1
Datasets and Preprocessing
457
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).
458
459
480
We provide an additional description of each small realworld dataset:
461
FLCHAIN dataset: A dataset containing information on patients with monoclonal gam- mopathy 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.
462
463
464
465
NWTCO dataset: 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.
466
467
SUPPORT dataset: A dataset from multi-center study investigating the prognosis and treatment preferences of seriously ill hospitalized adults, with the goal of improving end-of- life care and informing decision-makin processes. Number of covariates: 22.
469
METABRIC dataset: A dataset comprising genomic and clinical data breast cancer patients, focused on uncovering novel molecular subtypes for more precise prognostication and personalized treatment strategies. Number of covariates: 9.
472
13

---

474
Covariate preprocessing follows Kvamme et al.) 2019), and includes standardising continuous
475
variables and one-hot encoding categorical variables.
476
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" id",, "study_id" _id"", and "dicom_ _id" with the patient table from MIMIC IV, and on "subject_id" with the admission table. For patients without recorded date of death, censoring dates are be determined as I 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 recorded date of death.
478
479
480
481
482
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 ss mean and standard deviation values. Finally, the talin:val:test split of
483
484
485
486
487
8:1:1 is done the patient level ensuring no images from a patient in the test set was found in the
training data.
489
SurvSVHN: In this semi-synthetic dataset, sample survival times using the beta-exponential
480
distribution. The beta distribution used to sample from the exponential uses a fixed value of 500
491
for both shape parameters. We follow Petersen et al. [02T]] by cropping the centered multi-digit
492
numbers with a boundary of 30%, resizing it to resolution of 64x64, and then selecting 54 x 54
493
pixels at a random location. For survSVHN the valinet split is provided by Netzer et al. as is 230,755:5,000: 13,0688,
2011
484
B.2
Model Architecture and Hy Hyperparameters
496
For the smaller real-world datasets, the hazard function fo is small fixed Multi-layer Perceptron
487
network with hidden layer and 64 hidden nodes. We also apply fixed dropout rate of 1. Learning
438 rate, weight decay, batch set size and risk set size were found using a grid search across the possible
values in 4
Table 4: Hyperparameter values for small real-world datasets.
Hyperparameter
Values
Learning rate
[0.1, 0.01, 0.001, le-4]
Weight decay
[0.1, 0.01, 0.001, Ic-4, 1e-5, 0]
(Batch size, risk set size)
[322, 8), (16, 16), (8, 32), (4, 64), (1, 256)]
500
For imaging datasets, we fix learning rate and weight decay for both CPL and Diffisurv For both SurvSVHN and MIMIC IV CXR, we use fixed learning rate of 10-4 and weight decay of 10-5, We also used early stopping with patience of 20 epochs and a maximum of 100,000 training steps.
501
502
SurvSVHN: As per Petersen al. [2021] the model consists of four convolutional layers (with 5x5 kernel size and 32, 64, 128, 256 filters), each followed by ReLU and max- (2x2 stride). The architecture concludes with a fully connected layer of 64 units, another ReLU, and one-unit output layer.
504
506
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-BO BO and the linear prediction layer.
508
510
For the results in Table I we keep fixed batch size of 100. We also provide a comparison where the number of values is fixed in each batch in Table
511
Note that during evaluation, the sorting network is not used since only need 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.
513
514
515
Further implementation details and the best hyperparameters for each dataset are provided at anon git com
14

---

517
B.3
Compute Requirements
518
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 min- utes. However, the larger imaging datasets require more significant computational power. In the most demanding case, the MIMIC IV CXR experiments, run
519
520
1 IGB NVIDIA GeForce GTX 1080 Ti, took roughly 18.5 hours per experiment. Both Diftsurv 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
522
523
524
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.
526
527
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 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 (Fquation 15) and binary cross entropy (Fquation 16). The possible permutation matrix generation possible to precompute off GPU on the dataloader, so is not included in the timing. In T 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.
528
529
530
532
533
534
536
537
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 Diffsury Bitonic variant being marginally faster than other methods. In the context of full model runs, the difference between Diffsury and CPL in terms of training times is minimal; it the model architecture that have a dominant effect on compute time.
538
539
540
542
Table 5: Isolated compute time for different methods, with various batch sizes and risk set sizes over 100 trials. Mean time and 95% confidence intervals provided in milliseconds.
Batch Size, Risk Set Size
Method
512,2 2
128, 8
32. 32
8, 128
Diffsurv: Odd-Even
12.37 0.15
25.63 + 0.09
107.71 + 0.38
305.91 3.66
Diffsurv: Bitonic
3.85 + 0.03
17.29 = 0.08
55.49 + 0.47
92.32 + 0.80
CPL: Breslow
853.32 11.39 11.39
281.22 1.20
94.08 + 0.15
25.45 + 0.46
CPL: Efron
1729.87 + 5.54
494,83 + 4.27
164.73 + 1.40
51.69 + 0.76
CPL: Ranked List
719.42 + 4.99
225.69 0.94
74.36 0.18
18.19 0.31
C
Sorting Networks
543
There are multiple different types of sorting networks each with varying complexity. The ability to implement networks with the divide-and-conquee 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 Onloggnn verses the (n²) in Odd-Even networks.
545
547
549
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.
551
552
D Calibration of Predicted Permutations
554
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
15

---

Figure 4: Example sorting networks of size 8; (a) Odd-Even, (b) Bitonic.
Table 6: Calibration of rank probabilities from predicted permutation matrices on Sur'SVHN, keeping trails with
the number of events per batch equal. Mean (and standard deviation) Brier scores different seeds. Corresponds with C-index results in Table
Batch Size, Risk Set Size
Method
512, 2
128, 8
32, 32
8, 128
CPL: Breslow
0.081 (0.002)
0.181 (0.001)
0.056 (0.000)
0.014 (0.000)
CPL: Efron
0.081 (0.002)
0.180 (0.001)
0.055 (0.000)
0.014 (0.000)
CPL: Ranked List
0.066 (0.001)
0.185 (0.001)
0.053 (0.000)
0.012 (0.000)
Diffsurv: Bitonic
0.059 (0.001)
0.165 (0.001)
0.043 (0.000)
0.010 (0.000)
Diffsurr: Odd-Even
0.056 (0.001)
0.158 (0.001)
0.039 (0.000)
0.009 (0.000)
556
to incorrect risk assessments and treatment decisions, potentially resulting suboptimal patient care and even adverse clinical consequences. We focus on the calibration of predicted individual rankings, as demonstrated in Figure and Table Specifically, we qualitatively illustrate in Figure that for model with a risk set size of 8, both discrete predicted ranks and ranking probabilities accurately calibrated for the Diffsury approach. To perform quantitative comparison with baseline methods, we need derive ranking probabilities for the CPL model. Based the assumptions in Raykar et al. [2007], we assume that the probability of correct pairwise ordering for the CPL adheres 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
558
559
580
582
563
565
differentiable sorting operator. By subsequently calculating Brier scores for the rank probabilities in the predicted permutation matrices survSVHN, 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
567
A
B
Permutation probability
6
Predicted rank
0.20
5
4
3
0.15
2
Perfect calibration
0.10
2
7
3
4
Ground truth rank
Ground truth rank
Figure 5: Calibration visualization of the Diffsurv (Bitonic) predicted ranking for group of subjects on the SurvSVHN dataset.
16

---

589
E
Effect of top-k loss variants on predictive performance
570
The Top-K risk prediction variants for Diffisur 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. 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
572
578
Table 7: C-index for real-world and semi-synthetic datasets for Top-k loss variants
FLCHAIN
NWTCO
SUPPORT
METABRIC
MIMIC IV CXR
Top 10% prediction
Cox Partial Likelihood
(.010)
.725 (.014)
.650 (.008)
.691 (.013)
.760 (.002)
CPL-TopK (Variant I)
.794 (.011)
.715 (.011)
.602 (.003)
.644 (.016)
.751 (.002)
CPL-TopK (Variant II)
.798 (.011)
.722 (.016)
.600 (.010)
.588 (.008)
.753 (.002)
Diffsurv
.794 (.011)
.696 (.025)
.634 (.008)
.649 (.016)
.761 (.002)
Diffsurv-TopK
.783 (.011)
.689 (.022)
.596 (.005)
.639 (.023)
.754 (.001)
F Additional Results
574
575
In Table and T 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 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.
578
577
578
580
Table 8: Additional Results for MIMIC IV CXR. Mean and standard deviation of the C-index for different methods and batch risk sizes. Bold indicates significantly higher result with t-test and p < 0.01.f Most significant across all Batch Size, Risk Set Sizes.
Method
Batch Size, Risk Set Size
64, 2
4, 32
16, 8
1, 128
Diffsurv: Bitonic
0.761 (0.001)
0.761 (0.000)
0.763 (0.001)
0.761 (0.002)
Diffsurv: Odd-Even
0.761 (0.002)
0.756 (0.002)
0.761 (0.001)
0.749 (0.001)
CPL: Ranked List
0.760 (0.002)
0.755 (0.002)
0.758 (0.003)
0.755 (0.002)
Table 9: Additional results for survSVNH keeping the number of events per batch equal. Mcan (and standard deviation) over trials with different seeds. Metrie is C-index. Bold indicates significantly higher result with t-test and < 0.01.
Batch Size, Risk Set Size
Method
512, 2
128, 8
32, 32
8, 128
Diffsurv: Odd-Even
0.934 (0.001)
0.940 (0.001) 0.942 (0.001)
0.941 (0.001)
0.933 (0.002)
Diffsurv: Bitonic
0.931 (0.001)
0.940 (0.00166)
0.928 (0.001)
CPL: Breslow
0.905 (0.001)
0.897 (0.001)
0.910 (0.002)
0,919 (0.001)
CPL Efron
0.904 (0.002)
0.898 (0.002)
0.909 (0.003)
0.918 (0.003)
CPL: Ranked List
0.921 (0.001)
0.922 (0.003)
0.921 (0.001)
0.917 (0.003)
17


Reviews:
[{"review_id": "1cs8x0bqST", "reviewer": "Reviewer_GX1f", "summary": "The 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.", "questions": "1. 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", "limitations": "1. 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.", "rating": 5, "confidence": 4, "soundness": 2, "presentation": 2, "contribution": 2, "strengths": "By 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.", "weaknesses": "1. 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", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your feedback. We value your insights and aim to address each concern you've raised. Here are our responses and clarifications:\n\n---\n\n**Weakness 1 and Limitation 2 (Predicting the time of event vs ranking and calibration scores)**\n\nWhile estimating absolute timings for event occurrences is important in certain applications, many do not require this specific information. Decision-making through survival analysis frequently focuses on relative comparisons rather than absolute values. For instance, medical professionals may be more concerned with whether a particular observation will increase or decrease survival time rather than with the exact absolute value. Additionally, in some applications, the ranking itself is the primary objective of survival analysis. This may be evident in scenarios like assigning a limited discrete resource, such as a vaccine, to a subset of individuals most at risk.\n\nImportantly, we do not learn a hazard function as in CPL, but rather a ranking function. While it may be possible to recover a hazard function using adaptations of methods such as Breslow's estimator, we see this as future work.\n\nWhile we can’t assess commonly used calibration scores like IBS in the ranking setting, we now report Brier scores for the permutation ranking (see point model calibration in global response and attached PDF) and show that Diffsurv scores are better than the CPL baseline. \n\n---\n\n**Weakness 2 (Meaningful ordering within top-k?)**\n\nThe top-k task proposed in our paper intentionally does not focus on the order within the top-k, and we believe that this task has significant real-world clinical applications. For example, in scenarios where the objective is to identify the top-k highest-risk individuals for preventive measures (e.g., a vaccine with limited supply), the order within this group is irrelevant, as all selected individuals will be contacted. Therefore, we neither optimise for nor evaluate the meaningfulness of the predicted ordering of individuals within the top-k grouping. Instead, our results demonstrate that the models can accurately identify the top-k highest-risk individuals and that our method consistently outperforms the baselines in this context.\n\n---\n\n**Weakness 3 (Necessity/motivation for proposed method)**\n\nIntroducing differentiable sorting networks doesn't target the memory needs of large risk sets. By using these networks, our method integrates transitivity as an inductive bias into survival ranking, a novel approach to survival analysis with at least competitive results to established baselines. More importantly, Diffsurv establishes a framework that is also straightforward to build upon, such as in a top-k setting enabled by predicted permutation matrices from the differentiable sorting method. We note that established methods like CPL can’t easily be extended to predict permutation matrices directly, and we now show (see attached PDF in global response with new calibration experiments) that even if we extend CPL to do that, they exhibit worse calibration in predicted rankings compared to our method. \n\n---\n\n**Weakness 4 (Explanation of sorting networks)**\n\nWe have now expanded our description of sorting networks in Appendix E and added more references to traditional sorting networks (in addition to the papers introducing differentiable sorting networks).\n\n---\n\n**Weakness 5 & Question 3 (Comparison with baselines)**\n\nWe refer to Cox Partial Likelihood (CPL) as a loss function, without a specific covariate interaction, $h_\\theta(x_i)$. Cox Regression refers to the scenario where linear interaction is maintained, while in CPL baselines, $h_\\theta$ is parameterised with a neural network, explaining the better C-index performance. CPL is the default loss function in many recent deep learning survival analyses [1, 2]. We train our CPL baselines with the same architectures as Deepsurv, tuning the hyperparameters extensively. Other baselines, like DeepHit, are not expected to outperform CPL in this setting due to the non-proportional hazards affecting the time-dependent C-index but not the standard C-index. See our response to Weakness 1 regarding time dependency.\n\n---\n\n**Question 1 (Risk set size vs Batch size)**\n\nWe follow Petersen et al., 2021 [3] and train the model in batches of risk sets, i.e. each minibatch has shape (minibatch size, risk set size, feature size). Thus, for minibatch size 8 and risk set size 32, 256 samples are in the batch. This is described in more detail in section D.1 in the appendix.\n\n---\n\n**Question 2 (Performance with increasing batch size)**\n\nThe performance of both CPL and Diffsurv improves with larger risk set sizes, as shown in Table 1, but Diffsurv benefits more from larger risk set sizes. We believe this is because our method benefits more from the inherent transitivity in these larger risk sets due to the introduction of differentiable sorting methods. Please also see L275-281 in the manuscript for a more detailed explanation of this phenomenon and Table 2 with another experiment investigating this in more detail.\n\n---\n\n**Question 4 (Transitivity ratio can’t be infinity?)**\n\nWe don’t report a transitivity ratio of infinity. Instead, Table 2 shows that with an infinite number of quantiles (no discretisation as defined on L274), the transitivity ratio is .991.\n\n---\n\n**Limitation 1 (Time-varying covariates)**\n\nWe recognise the significance of models addressing time-varying hazards. While our method can be expanded in future work, our current focus is on fixed-time ranking and introducing differentiable sorting networks to risk analysis.\n\n---\n\n[1] Buergel et al. 2022. ‘Metabolomic Profiles Predict Individual Multidisease Outcomes’. Nature Medicine.\n\n[2] Carr et al. 2021. ‘Longitudinal Patient Stratification of Electronic Health Records with Flexible Adjustment for Clinical Outcomes’. PMLR.\n\n[3] Petersen et al. 2022. Differentiable sorting networks for scalable sorting and ranking supervision. ICML"}, {"review_id": "TjZESGdMhc", "reviewer": "Reviewer_RSfV", "summary": "This 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)", "questions": "- 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?", "limitations": "- This paper is limited to proportional hazards assumptions, but it is clearly acknowledged by the authors", "rating": 7, "confidence": 3, "soundness": 4, "presentation": 3, "contribution": 4, "strengths": "- 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)", "weaknesses": "- 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?", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We sincerely appreciate your thorough evaluation and positive feedback on our manuscript. Your insights into our work's potential impact and originality are highly encouraging. We acknowledge your concerns, especially regarding the clarity of notations and the handling of the top-k task. In this rebuttal, we strive to address each point further to enhance the quality and clarity of our contribution.\n\n---\n\n**Weakness 1 (Clarity)**\n\nHere, we reproduce the notation from Kvamme et al. [1]. In theory, the sum applies to all such j as in the C-index. In practice, we must sample a limited number j’s according to the minibatch size. This is done randomly from the set of possible j’s.\n\n---\n\n**Weakness 2 and Question 2 (Top-k task)**\n\nWe thank the reviewer for highlighting the potential influence of the set size on the top-k-score due to partial ordering. Indeed, a large set size, stemming from partial ordering, can theoretically lead to an arbitrarily high top-k-score. However, a few points to note:\n* Our primary emphasis is on relative scores between models. While individual scores could be elevated due to partial orderings, the relative differences between models remain and can be used to assess their comparative performances.\n* Empirically, we observed that the scores manifest meaningful differences across methods and real-world datasets, including those with substantial censoring. These variations support the practical effectiveness of the metric.\n\n---\n\n**Question 2 (Runtime analysis)**\n\nThank you for the question on runtime differences between our proposed method and previous approaches, which we have now analysed in more detail in the global response and attached PDF.\n\n---\n\n[1] Kvamme, Havard, Ørnulf Borgan, and Ida Scheel. \"Time-to-Event Prediction with Neural Networks and Cox Regression.\" Journal of Machine Learning Research 20 (2019): 1-30."}, {"review_id": "JjgzVWDs0K", "reviewer": "Reviewer_Rdgo", "summary": "The 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.", "questions": "-\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\n", "limitations": "-\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. ", "rating": 5, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "Extending a differentiable sorting method to account for censoring and applying it to train survival models is novel.", "weaknesses": "-\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", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We appreciate your feedback and hold your insights in high regard. Our goal is to respond to every concern you have mentioned:\n\n---\n\n**Weakness 1 (Writing)**\n\nThank you for the feedback on typographical and grammatical errors. We've addressed these and ensured greater consistency in notation and equation references.\n\n---\n\n**Weakness 2 and Question 1 (Pairwise vs Listwise and Formal definition of transitivity)**\n\nThe motivations for listwise over pairwise methods and their connection to transitivity warrant clarity. Recall that Cox Partial Likelihood (CPL) is a listwise approach, as indicated on L123. Let's highlight the benefits of listwise methods and then differentiate between differentiable sorting and CPL.\n\nPairwise sorting methods handle ranking pairs individually. They establish local orderings (A over B, B over C) but may not ensure global consistency (A over C) throughout the set, as outlined in L113-115 and L266-274. Such limitations arise from focusing solely on local pairings. In some scenarios, these local orders can create inconsistencies.\n\nIncreasing the minibatch size to include more pairs poses scalability challenges. Listwise methods evaluate comprehensive relationships among data points, scaling factorially. Ensuring pairwise order for every set may imply transitivity, but it's not always the case due to inherent inconsistencies when solely considering pairs. Refer to [1] for more on listwise motivations.\n\nComparing CPL and Diffsurv, while both are listwise, they have distinct principles. CPL follows a top-one approach. Each term in the product ensures that one patient with an earlier observed event must higher rank than a subset of those observed later. Diffsurv instead considers the order of every patient in the set, and any incorrect ordering within each layer of the relaxed sorting network propagates to a higher predicted probability of permuting to an incorrect rank.\n\n---\n\n**Weakness 3 and Question 2 (Transitivity in the semi-synthetic experiment)**\n\nIt is important to highlight the role of transitivity in larger risk set sizes, which directly contributes to the higher performance of our method compared to the CPL baseline. The experiments shown in Table 1 and the details in lines L275-281 of the manuscript support this claim by demonstrating that both CPL and Diffsurv improve with larger risk set sizes. However, Diffsurv, due to the introduction of differentiable sorting methods, benefits more from the inherent transitivity in these larger risk sets, resulting in a more substantial performance delta.\n\nFurthermore, our results validate the significance of listwise sorting in the semi-synthetic survSVHN dataset. Table 2 provides an in-depth exploration of this, revealing that as the degree of transitivity within the data increases, Diffsurv consistently outperforms CPL. This divergence emphasises the importance of the differentiable sorting network in capturing the transitivity in the ranking task.\n\n---\n\n**Weakness 4 (Benchmarking)**\n\nOur proposed method focuses on survival ranking, and we don’t learn a hazard function as in CPL, but rather a ranking function. Clinical decisions often hinge on relative comparisons, such as whether an observation affects survival time, rather than precise event timing. Additionally, there are specific applications in survival analysis where ranking is the primary objective, such as resource allocation among patients. In line with recent works in deep learning survival analysis, we chose CPL as the baseline and rigorously tuned the hyperparameters to ensure a fair comparison.\n\nSince our method is geared towards survival ranking, predicting event times and thus scores like IBS are not applicable in this situation. However, we recognise the value of assessing model calibration and appreciate your suggestion to include Brier scores for predictive power. We have incorporated the Brier scores for the permutation ranking in the attached PDF of our global response. Our findings demonstrate that Diffsurv outperforms the CPL baseline. \n\n---\n\n**Weakness 5 (Mathematical notation)**\n\nWe agree that adding the mathematical notation introduced in L73-77 and relating that to the possible ranks in L187-191 improves clarity and have adapted this section accordingly.\n\n---\n\n**Question 3 (Ranking loss for right-censored samples)**\n\nIn response to the concerns raised about the ranking of samples that cannot be directly compared, it is essential to highlight the design of our loss function, which has been constructed to solve this issue. The loss function (Equations 13 and 14) accounts for the potential challenges of incorporating right-censored samples. For an early right-censored sample $i$, $Q_{pi}$ is 1 almost everywhere, effectively sidestepping the risk comparison issue. The binary cross-entropy remains identical whether the model predicts uniform probability for all possible ranks or concentrates the probability mass on an early rank, inferred from the covariates. This ensures that the predicted rank falls within the possible ranks, considering censoring, without further influencing which rank within those possible ranks is preferable. We have now clarified this point in the methods section.\n\n---\n\n**Question 4 (Censoring in semi-synthetic data)**\n\nTo clarify, our semi-synthetic data generation uses only independent censoring. Is the question whether dependent censoring would favour our method in the benchmarks? \n\n---\n\n**Limitation 1 (Runtime of calculating Qp)**\n\nWe now report both theoretical time complexities of the differentiable sorting networks and empirical runtime results, including calculation of the $Q_p$ matrix in appendix section B.3 and show that the distinction in training times between Diffsurv and CPL is insignificant. Please also see our global response for further details on the runtime analyses. \n\n---\n\n[1] Cao et al. 2007. ‘Learning to Rank: From Pairwise Approach to Listwise Approach’. ICML.\n"}, {"review_id": "YTLNValkt9", "reviewer": "Reviewer_e7KQ", "summary": "This 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.", "questions": "- 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?", "limitations": "The 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.", "rating": 5, "confidence": 5, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "- 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", "weaknesses": "- 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\n\n", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your thorough review of our manuscript. We've addressed each point you raised and provided further clarifications where necessary:\n\n**Weakness 1 (Theoretical and practical advantages of the method)**\n\n*Theoretical Advantages*:\n\nDiffsurv integrates differentiable sorting networks, thereby introducing a transitive inductive bias into the survival ranking model. We have shown that this novel survival analysis approach performs competitively compared to CPL. More significantly, the ordering supervision enables straightforward implementation of new tasks, such as top-k risk prediction. In contrast, CPL cannot be readily adapted to predict permutation matrices directly. As we demonstrate in the global response and new experiments, attempts to adapt CPL for this purpose result in poorer calibration in the predicted rankings compared to Diffsurv.\n\n*Experimental Advantages*:\n\nOur experimental findings, detailed in Table 1 and lines L275-281, demonstrate Diffsurv outperforms CPL, particularly as transitivity increases within the data. This improvement is pronounced in larger risk set sizes, where our method leverages the inherent transitivity due to the introduction of sorting methods. Further experiments (Table 3) reveal consistent outperformance of CPL by Diffsurv in the top-k setting with significant clinical applications, like identifying top-k highest-risk individuals for specific treatment of preventative measures (see also our reply to Question 1 below).\n\n---\n\n**Weakness 2 (Real-world benefits)**\n\nWe acknowledge in the real-world risk stratification benchmark datasets, Diffsurv is only marginally better than the CPL baseline. We have introduced a novel approach to survival analysis using differentiable sorting algorithms that is at least competitive with established methods, but importantly also enables the development of novel downstream tasks using ordering supervision such as top-k risk prediction.\n\n---\n\n**Weakness 3 (Potential of Diffsurv in non-PH settings)**\n\nWe agree extending Diffsurv to time-dependent predictions without PH assumptions similar to CoxTime or DeepHit is a promising approach that will likely show a clearer performance benefit compared to the baselines and aim to address this in future work.\n\n---\n\n**Question 1 (Significance of top-k risk prediction task)**\n\nWe believe the top-k risk prediction task has important clinical applications: Whereas C-index accounts for risk stratification and measures the proportion of concordant pairs in all comparable pairs, the top-k risk prediction task as introduced in the manuscript tries to identify the set of top-k highest risk individuals, regardless of the orderings of individuals within the set of top-k highest individuals. The value of the top-k risk prediction task emerges in real-world clinical scenarios where resources like treatments or preventative interventions, such as vaccines, are constrained. In these situations, the primary concern is identifying high-risk individuals most likely to benefit from interventions, rather than predicting relative risk within that top tier. Risk ordering within the top-k individuals is secondary and may not affect immediate clinical decisions. Note this task is straightforward to implement with Diffsurv due to the predicted permutation matrix and thus showcases the advantages of including ordering supervision in a risk stratification model.\n\n---\n\n**Question 2 (Pairwise independence and transitivity in Kvamme et al.)**\n\nWe acknowledge that Kvamme et al. do not explicitly assume pairwise independence. However, it's worth noting optimisation with a risk set size of 1 inherently assumes pairwise independence. While Kvamme et al. demonstrate this is not an issue in their work, our empirical results contradict this. In our experiments, we observe benefits for larger risk sets, both for CPL and Diffsurv, though stronger for Diffsurv. The following differences in our study might explain this discrepancy:\n\n* We evaluate C-Index, while Kvamme uses MPLL, a lower bound of C-Index [1]\n* Our dataset and model complexity make optimisation more challenging, which might lead to our observed benefits for bigger risk sets\n* We have evaluated larger risk sets compared to Kvamme et al.'s study\n\n---\n\n**Question 3: (Top-K prediction changes to discrimination and calibration)**\n\nWhile we did not include these results in the manuscript, we have noticed all model variants trained using top-k objectives perform slightly worse in risk stratification when evaluated using C-index. \n\n---\n\n**Question 4 (Computational cost analysis)**\n\nWe have now analysed the time complexity of the methods and run additional experiments to measure runtimes, please see the global response and the attached PDF. \n\n---\n\n**Question 5 (Baselines)**\n\nWe believe there is no stronger baseline than a tuned deep neural network with a CPL objective. While other models like DeepHit tend to outperform CPL in $\\text{C-Index}^\\text{td}$, for the ranking setting (non time-dependent) C-Index is the appropriate metric. We still include RSF as it is a commonly used baseline in the literature and is meaningfully different from Cox regression and neural networks trained with CPL.\n\n**Limitation 1 (Theoretical analysis of Diffsurv)**\n\nWe agree CPL has some favourable statistical properties and appreciate the concern for a more thorough analysis of the theoretical properties of our proposed method. We would like to point out these align closely with the properties of sorting networks and differentiable sorting networks, which have been studied extensively in prior work [2, 3, 4].\n\n---\n\n[1] Steck et al. 2007.\"On ranking in survival analysis: Bounds on the concordance index.\" NeurIPS.\n\n[2] Knuth, D. E. 1997. The Art of Computer Programming. Addison–Wesley.\n\n[3] Petersen et al. 2021. \"Differentiable sorting networks for scalable sorting and ranking supervision.\" ICML.\n\n[4] Petersen et al. 2021. \"Monotonic Differentiable Sorting Networks.\" ICML.\n"}]

Author rebuttal:
We sincerely appreciate the thorough and constructive feedback provided by all four reviewers. 
As suggested by multiple reviewers, we have further investigated the runtime (both theoretically and empirically) and calibration of our proposed survival ranking method.

---

**Theoretical and empirical analysis of runtime**

Diffsurv and the baseline models using CPL use the same neural network architectures. The operations relating to these networks are the main bottleneck for all methods, particularly for the imaging datasets. Only minor differences in runtime are due to the respective Cox Partial Likelihood (CPL) and differentiable sorting operations.

*Theoretical Time Complexities*: First, we highlight the number of operations required for the differentiable sorting networks: Odd-Even as $\mathcal{O}(n^2)$ and Bitonic sorting networks as $\mathcal{O}(nlog^2n)$ as discussed in Appendix C and [1]. The time complexity of our primary baseline, CPL, is $\mathcal{O}(nlogn)$ as it also requires sorting the event times.

*Benchmarking Process*: We have conducted experiments to understand these differences (see results in the attached PDF). We measured the time taken for a forward and backward pass on an NVIDIA GeForce GTX 1080 Ti, utilising randomly generated logits. This experiment provides a more isolated measure of the compute times, allowing for a focused comparison between the methods. For Diffsurv, our timing includes the predicted permutation matrix computation via differentiable sorting networks and the subsequent masking (Equation 13) and binary cross-entropy (Equation 14). Importantly, we precomputed the potential permutation matrix ($Q_{p}$) generation off-GPU (same as during minibatch training), so it isn't included in this benchmark.

*Results*: As documented in Table 1, our Diffsurv methods, especially Bitonic, have a noticeable edge over CPL methods regarding compute time across various batch and risk set sizes. It's interesting to observe that as the risk set size goes up, CPL methods show a reduced compute time, while Diffsurv's Odd-Even variant sees a substantial increase, particularly when moving from risk set sizes of 32 to 128.

*Optimization Possibilities*: It's crucial to note that current implementations of CPL methods compute over batches using a straightforward 'for' loop since they don't support batch-parallel computation. There's significant room for further optimisation here.

*Overall Observations*: In comprehensive model runs, the difference in training times between Diffsurv and CPL is insignificant. The compute time is primarily driven by the model architecture. Specifically, we observed similar convergence times across methods, with the Diffsurv Bitonic variant being slightly quicker.

We recognise the importance of computational efficiency in practical deployments and point to [1] for further analysis of differentiable sorting network compute time.

---

**Model calibration**

It is important to note that we do not learn a hazard function as in CPL, but rather a ranking function, and thus can’t assess predictive model calibration using commonly used scores like IBS. We acknowledge that model calibration in survival analysis models is essential for ensuring that the predicted probabilities of outcomes align closely with the true probabilities. 

We include a new analysis (see attached PDF), focussing on the calibration of predicted individual rankings. Specifically, we first qualitatively illustrate in Figure 1 in the attached PDF that discrete predicted ranks and ranking probabilities are accurately calibrated for a model with a small risk set size for the Diffsurv approach. To perform a quantitative comparison with baseline methods, we need to derive ranking probabilities for the CPL model. Based on prior work [2], 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 from a pretrained model 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 and ground truth permutations derived from the hazards in the survSVHN dataset, we analyse 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 (Table 2 in the attached PDF). 

---

[1] Petersen, F., Borgelt, C., Kuehne, H., & Deussen, O. (2021, July). Differentiable sorting networks for scalable sorting and ranking supervision. In International Conference on Machine Learning (pp. 8546-8555). PMLR.

[2] Steck, Harald, et al. "On ranking in survival analysis: Bounds on the concordance index." Advances in neural information processing systems 20 (2007).

Meta-review: This paper proposes the use of differentiable sorting networks to rank right censored patients by risk (for the problem of survival analysis). This is achieved via the construction of a possible permutation matrix. The matrix captures whether its feasible for a point i to permute to rank j (and if so its set to 1) and otherwise set to 0. The numbers are setup to respect the ranking of patients in the censored and event set. The models is trained using a sorting network that enables models to learn to sort risks. The model is trained on the surSVHN dataset and the as well as standard benchmarks in survival analysis. The method appears to do reasonably well compared to learning with the Cox proportional likelihood based method when evaluated on c-index.

Overall I find this is a modest paper that I lean towards accepting because it provides a self-contained approach to ranking right censored data that sets itself apart from traditional methods for ranking. While I do think it has differences (both pros and cons) vs traditional methods for survival analysis I strongly recommend the authors make this transparent to readers at the outset (e.g. this method produces no survival curves). In addition, my understanding based on this work (though this is never explicitly stated) is that this model is expected to rank data when the outcome process is censored (conditionall) at random (given the covariates) -- it is worth making this assumption explicit so that practitioners are aware of the censorship domain this model was intended to work well in.
Decision: Accept (poster)