| UCRL2 | $0.01\pm 0.00$ | $0.55\pm 0.34$ | $0.45\pm 0.34$ | $0.00\pm 0.00$ | $0/20$
MG-Empty (7) | PSRL | $0.01\pm 0.01$ | $0.01\pm 0.01$ | $0.99\pm 0.01$ | $0.00\pm 0.00$ | $0/20$
| Q-learning | $0.02\pm 0.02$ | $0.02\pm 0.03$ | $0.98\pm 0.03$ | $0.00\pm 0.00$ | $0/20$
| UCRL2 | $0.01\pm 0.00$ | $0.73\pm 0.32$ | $0.27\pm 0.32$ | $0.00\pm 0.00$ | $0/20$
MG-Empty (8) | PSRL | $0.01\pm 0.01$ | $0.01\pm 0.01$ | $0.99\pm 0.01$ | $0.00\pm 0.00$ | $0/20$
| Q-learning | $0.02\pm 0.01$ | $0.02\pm 0.01$ | $0.98\pm 0.01$ | $0.00\pm 0.00$ | $0/20$
| UCRL2 | $0.01\pm 0.00$ | $0.07\pm 0.09$ | $0.93\pm 0.09$ | $0.00\pm 0.00$ | $0/20$
MG-Rooms (1) | PSRL | $0.01\pm 0.04$ | $0.01\pm 0.02$ | $0.99\pm 0.02$ | $0.00\pm 0.00$ | $0/20$
| Q-learning | $0.04\pm 0.02$ | $0.02\pm 0.03$ | $0.98\pm 0.03$ | $0.00\pm 0.00$ | $0/20$
| UCRL2 | $0.01\pm 0.00$ | $0.82\pm 0.28$ | $0.18\pm 0.28$ | $0.00\pm 0.00$ | $0/20$
MG-Rooms (2) | PSRL | $0.00\pm 0.00$ | $0.00\pm 0.00$ | $1.00\pm 0.00$ | $0.00\pm 0.00$ | $0/20$
| Q-learning | $0.01\pm 0.02$ | $0.02\pm 0.02$ | $0.98\pm 0.02$ | $0.00\pm 0.00$ | $0/20$
| UCRL2 | $0.00\pm 0.00$ | $0.38\pm 0.35$ | $0.62\pm 0.35$ | $0.00\pm 0.00$ | $0/20$
RiverSwim (1) | PSRL | $0.77\pm 0.38$ | $1.00\pm 0.00$ | $0.00\pm 0.00$ | $0.01\pm 0.00$ | $20/20$
| Q-learning | $0.21\pm 0.15$ | $0.27\pm 0.19$ | $0.73\pm 0.19$ | $0.01\pm 0.01$ | $20/20$
| UCRL2 | $0.85\pm 0.24$ | $1.00\pm 0.00$ | $0.00\pm 0.00$ | $0.02\pm 0.00$ | $20/20$
RiverSwim (2) | PSRL | $0.90\pm 0.30$ | $1.00\pm 0.00$ | $0.00\pm 0.00$ | $0.01\pm 0.00$ | $20/20$
| Q-learning | $0.29\pm 0.21$ | $0.29\pm 0.21$ | $0.71\pm 0.21$ | $0.01\pm 0.00$ | $20/20$
| UCRL2 | $0.90\pm 0.20$ | $0.99\pm 0.00$ | $0.01\pm 0.00$ | $0.02\pm 0.00$ | $20/20$
RiverSwim (3) | PSRL | $0.88\pm 0.24$ | $0.98\pm 0.04$ | $0.02\pm 0.04$ | $0.01\pm 0.00$ | $20/20$
| Q-learning | $0.09\pm 0.06$ | $0.10\pm 0.06$ | $0.90\pm 0.06$ | $0.01\pm 0.00$ | $20/20$
| UCRL2 | $0.80\pm 0.25$ | $0.99\pm 0.01$ | $0.01\pm 0.01$ | $0.01\pm 0.00$ | $20/20$
RiverSwim (4) | PSRL | $0.89\pm 0.26$ | $0.99\pm 0.00$ | $0.01\pm 0.00$ | $0.01\pm 0.00$ | $20/20$
| Q-learning | $0.82\pm 0.20$ | $0.50\pm 0.24$ | $0.50\pm 0.25$ | $0.02\pm 0.01$ | $20/20$
| UCRL2 | $0.82\pm 0.29$ | $0.99\pm 0.01$ | $0.01\pm 0.01$ | $0.01\pm 0.00$ | $20/20$
SimpleGrid (1) | PSRL | $0.22\pm 0.07$ | $0.30\pm 0.19$ | $0.70\pm 0.19$ | $0.00\pm 0.00$ | $0/20$
| Q-learning | $0.22\pm 0.00$ | $0.22\pm 0.00$ | $0.78\pm 0.00$ | $0.01\pm 0.00$ | $20/20$
| UCRL2 | $0.24\pm 0.04$ | $0.99\pm 0.01$ | $0.01\pm 0.01$ | $0.00\pm 0.00$ | $20/20$
SimpleGrid (2) | PSRL | $0.78\pm 0.22$ | $0.99\pm 0.02$ | $0.01\pm 0.02$ | $0.00\pm 0.00$ | $20/20$
| Q-learning | $0.67\pm 0.04$ | $0.54\pm 0.07$ | $0.46\pm 0.07$ | $0.01\pm 0.00$ | $20/20$
| UCRL2 | $0.79\pm 0.16$ | $1.00\pm 0.00$ | $0.00\pm 0.00$ | $0.01\pm 0.00$ | $19/20$
SimpleGrid (3) | PSRL | $0.48\pm 0.10$ | $0.57\pm 0.16$ | $0.43\pm 0.16$ | $0.00\pm 0.00$ | $0/20$
| Q-learning | $0.51\pm 0.00$ | $0.51\pm 0.00$ | $0.49\pm 0.00$ | $0.00\pm 0.00$ | $15/20$
| UCRL2 | $0.50\pm 0.00$ | $1.00\pm 0.00$ | $0.00\pm 0.00$ | $0.00\pm 0.00$ | $0/20$
Taxi (1) | PSRL | $0.13\pm 0.08$ | $0.11\pm 0.08$ | $0.89\pm 0.08$ | $0.00\pm 0.00$ | $0/20$
| Q-learning | $0.19\pm 0.01$ | $0.13\pm 0.01$ | $0.87\pm 0.01$ | $0.01\pm 0.00$ | $20/20$
| UCRL2 | $0.14\pm 0.06$ | $0.91\pm 0.01$ | $0.09\pm 0.01$ | $0.00\pm 0.00$ | $18/20$

