SurvTRACE: Transformers for Survival Analysis with Competing Events

TL;DR

SurvTRACE is a transformer-based model for survival analysis with competing events that does not assume a specific survival distribution, incorporates multi-task learning, and provides interpretability through attention mechanisms, validated on large clinical datasets.

Contribution

This work introduces SurvTRACE, a novel transformer model for survival analysis with competing risks that leverages multi-task learning and interpretability, without assuming a specific survival distribution.

Findings

Outperforms existing methods on METABRIC, SUPPORT, and SEER datasets.

Effectively handles competing risks and implicit confounders.

Provides interpretable attention-based covariate importance.

Abstract

In medicine, survival analysis studies the time duration to events of interest such as mortality. One major challenge is how to deal with multiple competing events (e.g., multiple disease diagnoses). In this work, we propose a transformer-based model that does not make the assumption for the underlying survival distribution and is capable of handling competing events, namely SurvTRACE. We account for the implicit \emph{confounders} in the observational setting in multi-events scenarios, which causes selection bias as the predicted survival probability is influenced by irrelevant factors. To sufficiently utilize the survival data to train transformers from scratch, multiple auxiliary tasks are designed for multi-task learning. The model hence learns a strong shared representation from all these tasks and in turn serves for better survival analysis. We further demonstrate how to inspect the covariate relevance and importance through interpretable attention mechanisms of SurvTRACE, which suffices to great potential in enhancing clinical trial design and new treatment development. Experiments on METABRIC, SUPPORT, and SEER data with 470k patients validate the all-around superiority of our method.