Differentiable sorting for censored time-to-event data.

Andre Vauvelle 1 2 Benjamin Wild 3 Roland Eils 3 Spiros Denaxas 1

Abstract

becomes particularly noticeable when the study’s endpoint
can be inferred from the observed covariates, such as age.
In such cases, the predicted event times are likely to biased
towards the censoring event time, thereby neglecting theactual event of interest (Kvamme & Borgan, 2023).

Survival analysis is a crucial semi-supervised task
in machine learning with significant real-world
applications, especially in healthcare. It is known
that survival analysis can be reduced to a ranking
task and be learnt with ordering supervision. Differentiable sorting methods have been shown tobe effective in this area but are unable to handle
censored orderings. To address this, we propose
Diffsurv, which predicts matrices of possible permutations that accommodate the ranking uncertainty caused by censored samples. Our experiments reveal that Diffsurv matches or outperforms
established baselines in various semi-simulated
and real-world risk prediction scenarios.

The Cox Proportional Hazards model is widely used forhandling censored data in survival analysis (Cox, 1972). The
model optimizes a partial likelihood function over ranked
data, considering only the order of events, not their exact
time of occurrence. As such, Cox’s partial likelihood serves
as a ranking loss, learning from the order of patients basedon their hazard of experiencing an event, not their exact
survival time.
Raykar et al. (2007) showed that Cox’s partial likelihood(CPL) and ranking losses can be directly equated, withboth 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 Lee et al. (2018).

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

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

The concept of censoring is a distinguishing characteristicthat sets survival analysis apart from conventional machinelearning approaches. Censoring refers to situations wherean event time remains unobserved because a patient might
not have undergone the event by the time of data collection.
This can be due to a variety of reasons such as the study
period ending before all events of interest have occurred or
subjects leaving the study.

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 withordering supervision (Grover et al., 2019; Blondel et al.,2020; Petersen et al., 2021). This method incorporates asorting algorithm into the network architecture, where theorder of the samples is known, but their exact values are unsupervised. With this, we introduce Diffsurv, an extensionof differentiable sorting methods that enables end-to-end
training of survival models with censored data. Briefly, ourcontributions are summarised:

Overlooking censoring can skew predictions towards the
censoring event, rather than the event of interest. This bias
1

University College London 2 BenevolentAI 3 Berlin Institute of Health. Correspondence to: Andre Vauvelle <andre.vauvelle.19@ucl.ac.uk>.
Published at the Differentiable Almost Everything Workshop of the
40 th International Conference on Machine Learning, Honolulu,
Hawaii, USA. July 2023. Copyright 2023 by the author(s).

1

CNN

Differentiable sorting for censored time-to-event data.

z1
z2
z3
z4
z5
z6
z7

P =

Qp =
P1 ·P2 ·P3 ·P4 ·P5 ·P6 ·P7 = P.5 .3 .1 .1 .0 .0 .0
.4 .3 .1 .1 .1 .0 .0
.1 .2 .4 .1 .1 .1 .0
.0 .1 .2 .4 .2 .1 .0
.0 .1 .1 .2 .3 .2 .1
.0 .0 .1 .1 .2 .3 .3
.0 .0 .0 .0 .1 .3 .6

1 1 1 1 1 1 1
1 1 0 0 0 0 0 
0 1 1 1 1 1 1 0 1 1 1 1 1 1 0 1 1 1 1 0 0 
0 0 1 1 1 1 1 
0 0 1 1 1 1 1

DiffsurvPn
p = j=1 (Qp ◦ P)i,j Survival
TopK

L=

Pn

i=1 yi log(pi )···Figure 1. Differentiable Sorting for Censored Time-to-Event Data. Inputs, in this case, SVHN images, are transformed into scalar values
through a neural network. A differentiable permutation matrix, P , is computed using sorting networks. The model can be optimized fordownstream 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.

• Our primary contribution is the extension differentiablesorting methods to account for censoring by introducing the concept of possible permutation matrices. (Section 2.3)

it possible to interpret each element Pij of the predicted
permutation matrix as the predicted probability of permuting
from a randomly assigned rank i to a true rank j.
There are multiple methods of relaxing sorting algorithms
to produce P , we will follow Petersen et al. (2021) byusing differentiable sorting networks. Sorting networks
are a family of sorting algorithms that consist of two basic
components: wires and conditional swaps. Wires carryvalues to be compared at conditional swaps, if one value isbigger than the other, the values carried forward are swappedaround. For a random sample of patients to be ordered,each layer of the sorting network can be considered an
independent permutation matrix Pl with elements given by• We demonstrate that differentiable sorting of censoreddata enables the development of new methods withpractical applications, using the example of end-to-end
learning for top-k risk stratification. (Section 2.3)
• We investigate the role of transitivity in survival analysis and show, through experiments with semi-simulateddata, that differentiable sorting networks can benefit
from this inherent property of the data. (Section 3)• We empirically demonstrate that our new differentiable sorting method matches or improves risk ranking
performance across multiple semi-simulated and realworld censored datasets. (Section 3)

