SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Differentiable sorting for censored time-to-event data.
Andre Vauvelle $^{12*}$ , Benjamin Wild $^{3*}$ , Roland Eils $^{3}$ , Spiros Denaxas $^{1}$
University College London<sup>1</sup>, BenevolentAI<sup>2</sup>, Berlin Institute of Health<sup>3</sup>
{andre.vauvelle.19,s.denaxas}@ucl.ac.uk
{benjamin.wild,roland.eils}@bih-charite.de
Survival analysis is a crucial semi-supervised task in machine learning with significant real-world applications, especially in healthcare. The most common approach to survival analysis, Cox's partial likelihood, can be interpreted as a ranking model optimized on a lower bound of the concordance index. We follow these connections further, with listwise ranking losses that allow for a relaxation of the pairwise independence assumption. Given the inherent transitivity of ranking, we explore differentiable sorting networks as a means to introduce a stronger transitive inductive bias during optimization. Despite their potential, current differentiable sorting methods cannot account for censoring, a crucial aspect of many real-world datasets. We propose a novel method, Diffsurv, to overcome this limitation by extending differentiable sorting methods to handle censored tasks. Diffsurv predicts matrices of possible permutations that accommodate the label uncertainty introduced by censored samples. Our experiments reveal that Diffsurv outperforms established baselines in various simulated and real-world risk prediction scenarios. Furthermore, we demonstrate the algorithmic advantages of Diffsurv by presenting a novel method for top-k risk prediction that surpasses current methods. In conclusion, Diffsurv not only provides a novel framework for survival analysis through differentiable sorting, but also significantly impacts real-world applications by improving risk stratification and offering a methodological foundation for developing predictive models in healthcare and beyond.
Survival analysis plays a pivotal role in many realworld machine learning applications, spanning fields such as reliability engineering, marketing, and insurance, with a particularly significant impact in healthcare. The goal of survival analysis is to predict the time until the occurrence of an event of interest, such as death, based on a set of covariates. In clinical studies, these include demographic variables such as sex and age, but may also encompass more complex data modalities such as medical images.
The concept of censoring is a distinguishing characteristic that sets survival analysis apart from conventional machine learning approaches. Particularly prevalent in observational datasets, it refers to situations where event times remain 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.
Overlooking censoring can skew predictions towards the censoring event, rather than the event of interest. This bias becomes particularly noticeable when the study's endpoint can be inferred from the observed covariates - age being a notable example. In such cases, the predicted event times are likely
to biased towards the censoring event time, thereby neglecting the actual event of interest [Kvamme and Borgan, 2019].
The Cox Proportional Hazards model is widely used for handling 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 based on 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, 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 Lee et al. [2018].
However, this relation operates under the assumption of pairwise independence. This simplification, while practical, can de-emphasize the transitive properties inherent in survival data. As shown by Goldstein and Langholz [1992], larger risk set sizes can lead to more efficient estimators, suggesting potential benefits in considering listwise ranking losses Cao et al. [2007]. These losses optimize over lists of values rather than individual pairs, thereby better capturing the transitive dynamics of the data. Despite the similarities, listwise losses have remained largely unexplored within the field of survival analysis. This could be partly due to an uncertainty around how to handle censoring.
We propose a new approach that takes advantage of recent developments in continuous relaxations of sorting operations, allowing end-to-end training of neural networks with ordering supervision [Grover et al., 2019, Blondel et al., 2020, Petersen et al., 2021]. This method incorporates a sorting algorithm into the network architecture, where the order of the samples is known, but their exact values are unsupervised. With this, we introduce Diffsurv, an extension of differentiable sorting methods that enables end-to-end training of survival ranking models with censored data.
Briefly, our contributions are summarised:
- Our primary contribution is the extension of differentiable sorting methods to account for censoring by introducing the concept of possible permutation matrices. (Section 3.1)
- We empirically demonstrate that our new differentiable sorting method matches or improves risk ranking performance across multiple semi-simulated and real-worlds censored datasets. (Section 4)
- We investigate the role of transitivity in survival analysis and show, through experiments with semi-simulated data, that differentiable sorting networks can benefit from this inherent property of the data. (Section 4)
- We demonstrate that differentiable sorting of censored data enables the development of new methods with practical applications, using the example of end-to-end learning for top-k risk stratification. (Section 3.2)
2 Survival Analysis and its Relation to Ranking
A dataset with censored event times is summarized as  $\mathcal{D} = \{t_i, \boldsymbol{x}_i, \delta_i\}_{i=1}^N$ , where  $N$  is the total number of patients. For a patient  $i$ , the time-to-event  $t_i$  is the minimum of the true survival time  $t_i^*$  and the censoring time  $c_i^*$ , with  $\delta_i$  indicating whether an event ( $t_i^* \leq c_i^*$ ,  $\delta_i = 1$ ) or censoring ( $t_i^* > c_i^*$ ,  $\delta_i = 0$ ) was observed. Covariates are  $\boldsymbol{x}_i \in \mathbb{R}^d$  representing a  $d$ -dimensional vector 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:
$$
\mathcal {L} (\theta) := \prod_ {i: \delta_ {i} = 1} \frac {f _ {\theta} \left(\boldsymbol {x} _ {i}\right)}{\sum_ {j : t _ {j} > t _ {i}} f _ {\theta} \left(\boldsymbol {x} _ {j}\right)}, \tag {1}
$$
36   
37   
38   
39   
40   
41   
42   
43   
44   
45   
46   
47   
48   
49   
50   
51   
52   
53   
54   
55   
56   
57   
58   
59   
60   
61   
62   
63   
64   
65   
66   
67   
68   
69   
70   
71   
72   
73   
74   
75   
76   
77   
78   
79   
80
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 for downstream tasks, such as risk stratification and top-k highest risk prediction, by using the matrix  $Q_{p}$  of possible permutations based on the observed events and censoring.
where  $f_{\theta}$  is the hazard function, a score prediction function estimating the probability of an event at a particular time, given input features  $\pmb{x}_i$ . The product only includes uncensored patients, whereas the denominator term also includes censored patients with  $t_j > t_i$ .
Reflecting the structure of survival data, the Cox Partial Likelihood (CPL) model compares individuals still "at risk" at each time point, similar to a nested case-control study. This directly shapes the likelihood equation in CPL, with the numerator representing the hazard function for the 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_{\theta} = \exp (\theta \cdot \pmb {x}_i)$  to relax the linear covariate interaction and proportional hazards assumptions. Introducing neural networks  $h_\theta$  to handle the non-linearity, adjusting  $f_{\theta}$  to be  $f_{\theta} = \exp (h_{\theta}(\pmb {x}_i))$  , and to manage non-proportional hazards, they set  $f_{\theta} = \exp (h_{\theta}(\pmb {x}_i,t_i))$
2.1 Pairwise Independence
Both of these previous works note that the risk set  $\mathcal{R} = \{j : t_j > t_i\}$  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  $\tilde{\mathcal{R}}$ , such that  $|\tilde{\mathcal{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
$$
\mathcal {L} (\theta) = \prod_ {i: \delta_ {i} = 1} \frac {f _ {\theta} \left(\boldsymbol {x} _ {i}\right)}{f _ {\theta} \left(\boldsymbol {x} _ {i}\right) + f _ {\theta} \left(\boldsymbol {x} _ {J (i)}\right)}, \text {w h e r e} J (i) \in \mathcal {R} \backslash \{i \}. \tag {2}
$$
Further, take the mean log partial likelihood to be
$$
\log = \frac {1}{n _ {e}} \sum_ {i: \delta_ {i} = 1} \log \left(1 + \exp \left[ h _ {\theta} \left(\boldsymbol {x} _ {J (i)}\right) - h _ {\theta} \left(\boldsymbol {x} _ {i}\right) \right]\right), \text {w h e r e} J (i) \in \mathcal {R} \backslash \{i \}, \tag {3}
$$
where  $n_e$  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-censored data. It is defined as
$$
\text {c - i n d e x} := \frac {1}{n} \sum_ {i: \delta_ {i} = 1} \mathbb {1} \left(f \left(\boldsymbol {x} _ {i}\right) <   f \left(\boldsymbol {x} _ {j}\right)\right), j \in \mathcal {R} \backslash \{i \}. \tag {4}
$$
Raykar et al. [2007] first showed that the Cox's partial likelihood is approximately equivalent to maximizing the concordance index or c-index and that closer bounds can be found by minimizing the
81   
82   
83   
84   
85   
86   
87   
88   
89   
90   
91   
92   
93   
94   
95   
96   
97
general ranking loss, with acceptable pairs  $\mathcal{A} = \{(i,j):\delta_i = 1\wedge t_j > t_i\}$  and
$$
\text {r a n k i n g - l o s s} := \frac {1}{| \mathcal {A} |} \sum_ {(i, j) \in \mathcal {A}} \phi \left(f _ {\theta} \left(\boldsymbol {x} _ {i}\right) - f _ {\theta} \left(\boldsymbol {x} _ {j}\right)\right), \tag {5}
$$
where  $\phi$  is a function that relaxes the non-differentiable  $\mathbb{1}$  of the c-index. From Equation 3, it can be seen that  $\phi : x \to -\log(1 + \exp(-x)) = \log(\sigma(x))$ . Here, we have shown that the simplifications to the partial likelihood made by Klamme et al. [2019] are equivalent to using the log-sigmoid ranking loss.
The key difference between ranking and partial likelihood losses comes when considering the assumption that it is reasonable to take a constant sample size of 2 (one pair in the risk set) in the partial likelihood. This effectively introduces the assumption that each pair  $(i,j)$  is independent of any other pair. However, this assumption seems puzzling given the inherent transitivity of ranking (if  $i > j$  and  $j > k$  then  $i > k$ ).
2.2 Listwise Ranking and Differentiable Sorting Networks
[Cao et al., 2007] first proposed the notion of a listwise ranking loss, arguing its benefits in situations where the entire order of items is of importance. This approach treats the ranking problem as a permutation problem, with the aim of learning a model that can provide the optimal permutation of an entire list of items, rather than considering independent pairs. Indeed, this idea is closely related to the partial likelihood from Cox's original work. As pointed out by Wang and Yang [2022], the Top-1 probability method originally proposed by Cao et al. [2007] and extended to ListMLE by Xia et al. [2008] takes the same form as CPL.
More recent works closely related to listwise ranking have begun to explore combining traditional sorting algorithms with differentiable sorting functions [Petersen et al., 2021]. Sorting networks are a family of sorting algorithms that consist of two basic components: wires and conditional swaps [Batcher, 1968, Knuth, 1998]. Wires carry values to be compared at conditional swaps. If one value is bigger than the other then the values carried forward are swapped around. This allows construction of sorting networks that can provably sort a list of values. Conditional swaps are exactly the min and max operators that ensure that with inputs  $\{a,b\}$  and outputs  $a^{*} \leq b^{*}$ ,  $a^{*} = \min(a,b)$  and  $b^{*} = \max(a,b)$ . Examples of odd even and bitonic networks are shown in Appendix C.
In order to train models based on ordering information alone, differences between predicted and true orderings must be backpropagated through sorting algorithms. However, they often require the use of non-differentiable max and min operators. These are analogous to the non-differentiable indicator function that was discussed earlier in the c-index equation 4. Differentiable sorting methods, similar to ranking losses, rely on approximating these operators with smooth alternatives [Grover et al., 2019].
Petersen et al. [2021] propose combining traditional sorting networks and differentiable sorting functions. Consider that when  $a$  asymptotically approaches  $b$ , the transition point where it surpasses  $b$  is non-continuous and therefore non-differentiable. Just as previously shown in the ranking loss, such operations can be made differentiable using the logistic relaxation
$$
\min  _ {\sigma} (a, b) = a \cdot \sigma (b - a) + b \cdot \sigma (a - b) \text {a n d} \max  _ {\sigma} (a, b) = a \cdot \sigma (a - b) + b \cdot \sigma (b - a). \tag {6}
$$
If an inverse temperature parameter  $\beta > 0$  is introduced such that  $\sigma : x \to \frac{1}{1 + e^{-\beta x}}$ , then as  $\beta \to \infty$  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  $\sigma : x \to \frac{1}{\pi} \arctan(\beta x) + \frac{1}{2}$ .
For an input list to be ordered, each layer of the sorting network can be considered an independent permutation matrix  $P_{l}$  with elements given by
$$
P _ {l, i i} = P _ {l, j j} = \sigma \left(a _ {j} - a _ {i}\right) \text {a n d} P _ {l, i j} = P _ {l, j i} = 1 - \sigma \left(a _ {j} - a _ {i}\right), \tag {7}
$$
where  $a$  signifies intermediate values being compared. The first layer is input with  $z_{i} = h_{\theta}(\pmb{x}_{i})$ , each vector of covariates or images being processed independently by the same neural network. The
107   
108   
109   
110   
111   
112   
113   
114   
115   
116   
117   
118   
119   
120   
121   
122   
123   
124   
125   
126   
127   
128   
129   
130   
131   
132   
133   
134   
135   
136   
137   
138   
139   
140   
141   
142   
143   
144   
145   
146   
147   
148   
149   
150
indices being compared at each layer are determined by the sorting network and the final predicted probability matrix is the product of each layer of sorting operations,
$$
\boldsymbol {P} = \left(\prod_ {l = 1} ^ {n} \boldsymbol {P} _ {l} ^ {\intercal}\right) ^ {\intercal}. \tag {8}
$$
Where  $P$ , is the final doubly-stochastic permutation matrix, doubly-stochastic meaning that the rows and columns both sum to 1. It is possible to interpret each element  $P_{ij}$  of the predicted permutation matrix as the predicted probability of permuting from a randomly assigned rank  $i$  to a true rank  $j$ . Finally, we can define a loss by minimizing the cross-entropy between the ground truth orders represented by true permutation matrix  $Q$  and predicted permutation matrix  $P$  as
$$
\mathcal {L} := \sum_ {c = 1} ^ {n} \left(\frac {1}{n} \operatorname {C r o s s E n t r o p y} \left(\boldsymbol {P} _ {c}, \boldsymbol {Q} _ {c}\right)\right), \tag {9}
$$
where  $P_{c}$  and  $Q_{c}$  denote the c-th columns of their respective matrices.
2.3 Differentiable Sorting Networks Relation to Ranking and Partial Likelihood
It is possible to directly relate differentiable sorting networks with ranking losses and partial likelihood. Expanding out the cross entropy loss, we find
$$
\mathcal {L} = \sum_ {c = 1} ^ {n} \left(\frac {1}{n} \sum_ {i = 1} ^ {n} q _ {i c} \log \left(p _ {i c}\right)\right), \tag {10}
$$
where  $q_{ic} = 1$  only when  $i$  is the true rank otherwise 0. Each  $p_{ic}$  is always a function of the difference in pairs of inputs  $x_{i}$  and  $x_{j}$ . This is complicated by the products of intermediate values  $a$  introduced by the sorting network but denoted as
$$
p _ {i c} = \prod_ {\left(a _ {i}, a _ {j}\right) \in \mathcal {P} _ {l}: l = 1} ^ {n} \sigma \left(a _ {i} - a _ {j}\right) \tag {11}
$$
where  $\mathcal{P}_l$  to denotes the set of comparisons to be made at each layer of the sorting network. With  $n = 2$  and  $\beta = 1$ , a sorting network only requires a single relaxed conditional swap and the loss returns to the same recognisable log-sigmoid ranking loss in Equation 5, and Cox negative log partial likelihood in Equation 3.
3.1 Diffsurv: Handling Censoring with Possible Permutation Matrices
For risk sets of size 2, given proper case-control sampling, it will always be possible to define a single ground truth permutation matrix  $\mathbf{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  $\mathbf{Q}$ . We cannot determine the exact rank of patients who are censored before another who experienced an event. It is only possible to know the range of possible ranks to which a patient should belong. In Figure 2, we provide an illustration demonstrating the possible ranks for a number of censored and uncensored events.
Though we no longer have access to a single permutation matrix, we may instead consider the set of all possible permutation matrices,  $\mathcal{Q} = \{Q_1, Q_1, \ldots, Q_\kappa\}$ . In the best case, all values are uncensored and  $|\mathcal{Q}| = 1$  and in the worse case, when all patients are censored  $|\mathcal{Q}| = n!$ . Our primary contribution is to extend differentiable sorting methods to censored ranks by discriminating between possible and impossible permutations.
We introduce a more computationally tractable representation of  $\mathcal{Q}$  by defining the possible permutation matrix,  $Q_{p}$ , which is the element-wise maximum of every permutation in  $\mathcal{Q}$ ,
$$
\boldsymbol {Q} _ {p i j} = \max  \left\{\boldsymbol {Q} _ {1 i j}, \boldsymbol {Q} _ {2 i j}, \dots , \boldsymbol {Q} _ {\kappa i j} \right\}. \tag {12}
$$
For survival analysis, it is possible to determine  $Q_{p}$  in linear time given a sorted list of event times  $t_{i}$  and event indicators  $\delta_{i}$ . We will consider higher ranks to correspond with a smaller time-to-event. Let us consider two scenarios:
153   
154   
155   
156   
157   
158   
159   
160   
161   
162   
163   
164   
165   
166   
167   
168   
169   
170   
171   
172   
173   
174   
175   
176   
177   
178   
179   
180   
181   
182   
183   
184   
185   
186   
187
Figure 2: For an example case (a) with two events  $(\odot, e_1$  and  $e_5)$  and multiple censored samples  $(\odot, c_1, c_3, c_4, c_6, c_7)$  the uncertainty in the possible permuted rankings (b) due to censoring is taken into account to derive the possible permutation matrix  $Q_p(\mathbf{c})$ .
1. For a right-censored individual  $i$  (i.e.,  $\delta_i = 0$ ), the possible ranks must be lower than the ranks of preceding uncensored events. We can express the set of possible ranks  $\mathcal{R}_i$  as:
$$
\mathcal {R} _ {i} = \{r \mid r <   \operatorname {r a n k} (j), \forall j: \delta_ {j} = 1, t _ {j} <   t _ {i} \} \tag {13}
$$
This set includes all ranks  $r$  that are less than the rank of any uncensored patient  $j$  with an event time  $t_j$  preceding the censoring time  $t_i$  of the individual  $i$ .
2. For an uncensored individual  $i$  (i.e.,  $\delta_i = 1$ ), the possible rank must be lower than all preceding uncensored events and higher than the ranks of all subsequent events. We can define the set of possible ranks as:
$$
\mathcal {R} _ {i} = \{r \mid r <   \operatorname {r a n k} (j), \forall j: \delta_ {j} = 1, t _ {j} <   t _ {i} \} \cap \{r \mid r > \operatorname {r a n k} (j), \forall j: t _ {j} > t _ {i} \} \tag {14}
$$
This set includes all ranks  $r$  that satisfy both conditions: being less than the rank of any preceding uncensored patient  $j$  with  $t_j < t_i$  and greater than the rank of any subsequent patient  $j$  with  $t_j > t_i$ .
With these observations, it's straightforward to construct  $Q_{p}$ . If it's feasible for patient  $i$  to permute to rank  $j$ , i.e.  $j \in \mathcal{R}_i$ , then  $Q_{pij} = 1$ , otherwise  $Q_{pij} = 0$ . See Figure 2 (c) for a visual representation of  $Q_{p}$ .
Given the possible permutation matrix  $Q_p$  and the predicted permutation matrix  $P$ , the vector of probabilities  $p$  of a value being ranked within the set of possible ranks can be computed. Although the ground truth probabilities are unknown, the range of possible ranks is known, and the model can be optimized to maximize the sum of the predicted probabilities of all possible ranks for each sample. Noted here as the column-sum of the element-wise product  $\circ$ , between  $Q_p$  and  $P$ .
$$
\boldsymbol {p} = \sum_ {j = 1} ^ {n} \left(\mathbf {Q} _ {p} \circ \mathbf {P}\right) _ {i, j}. \tag {15}
$$
The binary cross-entropy loss can then be easily applied
$$
\mathcal {L} = \sum_ {i = 1} ^ {n} - y _ {i} \log \left(p _ {i}\right) - \left(1 - y _ {i}\right) \log \left(1 - p _ {i}\right) \tag {16}
$$
where  $y_{i}$  indicates whether set of predicted ranks is possible or impossible.
Equation 15 accounts for the potential challenges of incorporating right-censored samples. The binary cross-entropy remains identical whether the model predicts uniform probability for all possible ranks or concentrates the probability mass on a single rank possible rank.
The introduction of the possible permutation matrix can be used in conjunction with any differentiable sorting method that outputs a doubly-stochastic permutation matrix. This includes methods such
188   
189   
190   
191   
192   
193   
194
195   
196   
197   
198   
199   
200   
201   
202   
203   
204   
205
as SinkhornSort from Cuturi et al. [2019]. Though, in this paper, we will restrict our focus to the discussed differentiable sorting networks. We refer to the use of differentiable sorting networks and the possible permutation matrix as Diffsurv.
3.2 Top-K risk prediction