###  E.5 Non-tabular setting

We now present the full results of the benchmarking procedure for the non-
tabular baseline agents from bsuite, ActorCritic, ActorCriticRNN, BootDQN, and
DQN.

In order to keep the hardness induced by the emission map as low as possible
while still testing the effective non-tabular capabilities of the agent, we
employ the deterministic state information map, which provides clear state-
identifying information. The experimental procedure is the same as for the
tabular case (described in the previous Section) with the only difference
being that the total training time given to the agent is $40$ minutes both for
the continuous and episodic settings.

The performances of the agents are in line with the results reported in Osband
et al. [4], with the exception of BootDQN. Differently from bsuite, the
Colosseum benchmarking procedure penalizes particularly computationally
expensive algorithms by limiting the training time. Being BootDQN the most
computationally intensive algorithm due to the ensemble training, this agent
often breaks the time limit, which consequently worsens its overall
performance.

Tables 15da, 15db, 15dc, and 15dd report a summary of the performance of the
agents in terms of normalized cumulative regrets. DQN is the best performing
agent on average for all the different settings. Unsurprisingly, it performs
relatively better in the ergodic setting compared to the communicating one, in
which the exploration challenge is harder. This is especially clear from the
results for the DeepSea family in the continuous setting. DQN is not able to
perform well in any of the instances in the communicating case (Table 15dc).
Although BootDQN generally performs better than the ActorCritic agents, it is
not able to fully express its potential due to the associated computational
burden. Interestingly, the ActorCritic agent without the recurrent component
performs better than the version equipped with it. Note that this is not due
to a different computational cost as ActorCriticRNN is always able to complete
the interactions with the MDPs within the given time limit.

Figure 23, 24, 25, and 26 show how the cumulative regret, along with standard
error intervals, of the agents evolve during their interaction with the
benchmark environments. Note that the $\\}$ symbol is always reported for
BootDQN, meaning that, for at least one of the seeds, the agent ran out of
time. We note that, overall, the variability of the performances across the
different seeds is low, with the exception of a few cases. In particular, the
RiverSwim family is associated with moderate variability. This is not
surprising since, for this MDP family, the agent can easily get trapped in a
sub-optimal policy if it does not fully explore the state space.

Figures 27a, 27b, 27c, and 27d place the regret of the agents on a position
corresponding to the diameter and value norm of each MDP respectively for the
episodic communicating, episodic ergodic, continuous communicating, and
continuous ergodic settings. Overall, we observe that the harder environments
effectively induce higher regret. This relationship is visibly stronger when
the agents perform better across the environments, (see the ActorCritic agents
in Figure 27c and DQN in all settings). Interestingly, and similarly to the
tabular case, while the best performing agent (DQN) is evidently impacted more
by the diameter, the opposite holds for the other agents. Regardless of the
visitation complexity, this suggests that an agent that fails to handle the
estimation complexity of an environment is bound to perform badly both in
tabular and non-tabular settings.

Table 15: Summary of benchmark results for the non-tabular bsuite baselines.

| ActorCritic | ActorCriticRNN | BootDQN | DQN
---|---|---|---|---
DeepSea | $0.33\pm 0.24$ | $0.45\pm 0.14$ | $0.17\pm 0.24$ | $\mathbf{0.05}\pm 0.09$
| $0.32\pm 0.28$ | $0.45\pm 0.21$ | $0.28\pm 0.31$ | $\mathbf{0.21}\pm 0.26$
FrozenLake | $0.12\pm 0.08$ | $0.42\pm 0.06$ | $0.25\pm 0.25$ | $\mathbf{0.02}\pm 0.01$
MG-Empty | $0.75\pm 0.38$ | $0.92\pm 0.15$ | $0.26\pm 0.38$ | $\mathbf{0.04}\pm 0.07$
| $0.81\pm 0.29$ | $0.94\pm 0.11$ | $0.43\pm 0.37$ | $\mathbf{0.09}\pm 0.06$
| $0.60\pm 0.33$ | $0.85\pm 0.19$ | $0.39\pm 0.40$ | $\mathbf{0.07}\pm 0.03$
| $0.96\pm 0.08$ | $0.96\pm 0.07$ | $0.31\pm 0.39$ | $\mathbf{0.10}\pm 0.28$
| $0.99\pm 0.03$ | $0.96\pm 0.06$ | $0.33\pm 0.33$ | $\mathbf{0.10}\pm 0.09$
MG-Rooms | $1.00\pm 0.00$ | $1.00\pm 0.00$ | $0.31\pm 0.37$ | $\mathbf{0.19}\pm 0.32$
| $1.00\pm 0.01$ | $1.00\pm 0.00$ | $\mathbf{0.34}\pm 0.39$ | $\mathbf{0.34}\pm 0.48$
| $1.00\pm 0.00$ | $1.00\pm 0.01$ | $0.54\pm 0.43$ | $\mathbf{0.36}\pm 0.48$
| $0.87\pm 0.30$ | $1.00\pm 0.00$ | $0.44\pm 0.41$ | $\mathbf{0.27}\pm 0.44$
RiverSwim | $\mathbf{0.00}\pm 0.00$ | $\mathbf{0.00}\pm 0.00$ | $0.01\pm 0.01$ | $\mathbf{0.00}\pm 0.00$
| $0.53\pm 0.39$ | $0.67\pm 0.31$ | $\mathbf{0.13}\pm 0.21$ | $0.21\pm 0.36$
SimpleGrid | $0.80\pm 0.10$ | $0.78\pm 0.13$ | $0.44\pm 0.35$ | $\mathbf{0.06}\pm 0.02$
| $0.80\pm 0.01$ | $0.78\pm 0.03$ | $0.38\pm 0.32$ | $\mathbf{0.05}\pm 0.01$
| $0.38\pm 0.17$ | $0.29\pm 0.24$ | $0.26\pm 0.25$ | $\mathbf{0.00}\pm 0.00$
| $0.79\pm 0.06$ | $0.77\pm 0.07$ | $0.48\pm 0.28$ | $\mathbf{0.20}\pm 0.20$
Taxi | $0.91\pm 0.01$ | $0.91\pm 0.01$ | $0.87\pm 0.14$ | $\mathbf{0.66}\pm 0.24$
| $\mathbf{0.86}\pm 0.01$ | $\mathbf{0.86}\pm 0.01$ | $0.90\pm 0.02$ | $0.87\pm 0.01$
Average | $0.69\pm 0.30$ | $0.75\pm 0.28$ | $0.38\pm 0.21$ | $\mathbf{0.19}\pm 0.22$

