Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

Tabular Foundation Models Can Do Survival Analysis

Da In Kim, Wei Siang Lai, Kelly W. Zhang

TL;DR

This work tackles right-censoring in survival analysis by reframing time-to-event tasks as a series of binary classifications obtained from discretized time intervals. Using this static and dynamic formulation, pretrained tabular foundation models can perform survival analysis via in-context learning without task-specific training, with consistency guarantees under conditional censoring. The authors prove that minimizing the population binary cross-entropy losses recovers true survival probabilities and demonstrate strong empirical gains across 53 real-world datasets, notably with MITRA and TabPFN, outperforming classical and deep baselines on multiple metrics. The approach offers a data-efficient, scalable pathway to leverage general-purpose tabular models for time-to-event prediction and opens avenues for further enhancements like ranking-focused training and continuous-time extensions.

Abstract

While tabular foundation models have achieved remarkable success in classification and regression, adapting them to model time-to-event outcomes for survival analysis is non-trivial due to right-censoring, where data observations may end before the event occurs. We develop a classification-based framework that reformulates both static and dynamic survival analysis as a series of binary classification problems by discretizing event times. Censored observations are naturally handled as examples with missing labels at certain time points. This classification formulation enables existing tabular foundation models to perform survival analysis through in-context learning without explicit training. We prove that under standard censoring assumptions, minimizing our binary classification loss recovers the true survival probabilities as the training set size increases. We demonstrate through evaluation across $53$ real-world datasets that off-the-shelf tabular foundation models with this classification formulation outperform classical and deep learning baselines on average over multiple survival metrics.

Tabular Foundation Models Can Do Survival Analysis

TL;DR

This work tackles right-censoring in survival analysis by reframing time-to-event tasks as a series of binary classifications obtained from discretized time intervals. Using this static and dynamic formulation, pretrained tabular foundation models can perform survival analysis via in-context learning without task-specific training, with consistency guarantees under conditional censoring. The authors prove that minimizing the population binary cross-entropy losses recovers true survival probabilities and demonstrate strong empirical gains across 53 real-world datasets, notably with MITRA and TabPFN, outperforming classical and deep baselines on multiple metrics. The approach offers a data-efficient, scalable pathway to leverage general-purpose tabular models for time-to-event prediction and opens avenues for further enhancements like ranking-focused training and continuous-time extensions.

Abstract

While tabular foundation models have achieved remarkable success in classification and regression, adapting them to model time-to-event outcomes for survival analysis is non-trivial due to right-censoring, where data observations may end before the event occurs. We develop a classification-based framework that reformulates both static and dynamic survival analysis as a series of binary classification problems by discretizing event times. Censored observations are naturally handled as examples with missing labels at certain time points. This classification formulation enables existing tabular foundation models to perform survival analysis through in-context learning without explicit training. We prove that under standard censoring assumptions, minimizing our binary classification loss recovers the true survival probabilities as the training set size increases. We demonstrate through evaluation across real-world datasets that off-the-shelf tabular foundation models with this classification formulation outperform classical and deep learning baselines on average over multiple survival metrics.
Paper Structure (56 sections, 3 theorems, 72 equations, 4 figures, 4 tables)

This paper contains 56 sections, 3 theorems, 72 equations, 4 figures, 4 tables.

Table of Contents

  1. Introduction
  2. Problem statement
  3. Static setting
  4. Dynamic setting
  5. Methodology: Time-to-event modeling as supervised learning
  6. Static setting
  7. Dynamic setting
  8. Theory: Consistent survival models via loss minimization
  9. Related work
  10. Experimental results
  11. Evaluation metrics
  12. Models
  13. Static results
  14. Dynamic results
  15. Discussion
  16. ...and 41 more sections

Key Result

Theorem 3.1

Let $p$ be any binary prediction model for $Y_{i,k}$ given $(X_{i,0}, t_k)$. Under Assumption assump:censoring (Conditionally Independent Censoring), the population loss $\ell_{\mathrm{static}}(p)$ is minimized if and only if, for each $k \in \{ 1, \dots, K-1 \}$, for all $x$ in the support of $X_{i,0}$ satisfying the positivity condition $\mathbb{P}(C_i \geq t_k \mid X_{i,0} = x) > 0$.

Figures (4)

  • Figure 1: Transforming of survival data into a binary classification dataset (static setting). The pink panel (top left) illustrates the original time-to-event data $(X_{i,0}, T_i, C_i, \delta_i)$, where event times $T_i$ and censoring times $C_i$ are continuous. For second example, the event remains latent ($T_2=?$) after the censoring time $C_2$. The blue panel (top right) displays the discretization of time into $5$ intervals $\{\tau_k \}_{k=1}^5$, where binary labels $Y_{i,k} = \mathbbm{1} (T_i \leq t_k)$ are defined at each timestep $t_k$ for $k=1, 2, 3, 4$. The table(bottom) shows the resulting classification dataset: each patient is transformed into at most $4$ tuples, where labels are only constructed for time steps preceding the censoring time ($t_k < C_i$), highlighting the impact of right-censoring on the training set.
  • Figure 2: Performance heterogeneity in the static setting. Distribution of C-index, Integrated AUC, and IBS for the top four models across the $43$ static datasets. Individual points represent specific datasets, while violin plots and vertical bars indicate the density, median, and interquartile range. Results are shown for $K=5$ discretization bins. MITRA and XGBoost exhibit the highest median performance across all metrics. Notably, MITRA maintains the best (or close to the best) lower-performing quartile, $25$th percentile for C-index/Integrated AUC and $75$th percentile for IBS, compared to all other evaluated models.
  • Figure 3: Correlation between classification and survival metrics (static). Each point represents a classification model evaluated on $1$ of the $43$ datasets. Results are shown for the test set and averaged over time-discretization granularities $K \in \{4, 5, 10, 15, 20\}$. We observe a strong correlation ($r = 0.89$) between the binary cross entropy loss and the integrated Brier score. This empirically supports our Theorem \ref{['thm:minimizer-static']} that a model that perfectly minimizes the loss necessarily means recovers the true survival probabilities.
  • Figure 4: Performance over time in the dynamic setting. Timepoint-wise performance across the five longitudinal datasets at timepoints $t_1, t_2, t_3$ ($K=5$ discretization bins). MITRA consistently outperforms both traditional landmarking baselines, joint modeling, and dynamic neural network-based survival models across most timepoints for all metrics. Notably, MITRA achieves a great lead over all other models on all three metrics at $t_2$. Baseline models are represented by dashed lines.

Theorems & Definitions (6)

  • Theorem 3.1: Consistent static survival model via loss minimization
  • proof
  • Lemma 1.1
  • proof
  • Theorem 1.2: Consistent dynamic survival model via loss minimization
  • proof