Counterfactual Understanding via Retrieval-aware Multimodal Modeling for Time-to-Event Survival Prediction

TL;DR

This paper introduces CURE, a multimodal framework for personalized time-to-event survival prediction that leverages comprehensive patient data and latent subgroup retrieval to improve accuracy over existing methods.

Contribution

CURE advances survival modeling by integrating multimodal data with cross-attention and mixture-of-experts, enabling patient-specific subgroup retrieval for better predictions.

Findings

CURE outperforms baselines on METABRIC and TCGA-LUAD datasets.

CURE achieves higher $C^{td}$ and lower IBS scores.

The framework effectively captures heterogeneity and treatment effects.

Abstract

This paper tackles the problem of time-to-event counterfactual survival prediction, aiming to optimize individualized survival outcomes in the presence of heterogeneity and censored data. We propose CURE, a framework that advances counterfactual survival modeling via comprehensive multimodal embedding and latent subgroup retrieval. CURE integrates clinical, paraclinical, demographic, and multi-omics information, which are aligned and fused through cross-attention mechanisms. Complex multi-omics signals can be adaptively refined using a mixture-of-experts architecture, emphasizing the most informative omics components. Building upon this representation, CURE implicitly retrieves patient-specific latent subgroups that capture both baseline survival dynamics and treatment-dependent variations. Experimental results on METABRIC and TCGA-LUAD datasets demonstrate that proposed CURE model consistently outperforms strong baselines in survival analysis, evaluated using the Time-dependent Concordance Index ($C^{td}$) and Integrated Brier Score (IBS). These findings highlight the potential of CURE to enhance multimodal understanding and serve as a foundation for future treatment recommendation models. All code and related resources are publicly available to facilitate the reproducibility https://github.com/L2R-UET/CURE.