Pl,ii = Pl,jj = σ(zj −zi ) and Pl,ij = Pl,ji = 1−σ(zj −zi ).
(1)
These elements represent conditional swaps between twopatient risk values (zi , zj ) and use a differentiable relaxationof the step function such as the logistic-sigmoid, whereσ : x → 1+e1−βx . The inverse temperature parameter β > 0
is introduced so when β → ∞ the functions tend to theexact min and max functions. The indices being comparedare determined by the sorting network and the final predictedprobability matrix Q
is the product of each layer of sorting⊺n
operations, P = ( l=1 Pl⊺ ) . For the base case, n = 2,
Diffsurv is equivalent to the pairwise ranking loss and Cox
partial likelihood. Further details on the relations between
Diffsurv and baselines is in Appendix A.2.

2. Methods
2.1. Censored data
A dataset with censored event times is summarized as D ={ti , xi , δi }Ni=1 , where N is the total number of patients. Fora patient i, the time-to-event ti is the minimum of the true
survival time t∗i and the censoring time c∗i , with δi indicating
whether an event (t∗i ≤ c∗i , δi = 1) or censoring (t∗i > c∗i ,δi = 0) was observed. Covariates are xi ∈ Rd representing
a 1-dimensional vector of size d but the methods discussed
here also generalise to higher dimensional tensors such asimage data.

2.3. Diffsurv: Handling Censoring with Possible
Permutation Matrices

2.2. Differentiable Sorting

For risk sets of size 2, given proper case-control sampling,
it will always be possible to define a single ground truth
permutation matrix Q. However, when venturing to higher
risk set sizes, differentiable sorting methods can no longer
handle censoring since there is not a single ground truth permutation matrix Q. We cannot determine the exact rank ofIn order to train models based on ordering information using
differentiable sorting algorithms (Petersen, 2022), we canminimize the cross-entropy between the ground truth orders
represented by true permutation matrix Q and a doublystochastic predicted permutation matrix P . This makes
2

Differentiable sorting for censored time-to-event data.

a)2. For uncensored values, we know that their rank must
be lower than all preceding uncensored events, andhigher than the ranks of all subsequent events.

time
c1

e2 c3

c4

e5 c6

c7

possible permuted ranksb)With these observations, it’s straightforward to construct
Qp . If it’s feasible for patient i to permute to rank j, then
Qpij = 1, otherwise Qpij = 0. See Figure 2 (c) for a visual
representation of Qp .

c1
e2
c3
c4
e5
c6
c7

c)

c1
1

1
0

Qp = 
0
0

0
0

e2
1
1
1
1
1
0
0

c3
1
0
1
1
1
1
1

c4
1
0
1
1
1
1
1

e5
1
0
1
1
1
1
1

c6
1
0
1
1
0
1
1

Given the possible permutation matrix Qp and the predicted
permutation matrix P, the vector of probabilities p of avalue being ranked within the set of possible ranks can becomputed. Although the ground truth probabilities are unknown, the range of possible ranks is known, and the model
can be optimized to maximize the sum of the predicted probabilities of all possible ranks for each sample. Noted here
as the column-sum of the element-wise product ◦, between
Qp and P.
nX(Qp ◦ P)i,j .
(3)
p=

c7
1  c10  e2
1
 c3
1
 c4
0  e5

1
c6
1
c7

j=1

The binary cross-entropy loss can then be easily applied

Figure 2. For an example case (a) with two events ( , e1 and e5 )and multiple censored samples ( , c1 , c3 , c4 , c6 , c7 ) the uncertainty
in the possible permuted rankings (b) due to censoring is taken
into account to derive the possible permutation matrix Qp (c).

L=

n
X

−yi log(pi ) − (1 − yi ) log(1 − pi )

(4)

i=1

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

patients who are censored before another who experienced
an event. It is only possible to know the range of possible
ranks for which a patient should belong. In Figure 2, we
provide an illustration demonstrating the possible ranks fornumber of censored and uncensored events.

The introduction of the possible permutation matrix can beused in conjunction with any differentiable sorting methodthat outputs a doubly-stochastic permutation matrix. This
includes methods such as SinkhornSort from Cuturi et al.
(2019). Though, in this paper, we will restrict our focus thediscussed differentiable sorting networks.

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

Finally, we demonstrate how the algorithmic supervision ofsorting algorithms enables the development of novel methods in survival analysis, using the example of top-k risk
prediction. In practical settings, it is often not necessary torank all samples correctly. Rather, it is essential to identify
the samples with the highest risk, such as by a healthcare
provider, to prioritize care and interventions. With Diffsurv, top-k risk prediction is straightforward to implement
by optimizing possible permutations within the top-k ranks,
whereby Qp is adjusted such that only the top-k patient’s
possible permutations are set to 1. Further details in Appendix B.

We introduce a more computationally tractable representation of Q by defining the possible permutation matrix, Qp ,which is the element-wise maximum of every permutation
in Q,
Qqij = max{Q1ij , Q2ij , . . . , Qκij }.
(2)
For survival analysis, it is possible to determine Qp in linear
time given a sorted list of event times ti and event indicators
δi . We will consider higher ranks to correspond with asmaller time-to-event. Let us consider two scenarios:

3. Experiments
In our experiments, we aim to assess the performance ofDiffsurv and compare it against the conventional Cox Partial Likelihood (CPL) methods. Initially, we focus on confirming the importance of taking a listwise approach andevaluating the ability of differentiable sorting networks to

1. For right-censored values, the only certainty is that
their rank must be lower than the ranks of preceding
uncensored events.
3