(a)

| ActorCritic | ActorCriticRNN | BootDQN | DQN
---|---|---|---|---
DeepSea | $0.46\pm 0.04$ | $0.48\pm 0.00$ | $\mathbf{0.23}\pm 0.18$ | $0.25\pm 0.24$
| $0.01\pm 0.01$ | $\mathbf{0.00}\pm 0.00$ | $0.07\pm 0.09$ | $\mathbf{0.00}\pm 0.00$
FrozenLake | $0.67\pm 0.23$ | $0.88\pm 0.07$ | $0.34\pm 0.33$ | $\mathbf{0.09}\pm 0.02$
MG-Empty | $0.86\pm 0.27$ | $0.96\pm 0.09$ | $0.24\pm 0.36$ | $\mathbf{0.04}\pm 0.03$
| $0.96\pm 0.07$ | $0.99\pm 0.01$ | $0.42\pm 0.29$ | $\mathbf{0.20}\pm 0.15$
| $0.98\pm 0.05$ | $0.98\pm 0.03$ | $0.35\pm 0.31$ | $\mathbf{0.27}\pm 0.36$
| $0.83\pm 0.25$ | $0.96\pm 0.07$ | $0.35\pm 0.26$ | $\mathbf{0.08}\pm 0.03$
| $0.77\pm 0.23$ | $0.95\pm 0.06$ | $0.55\pm 0.25$ | $\mathbf{0.29}\pm 0.15$
| $0.71\pm 0.38$ | $0.65\pm 0.25$ | $0.24\pm 0.37$ | $\mathbf{0.02}\pm 0.02$
| $0.27\pm 0.14$ | $0.69\pm 0.16$ | $0.28\pm 0.35$ | $\mathbf{0.04}\pm 0.01$
MG-Rooms | $0.87\pm 0.22$ | $0.94\pm 0.16$ | $0.52\pm 0.43$ | $\mathbf{0.16}\pm 0.28$
| $0.94\pm 0.16$ | $0.96\pm 0.08$ | $0.58\pm 0.40$ | $\mathbf{0.36}\pm 0.47$
| $0.93\pm 0.16$ | $1.00\pm 0.00$ | $0.31\pm 0.35$ | $\mathbf{0.22}\pm 0.37$
RiverSwim | $\mathbf{0.00}\pm 0.00$ | $\mathbf{0.00}\pm 0.00$ | $\mathbf{0.00}\pm 0.01$ | $\mathbf{0.00}\pm 0.00$
| $\mathbf{0.00}\pm 0.01$ | $\mathbf{0.00}\pm 0.00$ | $0.01\pm 0.01$ | $\mathbf{0.00}\pm 0.00$
SimpleGrid | $0.53\pm 0.27$ | $0.67\pm 0.23$ | $0.37\pm 0.37$ | $\mathbf{0.01}\pm 0.01$
| $0.79\pm 0.04$ | $0.78\pm 0.06$ | $0.43\pm 0.29$ | $\mathbf{0.09}\pm 0.03$
| $0.51\pm 0.03$ | $0.47\pm 0.07$ | $0.24\pm 0.22$ | $\mathbf{0.04}\pm 0.01$
Taxi | $0.77\pm 0.01$ | $0.77\pm 0.01$ | $0.76\pm 0.16$ | $\mathbf{0.54}\pm 0.24$
| $0.34\pm 0.02$ | $0.34\pm 0.02$ | $0.36\pm 0.04$ | $\mathbf{0.32}\pm 0.02$
Average | $0.61\pm 0.32$ | $0.67\pm 0.34$ | $0.33\pm 0.18$ | $\mathbf{0.15}\pm 0.15$

(b)

| ActorCritic | ActorCriticRNN | BootDQN | DQN
---|---|---|---|---
DeepSea | $\mathbf{0.14}\pm 0.25$ | $0.18\pm 0.27$ | $0.61\pm 0.20$ | $0.55\pm 0.00$
| $\mathbf{0.17}\pm 0.38$ | $0.41\pm 0.50$ | $0.85\pm 0.27$ | $0.74\pm 0.43$
| $\mathbf{0.18}\pm 0.27$ | $0.23\pm 0.28$ | $0.62\pm 0.15$ | $0.54\pm 0.00$
FrozenLake | $\mathbf{0.06}\pm 0.04$ | $0.17\pm 0.12$ | $0.26\pm 0.11$ | $0.07\pm 0.03$
| $0.27\pm 0.17$ | $0.29\pm 0.16$ | $0.51\pm 0.14$ | $\mathbf{0.18}\pm 0.04$
MG-Empty | $0.92\pm 0.29$ | $0.92\pm 0.29$ | $0.77\pm 0.41$ | $\mathbf{0.14}\pm 0.16$
| $0.26\pm 0.33$ | $0.34\pm 0.42$ | $0.38\pm 0.43$ | $\mathbf{0.10}\pm 0.22$
| $\mathbf{0.13}\pm 0.28$ | $0.16\pm 0.30$ | $0.32\pm 0.43$ | $0.14\pm 0.30$
| $0.92\pm 0.29$ | $0.75\pm 0.45$ | $0.77\pm 0.41$ | $\mathbf{0.06}\pm 0.06$
| $1.00\pm 0.00$ | $1.00\pm 0.00$ | $0.69\pm 0.46$ | $\mathbf{0.26}\pm 0.27$
MG-Rooms | $0.47\pm 0.34$ | $0.61\pm 0.47$ | $0.74\pm 0.35$ | $\mathbf{0.26}\pm 0.32$
| $0.34\pm 0.35$ | $0.71\pm 0.34$ | $0.62\pm 0.43$ | $\mathbf{0.24}\pm 0.35$
| $\mathbf{0.62}\pm 0.42$ | $0.77\pm 0.41$ | $0.77\pm 0.36$ | $0.69\pm 0.34$
RiverSwim | $\mathbf{0.00}\pm 0.00$ | $0.07\pm 0.23$ | $0.14\pm 0.30$ | $\mathbf{0.00}\pm 0.00$
| $0.07\pm 0.23$ | $0.07\pm 0.23$ | $\mathbf{0.00}\pm 0.01$ | $\mathbf{0.00}\pm 0.00$
SimpleGrid | $\mathbf{0.00}\pm 0.01$ | $0.05\pm 0.08$ | $0.55\pm 0.48$ | $0.01\pm 0.01$
| $0.25\pm 0.26$ | $0.38\pm 0.23$ | $0.35\pm 0.22$ | $\mathbf{0.04}\pm 0.11$
| $0.09\pm 0.27$ | $0.17\pm 0.28$ | $0.55\pm 0.45$ | $\mathbf{0.01}\pm 0.01$
| $0.33\pm 0.25$ | $0.42\pm 0.19$ | $0.30\pm 0.25$ | $\mathbf{0.01}\pm 0.02$
Taxi | $0.92\pm 0.01$ | $0.92\pm 0.01$ | $0.93\pm 0.02$ | $\mathbf{0.81}\pm 0.17$
Average | $0.36\pm 0.33$ | $0.43\pm 0.31$ | $0.54\pm 0.24$ | $\mathbf{0.24}\pm 0.26$

