🤖 machine learning

Staying Alive: Uncensored Survival Analysis with Tabular Foundation Models

This paper introduces a training-free survival regression method that leverages Tabular Foundation Models to iteratively impute right-censored data and construct an Accelerated Failure Time model, achieving competitive performance against traditional trained models on standard benchmarks.

Original authors: Mariana Vargas Vieyra

Published 2026-06-03

📖 5 min read🧠 Deep dive

CC BY 4.0

Original authors: Mariana Vargas Vieyra

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

The Big Picture: Predicting "When" Without Training

Imagine you are a doctor trying to predict how long a patient will survive after a diagnosis, or a car mechanic guessing when a specific part will break. This is called Survival Analysis. The tricky part is that often, you don't get to see the final event. Maybe the patient moves away, or the study ends before the car part breaks. In statistics, this is called right-censoring—you know the event hasn't happened yet, but you don't know when it will.

Usually, to solve this, you need to build a custom model from scratch for every new dataset, which is like hiring a new architect to design a house every time you buy a plot of land.

This paper introduces a new way to do this using Tabular Foundation Models (TFMs). Think of a TFM as a "super-smart, pre-trained detective" that has already read millions of different datasets. It can look at a new list of facts (a table of data) and make a prediction in a single glance, without needing to be retrained. The authors wanted to see if this "super-detective" could solve the "missing time" puzzle of survival analysis without any extra training.

The Problem: The Detective Can't See the End

The problem is that these super-detectives (TFMs) are used to seeing the full story. If you ask them, "How long will this patient live?" but you only give them data up to the point where they left the study, the detective gets confused. It expects a definite answer, not a "we don't know yet."

If you just ignore the patients with missing end-times (censored data), the detective will get biased. It will think everyone dies sooner than they actually do, because it's only looking at the people who died early and ignoring the ones who are still alive.

The Solution: The "Fill-in-the-Blanks" Game

The authors created a method called TabSA (Tabular Survival Analysis) that lets the detective solve the puzzle in two clever steps:

Step 1: The "Guess the Average" (The AFT Model)

First, they use the detective to guess the "average log-time" for everyone based on their features (like age, health markers, etc.). This is like the detective looking at a crowd and guessing, "Based on what I've seen before, this group usually lasts about X years."

However, they only let the detective look at the people who did finish the study to make this guess. To make it work, they only have to tune one single number (a scale parameter). It's like telling the detective, "You're good at guessing, just tell me how much to stretch or shrink your guesses."

Step 2: The "Imagination Game" (Iterative Imputation)

This is the magic part. Since the detective can't see the end of the story for the censored patients, the authors play a game of "filling in the blanks."

The First Guess: They start by guessing the missing end-times for the censored patients using a standard statistical trick (like a rough estimate based on who else is still alive).
The Loop: They feed these guesses back to the detective. The detective looks at the whole group (including the guesses) and says, "Actually, based on this new info, my guess for that person should be a bit longer."
Refining: They update the guesses and ask the detective again. They repeat this loop over and over (like refining a sketch) until the guesses stop changing.

This process is inspired by an old statistical method called the Buckley-James estimator, but instead of doing complex math calculations, they let the "super-detective" (the TFM) do the heavy lifting of figuring out what the missing times should be.

The Results: How Did They Do?

The authors tested this on five real-world datasets (heart attacks, breast cancer, critical care, etc.) and compared their method against:

Classical Models: The traditional "custom architects" that require heavy training.
Other Zero-Shot Methods: Other ways to use these detectives without training.

The Findings:

Ranking Power: Their method (TabSA-BJ) was excellent at ranking patients. It could correctly say, "Patient A will likely live longer than Patient B" almost as well as the expensive, fully trained models.
Calibration: While it was great at ranking, it was slightly less perfect at predicting the exact probability of survival at a specific time compared to a method that breaks time into buckets (discretization).
No Training Needed: The biggest win is that they achieved these results without retraining the model on the specific dataset. They just fed the data in, ran the "fill-in-the-blanks" loop, and got a result.

The Takeaway

The paper shows that we don't always need to build a new model from scratch to predict survival times. By using a pre-trained "foundation model" as a smart engine and playing a simple game of "guess and refine" to handle missing data, we can get results that compete with traditional, heavily trained statistical models.

It's like having a master chef who has tasted every dish in the world. Instead of teaching them a new recipe from scratch, you just give them the ingredients and ask, "Based on your experience, how long will this stew take to cook?" and let them adjust their guess until it feels right.

Technical Summary: Staying Alive: Uncensored Survival Analysis with Tabular Foundation Models