Differentiable sorting for censored time-to-event data.

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.
2†

Risk set size

4

8

16

Table 3. Results for real-world datasets. Mean (standard deviation).
Survival metric is c-index, Top 10% metric is top-k-score.

32

Cox Regression
Random Survival Forest
Cox Partial Likelihood
Diffsurv

Table 2. Results for training on survSVHN while increasing transitivity. Metric is c-index. Mean performance over 3 trails withdifferent seeds. Restricted to a fixed batch size and risk set size of
32.
2
.0

4
.374

8
.657

16
.819

32
.908

64
.954

128
.975

∞
.991

Diffsurv: Bitonic
Diffsurv: Odd Even

.643
.646

.803
.802

.882
.883

.922
.923

.933
.935

.939
.939

.939
.940

.939
.941

CPL: Ranked List
CPL: Efron
CPL: Breslow

.651
.647
.648

.803
.801
.801

.880
.871
.871

.909
.898
.898

.916
.904
.904

.920
.905
.909

.921
.909
.907

.920
.908
.910

NWTCO

6,524
69.9%

4,028
85.8%

SUPPORT METABRIC MIMIC CXR
8,873
32.0%

1,904
42.1%

377,11060.9%.750 (.083).789 (.011).794 (.013).793 (.009).692 (.021).691 (.024).709 (.015).703 (.026).598 (.010).614 (.009).642 (.006).645 (.002).628 (.013).641 (.012).698 (.011).684 (.011).760 (.002).763 (.001).460 (.013).469 (.007).460 (.009).452 (.011).482 (.019).390 (.068).413 (.061).413 (.054).395 (.082).421 (.065).280 (.023).479 (.016).479 (.035).296 (.015).508 (.027).249 (.065).527 (.083).487 (.058).331 (.102).533 (.092).390 (.010).408 (.008).406 (.006).412 (.002).412 (.009)

Survival

Diffsurv.918 (.003) .934 (.002) .940 (.001) .943 (.002) .941 (.002)
Cox Partial Likelihood .913 (.002) .925 (.002) .931 (.002) .933 (.002) .930 (.003)

Number of Quantiles
Transitivity Ratio

FLCHAIN
Size
Censored Proportion

Top 10% prediction
Cox Partial Likelihood
CPL-TopK (Variant I)
CPL-TopK (Variant II)
Diffsurv
Diffsurv-TopK

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, withDiffsurv consistently outperforming CPL. Table 2 sheds
light on a potential reason for this divergence. As the number of quantiles is increased, thereby enhancing the degreeof transitivity within the data, Diffsurv-based methods start
to surpass CPL methods. This finding underscores the roleof transitivity in survival data and validates Diffsurv’s effectiveness in encapsulating this inherent property.

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. Full details on experimental
setup including: datasets, baselines, network architectures,training, evaluation and compute requirements can be foundin the Appendix D. In results tables, bold indicates significant improvement (t-test, p ≤ 0.01).

Real-world datasets: 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 et al., 2019).
Further details in Appendix E.1.

Semi-synthetic survSVHN: Based on the Street View
House Numbers (SVHN) dataset (Netzer et al., 2011), we
simulate survival times akin to survMNIST (Pölsterl, 2019).
The increased complexity of SVHN over MNIST offers
a testbed which is better able to discern the performancedifferences between methods. Each house number parameterizes time function from which survival times are sampled.
See Appendix E.1 and Figure 3.

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

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, wegroup λi into distinct hazard quantiles. Each quantile encompasses a set of house numbers associated with a similar
hazard level. We then calculate the transitivity ratio, defined
triplets
as # of transitive, where a sampled triplet is considered# of tripletstransitive if (λi > λj > λk ).

4. Conclusion
Diffsurv represents a new direction for survival analysis.
Extending differentiable sorting methods to censored data
notably improves survival predictions, matching or outperforming established CPL methods across various datasets.
Our experiments highlight the role of transitivity in ranking and survival data, with sorting networks promoting atransitive inductive bias. Beyond survival analysis, the introduction of the possible permutations carries potential forother tasks that involve ranking based on limited order information. Future work could explore its applicability to
non-proportional hazards and the impact of ties, as well as
the efficiency in recovering survival functions using methods like Breslow’s estimator.

This methodology provides us with a means to control the
degree of transitivity in our data. At one extreme, we might
categorize data into only two groups, representing the lower
and upper halves of house numbers, which results in a transitivity ratio of 0. At the other extreme, each house number
could constitute its own unique category (indicated by ∞),leading to high transitivity.
Our results, summarized in Table 1, align with the expectations laid out in Appendix A.2: both Diffsurv and CPL
4

Differentiable sorting for censored time-to-event data.

References

Tests. JAMA, 247(18):2543–2546, May 1982. ISSN
0098-7484. doi: 10.1001/jama.1982.03320430047030.
URL https://doi.org/10.1001/jama.1982.
03320430047030.

Antolini, L., Boracchi, P., and Biganzoli, E.
Atime-dependent discrimination index for survival
data.
Statistics in Medicine, 24(24):3927–3944,
2005. ISSN 1097-0258. doi: 10.1002/sim.2427.
https://onlinelibrary.wiley.
URL
eprint:com/doi/abs/10.1002/sim.2427.
https://onlinelibrary.wiley.com/doi/pdf/10.1002/sim.2427.