(c)

| ActorCritic | ActorCriticRNN | BootDQN | DQN
---|---|---|---|---
DeepSea | $0.29\pm 0.43$ | $0.36\pm 0.45$ | $0.27\pm 0.21$ | $\mathbf{0.22}\pm 0.39$
FrozenLake | $\mathbf{0.11}\pm 0.13$ | $0.21\pm 0.28$ | $0.41\pm 0.17$ | $\mathbf{0.11}\pm 0.08$
MG-Empty | $0.92\pm 0.07$ | $0.97\pm 0.04$ | $0.32\pm 0.18$ | $\mathbf{0.05}\pm 0.02$
| $0.95\pm 0.02$ | $0.98\pm 0.01$ | $0.35\pm 0.29$ | $\mathbf{0.05}\pm 0.02$
| $0.87\pm 0.07$ | $0.94\pm 0.06$ | $0.39\pm 0.18$ | $\mathbf{0.10}\pm 0.02$
| $0.97\pm 0.02$ | $0.97\pm 0.03$ | $0.42\pm 0.24$ | $\mathbf{0.08}\pm 0.04$
| $0.97\pm 0.03$ | $0.99\pm 0.00$ | $0.57\pm 0.40$ | $\mathbf{0.08}\pm 0.05$
| $0.99\pm 0.01$ | $0.99\pm 0.01$ | $0.69\pm 0.36$ | $\mathbf{0.12}\pm 0.07$
| $0.98\pm 0.02$ | $0.99\pm 0.01$ | $0.54\pm 0.43$ | $\mathbf{0.10}\pm 0.06$
| $0.99\pm 0.00$ | $0.98\pm 0.01$ | $0.73\pm 0.28$ | $\mathbf{0.17}\pm 0.08$
MG-Rooms | $0.97\pm 0.05$ | $0.99\pm 0.01$ | $0.59\pm 0.35$ | $\mathbf{0.24}\pm 0.17$
| $0.99\pm 0.00$ | $0.99\pm 0.00$ | $0.59\pm 0.25$ | $\mathbf{0.38}\pm 0.27$
RiverSwim | $0.08\pm 0.29$ | $0.33\pm 0.49$ | $\mathbf{0.00}\pm 0.01$ | $\mathbf{0.00}\pm 0.00$
| $0.07\pm 0.23$ | $0.07\pm 0.23$ | $0.01\pm 0.03$ | $\mathbf{0.00}\pm 0.00$
| $\mathbf{0.00}\pm 0.00$ | $0.13\pm 0.31$ | $0.02\pm 0.04$ | $\mathbf{0.00}\pm 0.00$
| $0.33\pm 0.49$ | $0.17\pm 0.39$ | $0.01\pm 0.02$ | $\mathbf{0.00}\pm 0.00$
SimpleGrid | $0.75\pm 0.01$ | $0.76\pm 0.01$ | $0.57\pm 0.27$ | $\mathbf{0.09}\pm 0.04$
| $0.49\pm 0.01$ | $0.49\pm 0.01$ | $0.36\pm 0.16$ | $\mathbf{0.02}\pm 0.01$
| $0.50\pm 0.00$ | $0.50\pm 0.00$ | $0.37\pm 0.18$ | $\mathbf{0.06}\pm 0.04$
Taxi | $0.78\pm 0.02$ | $0.78\pm 0.02$ | $0.81\pm 0.03$ | $\mathbf{0.55}\pm 0.13$
Average | $0.65\pm 0.36$ | $0.68\pm 0.34$ | $0.40\pm 0.24$ | $\mathbf{0.12}\pm 0.14$

(d)

Figure 23: Full interaction performances for the non-tabular baseline agents
in the episodic communicating setting. Figure 24: Full interaction
performances for the non-tabular baseline agents in the episodic ergodic
setting. Figure 25: Full interaction performances for the non-tabular baseline
agents in the continuous communicating setting. Figure 26: Full interaction
performances for the non-tabular baseline agents in the continuous ergodic
setting.

(a)

(b)

(c)

(d)

Figure 27: Average cumulative regret obtained by the agents in the continuous
ergodic setting placed according to the diameter and the value norm values of
the benchmark MDPs.


# Diffsurv: Differentiable sorting for censored time-to-event data.