Problem Statement
Survival Analysis (SA) is a statistical framework used to model the time until an event of interest occurs, with applications ranging from healthcare to churn prediction. A central challenge in SA is right-censoring, where the event time for some subjects is unobserved because they leave the study or reach the observation horizon before the event occurs. While classical survival models naturally handle censored data, Tabular Foundation Models (TFMs)—which excel at in-context classification and regression without dataset-specific training—struggle with this task. TFMs typically expect fully observed prediction targets, making direct application to survival regression difficult. Existing attempts to adapt TFMs, such as discretizing time into bins for classification, introduce a trade-off between resolution and effective context size, often degrading performance with coarser binning.

Methodology
The authors propose TabSA, a training-free framework that leverages frozen TFMs to perform survival regression by framing the task as an in-context prediction problem. The method operates under an Accelerated Failure Time (AFT) model assumption, where the logarithm of the event time is modeled linearly. The core innovation is an iterative, non-parametric imputation strategy for censored data, inspired by the Buckley-James estimator.

The proposed algorithm proceeds as follows:

Model Construction: The TFM is used as a regressor backbone to predict the mean log-time ( $\hat{\mu}$ ) for subjects based on in-context examples.
Parameter Estimation: A single scalar parameter, the scale parameter $\sigma$ , is estimated via Maximum Likelihood Estimation (MLE) on the training set. This is the only parameter requiring fitting; the TFM backbone remains frozen.
Iterative Imputation (TabSA-BJ): To address right-censoring without discarding data, the method iteratively imputes pseudo-targets for censored subjects:
- Initialization: Pseudo-targets are initialized using a data-driven warm start based on the Kaplan-Meier jackknife estimator.
- Iteration: In each step, the TFM predicts the mean log-time for censored subjects using a stochastic context of sub-sampled uncensored and censored subjects (to prevent degenerate self-prediction).
- Update: Unlike the standard Buckley-James estimator which imputes the conditional mean, this method updates pseudo-targets by drawing from a truncated distribution ( $p(\log T_i | T_i > C_i)$ ), strictly enforcing the constraint that the imputed time must exceed the censoring time.
- Refinement: The scale parameter $\sigma$ is re-estimated via MLE on the updated dataset. This loop continues until the pseudo-targets stabilize.
Inference: At test time, the model performs a forward pass to obtain mean log-times, which are combined with the final $\sigma$ to construct the survival function.

Key Contributions

Zero-Shot Survival Regression: The paper introduces a method to perform survival regression using TFMs without dataset-specific training, requiring only the fitting of a single scalar parameter ( $\sigma$ ).
Non-Parametric In-Context Imputation: The authors propose a novel iterative mechanism to handle right-censoring by leveraging the TFM as a non-parametric estimator for imputing censored data, avoiding the resolution-performance trade-offs associated with time discretization.
AFT Integration: The framework successfully constructs an AFT model using a TFM backbone, bridging the gap between modern foundation models and classical survival analysis structures.

Experimental Results
The method was evaluated on five standard survival analysis benchmarks (WHAS500, GBSG, METABRIC, SUPPORT, FLCHAIN) covering cardiovascular disease, oncology, and critical care. The proposed method (TabSA-BJ) was compared against classical trained models (Cox PH, Weibull AFT, Log-Normal AFT, Random Survival Forest) and a zero-shot competitor (TabSA-Bin, which uses time discretization).

Discrimination (C-Index): TabSA-BJ achieved competitive C-Index scores, often matching or outperforming classical trained models and significantly outperforming simpler zero-shot variants (TabSA-CCA, TabSA-PO). It demonstrated particular strength on datasets with higher censoring rates.
Calibration (IBS): While TabSA-BJ showed strong discrimination, its Integrated Brier Score (IBS) was generally lower than that of TabSA-Bin and classical methods. The authors attribute this to the nature of the methods: TabSA-Bin optimizes survival probabilities at fixed quantiles (better for calibration), whereas TabSA-BJ operates on the log-time scale, preserving rank structure (better for discrimination).
Convergence: Analysis showed that the scale parameter $\sigma$ converges to stable values (often near 1), suggesting that the TFM backbone captures a substantial fraction of the variance in log-time without explicit survival-specific training.

Significance and Claims
The paper claims that Tabular Foundation Models encode sufficient prior knowledge about regression structures to generalize to survival tasks with minimal adaptation. By introducing a training-free, iterative imputation strategy, the work demonstrates that TFMs can effectively handle right-censored data, offering a promising direction for zero-shot statistical modeling in domains where labeled event-time data is scarce or where rapid deployment without retraining is required. The authors position TabSA-BJ and TabSA-Bin as complementary mechanisms, with the former excelling in discrimination and the latter in calibration.