Johnson, A. E. W., Pollard, T. J., Berkowitz, S. J.,
Greenbaum, N. R., Lungren, M. P., Deng, C.-y.,Mark, R. G., and Horng, S. MIMIC-CXR, a deidentified publicly available database of chest radiographs with free-text reports. Scientific Data, 6(1):317, December 2019. ISSN 2052-4463. doi: 10.1038/
s41597-019-0322-0. URL https://www.nature.
com/articles/s41597-019-0322-0. Number:
1 Publisher: Nature Publishing Group.

Blondel, M., Teboul, O., Berthet, Q., and Djolonga, J. Fastdifferentiable sorting and ranking. In International Conference on Machine Learning, pp. 950–959. PMLR, 2020.
Cao, Z., Qin, T., Liu, T.-Y., Tsai, M.-F., and Li, H. Learning to rank: from pairwise approach to listwise approach. In Proceedings of the 24th international conference on Machine learning, pp. 129–136, Corvalis Oregon
USA, June 2007. ACM. ISBN 978-1-59593-793-3. doi:
10.1145/1273496.1273513. URL https://dl.acm.
org/doi/10.1145/1273496.1273513.

Katzman, J. L., Shaham, U., Cloninger, A., Bates, J.,
Jiang, T., and Kluger, Y. DeepSurv: personalized treatment recommender system using a Cox proportional
hazards deep neural network. BMC Medical Research
Methodology, 18(1):24, February 2018. ISSN 14712288. doi: 10.1186/s12874-018-0482-1. URL https://doi.org/10.1186/s12874-018-0482-1.

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

Kvamme, H. and Borgan, O. The brier score under administrative censoring: Problems and a solution. Journal of
Machine Learning Research, 24:1–26, 2023. Submitted
12/19; Revised 11/22; Published 1/23.

Cuturi, M., Teboul, O., and Vert, J.-P. Differentiable RankKvamme, H., Borgan, Ø., and Scheel, I. Time-to-event preing and Sorting using Optimal Transport. In Advances in
diction with neural networks and cox regression. Journal
Neural Information Processing Systems, volume 32. Curof Machine Learning Research, 20:1–30, 2019.
ran Associates, Inc., 2019. URL https://papers.
nips.cc/paper_files/paper/2019/hash/
Lee, C., Zame, W., Yoon, J., and Schaar, M. V. D. DeepHit:
d8c24ca8f23c562a5600876ca2a550ce-Abstract.
A Deep Learning Approach to Survival Analysis withhtml.
Competing Risks. Proceedings of the AAAI Conferenceon Artificial Intelligence, pp. 8, 2018.
Davidson-Pilon, C.
lifelines: survival analysis in
Python. Journal of Open Source Software, 4(40):
Netzer, Y., Wang, T., Coates, A., Bissacco, A., Wu,
1317, August 2019. ISSN 2475-9066. doi: 10.21105/
B., and Ng, A. Y.
Reading digits in naturaljoss.01317.
URL https://joss.theoj.org/
images with unsupervised feature learning.
Inpapers/10.21105/joss.01317.
NIPS Workshop on Deep Learning and Unsupervised Feature Learning 2011, 2011. URL http:Goldstein, L. and Langholz, B. Asymptotic Theory for//ufldl.stanford.edu/housenumbers/
Nested Case-Control Sampling in the Cox Regression
nips2011_housenumbers.pdf.
Model. The Annals of Statistics, 20(4):1903–1928,
December 1992. ISSN 0090-5364, 2168-8966. doi:Petersen, F. Learning with Differentiable Algorithms. PhD
10.1214/aos/1176348895. Publisher: Institute of Mathethesis, University of Konstanz, 2022.
matical Statistics.
Petersen, F., Borgelt, C., Kuehne, H., and Deussen, O.
Grover, A., Wang, E., Zweig, A., and Ermon, S. Stochastic
Differentiable sorting networks for scalable sorting andoptimization of sorting networks via continuous relaxranking supervision. In Meila, M. and Zhang, T.
ations. In 7th International Conference on Learning(eds.),
Proceedings of the 38th International ConferRepresentations, ICLR 2019, New Orleans, LA, USA,ence
on
Machine Learning, ICML 2021, 18-24 JulyMay 6-9, 2019. OpenReview.net, 2019. URL https:
2021,
VirtualEvent, volume 139 of Proceedings of Ma//openreview.net/forum?id=H1eSS3CcKX.
chine Learning Research, pp. 8546–8555. PMLR, 2021.
Harrell, Jr, F. E., Califf, R. M., Pryor, D. B., Lee, K. L.,
URL http://proceedings.mlr.press/v139/
and Rosati, R. A. Evaluating the Yield of Medical
petersen21a.html.
5Differentiable sorting for censored time-to-event data.

