SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Finally, we demonstrate how the algorithmic supervision of sorting algorithms enables the development of novel methods in survival analysis, using the example of top-k risk prediction. In practical settings, it is often not necessary to rank all samples correctly. Rather, it is essential to identify the samples with the highest risk, such as by a healthcare provider, to prioritize care and interventions.
With Diffsurv, top-k risk prediction is straightforward to implement by modifying the loss such that the negative log-likelihood of predicted top-k ranks in  $\pmb{P}$  is maximised for individuals with a possible permutation to any top-k rank according to  $Q_{p}$ .
First, let's denote  $\mathcal{T}_k$  as the set of values with a possible permutation to a top-k rank, derived from the ground truth possible permutation matrix  $Q_{p}$ :
$$
\mathcal {T} _ {k} = \{i | \sum_ {j = 1} ^ {k} Q _ {p i j} > 0 \} \tag {17}
$$
Importantly, due to the uncertainty introduced by censoring, the set of individuals with a possible permutation to a top  $k$  rank  $\mathcal{T}_k$  can be arbitrarily large. For example, in case all individuals are censored,  $\mathcal{T}_k$  is the set of all individuals. Then, the top-k loss is described as:
$$
\mathcal {L} _ {\text {t o p - k}} = - \sum_ {i \in \mathcal {T} _ {k}} \log \left(\sum_ {j = 1} ^ {k} P _ {i j}\right). \tag {18}
$$
This loss is minimized when the model correctly predicts a top-k rank for the indices in  $\mathcal{T}_k$ . This represents the individuals with possible permutations to the top-k highest risk ranks. Importantly, this loss function is optimized for the identification of potential top-k high-risk individuals, without considering the specific order within these top-k ranks. To establish a baseline for comparison with Diffsurv's top-k risk prediction, we also introduce two variants of the Cox Partial Likelihood method. In the first variant, we adjust the likelihood term so that the product considers only the set of patients who have a potential permutation to a top-k rank, according to the matrix of possible permutations  $Q_{p}$ :
$$
\mathcal {L} _ {\mathrm {C P L} - \mathrm {I}} = \prod_ {i: i \in \mathcal {T} _ {k}} \frac {f _ {\theta} \left(\boldsymbol {x} _ {i}\right)}{\sum_ {j : t _ {j} > t _ {i}} f _ {\theta} \left(\boldsymbol {x} _ {j}\right)} \tag {19}
$$
In the second variant, we further limit the set of patients to those who have both experienced an event and have a possible permutation to a top-k rank:
$$
\mathcal {L} _ {\mathrm {C P L} _ {\mathrm {I I}}} = \prod_ {i: \delta_ {i} = 1 \wedge i \in \mathcal {T} _ {k}} \frac {f _ {\theta} \left(\boldsymbol {x} _ {i}\right)}{\sum_ {j : T _ {j} > T _ {i}} f _ {\theta} \left(\boldsymbol {x} _ {j}\right)}. \tag {20}
$$
Note that the denominator term is unchanged in both variants and considers only comparable pairs and includes censored patients  $T_{j} > 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:
$$
\mathcal {P} _ {k} = \{i | \operatorname {r a n k} \left(f _ {\theta} \left(x _ {i}\right)\right) \geq k \} \tag {21}
$$
We can then define the fraction of how many of these individuals are in the set of possible top-k highest ranks  $\mathcal{T}_k$  to evaluate the top-k risk prediction performance:
$$
\text {t o p - k - s c o r e} = \frac {\left| \mathcal {P} _ {k} \cap \mathcal {T} _ {k} \right|}{\left| \mathcal {P} _ {k} \right|} \tag {22}
$$
In our experiments, we aim to assess the performance of Diffsurv and compare it against the conventional Cox Partial Likelihood (CPL) methods. Initially, we focus on confirming the importance
213   
214   
215   
216   
217   
218   
219   
220   
221   
222   
223   
224   
225   
226   
227   
228   
229   
230   
231   
232   
233   
234   
235   
236   
237   
238   
239   
240   
241   
242   
243   
244   
245
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.
Baselines: We compare Diffsurv primarily against Cox's Partial Likelihood, using the Ranked List implementation from pycox [Kvamme et al., 2019]. We include Efron and Breslow estimates of CPL from Yang et al. [2022] for survSVHN. For smaller datasets, we add non-deep learning baselines: Lifelines' Cox Regression [Davidson-Pilon, 2019] and sksurv's Random Survival Forests [Pölsterl, 2020]. We do not compare with DeepHit [Lee et al., 2018] since we do not model non-proportional hazards. For an extended discussion, see Appendix A.
Network Architectures: For both CPL and Diffsurv, we use a fixed neural network architecture depending on the dataset. Small datasets utilize a single-layer neural network, survSVHN uses a ConvNet architecture as in Petersen et al. [2021], and MIMIC IV CXR uses EfficientNet-B0 [Tan and Le, 2020].
Training and Evaluation: We employ AdamW for optimization. Validation approach varies: for smaller datasets, we apply nested 5-fold cross-validation, while for imaging datasets we use train:val:test splits. We performed hyperparameter tuning for learning rate, weight decay, batch size, and risk set size. In the case of imaging datasets, we maintained fixed values for learning rate and weight decay. As in Petersen et al. [2021], we determine steepness as a function of the risk set size  $n$ ,  $\beta = 2n$  for odd-even and  $\beta = (\log_2 n)(1 + \log_2 n)$  for bitonic. The type of sorting network can either be bitonic or odd-even and is determined during hyperparameter tuning. Further details on the experimental setup, including compute time, are provided in Appendix B and at https://github.com/andre-vauvelle/diffsurv.
Semi-synthetic survSVHN: Based on the Street View House Numbers (SVHN) dataset [Netzer et al., 2011], we simulate survival times akin to survMNIST Pölsterl [2019]. The increased complexity of SVHN over MNIST offers a testbed which is better able to discern the performance differences between methods. Each house number parameterizes a beta-exponential time function for survival times. Risk parameters or hazards  $\lambda_{i}$  are calculated as the logarithm of house numbers, standardized and scaled for a mean survival time of 30. We introduce censoring by randomly selecting  $30\%$  of house numbers and replacing true times with values sampled uniformly between  $(0,t_i]$  (See Figure 3).
We can examine the implications of inherent transitivity within the data. Instead of parameterizing a time function based on unique hazards derived from house numbers, we group  $\lambda_{i}$  into distinct hazard quantiles. Each quantile encompasses a set of house numbers associated with a similar hazard level. We then calculate the transitivity ratio, defined as  $\frac{\# \text{of transitive triplets}}{\# \text{of triplets}}$ , where a sampled triplet is considered transitive if  $(\lambda_{i} > \lambda_{j} > \lambda_{k})$ .
This methodology provides us with a means to control the degree of transitivity in our data. At one extreme, we might categorize data into only two groups, representing the lower and upper halves of house numbers, which results in a transitivity ratio of 0. At the other extreme, each house number could constitute its own unique category (indicated by  $\infty$ ), leading to high transitivity.
Our results, summarized in Table 1, align with the expectations laid out in Section 2.3: both Diffsurv and CPL methods perform similarly when the risk set size is at its minimum ( $n = 2$ ). However, with the expansion of the risk set size, the performance of the two methods diverges, with Diffsurv consistently outperforming CPL. Table 2 sheds light on a potential reason for this divergence. As
248   
249   
250   
251   
252   
253   
254   
255   
256   
257   
258   
259   
260   
261   
262   
263   
264   
265   
266   
267   
268   
269   
270   
271   
272   
273   
274   
275   
276   
277   
278   
279   
280   
281   
282   
283   
284   
285   
286   
287   
288   
289
Table 1: Results for training on survSVHN with increasing risk set size. Mean (standard deviation) c-index over 5 trails with different seeds.  ${}^{ \dagger  }$  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)	.941 (.002)
Cox Partial Likelihood	.913 (.002)	.925 (.002)	.931 (.002)	.933 (.002)	.930 (.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 \leq 0.01$ ). Restricted to a fixed batch size and risk set size of 32.
Number of Quantiles	2	4	8	16	32	64	128	∞
Transitivity Ratio	.0	.374	.657	.819	.908	.954	.975	.991
Diffsurv: Bitonic	.643	.803	.882	.922	.933	.939	.939	.939
Diffsurv: Odd Even	.646	.802	.883	.923	.935	.939	.940	.941
CPL: Ranked List	.651	.803	.880	.909	.916	.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 effectiveness in encapsulating this inherent property. Despite the strong performance of Diffsurv, the C-index for the ground truth risks is 0.980, which is still far above 0.943 for Diffsurv, highlighting the challenging nature of the survSVHN dataset.
Real-world datasets: We assess our methods on several public datasets: Four small, popular 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 B.1.
The results presented in Table 3 demonstrate that Diffsurv achieves equal to or better performance on all datasets analyzed. Additionally, when Diffsurv is optimized for predicting the top  $10\%$  of highest-risk individuals, it matches or outperforms Cox's partial likelihood on the real-world datasets.
Diffsurv introduces a new perspective in survival analysis with censored data, highlighting the relations between survival analysis and the listwise ranking. Our experiments show the effectiveness of differentiable sorting methods for improving survival analysis predictions. Notably, Diffsurv matches or surpasses the performance of the established CPL methods across all examined datasets.
Crucially, Diffsurv sheds light on the importance of transitivity in ranking and survival data, revealing that methods sensitive to this inherent property, such as Diffsurv, show improved performance over those that are not. This insight underscores the value of a listwise approach in dealing with survival data and encourages further exploration for methods that promote a transitive inductive bias.
Moreover, Diffsurv provides a foundation for the development of innovative methods, including the top-k risk stratification method introduced in this work. Beyond survival analysis, the introduction of the possible permutations carries potential for other tasks that involve ranking based on limited order information. The utilization of specialized sorting networks, such as splitter selection networks as in Petersen et al. [2022b], could further leverage partial order information.
Though promising, this work is not without limitations. Future research could focus on extending its applicability to non-proportional hazards and understanding the impact of ties. Moreover, investigating 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.
290   
291   
292   
293   
294   
295   
296   
297   
298   
299   
300   
301   
302   
303   
304   
305   
306   
307   
308   
309   
310   
311   
312   
313   
314   
315   
316   
317   
318   
319
Table 3: Results for real-world and semi-synthetic datasets. Mean (standard deviation). Survival metric is c-index, Top  $10\%$  metric is top-k-score. Bold indicates significant improvement (t-test,  $p \leq 0.01$ ).
FLCHAIN	NWTCO	SUPPORT	METABRIC	MIMIC IV CXR	survSVHN
Size	6,524	4,028	8,873	1,904	377,110	248,823
Censored Proportion	69.9%	85.8%	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 (.015)	.642 (.006)	.698 (.011)	.760 (.002)	.933 (.002)
Diffsurv	.793 (.009)	.703 (.026)	.645 (.002)	.684 (.011)	.763 (.001)	.943 (.002)
Top 10% prediction
Cox Partial Likelihood	.460 (.013)	.390 (.068)	.280 (.023)	.249 (.065)	.390 (.010)	-
CPL-TopK (Variant I)	.469 (.007)	.413 (.061)	.479 (.016)	.527 (.083)	.408 (.008)	-
CPL-TopK (Variant II)	.460 (.009)	.413 (.054)	.479 (.035)	.487 (.058)	.406 (.006)	-
Diffsurv	.452 (.011)	.395 (.082)	.296 (.015)	.331 (.102)	.412 (.002)	-
Diffsurv-TopK	.482 (.019)	.421 (.065)	.508 (.027)	.533 (.092)	.412 (.009)	-
Furthermore, it is important to note that Diffsurv is a survival ranking method and thus can not be used to directly estimate the expected duration until an event occurs.
Overall, Diffsurv constitutes a meaningful advancement in survival analysis, showcasing its significant potential for enhancing risk prediction in real-world use cases. It not only demonstrates promising performance improvements, but also introduces new directions for future research, thereby making a valuable contribution to the field.
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 the anonymous reviewers for their thoughtful feedback and fruitful discussions, which have been instrumental in enhancing the manuscript.
320   
321   
322   
323   
324   
325   
326   
327   
328   
329   
330
331   
332   
333   
334   
335   
336   
337   
338   
339   
340   
341   
342   
343   
344   
345   
346   
347   
348   
349   
350   
351   
352   
353   
354   
355   
356   
357   
358   
359   
360   
361   
362   
363   
364   
365   
366   
367   
368   
369   
370   
371   
372   
373   
374   
375   
376   
377   
378   
379   
380
381   
382   
383   
384   
385   
386   
387   
388   
389   
390   
391   
392   
393   
394   
395   
396   
397   
398   
399   
400   
401   
402   
403   
404   
405   
406   
407   
408   
409   
410   
411   
412   
413   
414   
415   
416   
417   
418   
419   
420   
421   
422   
423   
424   
425
A Non-proportional Hazards
Our current implementation of Diffsurv operates under the proportional hazards assumption. While this may not fully capture the intricacies of some survival analysis problems—particularly those involving non-proportional hazards—it does not necessarily limit the model's effectiveness in scenarios where the goal is to assess cumulative risk from a fixed index date or from the date of an imaging study. This aligns with the necessity of time-dependent modifications to the C-index for non-proportional models as indicated by Antolini et al. [2005].
If we are primarily interested in understanding the cumulative hazard of an event occurring rather than tracking changes in the hazard over time, the assumption of proportional hazards becomes less pivotal. As such, Diffsurv and CPL remain valuable tools for these cases.
Despite the current limitation of Diffsurv to proportional hazards, it is conceivable that an extension to accommodate non-proportional hazards could be developed, similar to adaptations made for the CPL method.
For instance, as we briefly mentioned earlier, continuous-time extensions of partial likelihood can be used to enable non-proportional hazards Kvamme et al. [2019]. Implemented by directly modeling temporal covariates as  $f_{\theta} = \exp (h_{\theta}(x_i,T_i))$ .
Another class of methods focuses on discretizing the time-to-event variable and modeling the probability mass function (PMF) of event times. For instance, the DeepHit model Lee et al. [2018] employs a neural network architecture to learn the relationships between input features and discretized time-to-event outcomes. Time discretization facilitates modeling of non-proportional hazards but introduces two significant challenges: 1) sensitivity to the choice of time intervals, which can affect the model's accuracy and interpretability, and 2) increased computational complexity, as predictions must be made for each time interval. These models can be computationally expensive, especially for deep learning-based models like DeepHit, making them less suitable for high-dimensional and large-scale datasets, such as the imaging dataset used in this study.
Several future work proposals arise from these observations. First, differentiable sorting could explore the approach of directly modeling temporal covariates, resulting in a time-parameterized predicted permutation matrix. Second, extending Diffsurv to discrete time could be achieved by parameterizing a predicted permutation matrix for each time discretization.
B Training and evaluation
B.1 Datasets and Preprocessing
As in Goldstein and Langholz [1992] and Kvamme et al. [2019], we ensure that each risk set contains a valid risk set by sampling controls for a given case. Each batch consists of a number of risk sets such that the input data has shape (batch size, risk set size, covariate shape).
We provide an additional description of each small realworld dataset:
- FLCHAIN dataset: A dataset containing information on patients with monoclonal gammopathy of undetermined significance (MGUS), focusing on serum free light chain (FLC) levels to study their prognostic significance in predicting disease progression. Number of covariates: 8.
- 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 of seriously 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 uncovering novel molecular subtypes for more precise prognostication and personalized treatment strategies. Number of covariates: 9.
426   
427   
428   
429   
430   
431   
432   
433   
434   
435   
436   
437   
438   
439   
440   
441   
442   
443   
444   
445   
446   
447   
448   
449   
450   
451   
452   
453   
454   
455   
456   
457   
458   
459   
460   
461   
462   
463   
464   
465   
466   
467   
468   
469   
470   
471   
472   
473
Covariate preprocessing follows [Kvamme et al., 2019], and includes standardising continuous variables and one-hot encoding categorical variables.
MIMIC IV CXR: For the survival task, we extract death events from the MIMIC IV dataset. This is done by merging the data on "subject_id", "study_id", and "dicom_id" with the patient table from MIMIC IV, and on "subject_id" with the admission table. For patients without a recorded date of death, censoring dates are be determined as 1 year after the last recorded discharge date for each patient. We exclude 29,345 images without any matches in the MIMIC IV patient table, 19,337 images taken after the latest found discharge date and 55 images taken after a recorded date of death. Time to event is calculated as the number of days from the image study date to either date of death or the censoring date. For MIMIC IV CXR, images undergo several standard transformations: a random horizontal flip and a 15-degree rotation, resizing to  $230 \times 230$  pixels, a  $224 \times 224$  pixel center crop, and conversion to grayscale with three output channels. The data is then transformed into tensors and normalized using ImageNet's mean and standard deviation values. Finally, the train:val:test split of 8:1:1 is done at the patient level ensuring no images from a patient in the test set was found in the training data.
survSVHN: In this semi-synthetic dataset, we sample survival times using the beta-exponential distribution. The beta distribution used to sample from the exponential uses a fixed value of 500 for both shape parameters. We follow Petersen et al. [2021] by cropping the centered multi-digit numbers with a boundary of  $30\%$ , resizing it to a resolution of  $64 \times 64$ , and then selecting  $54 \times 54$  pixels at a random location. For survSVHN the train:val:test split is provided by Netzer et al. [2011] as is 230,755:5,000:13,068.
B.2 Model Architecture and Hyperparameters
For the smaller real-world datasets, the hazard function  $f_{\theta}$  is a small fixed Multi-layer Perceptron network with 1 hidden layer and 64 hidden nodes. We also apply a fixed dropout rate of 0.1. Learning rate, weight decay, batch set size and risk set size were found using a grid search across the possible values in Table 4.
Table 4: Hyperparameter values for small real-world datasets.
Hyperparameter	Values
Learning rate	[0.1, 0.01, 0.001, 1e-4]
Weight decay	[0.1, 0.01, 0.001, 1e-4, 1e-5, 0]
(Batch size, risk set size)	[(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 is fixed in each batch in Table 9.
Note that during evaluation, the sorting network is not used since we only need to evaluate the ranks of the trained risk scores. Similarly, case-control sampling is not used. We measure the ranking performance of the models using the concordance index.
Further implementation details and the best hyperparameters for each dataset are provided at anon@ git.com.
474   
475   
476   
477   
478   
479   
480   
481   
482   
483   
484   
485   
486   
487   
488   
489   
490   
491   
492   
493   
494   
495   
496   
497   
498   
499   
500   
501   
502   
503   
504   
505   
506   
507   
508   
509   
510   
511   
512   
513   
514   
515   
516
B.3 Compute Requirements