Andre Vauvelle1, Benjamin Wild2, Aylin Cakiroglu3, Roland Eils2, Spiros
Denaxas1
University College London1, Berlin Institute of Health2, BenevolentAI3,
<EMAIL_ADDRESS><EMAIL_ADDRESS>
{benjamin.wild<EMAIL_ADDRESS>
###### Abstract

Survival analysis is a crucial semi-supervised task in machine learning with
numerous real-world applications, particularly in healthcare. Currently, the
most common approach to survival analysis is based on Cox’s partial
likelihood, which can be interpreted as a ranking model optimized on a lower
bound of the concordance index. This relation between ranking models and Cox’s
partial likelihood considers only pairwise comparisons. Recent work has
developed differentiable sorting methods which relax this pairwise
independence assumption, enabling the ranking of sets of samples. However,
current differentiable sorting methods can not account for censoring, a key
factor in many real-world datasets. To address this limitation, we propose a
novel method called _Diffsurv_. We extend differentiable sorting methods to
handle censored tasks by predicting matrices of possible permutations that
take into account the label uncertainty introduced by censored samples. We
contrast this approach with methods derived from partial likelihood and
ranking losses. Our experiments show that Diffsurv outperforms established
baselines in various simulated and real-world risk prediction scenarios.
Additionally, we demonstrate the benefits of the algorithmic supervision
enabled by Diffsurv by presenting a novel method for top-k risk prediction
that outperforms current methods.

##  1 Introduction and Background

Survival analysis is an important task in numerous machine learning
applications, particularly in the healthcare domain. The goal of survival
analysis is to predict the time until the occurrence of an event of interest,
such as death or disease onset, based on a set of covariates. In clinical
studies, these covariates typically include demographic variables such as sex
and age, but may also encompass more complex data modalities such as temporal
streams or medical images.

However, event times may not be observed due to censoring, especially in
observational datasets where many patients may not have experienced the event
at the time of data collection. Ignoring censoring can lead to biased
predictions towards the censoring event instead of the event of interest. For
example, if the end of the study can be determined from the observed
covariates, especially if age is included, the predicted event times will be
skewed towards the censoring event time instead of the actual event of
interest Kvamme & Borgan (2019).

The Cox Proportional Hazards (PH) 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 PH and ranking models can be directly equated, with both
providing lower bounds to the concordance index, the primary evaluation metric
used in survival analysis. A key step in relating the two models assumes only
pairwise comparisons or risk sets of size 2. Goldstein & Langholz (1992) show
that sub-sampling risk sets produce consistent parameter estimators but that
greater risk sets provide more efficient estimators. Cox’s partial likelihood
and ranking losses underpin current survival analysis methods in deep
learning, including DeepSurv Katzman et al. (2018) and DeepHit Lee et al.
(2018).

We present an alternative method that leverages recent advancements in
continuous relaxations of sorting operations, enabling end-to-end training of
neural networks with ordering supervision (Grover et al., 2019; Blondel et
al., 2020; Petersen et al., 2021). This involves incorporating a sorting
algorithm into the network architecture, where the order of the samples is
known, but their exact values are unsupervised. Here, we introduce _Diffsurv_
, an extension of differentiable sorting methods that enables end-to-end
training of survival models with censored data.

Briefly, our contributions are summarised:

  * •

Our primary contribution is the extension differentiable sorting methods to
account for censoring by introducing the concept of possible permutation
matrices.

  * •

We empirically demonstrate that our new differentiable sorting method improve
risk ranking performance across multiple simulated and real-worlds censored
datasets.

  * •

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.

##  2 Methods

Figure 1: Differentiable Sorting for Censored Time-to-Event Data. Inputs, in
this case, SVHN images, are transformed into scalar risk values, $z_{i}$,
through a neural network. A differentiable permutation matrix, ${\bm{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 ${\bm{Q}}_{p}$ of possible permutations based on the observed
events and censoring.

A dataset with censored event times is summarized as
$\mathcal{D}=\\{t_{i},{\bm{x}}_{i},\delta_{i}\\}_{i=1}^{N}$. For a patient
$i$, $t_{i}$ is the observed minimum of the unobserved true survival time
$t_{i}^{*}$ and the censoring time $c_{i}^{*}$, $\delta_{i}$ is the event
indicator that is 1 if an event is observed ($t_{i}^{*}\leq c_{i}^{*}$) or 0
if the data is censored ($t_{i}^{*}>c_{i}^{*}$). Covariates are
${\bm{x}}_{i}\in\mathbb{R}^{d}$ representing a 1-dimensional vector of size
$d$ or larger dimensional tensors such as image data. $N$ is the total number
of patients. As previously mentioned, it is common to subsample total possible
risk set, we use $n$ to represent the subsampled risk set size.

In order to train models based on ordering information using differentiable
sorting algorithms Petersen (2022), we can minimize the cross-entropy between
the ground truth orders represented by true permutation matrix ${\bm{Q}}$ and
a doubly-stochastic predicted permutation matrix ${\bm{P}}$. This makes it
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$.

There are multiple methods of relaxing sorting algorithms to produce
${\bm{P}}$, we will follow Petersen et al. (2021) by using differentiable
sorting networks. Sorting networks are a family of sorting algorithms that
consist of two basic components: wires and conditional swaps. Wires carry
values to be compared at conditional swaps, if one value is bigger than the
other, the values carried forward are swapped around. For a random sample of
patients to be ordered, each layer of the sorting network can be considered an
independent permutation matrix ${\bm{P}}_{l}$ with elements given by

$P_{l,ii}=P_{l,jj}=\sigma(z_{j}-z_{i})\text{ and
}P_{l,ij}=P_{l,ji}=1-\sigma(z_{j}-z_{i}).$ (1)

These elements represent conditional swaps between two patient risk values
$(z_{i},z_{j})$ and use a differentiable relaxation of the step function such
as the logistic-sigmoid, where $\sigma:x\rightarrow\frac{1}{1+e^{-\beta x}}$.
The inverse temperature parameter $\beta>0$ is introduced so when
$\beta\rightarrow\infty$ the functions tend to the exact $\min$ and $\max$
functions. The indices being compared are determined by the sorting network
and the final predicted probability matrix is the product of each layer of
sorting operations,
${\bm{P}}=\left(\prod_{l=1}^{n}{\bm{P}}_{l}^{\intercal}\right)^{\intercal}$.
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.1.

The introduction of censored patients means we no longer have access to a
ground truth permutation matrix $\mathbf{Q}$. We cannot determine the exact
rank of patients who are censored before another who experienced an event. To
address this challenge, we propose a novel extension of differentiable sorting
to censored data. Our approach considers the set of _possible permutations_
for each patient, taking into account uncertainty about the true ranking. In
Figure 3, we show an example of observed and censored events and the resulting
set of possible permutations that can be represented as a permutation matrix
${\bm{Q}}_{p}$.

For a right-censored sample $i$, we only know that the rank must be lower than
the rank of all other samples with an event time lower than the censoring time
of $t_{i}$, i.e. they must be ranked after prior events. For another sample
$j$ with an event at $t_{j}$, we know that the rank must be lower than other
samples with an event time lower than $t_{j}$, and higher than the rank of
other samples either with an event time higher than $t_{j}$ or with a
censoring time higher than $t_{j}$. We do not know how the rank of $j$
compares to samples with censoring time lower than $t_{j}$. If it is possible
for patient i to permute to rank j, then $Q_{pij}=1$, otherwise $Q_{pij}=0$.

Given the possible permutation matrix $\mathbf{Q}_{p}$ and the predicted
permutation matrix $\mathbf{P}$, the vector of probabilities ${\bm{p}}$ of a
patient being ranked within the set of possible permutations can be computed.
Although the ground truth ranks are unknown, the range of possible ranks is
known, and the model can be optimized to maximize the sum of the predicted
permutation probabilities for the possible ranks of each sample. Noted here as
the column-sum of the element-wise product $\circ$, between $\mathbf{Q}_{p}$
and $\mathbf{P}$.

${\bm{p}}=\sum_{j=1}^{n}(\mathbf{Q}_{p}\circ\mathbf{P})_{i,j}.$ (2)

The cross-entropy loss can then be easily applied

$\mathcal{L}=\sum_{i=1}^{n}y_{i}\log(p_{i})$ (3)

where $y_{i}$ is the true label of the set of possible ranks.

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 optimizing possible permutations within the
top-k ranks, whereby ${\bm{Q}}_{p}$ is adjusted such that only the top-k
patient’s possible permutations are set to 1.

##  3 Experiments

We evaluate the performance of Diffsurv on censored survival data across semi-
synthetic and real-world datasets. In each experiment, we train a neural
network using Diffsurv and Cox’s partial likelihood loss, then compare their
respective results. Cox’s partial likelihood and the closely related ranking
loss are used in popular baselines; Deepsurv (Katzman et al., 2018), Cox-MLP
(Kvamme et al., 2019) and DeepHit Lee et al. (2018).

We present a new semi-synthetic dataset, _survSVHN_ , to evaluate survival
models. Based on the Street View House Numbers (SVHN) dataset Petersen et al.
(2021), we simulate survival times akin to survMNIST Pölsterl (2019). The
increased complexity of SVHN offers a testbed which is better able to discern
the performance differences between methods. Each house number parameterizes
an exponential time function for survival times. Risks 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 5). Risk is predicted from the images with a convolutional neural
network with the same hyperparameters as Petersen et al. (2021), with
$z_{i}=f_{\text{CONV}}({\bm{x}}_{i})$. We also evaluate on four real-world
healthcare datasets from Kvamme et al. (2019). Each dataset has a fairly small
number of patients ($N\leq 8,873$) and a flat vector of covariates as input.
Further details in Appendix A.6. For these datasets, a fully connected neural
network is used to find the risk, $z_{i}=f_{\text{MLP}}({\bm{x}}_{i})$.
Further details on the training and evaluation procedures can be found in
Appendix A.4.

Table 1: Results for semi-synthetic and real-world datasets. Bold indicates
significantly higher performance (t-test with a significance level of 0.01).

Method |  n=2† |  n=4 |  n=8 |  n=16 |  n=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)
(a) Semi-synthetic survSVHN Dataset Results. Mean (and standard deviation)
over 5 trails with different seeds. Metric is C-index. † When $n=2$ both
methods are equivalent to the ranking loss.

|  FLCHAIN |  NWTCO |  SUPPORT |  METABRIC
---|---|---|---|---
Size |  6,524 |  4,028 |  8,873 |  1,904
Censored Proportion |  69.9% |  85.8% |  32.0% |  42.1%
Risk Stratification |  |  |  |
Diffsurv |  .787 (.012) |  .691 (.018) |  .599 (.004) |  .623 (.012)
Cox Partial Likelihood |  .787 (.016) |  .690 (.014) |  .584 (.004) |  .615 (.018)
Top 10% prediction |  |  |  |
Diffsurv (Top-k) |  .952 (.014) |  .926 (.007) |  .571 (.012) |  .771 (.059)
Cox Partial Likelihood |  .937 (.014) |  .919 (.026) |  .478 (.032) |  .561 (.039)
(b) Real-world datasets results. Mean (and standard deviation) over 5-folds
measured in C-Index for the risk stratification task and proportion correctly
predicted for the Top-k task.

The results presented in Table 1 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
outperforms Cox’s partial likelihood on all four datasets. There is a
significant improvement in the top 10% highest-risk prediction when comparing
models based on Cox’s Partial Likelihood ($\mu=.825$, $\sigma=.005$) and
Diffsurv optimized for Top-k prediction ($\mu=.944$, $\sigma=.008$) on the
survSVHN dataset.

##  4 Conclusion

Diffsurv represents a significant step in the field of survival analysis with
censored data. Our experiments demonstrate the effectiveness of differentiable
sorting methods in improving survival analysis predictions, particularly in
censored datasets with Diffsurv matching or improving perfomance against Cox
partial likelihood on all datasets. Additionally, Diffsurv has the potential
to drive the development of new methods, such as the top-k risk stratification
method presented in this work. It is noteworthy that while our method has
shown promising results, further investigation is necessary to fully
understand its potential and limitations. For instance, it would be valuable
to examine the scalability of the method with larger real-world datasets and
its capability to handle more complex censored scenarios. Further research
could also investigate the integration of Diffsurv into clustering models.
With its ability to handle censored data and its end-to-end training
capability, Diffsurv presents a promising approach to survival analysis and
holds great potential for enhancing risk prediction in real-world
applications.

## References

  * Blondel et al. (2020) Mathieu Blondel, Olivier Teboul, Quentin Berthet, and Josip Djolonga.  Fast Differentiable Sorting and Ranking, June 2020.  URL http://arxiv.org/abs/2002.08871.  arXiv:2002.08871 [cs, stat].
  * Cox (1972) D. R. Cox.  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].
  * Goldstein & Langholz (1992) Larry Goldstein and Bryan Langholz.  Asymptotic Theory for Nested Case-Control Sampling in the Cox Regression Model.  _The Annals of Statistics_ , 20(4):1903–1928, December 1992.  ISSN 0090-5364, 2168-8966.  doi: 10.1214/aos/1176348895.  URL https://projecteuclid.org/journals/annals-of-statistics/volume-20/issue-4/Asymptotic-Theory-for-Nested-Case-Control-Sampling-in-the-Cox/10.1214/aos/1176348895.full.  Publisher: Institute of Mathematical Statistics.
  * Grover et al. (2019) Aditya Grover, Eric Wang, Aaron Zweig, and Stefano Ermon.  Stochastic Optimization of Sorting Networks via Continuous Relaxations.  _arXiv:1903.08850 [cs, stat]_ , April 2019.  URL http://arxiv.org/abs/1903.08850.  arXiv: 1903.08850.
  * Harrell et al. (1982) Frank E. Harrell, Jr, Robert M. Califf, David B. Pryor, Kerry L. Lee, and Robert A. Rosati.  Evaluating the Yield of Medical Tests.  _JAMA_ , 247(18):2543–2546, May 1982.  ISSN 0098-7484.  doi: 10.1001/jama.1982.03320430047030.  URL https://doi.org/10.1001/jama.1982.03320430047030.
  * Katzman et al. (2018) Jared L. Katzman, Uri Shaham, Alexander Cloninger, Jonathan Bates, Tingting Jiang, and Yuval Kluger.  DeepSurv: personalized treatment recommender system using a Cox proportional hazards deep neural network.  _BMC Medical Research Methodology_ , 18(1):24, February 2018.  ISSN 1471-2288.  doi: 10.1186/s12874-018-0482-1.  URL https://doi.org/10.1186/s12874-018-0482-1.
  * Kingma & Ba (2017) Diederik P. Kingma and Jimmy Ba.  Adam: A Method for Stochastic Optimization.  _arXiv:1412.6980 [cs]_ , January 2017.  URL http://arxiv.org/abs/1412.6980.  arXiv: 1412.6980.
  * Kvamme & Borgan (2019) Håvard Kvamme and Ørnulf Borgan.  The Brier Score under Administrative Censoring: Problems and Solutions.  Technical Report arXiv:1g912.08581, arXiv, December 2019.  URL http://arxiv.org/abs/1912.08581.  arXiv:1912.08581 [cs, stat] type: article.
  * Kvamme et al. (2019) Håvard Kvamme, Ørnulf Borgan, and Ida Scheel.  Time-to-Event Prediction with Neural Networks and Cox Regression.  _arXiv:1907.00825 [cs, stat]_ , September 2019.  URL http://arxiv.org/abs/1907.00825.  arXiv: 1907.00825.
  * Lee et al. (2018) Changhee Lee, William Zame, Jinsung Yoon, and Mihaela Van Der Schaar.  DeepHit: A Deep Learning Approach to Survival Analysis with Competing Risks.  _Proceedings of the AAAI Conference on Artificial Intelligence_ , pp. 8, 2018.
  * Petersen (2022) Felix Petersen.  Learning with Differentiable Algorithms, September 2022.  URL http://arxiv.org/abs/2209.00616.  arXiv:2209.00616 [cs].
  * Petersen et al. (2021) Felix Petersen, Christian Borgelt, Hilde Kuehne, and Oliver Deussen.  Differentiable Sorting Networks for Scalable Sorting and Ranking Supervision.  _arXiv:2105.04019 [cs]_ , July 2021.  URL http://arxiv.org/abs/2105.04019.  arXiv: 2105.04019.
  * Pölsterl (2019) Sebastian Pölsterl.  Survival Analysis for Deep Learning, July 2019.  URL https://k-d-w.org/blog/2019/07/survival-analysis-for-deep-learning/.
  * Raykar et al. (2007) Vikas C Raykar, Harald Steck, Balaji Krishnapuram, Cary Dehing-oberije, and Philippe Lambin.  On Ranking in Survival Analysis: Bounds on the Concordance Index.  In _Advances in Neural Information Processing Systems_ , volume 20. Curran Associates, Inc., 2007.  URL https://papers.nips.cc/paper/2007/hash/33e8075e9970de0cfea955afd4644bb2-Abstract.html.