Petersen, F., Borgelt, C., Kühne, H., and Deussen, O. Monotonic differentiable sorting networks. In Proc 10th International Conference on Learning Representations (ICLR
2022), 2022a.
Petersen, F., Kuehne, H., Borgelt, C., and Deussen, O. Differentiable top-k classification learning. In International
Conference on Machine Learning, pp. 17656–17668.
PMLR, 2022b.
Pölsterl, S. scikit-survival: A library for time-to-event analysis built on top of scikit-learn. Journal of Machine
Learning Research, 21(212):1–6, 2020. URL http://jmlr.org/papers/v21/20-729.html.
Pölsterl, S. Survival Analysis for Deep Learning, July 2019.
URL
https://k-d-w.org/blog/2019/07/
survival-analysis-for-deep-learning/.
Raykar, V. C., Steck, H., Krishnapuram, B., Dehingoberije, C., and Lambin, P. On Ranking in Survival
Analysis: Bounds on the Concordance Index. In
Advances in Neural Information Processing Systems,
volume 20. Curran Associates, Inc., 2007. URL https://papers.nips.cc/paper/2007/hash/
33e8075e9970de0cfea955afd4644bb2-Abstract.
html.
Tan, M. and Le, Q. Efficientnet: Rethinking model scaling for convolutional neural networks. In International
conference on machine learning, pp. 6105–6114. PMLR,
2019.
Yang, X., Abraham, L., Kim, S., Smirnov, P., Ruan, F.,
Haibe-Kains, B., and Tibshirani, R. Fastcph: Efficientsurvival analysis for neural networks. In NeurIPS 2022
Workshop on Learning from Time Series for Health, 2022.

6

Differentiable sorting for censored time-to-event data.

A. Survival Analysis and its Relation to Ranking
A dataset with censored event times is summarized as D = {ti , xi , δi }Ni=1 , where N is the total number of patients. Fora patient i, the time-to-event ti is the minimum of the true survival time t∗i and the censoring time c∗i , with δi indicating
whether an event (t∗i ≤ c∗i , δi = 1) or censoring (t∗i > c∗i , δi = 0) was observed. Covariates are xi ∈ Rd representing a
1-dimensional vector of size d but the methods discussed here also generalise to higher dimensional tensors such as image
data.
The widely-used method for addressing censoring in survival analysis is the Cox Partial Likelihood (CPL) model, introduced
by Cox (1972). The CPL is designed to maximize the following general form:Y

L(θ) :=

i:δi =1

Pfθ (xi ),j:tj >ti fθ (xj )(5)where fθ is the hazard function, a score prediction function estimating the probability of an event at a particular time, given
input features xi . The product only includes uncensored patients, whereas the denominator term also includes censoredpatients with tj > ti .
Reflecting the structure of survival data, the Cox Partial Likelihood (CPL) model compares individuals still ”at risk” at eachtime 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.
Extensions of the Cox model, like (Katzman et al., 2018) and (Kvamme et al., 2019), have modified fθ = exp(θ · xi ) torelax the linear covariate interaction and proportional hazards assumptions. Introducing neural networks hθ to handle thenon-linearity, adjusting fθ to be fθ = exp(hθ (xi )), and to manage non-proportional hazards, they set fθ = exp(hθ (xi , ti )).
A.1. Pairwise Independence
Both of these previous works note that the risk set R = {j : tj > ti } is intractable for deep learning applications as it
considers all comparable patients. To mitigate memory constraints, we can sample a fixed-size risk set, denoted as R̃, suchthat |R̃| = n < N . Kvamme et al. (2019) go further, arguing it is reasonable to take a constant sample size of 1 and include
the individual i in the risk set (such that n=2). This leads to the simplified loss of the form
L(θ) =

Yi:δi =1fθ (xi ), j ∈ R \ {i}.
fθ (xi ) + fθ (xj )(6)Further, take the mean log partial likelihood to beloss =1 X
log(1 + exp[hθ (xj ) − hθ (xi )]), j ∈ R \ {i},ne(7)i:δi =1

where ne is the number of non-censored events. In this simplified form, it can be seen that the partial likelihood only
considers the pairwise relative ordering or ranking of survival times.
The concordance index or c-index Harrell et al. (1982) is a commonly used as an evaluation for survival analysis methods
and is a generalization of the Area Under the Receiver Operating Characteristic Curve (AUROC) that handles right-censoreddata. It is defined as
1 X
c-index :=1(f (xi ) < f (xj )), j ∈ R \ {i}.
(8)ni:δi =1

Raykar et al. (2007) first showed that the Cox’s partial likelihood is approximately equivalent to maximizing the concordanceindex or c-index and that closer bounds can be found by minimizing the general ranking loss, with acceptable pairs
A = {(i, j) : δi = 1 ∧ tj > ti } andranking-loss :=1 X
φ(fθ (xi ) − fθ (xj )),|A|(i,j)∈A

7(9)

Differentiable sorting for censored time-to-event data.

where φ is a function that relaxes the non-differentiable 1 of the c-index. From Equation (7) it can be seen that φ : x →− log(1 + exp(−x)) = log(σ(x)). Here, we have shown that the simplifications to the partial likelihood made by Kvamme
et al. (2019) are equivalent to using the log-sigmoid ranking loss.
The key difference between ranking and partial likelihood losses comes when considering the assumption that it is reasonable
to take a constant sample size of 2 (one pair in the risk set) in the partial likelihood. This effectively introduces theassumption that each pair (i, j) is independent of any other pair. However, this assumption seems puzzling given the inherent
transitivity of ranking (if i > j and j > k then i > k).
A.2. Differentiable Sorting Networks Relation to Ranking and Partial Likelihood
It is possible to directly relate differentiable sorting networks with ranking losses and partial likelihood. Expanding out thecross entropy loss out we find
!
n
n
X
1X
qic log(pic ) ,
(10)
L=
n i=1
c=1

