Tabular Foundation Models Can Do Survival Analysis

TL;DR

This paper introduces a novel classification-based framework that enables existing tabular foundation models to perform survival analysis by discretizing event times and handling censored data naturally, achieving superior results on real-world datasets.

Contribution

It reformulates survival analysis as a series of binary classification problems, allowing foundation models to perform survival tasks without explicit training.

Findings

Outperforms classical and deep learning baselines on 53 datasets

Effectively handles censored data as missing labels

Proves consistency of survival probability estimation with increasing data

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.