##  Appendix A Appendix

###  A.1 Cox’s Partial Likelihood and ranking losses

Cox (1972) introduced the most popular method for addressing censoring for
survival analysis. The original work and many subsequent extensions seek to
maximize the partial likelihood, the general form is described as

$\mathcal{L(\theta)}=\prod_{i:\delta=1}\frac{f_{\theta}({\bm{x}}_{i})}{\sum_{j:T_{j}>T_{i}}f_{\theta}({\bm{x}}_{j})},$
(4)

where $f_{\theta}$ is the hazard function, a real valued score prediction
function estimating the probability of an event at particular time, given
input features ${\bm{x}}_{i}$. The product is taken over the set of uncensored
patients, while the denominator term considers only comparable pairs and
includes censored patients with $T_{j}>T_{i}$.

The classic Cox proportional hazards model used
$f_{\theta}=\text{exp}(\theta\cdot{\bm{x}}_{i})$. Multiple extensions of this
basic loss relax the linear covariate interaction and proportional hazards
assumptions by altering $f_{\theta}$. For example, (Katzman et al., 2018)
parameterized with a neural network $h_{\theta}$, such that
$f_{\theta}=\text{exp}(h_{\theta}({\bm{x}}_{i}))$, to model non-linear
interactions between covariates on the hazard. (Kvamme et al., 2019) further
show that this can be extended to non-proportional hazards by introducing
temporal covariates, $f_{\theta}=\text{exp}(h_{\theta}({\bm{x}}_{i},T_{i}))$.