where qic = 1 only when i is the true rank otherwise 0. Each pic is always a function of the difference in pairs of inputs xi
and xj . This is complicated by the products of intermediate values a introduced by the sorting network but denoted as
pic =

n
Y

σ(ai − aj )(11)(ai ,aj )∈Pl :l=1

where Pl to denotes the set of comparisons to be made at each layer of the sorting network. With n = 2 and β = 1, a sorting
network only requires a single relaxed conditional swap and the loss returns to the same recognisable log-sigmoid ranking
loss in Equation (9), and Cox negative log partial likelihood in Equation (7).
Other relaxation of the step function can also be considered, Petersen et al. (2022a) show that the Cauchy distribution
preserves monotonicity which is desirable for optimization. Given this, we use the Cauchy distribution as our relaxation for
all experiments, where σ : x → π1 arctan(βx) + 21 .

B. Top-K risk prediction
Finally, we demonstrate how the algorithmic supervision of sorting algorithms enables the development of novel methods
in survival analysis, using the example of top-k risk prediction. In practical settings, it is often not necessary to rank all
samples correctly. Rather, it is essential to identify the samples with the highest risk, such as by a healthcare provider, toprioritize care and interventions.
With Diffsurv, top-k risk prediction is straightforward to implement by modifying the loss such that the negative loglikelihood of predicted top-k ranks in P is maximised for individuals with a possible permutation to any top-k rank
according to Qp .
First, let’s denote Tk as the set of values with a possible permutation to a top-k rank, derived from the ground truth possible
permutation matrix Qp :k
XQpij > 0}(12)
Tk = {i|j=1Importantly, due to the uncertainty introduced by censoring, the set of individuals with a possible permutation to a top krank Tk can be arbitrarily large. For example, in case all individuals are censored, Tk is the set of all individuals. Then, thetop-k loss is described as:


k
X
X
Pij  .
log 
Ltop-k = −
(13)
j=1

i∈Tk

This loss is minimized when the model correctly predicts a top-k rank for the indices in Tk . This represents the individuals
with possible permutations to the top-k highest risk ranks. Importantly, this loss function is optimized for the identificationof potential top-k high-risk individuals, without considering the specific order within these top-k ranks. To establish abaseline for comparison with Diffsurv’s top-k risk prediction, we also introduce two variants of the Cox Partial Likelihood
8

Differentiable sorting for censored time-to-event data.

method. In the first variant, we adjust the likelihood term so that the product considers only the set of patients who have apotential permutation to a top-k rank, according to the matrix of possible permutations Qp :Yf (x )P θ i
LCPL I =(14)
j:tj >ti fθ (xj )i:i∈TkIn the second variant, we further limit the set of patients to those who have both experienced an event and have a possible
permutation to a top-k rank:Yfθ (xi )P
LCPL II =.
(15)
j:Tj >Ti fθ (xj )i:δi =1∧i∈Tk

Note that the denominator term is unchanged in both variants and considers only comparable pairs and includes censoredpatients Tj > T i. Evaluation of top-k risk prediction is also complicated by the uncertainty due to censoring. For Diffsurv
and both variants of the Cox Partial Likelihood, we can first define the set of individuals predicted to be within the top-k
highest risk:
Pk = {i|rank(fθ (xi )) ≥ k}(16)

We can then define the fraction of how many of these individuals are in the set of possible top-k highest ranks Tk to evaluate
the top-k risk prediction performance:|Pk ∩ Tk |top-k-score =
(17)
|Pk |
In this work this only experiment with Odd Even and Bitonic sorting networks. Yet, it is worth highlighting recent
developments which utilize specialized sorting networks, such as splitter selection networks as in Petersen et al. (2022b).

C. Non-proportional Hazards
Our current implementation of Diffsurv operates under the proportional hazards assumption. While this may not fully
capture the intricacies of some survival analysis problems—particularly those involving non-proportional hazards—it does
not necessarily limit the model’s effectiveness in scenarios where the goal is to assess cumulative risk from a fixed index
date or from the date of an imaging study. This aligns with the necessity of time-dependent modifications to the C-index for
non-proportional models as indicated by Antolini et al. (2005).
If we are primarily interested in understanding the cumulative hazard of an event occurring rather than tracking changes
in the hazard over time, the assumption of proportional hazards becomes less pivotal. As such, Diffsurv and CPL remainvaluable tools for these cases.
Despite the current limitation of Diffsurv to proportional hazards, it is conceivable that an extension to accommodate
non-proportional hazards could be developed, similar to adaptations made for the CPL method.
For instance, as we briefly mentioned earlier, continuous-time extensions of partial likelihood can be used to enable nonproportional hazards (Kvamme et al., 2019). Implemented by directly modeling temporal covariates as fθ = exp(hθ (xi , Ti )).
Another class of methods focuses on discretizing the time-to-event variable and modeling the probability mass function(PMF) of event times. For instance, the DeepHit model Lee et al. (2018) employs a neural network architecture to learn therelationships between input features and discretized time-to-event outcomes. Time discretization facilitates modeling of
non-proportional hazards but introduces two significant challenges: 1) sensitivity to the choice of time intervals, which canaffect the model’s accuracy and interpretability, and 2) increased computational complexity, as predictions must be made foreach time interval. These models can be computationally expensive, especially for deep learning-based models like DeepHit,making them less suitable for high-dimensional and large-scale datasets, such as the imaging dataset used in this study.
Several future work proposals arise from these observations. First, differentiable sorting could explore the approach of
directly modeling temporal covariates, resulting in a time-parameterized predicted permutation matrix. Second, extending
Diffsurv to discrete time could be achieved by parameterizing a predicted permutation matrix for each time discretization.

