Tabular Foundation Models Can Do Survival Analysis

Technical Summary: Tabular Foundation Models Can Do Survival Analysis

Problem Statement
Survival analysis, or time-to-event analysis, aims to model when an event of interest occurs rather than merely whether it occurs. A defining challenge in this domain is right-censoring, where the observation period ends before the event happens, leaving the true event time unknown. While Tabular Foundation Models (TFMs) have achieved remarkable success in classification and regression tasks—often outperforming methods like XGBoost in small-data regimes—adapting them to survival analysis is non-trivial. Standard regression architectures cannot natively account for censoring, and existing survival-specific models often require training from scratch on single datasets, limiting their utility in low-data settings. Furthermore, while TFMs are evaluated on accuracy or mean squared error, survival modeling requires well-calibrated probabilities to characterize risk dynamics over time.

Methodology
The authors propose a framework that reformulates survival analysis as a series of binary classification problems to leverage off-the-shelf TFMs via in-context learning, without requiring survival-specific pretraining, fine-tuning, or architectural changes.

1. Temporal Discretization: The continuous time horizon is partitioned into $K$ discrete intervals.
2. Cumulative Failure Formulation: Unlike traditional approaches that model the discrete-time hazard rate (the probability of an event occurring in a specific interval given survival up to that point), this method directly models the cumulative failure probability. For each time point $t_k$, a binary label $Y_{i,k}$ is defined as $1(T_i \le t_k)$, indicating whether the event has occurred by time $t_k$.
3. Handling Censoring: Censored observations are treated as examples with missing labels at time points beyond the censoring time. The training objective is a standard binary cross-entropy loss, restricted only to observed labels (where $t_k < C_i$).
4. In-Context Learning: The reformulated dataset (features, time points, and binary labels) is fed into a TFM (e.g., MITRA, TabPFN) or a standard classifier (XGBoost). The model learns to predict the probability of failure by a specific time $t_k$ given the covariates.
5. Dynamic Extension: For dynamic settings with time-varying covariates, the model takes the history of observations up to time $t_k$ and predicts the cumulative failure probability for a future horizon $\Delta$, conditioned on the subject having survived up to $t_k$.

Theoretical Guarantees
The paper provides a theoretical proof (Theorem 4.2) demonstrating that under the standard assumption of conditionally independent censoring, minimizing the proposed binary cross-entropy loss recovers the true underlying survival probabilities as the training set size approaches infinity. This establishes that the classification formulation is statistically consistent with the survival objective.

Key Contributions

1. Classification Formulation: The authors introduce a specific classification formulation that models cumulative failure probabilities directly. Empirically, this outperforms formulations based on discrete hazard rates (which suffer from multiplicative error accumulation as the number of bins increases) and multi-class formulations.
2. Consistency Proof: They prove that minimizing the binary classification loss yields consistent survival probability estimates under standard censoring assumptions.
3. Empirical Evaluation: The approach is evaluated across 48 real-world datasets (43 static and 5 dynamic) from the SurvSet repository. The study compares off-the-shelf TFMs using this formulation against classical baselines (CoxPH, Random Survival Forests) and deep learning models (DeepSurv, DeepHit, Joint Models).

Results

• Static Setting: Across 43 datasets, TFMs using the cumulative failure (CF) formulation (specifically MITRA-CF and XGBoost-CF) achieved the best overall performance across multiple metrics (C-Index, Integrated AUC, Integrated Brier Score). The CF formulation proved more robust to the number of discretization bins ($K$) compared to hazard-based (DH) approaches, which degraded as $K$ increased due to error accumulation.
• Dynamic Setting: In 5 dynamic datasets, MITRA-CF outperformed all baselines, including landmarked CoxPH/RSF and Dynamic DeepHit. The results suggest that the temporal priors embedded in TFMs are particularly effective for modeling longitudinal data complexities.
• Metric Correlation: The study found a strong correlation ($r=0.89$) between the binary cross-entropy loss and the Integrated Brier Score, supporting the theoretical claim that optimizing the classification objective directly improves survival probability calibration.

Significance and Claims
The paper claims that survival analysis capabilities can be elicited from off-the-shelf tabular foundation models through a simple binary classification reformulation. This approach offers a computationally efficient alternative to training specialized survival models from scratch, particularly in small-data regimes where foundation models excel.

The authors distinguish their work from concurrent efforts (Survival In-Context and SurvivalPFN) that pretrain transformers specifically on synthetic survival tasks. In contrast, this work demonstrates that general-purpose TFMs, without any survival-specific pretraining or architectural modification, can achieve state-of-the-art performance on survival tasks when provided with the correct classification formulation. The method is flexible regarding the number of time bins and applicable to both static and dynamic settings, though it remains limited to scenarios satisfying the conditionally independent censoring assumption.