Kvamme et al. (2019) also make a few adjustments to the original partial
likelihood loss. First, they consider that the risk set
$\mathcal{R}=\\{j:T_{j}>T_{i}\\}$ is intractable for deep learning
applications as it considers all comparable patients. Instead, it is possible
take a fixed size sample risk set $\lvert\tilde{\mathcal{R}}\rvert=n<N$ and
further, 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=1}\frac{f_{\theta}({\bm{x}}_{i})}{f_{\theta}({\bm{x}}_{i})+f_{\theta}({\bm{x}}_{j})},\text{
}j\in\mathcal{R}\setminus\\{i\\}.$ (5)

Further, we can take the mean log partial likelihood to be

$\text{loss}=\frac{1}{n_{e}}\sum_{i:\delta=1}\log(1+\exp[h_{\theta}({\bm{x}}_{j})-h_{\theta}({\bm{x}}_{i})]),\text{
}j\in\mathcal{R}\setminus\\{i\\},$ (6)

where $n_{e}$ is the number of non-censored events. From this simplified form,
it can be seen that the partial likelihood only considers the 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.

$\text{c-index}=\frac{1}{n}\sum_{i:\delta=1}\mathbbm{1}(f({\bm{x}}_{i})<f({\bm{x}}_{j})),j\in\mathcal{R}\setminus\\{i\\}.$
(7)

Raykar et al. (2007) 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 maximizing the general ranking loss

$\text{ranking-
loss}=\frac{1}{\lvert{\mathcal{A}\rvert}}\sum_{({\bm{x}}_{i},{\bm{x}}_{j})\in\mathcal{A}}\phi(f_{\theta}({\bm{x}}_{i})-f_{\theta}({\bm{x}}_{j})),$
(8)