D. Experiments
Baselines: We compare Diffsurv primarily against Cox’s Partial Likelihood, using the Ranked List implementation frompycox (Kvamme et al., 2019). We include Efron and Breslow estimates of CPL from Yang et al. (2022) for survSVHN.
9fλ = λe−λx

λ = 1.5λ=1
λ = 0.5

h1 =16∼ t∗1=32.5
t∗2 =12.9

h2 =28
Images from SVHNt∗0 =61.2

Parameterize time function
with log normal house numbers

Sample fromtime functionc∗0 =23.2c∗1 =29.3c∗2 =10.1
Sample censoringtime∗
ti = min(t∗i , ci )h0 =1c∗ ∼ Uniform(0, t∗ )Differentiable sorting for censored time-to-event data.

t0 =23.2
t1 =32.5
t2 =12.9
Replace 30% oftrue times withcensoring time

Figure 3. Visual abstract of the survSVHN dataset.

For smaller datasets, we add non-deep learning baselines: Lifelines’ Cox Regression (Davidson-Pilon, 2019) and sksurv’s
Random Survival Forests (Pölsterl, 2020). We do not compare with DeepHit (Lee et al., 2018), since we do not model
non-proportional hazards. For an extended discussion see Appendix C.
Network Architectures: For both CPL and Diffsurv, we use a fixed neural network architecture depending on the dataset.
Small datasets utilize a single-layer neural network, survSVHN uses a ConvNet architecture as in Petersen et al. (2021), andMIMIC IV CXR uses EfficientNet-B0 (Tan & Le, 2019).
Training and Evaluation: We employ AdamW for optimization. Validation approach varies: for smaller datasets, we apply
nested 5-fold cross-validation, while for imaging datasets we use train:val:test splits. We performed hyperparameter tuning
for learning rate, weight decay, batch size, and risk set size. In the case of imaging datasets, we maintained fixed values
for learning rate and weight decay. As in Petersen et al. (2021), we determine steepness as a function of the risk set sizen, β = 2n for odd-even and β = (log2 n)(1 + log2 n) for bitonic. The type of sorting network can either be bitonic or
odd-even and is determined during hyperparameter tuning. Further details on the experimental setup, including computetime, are provided in Appendix E and at anon@git.com.

E. Training and evaluation
E.1. Datasets and Preprocessing
As in Goldstein & Langholz (1992) and Kvamme et al. (2019), we ensure that each risk set contains a valid risk set bysampling controls for a given case. Each batch consists of a number of risk sets such that the input data has shape (batch
size, risk set size, covariate shape).
We provide an additional description of each small realworld dataset:• FLCHAIN dataset: A dataset containing information on patients with monoclonal gammopathy of undeterminedsignificance (MGUS), focusing on serum free light chain (FLC) levels to study their prognostic significance in predicting
disease progression. Number of covariates: 8.
• NWTCO dataset: A dataset from a series of clinical trials on the treatment and outcomes of children with Wilms’
tumor, a type of kidney cancer, aiming to improve understanding of tumor biology and optimize treatment strategies.
Number of covariates: 9.
• SUPPORT dataset: A dataset from a multi-center study investigating the prognosis and treatment preferences ofseriously ill hospitalized adults, with the goal of improving end-of-life care and informing decision-making processes.
Number of covariates: 22.
• METABRIC dataset: A dataset comprising genomic and clinical data on breast cancer patients, focused on uncoveringnovel molecular subtypes for more precise prognostication and personalized treatment strategies. Number of covariates:
9.
Covariate preprocessing follows (Kvamme et al., 2019), and includes standardising continuous variables and one-hot
encoding categorical variables.
MIMIC IV CXR: For the survival task, we extract death events from the MIMIC IV dataset. This is done by merging
the data on ”subject id”, ”study id”, and ”dicom id” with the patient table from MIMIC IV, and on ”subject id” with the10

Differentiable sorting for censored time-to-event data.

