Tabular Foundation Models for Clinical Survival Analysis via Survival-Aware Adaptation

Technical Summary: Tabular Foundation Models for Clinical Survival Analysis via Survival-Aware Adaptation

Problem Statement
Predicting time-to-event outcomes, such as mortality or disease progression, is a fundamental task in clinical decision-making, typically addressed through survival analysis. While classical statistical methods (e.g., Cox proportional hazards) and deep learning models (e.g., DeepSurv, DeepHit) are widely used, they often require task-specific training, extensive hyperparameter tuning, and large amounts of labeled data—resources that are frequently scarce in real-world clinical settings. Furthermore, these models must handle right-censored data, where the event of interest is not observed for all patients.

Recent advances in tabular foundation models (FMs), such as TabPFN, TabDPT, and TabICL, offer a promising alternative by learning general-purpose representations from synthetic tabular tasks, enabling strong performance via in-context learning without target dataset training. However, their applicability to clinical survival analysis remains underexplored. Existing FMs are primarily designed for discrete classification or regression and do not natively support censored time-to-event prediction. Prior attempts to adapt them often involve reformulating survival as a sequence of binary classification tasks over discretized time intervals, which introduces computational overhead, class imbalance, and inconsistencies in temporal modeling.

Methodology
The authors propose a lightweight adaptation approach to apply tabular foundation models to clinical survival analysis. The core strategy involves keeping the pretrained backbone of the foundation model fixed and training only a lightweight, survival-aware prediction head. The study evaluates two paradigms:

1. Zero-Shot In-Context Learning: The pretrained model is applied directly to a discretized survival task. For each time bin, the model predicts the probability of an event occurring using the training cohort as in-context support. The survival function is reconstructed from these probabilities, with post-processing to enforce monotonicity.
2. Survival-Aware Adaptation (Proposed): A Multi-Task Logistic Regression (MTLR) head is attached to the pretrained representations. Unlike discretized classification, MTLR models the survival distribution as a joint probability over valid survival sequences. This approach:
   • Encodes survival outcomes as binary sequences representing the transition from survival to event.
   • Defines a log-linear model over valid sequences, naturally handling censoring by marginalizing over sequences consistent with survival beyond the censoring time.
   • Enforces monotonicity by construction without requiring temporal expansion or post-hoc correction.

The authors evaluate three representative architectures (TabPFN, TabDPT, TabICL) across a diverse set of public survival benchmarks (SUPPORT, METABRIC, GBSG, WHAS, FLCHAIN, SEER, VETERANS) and two large-scale ICU cohorts (MIMIC-IV and eICU).

Key Contributions

• Adaptation Framework: The paper introduces a parameter-efficient method to extend tabular foundation models to censored time-to-event prediction by training survival-aware heads (specifically MTLR) on top of frozen pretrained representations.
• Comprehensive Evaluation: The study provides a rigorous comparison against strong baselines, including classical statistical models, tree-based methods, and deep neural survival models, across small-to-medium benchmarks and large-scale electronic health record (EHR) datasets.
• Empirical Validation: The work demonstrates that combining pretrained tabular representations with survival-aware objectives yields competitive or superior performance compared to state-of-the-art baselines, particularly in large, heterogeneous datasets.

Results
The experimental results indicate that survival-aware adaptation is the most consistent strategy across datasets:

• Performance: On the MIMIC-IV dataset, the TabDPT-FT-MTLR model achieved a C-index of 0.856, representing a +1.4% relative improvement over the best non-FM baseline (DeepSurv, 0.844) and a +6.7% improvement over the best zero-shot model (0.802). Similarly, on eICU, TabICL-FT-MTLR achieved a C-index of 0.797, outperforming DeepSurv (0.784) and the best zero-shot model (0.749).
• Zero-Shot vs. Adaptation: While zero-shot inference is competitive, particularly in smaller or cleaner datasets, it is consistently outperformed by survival-aware adaptation. The authors attribute this to the limitations of zero-shot discretized reformulations, which treat time bins independently and suffer from class imbalance.
• Representation Quality: Pretrained tabular backbones consistently outperformed standard MLP encoders when paired with the same MTLR objective, suggesting that foundation models encode useful inductive biases for tabular data that transfer effectively to survival tasks.
• Risk Stratification: In clinical utility analysis, survival-aware models produced clearer and more stable separation of low, medium, and high-risk patient groups compared to zero-shot models. The adapted models showed earlier divergence of risk groups and fewer curve crossings, indicating more reliable risk ranking over time.

Significance and Claims
The paper claims that tabular foundation models, when combined with lightweight survival-aware adaptation, provide a practical and effective alternative to traditional survival modeling approaches. The significance of this work lies in:

• Efficiency: The approach eliminates the need for extensive task-specific training of the backbone, simplifying deployment and reducing the data requirements for high-performance survival prediction.
• Robustness: The method demonstrates strong generalizability across diverse clinical domains and data regimes, from small public benchmarks to large-scale, heterogeneous ICU cohorts.
• Clinical Utility: By enabling more accurate risk stratification and earlier identification of high-risk patients, the approach supports more timely and actionable decision-making in critical care settings.

The authors conclude that while zero-shot inference serves as a strong baseline in low-resource settings, survival-aware adaptation becomes increasingly beneficial as dataset size and complexity grow, offering a principled way to handle censoring and temporal dependencies without modifying the pretrained foundation model.