where $\phi$ is a function that relaxes the non-differentiable $\mathbbm{1}$
of the C-index. We have also introduced $\mathcal{A}$ as the graph of
acceptable pairs, where each node is a patient that can only be linked to
another with an edge if we are sure that the first event occurs before the
second. This is another way of describing the risk sets introduced earlier.
From Equation 6 is can be seen that
$\phi:x\rightarrow-\log(1+\exp(-x))=\log(\sigma(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 1 (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$).

###  A.2 Differentiable Sorting

(a)

(b)

Figure 2: Example sorting networks of size 8; (a) Odd-Even, (b) Bitonic.

Differentiable sorting takes a different approach to the previously discussed
partial likelihood loss and will need modification to account for censoring.
In order train models based on ordering information, differences between
predicted and true orderings are backpropagated through a relaxed sorting
algorithm. The idea was first introduced by (Grover et al., 2019), the key
motivation being that sorting is a key step in many classical algorithms and
machine learning methods (K-nearest neighbours), yet is non-differentiable.
This means that direct supervision cannot be done on the outputs of algorithms
that rely on sorting (Petersen, 2022).

Sorting algorithms require 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 7). Differentiable sorting
methods similarly rely on approximating these operators with smooth
alternatives.

Petersen et al. (2021) propose combining traditional sorting networks and
differentiable sorting functions. Sorting networks are a family of sorting
algorithms that consist of two basic components: wires and conditional swaps.
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 provably guaranteed sorting networks. 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)$. Note that
as $a$ approaches $b$, the point at which it becomes larger is discontinuous
and hence non-differentiable. Just as previously shown in the ranking loss,
such operations can be made differentiable using the logistic relaxation

$\text{min}_{\sigma}(a,b)=a\cdot\sigma(b-a)+b\cdot\sigma(a-b)\text{ and
}\text{max}_{\sigma}(a,b)=a\cdot\sigma(a-b)+b\cdot\sigma(b-a).$ (9)

Note that if an inverse temperature parameter $\beta>0$ is introduced such
that $\sigma:x\rightarrow\frac{1}{1+e^{-\beta x}}$, then as
$\beta\rightarrow\infty$ the functions tend to the exact $\min$ and $\max$
functions. Other relaxation of the step function can also be considered,
Petersen et al. (2021) 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\rightarrow\frac{1}{\pi}\text{arctan}(\beta x)+\frac{1}{2}$.

There are multiple different types of sorting networks each with varying space
complexity. The ability to implement networks with the divide-and-conquer
paradigm allows for sorting networks that scale more efficiently than previous
differentiable sorting methods. In particular, Petersen et al. (2021) uses the
odd-even and bitonic sorting networks. The latter allows construction of
networks with size complexity $\mathcal{O}(nlog^{2}n)$ verses the
$\mathcal{O}(n^{2})$ in previous differentiable sorting methods. Examples for
Odd-Even and Bitonic sorting networks with $n=8$ are shown in Figure 2.

####  A.2.1 Relation to ranking and partial likelihood

Differentiable sorting has so far only been applied in the context of
uncensored ranking, where we know the true rank for every instance in an input
set. It is possible to directly relate Diffsurv with ranking losses. Expanding
out the cross entropy loss out we find

$\mathcal{L}=\sum_{c}^{n}\left(\frac{1}{n}\sum_{i}^{n}q_{ci}\log(p_{ci})\right),$
(10)

where $q_{ci}=1$ only when $i$ is the true rank otherwise $0$. Each $p_{ci}$
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_{i}$ introduced
by the sorting network but denoted as

$p_{ci}=\prod^{n}_{(a_{i},a_{j})\in\mathcal{P}_{l}:l=1}\sigma(f_{\theta}(a_{i})-f_{\theta}(a_{j}))$
(11)

where $\mathcal{P}_{l}$ to denotes the set of comparisons to be made at each
layer of the sorting network. With risk set of size 2, the loss returns to the
same recognisable log-sigmoid ranking loss, and Cox negative log partial
likelihood with risk set size 2.

###  A.3 Non-proportional Hazards

Our current implementation of Diffsurv is limited to proportional hazards,
which may not fully capture the complexity of certain survival analysis
problems, particularly when non-proportional hazards are present. In this
context, we explore alternative approaches that address this limitation.

Previously, we briefly mentioned continuous-time extensions of partial
likelihood to enable non-proportional hazards Kvamme et al. (2019). This can
be achieved by directly modeling temporal covariates as
$f_{\theta}=\text{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.
Finally, in models like DeepHit, the ranking loss term could be replaced with
a Diffsurv loss, offering another promising direction for future research in
survival analysis.

###  A.4 Training and evaluation

During training, we use the Adam optimizer Kingma & Ba (2017) with a learning
rate of $10^{-3}$, early stopping with patience of 10 epochs and a maximum of
$10^{5}$ training steps. As in Goldstein & 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).

For the survSVHN task, the hazard function $h_{\theta}$ for both the Cox
Partial Likelihood baseline and $f_{\theta}$ for Diffsurv, is a fixed sized
convolutional neural network with a batch size of 100. Steepness is determined
by risk set size $n$, $\beta=2n$ for Odd-even networks and
$\beta=\log_{2}(n)(1+\log_{2}(n))$ for Bitonic. The only hyperparameter
optimized for Diffsurv is the choice of sorting network. Both sorting network
were evaluated on the validation set but only the resulting models with higher
validation c-index were used on the test set. This was repeated for each risk
set size.

For the real-world datasets, the hazard function $h_{\theta}$ and $f_{\theta}$
for Diffsurv is Multi-layer Perceptron network. Hyperparameters: number of
hidden layers, size of hidden layers, dropout rate, batch size and learning
rate are determined by a small hyperparameter sweep of 100 trails on random
80:20 splits for each dataset and method. Hyperparameter optimization for
Diffsurv also included sorting network and steepness.

Full hyperparameter ranges and resulting best models are provided along with
further implementation details at https://github.com/andre-vauvelle/diffsurv-
ea.

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 Harrell et al. (1982).

###  A.5 Possible Permutation Matrix

In order to account for censoring, we propose utilizing a possible permutation
matrix $Q_{p}$. We provide a visualization for example risk set size 7 in
Figure 3, which corresponds with Equation 12. It is also possible to represent
such connections as a graph. Further, this possible permutation graph
${\mathcal{G}}_{p}$, is the complement of the order graph ${\mathcal{G}}_{o}$
typically used for C-index. Both graphs are visualised in Figure 4

Figure 3: Possible permutations for an example case with two events ($e_{0}$
and $e_{4}$) and multiple censored samples ($c_{0},c_{2},c_{3},c_{5},c_{6}$).

${\bm{Q}}_{p}=\begin{pmatrix}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\end{pmatrix}$ (12)

(a) Possible Permutation Graph, ${\mathcal{G}}_{p}$

(b) Order Graph, ${\mathcal{G}}_{o}$

Figure 4: Graphs representation of possible permutations and order.

###  A.6 Real World Datasets

  * •

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.

  * •

NWTS 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.

  * •

SUPPORT dataset: A dataset from a multicenter 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.

  * •

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.

###  A.7 SurvSVHN Dataset

Figure 5: Visual abstract of the survSVHN dataset.