admission table. For patients without a recorded date of death, censoring dates are be determined as 1 year after the lastrecorded discharge date for each patient. It is important to note that the time-shift anonymization protocol employed in the
MIMIC-IV dataset, which, while preserving patient confidentiality, maintains the internal consistency of time intervals foreach patient’s data.
We exclude 29,345 images without any matches in the MIMIC IV patient table, 19,337 images taken after the latest founddischarge date and 55 images taken after a recorded date of death. Time to event is calculated as the number of days fromthe image study date to either date of death or the censoring date. For MIMIC IV CXR, images undergo several standardtransformations: 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 normalizedusing ImageNet’s mean and standard deviation values. Finally, the train:val:test split of 8:1:1 is done at the patient level
ensuring no images from a patient in the test set was found in the training data.
survSVHN: In this semi-synthetic dataset, each house number parameterizes a beta-exponential time function from which we
can sample survival times. The risk parameters or hazards λi are calculated as the logarithm of house numbers, standardized
and scaled for a mean survival time of 30. We introduce censoring by randomly selecting 30% of house numbers andreplacing true times with values sampled uniformly between (0, ti ] (See Figure 3). A beta distribution used to sample fromthe house number parameterized exponential time function uses a fixed value of 500 for both shape parameters. We follow
Petersen et al. (2021) by cropping the centered multi-digit numbers with a boundary of 30%, resizing it to a resolution of
64×64, and then selecting 54 × 54 pixels at a random location. For survSVHN the train:val:test split is provided by Netzer
et al. (2011) and is 230,755:5,000:13,068.
E.2. Model Architecture and Hyperparameters
For the smaller real-world datasets, the hazard function fθ is small fixed Multi-layer Perceptron network with 1 hidden layer
and 64 hidden nodes. We also apply a fixed dropout rate of 0.1. Learning rate, weight decay, batch set size and risk set size
were found using a grid search across the possible values in Table 4.

Table 4. Hyperparameter values for small real-world datasets.

Hyperparameter

ValuesLearning rate
Weight decay(Batch size, risk set size)[0.1, 0.01, 0.001, 1e-4][0.1, 0.01, 0.001, 1e-4, 1e-5, 0][(32, 8), (16, 16), (8, 32), (4, 64), (1, 256)]

For imaging datasets, we fix learning rate and weight decay for both CPL and Diffsurv. For both survSVHN and MIMIC IV
CXR, we use a fixed learning rate of 10−4 and weight decay of 10−5 . We also used early stopping with a patience of 20
epochs and a maximum of 100,000 training steps.
survSVHN: As per Petersen et al. (2021), the model consists of four convolutional layers (with a 5x5 kernel size and 32, 64,
128, 256 filters), each followed by ReLU and max-pooling (2x2 stride). The architecture concludes with a fully connected
layer of 64 units, another ReLU, and a one-unit output layer.
MIMIC IV CXR: Here, we use EfficientNet-B0 with an added linear layer for single output. We first train the linear
prediction layer alone for the initial 2,000 steps. After this, we continue training, this time including both the EfficientNet-B0
and the linear prediction layer.
For the results in Table 1, we keep a fixed batch size of 100. We also provide a comparison where the number of values isfixed in each batch in Table 6.
Note that during evaluation, the sorting network is not used since we only need to evaluate the ranks of the trained risk scores.
Similarly, case-control sampling is not used. We measure the ranking performance of the models using the concordanceindex.
Further implementation details and the best hyperparameters for each dataset are provided at anon@git.com.
11Differentiable sorting for censored time-to-event data.

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

F. Sorting Networks
There are multiple different types of sorting networks each with varying complexity. The ability to implement networks
with the divide-and-conquer paradigm allows for sorting networks that scale more efficiently. Examples for Odd-Even and
Bitonic sorting networks with n = 8 are shown in Figure 4 and Figure 5. The latter allows construction of networks withsize complexity O(nlog 2 n) verses the O(n2 ) in Odd-Even networks.
It is worth emphasising these are not neural networks. They are called ”networks” because they are typically represented as
diagrams that show how the items are compared and swapped as they are being sorted. Differentiable sorting networks do
not introduce any additional parameters that need to be updated during optimization.

Figure 4. Example Odd-even sorting network of size 8.

Figure 5. Example Bitonic sorting network of size 8.

G. Additional Results
In Table 5 and Table 6, 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 adecrease 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.

12Differentiable sorting for censored time-to-event data.

Table 5. Additional Results for MIMIC IV CXR. Mean and standard deviation of the C-index for different methods and batch risk set
sizes. Bold indicates a significantly higher result with t-test and p ≤ 0.01.† Most significant across all Batch Size, Risk Set Sizes.

Method64, 2Diffsurv: Bitonic
Diffsurv: Odd-Even
CPL: Ranked List

0.761 (0.001)
0.761 (0.002)
0.760 (0.002)

Batch Size, Risk Set Size
4, 32
16, 8
0.761 (0.000)
0.756 (0.002)
0.755 (0.002)

0.763† (0.001)
0.761 (0.001)
0.758 (0.003)

1, 1280.761 (0.002)
0.749 (0.001)
0.755 (0.002)

Table 6. Additional results for survSVNH keeping the number of events per batch equal. Mean (and standard deviation) over 5 trials withdifferent seeds. Metric is C-index. Bold indicates a significantly higher result with t-test and p ≤ 0.01.

Method512, 2Batch Size, Risk Set Size
128, 832, 328, 128Diffsurv: Odd-Even
Diffsurv: Bitonic

0.934 (0.001)
0.931 (0.001)

0.940 (0.001)
0.942 (0.001)

0.941 (0.001)
0.940 (0.00166)

0.933 (0.002)
0.928 (0.001)

CPL: Breslow
CPL Efron
CPL: Ranked List

0.905 (0.001)
0.904 (0.002)
0.921 (0.001)

0.897 (0.001)
0.898 (0.002)
0.922 (0.003)

0.910 (0.002)
0.909 (0.003)
0.921 (0.001)

0.919 (0.001)
0.918 (0.003)
0.917 (0.003)

13