SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Experiments on smaller real-world datasets are compact enough to facilitate effective training on a CPU, with each variant, including the CPH baselines, completing per experiment in less than 20 minutes. However, the larger imaging datasets require more significant computational power. In the most demanding case, the MIMIC IV CXR experiments, run on an 11GB NVIDIA GeForce GTX 1080 Ti, took roughly 18.5 hours per experiment. Both Diffsurv and CPH methods exhibited comparable run times; however, the Bitonic variant was the fastest, with a lead of approximately 6 minutes. All neural network baselines were implemented using PyTorch and PyTorch Lightning. Although measures were taken to reduce compute time and complexity, such as using half-precision and distributed data parallel (DDP) training strategies, the overall training times are far from optimized.
To understand the runtime differences of Diffsurv and Cox Partial Likelihood variants in isolation, we also run an additional experiment. For each method, we compute and time over 100 trials for a forward and backward pass on a NVIDIA GeForce GTX 1080 Ti using randomly generated logits. For diffsurv, this includes computing the predicted permutation matrix with differentiable sorting networks and applying the masking (Equation 15) and binary cross entropy (Equation 16). The possible permutation matrix generation is possible to precompute off GPU on the dataloader, so is not included in the timing. In Table 5, the Diffsurv methods, particularly Bitonic, consistently outperform CPL methods in compute time almost across all batch sizes and risk set sizes. As risk set size increases, CPL methods exhibit a decreasing trend in compute time, while Diffsurv's Odd-Even method experiences a notable rise, especially from risk set sizes of 32 to 128.
It is worth noting that CPL methods are currently computed over batches using a simple for loop, as the present implementations do not support batch parallel computation. However, there's potential for further optimization. For similar batch sizes and risk set sizes, we noted very similar overall convergence times with Diffsurv Bitonic variant being marginally faster than other methods. In the context of full model runs, the difference between Diffsurv and CPL in terms of training times is minimal; it is the model architecture that have a dominant effect on compute time.
Table 5: Isolated compute time for different methods, with various batch sizes and risk set sizes over 100 trials. Mean time and  $95\%$  confidence intervals are provided in milliseconds.
Method	Batch Size, Risk Set Size
512, 2	128, 8	32, 32	8, 128
Diffsurv: Odd-Even	12.37 ± 0.15	25.63 ± 0.09	107.71 ± 0.38	305.91 ± 3.66
Diffsurv: Bitonic	3.85 ± 0.03	17.29 ± 0.08	55.49 ± 0.47	92.32 ± 0.80
CPL: Breslow	853.32 ± 11.39	281.22 ± 1.20	94.08 ± 0.15	25.45 ± 0.46
CPL: Efron	1729.87 ± 5.54	494.83 ± 4.27	164.73 ± 1.40	51.69 ± 0.76
CPL: Ranked List	719.42 ± 4.99	225.69 ± 0.94	74.36 ± 0.18	18.19 ± 0.31
C Sorting Networks
There are multiple different types of sorting networks each with varying complexity. The ability to implement networks with the divide-and-conquer paradigm allows for sorting networks that scale more efficiently. Examples for Odd-Even and Bitonic sorting networks with  $n = 8$  are shown in Figure 4. The latter allows construction of networks with size complexity  $\mathcal{O}(n\log^2 n)$  verses the  $\mathcal{O}(n^2)$  in Odd-Even networks.
It is worth emphasising these are not neural networks. They are called "networks" because they are typically represented as diagrams that show how the items are compared and swapped as they are being sorted. Differentiable sorting networks do not introduce any additional parameters that need to be updated during optimization.
D Calibration of Predicted Permutations
Model calibration in survival analysis models is essential for ensuring that the predicted probabilities of outcomes align closely with the true probabilities. An improperly calibrated model may lead
517   
518   
519   
520   
521   
522   
523   
524   
525   
526   
527   
528   
529   
530   
531   
532   
533   
534   
535   
536   
537   
538   
539   
540   
541   
542
Figure 4: Example sorting networks of size 8; (a) Odd-Even, (b) Bitonic.
Table 6: Calibration of rank probabilities from predicted permutation matrices on survSVHN, keeping the number of events per batch equal. Mean (and standard deviation) Brier scores over 5 trails with different seeds. Corresponds with C-index results in Table 9.
Method	Batch Size, Risk Set Size
512, 2	128, 8	32, 32	8, 128
CPL: Breslow	0.081 (0.002)	0.181 (0.001)	0.056 (0.000)	0.014 (0.000)
CPL: Efron	0.081 (0.002)	0.180 (0.001)	0.055 (0.000)	0.014 (0.000)
CPL: Ranked_List	0.066 (0.001)	0.185 (0.001)	0.053 (0.000)	0.012 (0.000)
Diffsurv: Bitonic	0.059 (0.001)	0.165 (0.001)	0.043 (0.000)	0.010 (0.000)
Diffsurv: Odd-Even	0.056 (0.001)	0.158 (0.001)	0.039 (0.000)	0.009 (0.000)
to incorrect risk assessments and treatment decisions, potentially resulting in suboptimal patient care and even adverse clinical consequences. We focus on the calibration of predicted individual rankings, as demonstrated in Figure 5 and Table 6. Specifically, we qualitatively illustrate in Figure 5 that for a model with a risk set size of 8, both discrete predicted ranks and ranking probabilities are accurately calibrated for the Diffsurv approach. To perform a quantitative comparison with baseline methods, we need to derive ranking probabilities for the CPL model. Based on the assumptions in Raykar et al. [2007], we assume that the probability of correct pairwise ordering for the CPL adheres to the logistic function. We thus compute permutation matrices using differential sorting networks, employing predicted partial log hazards as inputs and the logistic sigmoid function as the differentiable sorting operator. By subsequently calculating Brier scores for the rank probabilities in the predicted permutation matrices survSVHN, we analyze various combinations of batch size and risk set size. Our findings show that the Diffsurv models consistently exhibit the lowest Brier scores across all settings (refer to Table 6).
Figure 5: Calibration visualization of the Diffsurv (Bitonic) predicted ranking for a group of 8 subjects on the survSVHN dataset.
556   
557   
558   
559   
560   
561   
562   
563   
564   
565   
566   
567   
568
E Effect of top-k loss variants on predictive performance
The Top-K risk prediction variants for Diffsurv and the Cox Partial Likelihood, as introduced in Section 3.2, inherently bias the model to recognize only the top-k individuals with the highest risk. In this section, we explore the extent to which these specific variants influence the predictive performance of the models across the entire set of individuals. These results are shown in Table 7.
Table 7: C-index for real-world and semi-synthetic datasets for Top-k loss variants
Top 10% prediction	FLCHAIN	NWTCO	SUPPORT	METABRIC	MIMIC IV CXR
Cox Partial Likelihood	.796 (.010)	.725 (.014)	.650 (.008)	.691 (.013)	.760 (.002)
CPL-TopK (Variant I)	.794 (.011)	.715 (.011)	.602 (.003)	.644 (.016)	.751 (.002)
CPL-TopK (Variant II)	.798 (.011)	.722 (.016)	.600 (.010)	.658 (.008)	.753 (.002)
Diffsurv	.794 (.011)	.696 (.025)	.634 (.008)	.649 (.016)	.761 (.002)
Diffsurv-TopK	.783 (.011)	.689 (.022)	.596 (.005)	.639 (.023)	.754 (.001)
In Table 8 and Table 9, additional results for the MIMIC IV CXR and survSVHN imaging datasets are provided. Here, we maintain a constant total number of samples in each batch, which means that an increase in risk set size is compensated by a decrease in batch size. These results offer further understanding of the balance required between these two variables. We observed, while larger risk set sizes generally improve performance for both Diffsurv and CPH, the benefits tend to taper off as training can become more unstable and noisy with smaller batch sizes.
Table 8: 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 \leq 0.01$ . Most significant across all Batch Size, Risk Set Sizes.
Method	Batch Size, Risk Set Size
64, 2	4, 32	16, 8	1, 128
Diffsurv: Bitonic	0.761 (0.001)	0.761 (0.000)	0.763† (0.001)	0.761 (0.002)
Diffsurv: Odd-Even	0.761 (0.002)	0.756 (0.002)	0.761 (0.001)	0.749 (0.001)
CPL: Ranked List	0.760 (0.002)	0.755 (0.002)	0.758 (0.003)	0.755 (0.002)
Table 9: Additional results for survSVNH keeping the number of events per batch equal. Mean (and standard deviation) over 5 trials with different seeds. Metric is C-index. Bold indicates a significantly higher result with t-test and  $p \leq 0.01$ .
Method	Batch Size, Risk Set Size
512, 2	128, 8	32, 32	8, 128
Diffsurv: Odd-Even	0.934 (0.001)	0.940 (0.001)	0.941 (0.001)	0.933 (0.002)
Diffsurv: Bitonic	0.931 (0.001)	0.942 (0.001)	0.940 (0.00166)	0.928 (0.001)
CPL: Breslow	0.905 (0.001)	0.897 (0.001)	0.910 (0.002)	0.919 (0.001)
CPL Efron	0.904 (0.002)	0.898 (0.002)	0.909 (0.003)	0.918 (0.003)
CPL: Ranked List	0.921 (0.001)	0.922 (0.003)	0.921 (0.001)	0.917 (0.003)
569   
570   
571   
572   
573