Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

Adaptive Transformer Modelling of Density Function for Nonparametric Survival Analysis

Xin Zhang, Deval Mehta, Yanan Hu, Chao Zhu, David Darby, Zhen Yu, Daniel Merlo, Melissa Gresle, Anneke Van Der Walt, Helmut Butzkueven, Zongyuan Ge

TL;DR

This paper introduces UniSurv, a nonparametric survival model based on Transformer architecture that jointly handles static and dynamic data, missing values, and censoring while producing unimodal, calibrated PDFs. It leverages a novel Margin-Mean-Variance loss to guide the learned distribution, improves censoring prediction, and supports both static and longitudinal inputs through dedicated extraction branches. Comprehensive experiments on static and dynamic datasets show UniSurv achieves superior or competitive C-index, MAE-H, MAE-U, and mAUC, with robust PDF visualizations and insightful ablations. The approach offers a flexible, end-to-end framework for time-to-event prediction with practical implications for healthcare and other domains requiring reliable survival analysis.

Abstract

Survival analysis holds a crucial role across diverse disciplines, such as economics, engineering and healthcare. It empowers researchers to analyze both time-invariant and time-varying data, encompassing phenomena like customer churn, material degradation and various medical outcomes. Given the complexity and heterogeneity of such data, recent endeavors have demonstrated successful integration of deep learning methodologies to address limitations in conventional statistical approaches. However, current methods typically involve cluttered probability distribution function (PDF), have lower sensitivity in censoring prediction, only model static datasets, or only rely on recurrent neural networks for dynamic modelling. In this paper, we propose a novel survival regression method capable of producing high-quality unimodal PDFs without any prior distribution assumption, by optimizing novel Margin-Mean-Variance loss and leveraging the flexibility of Transformer to handle both temporal and non-temporal data, coined UniSurv. Extensive experiments on several datasets demonstrate that UniSurv places a significantly higher emphasis on censoring compared to other methods.

Adaptive Transformer Modelling of Density Function for Nonparametric Survival Analysis

TL;DR

This paper introduces UniSurv, a nonparametric survival model based on Transformer architecture that jointly handles static and dynamic data, missing values, and censoring while producing unimodal, calibrated PDFs. It leverages a novel Margin-Mean-Variance loss to guide the learned distribution, improves censoring prediction, and supports both static and longitudinal inputs through dedicated extraction branches. Comprehensive experiments on static and dynamic datasets show UniSurv achieves superior or competitive C-index, MAE-H, MAE-U, and mAUC, with robust PDF visualizations and insightful ablations. The approach offers a flexible, end-to-end framework for time-to-event prediction with practical implications for healthcare and other domains requiring reliable survival analysis.

Abstract

Survival analysis holds a crucial role across diverse disciplines, such as economics, engineering and healthcare. It empowers researchers to analyze both time-invariant and time-varying data, encompassing phenomena like customer churn, material degradation and various medical outcomes. Given the complexity and heterogeneity of such data, recent endeavors have demonstrated successful integration of deep learning methodologies to address limitations in conventional statistical approaches. However, current methods typically involve cluttered probability distribution function (PDF), have lower sensitivity in censoring prediction, only model static datasets, or only rely on recurrent neural networks for dynamic modelling. In this paper, we propose a novel survival regression method capable of producing high-quality unimodal PDFs without any prior distribution assumption, by optimizing novel Margin-Mean-Variance loss and leveraging the flexibility of Transformer to handle both temporal and non-temporal data, coined UniSurv. Extensive experiments on several datasets demonstrate that UniSurv places a significantly higher emphasis on censoring compared to other methods.
Paper Structure (34 sections, 15 equations, 6 figures, 5 tables)

This paper contains 34 sections, 15 equations, 6 figures, 5 tables.

Table of Contents

  1. Introduction
  2. Literature
  3. Method
  4. Survival Notation
  5. Model Description
  6. Static and Dynamic Extractions
  7. Transformer Encoder
  8. Loss Function
  9. Experiments
  10. Datasets
  11. Static Datasets
  12. Dynamic Datasets
  13. Evaluation Metrics
  14. Concordance
  15. MAE-Uncensored
  16. ...and 19 more sections

Figures (6)

  • Figure 1: The illustration of (a) the architecture of UniSurv model and (b) a schematic representation of the UniSurv during training and testing stages
  • Figure 2: The illustration of reaction tensor representation of a single individual in MSReactor
  • Figure 3: The difference between the estimated lifetime $\hat{\mu}^i$ (blue dot) and the true censoring time $T^i$ (green square) of TDSM, DeepHit and UniSurv-s in METABRIC. Each red line indicates the difference if $\hat{\mu}^i<T^i$ for individual $i$, which is conversely not displayed in the opposite scenario
  • Figure 4: The time-dependent AUC. The dashed line shows mAUC corresponding to each colored curve
  • Figure 5: The estimated PDFs by DDH and UniSurv for five randomly selected uncensoring cases in MSReactor. Each color represents an individual
  • ...and 1 more figures